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:
```text
<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:
```text
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
所有版本使用相同的监督损失:
```python
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 |
共同配置:
```yaml
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
正式训练同时保存:
```text
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 是否匹配。例如:
```bash
python train_predictor_disca.py \
--resume checkpoints/predictor_disca_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab/training_latest.pt
```
正式云端监控使用:
```bash
export SWANLAB_API_KEY='<token>'
```
token 只放在进程环境或用户登录缓存中,不写入仓库、配置或日志。
## 7. 在线训练集评测
统一评测 100 个训练 case × 4 actions,共 400 个 125-frame 视频:
```text
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 使用全新的优化状态:
```text
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`:
```text
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 停止:
```text
checkpoints/predictor_v4_rollout_fppp_blocks1-52_fp32master_8gpu_flr1e-5_blr1e-6_schedsteps2000_swarmup50_bwarmup-none_swanlab
```
使用第一阶段学习率的正式入口改用明确 LR 命名,避免用含义模糊的 `stage1lr`:
```text
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:
```text
Chunk 0:F-F-F-F,只建立 self-generated history 和 previous-chunk hidden,不训练
Chunk 1~7:F-P-P-P,三个 Predictor step 都计算 loss
```
在训练 chunk 内:
```text
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
包含:
```text
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 的核心调用是:
```python
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 实现前必须
改为:
```python
# 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 且可微:
```python
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 调用:
```text
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:
```python
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 执行:
```python
(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:
```python
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:
```python
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%,证明递归梯度链路
和优化器能够学习。该实验输出到:
```text
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
bash tools/launch_predictor_v4_training.sh
```
v4 F-P-P-P on-policy rollout:
```bash
bash tools/launch_predictor_v4_rollout_training.sh
```
resume:
```bash
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
bash tools/launch_predictor_v4_rollout_overfit8_training.sh
```
predictor-disca `[1,52]`:
```bash
bash tools/launch_predictor_disca_training.sh
```
predictor-disca `[0,53]`:
```bash
bash tools/launch_predictor_disca_training.sh \
checkpoints/predictor_disca_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab \
0,53
```