text_engine.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684
  1. import os
  2. # 禁用 FlashAttention 和 FLA,解决沐曦显卡共享内存不足问题
  3. os.environ["PYTORCH_NO_FLASH"] = "1"
  4. os.environ["FLASH_ATTENTION_ENABLED"] = "0"
  5. os.environ["USE_FLASH_ATTENTION"] = "0"
  6. os.environ["TORCH_FLASH_ATTN"] = "0"
  7. # 禁用 torch.compile,避免每个任务 fork 几十个 inductor worker
  8. os.environ["PT2_COMPILE"] = "0"
  9. os.environ["TORCHINDUCTOR_MAX_WORKERS"] = "1"
  10. # 解决 PyTorch 显存碎片化问题(避免 reserved unallocated 占用大量显存)
  11. os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
  12. # CUDA_VISIBLE_DEVICES 由 docker exec 层设置(remote_executor.py),此处不再覆盖
  13. # 单 GPU 模式: "3" (物理 GPU 3 → 逻辑 cuda:0)
  14. # 多 GPU 模式: "2,3" (物理 GPU 2,3 → 逻辑 cuda:0,1)
  15. # 启用 MPS 多进程服务,允许与 VLLM 共享 GPU
  16. os.environ["MACA_MPS_MODE"] = "1"
  17. import asyncio
  18. import json
  19. import logging
  20. from pathlib import Path
  21. from typing import Any
  22. from types import SimpleNamespace
  23. # 确定数据目录:优先用 DATA_DIR 环境变量,否则从 .env 文件读取,最后兜底
  24. def _resolve_data_dir() -> Path:
  25. v = os.environ.get("DATA_DIR") or os.environ.get("COMPUTE_NODE_REMOTE_DATA_DIR")
  26. if v:
  27. return Path(v)
  28. # 从 .env 文件读取 DATA_DIR(pydantic-settings 加载 .env 但不导出到 os.environ)
  29. env_file = Path(__file__).resolve().parent.parent.parent / ".env"
  30. if env_file.exists():
  31. for line in env_file.read_text():
  32. if line.strip().startswith("DATA_DIR="):
  33. return Path(line.split("=", 1)[1].strip())
  34. return Path("/root/Fine-tuning/backend/data")
  35. _data_dir = _resolve_data_dir()
  36. settings = SimpleNamespace(
  37. data_dir=_data_dir,
  38. processed_dir=_data_dir / "processed",
  39. adapters_dir=_data_dir / "adapters",
  40. models_dir=_data_dir / "models",
  41. )
  42. logger = logging.getLogger(__name__)
  43. from app.engines.base import BaseEngine
  44. class TextEngine(BaseEngine):
  45. """文本模型训练引擎 (LLaMA/Qwen/ChatGLM 等因果语言模型)。"""
  46. def __init__(self):
  47. self._tokenizer = None
  48. self._model = None
  49. async def load_model(self, model_id: str, **kwargs: Any) -> None:
  50. """下载并加载基础模型。GPU 加载超时直接报错。"""
  51. import torch
  52. from transformers import AutoModelForCausalLM, AutoTokenizer
  53. # 远程节点不查数据库,直接扫描本地模型目录
  54. local_path = str(settings.models_dir / model_id.replace("/", "_"))
  55. # 如果本地没有,从 HF 下载
  56. if not (Path(local_path) / "config.json").exists():
  57. ms_path = settings.models_dir / model_id
  58. if (ms_path / "config.json").exists():
  59. local_path = str(ms_path)
  60. else:
  61. from huggingface_hub import snapshot_download
  62. snapshot_download(
  63. repo_id=model_id,
  64. local_dir=local_path,
  65. local_dir_use_symlinks=False,
  66. )
  67. quantization = kwargs.get("quantization", None)
  68. gpu_timeout = int(os.environ.get("GPU_LOAD_TIMEOUT", "30"))
  69. # 记录 GPU 状态
  70. logger.info(f"CUDA available: {torch.cuda.is_available()}")
  71. logger.info(f"CUDA device count: {torch.cuda.device_count()}")
  72. if torch.cuda.is_available():
  73. for i in range(torch.cuda.device_count()):
  74. logger.info(f"GPU {i}: {torch.cuda.get_device_name(i)}")
  75. logger.info(f"GPU {i} memory: {torch.cuda.get_device_properties(i).total_memory / (1024**3):.2f} GB")
  76. else:
  77. raise RuntimeError("No GPU detected! Training requires GPU.")
  78. # DDP 模式: LOCAL_RANK 由 torchrun 设置;单 GPU 模式默认为 0
  79. local_rank = int(os.environ.get("LOCAL_RANK", "0"))
  80. device_map = {"": local_rank}
  81. load_kwargs: dict[str, Any] = {
  82. "dtype": torch.float16,
  83. "device_map": device_map,
  84. "low_cpu_mem_usage": True,
  85. "use_safetensors": True,
  86. "attn_implementation": "sdpa",
  87. }
  88. if quantization == "4bit" or quantization == "qlora":
  89. # 沐曦 GPU 不支持 bitsandbytes/HQQ,直接 fp16 + LoRA
  90. load_kwargs["torch_dtype"] = torch.float16
  91. logger.info("4-bit quantization not supported on this GPU; "
  92. "falling back to fp16 + LoRA")
  93. elif quantization == "8bit":
  94. # 沐曦 GPU 不支持 bitsandbytes,直接 fp16 + LoRA
  95. load_kwargs["torch_dtype"] = torch.float16
  96. logger.info("8-bit quantization not supported on this GPU; "
  97. "falling back to fp16 + LoRA")
  98. self._tokenizer = AutoTokenizer.from_pretrained(local_path, trust_remote_code=True)
  99. if self._tokenizer.pad_token is None:
  100. self._tokenizer.pad_token = self._tokenizer.eos_token
  101. # GPU 加载:用超时包装,避免 MetaX 驱动无限重试卡死
  102. model_load_result = [None]
  103. load_error = [None]
  104. def _load_on_gpu():
  105. try:
  106. model_load_result[0] = AutoModelForCausalLM.from_pretrained(local_path, **load_kwargs)
  107. except Exception as e:
  108. load_error[0] = e
  109. load_thread = __import__("threading").Thread(target=_load_on_gpu, daemon=True)
  110. load_thread.start()
  111. load_thread.join(timeout=gpu_timeout)
  112. if load_thread.is_alive():
  113. raise RuntimeError(
  114. f"GPU model loading timed out after {gpu_timeout}s. "
  115. f"This is usually caused by GPU resource conflict (e.g., VLLM occupying the GPU). "
  116. f"Set GPU_LOAD_TIMEOUT env var to adjust timeout."
  117. )
  118. if load_error[0] is not None:
  119. raise RuntimeError(f"GPU model loading failed: {load_error[0]}")
  120. self._model = model_load_result[0]
  121. logger.info(f"Loaded model on GPU: {model_id}")
  122. def get_peft_config(self, method: str, params: dict[str, Any]) -> Any:
  123. """根据 PEFT 方法返回对应的配置对象。"""
  124. from app.peft import (
  125. build_adalora_config,
  126. build_lora_config,
  127. build_qlora_config,
  128. )
  129. builders = {
  130. "lora": build_lora_config,
  131. "qlora": build_qlora_config,
  132. "adalora": build_adalora_config,
  133. }
  134. builder = builders.get(method, build_lora_config)
  135. return builder(params)
  136. async def preprocess_dataset(
  137. self,
  138. dataset_path: str,
  139. output_path: str,
  140. task_type: str = "sft",
  141. template: str = "alpaca",
  142. **kwargs: Any,
  143. ) -> str:
  144. """将数据集预处理为训练格式。"""
  145. from app.preprocessors import preprocess_file
  146. processed = preprocess_file(dataset_path, output_path, task_type, template)
  147. logger.info(f"Preprocessed {len(processed)} samples for {task_type}/{template}")
  148. return output_path
  149. async def train(
  150. self,
  151. job_id: str,
  152. dataset_path: str,
  153. peft_config: Any,
  154. training_args: dict[str, Any],
  155. callbacks: list | None = None,
  156. ) -> str:
  157. """执行训练。"""
  158. from peft import get_peft_model
  159. from transformers import DataCollatorForSeq2Seq, TrainingArguments
  160. task_type = training_args.get("task_type", "sft")
  161. epochs = training_args.get("epochs", 3)
  162. batch_size = training_args.get("batch_size", 4)
  163. gradient_accumulation = training_args.get("gradient_accumulation", 4)
  164. learning_rate = training_args.get("learning_rate", 2e-4)
  165. max_seq_length = training_args.get("max_seq_length", 2048)
  166. warmup_ratio = training_args.get("warmup_ratio", 0.05)
  167. save_strategy = training_args.get("save_strategy", "epoch")
  168. deepspeed_config = training_args.get("deepspeed", None)
  169. # DDP 支持
  170. local_rank = int(os.environ.get("LOCAL_RANK", "0"))
  171. world_size = int(os.environ.get("WORLD_SIZE", "1"))
  172. is_ddp = world_size > 1
  173. # SFT 需要预先 tokenize;DPO/PPO 各自处理数据
  174. if task_type == "sft":
  175. dataset = self._tokenize_dataset(dataset_path, max_seq_length)
  176. elif task_type == "dpo":
  177. dataset = self._load_dataset_dpo(dataset_path)
  178. else:
  179. dataset = None # PPO 在后面单独处理
  180. # 计算总步数(DDP 模式下 Trainer 自动按 world_size 分发数据)
  181. if dataset is not None:
  182. dataset_len = len(dataset)
  183. else:
  184. # PPO: 从文件行数估算
  185. with open(dataset_path, "r", encoding="utf-8") as f:
  186. dataset_len = sum(1 for line in f if line.strip())
  187. effective_batch = batch_size * gradient_accumulation * world_size
  188. max_steps = max(1, (dataset_len * epochs) // effective_batch)
  189. # AdaLoRA 要求 total_step > 0(通过属性名判断而非 isinstance,避免导入路径问题)
  190. if hasattr(peft_config, "init_r") and hasattr(peft_config, "target_r"):
  191. peft_config.total_step = max_steps
  192. # PPO 需要先用 AutoModelForCausalLMWithValueHead 包装,再应用 PEFT(后面单独处理)
  193. if task_type != "ppo":
  194. self._model = get_peft_model(self._model, peft_config)
  195. self._model.print_trainable_parameters()
  196. output_dir = str(settings.adapters_dir / job_id)
  197. tr_args = TrainingArguments(
  198. output_dir=output_dir,
  199. num_train_epochs=epochs,
  200. max_steps=max_steps,
  201. per_device_train_batch_size=batch_size,
  202. gradient_accumulation_steps=gradient_accumulation,
  203. learning_rate=learning_rate,
  204. warmup_ratio=warmup_ratio,
  205. save_strategy=save_strategy,
  206. logging_strategy="steps",
  207. logging_steps=10,
  208. fp16=True,
  209. optim="adamw_torch",
  210. remove_unused_columns=False,
  211. report_to="none",
  212. gradient_checkpointing=True,
  213. dataloader_num_workers=4,
  214. dataloader_pin_memory=False,
  215. local_rank=local_rank if is_ddp else -1,
  216. ddp_find_unused_parameters=False if is_ddp else None,
  217. **({"deepspeed": deepspeed_config} if deepspeed_config else {}),
  218. )
  219. # 本地模式用 WebSocket 回调,远程模式用传入的文件日志回调
  220. # 用 is None 判断而非 falsy,因为 DDP 非 rank 0 传入空列表 [],不需要进度回调
  221. all_callbacks = callbacks if callbacks is not None else [_ProgressCallback(job_id)]
  222. if task_type == "sft":
  223. from transformers import Trainer
  224. trainer = Trainer(
  225. model=self._model,
  226. args=tr_args,
  227. train_dataset=dataset,
  228. data_collator=DataCollatorForSeq2Seq(self._tokenizer),
  229. callbacks=all_callbacks,
  230. )
  231. elif task_type == "dpo":
  232. from copy import deepcopy
  233. # 兼容旧版 transformers(缺少 MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES)
  234. import transformers.models.auto.modeling_auto as _ma
  235. if not hasattr(_ma, "MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES"):
  236. _ma.MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES = {}
  237. from trl import DPOConfig, DPOTrainer
  238. # 兼容旧版 transformers:Trainer.__init__ 不接受 tokenizer/processing_class
  239. from transformers import Trainer as _HFTrainer
  240. _orig_trainer_init = _HFTrainer.__init__
  241. if not getattr(_HFTrainer, "_patched_kwargs", False):
  242. def _patched_trainer_init(self, *args, **kwargs):
  243. kwargs.pop("tokenizer", None)
  244. kwargs.pop("processing_class", None)
  245. _orig_trainer_init(self, *args, **kwargs)
  246. _HFTrainer.__init__ = _patched_trainer_init
  247. _HFTrainer._patched_kwargs = True
  248. # 显式创建 reference model 并冻结,避免 AdaLora 多 adapter 冲突
  249. ref_model = deepcopy(self._model)
  250. ref_model.eval()
  251. for param in ref_model.parameters():
  252. param.requires_grad = False
  253. # 将 ref_model 上的 PEFT adapter 设为推理模式
  254. # AdaLora 只允许 1 个可训练 adapter,policy model 已有 1 个
  255. if hasattr(ref_model, "set_adapter"):
  256. try:
  257. ref_model.set_adapter("default", inference_mode=True)
  258. except Exception:
  259. pass
  260. elif hasattr(ref_model, "peft_config"):
  261. for adapter_name in list(ref_model.peft_config.keys()):
  262. try:
  263. ref_model.peft_config[adapter_name].inference_mode = True
  264. except Exception:
  265. pass
  266. base_trainer_kwargs = dict(
  267. output_dir=output_dir,
  268. num_train_epochs=epochs,
  269. max_steps=max_steps,
  270. per_device_train_batch_size=batch_size,
  271. gradient_accumulation_steps=gradient_accumulation,
  272. learning_rate=learning_rate,
  273. warmup_ratio=warmup_ratio,
  274. save_strategy=save_strategy,
  275. logging_steps=10,
  276. fp16=True,
  277. report_to="none",
  278. dataloader_num_workers=4,
  279. dataloader_pin_memory=False,
  280. )
  281. trainer = DPOTrainer(
  282. model=self._model,
  283. ref_model=ref_model,
  284. args=DPOConfig(**base_trainer_kwargs),
  285. train_dataset=dataset,
  286. tokenizer=self._tokenizer,
  287. )
  288. elif task_type == "ppo":
  289. import torch
  290. from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer
  291. ppo_epochs = training_args.get("ppo_epochs", 4)
  292. vf_coef = training_args.get("vf_coef", 0.1)
  293. kl_coef = training_args.get("kl_coef", 0.2)
  294. response_length = training_args.get("response_length", 512)
  295. reward_model_path = training_args.get("reward_model_path")
  296. reward_type = training_args.get("reward_type", "heuristic")
  297. # PPO 专用:仅 tokenize prompt
  298. ppo_dataset = self._tokenize_dataset_ppo(dataset_path, max_seq_length, response_length)
  299. # PPO 需要 AutoModelForCausalLMWithValueHead(添加 value head 用于评估动作价值)
  300. # 通过 peft_config 参数让 TRL 内部处理 PEFT 包装,返回的对象是 PreTrainedModelWrapper
  301. # 不能用 get_peft_model(会产生 PeftModel,PPOTrainer 不认)
  302. self._model = AutoModelForCausalLMWithValueHead.from_pretrained(
  303. self._model, peft_config=peft_config,
  304. )
  305. if hasattr(self._model, "print_trainable_parameters"):
  306. self._model.print_trainable_parameters()
  307. # TRL 0.9.x PPOConfig 只接受 PPO 专用参数,不支持 HuggingFace Trainer 参数
  308. # mini_batch_size 必须满足:batch_size % (mini_batch_size * gradient_accumulation_steps) == 0
  309. ppo_config = PPOConfig(
  310. learning_rate=learning_rate,
  311. batch_size=batch_size,
  312. mini_batch_size=1,
  313. gradient_accumulation_steps=gradient_accumulation,
  314. ppo_epochs=ppo_epochs,
  315. vf_coef=vf_coef,
  316. init_kl_coef=kl_coef,
  317. )
  318. # ref_model=None 让 PPOTrainer 自动创建冻结的 reference model(用于 KL 惩罚)
  319. trainer = PPOTrainer(
  320. config=ppo_config,
  321. model=self._model,
  322. ref_model=None,
  323. tokenizer=self._tokenizer,
  324. dataset=ppo_dataset,
  325. )
  326. dataloader = trainer.dataloader
  327. total_steps = len(dataloader) * epochs
  328. step_count = 0
  329. for epoch in range(epochs):
  330. for batch in dataloader:
  331. step_count += 1
  332. query_tensors = batch["input_ids"]
  333. # 生成回答
  334. response_tensors = []
  335. for query in query_tensors:
  336. query_tensor = torch.tensor(query).unsqueeze(0).to(self._model.device)
  337. gen_output = self._model.generate(
  338. query_tensor,
  339. max_new_tokens=response_length,
  340. do_sample=True,
  341. top_p=0.9,
  342. temperature=0.7,
  343. )
  344. response_tensors.append(gen_output[0][query_tensor.shape[-1]:])
  345. # 解码文本用于奖励计算
  346. responses_text = [
  347. self._tokenizer.decode(r, skip_special_tokens=True)
  348. for r in response_tensors
  349. ]
  350. prompts_text = [
  351. self._tokenizer.decode(q, skip_special_tokens=True)
  352. for q in query_tensors
  353. ]
  354. # 计算奖励
  355. if reward_type == "model" and reward_model_path:
  356. from transformers import AutoModelForSequenceClassification
  357. reward_model = AutoModelForSequenceClassification.from_pretrained(
  358. reward_model_path, device_map={"": 0}
  359. )
  360. reward_inputs = [p + r for p, r in zip(prompts_text, responses_text)]
  361. tokenized = self._tokenizer(
  362. reward_inputs, return_tensors="pt", padding=True, truncation=True
  363. ).to(self._model.device)
  364. with torch.no_grad():
  365. rewards = reward_model(**tokenized).logits.squeeze(-1).tolist()
  366. else:
  367. rewards = _compute_heuristic_reward(prompts_text, responses_text)
  368. reward_tensors = [torch.tensor(r, device=self._model.device) for r in rewards]
  369. # PPO 更新
  370. stats = trainer.step(query_tensors, response_tensors, reward_tensors)
  371. # 报告进度
  372. if step_count % 10 == 0:
  373. for cb in (all_callbacks or []):
  374. if hasattr(cb, "on_log"):
  375. cb.on_log(
  376. SimpleNamespace(),
  377. SimpleNamespace(
  378. epoch=epoch, global_step=step_count, max_steps=total_steps
  379. ),
  380. None,
  381. logs={
  382. "loss": stats.get("ppo/loss/total", 0),
  383. "learning_rate": stats.get("ppo/learning_rate", learning_rate),
  384. },
  385. )
  386. os.makedirs(output_dir, exist_ok=True)
  387. self._model.save_pretrained(output_dir)
  388. self._tokenizer.save_pretrained(output_dir)
  389. logger.info(f"PPO training completed for job {job_id}")
  390. return output_dir
  391. else:
  392. raise ValueError(f"Unsupported task_type: {task_type}. Supported: sft, dpo, ppo")
  393. try:
  394. trainer.train()
  395. self._model.save_pretrained(output_dir)
  396. self._tokenizer.save_pretrained(output_dir)
  397. logger.info(f"Training completed for job {job_id}")
  398. except Exception as e:
  399. logger.error(f"Training failed for job {job_id}: {e}")
  400. raise
  401. return output_dir
  402. def get_model_info(self, model_id: str) -> dict[str, Any]:
  403. """读取模型配置信息。"""
  404. import json
  405. from pathlib import Path
  406. # 同步查找模型路径(兼容 HF 和 ModelScope)
  407. candidates = [
  408. settings.models_dir / model_id.replace("/", "_"),
  409. settings.models_dir / model_id,
  410. ]
  411. config_path = None
  412. for p in candidates:
  413. if (p / "config.json").exists():
  414. config_path = p / "config.json"
  415. break
  416. if not config_path:
  417. # 最后尝试扫描
  418. model_name = model_id.split("/")[-1]
  419. for cp in settings.models_dir.rglob("config.json"):
  420. if model_name in str(cp.parent):
  421. config_path = cp
  422. break
  423. if config_path.exists():
  424. with open(config_path) as f:
  425. config = json.load(f)
  426. return {
  427. "model_type": config.get("model_type", "causal_lm"),
  428. "context_length": config.get("max_position_embeddings", config.get("max_sequence_length", 2048)),
  429. "hidden_size": config.get("hidden_size", 0),
  430. "num_layers": config.get("num_hidden_layers", 0),
  431. }
  432. return {"model_type": "causal_lm", "context_length": 2048}
  433. def _tokenize_dataset_ppo(self, dataset_path: str, max_seq_length: int, response_length: int):
  434. """Tokenize PPO 数据集:仅 prompt(模型在训练中自己生成回答)。"""
  435. from datasets import Dataset as HFDataset
  436. data = []
  437. with open(dataset_path, "r", encoding="utf-8") as f:
  438. for line in f:
  439. line = line.strip()
  440. if line:
  441. item = json.loads(line)
  442. if "prompt" not in item:
  443. item["prompt"] = item.get("question", item.get("query", item.get("text", item.get("input", ""))))
  444. if isinstance(item["prompt"], (list, dict)):
  445. item["prompt"] = json.dumps(item["prompt"], ensure_ascii=False)
  446. item["prompt"] = str(item["prompt"])
  447. data.append(item)
  448. hf_dataset = HFDataset.from_list(data)
  449. def tokenize_fn(batch):
  450. raw_prompts = batch.get("prompt", [])
  451. prompts = [str(v) if v is not None else "" for v in raw_prompts]
  452. # 仅 tokenize prompt,预留 response_length 空间给生成的回答
  453. tokenized = self._tokenizer(
  454. prompts,
  455. truncation=True,
  456. max_length=max_seq_length - response_length,
  457. padding=False,
  458. )
  459. return tokenized
  460. tokenized_dataset = hf_dataset.map(
  461. tokenize_fn,
  462. batched=True,
  463. remove_columns=hf_dataset.column_names,
  464. )
  465. return tokenized_dataset
  466. def _tokenize_dataset(self, dataset_path: str, max_seq_length: int):
  467. """Tokenize 处理后的 JSONL 数据集。"""
  468. from datasets import Dataset as HFDataset
  469. data = []
  470. with open(dataset_path, "r", encoding="utf-8") as f:
  471. for line in f:
  472. line = line.strip()
  473. if line:
  474. item = json.loads(line)
  475. # 兼容多种列名 → 统一映射为 prompt / completion
  476. if "prompt" not in item:
  477. item["prompt"] = item.get("question", item.get("query", item.get("text", item.get("input", ""))))
  478. if "completion" not in item:
  479. item["completion"] = item.get("answer", item.get("response", item.get("target", item.get("output", ""))))
  480. # 确保 prompt 和 completion 是字符串
  481. if isinstance(item["prompt"], (list, dict)):
  482. item["prompt"] = json.dumps(item["prompt"], ensure_ascii=False)
  483. item["prompt"] = str(item["prompt"])
  484. if isinstance(item["completion"], (list, dict)):
  485. item["completion"] = json.dumps(item["completion"], ensure_ascii=False)
  486. item["completion"] = str(item["completion"])
  487. data.append(item)
  488. hf_dataset = HFDataset.from_list(data)
  489. def tokenize_fn(batch):
  490. def _to_str(v):
  491. if isinstance(v, (list, dict)):
  492. return json.dumps(v, ensure_ascii=False)
  493. return str(v) if v is not None else ""
  494. raw_prompts = batch.get("prompt", [])
  495. raw_completions = batch.get("completion", [])
  496. prompts = [_to_str(v) for v in raw_prompts]
  497. completions = [_to_str(v) for v in raw_completions]
  498. if not prompts:
  499. return {"input_ids": [], "attention_mask": [], "labels": []}
  500. full_texts = [f"{p}\n{c}" for p, c in zip(prompts, completions)]
  501. tokenized = self._tokenizer(
  502. full_texts, truncation=True, max_length=max_seq_length, padding=False,
  503. )
  504. tokenized["labels"] = list(tokenized["input_ids"])
  505. return tokenized
  506. tokenized_dataset = hf_dataset.map(
  507. tokenize_fn,
  508. batched=True,
  509. remove_columns=["prompt", "completion"],
  510. )
  511. return tokenized_dataset
  512. def _load_dataset_dpo(self, dataset_path: str):
  513. """加载 DPO 数据集,保留 prompt/chosen/rejected 原始文本,由 DPOTrainer 内部 tokenize。"""
  514. from datasets import Dataset as HFDataset
  515. data = []
  516. with open(dataset_path, "r", encoding="utf-8") as f:
  517. for line in f:
  518. line = line.strip()
  519. if line:
  520. item = json.loads(line)
  521. prompt = item.get("prompt", item.get("instruction", item.get("input", "")))
  522. chosen = item.get("chosen", item.get("positive", ""))
  523. rejected = item.get("rejected", item.get("negative", ""))
  524. data.append({
  525. "prompt": str(prompt) if prompt else "",
  526. "chosen": str(chosen) if chosen else "",
  527. "rejected": str(rejected) if rejected else "",
  528. })
  529. return HFDataset.from_list(data)
  530. try:
  531. from transformers import TrainerCallback as _TrainerCallbackBase
  532. except ImportError:
  533. _TrainerCallbackBase = object # 151 主节点无 transformers,仅做占位
  534. class _ProgressCallback(_TrainerCallbackBase):
  535. """自定义训练进度回调,通过 WebSocket 发送进度。"""
  536. def __init__(self, job_id: str):
  537. super().__init__()
  538. self.job_id = job_id
  539. def on_log(self, args, state, control, logs=None, **kwargs):
  540. if logs and "loss" in logs:
  541. asyncio.create_task(
  542. send_progress(
  543. self.job_id,
  544. epoch=int(state.epoch or 0),
  545. step=state.global_step,
  546. total_steps=state.max_steps or 0,
  547. loss=logs["loss"],
  548. learning_rate=logs.get("learning_rate", 0),
  549. )
  550. )
  551. def on_epoch_end(self, args, state, control, **kwargs):
  552. asyncio.create_task(
  553. send_epoch_done(self.job_id, epoch=int(state.epoch or 0), eval_loss=None, eval_accuracy=None)
  554. )
  555. def on_train_end(self, args, state, control, **kwargs):
  556. asyncio.create_task(
  557. send_completed(
  558. self.job_id,
  559. total_time_seconds=getattr(state, "train_runtime", 0),
  560. adapter_path=str(settings.adapters_dir / self.job_id),
  561. )
  562. )
  563. # 全局单例
  564. text_engine = TextEngine()
  565. def _compute_heuristic_reward(prompts: list[str], responses: list[str]) -> list[float]:
  566. """启发式奖励函数:无需额外奖励模型即可用于 PPO 训练。
  567. 评分维度:长度合理性 + 非空 + 重复度惩罚。
  568. """
  569. rewards = []
  570. for _prompt, response in zip(prompts, responses):
  571. reward = 0.0
  572. resp_len = len(response.split())
  573. # 长度评分:20-200 词为佳
  574. if 20 <= resp_len <= 200:
  575. reward += 0.5
  576. elif resp_len < 5:
  577. reward -= 1.0
  578. elif resp_len > 500:
  579. reward -= 0.5
  580. # 非空奖励
  581. if response.strip():
  582. reward += 0.2
  583. # 重复度惩罚(trigram 重复率过高)
  584. words = response.split()
  585. if len(words) > 10:
  586. trigrams = set(tuple(words[i:i+3]) for i in range(len(words) - 2))
  587. if len(trigrams) < len(words) * 0.3:
  588. reward -= 0.5
  589. rewards.append(reward)
  590. return rewards