import os # 禁用 FlashAttention 和 FLA，解决沐曦显卡共享内存不足问题 os.environ["PYTORCH_NO_FLASH"] = "1" os.environ["FLASH_ATTENTION_ENABLED"] = "0" os.environ["USE_FLASH_ATTENTION"] = "0" os.environ["TORCH_FLASH_ATTN"] = "0" # 禁用 torch.compile，避免每个任务 fork 几十个 inductor worker os.environ["PT2_COMPILE"] = "0" os.environ["TORCHINDUCTOR_MAX_WORKERS"] = "1" # 解决 PyTorch 显存碎片化问题（避免 reserved unallocated 占用大量显存） os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # 限制训练只用 GPU 3（GPU 0/1 被 VLLM 占用，GPU 2 已占用） # CUDA_VISIBLE_DEVICES 将物理 GPU 3 映射为容器内的 cuda:0 # device_map 中使用相对编号 0（对应物理 GPU 3） os.environ["CUDA_VISIBLE_DEVICES"] = "3" # 启用 MPS 多进程服务，允许与 VLLM 共享 GPU os.environ["MACA_MPS_MODE"] = "1" import asyncio import json import logging from pathlib import Path from typing import Any # 远程训练节点没有 pydantic-settings/数据库，直接用环境变量 from types import SimpleNamespace _data_dir = Path(os.environ.get("COMPUTE_NODE_REMOTE_DATA_DIR", "/root/Fine-tuning/backend/data")) settings = SimpleNamespace( data_dir=_data_dir, processed_dir=_data_dir / "processed", adapters_dir=_data_dir / "adapters", models_dir=_data_dir / "models", ) logger = logging.getLogger(__name__) from app.engines.base import BaseEngine class TextEngine(BaseEngine): """文本模型训练引擎 (LLaMA/Qwen/ChatGLM 等因果语言模型)。""" def __init__(self): self._tokenizer = None self._model = None async def load_model(self, model_id: str, **kwargs: Any) -> None: """下载并加载基础模型。GPU 加载超时直接报错。""" import torch from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig # 远程节点不查数据库，直接扫描本地模型目录 local_path = str(settings.models_dir / model_id.replace("/", "_")) # 如果本地没有，从 HF 下载 if not (Path(local_path) / "config.json").exists(): ms_path = settings.models_dir / model_id if (ms_path / "config.json").exists(): local_path = str(ms_path) else: from huggingface_hub import snapshot_download snapshot_download( repo_id=model_id, local_dir=local_path, local_dir_use_symlinks=False, ) quantization = kwargs.get("quantization", None) gpu_timeout = int(os.environ.get("GPU_LOAD_TIMEOUT", "30")) # 记录 GPU 状态 logger.info(f"CUDA available: {torch.cuda.is_available()}") logger.info(f"CUDA device count: {torch.cuda.device_count()}") if torch.cuda.is_available(): for i in range(torch.cuda.device_count()): logger.info(f"GPU {i}: {torch.cuda.get_device_name(i)}") logger.info(f"GPU {i} memory: {torch.cuda.get_device_properties(i).total_memory / (1024**3):.2f} GB") else: raise RuntimeError("No GPU detected! Training requires GPU.") # CUDA_VISIBLE_DEVICES=3 已将物理 GPU 3 映射为逻辑 GPU 0 device_map = {"": 0} load_kwargs: dict[str, Any] = { "dtype": torch.float16, "device_map": device_map, "low_cpu_mem_usage": True, "use_safetensors": True, "attn_implementation": "sdpa", } if quantization == "4bit" or quantization == "qlora": load_kwargs["quantization_config"] = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16, ) elif quantization == "8bit": load_kwargs["quantization_config"] = BitsAndBytesConfig( load_in_8bit=True, ) self._tokenizer = AutoTokenizer.from_pretrained(local_path, trust_remote_code=True) if self._tokenizer.pad_token is None: self._tokenizer.pad_token = self._tokenizer.eos_token # GPU 加载：用超时包装，避免 MetaX 驱动无限重试卡死 model_load_result = [None] load_error = [None] def _load_on_gpu(): try: model_load_result[0] = AutoModelForCausalLM.from_pretrained(local_path, **load_kwargs) except Exception as e: load_error[0] = e load_thread = __import__("threading").Thread(target=_load_on_gpu, daemon=True) load_thread.start() load_thread.join(timeout=gpu_timeout) if load_thread.is_alive(): raise RuntimeError( f"GPU model loading timed out after {gpu_timeout}s. " f"This is usually caused by GPU resource conflict (e.g., VLLM occupying the GPU). " f"Set GPU_LOAD_TIMEOUT env var to adjust timeout." ) if load_error[0] is not None: raise RuntimeError(f"GPU model loading failed: {load_error[0]}") self._model = model_load_result[0] logger.info(f"Loaded model on GPU: {model_id}") def get_peft_config(self, method: str, params: dict[str, Any]) -> Any: """根据 PEFT 方法返回对应的配置对象。""" from app.peft import ( build_adalora_config, build_lora_config, build_qlora_config, ) builders = { "lora": build_lora_config, "qlora": build_qlora_config, "adalora": build_adalora_config, } builder = builders.get(method, build_lora_config) return builder(params) async def preprocess_dataset( self, dataset_path: str, output_path: str, task_type: str = "sft", template: str = "alpaca", **kwargs: Any, ) -> str: """将数据集预处理为训练格式。""" from app.preprocessors import preprocess_file processed = preprocess_file(dataset_path, output_path, task_type, template) logger.info(f"Preprocessed {len(processed)} samples for {task_type}/{template}") return output_path async def train( self, job_id: str, dataset_path: str, peft_config: Any, training_args: dict[str, Any], callbacks: list | None = None, ) -> str: """执行训练。""" from peft import get_peft_model from transformers import DataCollatorForSeq2Seq, TrainingArguments task_type = training_args.get("task_type", "sft") epochs = training_args.get("epochs", 3) batch_size = training_args.get("batch_size", 4) gradient_accumulation = training_args.get("gradient_accumulation", 4) learning_rate = training_args.get("learning_rate", 2e-4) max_seq_length = training_args.get("max_seq_length", 2048) warmup_ratio = training_args.get("warmup_ratio", 0.05) save_strategy = training_args.get("save_strategy", "epoch") deepspeed_config = training_args.get("deepspeed", None) dataset = self._tokenize_dataset(dataset_path, max_seq_length) self._model = get_peft_model(self._model, peft_config) self._model.print_trainable_parameters() output_dir = str(settings.adapters_dir / job_id) tr_args = TrainingArguments( output_dir=output_dir, num_train_epochs=epochs, per_device_train_batch_size=batch_size, gradient_accumulation_steps=gradient_accumulation, learning_rate=learning_rate, warmup_ratio=warmup_ratio, save_strategy=save_strategy, logging_strategy="steps", logging_steps=10, fp16=True, optim="adamw_torch", remove_unused_columns=False, report_to="none", gradient_checkpointing=True, dataloader_num_workers=4, dataloader_pin_memory=False, **({"deepspeed": deepspeed_config} if deepspeed_config else {}), ) # 本地模式用 WebSocket 回调，远程模式用传入的文件日志回调 all_callbacks = callbacks if callbacks else [_ProgressCallback(job_id)] if task_type == "sft": from transformers import Trainer trainer = Trainer( model=self._model, args=tr_args, train_dataset=dataset, data_collator=DataCollatorForSeq2Seq(self._tokenizer), callbacks=all_callbacks, ) elif task_type == "dpo": from trl import DPOConfig, DPOTrainer base_trainer_kwargs = dict( output_dir=output_dir, num_train_epochs=epochs, per_device_train_batch_size=batch_size, gradient_accumulation_steps=gradient_accumulation, learning_rate=learning_rate, warmup_ratio=warmup_ratio, save_strategy=save_strategy, logging_steps=10, fp16=True, report_to="none", dataloader_num_workers=4, dataloader_pin_memory=False, ) trainer = DPOTrainer( model=self._model, args=DPOConfig(**base_trainer_kwargs), train_dataset=dataset, processing_class=self._tokenizer, ) elif task_type == "ppo": from transformers import Trainer logger.warning( "PPO mode: falling back to SFT Trainer. " "PPO requires a dedicated reward model setup. " "Current implementation trains as supervised fine-tuning." ) trainer = Trainer( model=self._model, args=tr_args, train_dataset=dataset, data_collator=DataCollatorForSeq2Seq(self._tokenizer), callbacks=all_callbacks, ) else: from transformers import Trainer raise ValueError(f"Unsupported task_type: {task_type}. Supported: sft, dpo, ppo") try: trainer.train() self._model.save_pretrained(output_dir) self._tokenizer.save_pretrained(output_dir) logger.info(f"Training completed for job {job_id}") except Exception as e: logger.error(f"Training failed for job {job_id}: {e}") raise return output_dir def get_model_info(self, model_id: str) -> dict[str, Any]: """读取模型配置信息。""" import json from pathlib import Path # 同步查找模型路径（兼容 HF 和 ModelScope） candidates = [ settings.models_dir / model_id.replace("/", "_"), settings.models_dir / model_id, ] config_path = None for p in candidates: if (p / "config.json").exists(): config_path = p / "config.json" break if not config_path: # 最后尝试扫描 model_name = model_id.split("/")[-1] for cp in settings.models_dir.rglob("config.json"): if model_name in str(cp.parent): config_path = cp break if config_path.exists(): with open(config_path) as f: config = json.load(f) return { "model_type": config.get("model_type", "causal_lm"), "context_length": config.get("max_position_embeddings", config.get("max_sequence_length", 2048)), "hidden_size": config.get("hidden_size", 0), "num_layers": config.get("num_hidden_layers", 0), } return {"model_type": "causal_lm", "context_length": 2048} def _tokenize_dataset(self, dataset_path: str, max_seq_length: int): """Tokenize 处理后的 JSONL 数据集。""" from datasets import Dataset as HFDataset data = [] with open(dataset_path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if line: item = json.loads(line) # 兼容多种列名 → 统一映射为 prompt / completion if "prompt" not in item: item["prompt"] = item.get("question", item.get("query", item.get("text", item.get("input", "")))) if "completion" not in item: item["completion"] = item.get("answer", item.get("response", item.get("target", item.get("output", "")))) # 确保 prompt 和 completion 是字符串 if isinstance(item["prompt"], (list, dict)): item["prompt"] = json.dumps(item["prompt"], ensure_ascii=False) item["prompt"] = str(item["prompt"]) if isinstance(item["completion"], (list, dict)): item["completion"] = json.dumps(item["completion"], ensure_ascii=False) item["completion"] = str(item["completion"]) data.append(item) hf_dataset = HFDataset.from_list(data) def tokenize_fn(batch): def _to_str(v): if isinstance(v, (list, dict)): return json.dumps(v, ensure_ascii=False) return str(v) if v is not None else "" raw_prompts = batch.get("prompt", []) raw_completions = batch.get("completion", []) prompts = [_to_str(v) for v in raw_prompts] completions = [_to_str(v) for v in raw_completions] if not prompts: return {"input_ids": [], "attention_mask": [], "labels": []} full_texts = [f"{p}\n{c}" for p, c in zip(prompts, completions)] tokenized = self._tokenizer( full_texts, truncation=True, max_length=max_seq_length, padding=False, ) tokenized["labels"] = list(tokenized["input_ids"]) return tokenized tokenized_dataset = hf_dataset.map( tokenize_fn, batched=True, remove_columns=["prompt", "completion"], ) return tokenized_dataset class _ProgressCallback: """自定义训练进度回调，通过 WebSocket 发送进度。""" def __init__(self, job_id: str): self.job_id = job_id def on_log(self, args, state, control, logs=None, **kwargs): if logs and "loss" in logs: asyncio.create_task( send_progress( self.job_id, epoch=int(state.epoch or 0), step=state.global_step, total_steps=state.max_steps or 0, loss=logs["loss"], learning_rate=logs.get("learning_rate", 0), ) ) def on_epoch_end(self, args, state, control, **kwargs): asyncio.create_task( send_epoch_done(self.job_id, epoch=int(state.epoch or 0), eval_loss=None, eval_accuracy=None) ) def on_train_end(self, args, state, control, **kwargs): asyncio.create_task( send_completed( self.job_id, total_time_seconds=getattr(state, "train_runtime", 0), adapter_path=str(settings.adapters_dir / self.job_id), ) ) def on_train_begin(self, args, state, control, **kwargs): pass def on_step_begin(self, args, state, control, **kwargs): pass def on_step_end(self, args, state, control, **kwargs): pass def on_evaluate(self, args, state, control, metrics=None, **kwargs): pass def on_save(self, args, state, control, **kwargs): pass def on_predict(self, args, state, control, metrics=None, **kwargs): pass def on_init_end(self, args, state, control, **kwargs): pass def on_epoch_begin(self, args, state, control, **kwargs): pass # 全局单例 text_engine = TextEngine()