| |
| """gc25d 双向 SFT（SPEC §6–§7）：Wan2.2-TI2V-5B 全参数微调 + 零初始化 gc_sem_cond（逐像素 parts id 条件），latent 与文本都预计算。 |
| |
| ActionRoPE `actionrope/train.py` 的 fork（结构逐段对应，便于 diff）：AropeTrainModule → GcTrainModule（model.gc_forward）， |
| arope / 动作臂 / CNS / 画布 / 权重图全部去掉，换成 sem 条件；数据集换成 dataset.GcLatentDataset（三种 dropout 在 worker 里抽）。 |
| |
| # 单进程（smoke，不用 DeepSpeed；AdamW 直接更新 bf16 参数、无 fp32 主权重，数字不代表训练效果） |
| CUDA_VISIBLE_DEVICES=2 .venv/bin/python code/gc25d/train.py --limit 8 --max_steps 4 --output outputs/_smoke/x |
| |
| # 多卡：accelerate + DeepSpeed ZeRO-2（启动器 scripts/train_sft.sh 选 / 生成 configs/accelerate_zero2*.yaml） |
| OUT=outputs/sft_v10 bash scripts/train_sft.sh |
| |
| 与 ActionRoPE 的差别（都在 GcTrainModule / main 里） |
| -------------------------------------------------- |
| * 前向 = model.gc_forward（plain RoPE、逐 cell 文本、patch embedding 后加 gc_sem_cond(sem_ids)），aux 带回 sem_rms / x_rms |
| （sem 项与 patch embedding 的 RMS；proj_out 零初始化 ⇒ step 1 的 sem_rms 恰为 0，之后应单调长起来）。 |
| * loss = MSE(帧 1..20) × training_weight(t)（权重图恒为 1，所以 ActionRoPE 的 loss_w 就是 mse，不再单列）。 |
| * 两个 AdamW 参数组：`dit.gc_sem_cond.*` 用 --cond_lr（默认 1e-4 = 10×；零初始化的输入通道在 lr 1e-5 下学得极慢， |
| train_recipe §6），其余 --lr；wd 相同。LambdaLR 的 warmup 系数同时乘到两组上。 |
| * 验证：固定 val clip × 固定 t × 固定噪声，每个 (clip, t) 跑三遍：cond（真 ids + 真文本）/ uncond（ids 全 255）/ |
| nulltext（NULL_PROMPT × 21）；记 mse 与两个 gap（uncond − cond、nulltext − cond）—— gap 是条件有没有被用上的直接读数。 |
| * 导出 = gcdistill.gc_dit.save_gc（831 个 bf16 键 = 825 Wan + 6 gc_sem_cond，键不带 dit.，metadata 带 vocab 等）， |
| 写完用 check_gc_ckpt 体检；save_state 走 distill.runlib.save_state_rotating（state_A / state_B 两槽 + state 软链）。 |
| * 热启动 --resume X.safetensors = model.load_gc_weights：gc25d 导出（词表必须相同）或普通 Wan / plain 导出 |
| （缺全部 gc_sem_cond.* ⇒ 保持零初始化）。--resume <dir> = accelerator.load_state 续训，**数据位置精确接上**： |
| trainer_state.json 记 (epoch, 本 epoch 已消费的 micro batch 数, 几何)，续训用 skip_first_batches 跳过已消费的部分 |
| （ActionRoPE 从 epoch 开头重放）；导出 metadata 的 init_ckpt 沿用首次启动的热启动权重，续训另记 resumed_from / resumed_step。 |
| * 分辨率 = geometry（env GC_RES，默认 832x480；scripts/train_sft.sh 透传）：--data_root 的 res 必须等于它（dataset 构造时查）， |
| 导出 metadata / trainer_state 记 res；dir 续训的存档 res 必须相同；safetensors 热启动的 res 不同只打印提示（跨分辨率热启动合法）。 |
| * 数据顺序用 accelerate 的 SeedableRandomSampler（data_seed = --seed）：每个 epoch 的排列 = f(seed, epoch)，各 rank 同一排列再分片； |
| 不用它时排列来自 sampler generator 的累积状态，load_state 不存（只给 IterableDataset 存），续训没法接上。 |
| * step_time 是整步墙钟（含等数据），data_time 单列等 DataLoader 的部分（epoch 开头重新 fork worker 也在这里）。 |
| |
| 学习率调度不交给 accelerate.prepare：accelerate 包装后的 scheduler 每步会 step num_processes 次 |
| （它假设 scheduler 是按样本数建的），warmup 会被悄悄压缩；这里手动 step，一步就是一步。 |
| |
| 随机数：set_seed(seed, device_specific=True) 让每个 rank 的 torch / CUDA RNG 都不同 —— 噪声 ε（randn_like）各 rank 独立； |
| t 由各 rank 的 cpu_gen 抽；dropout 抽签按 (seed, epoch, index) 在 dataset 里定（与 rank / worker 无关）。数据分片不受影响： |
| SeedableRandomSampler 用显式 data_seed，不吃 device_specific 的种子。不加 device_specific 时 8 张卡每步抽到同一份 ε。 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import os |
| import platform |
| import socket |
| import sys |
| import time |
| import warnings |
| import zlib |
|
|
| os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True") |
|
|
| import torch |
| import torch.nn as nn |
| from accelerate import Accelerator |
| from accelerate.utils import DataLoaderConfiguration, DistributedType, set_seed |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| ROOT = os.path.dirname(os.path.dirname(HERE)) |
| AROPE_ROOT = os.environ.get("AROPE_ROOT", os.path.join(ROOT, "code", "arope")) |
| for _p in (HERE, AROPE_ROOT): |
| if _p not in sys.path and os.path.isdir(_p): |
| sys.path.append(_p) |
| import geometry as GEO |
| import model as GM |
| from dataset import CELLS, UNKNOWN_ID, GcLatentDataset, gc_collate |
| from text import NULL_PROMPT |
| from gcdistill.gc_dit import GC_DIT_KEYS, GC_META, check_gc_ckpt, cond_spec, save_gc |
| from distill.runlib import CsvLog, save_state_rotating |
| from diffsynth.diffusion.flow_match import FlowMatchScheduler |
|
|
| MODEL_DIR = os.path.join(ROOT, "models/Wan2.2-TI2V-5B") |
| DATA_ROOT = os.path.join(ROOT, "data/gc25d_v10") |
| DIT_FILES = [os.path.join(MODEL_DIR, f"diffusion_pytorch_model-0000{i}-of-00003.safetensors") for i in (1, 2, 3)] |
| |
| DIT_PREFIX = "dit." |
| COND_PREFIX = cond_spec("ids")["prefix"] |
| VAL_VARIANTS = ("cond", "uncond", "nulltext") |
|
|
|
|
| |
| |
| |
|
|
| def parse_args(argv=None): |
| ap = argparse.ArgumentParser(description="gc25d 双向 SFT（SPEC §7）") |
| ap.add_argument("--vocab", choices=tuple(GM.VOCABS), default="parts", |
| help="条件词表：parts（v10 默认，K=20）/ ids（旧 gc25d，K=10，parts 经 collapse_parts 折叠）") |
| ap.add_argument("--cond_type", choices=("ids",), default="ids", help="条件类型（gc25d v10 只有语义 id；记进 metadata）") |
| ap.add_argument("--cond_key", choices=("ids",), default="ids", help="cond/*.npz 读哪个键（记进 metadata）") |
| ap.add_argument("--text_dropout", type=float, default=0.1, help="逐样本：21 个 cell 全换成 NULL_PROMPT") |
| ap.add_argument("--cond_dropout", type=float, default=0.05, help="逐样本：ids 全 255（与 text dropout 独立）") |
| ap.add_argument("--mask_p", type=float, default=0.5, help="逐样本：世界锚定的 UNKNOWN 遮挡（cond 未 drop 时）") |
| ap.add_argument("--include_dropped", action="store_true", help="训练集保留 clips.jsonl 里 drop=true 的 clip（默认剔除）") |
| ap.add_argument("--first_frame_cond", action=argparse.BooleanOptionalAction, default=True) |
|
|
| ap.add_argument("--data_root", default=DATA_ROOT, help="latents/ cond/ meta/ index/clips.jsonl text_table_gc25d.pt") |
| ap.add_argument("--text_table", default=None, help="默认 <data_root>/text_table_gc25d.pt") |
| ap.add_argument("--model_dir", default=MODEL_DIR) |
| ap.add_argument("--limit", type=int, default=None, help="smoke 用：只取这么多训练 clip") |
|
|
| ap.add_argument("--lr", type=float, default=1e-5) |
| ap.add_argument("--cond_lr", type=float, default=1e-4, help=f"{COND_PREFIX}* 参数组的 lr") |
| ap.add_argument("--weight_decay", type=float, default=0.01) |
| ap.add_argument("--warmup_steps", type=int, default=500, help="线性 warmup 步数（0 ⇒ 常数 lr），之后常数") |
| ap.add_argument("--grad_clip", type=float, default=1.0, |
| help="梯度裁剪阈值；DeepSpeed 下写进 deepspeed_config.gradient_clipping（覆盖 yaml），" |
| "日志里的 grad_norm 是裁剪前的全局范数") |
| ap.add_argument("--max_steps", type=int, default=20000) |
| ap.add_argument("--batch_size", type=int, default=1, help="每卡 micro batch（一个 micro batch 共用一个 t）") |
| ap.add_argument("--grad_accum", type=int, default=2, |
| help="梯度累积步数；step 计的是优化器更新数。DeepSpeed 下必须与 accelerate yaml 的 " |
| "gradient_accumulation_steps 一致（configs/accelerate_zero2_ga2.yaml 是 2）") |
| ap.add_argument("--num_workers", type=int, default=4) |
| ap.add_argument("--gradient_checkpointing", action=argparse.BooleanOptionalAction, default=True) |
|
|
| ap.add_argument("--save_every", type=int, default=1000) |
| ap.add_argument("--save_state", action=argparse.BooleanOptionalAction, default=True, |
| help="每次导出权重时同时 accelerator.save_state（含优化器分片，8 卡约 70 GB）；" |
| "state_A / state_B 两个槽轮流写，<output>/state 是指向最新完整槽的软链") |
| ap.add_argument("--val_every", type=int, default=500, help="0 ⇒ 不验证") |
| ap.add_argument("--val_n", type=int, default=16) |
| ap.add_argument("--val_timesteps", default="200,500,800") |
| ap.add_argument("--val_uncond", action=argparse.BooleanOptionalAction, default=True, |
| help="验证时每个 (clip, t) 再跑 uncond（ids 全 255）与 nulltext（NULL_PROMPT×21）两遍") |
| ap.add_argument("--log_every", type=int, default=1) |
| ap.add_argument("--output", required=True, help="outputs/<run>") |
| ap.add_argument("--resume", default=None, |
| help="*.safetensors ⇒ 只热启动权重（step 从 0 起；普通 Wan 导出 ⇒ gc_sem_cond 零初始化）；" |
| "目录 ⇒ accelerator.load_state（续训）") |
| ap.add_argument("--seed", type=int, default=0) |
| args = ap.parse_args(argv) |
|
|
| if args.text_table is None: |
| args.text_table = os.path.join(args.data_root, "text_table_gc25d.pt") |
| args.val_timesteps = [float(t) for t in args.val_timesteps.split(",") if t.strip()] |
| if not os.path.isabs(args.output): |
| args.output = os.path.join(ROOT, args.output) |
| |
| args.init_ckpt = args.resume if args.resume and args.resume.endswith(".safetensors") else "" |
| args.resumed_from, args.resumed_step = "", 0 |
| return args |
|
|
|
|
| |
| |
| |
|
|
| def load_dit(model_dir: str, device: str = "cpu"): |
| """只加载 DiT（不挂 T5 / VAE），bf16。先落在 CPU，交给 accelerator.prepare 搬上卡。（= model.load_dit）""" |
| return GM.load_dit(model_dir, device) |
|
|
|
|
| def load_weights_into_dit(dit: nn.Module, path: str, vocab: str) -> dict: |
| """从导出的 safetensors 热启动（model.load_gc_weights：严格键集，只放行"全部 gc_sem_cond.* 缺失且模块仍零初始化"）。 |
| ckpt 带 gc_sem_cond 时词表必须与 --vocab 相同。""" |
| info = GM.load_gc_weights(dit, path) |
| if info["vocab"] != vocab: |
| raise RuntimeError(f"{path} 的 gc_sem_cond 词表是 {info['vocab']!r}，本次训练 --vocab {vocab!r}") |
| return {k: v for k, v in info.items() if k != "metadata"} | {"ckpt_metadata": info.get("metadata", {})} |
|
|
|
|
| def sem_param_groups(module: nn.Module, lr: float, cond_lr: float, weight_decay: float = 0.01, |
| prefix: str = COND_PREFIX) -> list[dict]: |
| """可训参数按名字前缀分两组：[其余（lr）, prefix*（cond_lr）]；prefix 组为空时只返回一组。wd 两组相同。 |
| prefix 默认 "dit.gc_sem_cond."（GcTrainModule / gcdistill.GcDiT 都把 WanModel 挂在 .dit 下）。""" |
| rest, cond = [], [] |
| for name, p in module.named_parameters(): |
| if p.requires_grad: |
| (cond if name.startswith(prefix) else rest).append(p) |
| groups = [{"params": rest, "lr": lr, "weight_decay": weight_decay}] |
| if cond: |
| groups.append({"params": cond, "lr": cond_lr, "weight_decay": weight_decay}) |
| return groups |
|
|
|
|
| def warmup_lambda(warmup_steps: int): |
| """LambdaLR 系数：min(1, (s+1)/warm) 线性 warmup 后常数（warm=0 ⇒ 恒 1）。""" |
| warm = max(int(warmup_steps), 0) |
| return lambda s: min(1.0, (s + 1) / warm) if warm > 0 else 1.0 |
|
|
|
|
| class GcTrainModule(nn.Module): |
| """把 DiT + 扩散 loss 包成一个 module，forward 直接返回 loss —— DeepSpeed 引擎要求前向在 module 内。 |
| gc_sem_cond 挂在 dit 上（install_sem_cond），是这个 module 的子模块 ⇒ 进 ZeRO 分片、拿得到梯度。""" |
|
|
| def __init__(self, dit, first_frame_cond: bool, gradient_checkpointing: bool): |
| super().__init__() |
| self.dit = dit |
| self.first_frame_cond = first_frame_cond |
| self.gradient_checkpointing = gradient_checkpointing |
| self.scheduler = FlowMatchScheduler("Wan") |
| self.scheduler.set_timesteps(1000, training=True) |
|
|
| def timestep_id_for(self, t: float) -> int: |
| return int(torch.argmin((self.scheduler.timesteps - t).abs())) |
|
|
| def forward(self, batch: dict, timestep_id: int, noise: torch.Tensor | None = None, train: bool = True): |
| x0 = batch["input_latents"] |
| timestep = self.scheduler.timesteps[timestep_id].view(1) |
| if noise is None: |
| noise = torch.randn_like(x0) |
| |
| x_t = self.scheduler.add_noise(x0, noise, timestep) |
| target = self.scheduler.training_target(x0, noise, timestep) |
| |
| x_t[:, :, 0:1] = x0[:, :, 0:1] |
|
|
| aux = {} |
| pred = GM.gc_forward( |
| self.dit, x_t, timestep.to(device=x0.device, dtype=x0.dtype), batch["context"], batch["sem_ids"], |
| first_frame_cond=self.first_frame_cond, context_ids=batch["context_ids"], |
| use_gradient_checkpointing=self.gradient_checkpointing and train, aux=aux, |
| ) |
| |
| pred, target = pred[:, :, 1:], target[:, :, 1:] |
| se = (pred.float() - target.float()) ** 2 |
| mse = se.mean() |
| loss = mse * self.scheduler.training_weight(timestep).to(mse.device) |
| stats = { |
| "mse": mse.detach(), |
| "timestep": float(timestep.item()), |
| "sem_rms": aux["sem_rms"].detach(), |
| "x_rms": aux["x_rms"].detach(), |
| } |
| return loss, stats |
|
|
|
|
| |
| |
| |
|
|
| def export_metadata(args, step: int, extra: dict | None = None) -> dict: |
| """step-N.safetensors 的 metadata（save_gc 先填 GC_META，再盖这些；全部转成字符串）。""" |
| meta = dict(GC_META) |
| meta.update({ |
| "vocab": args.vocab, "data": "v10", "data_root": args.data_root, "step": step, "res": GEO.RES_TAG, |
| "cond_type": args.cond_type, "cond_key": args.cond_key, "stage": "sft_bidirectional", |
| "text_dropout": args.text_dropout, "cond_dropout": args.cond_dropout, "mask_p": args.mask_p, |
| "lr": args.lr, "cond_lr": args.cond_lr, "weight_decay": args.weight_decay, "warmup_steps": args.warmup_steps, |
| "first_frame_cond": args.first_frame_cond, "seed": args.seed, |
| |
| "init_ckpt": args.init_ckpt, "resumed_from": args.resumed_from, "resumed_step": args.resumed_step, |
| "model_dir": args.model_dir, "trainer": "gc25d/train.py", |
| }) |
| meta.update(extra or {}) |
| return {k: str(v) for k, v in meta.items()} |
|
|
|
|
| def export_dit_weights(accelerator: Accelerator, model, path: str, metadata: dict) -> int: |
| """取完整权重（ZeRO-2 下参数本就不分片），主进程去掉 'dit.' 前缀写 safetensors（gcdistill.gc_dit.save_gc： |
| 断言 831 个 bf16 键、.tmp + rename、chmod 644），写完 check_gc_ckpt 体检。返回文件字节数（主进程）。""" |
| size = save_gc(accelerator, model, path, metadata) |
| if accelerator.is_main_process: |
| res = check_gc_ckpt(path) |
| if res["metadata"].get("vocab") != metadata["vocab"]: |
| raise RuntimeError(f"{path}: metadata vocab {res['metadata'].get('vocab')!r} ≠ {metadata['vocab']!r}") |
| print(f"[save] check_gc_ckpt OK：{res['n_tensors']} 个 {res['dtypes']} 张量，vocab={res['metadata']['vocab']}", flush=True) |
| return size |
|
|
|
|
| def to_device(batch: dict, device) -> dict: |
| return {k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} |
|
|
|
|
| def rotate_csv_if_header_changed(path: str, fields: list[str]) -> str | None: |
| """续训进同一个 OUT、而 CSV 列变了（代码升级）⇒ 旧文件改名 <name>.old<N>.csv，免得 CsvLog 往旧表头下追加错位的行。""" |
| if not (os.path.isfile(path) and os.path.getsize(path) > 0): |
| return None |
| with open(path, encoding="utf-8") as fh: |
| if fh.readline().rstrip("\r\n").split(",") == list(fields): |
| return None |
| stem, ext = os.path.splitext(path) |
| n = 0 |
| while os.path.exists(f"{stem}.old{n}{ext}"): |
| n += 1 |
| os.replace(path, f"{stem}.old{n}{ext}") |
| return f"{stem}.old{n}{ext}" |
|
|
|
|
| |
| |
| |
|
|
| def epoch_steps(n_micro: int, grad_accum: int) -> int: |
| """每 epoch 的优化器步数。accelerate 在 dataloader 末尾强制同步（end_of_dataloader），不足 GA 个 micro batch 的尾巴 |
| 也算一步 ⇒ 向上取整（每进程 25 个 micro、GA 2 ⇒ 13 步，最后一步只有 1 个 micro）。""" |
| return max(math.ceil(n_micro / max(grad_accum, 1)), 1) |
|
|
|
|
| def resume_position(ts: dict, geometry: dict) -> tuple[int, int, str]: |
| """dir 续训的数据位置 → (epoch, 该 epoch 开头要跳过的 micro batch 数, 说明)。 |
| |
| 数据顺序 = f(seed, epoch)（SeedableRandomSampler），dropout 抽签 = f(seed, epoch, index) ⇒ 几何（进程数 / GA / micro batch / |
| 每进程 micro 数 / 训练集大小 / seed）不变时，跳过已消费的 micro batch 就与不中断的训练逐样本相同。 |
| 几何变了 ⇒ 放弃当前 epoch 剩下的部分、从下一个 epoch 开头起（换排列、换抽签，不会重放同一份样本）。 |
| 旧存档（trainer_state 没有 epoch）⇒ 按 step 与当前几何的每 epoch 步数推 epoch（不重放 epoch 0 的抽签，顺序不保证）。""" |
| n_micro, ga = int(geometry["n_micro"]), int(geometry["grad_accum"]) |
| if "epoch" in ts: |
| epoch, used = int(ts["epoch"]), int(ts.get("epoch_micro", 0)) |
| if any(ts.get(k) != v for k, v in geometry.items()): |
| return epoch + 1, 0, "几何变了：从下一个 epoch 开头起" |
| if used >= n_micro: |
| return epoch + 1, 0, "精确（存档时正好在 epoch 末）" |
| return epoch, used, "精确" |
| spe = epoch_steps(n_micro, ga) |
| step = int(ts["step"]) |
| epoch = step // spe |
| return epoch, (step - epoch * spe) * ga, "旧存档：按 step 推算" |
|
|
|
|
| def init_ckpt_of_state(state_dir: str, ts: dict) -> str: |
| """dir 续训时接回首次启动的 init_ckpt：trainer_state.json 里记的；旧存档没记 ⇒ 读同一 run 的 config.json |
| （state_A/B 的上一级）里首次启动的 --resume（是 safetensors 才算）。""" |
| if "init_ckpt" in ts: |
| return ts["init_ckpt"] or "" |
| try: |
| with open(os.path.join(os.path.dirname(os.path.realpath(state_dir)), "config.json"), encoding="utf-8") as fh: |
| a = json.load(fh)["args"] |
| except (OSError, ValueError, KeyError): |
| return "" |
| r = a.get("init_ckpt") or a.get("resume") or "" |
| return r if str(r).endswith(".safetensors") else "" |
|
|
|
|
| |
| |
| |
|
|
| def val_variants(batch: dict, null_context: torch.Tensor | None, null_ids: torch.Tensor | None, uncond: bool) -> dict: |
| """一条 val 样本 → {cond, uncond, nulltext} 三份输入（uncond / nulltext 只换一种条件，与 infer 的两个 CFG 负支同形）。""" |
| out = {"cond": batch} |
| if uncond: |
| out["uncond"] = dict(batch, sem_ids=torch.full_like(batch["sem_ids"], UNKNOWN_ID)) |
| b = batch["context"].shape[0] |
| out["nulltext"] = dict(batch, context=null_context.expand(b, *null_context.shape[1:]), |
| context_ids=null_ids.expand(b, -1)) |
| return out |
|
|
|
|
| @torch.no_grad() |
| def run_validation(accelerator: Accelerator, model, module: GcTrainModule, val_ds, tids: list[int], seed: int, |
| null_context=None, null_ids=None, uncond: bool = True): |
| """固定 clip 集 × 固定 timestep 集 × 固定噪声（三个 variant 共用同一份噪声），各进程分摊 clip，主进程汇总。 |
| |
| ZeRO-2 的前向没有集合通信，所以各进程 clip 数不等也没关系；最后的 reduce 是集合操作，人人都要到。 |
| 返回 {t: {"mse", "mse_uncond", "mse_nulltext", "gap_uncond", "gap_nulltext", "n"}}（所有进程都拿到同样的值）。""" |
| model.eval() |
| device = accelerator.device |
| rank, world = accelerator.process_index, accelerator.num_processes |
| variants = VAL_VARIANTS if uncond else VAL_VARIANTS[:1] |
| acc = torch.zeros(len(tids), len(variants) + 1, device=device, dtype=torch.float64) |
| for i in range(rank, len(val_ds), world): |
| sample = val_ds[i] |
| batch = to_device(gc_collate([sample]), device) |
| vb = val_variants(batch, null_context, null_ids, uncond) |
| for j, tid in enumerate(tids): |
| g = torch.Generator(device=device) |
| g.manual_seed((zlib.crc32(sample["clip_id"].encode()) * 1009 + tid * 7 + seed) % (2 ** 31)) |
| noise = torch.randn(batch["input_latents"].shape, generator=g, device=device, dtype=batch["input_latents"].dtype) |
| for v, name in enumerate(variants): |
| _, stats = model(vb[name], tid, noise=noise, train=False) |
| acc[j, v] += stats["mse"].double() |
| acc[j, -1] += 1 |
| acc = accelerator.reduce(acc, reduction="sum") |
| model.train() |
| out = {} |
| for j, tid in enumerate(tids): |
| n = max(acc[j, -1].item(), 1.0) |
| r = {("mse" if name == "cond" else f"mse_{name}"): acc[j, v].item() / n for v, name in enumerate(variants)} |
| if uncond: |
| r["gap_uncond"] = r["mse_uncond"] - r["mse"] |
| r["gap_nulltext"] = r["mse_nulltext"] - r["mse"] |
| r["n"] = int(acc[j, -1].item()) |
| out[float(module.scheduler.timesteps[tid])] = r |
| return out |
|
|
|
|
| |
| |
| |
|
|
| def main(argv=None): |
| args = parse_args(argv) |
| |
| accelerator = Accelerator(gradient_accumulation_steps=args.grad_accum, |
| dataloader_config=DataLoaderConfiguration(use_seedable_sampler=True, data_seed=args.seed)) |
| is_main = accelerator.is_main_process |
| use_ds = accelerator.distributed_type == DistributedType.DEEPSPEED |
| set_seed(args.seed, device_specific=True) |
| if use_ds: |
| |
| |
| accelerator.state.deepspeed_plugin.deepspeed_config["gradient_clipping"] = float(args.grad_clip) |
| ga = accelerator.state.deepspeed_plugin.deepspeed_config.get("gradient_accumulation_steps") |
| if ga not in ("auto", args.grad_accum): |
| raise SystemExit(f"--grad_accum={args.grad_accum} 与 accelerate yaml 的 gradient_accumulation_steps={ga} 不一致，" |
| f"请换 configs/accelerate_zero2_ga{args.grad_accum}.yaml（scripts/train_sft.sh 会自动选 / 生成）") |
| accelerator.state.deepspeed_plugin.deepspeed_config["gradient_accumulation_steps"] = args.grad_accum |
| if is_main: |
| os.makedirs(args.output, exist_ok=True) |
| print(f"[env] world={accelerator.num_processes} distributed={accelerator.distributed_type} " |
| f"mixed_precision={accelerator.mixed_precision} device={accelerator.device} " |
| f"res={GEO.RES_TAG}（{GEO.RES_ENV}，latent {GEO.LAT_H}×{GEO.LAT_W}，{GEO.TOK_PER_CELL} token/cell）", flush=True) |
|
|
| |
| train_ds = GcLatentDataset(args.data_root, "train", text_table=args.text_table, vocab=args.vocab, |
| text_dropout=args.text_dropout, cond_dropout=args.cond_dropout, mask_p=args.mask_p, |
| limit=args.limit, seed=args.seed, include_dropped=args.include_dropped, verbose=is_main) |
| val_ds = None |
| if args.val_every > 0 and args.val_n > 0: |
| |
| val_ds = GcLatentDataset(args.data_root, "val", text_table=train_ds.text_table, vocab=args.vocab, |
| limit=args.val_n, seed=args.seed, verbose=is_main) |
| leak = set(train_ds.clip_ids()) & set(val_ds.clip_ids()) |
| runs_tr = {train_ds.rows[c]["run"] for c in train_ds.clip_ids()} |
| runs_va = {val_ds.rows[c]["run"] for c in val_ds.clip_ids()} |
| if leak or (runs_tr & runs_va): |
| raise SystemExit(f"train / val 泄漏：共用 clip {sorted(leak)[:3]}，共用 run {sorted(runs_tr & runs_va)[:3]}") |
| |
| dataloader = torch.utils.data.DataLoader( |
| train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True, collate_fn=gc_collate, |
| num_workers=args.num_workers, pin_memory=True, persistent_workers=False, |
| ) |
|
|
| |
| t0 = time.time() |
| dit = load_dit(args.model_dir, device="cpu") |
| GM.install_sem_cond(dit, args.vocab) |
| resume_info = None |
| if args.resume and args.resume.endswith(".safetensors"): |
| resume_info = load_weights_into_dit(dit, args.resume, args.vocab) |
| resume_info["res"] = GEO.ckpt_res(args.resume) |
| if is_main: |
| print(f"[resume] 热启动权重 {args.resume}: n_tensors={resume_info['n_tensors']} " |
| f"has_sem_cond={resume_info['has_sem_cond']} missing={len(resume_info['missing'])}", flush=True) |
| if resume_info["res"] != GEO.RES_TAG: |
| print(f"[resume] 注意：热启动 ckpt 的 res={resume_info['res']}，本次 {GEO.RES_ENV}={GEO.RES_TAG}" |
| f"（跨分辨率热启动；DiT / gc_sem_cond 对分辨率通用）", flush=True) |
| dit.train().requires_grad_(True) |
| module = GcTrainModule(dit, args.first_frame_cond, args.gradient_checkpointing) |
| enc = getattr(dit, GM.SEM_COND_NAME) |
| n_params = sum(p.numel() for p in dit.parameters()) |
| n_sem_params = sum(p.numel() for p in enc.parameters()) |
| if is_main: |
| print(f"[model] DiT 加载 {time.time() - t0:.0f}s，{n_params:,} 参数（gc_sem_cond {n_sem_params:,}，vocab={args.vocab}，" |
| f"in_channels={enc.in_channels}，zero_output={enc.is_zero_output()}）", flush=True) |
|
|
| groups = sem_param_groups(module, args.lr, args.cond_lr, args.weight_decay) |
| assert len(groups) == 2 and sum(len(g["params"]) for g in groups) == len(list(module.parameters())), \ |
| "参数组应恰为 [DiT, gc_sem_cond] 且覆盖全部参数" |
| optimizer = torch.optim.AdamW(groups, lr=args.lr, weight_decay=args.weight_decay) |
| lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, warmup_lambda(args.warmup_steps)) |
|
|
| model, optimizer, dataloader = accelerator.prepare(module, optimizer, dataloader) |
| device = accelerator.device |
| if is_main: |
| clip_eff = getattr(getattr(model, "optimizer", None), "clip_grad", None) if use_ds else args.grad_clip |
| print(f"[env] grad_clip={args.grad_clip}（生效值 {clip_eff}）", flush=True) |
| |
| module.dit.freqs = tuple(t.to(device) for t in module.dit.freqs) |
|
|
| |
| geometry = {"num_processes": accelerator.num_processes, "grad_accum": args.grad_accum, "batch_size": args.batch_size, |
| "n_micro": len(dataloader), "n_train": len(train_ds), "seed": args.seed} |
| start_step, start_epoch, skip_micro = 0, 0, 0 |
| if args.resume and not args.resume.endswith(".safetensors"): |
| with open(os.path.join(args.resume, "trainer_state.json"), encoding="utf-8") as fh: |
| ts = json.load(fh) |
| if ts.get("vocab", args.vocab) != args.vocab: |
| raise SystemExit(f"{args.resume} 是 vocab={ts.get('vocab')!r} 的存档，本次 --vocab {args.vocab!r}") |
| if ts.get("res", GEO.LEGACY_RES) != GEO.RES_TAG: |
| raise SystemExit(f"{args.resume} 是 res={ts.get('res', GEO.LEGACY_RES)} 的存档（缺失 = {GEO.LEGACY_RES}），" |
| f"本次 {GEO.RES_ENV}={GEO.RES_TAG}") |
| accelerator.load_state(args.resume) |
| start_step = int(ts["step"]) |
| with warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| for _ in range(start_step): |
| lr_scheduler.step() |
| start_epoch, skip_micro, how = resume_position(ts, geometry) |
| args.init_ckpt = init_ckpt_of_state(args.resume, ts) |
| args.resumed_from, args.resumed_step = os.path.realpath(args.resume), start_step |
| if is_main: |
| print(f"[resume] accelerator.load_state({args.resume} → {args.resumed_from})，从 step {start_step} 续训；" |
| f"数据从 epoch {start_epoch} 的第 {skip_micro} 个 micro batch 接上（{how}）；init_ckpt={args.init_ckpt or '（底模）'}", |
| flush=True) |
| if how.startswith("几何"): |
| print(f"[resume] 存档几何 { {k: ts.get(k) for k in geometry} } ≠ 本次 {geometry}", flush=True) |
|
|
| |
| if is_main: |
| cfg = { |
| "args": vars(args), |
| "res": GEO.RES_TAG, |
| "dataset": {"root": args.data_root, "res": train_ds.res, "train_total": train_ds.n_total, |
| "train_dropped_rows": train_ds.n_dropped, |
| "train_used": len(train_ds), "text_table_keys": len(train_ds.text_table), |
| "val_used": len(val_ds) if val_ds is not None else 0, |
| "val_clip_ids": val_ds.clip_ids() if val_ds is not None else [], |
| "geometry": geometry, "steps_per_epoch": epoch_steps(len(dataloader), args.grad_accum), |
| "start_epoch": start_epoch, "skip_micro": skip_micro}, |
| "model": {"dit_files": [os.path.join(args.model_dir, os.path.basename(f)) for f in DIT_FILES], |
| "n_params": n_params, "n_sem_cond_params": n_sem_params, "vocab": args.vocab, |
| "sem_cond_in_channels": enc.in_channels, "export_keys": GC_DIT_KEYS, |
| "param_groups": [{"lr": g["lr"], "n_params": sum(p.numel() for p in g["params"])} for g in groups], |
| "resume_info": resume_info, "start_step": start_step}, |
| "env": {"world_size": accelerator.num_processes, "distributed_type": str(accelerator.distributed_type), |
| "mixed_precision": accelerator.mixed_precision, "torch": torch.__version__, |
| "cuda": torch.version.cuda, "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, |
| "python": sys.executable, "platform": platform.platform(), "hostname": socket.gethostname(), |
| "argv": sys.argv, "started": time.strftime("%Y-%m-%d %H:%M:%S"), |
| "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), |
| "global_batch": accelerator.num_processes * args.batch_size * args.grad_accum}, |
| "deepspeed_config": accelerator.state.deepspeed_plugin.deepspeed_config if use_ds else None, |
| } |
| try: |
| import accelerate, deepspeed, diffsynth |
| cfg["env"]["accelerate"] = accelerate.__version__ |
| cfg["env"]["deepspeed"] = deepspeed.__version__ |
| except Exception: |
| pass |
| with open(os.path.join(args.output, "config.json" if start_step == 0 else f"config_resume{start_step}.json"), |
| "w", encoding="utf-8") as fh: |
| json.dump(cfg, fh, indent=2, ensure_ascii=False, default=str) |
|
|
| |
| train_log = val_log = tb = None |
| if is_main: |
| |
| log_files = {"train_log.csv": ["step", "loss", "mse", "timestep", "lr", "cond_lr", "grad_norm", "sem_rms", "x_rms", |
| "text_drop", "cond_drop", "mask_frac", "step_time", "data_time", "mem_gb", "epoch", |
| "clips"], |
| "val_log.csv": ["step", "timestep", "mse", "mse_uncond", "mse_nulltext", "gap_uncond", "gap_nulltext", "n"]} |
| for name, fields in log_files.items(): |
| old = rotate_csv_if_header_changed(os.path.join(args.output, name), fields) |
| if old: |
| print(f"[log] {name} 的列变了，旧文件改名 {old}", flush=True) |
| train_log = CsvLog(os.path.join(args.output, "train_log.csv"), log_files["train_log.csv"]) |
| val_log = CsvLog(os.path.join(args.output, "val_log.csv"), log_files["val_log.csv"]) |
| from torch.utils.tensorboard import SummaryWriter |
| tb = SummaryWriter(os.path.join(args.output, "tensorboard")) |
|
|
| |
| cpu_gen = torch.Generator().manual_seed(args.seed * 1000 + accelerator.process_index + start_step) |
| val_tids = [module.timestep_id_for(t) for t in args.val_timesteps] |
| null_context = null_ids = None |
| if val_ds is not None and args.val_uncond: |
| null_context, null_ids = train_ds.context_for([NULL_PROMPT] * CELLS) |
| null_context, null_ids = null_context.unsqueeze(0).to(device), null_ids.unsqueeze(0).to(device) |
|
|
| def do_validation(step): |
| if val_ds is None: |
| return |
| tv = time.time() |
| res = run_validation(accelerator, model, module, val_ds, val_tids, args.seed, null_context, null_ids, args.val_uncond) |
| if is_main: |
| for t, r in res.items(): |
| val_log.write({"step": step, "timestep": t, **r}) |
| for k, v in r.items(): |
| if k != "n": |
| tb.add_scalar(f"val/{k}_t{int(round(t))}", v, step) |
| means = {k: sum(r[k] for r in res.values()) / len(res) for k in next(iter(res.values())) if k != "n"} |
| for k, v in means.items(): |
| tb.add_scalar(f"val/{k}_mean", v, step) |
| print(f"[val] step {step} " + " ".join( |
| f"t={t:.0f}: mse={r['mse']:.4f}" + (f" uncond={r['mse_uncond']:.4f} null={r['mse_nulltext']:.4f}" if args.val_uncond else "") |
| + f" (n={r['n']})" for t, r in res.items()) |
| + " | mean " + " ".join(f"{k}={v:.4f}" for k, v in means.items()) + f" {time.time() - tv:.0f}s", flush=True) |
|
|
| data_pos = {"epoch": start_epoch, "epoch_micro": skip_micro} |
|
|
| def do_save(step): |
| path = os.path.join(args.output, f"step-{step}.safetensors") |
| export_dit_weights(accelerator, model, path, export_metadata(args, step)) |
| if args.save_state: |
| save_state_rotating(accelerator, args.output, {"step": step, "arm": "gc25d", "vocab": args.vocab, "res": GEO.RES_TAG, |
| "init_ckpt": args.init_ckpt, **data_pos, **geometry}) |
|
|
| |
| model.train() |
| step = start_step |
| |
| steps_per_epoch = epoch_steps(len(dataloader), args.grad_accum) |
| epoch = start_epoch |
| last_saved = start_step |
| step_times, data_times = [], [] |
| torch.cuda.reset_peak_memory_stats(device) |
| if is_main: |
| print(f"[train] {len(train_ds)} 个 clip，每进程 batch {args.batch_size} × GA {args.grad_accum}，{accelerator.num_processes} 进程" |
| f"（全局 batch {accelerator.num_processes * args.batch_size * args.grad_accum}），每进程每 epoch {len(dataloader)} 个 " |
| f"micro batch = {steps_per_epoch} 步，从 step {start_step}（epoch {epoch}，跳过 {skip_micro} 个 micro）到 {args.max_steps}", |
| flush=True) |
| micro_time = data_time = 0.0 |
| step_clips = [] |
| t_prev = time.time() |
| while step < args.max_steps: |
| train_ds.set_epoch(epoch) |
| if hasattr(dataloader, "set_epoch"): |
| dataloader.set_epoch(epoch) |
| epoch_loader, epoch_micro = dataloader, 0 |
| if skip_micro: |
| |
| epoch_loader, epoch_micro, skip_micro = accelerator.skip_first_batches(dataloader, skip_micro), skip_micro, 0 |
| epoch_loader.set_epoch(epoch) |
| for batch in epoch_loader: |
| if step >= args.max_steps: |
| break |
| t_step = time.time() |
| data_time += t_step - t_prev |
| epoch_micro += 1 |
| step_clips.extend(batch["clip_id"]) |
| batch = to_device(batch, device) |
| tid = int(torch.randint(0, 1000, (1,), generator=cpu_gen)) |
| lrs = lr_scheduler.get_last_lr() |
| |
| |
| with accelerator.accumulate(model): |
| loss, stats = model(batch, tid) |
| accelerator.backward(loss) |
| |
| grad_norm = accelerator.clip_grad_norm_(model.parameters(), args.grad_clip) if accelerator.sync_gradients else None |
| optimizer.step() |
| if accelerator.sync_gradients: |
| lr_scheduler.step() |
| optimizer.zero_grad(set_to_none=True) |
| torch.cuda.synchronize(device) |
| t_prev = time.time() |
| micro_time += t_prev - t_step |
| if not accelerator.sync_gradients: |
| continue |
| step += 1 |
| data_pos = {"epoch": epoch, "epoch_micro": epoch_micro} |
| |
| step_time, step_data_time, micro_time, data_time = micro_time + data_time, data_time, 0.0, 0.0 |
| clips, step_clips = ";".join(step_clips), [] |
| step_times.append(step_time) |
| data_times.append(step_data_time) |
| loss_mean = accelerator.reduce(loss.detach().float(), reduction="mean").item() |
| if not math.isfinite(loss_mean): |
| raise SystemExit(f"step {step}: loss = {loss_mean}（非有限），停训") |
| if is_main and (step % args.log_every == 0 or step == args.max_steps): |
| gn = float(grad_norm) if grad_norm is not None and math.isfinite(float(grad_norm)) else float("nan") |
| dr = batch["dropped"] |
| row = {"step": step, "loss": loss_mean, "mse": stats["mse"].item(), "timestep": stats["timestep"], |
| "lr": lrs[0], "cond_lr": lrs[-1], "grad_norm": gn, |
| "sem_rms": stats["sem_rms"].item(), "x_rms": stats["x_rms"].item(), |
| "text_drop": float(dr["text"].float().mean()), "cond_drop": float(dr["cond"].float().mean()), |
| "mask_frac": float(dr["mask"].float().mean()), "step_time": step_time, "data_time": step_data_time, |
| "mem_gb": torch.cuda.max_memory_allocated(device) / 1024 ** 3, "epoch": epoch, "clips": clips} |
| train_log.write(row) |
| for k in ("loss", "mse", "lr", "cond_lr", "grad_norm", "sem_rms", "x_rms", "step_time", "data_time", "mem_gb"): |
| tb.add_scalar(f"train/{k}", row[k], step) |
| tb.add_scalar("train/sem_over_x", row["sem_rms"] / max(row["x_rms"], 1e-12), step) |
| print(f"step {step}/{args.max_steps} loss {loss_mean:.4f} (mse {row['mse']:.4f} t {row['timestep']:.0f} " |
| f"sem_rms {row['sem_rms']:.3e} x_rms {row['x_rms']:.3f} drop t{int(row['text_drop'])}c{int(row['cond_drop'])}" |
| f"m{row['mask_frac']:.2f}) lr {row['lr']:.2e}/{row['cond_lr']:.2e} gnorm {gn:.3f} " |
| f"{step_time:.2f}s (data {step_data_time:.2f}s) mem {row['mem_gb']:.1f}GB", flush=True) |
|
|
| if args.val_every > 0 and step % args.val_every == 0: |
| do_validation(step) |
| if args.save_every > 0 and step % args.save_every == 0: |
| do_save(step) |
| last_saved = step |
| t_prev = time.time() |
| epoch += 1 |
|
|
| if last_saved != step: |
| do_save(step) |
| peak = torch.tensor(torch.cuda.max_memory_allocated(device) / 1024 ** 3, device=device) |
| peak_max = accelerator.reduce(peak, reduction="max").item() |
| if is_main: |
| tail = step_times[2:] or step_times |
| dtail = data_times[2:] or data_times |
| s_step, s_data = sum(tail) / max(len(tail), 1), sum(dtail) / max(len(dtail), 1) |
| print(f"[done] step {step}，本进程峰值显存 {peak.item():.1f} GB，各进程最大 {peak_max:.1f} GB，" |
| f"s/step {s_step:.2f}（其中等数据 {s_data:.2f}）", flush=True) |
| with open(os.path.join(args.output, "done.json"), "w", encoding="utf-8") as fh: |
| json.dump({"step": step, "start_step": start_step, "peak_mem_gb_main": peak.item(), "peak_mem_gb_max": peak_max, |
| "s_per_step_mean": s_step, "data_s_per_step_mean": s_data, "n_steps_timed": len(tail), |
| "data_pos": data_pos, "finished": time.strftime("%Y-%m-%d %H:%M:%S")}, fh, indent=2) |
| train_log.close() |
| val_log.close() |
| tb.close() |
| accelerator.wait_for_everyone() |
| accelerator.end_training() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
