import os # 禁用 FlashAttention，解决沐曦显卡共享内存不足问题 os.environ["PYTORCH_NO_FLASH"] = "1" os.environ["FLASH_ATTENTION_ENABLED"] = "0" os.environ["USE_FLASH_ATTENTION"] = "0" os.environ["TORCH_FLASH_ATTN"] = "0" import asyncio import json from pathlib import Path from typing import Any from app.config import get_settings from app.core.logging import logger from app.engines.base import BaseEngine settings = get_settings() 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: """下载并加载基础模型。""" import torch from transformers import AutoModelForCausalLM, AutoTokenizer # 优先从数据库获取实际路径（兼容 ModelScope 下载的目录结构） from app.services.model_service import resolve_model_path model_path = await resolve_model_path(model_id) if model_path: local_path = model_path else: local_path = str(settings.models_dir / model_id.replace("/", "_")) # 如果本地没有，从 HF 下载 if not (Path(local_path) / "config.json").exists(): 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 状态 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: logger.warning("No GPU detected! Training will run on CPU.") max_memory = {i: "4GB" for i in range(torch.cuda.device_count())} if torch.cuda.is_available() else None load_kwargs: dict[str, Any] = { "torch_dtype": torch.float16, "device_map": "auto", "low_cpu_mem_usage": True, "use_safetensors": True, "max_memory": max_memory, } if quantization == "4bit" or quantization == "qlora": load_kwargs["load_in_4bit"] = True load_kwargs["bnb_4bit_quant_type"] = "nf4" load_kwargs["bnb_4bit_use_double_quant"] = True load_kwargs["bnb_4bit_compute_dtype"] = torch.float16 elif quantization == "8bit": load_kwargs["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 self._model = AutoModelForCausalLM.from_pretrained(local_path, **load_kwargs) logger.info(f"Loaded model: {model_id}") def get_peft_config(self, method: str, params: dict[str, Any]) -> Any: """根据 PEFT 方法返回对应的配置对象。""" from app.peft import ( build_adalora_config, build_ia3_config, build_lora_config, build_prefix_tuning_config, build_qlora_config, ) builders = { "lora": build_lora_config, "qlora": build_qlora_config, "ia3": build_ia3_config, "adalora": build_adalora_config, "prefix_tuning": build_prefix_tuning_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], ) -> 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=False, dataloader_num_workers=1, dataloader_pin_memory=False, **({"deepspeed": deepspeed_config} if deepspeed_config else {}), ) callback = _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=[callback], ) else: from trl import ( DPOConfig, DPOTrainer, KTOConfig, KTOTrainer, ORPOConfig, ORPOTrainer, ) 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", ) if task_type == "dpo": trainer = DPOTrainer( model=self._model, args=DPOConfig(**base_trainer_kwargs), train_dataset=dataset, processing_class=self._tokenizer, ) elif task_type == "orpo": trainer = ORPOTrainer( model=self._model, args=ORPOConfig(**base_trainer_kwargs), train_dataset=dataset, processing_class=self._tokenizer, ) elif task_type == "kto": trainer = KTOTrainer( model=self._model, args=KTOConfig(**base_trainer_kwargs), train_dataset=dataset, processing_class=self._tokenizer, ) else: trainer = Trainer( model=self._model, args=tr_args, train_dataset=dataset, data_collator=DataCollatorForSeq2Seq(self._tokenizer), callbacks=[callback], ) 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" in item: if isinstance(item["prompt"], (list, dict)): item["prompt"] = json.dumps(item["prompt"], ensure_ascii=False) item["prompt"] = str(item["prompt"]) if "completion" in item: 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()