#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Qwen3-Omni 智能GPU/CPU Offloading系統 功能: 使用Transformers accelerate的自動offloading，避免手動設備分配問題 策略: 讓accelerate庫自動處理設備間的數據傳輸 """ import torch import gc import time import warnings import traceback import psutil from transformers import ( Qwen3OmniMoeForConditionalGeneration, Qwen3OmniMoeProcessor, ) from accelerate import init_empty_weights, load_checkpoint_and_dispatch warnings.filterwarnings("ignore") class SmartOffloadingRunner: """智能Offloading推理運行器""" def __init__(self, model_path: str = "/var/www/qwen_model_quantized"): self.model_path = model_path self.model = None self.processor = None self.device = None self.gpu_available = torch.cuda.is_available() if self.gpu_available: self.gpu_props = torch.cuda.get_device_properties(0) self.total_gpu_memory = self.gpu_props.total_memory / 1024**3 # 設置合理的GPU記憶體限制，預留緩衝 self.max_gpu_memory = min(self.total_gpu_memory * 0.85, 24.0) # 最多24GB else: self.max_gpu_memory = 0 def get_optimal_device_map(self): """獲取最佳設備映射""" if not self.gpu_available: print("🖥️ GPU不可用，使用CPU模式") return "cpu" print(f"🔍 GPU: {self.gpu_props.name} ({self.total_gpu_memory:.1f}GB)") print(f"📊 允許GPU使用: {self.max_gpu_memory:.1f}GB") # 使用accelerate的自動offloading device_map = "auto" return device_map def load_model_with_smart_offloading(self): """使用智能offloading載入模型""" print("🚀 Qwen3-Omni 智能GPU/CPU Offloading系統") print("=" * 60) # 記憶體狀態 cpu_memory = psutil.virtual_memory().available / 1024**3 print(f"💾 可用記憶體: CPU {cpu_memory:.1f}GB", end="") if self.gpu_available: print(f", GPU {self.total_gpu_memory:.1f}GB") else: print() print("\n📦 載入processor...") self.processor = Qwen3OmniMoeProcessor.from_pretrained( self.model_path, trust_remote_code=True ) # 設置tokenizer if self.processor.tokenizer.pad_token is None: self.processor.tokenizer.pad_token = self.processor.tokenizer.eos_token print("🧠 使用智能offloading載入模型...") start_time = time.time() # 獲取設備映射 device_map = self.get_optimal_device_map() # 載入模型 try: if device_map == "cpu": # 純CPU模式 self.device = "cpu" torch.set_num_threads(min(8, psutil.cpu_count())) self.model = Qwen3OmniMoeForConditionalGeneration.from_pretrained( self.model_path, torch_dtype=torch.float32, device_map="cpu", trust_remote_code=True, low_cpu_mem_usage=True, ) # 處理meta device has_meta = any(p.device.type == 'meta' for p in self.model.parameters()) if has_meta: print("⚠️ 處理meta device權重...") self.model = self.model.to_empty(device="cpu") print("✅ meta device權重已初始化到CPU") else: # GPU+CPU offloading模式 self.device = "cuda:0" # 設置記憶體限制 max_memory = {0: f"{self.max_gpu_memory}GB", "cpu": "60GB"} self.model = Qwen3OmniMoeForConditionalGeneration.from_pretrained( self.model_path, torch_dtype=torch.float16, device_map=device_map, max_memory=max_memory, trust_remote_code=True, low_cpu_mem_usage=True, offload_folder="./offload_cache", # offload到磁碟的臨時文件夾 offload_state_dict=True, ) self.model.eval() load_time = time.time() - start_time print(f"✅ 模型載入完成! 用時: {load_time:.1f}秒") # 顯示最終記憶體使用 print("📊 記憶體使用狀態:") print(f" CPU: {psutil.virtual_memory().used / 1024**3:.1f}GB") if self.gpu_available: gpu_allocated = torch.cuda.memory_allocated() / 1024**3 print(f" GPU: {gpu_allocated:.1f}GB") # 顯示設備分配摘要 if hasattr(self.model, 'hf_device_map'): gpu_layers = sum(1 for dev in self.model.hf_device_map.values() if str(dev).startswith('cuda')) cpu_layers = sum(1 for dev in self.model.hf_device_map.values() if str(dev) == 'cpu') print(f"🎯 設備分配: GPU層數={gpu_layers}, CPU層數={cpu_layers}") return True except Exception as e: print(f"❌ 載入失敗: {e}") print("🔄 回退到CPU模式...") return self.fallback_to_cpu() def fallback_to_cpu(self): """回退到CPU模式""" try: self.device = "cpu" torch.set_num_threads(6) # 不使用device_map，避免自動分配問題 self.model = Qwen3OmniMoeForConditionalGeneration.from_pretrained( self.model_path, torch_dtype=torch.float32, trust_remote_code=True, low_cpu_mem_usage=True, ) # 處理meta device has_meta = any(p.device.type == 'meta' for p in self.model.parameters()) if has_meta: print("⚠️ CPU模式處理meta device...") self.model = self.model.to_empty(device="cpu") print("✅ CPU模式載入完成") else: # 確保模型在CPU上 self.model = self.model.to("cpu") print("✅ CPU模式載入完成") self.model.eval() return True except Exception as e: print(f"❌ CPU模式也失敗: {e}") traceback.print_exc() return False def generate_response(self, prompt: str, max_tokens: int = 128) -> tuple: """生成回應""" start_time = time.time() # 準備輸入 inputs = self.processor.tokenizer( prompt, return_tensors="pt", max_length=2048, truncation=True ) # 確定主設備 main_device = "cuda:0" if (self.gpu_available and hasattr(self.model, 'hf_device_map')) else "cpu" # 將輸入移到主設備 if main_device == "cuda:0": inputs = {k: v.to(main_device) for k, v in inputs.items()} print(f"💭 生成中... (主設備: {main_device})") # 生成 with torch.no_grad(): outputs = self.model.generate( input_ids=inputs['input_ids'], attention_mask=inputs.get('attention_mask'), max_new_tokens=max_tokens, do_sample=False, # 使用greedy解碼避免採樣問題 num_beams=1, pad_token_id=self.processor.tokenizer.eos_token_id, eos_token_id=self.processor.tokenizer.eos_token_id, ) # 解碼 response = self.processor.tokenizer.decode( outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True ).strip() # 統計 gen_time = time.time() - start_time new_tokens = outputs.shape[1] - inputs['input_ids'].shape[1] tokens_per_sec = new_tokens / gen_time if gen_time > 0 else 0 # 清理 del inputs, outputs if self.gpu_available: torch.cuda.empty_cache() gc.collect() stats = { 'generation_time': gen_time, 'new_tokens': new_tokens, 'tokens_per_second': tokens_per_sec, 'main_device': main_device } return response, stats def run_tests(self): """運行測試""" test_prompts = [ "你好，請用一句話介紹你自己。", "什麼是人工智能？", ] print("\n🧪 智能Offloading測試...") print("-" * 50) total_tokens = 0 total_time = 0 for i, prompt in enumerate(test_prompts, 1): print(f"\n📝 測試 {i}/{len(test_prompts)}: {prompt}") try: response, stats = self.generate_response(prompt, max_tokens=80) print(f"⚡ 速度: {stats['tokens_per_second']:.2f} tokens/秒") print(f"📤 回應: {response}") total_tokens += stats['new_tokens'] total_time += stats['generation_time'] except Exception as e: print(f"❌ 測試失敗: {e}") print("🔍 詳細錯誤:") traceback.print_exc() # 性能總結 if total_time > 0: avg_speed = total_tokens / total_time print(f"\n📈 Offloading性能總結:") print(f" 平均速度: {avg_speed:.2f} tokens/秒") print(f" 總tokens: {total_tokens}") print(f" 總用時: {total_time:.2f}秒") # 最終記憶體狀態 print(f" 最終CPU記憶體: {psutil.virtual_memory().used / 1024**3:.1f}GB") if self.gpu_available: print(f" 最終GPU記憶體: {torch.cuda.memory_allocated() / 1024**3:.1f}GB") def cleanup(self): """清理資源""" if self.model is not None: del self.model if self.processor is not None: del self.processor if self.gpu_available: torch.cuda.empty_cache() gc.collect() # 清理offload緩存 import shutil import os if os.path.exists("./offload_cache"): shutil.rmtree("./offload_cache") print("🧹 資源清理完成") def main(): runner = SmartOffloadingRunner() try: # 載入模型 success = runner.load_model_with_smart_offloading() if success: # 運行測試 runner.run_tests() print("\n🎉 智能Offloading測試完成!") print("💡 提示: 使用accelerate自動offloading，GPU+CPU協同工作") else: print("💥 載入失敗") except Exception as e: print(f"❌ 執行失敗: {e}") traceback.print_exc() finally: runner.cleanup() if __name__ == "__main__": main()