训练1.3B模型问题

[Mikexu007]

大佬您好，我想全训练流程跑一下1.3B的模型。

1. 请问第一阶段直接将stage_1_init.yaml的参数“transformer_model_name_or_path“ 中的值改成“Wan-AI/Wan2.1-T2V-1.3B-Diffusers”就可以跑了吗？还需要做其他改变吗？
2. stage_1_init.yaml以及其他训练流程可否兼容Wan2.2？比如我想跑 Wan-AI/Wan2.2-TI2V-5B，不兼容的话要做什么改变？

[SHYuanBest]

@Mikexu007 感谢关注!

1. 是的，应该不需要改别的，因为wan2.1 14B和wan2.1 1.3B基本是一样的，只不过层数不一样（如果没记错
2. 暂时没太兼容，需要小改一下代码，主要集中在transformer和pipeline部分。具体参考如下链接：

https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py
https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py
(可选) https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan_i2v.py

[Mikexu007]

@SHYuanBest 大佬，我只有8张H100。我琢磨了一下”Training Workflow: stage_1_init.yaml -> stage_1_post.yaml -> stage_2_init.yaml -> stage_2_post.yaml -> stage_3_ode.yaml -> stage_3_post.yaml“， 我将stage_1_init.yaml的参数“transformer_model_name_or_path“ 中的值改成“Wan-AI/Wan2.1-T2V-1.3B-Diffusers”。但是发现后续的stage_1_post、stage_2_init ...stage_3_post，好像也没有用到stage_1_init的输出，所以除了stage_1_init训练阶段消耗GPU比较小，其他都是需要消耗很多GPU？
以下是claude code的回答
”- 每个 config 自带不同的 output_dir：

• ablation_stage_1_init
• ablation_stage_1_post
• ablation_stage_2_init
• ablation_stage_2_post
• ablation_stage_3_ode
• ablation_stage_3_post
• train_helios.py 会根据 config 中的 pretrained_model_name_or_path / transformer_model_name_or_path 直接加载指定模型
• 这些 config 没有设置：
  • load_checkpoints_custom: true
  • load_model_path: ...
  • load_dcp: true
• 所以它们不会自动加载前一个阶段你刚训练出来的 checkpoint

但是有一点需要区分：

• 这些 config 不是完全随机的独立训练，它们分别使用了不同的预训练起点：
  • stage_1_init 用的是 Wan2.1 Transformer
  • stage_1_post 用的是 Helios-Base/transformer_init
  • stage_2_post 用的是 Helios-Mid/transformer_init
  • stage_3_post 用的是 Helios-Distilled/transformer_ode
• 也就是说，它们更像是不同阶段/不同起点的实验组合，而不是“直接把上一阶段训练结果接着当成下一阶段输入”。“

[SHYuanBest]

但是发现后续的stage_1_post、stage_2_init ...stage_3_post，好像也没有用到stage_1_init的输出，所以除了stage_1_init训练阶段消耗GPU比较小，其他都是需要消耗很多GPU？ 以下是claude code的回答 ”- 每个 config 自带不同的 output_dir：

通过使用这个代码https://github.com/PKU-YuanGroup/Helios/blob/main/tools/merge_lora_for_helios.py 去merge lora，即可得到下一阶段的transformer_path。以训练14B的stage-1-init为例，stage-1-post的transformer（Helios-Base）是通过merge "stage-1-init的lora"+"wan-14B"得到的。

可以参考这里#46

[Mikexu007]

但是发现后续的stage_1_post、stage_2_init ...stage_3_post，好像也没有用到stage_1_init的输出，所以除了stage_1_init训练阶段消耗GPU比较小，其他都是需要消耗很多GPU？ 以下是claude code的回答 ”- 每个 config 自带不同的 output_dir：

通过使用这个代码https://github.com/PKU-YuanGroup/Helios/blob/main/tools/merge_lora_for_helios.py 去merge lora，即可得到下一阶段的transformer_path。以训练14B的stage-1-init为例，stage-1-post的transformer（Helios-Base）是通过merge "stage-1-init的lora"+"wan-14B"得到的。

可以参考这里#46

感谢巨佬回答！😆🎈🧧

[SHYuanBest]

transformer_init这个权重对应的是stage-1-init的merge权重