Instructions to use Cccccz/HY with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Diffusers
How to use Cccccz/HY with Diffusers:
pip install -U diffusers transformers accelerate
import torch from diffusers import DiffusionPipeline # switch to "mps" for apple devices pipe = DiffusionPipeline.from_pretrained("Cccccz/HY", dtype=torch.bfloat16, device_map="cuda") prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" image = pipe(prompt).images[0] - Notebooks
- Google Colab
- Kaggle
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
```
|