431 lines
17 KiB
Python
431 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
═══════════════════════════════════════════════════════════
|
||
Qwen2.5-7B 全参数 SFT 训练入口 · train.py
|
||
═══════════════════════════════════════════════════════════
|
||
签发: 铸渊 · ICE-GL-ZY001 · 国作登字-2026-A-00037559
|
||
|
||
V100 32G × 4 上的 Qwen2.5-7B 全参 SFT.
|
||
策略: DeepSpeed ZeRO-3 + 优化器 CPU offload + gradient checkpointing + fp16
|
||
|
||
启动方式 (由 start-training.sh 调用):
|
||
deepspeed --num_gpus=4 train.py
|
||
|
||
stdout 协议(被 watch-training-output.sh 解析):
|
||
ZY_PROGRESS step=N total=M loss=X lr=Y epoch=E total_epochs=TE thr=T
|
||
|
||
环境变量:
|
||
ZY_TRAIN_DATA 数据根 (默认 /data/guanghu)
|
||
ZY_MODEL_DIR 模型路径 (默认 $ZY_TRAIN_DATA/models/Qwen2.5-7B)
|
||
ZY_DATA_PATH SFT JSONL (默认 $ZY_TRAIN_DATA/processed/sft.jsonl)
|
||
ZY_OUTPUT_DIR checkpoint 输出 (默认 $ZY_TRAIN_DATA/checkpoints/qwen2_5_7b_sft)
|
||
ZY_DS_CONFIG DeepSpeed json (默认 server/training-agent/configs/ds_zero3_offload.json)
|
||
ZY_NUM_EPOCHS 默认 3
|
||
ZY_LR 默认 2e-5
|
||
ZY_MAX_SEQ_LEN 默认 2048
|
||
ZY_PER_DEVICE_BSZ 默认 1
|
||
ZY_GRAD_ACCUM 默认 16
|
||
ZY_SAVE_STEPS 默认 200
|
||
ZY_LOGGING_STEPS 默认 5
|
||
ZY_REPORT_EVERY_STEPS 默认 5 (ZY_PROGRESS 协议输出节流)
|
||
"""
|
||
from __future__ import annotations
|
||
import json
|
||
import math
|
||
import multiprocessing as _mp
|
||
import os
|
||
import sys
|
||
import time
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
# ── multi-rank × multi-proc fork 安全护栏 ──
|
||
# 必须在 import torch / transformers / datasets 之前生效:
|
||
# 1. 关掉 fast tokenizer 的 Rust 线程池, 避免 fork 后子进程死锁/abort
|
||
# (这是 deepspeed 多 rank 同时在 datasets.map 里 fork worker 时
|
||
# iflatmap_unordered 静默崩溃的最常见根因).
|
||
# 2. 把 multiprocessing 默认启动方式从 fork 切到 spawn, 即使后续有人
|
||
# 把 ZY_MAP_NUM_PROC 调高也安全.
|
||
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
||
try:
|
||
# force=True 会强制覆盖已有的 start method;
|
||
# 极少数情况下 (子解释器/已有活跃池) 仍会抛 RuntimeError, 此时降级跳过.
|
||
_mp.set_start_method("spawn", force=True)
|
||
except RuntimeError:
|
||
pass
|
||
|
||
import torch
|
||
from datasets import load_dataset
|
||
from transformers import (
|
||
AutoModelForCausalLM,
|
||
AutoTokenizer,
|
||
Trainer,
|
||
TrainingArguments,
|
||
TrainerCallback,
|
||
set_seed,
|
||
)
|
||
|
||
# ── 配置 ──
|
||
DATA_ROOT = Path(os.environ.get("ZY_TRAIN_DATA", "/data/guanghu"))
|
||
MODEL_DIR = Path(os.environ.get("ZY_MODEL_DIR", str(DATA_ROOT / "models" / "Qwen2.5-7B")))
|
||
DATA_PATH = Path(os.environ.get("ZY_DATA_PATH", str(DATA_ROOT / "processed" / "sft.jsonl")))
|
||
OUTPUT_DIR = Path(os.environ.get("ZY_OUTPUT_DIR", str(DATA_ROOT / "checkpoints" / "qwen2_5_7b_sft")))
|
||
DS_CONFIG = Path(os.environ.get("ZY_DS_CONFIG", str(Path(__file__).parent / "configs" / "ds_zero3_offload.json")))
|
||
|
||
NUM_EPOCHS = float(os.environ.get("ZY_NUM_EPOCHS", "3"))
|
||
LR = float(os.environ.get("ZY_LR", "2e-5"))
|
||
MAX_SEQ_LEN = int(os.environ.get("ZY_MAX_SEQ_LEN", "2048"))
|
||
PER_DEVICE_BSZ = int(os.environ.get("ZY_PER_DEVICE_BSZ", "1"))
|
||
GRAD_ACCUM = int(os.environ.get("ZY_GRAD_ACCUM", "16"))
|
||
SAVE_STEPS = int(os.environ.get("ZY_SAVE_STEPS", "200"))
|
||
LOGGING_STEPS = int(os.environ.get("ZY_LOGGING_STEPS", "5"))
|
||
REPORT_EVERY = int(os.environ.get("ZY_REPORT_EVERY_STEPS", "5"))
|
||
SEED = int(os.environ.get("ZY_SEED", "42"))
|
||
|
||
# datasets.map 的 worker 数. 默认 1 — 这是有意的:
|
||
# - deepspeed --num_gpus=N 起 N 个 rank, 每个 rank 都会调一次 build_dataset.
|
||
# 如果这里再 fork 出 cpu//2 个 map worker, 就是 N × (cpu//2) 个 fork 子进程
|
||
# 同时加载 fast tokenizer + 同时写 datasets cache, 极易触发
|
||
# `iflatmap_unordered` 静默崩溃.
|
||
# - 几万行 SFT 用 fast tokenizer 串行 tokenize 通常 < 1 分钟, 远小于一次训练代价.
|
||
# - 真有大数据集需要并行, 用 ZY_MAP_NUM_PROC=N 显式开 (此时务必只让 rank 0 跑 map,
|
||
# 由调用方保证, 这里不再做 multi-rank 同时并行 map 的兼容).
|
||
MAP_NUM_PROC = max(1, int(os.environ.get("ZY_MAP_NUM_PROC", "1")))
|
||
|
||
IGNORE_INDEX = -100
|
||
|
||
|
||
def is_main_process() -> bool:
|
||
return int(os.environ.get("LOCAL_RANK", "0")) == 0
|
||
|
||
|
||
def log(msg: str):
|
||
if is_main_process():
|
||
print(msg, flush=True)
|
||
|
||
|
||
# ── 数据加载 + 模板化 ──
|
||
|
||
def _resolve_role_anchors(tokenizer) -> dict[str, list[int]]:
|
||
"""解析 ChatML 的 assistant 段定位锚点.
|
||
|
||
返回:
|
||
im_start_id — <|im_start|> 的 token id
|
||
im_end_id — <|im_end|> 的 token id
|
||
asst_role_ids — "assistant\n" 编码后的 token id 序列 (作为内容前的角色头)
|
||
|
||
思路:
|
||
Qwen2.5 的 chat_template 把每段对话包成 <|im_start|>{role}\n{content}<|im_end|>\n.
|
||
<|im_start|> / <|im_end|> 是 special token, BPE 不会跨它们合并,
|
||
因此可以在 full_ids 上直接用 token-id 做线性扫描定位 assistant 段,
|
||
不依赖任何 "前缀对齐" 假设. 这是修复 "有效样本: 0" 故障的关键.
|
||
"""
|
||
im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
||
im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
|
||
if im_start_id is None or im_end_id is None or im_start_id == tokenizer.unk_token_id:
|
||
raise RuntimeError(
|
||
"[train] tokenizer 缺少 <|im_start|>/<|im_end|> special token, "
|
||
"确认这是 Qwen2.5 系列模型 (而不是兼容包装)."
|
||
)
|
||
# "assistant\n" 编码为 content token (不带任何 special token)
|
||
asst_role_ids = tokenizer.encode("assistant\n", add_special_tokens=False)
|
||
if not asst_role_ids:
|
||
raise RuntimeError("[train] 无法编码 'assistant\\n' 角色头, tokenizer 异常.")
|
||
return {
|
||
"im_start_id": im_start_id,
|
||
"im_end_id": im_end_id,
|
||
"asst_role_ids": asst_role_ids,
|
||
}
|
||
|
||
|
||
def _mask_assistant_segments(
|
||
full_ids: list[int],
|
||
im_start_id: int,
|
||
im_end_id: int,
|
||
asst_role_ids: list[int],
|
||
) -> tuple[list[int], int]:
|
||
"""在 full_ids 上线性扫描, 找出所有 assistant 段并返回 labels + 标记到的 token 数.
|
||
|
||
被标记的范围 = assistant 内容 + 闭合的 <|im_end|> (让模型学会自然停止).
|
||
若 max_seq_len 截断导致最后一段没有闭合的 im_end_id, 则标到序列末尾.
|
||
"""
|
||
n = len(full_ids)
|
||
labels = [IGNORE_INDEX] * n
|
||
role_len = len(asst_role_ids)
|
||
marked = 0
|
||
k = 0
|
||
while k < n:
|
||
if full_ids[k] != im_start_id:
|
||
k += 1
|
||
continue
|
||
# 检查 <|im_start|> 后是否是 "assistant\n"
|
||
if k + role_len < n and full_ids[k + 1 : k + 1 + role_len] == asst_role_ids:
|
||
content_start = k + 1 + role_len
|
||
# 找到内容结尾的 <|im_end|>
|
||
j = content_start
|
||
while j < n and full_ids[j] != im_end_id:
|
||
j += 1
|
||
content_end = min(j, n - 1) # 含 im_end (若存在), 否则到序列末尾
|
||
for p in range(content_start, content_end + 1):
|
||
labels[p] = full_ids[p]
|
||
marked += content_end + 1 - content_start
|
||
k = content_end + 1
|
||
else:
|
||
k += 1
|
||
return labels, marked
|
||
|
||
|
||
def build_dataset(tokenizer):
|
||
if not DATA_PATH.is_file():
|
||
raise FileNotFoundError(f"训练数据不存在: {DATA_PATH} · 请先跑 preprocess-corpus.py")
|
||
|
||
log(f"[train] 加载数据: {DATA_PATH}")
|
||
raw = load_dataset("json", data_files=str(DATA_PATH), split="train")
|
||
log(f"[train] 样本数: {len(raw)}")
|
||
|
||
anchors = _resolve_role_anchors(tokenizer)
|
||
im_start_id = anchors["im_start_id"]
|
||
im_end_id = anchors["im_end_id"]
|
||
asst_role_ids = anchors["asst_role_ids"]
|
||
|
||
def encode(example: dict[str, Any], idx: int | None = None) -> dict[str, list[int]]:
|
||
try:
|
||
msgs = example["messages"]
|
||
# apply_chat_template 一次性按 Qwen2.5 ChatML 格式化整段对话
|
||
full_ids = tokenizer.apply_chat_template(
|
||
msgs, tokenize=True, add_generation_prompt=False,
|
||
truncation=True, max_length=MAX_SEQ_LEN,
|
||
)
|
||
labels, _marked = _mask_assistant_segments(
|
||
full_ids, im_start_id, im_end_id, asst_role_ids
|
||
)
|
||
return {"input_ids": full_ids, "labels": labels, "attention_mask": [1] * len(full_ids)}
|
||
except Exception as e:
|
||
# 让 worker 子进程在崩之前留下"那一条样本是谁"的痕迹.
|
||
# iflatmap_unordered 在主进程只能看到 worker 的最终 traceback,
|
||
# 这里把样本指纹打到 stderr, 便于事后定位脏数据.
|
||
try:
|
||
msgs = example.get("messages")
|
||
preview = json.dumps(msgs, ensure_ascii=False)[:300] if msgs is not None else "<no messages>"
|
||
except Exception:
|
||
preview = "<unprintable>"
|
||
sys.stderr.write(
|
||
f"[train.encode] 样本编码失败 idx={idx} err={type(e).__name__}: {e}\n"
|
||
f"[train.encode] messages 预览: {preview}\n"
|
||
)
|
||
sys.stderr.flush()
|
||
raise
|
||
|
||
cols = raw.column_names
|
||
log(f"[train] tokenize map: num_proc={MAP_NUM_PROC} (默认 1, 用 ZY_MAP_NUM_PROC 覆盖)")
|
||
ds = raw.map(encode, with_indices=True, remove_columns=cols, num_proc=MAP_NUM_PROC)
|
||
# 过滤掉没有 assistant token 的样本 (理论上极少, 留作防御)
|
||
ds = ds.filter(lambda ex: any(l != IGNORE_INDEX for l in ex["labels"]))
|
||
log(f"[train] 有效样本: {len(ds)}")
|
||
|
||
# ── "自动门" 防呆守护 ──
|
||
# 若一条都没标到, 立即停, 不让 Trainer 在空数据集上裸奔.
|
||
# 同时打首条样本的诊断信息, 让下一次定位时 < 30 秒.
|
||
if len(ds) == 0:
|
||
sample = raw[0] if len(raw) > 0 else None
|
||
diag = ["[train] ❌ 有效样本为 0 — 标记环节失效, 终止训练."]
|
||
if sample is not None:
|
||
try:
|
||
fid = tokenizer.apply_chat_template(
|
||
sample["messages"], tokenize=True, add_generation_prompt=False,
|
||
truncation=True, max_length=MAX_SEQ_LEN,
|
||
)
|
||
preview = tokenizer.decode(fid[:200], skip_special_tokens=False)
|
||
diag.append(f"[train] 诊断: im_start_id={im_start_id} im_end_id={im_end_id} "
|
||
f"asst_role_ids={asst_role_ids}")
|
||
diag.append(f"[train] 诊断: 首条样本 token 数={len(fid)}")
|
||
diag.append(f"[train] 诊断: 首条样本前 200 token 解码=\n{preview}")
|
||
except Exception as e:
|
||
diag.append(f"[train] 诊断失败: {e}")
|
||
for line in diag:
|
||
log(line)
|
||
raise RuntimeError("有效样本为 0 — 见上方诊断.")
|
||
|
||
# 标注质量统计 — "门会回报它做了什么"
|
||
if is_main_process():
|
||
try:
|
||
sample_n = min(64, len(ds))
|
||
asst_tokens = 0
|
||
total_tokens = 0
|
||
for i in range(sample_n):
|
||
row = ds[i]
|
||
total_tokens += len(row["labels"])
|
||
asst_tokens += sum(1 for l in row["labels"] if l != IGNORE_INDEX)
|
||
ratio = asst_tokens / total_tokens if total_tokens else 0.0
|
||
log(f"[train] 标注统计 (前 {sample_n} 条采样): "
|
||
f"平均 assistant token 占比={ratio:.2%} "
|
||
f"(平均 {asst_tokens // sample_n} tok/样本)")
|
||
except Exception as e:
|
||
log(f"[train] 标注统计失败(非致命): {e}")
|
||
return ds
|
||
|
||
|
||
@dataclass
|
||
class PadCollator:
|
||
tokenizer: Any
|
||
pad_to_multiple_of: int = 8
|
||
|
||
def __call__(self, features: list[dict]) -> dict[str, torch.Tensor]:
|
||
max_len = max(len(f["input_ids"]) for f in features)
|
||
if self.pad_to_multiple_of > 1:
|
||
max_len = math.ceil(max_len / self.pad_to_multiple_of) * self.pad_to_multiple_of
|
||
pad_id = self.tokenizer.pad_token_id
|
||
if pad_id is None:
|
||
pad_id = self.tokenizer.eos_token_id
|
||
|
||
def _pad(seq: list[int], val: int) -> list[int]:
|
||
return seq + [val] * (max_len - len(seq))
|
||
|
||
input_ids = torch.tensor([_pad(f["input_ids"], pad_id) for f in features], dtype=torch.long)
|
||
labels = torch.tensor([_pad(f["labels"], IGNORE_INDEX) for f in features], dtype=torch.long)
|
||
attn = torch.tensor([_pad(f["attention_mask"], 0) for f in features], dtype=torch.long)
|
||
return {"input_ids": input_ids, "labels": labels, "attention_mask": attn}
|
||
|
||
|
||
# ── 心跳协议: Trainer Callback ──
|
||
|
||
class ZYProgressCallback(TrainerCallback):
|
||
"""每 REPORT_EVERY 步输出 stdout 协议行,被 watch-training-output.sh 解析."""
|
||
|
||
def __init__(self):
|
||
self.t0 = time.time()
|
||
self.last_step = -1
|
||
|
||
def on_log(self, args, state, control, logs=None, **kwargs): # noqa: D401
|
||
if not is_main_process() or not logs:
|
||
return
|
||
step = state.global_step or 0
|
||
if step == self.last_step:
|
||
return
|
||
if step > 0 and (step - self.last_step) < REPORT_EVERY and step != state.max_steps:
|
||
return
|
||
self.last_step = step
|
||
loss = logs.get("loss") if "loss" in logs else logs.get("train_loss")
|
||
lr = logs.get("learning_rate")
|
||
elapsed = max(time.time() - self.t0, 1e-6)
|
||
thr = step / elapsed if step > 0 else 0.0
|
||
epoch = math.floor(state.epoch or 0)
|
||
line = (
|
||
f"ZY_PROGRESS step={step} total={state.max_steps} "
|
||
f"epoch={epoch} total_epochs={int(args.num_train_epochs)} "
|
||
f"loss={loss if loss is not None else 'nan'} "
|
||
f"lr={lr if lr is not None else 'nan'} "
|
||
f"thr={thr:.4f}"
|
||
)
|
||
print(line, flush=True)
|
||
|
||
|
||
def main() -> int:
|
||
set_seed(SEED)
|
||
|
||
if not MODEL_DIR.is_dir():
|
||
log(f"❌ 模型目录不存在: {MODEL_DIR} · 请先跑 download-model.py")
|
||
return 2
|
||
if not DS_CONFIG.is_file():
|
||
log(f"❌ DeepSpeed 配置不存在: {DS_CONFIG}")
|
||
return 2
|
||
|
||
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
||
log(f"[train] 模型: {MODEL_DIR}")
|
||
log(f"[train] 数据: {DATA_PATH}")
|
||
log(f"[train] 输出: {OUTPUT_DIR}")
|
||
log(f"[train] DS: {DS_CONFIG}")
|
||
|
||
tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR), trust_remote_code=True, use_fast=True)
|
||
if tokenizer.pad_token is None:
|
||
tokenizer.pad_token = tokenizer.eos_token
|
||
|
||
train_ds = build_dataset(tokenizer)
|
||
|
||
log("[train] 加载模型 (fp16)...")
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
str(MODEL_DIR),
|
||
dtype=torch.float16,
|
||
trust_remote_code=True,
|
||
use_cache=False, # 与 gradient_checkpointing 冲突
|
||
)
|
||
model.gradient_checkpointing_enable()
|
||
if hasattr(model, "enable_input_require_grads"):
|
||
model.enable_input_require_grads()
|
||
|
||
args = TrainingArguments(
|
||
output_dir=str(OUTPUT_DIR),
|
||
num_train_epochs=NUM_EPOCHS,
|
||
per_device_train_batch_size=PER_DEVICE_BSZ,
|
||
gradient_accumulation_steps=GRAD_ACCUM,
|
||
learning_rate=LR,
|
||
warmup_ratio=0.03,
|
||
lr_scheduler_type="cosine",
|
||
weight_decay=0.0,
|
||
max_grad_norm=1.0,
|
||
fp16=True,
|
||
bf16=False,
|
||
gradient_checkpointing=True,
|
||
logging_steps=LOGGING_STEPS,
|
||
save_steps=SAVE_STEPS,
|
||
save_total_limit=3,
|
||
save_strategy="steps",
|
||
report_to=["tensorboard"],
|
||
deepspeed=str(DS_CONFIG),
|
||
ddp_find_unused_parameters=False,
|
||
dataloader_num_workers=2,
|
||
dataloader_pin_memory=True,
|
||
remove_unused_columns=False,
|
||
seed=SEED,
|
||
)
|
||
|
||
trainer = Trainer(
|
||
model=model,
|
||
args=args,
|
||
train_dataset=train_ds,
|
||
data_collator=PadCollator(tokenizer=tokenizer),
|
||
callbacks=[ZYProgressCallback()],
|
||
)
|
||
|
||
if is_main_process():
|
||
# 写一个供副将查询的训练 meta
|
||
try:
|
||
(OUTPUT_DIR / "training-meta.json").write_text(
|
||
json.dumps({
|
||
"model": str(MODEL_DIR),
|
||
"data": str(DATA_PATH),
|
||
"max_seq_len": MAX_SEQ_LEN,
|
||
"num_train_epochs": NUM_EPOCHS,
|
||
"per_device_batch": PER_DEVICE_BSZ,
|
||
"grad_accum": GRAD_ACCUM,
|
||
"effective_batch": PER_DEVICE_BSZ * GRAD_ACCUM * max(1, torch.cuda.device_count()),
|
||
"lr": LR,
|
||
"fp16": True,
|
||
"deepspeed_config": str(DS_CONFIG),
|
||
}, indent=2, ensure_ascii=False),
|
||
encoding="utf-8",
|
||
)
|
||
except Exception as e:
|
||
log(f"[train] meta 写入失败(非致命): {e}")
|
||
|
||
log("[train] 🔥 开始训练")
|
||
train_result = trainer.train()
|
||
|
||
log("[train] 训练结束 · 保存最终模型...")
|
||
trainer.save_model(str(OUTPUT_DIR / "final"))
|
||
trainer.save_state()
|
||
if is_main_process():
|
||
try:
|
||
tokenizer.save_pretrained(str(OUTPUT_DIR / "final"))
|
||
except Exception as e:
|
||
log(f"[train] tokenizer 保存失败(非致命): {e}")
|
||
|
||
log(f"[train] ✅ 完成 · global_step={trainer.state.global_step} · loss={train_result.training_loss:.4f}")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|