Diffusers
Safetensors
HY / train.md
Cccccz's picture
Upload batch 65: 500 files (0.01 GiB)
74da989 verified
|
Raw History Blame Contribute Delete
25.5 kB

Train

实验命名规范

从 2026-07-22 起,checkpoint 目录、独立日志、SwanLab experiment 和评测输出统一显式记录两组 学习率、训练总步数与两种 warmup:

<experiment>_flr<fusion_lr>_blr<blocks_lr>_schedsteps<max_steps>_swarmup<scheduler_warmup>_bwarmup<block_warmup>
  • flr:image fusion/residual 参数组的峰值学习率;
  • blr:两个 Predictor blocks 参数组的峰值学习率;
  • schedsteps:创建 LR scheduler 时使用的训练总步数,不代表某次运行实际走到的步数;
  • swarmup:全局 cosine scheduler 的 warmup steps;
  • bwarmup:blocks 参数组额外冻结/延迟启用的 warmup steps;旧 rollout 没有该独立参数时写 bwarmup-none;
  • ckptstep:仅用于评测输出,表示本次评测实际加载的 checkpoint step。

例如,计划训练 2000 steps、实际使用 step 300 权重进行评测时,训练目录保留 schedsteps2000,评测目录同时写 ckptstep300 与 schedsteps2000。目录内的 train_log.jsonl、 torchrun.log、train_config.json 和 swanlab/ 不重复拼接参数,它们继承父目录的完整命名。

当前主要配置如下:

实验组 flr blr schedsteps swarmup bwarmup
v4 / predictor-disca 离线 Stage 1 1e-4 1e-5 2000 100 100
早期整轨迹 FPPP rollout 1e-5 1e-6 2000 50 none
overfit8 FPPP rollout 1e-4 1e-5 2000 8 8
overfit8 FPPF rollout 1e-4 1e-5 300 8 8
旧 timestep-only sweep 1e-4 1e-5 200 8 8
detached timestep-only 1e-4 1e-5 1000 8 8

1. 监督目标

当前 v4 数据为每个 chunk 保存四步 Full-DiT Teacher hidden 和 velocity。正式 v4 与 predictor-disca 使用全部相邻监督 pair:

step 0 final hidden + step 1 input -> step 1 hidden / velocity
step 1 final hidden + step 2 input -> step 2 hidden / velocity
step 2 final hidden + step 3 input -> step 3 hidden / velocity

v1/v2/v3 的历史训练只使用 0→1 和 2→3,对应 F-P-F-P;从 v4 开始加入 1→2,使 Predictor 具备连续替代和任意选择后三步的能力。

2. Loss

所有版本使用相同的监督损失:

hidden_loss = mse(pred_hidden.float(), target_hidden.float())
velocity_loss = mse(pred_velocity.float(), target_velocity.float())
loss = 0.1 * hidden_loss + 1.0 * velocity_loss
  • hidden 和 velocity 的 forward 使用 BF16;
  • MSE 转为 FP32 计算;
  • velocity MSE 是主目标;
  • hidden MSE 是 final-layer 前 feature 的辅助约束;
  • frozen final layer 只冻结参数,forward 不能放进 no_grad(),否则 velocity loss 无法回传到 Predictor hidden。

3. 优化方法

当前正式训练使用 AdamW:

参数组 学习率
image fusion + residual head 1e-4
两个 Predictor blocks 1e-5

共同配置:

betas: [0.9, 0.95]
weight_decay: 0.01
grad_clip: 1.0
precision: bfloat16
loss_precision: float32
max_steps: 2000
scheduler: cosine
scheduler_warmup_steps: 100
block_warmup_steps: 100
gradient_checkpointing: true

前 100 optimizer steps 中 Predictor block 参数组的学习率为 0,只训练 fusion/residual;之后 block 组按较低学习率参与训练。两个参数组都经过 100-step warmup 和 cosine decay,step 2000 衰减到 0。

4. 各版本训练方式

v1

  • 入口:train_predictor.py;
  • 单 GPU、micro-batch 1;
  • gradient accumulation 4,有效 batch 4;
  • 2000 steps;
  • pairs:0→1、2→3;
  • source blocks:[1,52];
  • 每 500 steps 保存,默认保留 2 个快照。

旧机器上的 v1 实验结果不再纳入当前文档,这里只保留训练方法。

v2/v3

  • 入口:train_predictor_v2.py --predictor_version v2|v3;
  • 单 GPU,默认 batch 1、gradient accumulation 4;
  • pairs:0→1、2→3;
  • source blocks:[0,53] 或 [1,52];
  • 2000 steps、100-step block warmup;
  • SwanLab 可选。

v3 Dataset 额外加载前一 chunk 的 target-step final hidden,并排除 chunk 0。迁移后的当前目录未保留 v2/v3 正式 checkpoint,因此这里只记录方法,不将缺失产物写成可复现实验结果。

v4

  • 入口:train_predictor_v4.py;
  • 启动:tools/launch_predictor_v4_training.sh;
  • 8 卡 DDP;
  • 每卡 batch size 16,全局 batch size 128;
  • gradient accumulation 1;
  • 2800 records × 3 transitions = 8400 pairs;
  • 按 Context frames 分桶,减少动态 padding;
  • 2000 steps;
  • 每 100 steps 保存,保留 20 个 safetensors;
  • SwanLab 仅由 rank 0 上报。

已训练 blocks [0,53] 和 [1,52] 两版。

predictor-disca

  • 入口:train_predictor_disca.py;
  • 启动:tools/launch_predictor_disca_training.sh;
  • 8 卡 DDP、每卡 batch size 16、全局 batch size 128;
  • 8400 pairs、2000 steps;
  • 不加载 Context pre-feature 或 previous-chunk hidden;
  • 不按 Context frames 分桶,使用普通 DistributedSampler;
  • 每 100 steps 保存,保留 20 个 safetensors;
  • 当前正式版使用 blocks [0,53]。

5. 当前正式训练结果

下表是训练日志在 step 2000 记录窗口内的均值,不是完整数据集 validation loss:

版本 blocks total loss hidden MSE velocity MSE peak memory/GPU avg micro time
v4 [0,53] 0.011895 0.011744 0.010721 62.82 GiB 4.47 s
v4 [1,52] 0.011537 0.011732 0.010364 62.82 GiB 4.49 s
predictor-disca [0,53] 0.011053 0.011121 0.009941 31.18 GiB 1.66 s

训练 loss 不能直接替代 rollout 视频质量。predictor-disca 的训练 loss 更低,但其训练集视频指标明显 低于 v4,说明缺失 AR Context 后更容易拟合单步监督,却在连续在线 rollout 中累积了更大误差。

6. Checkpoint 与 resume

正式训练同时保存:

predictor_step_XXXXX.safetensors   # 仅模型推理权重
training_latest.pt                 # model、optimizer、scheduler、step、epoch、配置
train_config.json
train_log.jsonl
torchrun.log
swanlab/

resume 会恢复 model、optimizer、scheduler、global/micro step 和 epoch,并检查 source blocks 与 Predictor variant 是否匹配。例如:

python train_predictor_disca.py \
  --resume checkpoints/predictor_disca_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/training_latest.pt

正式云端监控使用:

export SWANLAB_API_KEY='<token>'

token 只放在进程环境或用户登录缓存中,不写入仓库、配置或日志。

7. 在线训练集评测

统一评测 100 个训练 case × 4 actions,共 400 个 125-frame 视频:

reference: Full-DiT F-F-F-F
reuse:     Full steps 0、3;steps 1、2 reuse
predictor: Full steps 0、3;steps 1、2 Predictor

视频指标包括 full-resolution PSNR、frame-mean PSNR、SSIM、AlexNet LPIPS、pixel MSE/RMSE/MAE 和 temporal-delta MSE。

方法 blocks PSNR SSIM LPIPS temporal delta MSE
Direct Reuse — 18.4639 0.6429 0.2083 0.01212
v4 [0,53] 21.9337 0.7606 0.1204 0.00865
v4 [1,52] 22.1090 0.7670 0.1168 0.00832
predictor-disca [0,53] 19.4084 0.6748 0.1816 0.01061

这些是 train-set/in-sample 指标,只能比较架构在已训练图文 pair 上的拟合和 rollout 行为,不能证明 对新图、新 prompt、新 seed 或新动作的泛化。validation/test 应使用 data.md 中互斥的 25/50 pair holdout。

8. v4 递归 F-P-P-P Self-Rollout 微调方案

8.1 定位与实现状态

这是一阶段离线 Teacher-forcing 训练之后的二阶段 rollout 微调。初始化应加载已有 v4 Predictor 权重,而不是重新从两个 Teacher blocks 开始训练。当前训练集视频结果中 [1,52] 最好,因此正式 实验优先从以下离线 v4 FP32-master 训练状态中的 model 初始化;optimizer/scheduler 不继承, rollout 使用全新的优化状态:

checkpoints/predictor_v4_prefeature_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/training_latest.pt

[0,53] 可以用同一套流程作为独立对照,但 source blocks 必须与初始 checkpoint 一致。rollout 训练已作为独立入口实现,不改变现有离线 train_predictor_v4.py:

train_predictor_v4_rollout.py
predictor_training/trajectory_dataset.py
predictor_training/rollout_utils.py
tools/launch_predictor_v4_rollout_training.sh

低学习率 [1,52] 基线输出到以下目录,并已按实验要求在 step 600 停止:

checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-5_blr1e-6_schedsteps2000_swarmup50_bwarmup-none_swanlab

使用第一阶段学习率的正式入口改用明确 LR 命名,避免用含义模糊的 stage1lr:

checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab

8.2 目标轨迹

每条 item 是一个完整的 (case_id, action_id) trajectory,而不是一个独立 transition:

Chunk 0:F-F-F-F,只建立 self-generated history 和 previous-chunk hidden,不训练
Chunk 1~7:F-P-P-P,三个 Predictor step 都计算 loss

在训练 chunk 内:

initial x0
  -> Full step 0: h0, v0
  -> Euler update: x1
  -> Predictor step 1: pred_h1, pred_v1
  -> differentiable Euler update: x2
  -> Predictor step 2: pred_h2, pred_v2
  -> differentiable Euler update: x3
  -> Predictor step 3: pred_h3, pred_v3
  -> differentiable Euler update: final chunk latent

同一 chunk 内 pred_h1 → pred_h2 → pred_h3 是递归 same-chunk anchor,三个 Predictor 调用保留 一张完整计算图。chunk 之间则同时 detach:

  • final chunk latent detach 后加入 history;
  • pred_h1/pred_h2/pred_h3 detach 后作为下一 chunk 的同 step previous-chunk hidden。

因此,v4 的 previous-chunk 输入不能只用最终 latent 代替。Chunk 0 虽不训练,也必须保存 Full-DiT step 1、2、3 的 final hidden,分别作为 Chunk 1 Predictor step 1、2、3 的 previous-chunk hidden。

8.3 轻量样本表与全在线构建

Self-rollout 不依赖之前构建的 F-F-F-F tensor 数据集。Dataset 只保存一张轻量样本表,每个 item 包含:

case_id, source_image, caption, action_id, pose, seed

100 个图文 pair × 4 actions 组成 400 条 trajectory。__getitem__ 不读取 step tensor、Teacher hidden/velocity、Context pre-feature 或 K/V;每个 rank 拿到样本描述后,使用当前正式 WorldPlay pipeline 在线完成:

  1. 按与数据构建/评测相同的规则将输入图像 resize/crop 到 480×832;
  2. 使用 pose_to_input(pose, 32) 生成完整 8-chunk viewmats、Ks 和 action;
  3. 用指定 seed 调用 prepare_latents,一次生成 [1,32,32,30,52] 初始噪声,再按 4 latent frames 切成 8 个 chunk;
  4. 从现有 case tensor 读取静态 image_condition_latent,在线构造 condition/mask,避免重复 VAE encode;
  5. 从离线数据集的 text_kv_exact_all54_bf16/ 读取全部 54 层精确 BF16 Text K/V,并按 text_valid_mask 去掉磁盘 padding;
  6. 从 Chunk 0 开始在线执行整条 F-F-F-F/F-P-P-P trajectory,并在线生成所有监督 target。

每一步仍按当前 pipeline 语义构造 65 通道输入:32 通道 current noisy latent 加 32 通道 image condition 和 1 通道 mask。只有 Chunk 0 的第一个 latent frame 含 image condition;Chunk 1~7 的 33 个 condition 通道为零。history prefill 同样拼接对应历史帧的固定 condition,不能只把 32 通道 self latent 送入 Teacher。

因此 rollout 训练不读取或生成持久化的 step_1/2/3_noisy_sample、Teacher target、Context pre-feature、previous-step hidden 或全层 Text K/V sidecar。Text K/V 与 Predictor policy 和 trajectory 状态无关,按 case 存在主离线数据集内;所有 policy-dependent 状态仍在当前参数、当前 self history 上 产生,语义是严格 on-policy。图像/caption/action 样本表不是 Teacher trajectory 数据,可以直接从 当前 100 个训练 pair 的来源清单生成;留出集不参与训练。

静态 Text K/V 构建时来自 prompt、vision 和 ByT5 编码,训练主循环不再加载这三个编码器,也不再 重复 text prefill。

8.4 v4 Context 的准确语义

每个 chunk 使用当前 detached、self-generated history 运行与正式推理相同的 select_aligned_memory_frames。选择结果只由相机轨迹决定;被选中的 latent 内容必须来自当前 self history。

随后冻结 Full Teacher 对这些 history latent 执行一次 54-layer AR vision prefill,得到精确 Context Vision K/V。该 cache:

  • 与当前 trajectory 的全层 Text K/V 合并;
  • 在当前 chunk 的四个 denoise steps 间只读共享;
  • Teacher target 使用全部 54 层 cache;
  • v4 Predictor 直接读取其 source blocks 对应的同一份 exact K/V;
  • 全部 detach,且不对 Context 建图。

这里不走 context_prefeature -> frozen projector -> K/V 路径。pre-feature 是离线 pair 训练的磁盘 压缩方案;在线 rollout 已经由 Teacher prefill 持有 exact K/V,行为应与当前推理 pipeline 一致。

8.5 单个 chunk 的 Teacher 与 Predictor 顺序

Chunk 0 在 torch.no_grad() 下执行四次 Full-DiT,并保存 final latent 和 Full hidden-by-step。 Chunk 1~7 的顺序固定为:

  1. 从 detached history 选择 Context,并由 Teacher 生成 exact AR K/V;
  2. Full Teacher 在 step 0 输入上产生 h0/v0,用正式 Euler scheduler 更新到 x1;
  3. 对 step 1~3,先在 current_student_latent.detach() 上运行 Full Teacher,获得 target_hidden/target_velocity;
  4. Predictor 使用相同 noisy latent、condition、相机参数和 Context K/V;
  5. step 1 的 anchor 是 h0,step 2 是 pred_h1,step 3 是 pred_h2;
  6. 每个 Predictor velocity 经可微 scheduler 更新到下一 latent;
  7. 三步 loss 求均值,当前 chunk 立即 backward;
  8. final latent 与三步 pred hidden detach 后加入下一 chunk 状态。

Teacher target 必须来自 Predictor 真正访问的状态。对应当前 Transformer API 的核心调用是:

with torch.no_grad():
    teacher_model_input = torch.cat(
        [current_student_latent.detach(), current_condition_33ch], dim=1
    )
    teacher_output = full_teacher(
        bi_inference=False,
        ar_txt_inference=False,
        ar_vision_inference=True,
        hidden_states=teacher_model_input,
        timestep=current_timestep,
        timestep_r=None,
        mask_type=task_type,
        action=current_action,
        viewmats=current_viewmats,
        Ks=current_Ks,
        kv_cache=exact_context_kv,
        cache_vision=False,
        rope_temporal_size=rope_temporal_size,
        start_rope_start_idx=start_rope_start_idx,
        return_final_hidden=True,
    )
    target_velocity = teacher_output[0]
    target_hidden = teacher_output[2]

不能读取原 F-F-F-F trajectory 中相同步的 target。Teacher 只用于构造 stop-gradient label,其输出 立即释放。Predictor 使用的 target_frame_condition 继续由冻结 Teacher 的 time_in(timestep) + action_in(action) 计算,不从旧 target 文件推断。

8.6 完整 timestep BPTT 的代码前置条件

当前 HYWorldPlayPredictorV2.forward(v4 继承该实现)使用 torch.no_grad() 包裹冻结 img_in。 这会阻断 x2/x3 -> img_in -> Predictor 的输入梯度,不满足完整 timestep BPTT。rollout 实现前必须 改为:

# img_in 参数保持 requires_grad=False,但 forward 不进入 no_grad。
current_img = self.img_in(target_model_input)

冻结参数不会产生 parameter gradient;当离线训练输入不需要梯度时也不会额外建图,但 rollout 的 step 2、3 可以对前一步 scheduler 输出求导。冻结 final layer 同理:参数不训练,forward 仍必须保留 对 pred_hidden 的梯度。Context projector 保持 no_grad(),因为 Context 本来就应 detach。

8.7 可微 FlowMatch Euler update

当前项目的 FlowMatchDiscreteScheduler.step 在 solver=euler 时核心公式已经是 out-of-place 且可微:

x_next = x.float() + pred_velocity.float() * (
    sigmas[step_index + 1] - sigmas[step_index]
)

但正式 pipeline 会把结果原位写回共享 latents,不应直接复制该容器逻辑。训练应实现无状态的 differentiable_euler_step:调用相同 set_timesteps(4) 后读取同一组 sigmas,每次只接收和返回当前 chunk local tensor,无 detach、no_grad 或 inplace slice assignment。其 FP32 update 和输出 dtype 必须与正式 scheduler 数值一致。

step 2、3 的 loss 应同时通过两类路径影响前面的 Predictor 调用:

hidden anchor path: pred_h1 -> Predictor step 2 -> pred_h2 -> Predictor step 3
latent path:        pred_v1 -> x2 -> Predictor step 2 -> pred_v2 -> x3 -> Predictor step 3

8.8 Loss 与归一化

每个 Predictor step 继续使用 v4 的 FP32 MSE:

step_loss = (
    mse(pred_velocity.float(), target_velocity.float())
    + 0.1 * mse(pred_hidden.float(), target_hidden.float())
)
chunk_loss = (step1_loss + step2_loss + step3_loss) / 3

一条 trajectory 有七个训练 chunk,所以每个 chunk 执行:

(chunk_loss / 7).backward()

不要再除以 world size,DDP 会对 8 个 rank 自动求平均。若以后额外累计多条 trajectory 才更新一次, 还需要再除以 gradient_accumulation_trajectories。日志应分别记录 step 1/2/3 hidden MSE、velocity MSE、chunk id 均值和整条 trajectory loss,不能只记录合并 total loss。

8.9 8 卡 DDP 与同步边界

每个 rank 持有一份冻结的 54-layer Teacher、一份 DDP v4 Predictor 和一条独立 trajectory。只包装 Predictor:

predictor = DDP(
    predictor,
    device_ids=[local_rank],
    output_device=local_rank,
    broadcast_buffers=False,
    gradient_as_bucket_view=True,
    find_unused_parameters=False,
)

400 条 trajectory 可被 8 个 rank 整除。DistributedSampler(..., drop_last=True) 每个 epoch 给每个 rank 50 条,调用 sampler.set_epoch(epoch) 保证 shuffle 一致且不同 rank 不重叠。micro-batch 固定为 每卡一条 trajectory,因此每个 optimizer step 的全局 batch 是 8 trajectories。

Chunk 1~6 只累计本地梯度,Chunk 7 backward 时执行一次 all-reduce:

optimizer.zero_grad(set_to_none=True)

for chunk_id in range(1, 8):
    sync_ctx = nullcontext() if chunk_id == 7 else predictor.no_sync()
    with sync_ctx:  # 必须同时覆盖三个 DDP forward 和当前 backward
        chunk_loss, final_latent, hidden_by_step = train_fppp_chunk(...)
        (chunk_loss / 7).backward()

    history.append(final_latent.detach())
    previous_hidden_by_step = {
        step: hidden.detach() for step, hidden in hidden_by_step.items()
    }
    del chunk_loss, final_latent, hidden_by_step

clip_grad_norm_(predictor.parameters(), 1.0)
optimizer.step()
lr_scheduler.step()

七个 chunk 中间不能 optimizer.step() 或 zero_grad()。现有 Predictor checkpoint 使用 use_reentrant=False,应继续保留;三个 DDP forward 合并一次 backward 以及七次 chunk backward 的 组合必须先通过单卡/DDP 小规模梯度一致性测试。

8.10 显存与训练配置

每个 rank 加载 Transformer Teacher 和 Predictor。image-condition latent 直接从 case tensor 读取; 全层 Text K/V 从主离线数据集读取并去掉 padding。训练不加载 prompt encoder、vision encoder、 ByT5、VAE、SR 或视频编码组件。可微 rollout 阶段只让 Teacher、Predictor、当前 Text/Context cache 和当前 chunk graph 驻留 GPU。

Teacher 全程 eval、requires_grad_(False) 和 no_grad();Predictor forward activation 使用 BF16 autocast,所有可训练 Predictor 参数及 AdamW moments 保持 FP32,loss 使用 FP32,并开启现有 non-reentrant gradient checkpointing。冻结参数和 Teacher 权重仍为 BF16。每个 chunk backward 后 立即释放三步图,history、previous hidden 和 Context K/V 全部 detach。

checkpoint 分成两种精度语义:training_latest.pt 保留 FP32 trainable master 参数、FP32 optimizer 状态和 scheduler,用于无损 resume;predictor_step_XXXXX.safetensors 显式转换为 BF16,只用于推理 和评测。resume 时会检查 checkpoint 的 FP32-master 配置,禁止把旧 BF16 rollout optimizer 静默载入 当前训练。

建议从已有 v4 权重开始一个全新的 optimizer/scheduler,不加载 Teacher-forcing training_latest.pt 中的 optimizer 状态。rollout checkpoint 自身仍保存 model、optimizer、scheduler、 global optimizer step、epoch 和 sampler/trajectory 进度,并每 100 optimizer steps 保存一次。

低学习率基线采用 fusion/residual 1e-5、blocks 1e-6 和 50-step warmup,运行至 step 600 后, 完整 epoch mean loss 只下降约 3%,因此停止。后续正式配置改为与第一阶段相同的 LR 策略: fusion/residual 峰值 1e-4,blocks 峰值 1e-5;fusion 前 100 steps warmup,blocks 前 100 steps 学习率为 0,之后两组共同 cosine decay 到 step 2000。这里的一个 step 指完整 8-rank trajectory 更新,不是一个 chunk。

每 100 steps 保存一次 checkpoint。训练入口持续断言每步 8 个 rank 的 trajectory 主键互异,并检查 step 2/3 中间 latent 的梯度非零。

正式启动前的 8 卡完整 8-chunk smoke test 已通过:fresh step 1 的 loss 为 0.013256, latent_grad_l1_min=1.63e-4,峰值显存 49.13 GiB/GPU;从该 checkpoint resume 到 step 2 后, latent_grad_l1_min=2.98e-4,峰值显存 50.60 GiB/GPU。训练 checkpoint 中 48 个可训练参数张量及 Adam moments 均为 FP32,推理 safetensors 中 80 个张量均为 BF16。

为区分学习率问题与梯度/模型容量问题,另固定随机种子 20260721 从 100 个训练 case 中选择 [4,11,17,28,58,80,94,96],每个 case 保留四种 action,共 32 条 trajectory。8 卡每个 epoch 恰好 4 steps;临时将 fusion warmup 和 block freeze 都缩短为 8 steps,LR 仍为 1e-4/1e-5。 截至 epoch 5,完整 overfit mean loss 从 0.01969 降至 0.01836,下降约 6.8%,证明递归梯度链路 和优化器能够学习。该实验输出到:

checkpoints/predictor_v4_rollout_fppp_blocks1-52_overfit8_seed20260721_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup8_bwarmup8_swanlab

8.11 分阶段验收

实现必须依次通过:

  1. Dataset 只有 400 条互异的 (case_id, action_id) 元数据,8 个 rank 每步读取不同主键;
  2. Chunk 0 全部由 Full Teacher 计算、无参数梯度,并产生 step 1/2/3 previous hidden;
  3. step 2 anchor 与 pred_h1、step 3 anchor 与 pred_h2 数值和对象来源正确;
  4. 每个 Teacher target 的输入等于当前 student latent 的 detached copy,不读取任何离线 FFFF 状态;
  5. self history 生成的 Context latent 正确,selected indices 与正式推理函数一致;
  6. 训练 Euler update 与正式 FlowMatchDiscreteScheduler.step 在四个 step 上误差满足 FP32 容差;
  7. 只保留 step 3 loss 时,x2/x3.grad 非零,且前两次 Predictor 调用对参数有非零贡献;
  8. Chunk 2 loss 不会回传到 Chunk 1 final latent 或 previous hidden;
  9. Chunk 1~6 不触发 all-reduce,Chunk 7 每个 optimizer step 只同步一次;
  10. 单卡小样本中,分 chunk backward 与保留七个 chunk 后一次 backward 的梯度在容差内一致;
  11. 先完成 1-trajectory forward、1-step backward、8 卡 2-step DDP smoke test,再做 20-step overfit;
  12. 无 NaN/Inf,step 1/2/3 loss 均能下降,显存峰值有日志记录。

最终评测必须使用 Chunk 0 F-F-F-F、Chunk 1~7 F-P-P-P,即 replace_steps=(1,2,3)。在同一训练集和 25/50 holdout 上对比:Full-DiT、Direct Reuse、rollout 微调前 v4 checkpoint 和微调后 checkpoint,并同时报告视频质量、端到端耗时和组件耗时。只有留出集 F-P-P-P 优于微调前权重,才能说明 self-rollout 微调带来泛化收益。

9. 训练入口速查

v4:

bash tools/launch_predictor_v4_training.sh

v4 F-P-P-P on-policy rollout:

bash tools/launch_predictor_v4_rollout_training.sh

resume:

bash tools/launch_predictor_v4_rollout_training.sh \
  checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab \
  checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/training_latest.pt

8-case overfit:

bash tools/launch_predictor_v4_rollout_overfit8_training.sh

predictor-disca [1,52]:

bash tools/launch_predictor_disca_training.sh

predictor-disca [0,53]:

bash tools/launch_predictor_disca_training.sh \
  checkpoints/predictor_disca_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab \
  0,53