''' 多卡GPU inference: # 使用所有可用GPU python generating.py \ --test_data_path /opt/tiger/PLLM1.0/dataset/processed_data/abstract_test.json \ --output_path /opt/tiger/PLLM1.0/dataset/generated_answers.json \ --base_model meta-llama/Llama-2-7b-hf \ --embedding_model all-mpnet-base-v2 \ --cluster_info_path /opt/tiger/PLLM1.0/clustering/cluster_loras/cluster_info.json \ --lora_base_path /opt/tiger/PLLM1.0/clustering/cluster_loras \ --max_length 256 \ --temperature 0.3 # 指定使用4张GPU python multi_gpu_generation.py \ --num_gpus 4 \ # ... 其他参数相同 ''' import json import torch import torch.multiprocessing as mp import numpy as np from sentence_transformers import SentenceTransformer from transformers import AutoTokenizer, AutoModelForCausalLM from peft import PeftModel from sklearn.metrics.pairwise import cosine_similarity from tqdm import tqdm import os from typing import List, Dict, Any, Tuple import argparse import gc import time import threading import queue from concurrent.futures import ThreadPoolExecutor, as_completed import fcntl import copy class MultiGPUAnswerGenerator: def __init__(self, base_model_name: str = "meta-llama/Llama-2-7b-hf", embedding_model: str = "all-mpnet-base-v2", cluster_info_path: str = "./cluster_loras/cluster_info.json", lora_base_path: str = "./cluster_loras", num_gpus: int = None): self.base_model_name = base_model_name self.embedding_model_name = embedding_model self.cluster_info_path = cluster_info_path self.lora_base_path = lora_base_path # 检测可用GPU数量 if num_gpus is None: self.num_gpus = torch.cuda.device_count() else: self.num_gpus = min(num_gpus, torch.cuda.device_count()) if self.num_gpus == 0: raise RuntimeError("没有检测到可用的GPU") print(f"🚀 使用 {self.num_gpus} 张GPU进行并行生成") for i in range(self.num_gpus): props = torch.cuda.get_device_properties(i) print(f" GPU {i}: {props.name} ({props.total_memory/1024**3:.1f}GB)") # 加载embedding模型 print(f"Loading embedding model: {embedding_model}") self.embedding_model = SentenceTransformer(embedding_model) # 加载聚类信息 self.cluster_info = self.load_cluster_info() self.cluster_centers = self.compute_cluster_centers() # 实时写入相关 self.write_lock = threading.Lock() self.results_queue = queue.Queue() self.writer_thread = None self.stop_writing = threading.Event() def load_cluster_info(self) -> Dict[str, Any]: """加载聚类信息""" print(f"Loading cluster info from {self.cluster_info_path}") with open(self.cluster_info_path, 'r', encoding='utf-8') as f: cluster_info = json.load(f) print(f"Loaded {cluster_info['n_clusters']} clusters") return cluster_info def compute_cluster_centers(self) -> np.ndarray: """计算聚类中心""" if 'cluster_centers' in self.cluster_info: centers = np.array(self.cluster_info['cluster_centers']) print(f"Using pre-computed cluster centers: {centers.shape}") return centers else: print("⚠️ No pre-computed cluster centers found") return None def extract_user_profile(self, user_data: Dict[str, Any]) -> str: """从用户数据中提取profile文本""" profile_text = "" if 'profile' in user_data and isinstance(user_data['profile'], list): profile_texts = [] for profile_item in user_data['profile']: if isinstance(profile_item, dict): if 'abstract' in profile_item: profile_texts.append(profile_item['abstract']) if 'title' in profile_item: profile_texts.append(profile_item['title']) profile_text = ' '.join(profile_texts) return profile_text def find_nearest_cluster(self, user_embedding: np.ndarray) -> int: """找到最近的聚类""" if self.cluster_centers is None: return 0 similarities = cosine_similarity([user_embedding], self.cluster_centers)[0] return np.argmax(similarities) def start_real_time_writer(self, output_path: str): """启动实时写入线程""" def writer_worker(): written_results = [] # 如果文件已存在，读取已有结果 if os.path.exists(output_path): try: with open(output_path, 'r', encoding='utf-8') as f: written_results = json.load(f) print(f"📂 恢复已有结果: {len(written_results)} 条记录") except: written_results = [] while not self.stop_writing.is_set() or not self.results_queue.empty(): try: # 等待新结果，超时1秒 new_results = self.results_queue.get(timeout=1.0) written_results.extend(new_results) # 实时写入文件 with self.write_lock: temp_path = output_path + '.tmp' with open(temp_path, 'w', encoding='utf-8') as f: json.dump(written_results, f, indent=2, ensure_ascii=False) # 原子性替换文件 os.replace(temp_path, output_path) print(f"💾 实时保存: 已写入 {len(written_results)} 条结果") except queue.Empty: continue except Exception as e: print(f"❌ 写入错误: {e}") self.writer_thread = threading.Thread(target=writer_worker, daemon=True) self.writer_thread.start() print("✅ 实时写入线程已启动") def stop_real_time_writer(self): """停止实时写入线程""" if self.writer_thread: self.stop_writing.set() self.writer_thread.join(timeout=10) print("✅ 实时写入线程已停止") def load_model_on_gpu(self, cluster_id: int, gpu_id: int) -> Tuple[Any, Any]: """在指定GPU上加载模型""" cluster_dir = os.path.join(self.lora_base_path, f"cluster_{cluster_id}") if not os.path.exists(cluster_dir): print(f"⚠️ Cluster {cluster_id} directory not found") return None, None try: print(f"🔧 GPU {gpu_id}: 加载 cluster_{cluster_id} 模型...") # 设置当前GPU torch.cuda.set_device(gpu_id) # 加载tokenizer tokenizer = AutoTokenizer.from_pretrained(cluster_dir) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" # 加载基础模型到指定GPU base_model = AutoModelForCausalLM.from_pretrained( self.base_model_name, torch_dtype=torch.float16, # 使用float16节省内存 device_map={f"": f"cuda:{gpu_id}"}, # 强制指定GPU low_cpu_mem_usage=True, max_memory={gpu_id: "70GB"} # 为H100预留一些内存 ) # 加载LoRA适配器 model = PeftModel.from_pretrained(base_model, cluster_dir) # 合并LoRA权重以提高推理速度 try: model = model.merge_and_unload() print(f"✅ GPU {gpu_id}: LoRA权重已合并") except Exception as e: print(f"⚠️ GPU {gpu_id}: LoRA权重合并失败: {e}") model.eval() print(f"✅ GPU {gpu_id}: cluster_{cluster_id} 模型加载完成") return model, tokenizer except Exception as e: print(f"❌ GPU {gpu_id}: 加载模型失败: {e}") return None, None def generate_batch_on_gpu(self, samples: List[Dict], cluster_id: int, gpu_id: int, max_length: int = 256, temperature: float = 0.3) -> List[Dict]: """在指定GPU上批量生成答案""" try: # 加载模型 model, tokenizer = self.load_model_on_gpu(cluster_id, gpu_id) if model is None or tokenizer is None: return self._create_empty_results(samples, cluster_id) results = [] batch_size = 4 # 每个GPU的batch size # 分批处理样本 for i in range(0, len(samples), batch_size): batch = samples[i:i+batch_size] batch_queries = [s['query'] for s in batch] try: # 编码输入 inputs = tokenizer( batch_queries, return_tensors="pt", padding=True, truncation=True, max_length=128 ) # 移动到指定GPU device = f"cuda:{gpu_id}" inputs = {k: v.to(device) for k, v in inputs.items()} # 生成答案 with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=max_length, temperature=temperature, do_sample=True, top_p=0.9, top_k=50, pad_token_id=tokenizer.eos_token_id, eos_token_id=tokenizer.eos_token_id, use_cache=True, repetition_penalty=1.1 ) # 解码答案 generated_texts = tokenizer.batch_decode(outputs, skip_special_tokens=True) # 构建结果 for j, sample in enumerate(batch): answer = generated_texts[j][len(batch_queries[j]):].strip() results.append({ 'user_id': sample['user_id'], 'query': sample['query'], 'profile_text': sample['profile_text'], 'nearest_cluster': cluster_id, 'generated_answer': answer, 'ground_truth': sample['ground_truth'], 'gpu_id': gpu_id # 记录使用的GPU }) # 每隔几个batch清理内存 if (i // batch_size) % 5 == 0: torch.cuda.empty_cache() except Exception as e: print(f"❌ GPU {gpu_id} batch {i}-{i+batch_size} 生成失败: {e}") # 添加空结果 for sample in batch: results.append({ 'user_id': sample['user_id'], 'query': sample['query'], 'profile_text': sample['profile_text'], 'nearest_cluster': cluster_id, 'generated_answer': "", 'ground_truth': sample['ground_truth'], 'gpu_id': gpu_id }) # 清理模型 del model, tokenizer torch.cuda.empty_cache() print(f"✅ GPU {gpu_id}: 完成 {len(results)} 个样本的生成") return results except Exception as e: print(f"❌ GPU {gpu_id}: 处理失败: {e}") return self._create_empty_results(samples, cluster_id, gpu_id) def _create_empty_results(self, samples: List[Dict], cluster_id: int, gpu_id: int = -1) -> List[Dict]: """创建空结果""" return [{ 'user_id': s['user_id'], 'query': s['query'], 'profile_text': s['profile_text'], 'nearest_cluster': cluster_id, 'generated_answer': "", 'ground_truth': s['ground_truth'], 'gpu_id': gpu_id } for s in samples] def split_samples_for_gpus(self, samples: List[Dict], cluster_id: int) -> List[Tuple]: """将样本分配给多个GPU""" if not samples: return [] # 计算每个GPU处理的样本数 samples_per_gpu = len(samples) // self.num_gpus remainder = len(samples) % self.num_gpus gpu_tasks = [] start_idx = 0 for gpu_id in range(self.num_gpus): # 为前几个GPU分配额外的样本（如果有余数） current_batch_size = samples_per_gpu + (1 if gpu_id < remainder else 0) end_idx = start_idx + current_batch_size if start_idx < len(samples): gpu_samples = samples[start_idx:end_idx] gpu_tasks.append((gpu_samples, cluster_id, gpu_id)) start_idx = end_idx return gpu_tasks def process_cluster_parallel(self, cluster_id: int, samples: List[Dict], max_length: int = 256, temperature: float = 0.3) -> List[Dict]: """并行处理单个聚类的样本""" print(f"\n🚀 并行处理 cluster_{cluster_id}: {len(samples)} 个样本") # 将样本分配给多个GPU gpu_tasks = self.split_samples_for_gpus(samples, cluster_id) if not gpu_tasks: return [] print(f"📋 任务分配:") for i, (gpu_samples, _, gpu_id) in enumerate(gpu_tasks): print(f" GPU {gpu_id}: {len(gpu_samples)} 个样本") all_results = [] # 使用线程池并行处理 with ThreadPoolExecutor(max_workers=self.num_gpus) as executor: future_to_gpu = {} # 提交任务到线程池 for gpu_samples, cluster_id, gpu_id in gpu_tasks: future = executor.submit( self.generate_batch_on_gpu, gpu_samples, cluster_id, gpu_id, max_length, temperature ) future_to_gpu[future] = gpu_id # 收集结果 for future in as_completed(future_to_gpu): gpu_id = future_to_gpu[future] try: results = future.result() all_results.extend(results) # 实时将结果放入队列 if results: self.results_queue.put(results) print(f"✅ GPU {gpu_id} 完成处理") except Exception as e: print(f"❌ GPU {gpu_id} 处理异常: {e}") print(f"✅ Cluster {cluster_id} 并行处理完成: {len(all_results)} 个结果") return all_results def convert_numpy_types(self, obj): """转换numpy类型为Python原生类型""" if isinstance(obj, np.integer): return int(obj) elif isinstance(obj, np.floating): return float(obj) elif isinstance(obj, np.ndarray): return obj.tolist() elif isinstance(obj, dict): return {key: self.convert_numpy_types(value) for key, value in obj.items()} elif isinstance(obj, list): return [self.convert_numpy_types(item) for item in obj] else: return obj def process_test_data(self, test_data_path: str, output_path: str, max_length: int = 256, temperature: float = 0.3): """处理测试数据的主函数""" print(f"📂 处理测试数据: {test_data_path}") # 启动实时写入线程 self.start_real_time_writer(output_path) try: # 加载测试数据 with open(test_data_path, 'r', encoding='utf-8') as f: test_data = json.load(f) # 加载或计算聚类分配 cluster_to_samples = self._load_cluster_assignments(test_data) # 按聚类处理 total_processed = 0 start_time = time.time() for cluster_id, samples in cluster_to_samples.items(): print(f"\n{'='*60}") print(f"处理聚类 {cluster_id}: {len(samples)} 个样本") cluster_start = time.time() # 并行处理聚类样本 cluster_results = self.process_cluster_parallel( cluster_id, samples, max_length, temperature ) total_processed += len(cluster_results) cluster_time = time.time() - cluster_start print(f"✅ Cluster {cluster_id} 完成，耗时: {cluster_time:.1f}s") print(f"📊 处理速度: {len(samples)/cluster_time:.1f} 样本/秒") print(f"📈 总进度: {total_processed}/{len(test_data)} ({total_processed/len(test_data)*100:.1f}%)") # 清理GPU内存 for gpu_id in range(self.num_gpus): torch.cuda.set_device(gpu_id) torch.cuda.empty_cache() gc.collect() # 等待写入完成 print("\n⏳ 等待所有结果写入完成...") time.sleep(2) # 给写入线程一些时间 total_time = time.time() - start_time print(f"\n🎉 全部处理完成！") print(f"⏱️ 总耗时: {total_time:.1f}s") print(f"⚡ 平均速度: {total_processed/total_time:.1f} 样本/秒") print(f"💾 结果保存到: {output_path}") except Exception as e: print(f"❌ 处理过程中出错: {e}") raise finally: # 停止实时写入线程 self.stop_real_time_writer() def _load_cluster_assignments(self, test_data): """加载或计算聚类分配""" assignment_path = "cluster_assignment.json" cluster_to_samples = {} if os.path.exists(assignment_path): print(f"📂 加载已有的聚类分配: {assignment_path}") with open(assignment_path, "r", encoding="utf-8") as f: assignments = {item["user_id"]: item["nearest_cluster"] for item in json.load(f)} for i, item in enumerate(test_data): user_id = item.get('name', f'user_{i}') query = item.get('input', '') if not query: continue profile_text = self.extract_user_profile(item) nearest_cluster = assignments.get(user_id, 0) if nearest_cluster not in cluster_to_samples: cluster_to_samples[nearest_cluster] = [] cluster_to_samples[nearest_cluster].append({ 'user_id': user_id, 'query': query, 'profile_text': profile_text, 'ground_truth': item.get('output', '') }) else: print("🔄 计算聚类分配...") assignments_list = [] for i, item in enumerate(tqdm(test_data, desc="Computing assignments")): user_id = item.get('name', f'user_{i}') query = item.get('input', '') if not query: continue profile_text = self.extract_user_profile(item) user_embedding = self.embedding_model.encode([profile_text])[0] nearest_cluster = self.find_nearest_cluster(user_embedding) assignments_list.append({"user_id": user_id, "nearest_cluster": int(nearest_cluster)}) if nearest_cluster not in cluster_to_samples: cluster_to_samples[nearest_cluster] = [] cluster_to_samples[nearest_cluster].append({ 'user_id': user_id, 'query': query, 'profile_text': profile_text, 'ground_truth': item.get('output', '') }) # 保存聚类分配结果 with open(assignment_path, "w", encoding="utf-8") as f: json.dump(assignments_list, f, ensure_ascii=False, indent=2) print(f"📊 聚类分布:") for cluster_id, samples in cluster_to_samples.items(): print(f" Cluster {cluster_id}: {len(samples)} 个样本") return cluster_to_samples def main(): # 设置多进程启动方法 mp.set_start_method('spawn', force=True) parser = argparse.ArgumentParser(description='多GPU并行聚类LoRA生成器') parser.add_argument('--test_data_path', type=str, required=True, help='测试数据文件路径') parser.add_argument('--output_path', type=str, default='./generated_answers.json', help='输出文件路径') parser.add_argument('--base_model', type=str, default='meta-llama/Llama-2-7b-hf', help='基础模型名称') parser.add_argument('--embedding_model', type=str, default='all-mpnet-base-v2', help='嵌入模型名称') parser.add_argument('--cluster_info_path', type=str, default='./cluster_loras/cluster_info.json', help='聚类信息文件路径') parser.add_argument('--lora_base_path', type=str, default='./cluster_loras', help='LoRA模型基础路径') parser.add_argument('--max_length', type=int, default=256, help='生成答案的最大长度') parser.add_argument('--temperature', type=float, default=0.3, help='生成温度') parser.add_argument('--num_gpus', type=int, default=None, help='使用的GPU数量（默认使用所有可用GPU）') args = parser.parse_args() # 创建输出目录 os.makedirs(os.path.dirname(args.output_path), exist_ok=True) try: # 初始化多GPU生成器 generator = MultiGPUAnswerGenerator( base_model_name=args.base_model, embedding_model=args.embedding_model, cluster_info_path=args.cluster_info_path, lora_base_path=args.lora_base_path, num_gpus=args.num_gpus ) # 处理测试数据 generator.process_test_data( args.test_data_path, args.output_path, max_length=args.max_length, temperature=args.temperature ) print(f"\n✅ 多GPU并行生成完成!") print(f"📁 结果文件: {args.output_path}") except Exception as e: print(f"❌ 程序执行失败: {e}") # 清理所有GPU内存 for gpu_id in range(torch.cuda.device_count()): torch.cuda.set_device(gpu_id) torch.cuda.empty_cache() raise if __name__ == "__main__": main()