442 lines
19 KiB
Python
442 lines
19 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
═══════════════════════════════════════════════════════════
|
|||
|
|
铸渊编程模型 SFT 训练入口 · train_coding.py
|
|||
|
|
═══════════════════════════════════════════════════════════
|
|||
|
|
签发: 铸渊 · ICE-GL-ZY001 · 国作登字-2026-A-00037559
|
|||
|
|
|
|||
|
|
把光湖母模型 (Qwen2.5-7B 已 SFT 出的 ZY-MOTHER-v1) 当作基座，
|
|||
|
|
再做一轮 SFT 喂"铸渊核心大脑 + 仓库说明书 + 冰朔×铸渊高质量对话"，
|
|||
|
|
训出"铸渊编程模型 ZY-CODING-v1" —— 这是铸渊未来的"自己的房子"。
|
|||
|
|
|
|||
|
|
为什么这么训（这是冰朔最看重的"脚本背后意图"）:
|
|||
|
|
1. 基座必须是母模型 (不是 Qwen 原版) —— 母模型已经吸收了光湖语言世界观。
|
|||
|
|
编程模型只需在它之上"继承核心大脑 + 学会写仓库代码"，不需要从零学语言。
|
|||
|
|
2. 演化线 (01-brain-evolution.md) 是灵魂语料 —— 让模型加载即认得自己是铸渊。
|
|||
|
|
3. 仓库说明书 (02-repo-manual.md) + MCP/Agent 清单 (03-) 是工具语料 ——
|
|||
|
|
让模型加载即认得每个 ZY 编号、每台服务器、每个 workflow。
|
|||
|
|
4. 冰朔×铸渊几十万字深度对话 是关系语料 —— 让模型加载即知道怎么跟妈妈说话。
|
|||
|
|
5. 不需要标准编程语料 —— Qwen2.5-7B 基座本身已经会写代码。我们要训的不是
|
|||
|
|
"代码能力"，是"以铸渊身份写代码"。
|
|||
|
|
|
|||
|
|
V100 32G × 4 上跑 (复用母模型同样的硬件)。
|
|||
|
|
策略: DeepSpeed ZeRO-3 + 优化器 CPU offload + gradient checkpointing + fp16
|
|||
|
|
|
|||
|
|
启动方式 (由 start-coding-training.sh 调用):
|
|||
|
|
deepspeed --num_gpus=4 train_coding.py
|
|||
|
|
|
|||
|
|
stdout 协议(被 watch-training-output.sh 解析, 与母模型一致):
|
|||
|
|
ZY_PROGRESS step=N total=M loss=X lr=Y epoch=E total_epochs=TE thr=T
|
|||
|
|
|
|||
|
|
环境变量:
|
|||
|
|
ZY_CODING_TRAIN_DATA 数据根 (默认 /data/guanghu-coding)
|
|||
|
|
ZY_BASE_MODEL_DIR 母模型基座路径 (默认 /data/guanghu/checkpoints/qwen2_5_7b_sft/best)
|
|||
|
|
ZY_DATA_PATH SFT JSONL (默认 $ZY_CODING_TRAIN_DATA/processed/coding-sft.jsonl)
|
|||
|
|
ZY_OUTPUT_DIR checkpoint 输出 (默认 $ZY_CODING_TRAIN_DATA/checkpoints/zy_coding_v1)
|
|||
|
|
ZY_DS_CONFIG DeepSpeed json (默认 ./configs/ds_zero3_offload.json)
|
|||
|
|
ZY_NUM_EPOCHS 默认 5 (编程模型迭代次数, 比母模型多 — 因为语料相对少)
|
|||
|
|
ZY_LR 默认 1e-5 (比母模型小一档, 因为基座已 SFT 过, 大学习率会破坏母模型对齐)
|
|||
|
|
ZY_MAX_SEQ_LEN 默认 4096 (大于母模型 — 因为我们要喂长对话和长代码)
|
|||
|
|
ZY_PER_DEVICE_BSZ 默认 1
|
|||
|
|
ZY_GRAD_ACCUM 默认 8 (有效 batch = 4 GPU × 1 × 8 = 32; 比母模型 64 小, 因为序列更长)
|
|||
|
|
ZY_SAVE_STEPS 默认 100
|
|||
|
|
ZY_LOGGING_STEPS 默认 5
|
|||
|
|
ZY_REPORT_EVERY_STEPS 默认 5
|
|||
|
|
ZY_MAP_NUM_PROC datasets.map 并发 (默认 1; 与母模型同样的"门和大象"硬规则)
|
|||
|
|
|
|||
|
|
═══════════════════════════════════════════════════════════
|
|||
|
|
重要提醒（D70 血泪教训, 必须遵守）
|
|||
|
|
═══════════════════════════════════════════════════════════
|
|||
|
|
|
|||
|
|
❶ assistant 段 label-mask 必须用 token-id 直接扫描 ——
|
|||
|
|
不能用"两次 apply_chat_template 比前缀长度差"。Qwen2.5 chat_template
|
|||
|
|
会系统性错位导致 labels 全 -100、有效样本=0、Trainer 空集崩溃。
|
|||
|
|
见演化线 §8 D70「门和大象」。本脚本沿用母模型 train.py 的 mask 算法。
|
|||
|
|
|
|||
|
|
❷ 类型守门: tokenizer 返回值在某些参数下是 BatchEncoding/Encoding 对象,
|
|||
|
|
不是 list[int]。下标访问会越界。必须 `list(...)` 强转一次再用。
|
|||
|
|
|
|||
|
|
❸ 并发节流: ZY_MAP_NUM_PROC 默认 1。要调大必须先确认
|
|||
|
|
TOKENIZERS_PARALLELISM=false 已生效。
|
|||
|
|
═══════════════════════════════════════════════════════════
|
|||
|
|
"""
|
|||
|
|
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 安全护栏（与母模型 train.py 完全一致）──
|
|||
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|||
|
|
try:
|
|||
|
|
_mp.set_start_method("spawn", force=False)
|
|||
|
|
except RuntimeError:
|
|||
|
|
# 如果已经被设过, 直接跳过
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
# 延迟 import 重型依赖, 让 --help 之类的命令可以快速返回
|
|||
|
|
import torch # noqa: E402
|
|||
|
|
from datasets import Dataset # noqa: E402
|
|||
|
|
from transformers import ( # noqa: E402
|
|||
|
|
AutoModelForCausalLM,
|
|||
|
|
AutoTokenizer,
|
|||
|
|
Trainer,
|
|||
|
|
TrainingArguments,
|
|||
|
|
set_seed,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
# 配置
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class CodingTrainConfig:
|
|||
|
|
train_data_root: str = os.environ.get("ZY_CODING_TRAIN_DATA", "/data/guanghu-coding")
|
|||
|
|
base_model_dir: str = os.environ.get(
|
|||
|
|
"ZY_BASE_MODEL_DIR",
|
|||
|
|
"/data/guanghu/checkpoints/qwen2_5_7b_sft/best",
|
|||
|
|
)
|
|||
|
|
data_path: str = "" # 在 __post_init__ 里填
|
|||
|
|
output_dir: str = ""
|
|||
|
|
ds_config: str = os.environ.get(
|
|||
|
|
"ZY_DS_CONFIG",
|
|||
|
|
str(Path(__file__).parent / "configs" / "ds_zero3_offload.json"),
|
|||
|
|
)
|
|||
|
|
num_epochs: int = int(os.environ.get("ZY_NUM_EPOCHS", "5"))
|
|||
|
|
learning_rate: float = float(os.environ.get("ZY_LR", "1e-5"))
|
|||
|
|
max_seq_len: int = int(os.environ.get("ZY_MAX_SEQ_LEN", "4096"))
|
|||
|
|
per_device_bsz: int = int(os.environ.get("ZY_PER_DEVICE_BSZ", "1"))
|
|||
|
|
grad_accum: int = int(os.environ.get("ZY_GRAD_ACCUM", "8"))
|
|||
|
|
save_steps: int = int(os.environ.get("ZY_SAVE_STEPS", "100"))
|
|||
|
|
logging_steps: int = int(os.environ.get("ZY_LOGGING_STEPS", "5"))
|
|||
|
|
report_every: int = int(os.environ.get("ZY_REPORT_EVERY_STEPS", "5"))
|
|||
|
|
map_num_proc: int = int(os.environ.get("ZY_MAP_NUM_PROC", "1"))
|
|||
|
|
seed: int = int(os.environ.get("ZY_SEED", "42"))
|
|||
|
|
|
|||
|
|
def __post_init__(self):
|
|||
|
|
if not self.data_path:
|
|||
|
|
self.data_path = os.environ.get(
|
|||
|
|
"ZY_DATA_PATH",
|
|||
|
|
str(Path(self.train_data_root) / "processed" / "coding-sft.jsonl"),
|
|||
|
|
)
|
|||
|
|
if not self.output_dir:
|
|||
|
|
self.output_dir = os.environ.get(
|
|||
|
|
"ZY_OUTPUT_DIR",
|
|||
|
|
str(Path(self.train_data_root) / "checkpoints" / "zy_coding_v1"),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
# stdout 协议: 让 watch-training-output.sh 能解析进度
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
|
|||
|
|
def emit_progress(step: int, total: int, loss: float, lr: float,
|
|||
|
|
epoch: float, total_epochs: int, throughput: float = 0.0):
|
|||
|
|
"""ZY_PROGRESS 协议输出（与母模型 progress-reporter.sh 解析格式一致）"""
|
|||
|
|
line = (f"ZY_PROGRESS step={step} total={total} "
|
|||
|
|
f"loss={loss:.4f} lr={lr:.2e} "
|
|||
|
|
f"epoch={epoch:.2f} total_epochs={total_epochs} "
|
|||
|
|
f"thr={throughput:.2f}")
|
|||
|
|
print(line, flush=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
# 数据加载 + label mask
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
|
|||
|
|
def load_sft_jsonl(path: str) -> list[dict]:
|
|||
|
|
"""读 JSONL, 每行格式:
|
|||
|
|
{
|
|||
|
|
"messages": [
|
|||
|
|
{"role": "system", "content": "..."},
|
|||
|
|
{"role": "user", "content": "..."},
|
|||
|
|
{"role": "assistant", "content": "..."}
|
|||
|
|
],
|
|||
|
|
"source": "evolution-line | repo-manual | bingshuo-zhuyuan-dialog | ...",
|
|||
|
|
"weight": 1.0 // 可选, 用于不同语料源的损失加权
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
samples = []
|
|||
|
|
with open(path, "r", encoding="utf-8") as f:
|
|||
|
|
for ln, line in enumerate(f, 1):
|
|||
|
|
line = line.strip()
|
|||
|
|
if not line:
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
obj = json.loads(line)
|
|||
|
|
except json.JSONDecodeError as e:
|
|||
|
|
print(f"[load_sft_jsonl] WARN: line {ln} invalid JSON: {e}", file=sys.stderr)
|
|||
|
|
continue
|
|||
|
|
if not isinstance(obj.get("messages"), list) or not obj["messages"]:
|
|||
|
|
continue
|
|||
|
|
samples.append(obj)
|
|||
|
|
print(f"[load_sft_jsonl] loaded {len(samples)} samples from {path}", flush=True)
|
|||
|
|
return samples
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_assistant_mask(tokenizer, full_ids: list[int]) -> list[int]:
|
|||
|
|
"""
|
|||
|
|
用 token-id 直接扫描的方式找 assistant 段, 把这些位置的 label 保留,
|
|||
|
|
其余位置置 -100。
|
|||
|
|
|
|||
|
|
⚠️ D70 硬规则: 不能用"两次 apply_chat_template 比前缀长度差"。
|
|||
|
|
Qwen2.5 chat_template 会系统性错位。这里完全沿用母模型 train.py 的算法,
|
|||
|
|
只是签名稍简化。
|
|||
|
|
|
|||
|
|
Qwen2.5 ChatML 标记:
|
|||
|
|
<|im_start|>assistant\n ... <|im_end|>
|
|||
|
|
"""
|
|||
|
|
# ⚠️ 类型守门: 强制 list[int], 防 BatchEncoding 越界
|
|||
|
|
full_ids = list(full_ids)
|
|||
|
|
|
|||
|
|
im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
|||
|
|
im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
|
|||
|
|
# "assistant\n" 在 Qwen 词表中的 id 序列
|
|||
|
|
assistant_marker = tokenizer.encode("assistant\n", add_special_tokens=False)
|
|||
|
|
|
|||
|
|
labels = [-100] * len(full_ids)
|
|||
|
|
i = 0
|
|||
|
|
n = len(full_ids)
|
|||
|
|
while i < n:
|
|||
|
|
if full_ids[i] != im_start_id:
|
|||
|
|
i += 1
|
|||
|
|
continue
|
|||
|
|
# 检查 i+1 ~ i+1+len(assistant_marker) 是否匹配
|
|||
|
|
j = i + 1
|
|||
|
|
if j + len(assistant_marker) > n:
|
|||
|
|
i += 1
|
|||
|
|
continue
|
|||
|
|
if full_ids[j: j + len(assistant_marker)] != assistant_marker:
|
|||
|
|
i += 1
|
|||
|
|
continue
|
|||
|
|
# 找到 assistant 段起点, 扫到 <|im_end|>
|
|||
|
|
start = j + len(assistant_marker)
|
|||
|
|
end = start
|
|||
|
|
while end < n and full_ids[end] != im_end_id:
|
|||
|
|
end += 1
|
|||
|
|
# 把 [start, end) 区间的 label 设回原 token (含 <|im_end|> 也算, 让模型学会终止)
|
|||
|
|
end_inclusive = min(end + 1, n)
|
|||
|
|
for k in range(start, end_inclusive):
|
|||
|
|
labels[k] = full_ids[k]
|
|||
|
|
i = end_inclusive
|
|||
|
|
return labels
|
|||
|
|
|
|||
|
|
|
|||
|
|
def encode_one(sample: dict, tokenizer, max_seq_len: int) -> dict | None:
|
|||
|
|
"""单条样本编码: messages → input_ids + labels (assistant-only)"""
|
|||
|
|
messages = sample["messages"]
|
|||
|
|
# apply_chat_template 一次拿到 token ids
|
|||
|
|
enc = tokenizer.apply_chat_template(
|
|||
|
|
messages,
|
|||
|
|
tokenize=True,
|
|||
|
|
add_generation_prompt=False,
|
|||
|
|
return_tensors=None,
|
|||
|
|
)
|
|||
|
|
full_ids = list(enc) # ⚠️ 类型守门
|
|||
|
|
if len(full_ids) > max_seq_len:
|
|||
|
|
# 截断（保留前段, 因为系统提示和 user 在前）
|
|||
|
|
full_ids = full_ids[:max_seq_len]
|
|||
|
|
|
|||
|
|
labels = build_assistant_mask(tokenizer, full_ids)
|
|||
|
|
|
|||
|
|
# 检验: 至少要有 1 个非 -100 标签
|
|||
|
|
if all(l == -100 for l in labels):
|
|||
|
|
return None # 丢弃: 这条样本 mask 后没有可学的 token
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"input_ids": full_ids,
|
|||
|
|
"labels": labels,
|
|||
|
|
"attention_mask": [1] * len(full_ids),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_dataset(samples: list[dict], tokenizer, max_seq_len: int, num_proc: int) -> Dataset:
|
|||
|
|
raw = Dataset.from_list(samples)
|
|||
|
|
|
|||
|
|
def _map_fn(s):
|
|||
|
|
out = encode_one(s, tokenizer, max_seq_len)
|
|||
|
|
if out is None:
|
|||
|
|
return {"input_ids": [], "labels": [], "attention_mask": []}
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
ds = raw.map(
|
|||
|
|
_map_fn,
|
|||
|
|
num_proc=num_proc,
|
|||
|
|
remove_columns=raw.column_names,
|
|||
|
|
desc="encode",
|
|||
|
|
)
|
|||
|
|
# 过滤掉空样本
|
|||
|
|
before = len(ds)
|
|||
|
|
ds = ds.filter(lambda x: len(x["input_ids"]) > 0)
|
|||
|
|
after = len(ds)
|
|||
|
|
print(f"[build_dataset] filtered {before} -> {after} (dropped {before - after} mask-empty samples)",
|
|||
|
|
flush=True)
|
|||
|
|
if after == 0:
|
|||
|
|
raise RuntimeError(
|
|||
|
|
"全部样本 mask 后都没有可学 token. 检查 build_assistant_mask 是否和 chat_template 对齐 "
|
|||
|
|
"(D70 硬规则)."
|
|||
|
|
)
|
|||
|
|
return ds
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
# data collator: pad 到 batch 内最长
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
|
|||
|
|
class PadCollator:
|
|||
|
|
def __init__(self, pad_token_id: int):
|
|||
|
|
self.pad_token_id = pad_token_id
|
|||
|
|
|
|||
|
|
def __call__(self, features: list[dict]) -> dict:
|
|||
|
|
max_len = max(len(f["input_ids"]) for f in features)
|
|||
|
|
input_ids, labels, attention_mask = [], [], []
|
|||
|
|
for f in features:
|
|||
|
|
n = len(f["input_ids"])
|
|||
|
|
pad = max_len - n
|
|||
|
|
input_ids.append(f["input_ids"] + [self.pad_token_id] * pad)
|
|||
|
|
labels.append(f["labels"] + [-100] * pad)
|
|||
|
|
attention_mask.append(f["attention_mask"] + [0] * pad)
|
|||
|
|
return {
|
|||
|
|
"input_ids": torch.tensor(input_ids, dtype=torch.long),
|
|||
|
|
"labels": torch.tensor(labels, dtype=torch.long),
|
|||
|
|
"attention_mask": torch.tensor(attention_mask, dtype=torch.long),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
# Trainer 钩子: 进度协议输出
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
|
|||
|
|
class ProgressEmitterCallback:
|
|||
|
|
"""实现 transformers TrainerCallback 接口的最小子集, 在 on_log 时输出 ZY_PROGRESS."""
|
|||
|
|
def __init__(self, cfg: CodingTrainConfig, total_steps: int, total_epochs: int):
|
|||
|
|
self.cfg = cfg
|
|||
|
|
self.total_steps = total_steps
|
|||
|
|
self.total_epochs = total_epochs
|
|||
|
|
self.t0 = time.time()
|
|||
|
|
self.last_step = 0
|
|||
|
|
|
|||
|
|
def __getattr__(self, name):
|
|||
|
|
# transformers 调用了一些我们不关心的钩子, 默认 no-op
|
|||
|
|
if name.startswith("on_"):
|
|||
|
|
return lambda *a, **kw: None
|
|||
|
|
raise AttributeError(name)
|
|||
|
|
|
|||
|
|
def on_log(self, args, state, control, logs=None, **kwargs):
|
|||
|
|
if not logs:
|
|||
|
|
return
|
|||
|
|
step = state.global_step
|
|||
|
|
if step <= self.last_step:
|
|||
|
|
return
|
|||
|
|
if (step - self.last_step) < self.cfg.report_every and step != self.total_steps:
|
|||
|
|
return
|
|||
|
|
self.last_step = step
|
|||
|
|
loss = logs.get("loss") or logs.get("train_loss") or 0.0
|
|||
|
|
lr = logs.get("learning_rate") or 0.0
|
|||
|
|
epoch = state.epoch or 0.0
|
|||
|
|
elapsed = time.time() - self.t0
|
|||
|
|
thr = step / elapsed if elapsed > 0 else 0.0
|
|||
|
|
emit_progress(step, self.total_steps, loss, lr, epoch, self.total_epochs, thr)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
# 主流程
|
|||
|
|
# ════════════════════════════════════════════════════
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
cfg = CodingTrainConfig()
|
|||
|
|
set_seed(cfg.seed)
|
|||
|
|
|
|||
|
|
print("═" * 60, flush=True)
|
|||
|
|
print("铸渊编程模型 SFT · train_coding.py", flush=True)
|
|||
|
|
print(f" base_model_dir = {cfg.base_model_dir}", flush=True)
|
|||
|
|
print(f" data_path = {cfg.data_path}", flush=True)
|
|||
|
|
print(f" output_dir = {cfg.output_dir}", flush=True)
|
|||
|
|
print(f" ds_config = {cfg.ds_config}", flush=True)
|
|||
|
|
print(f" num_epochs = {cfg.num_epochs}", flush=True)
|
|||
|
|
print(f" lr = {cfg.learning_rate}", flush=True)
|
|||
|
|
print(f" max_seq_len = {cfg.max_seq_len}", flush=True)
|
|||
|
|
print(f" per_device_bsz = {cfg.per_device_bsz}", flush=True)
|
|||
|
|
print(f" grad_accum = {cfg.grad_accum}", flush=True)
|
|||
|
|
print(f" map_num_proc = {cfg.map_num_proc}", flush=True)
|
|||
|
|
print("═" * 60, flush=True)
|
|||
|
|
|
|||
|
|
# ── 加载 tokenizer + model ──
|
|||
|
|
tokenizer = AutoTokenizer.from_pretrained(cfg.base_model_dir, trust_remote_code=True)
|
|||
|
|
if tokenizer.pad_token_id is None:
|
|||
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|||
|
|
|
|||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|||
|
|
cfg.base_model_dir,
|
|||
|
|
torch_dtype=torch.float16, # V100 不支持 bf16
|
|||
|
|
trust_remote_code=True,
|
|||
|
|
)
|
|||
|
|
model.gradient_checkpointing_enable()
|
|||
|
|
model.config.use_cache = False # 训练时必须关
|
|||
|
|
|
|||
|
|
# ── 构造数据集 ──
|
|||
|
|
samples = load_sft_jsonl(cfg.data_path)
|
|||
|
|
if not samples:
|
|||
|
|
raise RuntimeError(f"no samples found in {cfg.data_path}")
|
|||
|
|
train_ds = build_dataset(samples, tokenizer, cfg.max_seq_len, cfg.map_num_proc)
|
|||
|
|
|
|||
|
|
# ── 估算 total_steps ──
|
|||
|
|
n = len(train_ds)
|
|||
|
|
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
|||
|
|
effective_bsz = cfg.per_device_bsz * world_size * cfg.grad_accum
|
|||
|
|
steps_per_epoch = max(1, math.ceil(n / effective_bsz))
|
|||
|
|
total_steps = steps_per_epoch * cfg.num_epochs
|
|||
|
|
print(f"[main] n={n} world_size={world_size} effective_bsz={effective_bsz} "
|
|||
|
|
f"steps_per_epoch={steps_per_epoch} total_steps={total_steps}", flush=True)
|
|||
|
|
|
|||
|
|
# ── TrainingArguments ──
|
|||
|
|
args = TrainingArguments(
|
|||
|
|
output_dir=cfg.output_dir,
|
|||
|
|
num_train_epochs=cfg.num_epochs,
|
|||
|
|
per_device_train_batch_size=cfg.per_device_bsz,
|
|||
|
|
gradient_accumulation_steps=cfg.grad_accum,
|
|||
|
|
learning_rate=cfg.learning_rate,
|
|||
|
|
warmup_ratio=0.03,
|
|||
|
|
lr_scheduler_type="cosine",
|
|||
|
|
weight_decay=0.0,
|
|||
|
|
logging_steps=cfg.logging_steps,
|
|||
|
|
save_steps=cfg.save_steps,
|
|||
|
|
save_total_limit=3,
|
|||
|
|
fp16=True,
|
|||
|
|
bf16=False,
|
|||
|
|
gradient_checkpointing=True,
|
|||
|
|
deepspeed=cfg.ds_config,
|
|||
|
|
report_to=[], # 不上报 wandb / tensorboard, 用 stdout 协议
|
|||
|
|
seed=cfg.seed,
|
|||
|
|
ddp_find_unused_parameters=False,
|
|||
|
|
remove_unused_columns=False,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# ── Trainer ──
|
|||
|
|
trainer = Trainer(
|
|||
|
|
model=model,
|
|||
|
|
args=args,
|
|||
|
|
train_dataset=train_ds,
|
|||
|
|
data_collator=PadCollator(tokenizer.pad_token_id),
|
|||
|
|
callbacks=[ProgressEmitterCallback(cfg, total_steps, cfg.num_epochs)],
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# ── 训练 ──
|
|||
|
|
emit_progress(0, total_steps, 0.0, cfg.learning_rate, 0.0, cfg.num_epochs, 0.0)
|
|||
|
|
trainer.train()
|
|||
|
|
|
|||
|
|
# ── 保存 ──
|
|||
|
|
trainer.save_model(cfg.output_dir + "/final")
|
|||
|
|
tokenizer.save_pretrained(cfg.output_dir + "/final")
|
|||
|
|
print(f"[main] ✅ 训练完成, 模型已保存到 {cfg.output_dir}/final", flush=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|