[RFC]: MindSpeed-MM 生成类模型异构FSDP2统一框架-MindSpeed-MM-AtomGit

已关闭

[RFC]: MindSpeed-MM 生成类模型异构FSDP2统一框架 #576

chenpeizhe创建于 8月7日关闭于 8月29日

8月7日创建

状态（Status）： Draft
作者（Authors）： @chenpeizhe
创建日期（Created）： 2026-08-07
更新日期（Updated）： 2026-08-07
相关 Issue/PR： #2
#https://gitcode.com/Ascend/MindSpeed-MM/pull/2920
#https://gitcode.com/Ascend/MindSpeed-MM/pull/2921
#https://gitcode.com/Ascend/MindSpeed-MM/pull/2922
#https://gitcode.com/Ascend/MindSpeed-MM/pull/2924
#https://gitcode.com/Ascend/MindSpeed-MM/pull/2926
#https://gitcode.com/Ascend/MindSpeed-MM/pull/2928

1. 概述

1.1 简介

支持MindSpeed-MM社区FSDP2框架生成类模型统一框架，提供高效迁移途径及高效灵活训练范式

1.2 动机

当前MindSpeed-MM社区无FSDP2统一生成模型训练框架，迁移模型工作量大，相关特性可复用性差，同时不支持异构模型统一训练

1.3 目标

构建MindSpeed-MM社区FSDP2框架生成类模型统一框架，形成统一的训练流程，高度集成可复用，支持异构模型统一训练和个性化配置

2. 用例分析

共44个文件变更
+5986
-727
*

3. 方案设计

3.1 总体方案

1. 总体架构流程图

flowchart TD
    A["torchrun trainer.py + yaml"] --> B["配置解析<br/>Arguments 聚合 6 个 dataclass"]
    B --> C["Trainer.__init__"]
    C --> D["initialize<br/>插件注册 / 分布式 / parallel_state"]
    D --> E["① 模型构建<br/>ModelContainerHub.build"]
    E --> F["② features 装配<br/>recompute / activation offload"]
    F --> G["③ FSDP2 包裹<br/>ParallelApplier"]
    G --> H["④ meta 物化 + 权重加载"]
    H --> I["optimizer / dataloader<br/>checkpointer / scheduler"]
    I --> J["TrainEngine.train 训练循环"]

各环节代码位置：

环节	位置
配置解析	`trainer.py:489` → `params/argument.py:33`（parallel / model / data / training / tools / features 六节）
初始化	`trainer.py:141`：`import_plugin` 触发 `@model_register` 注册 → `init_process_group` → `init_parallel_state`
① 模型构建	`trainer.py:181` → `model_container_hub.py:61`（详见第 2 节）
② features 装配	`apply_features.py:92` `pre_fully_shard_apply`：recompute 替换 forward → activation offload hooks
③ FSDP2 包裹	`torch_parallelize.py:61` → `fully_shard_parallel.py:43`（详见第 5 节）
④ meta 物化	`trainer.py:208`：FSDP 包裹之后才 `to_empty` + 权重加载
训练循环	`train_engine.py:226`

两个关键时序约束：

features 先于 FSDP 包裹：recompute / activation offload 是 forward 级包装，若后于 fully_shard 执行，FSDP hook 会被包到 checkpoint 外层而失效；
meta 物化后于 FSDP 包裹：先在 meta device 上建立分片结构（DTensor 布局），再 to_empty 物化并加载权重。

1.1 整体训练逻辑分层视图

框架分六层，每层只认相邻层的契约；TrainEngine 对模型的全部假设只有三条（model(**batch) 返回 .loss、梯度同步委托、checkpoint 走 state_dict），这是 trainer.py 对接入新生成模型零改动的原因：

flowchart TD
    L1["① yaml 配置层<br/>parallel / model.models / features / training"]
    L2["② Trainer 基建层（模型无关）<br/>分布式 / optimizer / dataloader / checkpointer"]
    L3["③ 组件构建层 ModelContainerHub<br/>逐子模型 meta_init / setup_weights / freeze"]
    L4["④ 组件管理层 ModelContainer<br/>FSDP opt-out / checkpoint 策略 / train() 保护<br/>forward 委托"]
    L5["⑤ 数据流编排层 GenerativePipeline<br/>prepare_conditions → noising<br/>→ predict → compute_loss"]
    L6["⑥ 组件实例层<br/>ae / text_encoder / predictor<br/>（各自独立 FSDP 单元）"]
    L1 --> L2 --> L3 --> L4 --> L5 --> L6
    L2 -.->|"model(**batch) → .loss<br/>train_engine.py:152"| L4

分层职责说明：

③④ 组件侧：container 负责子组件的初始化（③构建期）与实例持有（④运行期），向 pipeline 提供组件实例；checkpoint 策略、freeze 保护、FSDP opt-out 都属于组件生命周期管理，留在容器；
⑤ 编排侧：子组件间的数据流（条件编码 → 加噪 → 预测 → loss）抽象在 pipeline 层，不写进各模型的 container；pipeline 是纯编排器（非 nn.Module），因此 named_modules 路径、FSDP 模式匹配、checkpoint 路径全部不受影响，TrainEngine 也无需感知；
② 基建侧：TrainEngine 只凭三条契约驱动训练，wan2.2（容器路径）与 ltx2（单模型路径）共用同一 TrainEngine——统一 pipeline 在基建层由此成立。

1.2 配置下发链路

yaml 各节经 Arguments 聚合、校验后分发到各执行器，配置面与执行器一一对应：

flowchart TD
    Y["yaml 六节<br/>parallel / model / data<br/>training / tools / features"] --> A["ConfigManager<br/>Arguments 聚合 + 校验"]
    A --> V1["ParallelArguments.model_post_init<br/>fsdp=auto 解析 / tp=1 强制"]
    A --> V2["compute_distributed_training<br/>跨节联动校验"]

    V1 --> P1["parallel.fsdp_plan<br/>apply_modules / ignored_modules<br/>混精 / prefetch / cpu_offload"]
    V1 --> P2["parallel 其他<br/>ring_attention_size<br/>ulysses_parallel_size / ep_plan"]
    A --> M["model.models.{ae,text_encoder,predictor}<br/>meta_init / dtype / freeze<br/>from_pretrained / load_format"]
    A --> F["features<br/>recompute_plan / activation_offload_plan"]
    A --> T["training / tools<br/>lr / iters / load / save<br/>profile / memory_profile"]

    P1 --> EX1["ParallelApplier<br/>FSDP2 包裹 + sub mesh"]
    P2 --> EX2["parallel_state<br/>mesh / CP 组构建"]
    M --> EX3["ModelContainerHub.build<br/>子模型构建 + 权重加载"]
    F --> EX4["FeaturesApplier<br/>recompute / activation offload"]
    T --> EX5["TrainEngine / Checkpointer<br/>训练循环 / 存取恢复"]

关键校验点：fully_shard_parallel_size: auto 的解析、tp=1 强制及跨节联动校验都在配置解析期完成（params/parallel_args.py），非法配置在训练开始前即报错。

2. 模型构建与权重初始化流程视图

2.1 构建路径总览

Trainer.get_foundation_model（trainer.py:171）按 yaml 是否含 model.models 分两条路：单模型走 ModelHub.build（modelhub.py:97）；Wan2.2 容器走 ModelContainerHub.build（model_container_hub.py:61），后者对每个子模型独立执行下图流程：

flowchart TD
    S["遍历 yaml model.models<br/>ae / text_encoder / predictor"] --> P1["pop freeze / meta_init<br/>显式传参 meta_init=子模型值<br/>model_container_hub.py"]
    P1 --> Q1{"meta_init ?"}
    Q1 -- "是" --> M1["init_empty_weights 上下文<br/>cls._from_config 建 meta 模型<br/>modelhub.py:40"]
    Q1 -- "否" --> M2["cls.from_pretrained<br/>直接加载真权重<br/>modelhub.py:48"]
    M1 --> Q2{"存在 training.load<br/>训练 ckpt ?"}
    M2 --> Q2
    Q2 -- "有" --> W3["跳过 setup_weights<br/>权重由 checkpointer 恢复"]
    Q2 -- "无" --> W1["setup_weights<br/>按 load_format 加载"]
    W1 --> Q3{"load_format 嗅探<br/>auto 时按目录内容分流<br/>resolve_load_format (convert.py)"}
    Q3 -- "diffusers" --> L2["load_diffusers_weights<br/>model_type 选择键映射 pipeline"]
    Q3 -- "hf" --> L3["load_hf_weights"]
    Q3 -- "dcp" --> Q5{"源目录是 DCP ?"}
    Q5 -- "是 / 有效缓存已有" --> L1["load_dcp_weights"]
    Q5 -- "否(diffusers/hf)" --> CV["在线转换为 DCP<br/>convert_*_to_dcp<br/>写入 conversion_dir"]
    CV --> L1
    L1 & L2 & L3 --> Z["setup_module_weights<br/>meta_init_utils.py:39<br/>加载成功则跳过随机初始化"]
    W3 --> Z
    Z --> F1["_apply_freeze<br/>冻结: eval + requires_grad_(False)"]
    F1 --> Q4{"非 meta 且冻结 ?"}
    Q4 -- "是" --> T1["立即 .to(NPU)<br/>_ms_mm_skip_to_device 可豁免"]
    Q4 -- "否" --> T2["留在 meta / CPU<br/>等 FSDP 包裹后统一物化"]
    T1 & T2 --> E["container_cls(built_models)<br/>组装 Wan2_2ModelContainer"]

2.2 三个子模型的构建方式对照

子模型	model_id / Bridge	meta_init	权重来源	构建与加载方式
ae (VAE)	`wan_video_vae` → `WanVideoVAEBridge`	false	diffusers vae 目录	`WanVideoVAE`（`DiffusersAEModel + BaseModel`）`__init__` 内直接 `AutoencoderKLWan.from_pretrained`；`setup_weights` 为显式 no-op（权重已在构造内加载，防止二次加载）；构建后即 `.to(NPU)`
text_encoder (UMT5)	`UMT5` → `UMT5Bridge`	true	hf 格式目录	meta 构建（`models/wan2_2/text_encoder.py` 的 `build_text_encoder`）→ `BaseModel.setup_weights` 默认实现物化加载（hf 直载），特有逻辑经钩子注入：`post_load` 重绑 `shared.weight` ↔ `encoder.embed_tokens.weight`
predictor (DiT)	`wan2_2_predictor` → `WanDiTFSDP2Bridge`	true	dcp / diffusers / hf 均可	meta 构建 `WanDiTFSDP2WithSetup`（`WanDiTFSDP2 + BaseModel`）→ `BaseModel.setup_weights` 默认实现：经 `resolve_load_format` 嗅探后统一走 `load_checkpoint_with_conversion`，diffusers 键映射由类属性 `weight_transform_model_type = "wan2_2"` 选择注册在 `WEIGHT_TRANSFORM_PIPELINES` 的 `Wan22DiffusersTransformPipeline`（`models/wan2_2/diffusers_ckpt_trans.py`，表格驱动）；`post_load` 做 dtype 重 cast；加载失败直接抛异常，不静默回退随机初始化

2.3 权重加载统一管线与在线 DCP 转换

模型侧的 setup_weights 只保留特有逻辑，模型无关的加载编排全部收敛到 checkpoint 层三个函数：

resolve_checkpoint_source(config)（convert.py）：裁决权重来源（checkpoint_path 优先于 from_pretrained）、提取声明格式（load_format / checkpoint_format / hub_backend）、解析转换缓存目录（checkpoint_conversion_dir → dcp_cache_dir → dcp_conversion_dir → checkpoint_path，按此优先级）；
resolve_load_format(path, declared)（convert.py）：归一化声明格式（auto/hf/diffusers/dcp，含别名），auto 时按目录内容嗅探（DCP → diffusers → hf）；
load_checkpoint_with_conversion(model, path, fmt, model_type=...)（convert.py）：统一加载入口。hf 直载；diffusers 经 load_diffusers_weights 加载，键映射 pipeline 由 model_type（WEIGHT_TRANSFORM_PIPELINES 注册键，如 "wan2_2"）选择，公共函数不绑定具体模型，新模型注册自己的 pipeline 并传入 model_type 即可复用。

为什么权重加载是独立阶段（setup_weights），而不是折进 from_pretrained/_from_config：

meta 路径物理约束：meta 构建在 ModelHub.build 的 init_empty_weights() 上下文内执行 _from_config，权重物化必须在退出上下文之后（先 materialize_meta_params 再加载），无法折进 _from_config；
training.load 门控在调用点：ModelContainerHub 只在 training_args.load is None 时调用 setup_weights（DCP 断点恢复会覆盖权重），而两个抽象 classmethod 的签名只收 ModelArguments，拿不到 training 配置；
语义清晰：from_pretrained/_from_config 保持"结构构造"语义，权重加载是两阶段生命周期（结构 → 物化+加载）的第二阶段。非 meta 的子模型（如 ae）可以不走该阶段、在构造内直载，两条路径并存不冲突。

setup_weights 的默认实现目前放在 BaseModel 上（base_model.py）；若希望保持 BaseModel 极简（仅两个抽象 classmethod），可演进为生成模型专用抽象类 GenerativeBaseModel(BaseModel) 承载该默认实现，调用点契约不变。

在线 DCP 转换：当声明 load_format: dcp 而源目录不是 DCP 时，框架先把源权重（diffusers 或 hf）加载进模型、再以 DCP iteration 0 落盘到转换缓存目录，随后从缓存目录加载；二次运行时检测到缓存目录已是 DCP，直接复用，不再重复转换。配置示例：

predictor:
  model_id: wan2_2_predictor
  from_pretrained: /home/weights/Wan2.2-T2V-A14B-Diffusers/transformer_2  # diffusers 源权重
  load_format: dcp                          # 声明走 DCP（触发在线转换）
  checkpoint_conversion_dir: /path/to/dcp_cache/transformer_2   # 转换缓存目录

转换产物保存位置：<checkpoint_conversion_dir>/iter_0000000/（DCP 分片）+ <checkpoint_conversion_dir>/latest_checkpointed_iteration.txt（tracker 文件）。若不配置 checkpoint_conversion_dir，可改用 checkpoint_path 指向缓存目录（此时它不再作为权重来源）。

缓存有效性判断与容错（dcp_checkpointer.py）：

looks_like_dcp_checkpoint_dir 只在目录含有真实 DCP 内容时返回 True：tracker 文件、.distcp 分片、或内含 .metadata 的 iter_* / release 子目录。被中断的转换留下的空壳 iter_* 目录不会被误判为可复用缓存，下次运行会自动重新转换（save 直接写入已存在的目录，无需手工清理）；
DistributedCheckpointer.load 在 tracker 文件缺失时不再默认 iteration=-1，而是自动发现唯一含 .metadata 的 iter_*（或 release）子目录进行加载；无有效候选或多个候选时抛出明确错误（提示 checkpoint 不完整或存在歧义），避免 iter_-000001/.metadata 之类的隐晦报错。

2.4 设计要点

同一容器内异构并存：ModelContainerHub.build 逐子模型向 ModelHub.build 显式传 meta_init 参数（可选关键字，缺省回落 training_args.init_model_with_meta_device，单模型路径行为不变），使 ae 走真权重直载、text_encoder / predictor 走 meta 初始化，互不干扰；
meta 物化两段式：子模型构建期 setup_weights 只做"权重内容"加载；meta 参数的"实体化"（to_empty）推迟到 FSDP 包裹之后（trainer.py:208），保证 DTensor 分片布局先建立。trainer 的初始化阶段（init_model_weights，随机初始化用）会跳过 setup_weights 已成功加载（_weights_loaded=True）的子模型，只对随机初始化兜底的模型执行——否则已加载的权重可能被重初始化逻辑破坏（如带 padding_idx 的 Embedding 会被 empty 张量替换）；
dtype 配置生效：modelhub._apply_configured_dtype 在模型配置显式声明 dtype 时按配置 cast（如 wan2.2 各子模型的 dtype: bf16），未配置时维持原 .float() 行为，无需模型侧再重写 float()；
冻结语义：freeze: true → eval() + 全参数 requires_grad_(False)（model_container_hub.py:42）；容器 train() 只切 predictor（container:185），checkpoint 只存 predictor（基类 checkpoint_submodels = ("predictor",)）；
调度器网格是数值契约：WanFlowMatchScheduler（pipeline 持有）的 sigma 网格与 DiffSynth v1.x 一致（shift=3.0、sigma_min=0.003/1.002），与官方 scheduler_config.json（flow_shift: 3.0）同口径；它不等同于 DiffSynth v2 Wan 模板（shift=5）和 diffusers FlowMatchEulerDiscreteScheduler（端点约定不同，实测 sigma 最大差 ~0.006）。对账脚本必须显式对齐该网格，禁止按上游模板"升级"（详见 models/wan2_2/diffusion.py 文件头注释）。

3. UML 用例视图

系统边界内为框架提供的核心能力（用例），系统边界外为使用者（训练工程师）。每个用例标注其实现代码位置。

flowchart LR
    actor(["训练工程师"])

    subgraph boundary["MindSpeed-MM FSDP2 训练框架"]
        UC1["UC1 配置异构子模型初始化<br/>各子模型独立 meta_init / dtype /<br/>freeze / from_pretrained / load_format"]
        UC2["UC2 自由选择 FSDP2 切分模块<br/>fsdp_plan.apply_modules /<br/>ignored_modules 模式配置"]
        UC3["UC3 配置显存优化策略<br/>recompute / activation_offload /<br/>cpu_offload"]
        UC4["UC4 配置混合精度与预取<br/>param_dtype / reduce_dtype /<br/>num_to_forward/backward_prefetch"]
        UC5["UC5 执行训练与断点恢复<br/>training.load / save (dcp)"]
        UC6["UC6 性能与显存剖析<br/>tools.profile / tools.memory_profile"]

        SYS1(("ModelContainerHub.build<br/>model_container_hub.py:61"))
        SYS2(("fully_shard_parallel_modules<br/>fully_shard_parallel.py:43"))
        SYS3(("FeaturesApplier<br/>apply_features.py:17"))
        SYS4(("set_modules_to_prefetch<br/>fully_shard_parallel.py:457"))
        SYS5(("TrainEngine + Checkpointer<br/>train_engine.py:226"))
        SYS6(("Profiling tools<br/>trainer.py get_dataloader 段"))
    end

    actor --> UC1
    actor --> UC2
    actor --> UC3
    actor --> UC4
    actor --> UC5
    actor --> UC6

    UC1 -.include.-> SYS1
    UC2 -.include.-> SYS2
    UC3 -.include.-> SYS3
    UC4 -.include.-> SYS4
    UC5 -.include.-> SYS5
    UC6 -.include.-> SYS6

3.1 用例说明

用例	用户配置入口（yaml）	框架处理
UC1 异构子模型初始化	`model.models.{ae,text_encoder,predictor}` 各自的 `meta_init` / `dtype` / `freeze` / `from_pretrained` / `load_format`	`ModelContainerHub.build` 逐子模型委托 `ModelHub.build(meta_init=...)`，经 `model_register` 取类构建；权重经 `BaseModel.setup_weights` 默认实现加载
UC2 自由选择切分模块	`parallel.fsdp_plan.apply_modules` / `ignored_modules`，支持 `predictor.blocks.{*}` 通配	后序遍历 `named_modules()` 模式匹配 → ignored 参数级排除 → 逐个 `fully_shard` → 子模型残余参数兜底 → 容器根级 opt-out
UC3 显存优化策略	`features.recompute_plan` / `activation_offload_plan` / `parallel.fsdp_plan.cpu_offload`	`FeaturesApplier.pre_fully_shard_apply`：recompute 替换模块 forward 为 checkpoint 包装；activation offload 用 saved_tensors_hooks 异步搬激活到 CPU
UC4 混合精度与预取	`param_dtype: bf16` / `reduce_dtype: fp32` / `num_to_*_prefetch`	`MixedPrecisionPolicy` + `CPUOffloadPolicy` 装配进每个 `fully_shard` 调用；prefetch 按子模型 owner 过滤（`_filter_prefetch_layers_by_owner`，fully_shard_parallel.py:383），禁止跨子模型预取
UC5 训练与恢复	`training.load` / `save_format: dcp`	`TrainEngine` 构造时按需 `load()`；`Wan2_2ModelContainer.state_dict()` 只导出 predictor 权重，冻结子模型永远从 `from_pretrained` 重载
UC6 剖析	`tools.profile.enable` / `tools.memory_profile.enable`	训练循环内按 `start_step/end_step` 采集 timeline 与 memory snapshot

4. 模型结构视图（类关系）

classDiagram
    class ModelContainer {
        +pipeline  ; 纯编排器（非 nn.Module），默认 None
        +get_sub_models()
        +state_dict()  ; 按 checkpoint_submodels 保存
        +load_state_dict()  ; 按子模型前缀分发，返回 _IncompatibleKeys
        +fully_shard() bool  ; 返回 True，opt-out 根级包裹
        +forward()  ; pipeline 非空时委托，否则抛 NotImplementedError
        +set_is_last_backward()
        +set_requires_all_reduce()
    }
    class Wan2_2ModelContainer {
        +ae: WanVideoVAE
        +text_encoder: _SimpleTextEncoder
        +predictor: WanDiTFSDP2
        +pipeline: WanT2VPipeline  ; 数据流编排委托给它
        +checkpoint_submodels = ("predictor",)  ; 只存 predictor
        +train()  ; 只切 predictor 训练态
    }
    ModelContainer <|-- Wan2_2ModelContainer

    class GenerativePipelineBase {
        <<纯编排器，非 nn.Module>>
        +prepare_conditions(container, batch)
        +noising(container, conditions, batch)
        +predict(container, inputs)
        +compute_loss(...)
        +forward(container, batch)  ; 编排四个阶段
    }
    class FlowMatchT2VPipeline {
        +scheduler  ; flow-match 调度器由 pipeline 持有
        位置: models/wan2_2/pipeline_wan2_2.py（与 WanT2VPipeline 同文件合并）
        组件角色约定: ae / text_encoder / predictor
    }
    class WanT2VPipeline {
        +task t2v/i2v 分支
        +_prepare_i2v_condition()
        +scheduler = WanFlowMatchScheduler
    }
    GenerativePipelineBase <|-- FlowMatchT2VPipeline
    FlowMatchT2VPipeline <|-- WanT2VPipeline
    Wan2_2ModelContainer o-- WanT2VPipeline

    class WanVideoVAE {
        +forward = encode  ; no_grad 下运行
    }
    class _SimpleTextEncoder {
        +model: UMT5EncoderModel  ; no_grad 下运行
        +forward = encode  ; 保证触发 FSDP2 hooks
        位置: models/wan2_2/text_encoder.py
    }
    class WanDiTFSDP2 {
        +patch_embedding: Conv3d
        +text_embedding: Sequential
        +time_embedding: Sequential
        +time_projection: Sequential
        +blocks: ModuleList~WanAttentionBlock x 40~
        +head: Head
        -_freqs  ; RoPE 频率每 forward 预计算一次，40 层共享
    }
    Wan2_2ModelContainer *-- WanVideoVAE
    Wan2_2ModelContainer *-- _SimpleTextEncoder
    Wan2_2ModelContainer *-- WanDiTFSDP2

    class WanAttentionBlock {
        +norm1: WanLayerNorm
        +self_attn: WanSelfAttention
        +norm3: WanLayerNorm 或 Identity
        +cross_attn: WanCrossAttention
        +norm2: WanLayerNorm
        +ffn: Sequential(Linear, GELU, Linear)
        +modulation: Parameter(1, 6, dim)  ; adaLN 6 组 shift/scale/gate
    }
    class WanSelfAttention {
        +q / k / v / o: Linear
        +norm_q / norm_k: WanRMSNorm
        +forward() QK-norm → RoPE → CP(Ulysses/Ring) → flash_attention → gather → o
    }
    class WanCrossAttention {
        +forward() 全文本序列 k_lens=None，无 RoPE / CP
    }
    class Head {
        +norm: WanLayerNorm
        +head: Linear
        +modulation: Parameter
    }
    WanDiTFSDP2 *-- WanAttentionBlock
    WanDiTFSDP2 *-- Head
    WanAttentionBlock *-- WanSelfAttention
    WanSelfAttention <|-- WanCrossAttention

结构要点：

ModelContainer.__init__ 用 setattr 而非 ModuleDict 挂子模型（model_container.py:19），保证 named_modules() 路径为 predictor.blocks.0 扁平形式，与 apply_modules 模式直接对应；
Wan2_2ModelContainer.train()（wan2_2_model_container.py）只调用 predictor.train(mode)，防止训练框架递归 model.train() 激活冻结 UMT5 中的 dropout；
注册点集中在 wan2_2_model_container.py：wan_video_vae → WanVideoVAEBridge、UMT5 → UMT5Bridge、wan2_2_predictor → WanDiTFSDP2Bridge、wan2_2_container → Wan2_2ModelContainer，靠 training.plugin 列表经 import_plugin 动态加载。

5. FSDP2 切分代码交互图

5.1 构建期：包裹流程

sequenceDiagram
    participant T as Trainer.get_model
    participant P as ParallelApplier<br/>(torch_parallelize.py:61)
    participant F as fully_shard_parallel_modules<br/>(fully_shard_parallel.py:43)
    participant M as module_name_match
    participant C as Wan2_2ModelContainer

    T->>P: model (Container, meta device)
    P->>F: apply_fsdp_modules
    F->>F: get_ignored_modules (:293)<br/>"ae" → ignored_params 全集
    F->>F: 后序遍历 named_modules (:208)<br/>保证子模块先于父模块包裹
    loop 每个 apply_modules 模式
        F->>M: "text_encoder" → _SimpleTextEncoder 整体<br/>"predictor.blocks.{*}" → 40 x WanAttentionBlock<br/>"predictor.head" → Head
        F->>F: fully_shard(module, **config)<br/>mp_policy bf16/fp32 + reshard_after_forward<br/>+ CPUOffloadPolicy + ignored_params
    end
    F->>F: 残余参数兜底 (:172-181)<br/>predictor 整体再包一层 FSDP<br/>(patch/text/time_embedding 等未命中参数)
    F->>C: container.fully_shard() → True
    F-->>P: return（跳过根级包裹）
    P->>P: set_modules_to_prefetch (:457)<br/>_filter_prefetch_layers_by_owner (:383)<br/>禁止跨子模型预取

三层模块选择语义：

apply_modules：命中的模块成为独立 FSDP 单元；
ignored_modules：参数级排除，ae 完全不进 FSDP 也不受 mp_policy 管理；
子模型残余兜底 + 容器根级 opt-out（ModelContainer.fully_shard() 返回 True，model_container.py:38）。

5.2 运行期：一个 train_step

sequenceDiagram
    participant E as TrainEngine.train_step<br/>(train_engine.py:130)
    participant C as Wan2_2ModelContainer<br/>forward 委托
    participant PL as WanT2VPipeline<br/>(pipeline_wan2_2.py)
    participant V as ae (VAE)
    participant X as text_encoder (UMT5)
    participant D as WanDiTFSDP2 (predictor)
    participant S as WanFlowMatchScheduler<br/>(pipeline 持有)

    E->>E: configure_hsdp_gradient_sync<br/>非末 micro-step 关闭 FSDP 梯度同步
    E->>C: model(**batch)
    C->>PL: pipeline.forward(container, **batch)
    Note over PL,V: prepare_conditions，torch.no_grad()
    PL->>V: ae(video) → latents
    Note over PL,X: 显式走 forward()<br/>以触发 FSDP2 unshard/reshard hooks
    PL->>X: text_encoder(prompt_ids, mask) → prompt
    Note over PL,S: noising
    PL->>S: q_sample(latents) → noised_latents, noise, timesteps
    Note over PL,D: predict（带梯度）
    PL->>D: predictor(noised_latents, t, prompt)
    loop 40 x WanAttentionBlock（各自独立 FSDP 单元）
        D->>D: unshard → x += gate * self_attn(norm1(x))<br/>→ x += cross_attn(norm3(x))<br/>→ x += gate * ffn(norm2(x)) → reshard
    end
    D-->>PL: model_output
    Note over PL,S: compute_loss
    PL->>S: training_losses → MSE x 时间步权重
    PL-->>E: Wan2_2ModelOutput(loss)
    E->>E: loss.backward() → clip_grad_norm<br/>→ optimizer.step → lr_scheduler.step

pipeline 是纯编排器（非 nn.Module），named_modules() 路径、FSDP 模式匹配、owner prefetch 过滤均不受影响；组件调用仍经各子模型 forward()，FSDP2 pre/post hooks 照常触发。

5.3 子模型独立切分（fsdp_plan.sub_plans）【本地暂存，未上库】

该特性（含配套的冻结 root reshard 恢复）已完成开发与验证，但按分阶段上库计划暂未包含在当前代码中，完整实现与测试已备份至 backups/20260728_feature_staging/（含恢复方式说明）。

设计摘要（恢复时参考）：sub_plans 允许 ModelContainer 的每个子模型独立选择是否切分、独立配置切分数（如 vae 不切、text_encoder 切 4 卡、DiT 全组 auto），通过按子模型 owner 解析有效 plan + ParallelState.get_fsdp_sub_mesh（2D (replica, shard) 子 mesh）实现；配套的一次性恢复钩子解决 torch FSDP2 lazy init 对 root 模块置空 post_forward_mesh_info 导致冻结 root 子模型参数永不 reshard 的问题（恢复后 text_encoder 稳态显存从全量降为 1/shard_size）。语义上"参数组 ≠ 数据组"，每卡数据与全局 batch 划分不变，冻结子模型与全局切分数值逐比特一致。

6. 容器与框架的契约小结

支撑"冻结异构子模型 + 单 predictor 训练"形态的框架契约：

契约	位置	作用
`fully_shard()` 返回 True	model_container.py	容器 opt-out 根级 FSDP 包裹，切分粒度完全交给 yaml
`forward` 委托 pipeline	model_container.py + pipelines/	容器只管组件生命周期；数据流编排（条件编码/加噪/预测/loss）统一在 pipeline 层，容器默认 forward 委托 `pipeline.forward(self, **batch)`，子类也可自行重写
`train()` 只训 predictor	wan2_2_model_container.py	冻结子模型保持 eval，避免 dropout 引入数值漂移
`state_dict()` 只存 predictor	model_container.py 基类 + `checkpoint_submodels = ("predictor",)`	checkpoint 只含可训练参数，ae/text_encoder 从预训练源重载
`setup_weights` 默认实现	base_model.py	生成模型权重加载统一编排（resolve_checkpoint_source → resolve_load_format → load_checkpoint_with_conversion → setup_module_weights），模型侧只留 `post_load` 钩子与 `weight_transform_model_type` 类属性
prefetch 按 owner 过滤	fully_shard_parallel.py:383	防止 predictor 末层预取 text_encoder 首层（异构串行执行下无意义且浪费显存）

异构 meta init 的核心机关在 ModelContainerHub.build：逐个子模型向 ModelHub.build 显式传 meta_init，使同一容器内 ae 可直接 from_pretrained 加载真权重，而 text_encoder / predictor 走 meta 初始化 + BaseModel.setup_weights（按 dcp / diffusers / hf 格式嗅探加载）。

新增一个生成模型的最小工作集：子模型构建类继承 BaseModel（实现 from_pretrained/_from_config，可选 post_load 钩子）+ 一个 ModelContainer 子类（声明 checkpoint_submodels、可选重写 train()）+ 一个 pipeline（flow-match T2V 家族直接复用 FlowMatchT2VPipeline 或按 GenerativePipelineBase 特化）+ yaml model.models 与 training.plugin。TrainEngine/Trainer 零改动。

7. 关键修改逻辑（演进记录）

本节记录框架从"能跑通"到"可上库"的关键修改及其动机，便于理解当前架构形态的来由。

7.1 统一生成模型 pipeline（阶段 A/B）

动机：Wan2_2ModelContainer 曾一身兼两职（组件管理 + 数据流编排），每接入一个生成模型都要重写整套 forward 编排。
修改：拆出 pipeline 层（pipelines/generative_pipeline.py 的 GenerativePipelineBase 四阶段接口 + FlowMatchT2VPipeline 通用 flow-match 数据流 + wan 特化 WanT2VPipeline），容器 forward 默认委托 pipeline.forward(self, **batch)；容器从 342 行瘦到 170 行，只保留组件生命周期管理。pipeline 是纯编排器（非 nn.Module），FSDP/checkpoint/TrainEngine 全部无感。注：FlowMatchT2VPipeline 与 WanT2VPipeline 已合并为单一文件 models/wan2_2/pipeline_wan2_2.py（通用类与 wan 特化类同文件，通用基类 GenerativePipelineBase 仍在 pipelines/ 下供后续模型复用）。
验证：重构前后 tiny 容器 forward 的 loss 与全部中间量逐比特相等。

7.2 权重加载统一（阶段 A + checkpoint 整改）

动机：模型侧曾各写一套 load_format 归一化/嗅探/加载路由（三处不一致的拷贝），checkpoint 层与 9e208c66 的 pipeline 体系存在重复造轮子。
修改：BaseModel.setup_weights 默认实现（resolve_checkpoint_source → resolve_load_format → load_checkpoint_with_conversion → setup_module_weights），模型侧只留 post_load 钩子与 weight_transform_model_type 类属性；diffusers 键映射为表格驱动——checkpoint 层提供通用 DiffusersKeyMapTransformPipeline 与参数化 convert_diffusers_key，映射表由各模型插件自带并注册进 WEIGHT_TRANSFORM_PIPELINES（wan2_2 的表与 pipeline 在 models/wan2_2/diffusers_ckpt_trans.py，checkpoint 层无模型知识，新增模型仅需自带表格 + 一行注册）；load_diffusers_weights 变为薄封装且 model_type 可选；convert_hf_to_dcp/convert_diffusers_to_dcp 合并为 _convert_to_dcp 薄封装；ModelHub.build 增加显式 meta_init 参数（替代全局标志改写，单模型路径行为不变）。
容错：looks_like_dcp_checkpoint_dir 要求真实 DCP 内容（空壳 iter_* 不误判复用）；DistributedCheckpointer.load 无 tracker 时自动发现唯一有效 iter_* 或抛出明确错误（修复 iter_-000001 隐晦报错）。

7.3 子模型分组切分（阶段 D1，fsdp_plan.sub_plans）【已暂存，见 7.7】

动机：显存敏感场景需要"TE 切 4 卡、DiT 全组、VAE 不切"的异构切分，而原框架只有全局单一 mesh。
状态：实现已完成并验证（FSDPPlanConfig.sub_plans + ParallelState.get_fsdp_sub_mesh 2D 子 mesh + 按 owner 解析有效 plan，一期仅冻结子模型），但按分阶段上库计划已从代码中回退，完整实现与测试备份于 backups/20260728_feature_staging/（详见 7.7 节）。

7.4 CP 正确性修复

去 Megatron 依赖：ring_context_parallel.compute_mask 曾惰性导入 from mindspeed.utils import batch_index（该模块硬依赖 megatron），改为模块内 vendored _batch_index（与上游逐行等价的 bisect 实现，等价性单测固化）。
恢复顺序：ulysses+ring 组合时原代码先 ulysses 还原后 ring gather（未按切分逆序），导致 All tensor operands to scatter/gather must have the same size；改为先 ring gather 再 ulysses 还原。ring 还原同时改用与 zigzag 切分配对的 load_balanced_gather_forward_split_backward（修复普通 gather 导致 token 顺序 0,3,1,2 错乱的隐性数值 bug，ring-only 场景同样受益）。
验证：CPU gloo 4 进程 mock 测试，u2r2 / u2r1 / u1r2 三种组合还原结果与原输入逐比特一致。

7.5 meta 初始化健壮性

问题：text_encoder 开 meta_init: true 后，trainer 初始化阶段对已加载权重的子模型仍无条件调 init_model_weights → _SimpleTextEncoder 无 init_weights 直接报错；且对含 padding_idx Embedding 的模型，该阶段的 empty 替换会销毁已加载权重。
修改：trainer 只对 _weights_loaded != True 的子模型执行 init_model_weights（随机初始化路径）；_SimpleTextEncoder 补守卫版 init_weights（loaded 时 no-op，否则委托 transformers 初始化）。

7.6 数据管线瘦身

动机：fsdp/data/ 新增约 4200 行中 ~40% 是 wan2.2 永不触达的拷贝（其他模型的 transform/processor/reader、conversation 模板等）。
修改：删除约 1600 行死代码（conversation.py 全文件、MaskGenerator + 19 个无用 transform、Opensora/CogVideoX/Reward processor、AvVideo/read_video_av、无用 utils/常量）；保留 wan2.2 实际使用项与 Wan 家族 VACEVideoProcessor；diffusion 推理代码与 i2v 链路按规划保留。fsdp 数据管线作为统一依赖独立演进，不回并老路径。wan2.2 新增的 5 个数据模块单独收进 fsdp/data/data_utils/generative_data_utils/（constants.py/data_transform.py/transform_pipeline.py/video_processor.py/video_reader.py），与既有 utils.py、func_utils/ 分离，引用路径已同步更新。

7.7 特性暂存说明（fsdp_subplan / cross_submodel_prefetch / condition_mbs）

以下三个特性已完成开发与验证，但按分阶段上库计划暂未包含在当前代码中，完整实现、测试与验证记录已备份至 backups/20260728_feature_staging/（MANIFEST.md 含逐项文件清单与恢复方式）：

fsdp_subplan（子模型独立 FSDP2 切分）：FSDPPlanConfig.sub_plans 按子模型 owner 覆盖 shard_size/reshard/cpu_offload/enabled；ParallelState.get_fsdp_sub_mesh（2D (replica, shard) 子 mesh）；一期仅冻结子模型。配套冻结 root reshard 恢复：解决 torch FSDP2 lazy init 对 root 模块置空 post_forward_mesh_info 导致冻结 root 参数永不 reshard 的问题（text_encoder 稳态显存 ~11.2GB → ~2.8GB/shard4）。
cross_submodel_prefetch（跨子模型边界预取）：pipeline 在 VAE 计算前对 text_encoder 发起异步 unshard、在 TE 编码前对 predictor root+block0 发起异步 unshard，掩盖跨子模型边界的 all-gather。
condition_mbs（条件阶段批处理）：training.condition_mbs 控制每 N 个 micro-step 对 VAE/TE 做一次批量条件编码并缓存逐槽消费（编码开销摊为 1/N、TE all-gather 频率降为 1/N；CP 整除时各 rank 只编码自己的槽位再交换，消除 CP 组内冗余编码）。

验证记录（备份时）：组编码与逐样本编码逐比特一致（含真实 UMT5 bf16、变长 padding mask）；预取幂等且后续前向零新增 all-gather；冻结 root reshard 恢复循环终态 SHARDED；全量 ut_fsdp 回归通过。

8. 规划中的能力（未实施，设计已评审）

8.1 encoder-dp：条件编码冗余消除与批量化（condition_mbs）【本地暂存，见 7.7】

问题：① 每 micro-step 单独做 VAE/TE 编码，单次计算小、kernel 稀疏、free 大；② CP=N 时 CP 组内 N 卡对同一样本冗余编码 N 次。

方案状态【本地暂存，未上库】：training.condition_mbs（条件阶段批处理 + 条件缓存）与 fsdp_plan.cross_submodel_prefetch（跨子模型边界预取）两个方案均已完成开发与验证，但按分阶段上库计划暂未包含在当前代码中，完整实现与测试备份于 backups/20260728_feature_staging/。设计摘要：每 N 个 micro-step 对 VAE/TE 做一次批量条件编码并缓存逐槽消费（编码开销摊为 1/N、TE all-gather 频率降为 1/N、CP 整除时各 rank 只编码自己槽位再交换）；pipeline 在 VAE/TE 计算期间对下一个子模型发起异步 unshard，掩盖边界 all-gather。验证记录：组编码与逐样本编码逐比特一致（含真实 UMT5 bf16、变长 padding mask）；预取幂等且后续前向零新增 all-gather。

8.2 HSDP：可训练子模块的分组切分（阶段 D2）

问题：D1 的 sub_plans 仅支持冻结子模型；可训练子模块用非全局 shard_size 时，副本组间梯度会发散。（注：D1 实现当前在 backups/20260728_feature_staging/ 暂存，未上库。）

方案（尚未实施）：D1 的 2D 子 mesh (dp_sub_replicate, dp_sub_shard) 本身就是 HSDP 形态——fully_shard 拿 2D mesh 自动做"组内 shard + 副本组间梯度 all-reduce"，梯度全局平均与全局 FSDP 数学等价：

flowchart LR
    subgraph G0["副本组 0（rank 0-3）"]
        A1["shard 0"] --- A2["shard 1"] --- A3["shard 2"] --- A4["shard 3"]
    end
    subgraph G1["副本组 1（rank 4-7）"]
        B1["shard 0"] --- B2["shard 1"] --- B3["shard 2"] --- B4["shard 3"]
    end
    A1 <-."反向: all-reduce（副本组间）".-> B1
    A2 <-.-> B2
    A3 <-.-> B3
    A4 <-.-> B4
    Note1["前向: 组内 all-gather 重建全参数<br/>反向: 组内 reduce-scatter + 组间 grad all-reduce"]

待办项（实施时）：① _validate_fsdp_sub_plans 放开 trainable: true（改显式声明校验）；② clip_grad_norm 按 mesh 分组归约（现逻辑对全局组单一 all-reduce，HSDP 副本会重复计数范数 → 裁剪系数错误，需按 grad.device_mesh 分组在各自 shard 维组归约）；③ 约束文档化（可训练子模块不得有 forward 中更新的 buffer）；④ 验证口径：前向逐比特、梯度/收敛容差（归约顺序差异属预期）、副本组权重不发散、clip 范数与全量参考一致、DCP 异构 mesh 回归。*

4. 测试设计

提供相关模型构建、权重加载、切分、训练等测试用例

5. 缺点和风险

说明潜在风险（Breaking Change、性能回退、复杂度提升、引入的安全问题）、负面影响（对现有功能/用户的冲击）、实现成本（代码量/维护成本/人力投入）、是否有API或版本兼容性、旧版本迁移方案问题等，给出应对措施。

6. 现有技术

参考其他项目/社区的类似设计，说明借鉴与差异。

7. 未解决问题

待社区讨论/决策的开放问题，如硬件适配范围、参数默认值等（需在RFC通过前解决）。

附录

参考资料链接
术语表
文档更新计划

欢迎加入社区，感谢您对社区的贡献 🎉!

ascend-robot成员

8月7日添加了label：rfc

C chenpeizhe成员

8月7日修改了issue 的描述

C chenpeizhe成员

8月7日修改了issue 的描述

C chenpeizhe成员

8月7日修改了issue 的描述

C chenpeizhe成员

8月11日关联了pull request：feat: support fsdp model container

yeqm成员

8月21日评论：

👋 您好，欢迎向 MindSpeed MM 提交 Issue！
我们已收到您的反馈，感谢你对开源社区的支持。🎉
📅 处理时效：维护团队将在 24 小时内查看并回复您的问题（工作日）。
🚨 紧急联系：如果您的问题非常紧急，可通过微信联系我们。
请确保 Issue 描述清晰，包含复现步骤和日志，这将帮助我们更快定位问题。谢谢！

ascend-robot成员

8月29日关闭了 issue

ascend-robot成员

8月29日添加了label：resolved