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
|
Download train.md from Cccccz/HY: direct link, hf CLI and curl.
- Browser
- Download file 25.5 kB
-
https://huggingface.co/Cccccz/HY/resolve/main/train.md
- Command line
-
hf download hf://Cccccz/HY/train.md
-
curl -L -o train.md https://huggingface.co/Cccccz/HY/resolve/main/train.md
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 | |
| ``` | |