Fine-tuning/backend/app/engines/text_engine.py

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"
# 限制训练只用 GPU 2 和 3（GPU 0/1 被 VLLM 占用）
# 沐曦 GPU 优先用 METAX_VISIBLE_DEVICES，同时设 CUDA_VISIBLE_DEVICES 兜底
os.environ["METAX_VISIBLE_DEVICES"] = "2,3"
os.environ["CUDA_VISIBLE_DEVICES"] = "2,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

        # 远程节点不查数据库，直接扫描本地模型目录
        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.")

        # 沐曦 MPS 模式下 device_map="auto" 会导致 tensor 分散到不同设备，
        # 引发跨 GPU 计算错误。强制用 "cuda:0"（对应 METAX_VISIBLE_DEVICES 的第一个）。
        visible_devices = os.environ.get("METAX_VISIBLE_DEVICES", "0")
        device_map = {
            "": int(visible_devices.split(",")[0]) if visible_devices else 0
        }

        load_kwargs: dict[str, Any] = {
            "torch_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["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

        # 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_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],
        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=False,
            dataloader_num_workers=0,
            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,
            )
        else:
            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",
            )

            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":
                from trl import ORPOConfig, ORPOTrainer

                trainer = ORPOTrainer(
                    model=self._model,
                    args=ORPOConfig(**base_trainer_kwargs),
                    train_dataset=dataset,
                    processing_class=self._tokenizer,
                )
            elif task_type == "kto":
                from trl import KTOConfig, KTOTrainer

                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=all_callbacks,
                )

        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()