变更摘要

本 PR 为 MindSpeed-MM 的 FSDP2 路径新增 Wan2.2 模型的完整建模支持，对原有的 modeling_wan2_2.py 进行了大规模重构，并将文本编码器和 VAE 组件拆分为独立模块。核心改动包括：抽取 _SimpleTextEncoder 与 build_text_encoder 至独立的 text_encoder.py；新增 vae.py 提供 WanVideoVAE 的分块编解码实现；在 WanModel 中引入 fp32_modulation / fp32_calculate 双精度路径、RoPE 频率预计算优化、以及文本填充掩码支持；同时清理了 WanSelfAttention / WanAttentionBlock 中的冗余参数，并修复了 NPU 上的混合精度和上下文并行通信问题。

主要改动

• 文本编码器模块化：将 _SimpleTextEncoder 和 build_text_encoder 从 modeling_wan2_2.py 中移出，新建 text_encoder.py。_SimpleTextEncoder 现在继承 BaseModel，支持权重加载、设备间迁移（CPU→NPU），并新增 post_load 方法处理 UMT5 的嵌入层绑定。
• 新增 VAE 模块：新建 vae.py，包含 DiffusersAEModel 基类与 WanVideoVAE。WanVideoVAE 实现带边界融合的分块编解码（tiled_encode / tiled_decode），支持两种潜在空间归一化模式（value_shift_scale / channel_specified_shift_scale），并直接通过 diffusers 加载权重。
• 引入双精度调制路径：在 WanAttentionBlock 和 Head 中新增 fp32_modulation 和 fp32_calculate 参数，提供三种前向计算模式：全 fp32 计算路径、仅调制参数保持 fp32 路径、以及原有的自动混合精度路径，以匹配参考实现的数值精度。
• RoPE 频率预计算优化：新增 precompute_rope_freqs_i 函数，在 WanModel.forward 中一次性为整个批次预计算 3D RoPE 频率并传入各注意力层，避免每层重复构建。rope_apply 相应简化，直接使用预展开的 (batch, seq, 1, c) 频率张量；rope_params 返回类型改为 complex64。
• 上下文并行通信修复与参数清理：Ring CP 路径的 gather 操作改用 load_balanced_gather_forward_split_backward；Ulysses 路径增加 cal_split_sizes 支持非对齐序列长度。移除了 WanSelfAttention、WanAttentionBlock 中的 layer_idx/num_layers 参数，以及 WanDiTFSDP2 中的 qk_norm_type、max_seq_len、seperated_timestep 等未使用或冗余的配置项。

代码审查

审查总结

审查覆盖的三个文件

按优先级统计

• P0: 1 个 — WanFlowMatchScheduler 导入路径指向不存在的 diffusion 模块，运行时必崩溃
• P1: 1 个 — WanCrossAttention 丢弃 context_lens，导致 cross-attention 对 padding 位置产生错误注意力
• P2: 1 个 — vae.py 中 decode 方法在 tiling 模式下调用 self.model.tiled_decode 而非 self.tiled_decode，与 encode 不一致
• P3: 2 个 — timestep dtype 对齐移除（低风险不一致）、autocast 作用域不完整（无当前功能影响）

总体风险评估

高风险。P0 导入路径错误会导致 get_wan_flow_match_scheduler() 首次调用时直接崩溃，阻塞所有依赖该调度器的推理/训练流程。P1 的 cross-attention bug 会静默地降低生成质量（当 batch 内文本长度不一致时），难以察觉但影响面广。建议在合入前修复 P0 和 P1，P2 也建议一并修复。

⛔ 需要修改

Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here

PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

💡 Tip:

• Committer can comment /approve or /lgtm
• Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

chenpeizhe, thanks for your pull request. All authors of the commits have signed the CLA. 👍

/approve