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"
# 解决 PyTorch 显存碎片化问题（避免 reserved unallocated 占用大量显存）
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
# CUDA_VISIBLE_DEVICES 由 docker exec 层设置（remote_executor.py），此处不再覆盖
# 单 GPU 模式: "3"  (物理 GPU 3 → 逻辑 cuda:0)
# 多 GPU 模式: "2,3" (物理 GPU 2,3 → 逻辑 cuda:0,1)
# 启用 MPS 多进程服务，允许与 VLLM 共享 GPU
os.environ["MACA_MPS_MODE"] = "1"

import asyncio
import json
import logging
from pathlib import Path
from typing import Any

from types import SimpleNamespace

# 确定数据目录：优先用 DATA_DIR 环境变量，否则从 .env 文件读取，最后兜底
def _resolve_data_dir() -> Path:
    v = os.environ.get("DATA_DIR") or os.environ.get("COMPUTE_NODE_REMOTE_DATA_DIR")
    if v:
        return Path(v)
    # 从 .env 文件读取 DATA_DIR（pydantic-settings 加载 .env 但不导出到 os.environ）
    env_file = Path(__file__).resolve().parent.parent.parent / ".env"
    if env_file.exists():
        for line in env_file.read_text():
            if line.strip().startswith("DATA_DIR="):
                return Path(line.split("=", 1)[1].strip())
    return Path("/root/Fine-tuning/backend/data")

_data_dir = _resolve_data_dir()
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.")

        # DDP 模式: LOCAL_RANK 由 torchrun 设置；单 GPU 模式默认为 0
        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
        device_map = {"": local_rank}

        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":
            # 沐曦 GPU 不支持 bitsandbytes/HQQ，直接 fp16 + LoRA
            load_kwargs["torch_dtype"] = torch.float16
            logger.info("4-bit quantization not supported on this GPU; "
                        "falling back to fp16 + LoRA")
        elif quantization == "8bit":
            # 沐曦 GPU 不支持 bitsandbytes，直接 fp16 + LoRA
            load_kwargs["torch_dtype"] = torch.float16
            logger.info("8-bit quantization not supported on this GPU; "
                        "falling back to fp16 + LoRA")

        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)

        # DDP 支持
        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
        world_size = int(os.environ.get("WORLD_SIZE", "1"))
        is_ddp = world_size > 1

        # SFT 需要预先 tokenize；DPO/PPO 各自处理数据
        if task_type == "sft":
            dataset = self._tokenize_dataset(dataset_path, max_seq_length)
        elif task_type == "dpo":
            dataset = self._load_dataset_dpo(dataset_path)
        else:
            dataset = None  # PPO 在后面单独处理

        # 计算总步数（DDP 模式下 Trainer 自动按 world_size 分发数据）
        if dataset is not None:
            dataset_len = len(dataset)
        else:
            # PPO: 从文件行数估算
            with open(dataset_path, "r", encoding="utf-8") as f:
                dataset_len = sum(1 for line in f if line.strip())
        effective_batch = batch_size * gradient_accumulation * world_size
        max_steps = max(1, (dataset_len * epochs) // effective_batch)

        # AdaLoRA 要求 total_step > 0（通过属性名判断而非 isinstance，避免导入路径问题）
        if hasattr(peft_config, "init_r") and hasattr(peft_config, "target_r"):
            peft_config.total_step = max_steps

        # PPO 需要先用 AutoModelForCausalLMWithValueHead 包装，再应用 PEFT（后面单独处理）
        if task_type != "ppo":
            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,
            max_steps=max_steps,
            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,
            local_rank=local_rank if is_ddp else -1,
            ddp_find_unused_parameters=False if is_ddp else None,
            **({"deepspeed": deepspeed_config} if deepspeed_config else {}),
        )

        # 本地模式用 WebSocket 回调，远程模式用传入的文件日志回调
        # 用 is None 判断而非 falsy，因为 DDP 非 rank 0 传入空列表 []，不需要进度回调
        all_callbacks = callbacks if callbacks is not None 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 copy import deepcopy

            from trl import DPOConfig, DPOTrainer

            # 显式创建 reference model 并冻结，避免 AdaLora 多 adapter 冲突
            ref_model = deepcopy(self._model)
            ref_model.eval()
            for param in ref_model.parameters():
                param.requires_grad = False
            # 将 ref_model 上的 PEFT adapter 设为推理模式
            # AdaLora 只允许 1 个可训练 adapter，policy model 已有 1 个
            if hasattr(ref_model, "set_adapter"):
                try:
                    ref_model.set_adapter("default", inference_mode=True)
                except Exception:
                    pass
            elif hasattr(ref_model, "peft_config"):
                for adapter_name in list(ref_model.peft_config.keys()):
                    try:
                        ref_model.peft_config[adapter_name].inference_mode = True
                    except Exception:
                        pass

            base_trainer_kwargs = dict(
                output_dir=output_dir,
                num_train_epochs=epochs,
                max_steps=max_steps,
                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,
                ref_model=ref_model,
                args=DPOConfig(**base_trainer_kwargs),
                train_dataset=dataset,
                processing_class=self._tokenizer,
            )
        elif task_type == "ppo":
            import torch
            from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer

            ppo_epochs = training_args.get("ppo_epochs", 4)
            vf_coef = training_args.get("vf_coef", 0.1)
            kl_coef = training_args.get("kl_coef", 0.2)
            response_length = training_args.get("response_length", 512)
            reward_model_path = training_args.get("reward_model_path")
            reward_type = training_args.get("reward_type", "heuristic")

            # PPO 专用：仅 tokenize prompt
            ppo_dataset = self._tokenize_dataset_ppo(dataset_path, max_seq_length, response_length)

            # PPO 需要 AutoModelForCausalLMWithValueHead（添加 value head 用于评估动作价值）
            # 通过 peft_config 参数让 TRL 内部处理 PEFT 包装，返回的对象是 PreTrainedModelWrapper
            # 不能用 get_peft_model（会产生 PeftModel，PPOTrainer 不认）
            self._model = AutoModelForCausalLMWithValueHead.from_pretrained(
                self._model, peft_config=peft_config,
            )
            if hasattr(self._model, "print_trainable_parameters"):
                self._model.print_trainable_parameters()

            # TRL 0.9.x PPOConfig 只接受 PPO 专用参数，不支持 HuggingFace Trainer 参数
            # mini_batch_size 必须满足：batch_size % (mini_batch_size * gradient_accumulation_steps) == 0
            ppo_config = PPOConfig(
                learning_rate=learning_rate,
                batch_size=batch_size,
                mini_batch_size=1,
                gradient_accumulation_steps=gradient_accumulation,
                ppo_epochs=ppo_epochs,
                vf_coef=vf_coef,
                init_kl_coef=kl_coef,
            )

            # ref_model=None 让 PPOTrainer 自动创建冻结的 reference model（用于 KL 惩罚）
            trainer = PPOTrainer(
                config=ppo_config,
                model=self._model,
                ref_model=None,
                tokenizer=self._tokenizer,
                dataset=ppo_dataset,
            )

            dataloader = trainer.dataloader
            total_steps = len(dataloader) * epochs
            step_count = 0

            for epoch in range(epochs):
                for batch in dataloader:
                    step_count += 1
                    query_tensors = batch["input_ids"]

                    # 生成回答
                    response_tensors = []
                    for query in query_tensors:
                        query_tensor = torch.tensor(query).unsqueeze(0).to(self._model.device)
                        gen_output = self._model.generate(
                            query_tensor,
                            max_new_tokens=response_length,
                            do_sample=True,
                            top_p=0.9,
                            temperature=0.7,
                        )
                        response_tensors.append(gen_output[0][query_tensor.shape[-1]:])

                    # 解码文本用于奖励计算
                    responses_text = [
                        self._tokenizer.decode(r, skip_special_tokens=True)
                        for r in response_tensors
                    ]
                    prompts_text = [
                        self._tokenizer.decode(q, skip_special_tokens=True)
                        for q in query_tensors
                    ]

                    # 计算奖励
                    if reward_type == "model" and reward_model_path:
                        from transformers import AutoModelForSequenceClassification

                        reward_model = AutoModelForSequenceClassification.from_pretrained(
                            reward_model_path, device_map={"": 0}
                        )
                        reward_inputs = [p + r for p, r in zip(prompts_text, responses_text)]
                        tokenized = self._tokenizer(
                            reward_inputs, return_tensors="pt", padding=True, truncation=True
                        ).to(self._model.device)
                        with torch.no_grad():
                            rewards = reward_model(**tokenized).logits.squeeze(-1).tolist()
                    else:
                        rewards = _compute_heuristic_reward(prompts_text, responses_text)

                    reward_tensors = [torch.tensor(r, device=self._model.device) for r in rewards]

                    # PPO 更新
                    stats = trainer.step(query_tensors, response_tensors, reward_tensors)

                    # 报告进度
                    if step_count % 10 == 0:
                        for cb in (all_callbacks or []):
                            if hasattr(cb, "on_log"):
                                cb.on_log(
                                    SimpleNamespace(),
                                    SimpleNamespace(
                                        epoch=epoch, global_step=step_count, max_steps=total_steps
                                    ),
                                    None,
                                    logs={
                                        "loss": stats.get("ppo/loss/total", 0),
                                        "learning_rate": stats.get("ppo/learning_rate", learning_rate),
                                    },
                                )

            os.makedirs(output_dir, exist_ok=True)
            self._model.save_pretrained(output_dir)
            self._tokenizer.save_pretrained(output_dir)
            logger.info(f"PPO training completed for job {job_id}")
            return output_dir
        else:
            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_ppo(self, dataset_path: str, max_seq_length: int, response_length: int):
        """Tokenize PPO 数据集：仅 prompt（模型在训练中自己生成回答）。"""
        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)
                    if "prompt" not in item:
                        item["prompt"] = item.get("question", item.get("query", item.get("text", item.get("input", ""))))
                    if isinstance(item["prompt"], (list, dict)):
                        item["prompt"] = json.dumps(item["prompt"], ensure_ascii=False)
                    item["prompt"] = str(item["prompt"])
                    data.append(item)

        hf_dataset = HFDataset.from_list(data)

        def tokenize_fn(batch):
            raw_prompts = batch.get("prompt", [])
            prompts = [str(v) if v is not None else "" for v in raw_prompts]
            # 仅 tokenize prompt，预留 response_length 空间给生成的回答
            tokenized = self._tokenizer(
                prompts,
                truncation=True,
                max_length=max_seq_length - response_length,
                padding=False,
            )
            return tokenized

        tokenized_dataset = hf_dataset.map(
            tokenize_fn,
            batched=True,
            remove_columns=hf_dataset.column_names,
        )
        return tokenized_dataset

    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

    def _load_dataset_dpo(self, dataset_path: str):
        """加载 DPO 数据集，保留 prompt/chosen/rejected 原始文本，由 DPOTrainer 内部 tokenize。"""
        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 = item.get("prompt", item.get("instruction", item.get("input", "")))
                    chosen = item.get("chosen", item.get("positive", ""))
                    rejected = item.get("rejected", item.get("negative", ""))
                    data.append({
                        "prompt": str(prompt) if prompt else "",
                        "chosen": str(chosen) if chosen else "",
                        "rejected": str(rejected) if rejected else "",
                    })
        return HFDataset.from_list(data)


try:
    from transformers import TrainerCallback as _TrainerCallbackBase
except ImportError:
    _TrainerCallbackBase = object  # 151 主节点无 transformers，仅做占位


class _ProgressCallback(_TrainerCallbackBase):
    """自定义训练进度回调，通过 WebSocket 发送进度。"""

    def __init__(self, job_id: str):
        super().__init__()
        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),
            )
        )


# 全局单例
text_engine = TextEngine()


def _compute_heuristic_reward(prompts: list[str], responses: list[str]) -> list[float]:
    """启发式奖励函数：无需额外奖励模型即可用于 PPO 训练。

    评分维度：长度合理性 + 非空 + 重复度惩罚。
    """
    rewards = []
    for _prompt, response in zip(prompts, responses):
        reward = 0.0
        resp_len = len(response.split())
        # 长度评分：20-200 词为佳
        if 20 <= resp_len <= 200:
            reward += 0.5
        elif resp_len < 5:
            reward -= 1.0
        elif resp_len > 500:
            reward -= 0.5
        # 非空奖励
        if response.strip():
            reward += 0.2
        # 重复度惩罚（trigram 重复率过高）
        words = response.split()
        if len(words) > 10:
            trigrams = set(tuple(words[i:i+3]) for i in range(len(words) - 2))
            if len(trigrams) < len(words) * 0.3:
                reward -= 0.5
        rewards.append(reward)
    return rewards