Diffusers
Safetensors
File size: 25,513 Bytes
74da989
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
# 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
```