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
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_h3detach 后作为下一 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 在线完成:
- 按与数据构建/评测相同的规则将输入图像 resize/crop 到 480×832;
- 使用
pose_to_input(pose, 32)生成完整 8-chunk viewmats、Ks 和 action; - 用指定 seed 调用
prepare_latents,一次生成[1,32,32,30,52]初始噪声,再按 4 latent frames 切成 8 个 chunk; - 从现有 case tensor 读取静态
image_condition_latent,在线构造 condition/mask,避免重复 VAE encode; - 从离线数据集的
text_kv_exact_all54_bf16/读取全部 54 层精确 BF16 Text K/V,并按text_valid_mask去掉磁盘 padding; - 从 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 的顺序固定为:
- 从 detached history 选择 Context,并由 Teacher 生成 exact AR K/V;
- Full Teacher 在 step 0 输入上产生
h0/v0,用正式 Euler scheduler 更新到x1; - 对 step 1~3,先在
current_student_latent.detach()上运行 Full Teacher,获得target_hidden/target_velocity; - Predictor 使用相同 noisy latent、condition、相机参数和 Context K/V;
- step 1 的 anchor 是
h0,step 2 是pred_h1,step 3 是pred_h2; - 每个 Predictor velocity 经可微 scheduler 更新到下一 latent;
- 三步 loss 求均值,当前 chunk 立即 backward;
- 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 分阶段验收
实现必须依次通过:
- Dataset 只有 400 条互异的
(case_id, action_id)元数据,8 个 rank 每步读取不同主键; - Chunk 0 全部由 Full Teacher 计算、无参数梯度,并产生 step 1/2/3 previous hidden;
- step 2 anchor 与
pred_h1、step 3 anchor 与pred_h2数值和对象来源正确; - 每个 Teacher target 的输入等于当前 student latent 的 detached copy,不读取任何离线 FFFF 状态;
- self history 生成的 Context latent 正确,selected indices 与正式推理函数一致;
- 训练 Euler update 与正式
FlowMatchDiscreteScheduler.step在四个 step 上误差满足 FP32 容差; - 只保留 step 3 loss 时,
x2/x3.grad非零,且前两次 Predictor 调用对参数有非零贡献; - Chunk 2 loss 不会回传到 Chunk 1 final latent 或 previous hidden;
- Chunk 1~6 不触发 all-reduce,Chunk 7 每个 optimizer step 只同步一次;
- 单卡小样本中,分 chunk backward 与保留七个 chunk 后一次 backward 的梯度在容差内一致;
- 先完成 1-trajectory forward、1-step backward、8 卡 2-step DDP smoke test,再做 20-step overfit;
- 无 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