feat: add model container-MindSpeed-MM-AtomGit

Merged

feat: add model container #2921

chenpeizheAdd on Aug 4Merged fromchenpeizhe/MindSpeed-MM_wan22fsdp_pr2： master into Ascend/MindSpeed-MM： master

feat: add model container #2921

Merged

chenpeizheAdd on Aug 4

Merged fromchenpeizhe/MindSpeed-MM_wan22fsdp_pr2 masterintoAscend/MindSpeed-MM master

Total: 5 files changed+776-0

Amindspeed_mm/fsdp/models/model_container.py+135-0

@@ -0,0 +1,135 @@
	1	+import logging
	2	+from typing import Dict, Optional, Tuple
	3	+
	4	+import torch.nn as nn
	5	+from torch.nn.modules.module import _IncompatibleKeys
	6	+
	7	+
	8	+logger = logging.getLogger(__name__)
	9	+
	10	+
	11	+class ModelContainer(nn.Module):
	12	+ """Lightweight container for top-level sub-models (own meta-init/parallel/freeze policy), registered as direct
	13	+ attributes for flat FSDP pattern paths; ``forward`` delegates to an attached ``pipeline`` unless overridden.
	14	+ ``checkpoint_submodels`` selects sub-models entering checkpoints (default: trainable ones); ``legacy_flat_key_target`` receives prefix-less keys."""
	15	+
	16	+ checkpoint_submodels: Optional[Tuple[str, ...]] = None
	17	+ legacy_flat_key_target: Optional[str] = None
	18	+ pipeline = None
	19	+
	20	+ def __init__(self, models: Dict[str, nn.Module]):
	21	+ super().__init__()
	22	+ for name, model in models.items():
	23	+ setattr(self, name, model)
	24	+ self._model_names = list(models.keys())
	25	+
	26	+ def forward(self, args, *kwargs):
	27	+ if self.pipeline is not None:
	28	+ return self.pipeline.forward(self, args, *kwargs)
	29	+ raise NotImplementedError(
	30	+ "ModelContainer.forward() must be overridden by a subclass."
	31	+ )
	32	+
	33	+ def train(self, mode: bool = True):
	34	+ """Sub-models follow ``mode``; fully frozen sub-models are forced back to eval (e.g. to keep frozen dropout off)."""
	35	+ super().train(mode)
	36	+ for name in self._model_names:
	37	+ sub_model = getattr(self, name, None)
	38	+ if sub_model is not None and self._is_fully_frozen(sub_model):
	39	+ sub_model.eval()
	40	+ return self
	41	+
	42	+ @staticmethod
	43	+ def _is_fully_frozen(module: nn.Module) -> bool:
	44	+ params = list(module.parameters())
	45	+ return len(params) > 0 and all(not p.requires_grad for p in params)
	46	+
	47	+ def get_sub_models(self) -> Dict[str, nn.Module]:
	48	+ """Return a shallow copy of the sub-model mapping."""
	49	+ return {name: getattr(self, name) for name in self._model_names}
	50	+
	51	+ def get_checkpoint_submodel_names(self):
	52	+ """Names of sub-models whose weights go into training checkpoints."""
	53	+ if self.checkpoint_submodels is not None:
	54	+ return [name for name in self.checkpoint_submodels if name in self._model_names]
	55	+ return [
	56	+ name
	57	+ for name in self._model_names
	58	+ if not self._is_fully_frozen(getattr(self, name))
	59	+ ]
	60	+
	61	+ def state_dict(self, destination=None, prefix="", keep_vars=False):
	62	+ """Save the checkpointed sub-models' weights with the container prefix."""
	63	+ if destination is None:
	64	+ destination = {}
	65	+ for name in self.get_checkpoint_submodel_names():
	66	+ sub_model = getattr(self, name, None)
	67	+ if sub_model is None:
	68	+ continue
	69	+ sub_prefix = f"{prefix}{name}."
	70	+ for key, value in sub_model.state_dict(prefix="", keep_vars=keep_vars).items():
	71	+ destination[f"{sub_prefix}{key}"] = value
	72	+ return destination
	73	+
	74	+ def load_state_dict(self, state_dict, strict=True, **kwargs):
	75	+ """Dispatch keys to sub-models by ``<name>.`` prefix (flat keys go to ``legacy_flat_key_target``); with ``strict=True`` missing/unexpected keys raise."""
	76	+ sub_sds = {name: {} for name in self._model_names}
	77	+ missing_keys = []
	78	+ unexpected_keys = []
	79	+
	80	+ for key, value in state_dict.items():
	81	+ dispatched = False
	82	+ for name in self._model_names:
	83	+ sub_prefix = f"{name}."
	84	+ if key.startswith(sub_prefix):
	85	+ sub_sds[name][key[len(sub_prefix):]] = value
	86	+ dispatched = True
	87	+ break
	88	+ if not dispatched:
	89	+ if self.legacy_flat_key_target is not None and self.legacy_flat_key_target in sub_sds:
	90	+ sub_sds[self.legacy_flat_key_target][key] = value
	91	+ else:
	92	+ unexpected_keys.append(key)
	atomgit-botAug 4 🟡 Medium Priority model_container.py 第94-95行 `load_state_dict` 方法中： if self.legacy_flat_key_target is not None: sub_sds[self.legacy_flat_key_target][key] = value `sub_sds` 在第81行初始化为 `{name: {} for name in self._model_names}`。若 `legacy_flat_key_target` 的值（例如 `"predictor"`）不在 `self._model_names` 中，访问 `sub_sds[self.legacy_flat_key_target]` 会抛出 `KeyError`。 `Wan2_2ModelContainer` 设置 `legacy_flat_key_target = "predictor"`（wan2_2_model_container.py 第138行），若用户构造容器时 `models` 字典中不包含 `"predictor"` 键，然后调用 `load_state_dict` 加载包含扁平 key 的旧 checkpoint，将直接崩溃。触发条件：容器 `_model_names` 中不包含 `legacy_flat_key_target` 指定的名称时，加载含无前缀 key 的 state_dict。建议：在写入 `sub_sds[self.legacy_flat_key_target]` 前检查 `self.legacy_flat_key_target in sub_sds`，若不在则将该 key 加入 `unexpected_keys`。 Incorrect? C chenpeizhe Aug 5 Comments： Aug 5 Comments：
	93	+
	94	+ for name, sub_sd in sub_sds.items():
	95	+ sub_model = getattr(self, name, None)
	96	+ if sub_model is None or not sub_sd:
	97	+ continue
	98	+ ret = sub_model.load_state_dict(sub_sd, strict=False, **kwargs)
	99	+ missing_keys.extend(f"{name}.{key}" for key in ret.missing_keys)
	100	+ unexpected_keys.extend(f"{name}.{key}" for key in ret.unexpected_keys)
	101	+
	102	+ error_msgs = []
	103	+ if missing_keys:
	104	+ error_msgs.append(f"Missing key(s) in state_dict: {missing_keys}")
	105	+ if unexpected_keys:
	106	+ error_msgs.append(f"Unexpected key(s) in state_dict: {unexpected_keys}")
	107	+ if strict and error_msgs:
	108	+ raise RuntimeError(
	109	+ "Error(s) in loading state_dict for {}:\n\t{}".format(
	110	+ type(self).__name__, "\n\t".join(error_msgs)
	111	+ )
	112	+ )
	113	+ for msg in error_msgs:
	114	+ logger.info(msg)
	115	+ return _IncompatibleKeys(missing_keys, unexpected_keys)
	116	+
	117	+ # ------------------------------------------------------------------
	118	+ # FSDP2 integration helpers
	119	+ # ------------------------------------------------------------------
	120	+
	121	+ def fully_shard(self, fsdp_plan):
	122	+ """Opt out of root-level FSDP wrapping (container holds no trainable parameters; sub-models are wrapped beforehand)."""
	123	+ return True
	124	+
	125	+ def set_is_last_backward(self, is_last_step: bool):
	126	+ """Delegate to any wrapped sub-models."""
	127	+ for _name, module in self.named_modules():
	128	+ if hasattr(module, "set_is_last_backward") and module is not self:
	129	+ module.set_is_last_backward(is_last_step)
	130	+
	131	+ def set_requires_all_reduce(self, is_last_step: bool):
	132	+ """Delegate to any wrapped sub-models."""
	133	+ for _name, module in self.named_modules():
	134	+ if hasattr(module, "set_requires_all_reduce") and module is not self:
	135	+ module.set_requires_all_reduce(is_last_step)

Amindspeed_mm/fsdp/models/wan2_2/diffusers_ckpt_trans.py+56-0

@@ -0,0 +1,56 @@
	1	+# Copyright 2025 Bytedance Ltd. and/or its affiliates
	2	+"""Diffusers -> native WanModel key mapping tables.
	3	+"""
	4	+
	5	+# Exact key renames (old_key -> new_key)
	6	+DIFFUSERS_CONVERT_MAPPING = {
	7	+ "condition_embedder.text_embedder.linear_1.bias": "text_embedding.0.bias",
	8	+ "condition_embedder.text_embedder.linear_1.weight": "text_embedding.0.weight",
	9	+ "condition_embedder.text_embedder.linear_2.bias": "text_embedding.2.bias",
	10	+ "condition_embedder.text_embedder.linear_2.weight": "text_embedding.2.weight",
	11	+ "condition_embedder.time_embedder.linear_1.bias": "time_embedding.0.bias",
	12	+ "condition_embedder.time_embedder.linear_1.weight": "time_embedding.0.weight",
	13	+ "condition_embedder.time_embedder.linear_2.bias": "time_embedding.2.bias",
	14	+ "condition_embedder.time_embedder.linear_2.weight": "time_embedding.2.weight",
	15	+ "condition_embedder.time_proj.bias": "time_projection.1.bias",
	16	+ "condition_embedder.time_proj.weight": "time_projection.1.weight",
	17	+ "scale_shift_table": "head.modulation",
	18	+ "proj_out.bias": "head.head.bias",
	19	+ "proj_out.weight": "head.head.weight",
	20	+}
	21	+
	22	+# Sub-string replacements applied to every key
	23	+DIFFUSERS_STR_REPLACE_MAPPING = {
	24	+ "attn1.norm_q": "self_attn.norm_q",
	25	+ "attn1.norm_k": "self_attn.norm_k",
	26	+ "attn2.norm_q": "cross_attn.norm_q",
	H H htwangAug 6 后续可以考虑与新版本的权重转换方案做兼容 C chenpeizhe Aug 6 Comments： Aug 6 Comments：
	27	+ "attn2.norm_k": "cross_attn.norm_k",
	28	+ "attn1.to_q.": "self_attn.q.",
	29	+ "attn1.to_k.": "self_attn.k.",
	30	+ "attn1.to_v.": "self_attn.v.",
	31	+ "attn1.to_out.0.": "self_attn.o.",
	32	+ "attn2.to_q.": "cross_attn.q.",
	33	+ "attn2.to_k.": "cross_attn.k.",
	34	+ "attn2.to_v.": "cross_attn.v.",
	35	+ "attn2.to_out.0.": "cross_attn.o.",
	36	+ ".ffn.net.0.proj.": ".ffn.0.",
	37	+ ".ffn.net.2.": ".ffn.2.",
	38	+ "scale_shift_table": "modulation",
	39	+ ".norm2.": ".norm3.",
	40	+}
	41	+
	42	+
	43	+from mindspeed_mm.fsdp.checkpoint.convert import (
	44	+ DiffusersKeyMapTransformPipeline,
	45	+ WEIGHT_TRANSFORM_PIPELINES,
	46	+)
	atomgit-botAug 4 🔴 Critical diffusers_ckpt_trans.py 第43-46行从 `mindspeed_mm.fsdp.checkpoint.convert` 导入 `DiffusersKeyMapTransformPipeline`，但该模块（convert.py）中仅定义了 `WeightTransformPipeline` 和 `Qwen35WeightTransformPipeline`，不存在 `DiffusersKeyMapTransformPipeline` 类。全仓库搜索也未发现该类的定义。这会导致模块加载时抛出 `ImportError`，无法运行。触发条件：任何导入 `diffusers_ckpt_trans` 模块的代码路径都会立即失败。建议：在 `mindspeed_mm.fsdp.checkpoint.convert` 模块中添加 `DiffusersKeyMapTransformPipeline` 基类，实现基于 `CONVERT_MAPPING` 和 `STR_REPLACE_MAPPING` 的键转换逻辑；或将该基类与导入一起在本 PR 中提交。 Incorrect? C chenpeizhe Aug 5 Comments： Aug 5 Comments：
	47	+
	48	+
	49	+class Wan22DiffusersTransformPipeline(DiffusersKeyMapTransformPipeline):
	50	+ """Wan2.2 diffusers -> native transform pipeline, driven by the tables above."""
	51	+
	52	+ CONVERT_MAPPING = DIFFUSERS_CONVERT_MAPPING
	53	+ STR_REPLACE_MAPPING = DIFFUSERS_STR_REPLACE_MAPPING
	54	+
	55	+
	56	+WEIGHT_TRANSFORM_PIPELINES["wan2_2"] = Wan22DiffusersTransformPipeline

Amindspeed_mm/fsdp/models/wan2_2/diffusion.py+191-0

@@ -0,0 +1,191 @@

1 +from typing import Optional

2 +

3 +import torch

4 +from tqdm.auto import tqdm

5 +

6 +

7 +

8 +class WanFlowMatchScheduler():

9 +

10 + def __init__(

11 + self,

12 + num_inference_timesteps=None,

13 + num_train_timesteps=1000,

14 + shift=3.0,

15 + sigma_max=1.0,

16 + sigma_min=0.003 / 1.002,

17 + inverse_timesteps=False,

18 + extra_one_step=False,

19 + reverse_sigmas=False,

20 + guidance_scale=0.0,

21 + max_timestep_boundary=1.0,

22 + min_timestep_boundary=0.0,

23 + **kwargs

24 + ):

25 + self.num_inference_timesteps = num_inference_timesteps

26 + self.num_train_timesteps = num_train_timesteps

27 + self.shift = shift

28 + self.sigma_max = sigma_max

29 + self.sigma_min = sigma_min

30 + self.inverse_timesteps = inverse_timesteps

31 + self.extra_one_step = extra_one_step

32 + self.reverse_sigmas = reverse_sigmas

33 + self.guidance_scale = guidance_scale

34 + self.do_classifier_free_guidance = guidance_scale > 1.0

35 + if self.num_inference_timesteps is not None:

36 + self._set_timesteps(self.num_inference_timesteps, training=False)

37 + else:

38 + self._set_timesteps(self.num_train_timesteps, training=True)

39 +

40 + if not (0 <= min_timestep_boundary < max_timestep_boundary <= 1):

41 + raise ValueError("min_timestep_boundary and max_timestep_boundary must satisfy 0 <= min < max <= 1")

42 +

43 + self.max_timestep_boundary = max_timestep_boundary

44 + self.min_timestep_boundary = min_timestep_boundary

45 +

46 + def training_losses(

47 + self,

48 + model_output: torch.Tensor,

49 + x_start: Optional[torch.Tensor] = None,

50 + x_t: Optional[torch.Tensor] = None,

51 + noise: Optional[torch.Tensor] = None,

52 + mask: Optional[torch.Tensor] = None,

53 + t: Optional[torch.Tensor] = None,

54 + **kwargs

55 + ):

56 + reduction = kwargs.get("reduction", "mean")

57 + if noise is None or x_start is None or t is None:

58 + raise ValueError("noise, x_start and t are required for training_losses.")

59 + target = noise - x_start

60 + loss = torch.nn.functional.mse_loss(model_output.float(), target.float(), reduction=reduction)

atomgit-botAug 4

🟡 Medium Priority

diffusion.py 第60行 training_losses 方法中： target = noise - x_start noise 和 x_start 参数均声明为 Optional[torch.Tensor] = None（第52-53行），但函数体在未检查 None 的情况下直接执行减法运算。若调用方不传 noise 或 x_start，将抛出 TypeError: unsupported operand type(s) for -: 'NoneType' and 'NoneType'。

调用链：pipeline_wan2_2.py 第143-153行的 compute_loss 方法中，noise=inputs["noise"] 和 x_start=conditions["latents"] 均有值传递，所以在当前调用路径下不会触发。但该函数作为公共 API 被直接调用时，缺少输入校验会导致静默崩溃。

触发条件：直接调用 scheduler.training_losses(model_output=...) 而不传 noise 和 x_start。

建议：在 target = noise - x_start 之前添加 None 检查：if noise is None or x_start is None: raise ValueError(...)。

Incorrect?

chenpeizhe

Aug 5 Comments：

61 + loss *= self._training_weight(t)

62 + return loss

63 +

64 + def sample(

65 + self,

66 + model,

67 + latents,

68 + model_kwargs

69 + ):

70 + num_inference_steps = self.num_inference_timesteps or self.num_train_timesteps

71 + num_warmup_steps = len(self.timesteps) - num_inference_steps

72 +

73 + # for loop denoising to get clean latents

74 + with tqdm(total=num_inference_steps) as progress_bar:

75 + for i, t in enumerate(self.timesteps):

76 + latent_model_input = latents

77 + timestep = t.expand(latent_model_input.shape[0]).to(latents.device)

78 +

79 + noise_pred = model(

80 + latent_model_input,

81 + timestep,

82 + model_kwargs.get('prompt_embeds'),

83 + **model_kwargs

84 + )

85 +

86 + if self.do_classifier_free_guidance:

87 + noise_uncond = model(

88 + latent_model_input,

89 + timestep,

90 + model_kwargs.get('negative_prompt_embeds'),

91 + **model_kwargs

92 + )

93 + noise_pred = noise_uncond + self.guidance_scale * (noise_pred - noise_uncond)

94 +

95 + # compute the previous noisy sample x_t -> x_t-1

96 + latents = self._step(noise_pred, t, latents)

97 +

98 + if i == len(self.timesteps) - 1 or ((i + 1) > num_warmup_steps):

99 + progress_bar.update()

100 +

101 + return latents

102 +

103 + def q_sample(self, latents, noise=None, t=None, **kwargs):

104 + # Optional deterministic generator for cross-backend alignment.

105 + # When ``generator`` is provided, both noise and timestep are sampled

106 + # from this generator instead of the default RNG, which makes q_sample

107 + # independent of model-initialization side effects.

108 + generator = kwargs.get("generator", None)

109 +

110 + if noise is None:

111 + if generator is not None:

112 + noise = torch.randn(

113 + latents.shape,

114 + dtype=latents.dtype,

115 + device=latents.device,

116 + generator=generator,

117 + )

118 + else:

119 + noise = torch.randn_like(latents)

120 +

121 + # use first frame latents in noise

122 + model_kwargs = kwargs.get("model_kwargs", {})

123 + first_frame_latents = model_kwargs.get("first_frame_latents", None)

124 + if first_frame_latents is not None:

125 + noise[:, :, 0: 1] = first_frame_latents

126 +

127 + if t is None:

128 + max_timestep_boundary = int(self.max_timestep_boundary * self.num_train_timesteps)

129 + min_timestep_boundary = int(self.min_timestep_boundary * self.num_train_timesteps)

130 + if generator is not None:

131 + timestep_idx = torch.randint(

132 + min_timestep_boundary,

133 + max_timestep_boundary,

134 + (1,),

135 + device=generator.device,

136 + generator=generator,

137 + )

138 + else:

139 + timestep_idx = torch.randint(min_timestep_boundary, max_timestep_boundary, (1,))

140 + timestep_idx = timestep_idx.cpu()

141 + timestep = self.timesteps[timestep_idx].to(latents.device)

142 + else:

143 + timestep = t

144 + matches = (t.to("cpu") == self.timesteps).nonzero()

145 + # Fall back to the nearest grid point (same as ``_step``) when t is off-grid.

146 + if matches.numel() > 0:

147 + timestep_idx = matches[0]

148 + else:

149 + timestep_idx = torch.argmin((self.timesteps - t.to("cpu")).abs()).unsqueeze(0)

150 +

151 + sigma = self.sigmas[timestep_idx].to(latents.device)

152 + noised_latents = (1 - sigma) * latents + sigma * noise

153 + return noised_latents, noise, timestep

154 +

155 + def _set_timesteps(self, num_steps=100, training=False):

156 + sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min)

157 + if self.extra_one_step:

158 + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_steps + 1)[:-1]

159 + else:

160 + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_steps)

161 + if self.inverse_timesteps:

162 + self.sigmas = torch.flip(self.sigmas, dims=[0])

163 + self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)

164 + if self.reverse_sigmas:

165 + self.sigmas = 1 - self.sigmas

166 + self.timesteps = self.sigmas * self.num_train_timesteps

167 + if training:

168 + y = torch.exp(-2 * ((self.timesteps - num_steps / 2) / num_steps) ** 2)

169 + y_shifted = y - y.min()

170 + bsmntw_weighing = y_shifted * (num_steps / y_shifted.sum())

171 + self.linear_timesteps_weights = bsmntw_weighing

172 +

173 +

174 + def _step(self, model_output, timestep, sample):

175 + if isinstance(timestep, torch.Tensor):

176 + timestep = timestep.cpu()

177 + timestep_idx = torch.argmin((self.timesteps - timestep).abs())

178 + sigma = self.sigmas[timestep_idx]

179 + if timestep_idx + 1 >= len(self.timesteps):

180 + sigma_next = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0

181 + else:

182 + sigma_next = self.sigmas[timestep_idx + 1]

183 + prev_sample = sample + model_output * (sigma_next - sigma)

184 + return prev_sample

185 +

186 + def _training_weight(self, timestep):

187 + if timestep.numel() > 1:

188 + timestep = timestep[0]

atomgit-botAug 4

🟠 High Priority

diffusion.py 第179-181行 _training_weight 方法中：当 timestep 是一个 0-d tensor（标量 tensor，例如 batch_size=1 时 t 被索引后得到，或 timestep[0] 产生）时，len() 会抛出 TypeError: len() of a 0-d tensor。

该函数被 training_losses（第62行）调用，传入的 t 是 inputs["timesteps"]。对于 batch_size=1 的场景，t 可能是 0-d tensor，触发此崩溃。对于 batch_size>1，第一次调用 _training_weight 后 timestep[0] 产生 0-d tensor 赋值给局部变量 timestep，第二次调用同一实例时也会触发。

触发条件：训练时 batch_size=1，或者任何使得 t 为 0-d tensor 的场景。

建议：将 len(timestep) > 1 替换为 timestep.numel() > 1，避免 0-d tensor 上的 TypeError。

Suggestions

+    def _training_weight(self, timestep):
+        if timestep.numel() > 1:
-            timestep = timestep[0]
+            timestep = timestep[0]
+        timestep_idx = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs())
+        weights = self.linear_timesteps_weights[timestep_idx]
+        return weights

Apply suggestion

Incorrect?

chenpeizhe

Aug 5 Comments：

189 + timestep_idx = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs())

190 + weights = self.linear_timesteps_weights[timestep_idx]

191 + return weights

Amindspeed_mm/fsdp/models/wan2_2/pipeline_wan2_2.py+251-0

@@ -0,0 +1,251 @@
	1	+
	2	+"""Wan2.2 flow-matching T2V training pipeline (generic ``FlowMatchT2VPipeline`` + Wan2.2 specialization).
	3	+Data flow: condition encoding (grad-enabled for trainable sub-models) -> ``scheduler.q_sample`` noising -> predictor forward (grad) -> ``scheduler.training_losses``."""
	4	+
	5	+import torch
	6	+
	7	+from mindspeed_mm.fsdp.train.generative_pipeline import (
	8	+ GenerativePipelineBase,
	9	+ GenerativePipelineOutput,
	10	+ grad_context_for,
	11	+)
	atomgit-botAug 4 🔴 Critical pipeline_wan2_2.py 第8-11行从 `mindspeed_mm.fsdp.pipelines.generative_pipeline` 导入 `GenerativePipelineBase` 和 `GenerativePipelineOutput`，但 `mindspeed_mm/fsdp/pipelines/` 目录在仓库中完全不存在。`FlowMatchT2VPipeline` 继承自 `GenerativePipelineBase`（第18行），`compute_loss` 返回 `self.output_type(loss=loss)` 使用 `GenerativePipelineOutput`（第153行）。这导致模块加载时抛出 `ModuleNotFoundError` 或 `ImportError`。触发条件：任何导入 `pipeline_wan2_2` 的代码路径都会立即失败。建议：创建 `mindspeed_mm/fsdp/pipelines/generative_pipeline.py` 模块，定义 `GenerativePipelineBase` 基类和 `GenerativePipelineOutput` 数据类。`GenerativePipelineBase` 需要实现 `forward(self, container, args, *kwargs)` 编排方法；`GenerativePipelineOutput` 至少需要接受 `loss` 参数。 Incorrect? C chenpeizhe Aug 5 Comments： Aug 5 Comments：
	12	+from mindspeed_mm.fsdp.models.wan2_2.modeling_wan2_2 import (
	13	+ Wan2_2ModelOutput,
	14	+ get_wan_flow_match_scheduler,
	15	+)
	16	+
	17	+
	18	+class FlowMatchT2VPipeline(GenerativePipelineBase):
	19	+ """Flow-matching T2V training pipeline shared by video DiT models."""
	20	+
	21	+ # Sub-model role names on the container (default convention).
	22	+ ae_name = "ae"
	23	+ text_encoder_name = "text_encoder"
	24	+ predictor_name = "predictor"
	25	+
	26	+ # Output type wrapping the loss; overridden by model-specific pipelines.
	27	+ output_type = GenerativePipelineOutput
	28	+
	29	+ def __init__(self, scheduler):
	30	+ # Weight-free component only; the pipeline itself is not an nn.Module.
	31	+ self.scheduler = scheduler
	32	+
	33	+ # ------------------------------------------------------------------
	34	+ # Stage 1: condition encoding (grad-enabled for trainable sub-models)
	35	+ # ------------------------------------------------------------------
	36	+
	37	+ def prepare_conditions(
	38	+ self,
	39	+ container,
	40	+ video: torch.Tensor = None,
	41	+ prompt_ids: torch.Tensor = None,
	42	+ prompt_mask: torch.Tensor = None,
	43	+ video_mask: torch.Tensor = None,
	44	+ i2v_clip_feature: torch.Tensor = None,
	45	+ i2v_vae_feature: torch.Tensor = None,
	46	+ **kwargs,
	47	+ ) -> dict:
	48	+ ae = getattr(container, self.ae_name, None)
	49	+ text_encoder = getattr(container, self.text_encoder_name, None)
	50	+ with grad_context_for(ae):
	51	+ latents = self.encode_latents(container, video, kwargs)
	52	+ i2v_vae_feature = self.prepare_image_conditions(
	53	+ container, video, latents, i2v_vae_feature
	54	+ )
	55	+ with grad_context_for(text_encoder):
	56	+ prompt, prompt_mask = self.encode_prompt(
	57	+ container, video, latents, prompt_ids, prompt_mask, kwargs
	58	+ )
	59	+ return {
	60	+ "latents": latents,
	61	+ "prompt": prompt,
	62	+ "prompt_mask": prompt_mask,
	63	+ "i2v_clip_feature": i2v_clip_feature,
	64	+ "i2v_vae_feature": i2v_vae_feature,
	65	+ # Extra batch kwargs threaded through q_sample / predictor /
	66	+ # training_losses; hooks may update this dict in place.
	67	+ "extra_kwargs": kwargs,
	68	+ }
	69	+
	70	+ def encode_latents(self, container, video, extra_kwargs):
	71	+ """VAE-encode the video into latents (identity when no ae is attached)."""
	72	+ ae = getattr(container, self.ae_name, None)
	73	+ if ae is None:
	74	+ return video
	75	+ ae_dtype = next(ae.parameters()).dtype
	76	+ video_ae = video.to(ae_dtype) if video is not None else None
	77	+ encode_out = ae(video_ae)
	78	+ return encode_out[0] if isinstance(encode_out, tuple) else encode_out
	79	+
	80	+ def prepare_image_conditions(self, container, video, latents, i2v_vae_feature):
	81	+ """Hook for image-conditioning models; default is a pass-through."""
	82	+ return i2v_vae_feature
	83	+
	84	+ def encode_prompt(self, container, video, latents, prompt_ids, prompt_mask, extra_kwargs):
	85	+ """Text-encode the prompt (pass-through when no text encoder is attached)."""
	86	+ text_encoder = getattr(container, self.text_encoder_name, None)
	87	+ if text_encoder is None:
	88	+ return prompt_ids, prompt_mask
	89	+ if prompt_ids is None:
	90	+ return self.default_prompt(container, video, latents, prompt_mask)
	91	+ te_dtype = next(text_encoder.parameters()).dtype
	92	+ prompt_mask_te = (
	93	+ prompt_mask.to(te_dtype)
	94	+ if prompt_mask is not None and prompt_mask.is_floating_point()
	95	+ else prompt_mask
	96	+ )
	97	+ # Call through forward() so that FSDP2 pre/post hooks
	98	+ # (unshard/reshard) are triggered when the text encoder
	99	+ # is wrapped by FSDP2.
	100	+ prompt, prompt_mask = text_encoder(prompt_ids, prompt_mask_te, **extra_kwargs)
	101	+ return prompt, prompt_mask
	102	+
	103	+ def default_prompt(self, container, video, latents, prompt_mask):
	104	+ """Fallback prompt when a text encoder is attached but prompt_ids is None."""
	105	+ return None, prompt_mask
	106	+
	107	+ # ------------------------------------------------------------------
	108	+ # Stage 2: forward-diffusion noising
	109	+ # ------------------------------------------------------------------
	110	+
	111	+ def noising(self, container, conditions, **batch) -> dict:
	112	+ kwargs = conditions["extra_kwargs"]
	113	+ noised_latents, noise, timesteps = self.scheduler.q_sample(
	114	+ conditions["latents"],
	115	+ model_kwargs=kwargs,
	116	+ mask=batch.get("video_mask"),
	117	+ )
	118	+ inputs = dict(conditions)
	119	+ inputs.update(
	120	+ noised_latents=noised_latents,
	121	+ noise=noise,
	122	+ timesteps=timesteps,
	123	+ )
	124	+ return inputs
	125	+
	126	+ # ------------------------------------------------------------------
	127	+ # Stage 3: predictor forward (with gradients)
	128	+ # ------------------------------------------------------------------
	129	+
	130	+ def predict(self, container, **inputs):
	131	+ predictor = getattr(container, self.predictor_name)
	132	+ return predictor(
	133	+ inputs["noised_latents"],
	134	+ timestep=inputs["timesteps"],
	135	+ prompt=inputs["prompt"],
	136	+ prompt_mask=inputs["prompt_mask"],
	137	+ i2v_clip_feature=inputs["i2v_clip_feature"],
	138	+ i2v_vae_feature=inputs["i2v_vae_feature"],
	139	+ **inputs["extra_kwargs"],
	140	+ )
	141	+
	142	+ # ------------------------------------------------------------------
	143	+ # Stage 4: loss computation
	144	+ # ------------------------------------------------------------------
	145	+
	146	+ def compute_loss(self, container, conditions, inputs, model_output, **batch):
	147	+ loss = self.scheduler.training_losses(
	148	+ model_output=model_output,
	149	+ x_start=conditions["latents"],
	150	+ x_t=inputs["noised_latents"],
	151	+ noise=inputs["noise"],
	152	+ t=inputs["timesteps"],
	153	+ mask=batch.get("video_mask"),
	154	+ **inputs["extra_kwargs"],
	155	+ )
	156	+ return self.output_type(loss=loss)
	157	+
	158	+
	159	+class WanT2VPipeline(FlowMatchT2VPipeline):
	160	+ """Wan2.2 T2V/I2V training pipeline."""
	161	+
	162	+ output_type = Wan2_2ModelOutput
	163	+
	164	+ def __init__(self, task: str = "t2v"):
	165	+ super().__init__(scheduler=get_wan_flow_match_scheduler()())
	166	+ self.task = task
	167	+
	168	+ def encode_latents(self, container, video, extra_kwargs):
	169	+ ae = getattr(container, self.ae_name, None)
	170	+ if ae is None:
	171	+ return video
	172	+ ae_dtype = next(ae.parameters()).dtype
	173	+ video_ae = video.to(ae_dtype) if video is not None else None
	174	+ if self.task == "t2v":
	175	+ encode_out = ae(video_ae)
	176	+ latents = encode_out[0] if isinstance(encode_out, tuple) else encode_out
	177	+ elif self.task == "i2v":
	178	+ ae_kwargs = {
	179	+ k: v for k, v in extra_kwargs.items()
	180	+ if k not in ("first_frame", "use_cache")
	181	+ }
	182	+ encode_out = ae(video_ae, **ae_kwargs)
	183	+ latents = encode_out[0] if isinstance(encode_out, tuple) else encode_out
	184	+ i2v_results = (
	185	+ encode_out[1]
	186	+ if isinstance(encode_out, tuple) and len(encode_out) > 1
	187	+ else None
	188	+ )
	189	+ if i2v_results is not None:
	190	+ extra_kwargs.update(i2v_results)
	191	+ else:
	192	+ raise NotImplementedError(f"Task {self.task} is not implemented!")
	193	+ return latents
	194	+
	195	+ def prepare_image_conditions(self, container, video, latents, i2v_vae_feature):
	196	+ if (
	197	+ self.task == "i2v"
	198	+ and i2v_vae_feature is None
	199	+ and video is not None
	200	+ ):
	201	+ i2v_vae_feature = self._prepare_i2v_condition(container, video, latents)
	202	+ return i2v_vae_feature
	203	+
	204	+ def default_prompt(self, container, video, latents, prompt_mask):
	205	+ B = (
	206	+ video.shape[0]
	207	+ if video is not None
	208	+ else latents.shape[0]
	209	+ )
	210	+ device = (
	211	+ video.device
	212	+ if video is not None
	213	+ else latents.device
	214	+ )
	215	+ text_dim = getattr(container, self.predictor_name).text_dim
	216	+ prompt = torch.zeros(
	217	+ (B, 1, text_dim), device=device, dtype=torch.float32
	218	+ )
	219	+ return prompt, prompt_mask
	220	+
	221	+ def _prepare_i2v_condition(
	222	+ self, container, video: torch.Tensor, latents: torch.Tensor
	223	+ ) -> torch.Tensor:
	224	+ B, C, T, H, W = video.shape
	225	+ _, _, T_latent, H_latent, W_latent = latents.shape
	226	+
	227	+ first_frame = video[:, :, 0:1, :, :]
	228	+ zeros = torch.zeros(
	229	+ B, C, T - 1, H, W, device=video.device, dtype=video.dtype
	230	+ )
	231	+ vae_input = torch.cat([first_frame, zeros], dim=2)
	232	+
	233	+ ae = getattr(container, self.ae_name, None)
	234	+ if ae is None:
	235	+ raise RuntimeError("AE model is required for the I2V task but is not attached to the container.")
	236	+ with grad_context_for(ae):
	237	+ latent_condition = ae(vae_input)
	238	+ if isinstance(latent_condition, tuple):
	239	+ latent_condition = latent_condition[0]
	atomgit-botAug 4 🟡 Medium Priority pipeline_wan2_2.py 第230-234行 `_prepare_i2v_condition` 方法中： ae = getattr(container, self.ae_name, None) with torch.no_grad(): latent_condition = ae(vae_input) 第230行通过 `getattr(..., None)` 获取 `ae` 允许为 `None`，但第232行直接调用 `ae(vae_input)` 未检查 `None`。若在 I2V 任务模式下容器未注册 `ae` 子模型，将抛出 `TypeError: 'NoneType' object is not callable`。注意：父类 `FlowMatchT2VPipeline.encode_latents`（第69行）对 `ae` 做了 None 检查后提前返回，但子类 `WanT2VPipeline.encode_latents`（第166-167行）同样会检查。然而 `_prepare_i2v_condition` 被 `prepare_image_conditions`（第192-199行）直接调用，走的是独立路径。当 `prepare_image_conditions` 判断 `self.task == "i2v" and i2v_vae_feature is None and video is not None` 时进入此函数，此时 `ae` 若为 None 则崩溃。触发条件：I2V 任务模式下，容器未配置 `ae` 子模型。建议：在第230行之后添加 `if ae is None: raise RuntimeError("AE model is required for I2V task")` 或 `return None`。 Incorrect? C chenpeizhe Aug 5 Comments： Aug 5 Comments：
	240	+
	241	+ msk = torch.ones(
	242	+ B, T, H_latent, W_latent, device=latents.device, dtype=latents.dtype
	243	+ )
	244	+ msk[:, 1:] = 0
	245	+ msk = torch.cat(
	246	+ [msk[:, 0:1].repeat(1, 4, 1, 1), msk[:, 1:]], dim=1
	247	+ )
	248	+ msk = msk.view(B, (T + 3) // 4, 4, H_latent, W_latent).transpose(1, 2)
	249	+
	250	+ i2v_vae_feature = torch.cat([msk, latent_condition], dim=1)
	251	+ return i2v_vae_feature

Amindspeed_mm/fsdp/models/wan2_2/wan2_2_model_container.py+143-0

@@ -0,0 +1,143 @@

1 +from typing import Dict

2 +

3 +import torch

4 +import torch.nn as nn

5 +

6 +from mindspeed_mm.fsdp.models.base_model import GenerativeBaseModel

7 +from mindspeed_mm.fsdp.models.model_container import ModelContainer

8 +from mindspeed_mm.fsdp.models.wan2_2.modeling_wan2_2 import (

9 + WanDiTFSDP2,

10 + get_wan_video_vae,

11 +)

12 +from mindspeed_mm.fsdp.models.wan2_2.text_encoder import build_text_encoder

atomgit-botAug 4

🔴 Critical

变更行第 12 行：from mindspeed_mm.fsdp.models.wan2_2.text_encoder import build_text_encoder

build_text_encoder 实际定义在 mindspeed_mm/fsdp/models/wan2_2/modeling_wan2_2.py 第 161 行。在 wan2_2 目录下不存在 text_encoder.py 文件（该目录只有 diffusers_ckpt_trans.py、diffusion.py、modeling_wan2_2.py、pipeline_wan2_2.py 和 wan2_2_model_container.py）。

此导入在模块加载时会立即抛出 ModuleNotFoundError，导致整个模块无法被导入，进而使 UMT5Bridge._from_config（第 65 行调用 build_text_encoder(cfg)）无法执行。此为 P0 级阻断性问题。

建议：将导入路径修正为 mindspeed_mm.fsdp.models.wan2_2.modeling_wan2_2，或将 build_text_encoder 合并到已有的 modeling_wan2_2 导入语句中（第 8-11 行）。

Suggestions

-from mindspeed_mm.fsdp.models.wan2_2.text_encoder import build_text_encoder
+from mindspeed_mm.fsdp.models.wan2_2.modeling_wan2_2 import (
+    WanDiTFSDP2,
+    get_wan_video_vae,
+    build_text_encoder,
+)

Apply suggestion

Incorrect?

chenpeizhe

Aug 5 Comments：

13 +from mindspeed_mm.fsdp.models.wan2_2.pipeline_wan2_2 import WanT2VPipeline

14 +from mindspeed_mm.fsdp.utils.register import model_register

15 +from mindspeed_mm.fsdp.utils.dtype import get_dtype

16 +

17 +

18 +def _config_to_dict(config):

19 + return config.to_dict() if hasattr(config, "to_dict") else dict(config)

20 +

21 +

22 +def _resolve_torch_dtype(value):

23 + if value in (None, ""):

24 + return None

25 + if isinstance(value, torch.dtype):

26 + return value

27 + return get_dtype(str(value).strip().lower())

28 +

29 +

30 +# ---------------------------------------------------------------------------

31 +# Bridge classes for sub-models that are NOT registered in model_register

32 +# by the original codebase. These bridges allow ModelHub._build_custom_model

33 +# to construct ae / text_encoder / predictor independently in the new

34 +# model-list mode while keeping the legacy single-model path untouched.

35 +# ---------------------------------------------------------------------------

36 +

37 +@model_register.register("wan_video_vae")

38 +class WanVideoVAEBridge(GenerativeBaseModel):

39 + """Bridge to build WanVideoVAE (ae) in model-list mode."""

40 +

41 + @classmethod

42 + def _from_config(cls, config):

43 + cfg = _config_to_dict(config)

44 + cfg.pop("freeze", None)

45 + cfg.pop("meta_init", None)

46 + target_dtype = _resolve_torch_dtype(cfg.get("dtype", None))

47 + if target_dtype is not None:

48 + cfg["dtype"] = target_dtype

49 + return get_wan_video_vae()(**cfg)

50 +

51 + @classmethod

52 + def from_pretrained(cls, config):

53 + return cls._from_config(config)

54 +

55 +

56 +@model_register.register("UMT5")

57 +class UMT5Bridge(GenerativeBaseModel):

58 + """Bridge to build UMT5 text encoder in model-list mode."""

59 +

60 + @classmethod

61 + def _from_config(cls, config):

62 + cfg = _config_to_dict(config)

63 + cfg.pop("freeze", None)

64 + cfg.pop("meta_init", None)

65 + return build_text_encoder(cfg)

66 +

67 + @classmethod

68 + def from_pretrained(cls, config):

69 + return cls._from_config(config)

70 +

71 +

72 +class WanDiTFSDP2WithSetup(WanDiTFSDP2, GenerativeBaseModel):

73 + """WanDiTFSDP2 subclass keeping only the Wan-specific dtype re-cast and init-skipping guard on top of the ``GenerativeBaseModel`` default ``setup_weights``."""

74 +

75 + weight_transform_model_type = "wan2_2"

76 +

77 + @classmethod

78 + def from_pretrained(cls, config):

79 + raise NotImplementedError(

80 + "WanDiTFSDP2WithSetup is constructed via WanDiTFSDP2Bridge."

81 + )

82 +

83 + @classmethod

84 + def _from_config(cls, config):

85 + raise NotImplementedError(

86 + "WanDiTFSDP2WithSetup is constructed via WanDiTFSDP2Bridge."

87 + )

88 +

89 + def post_load(self, loaded: bool):

90 + target_dtype = getattr(self, "_wan_target_dtype", None)

91 + if target_dtype is not None and target_dtype != torch.float32:

92 + self.to(dtype=target_dtype)

93 +

94 + def init_weights(self):

95 + if getattr(self, "_weights_loaded", False):

96 + return

97 + super().init_weights()

98 +

99 +

100 +@model_register.register("wan2_2_predictor")

101 +class WanDiTFSDP2Bridge(GenerativeBaseModel):

102 + """Bridge to build the native WanDiTFSDP2 predictor in model-list mode."""

103 +

104 + @classmethod

105 + def _from_config(cls, config):

106 + cfg = _config_to_dict(config)

107 + cfg.pop("freeze", None)

108 + cfg.pop("meta_init", None)

109 + target_dtype = _resolve_torch_dtype(cfg.pop("dtype", None))

110 + model_backend = cfg.pop("model_backend", None)

111 + if model_backend not in (None, "", "native"):

112 + raise ValueError(

113 + "Wan2.2 FSDP2 ModelContainer only supports the native predictor. "

114 + "Remove 'model_backend' or set it to 'native'."

115 + )

116 + model = WanDiTFSDP2WithSetup(**cfg)

117 + if target_dtype is not None:

118 + model._wan_target_dtype = target_dtype

119 + model.to(dtype=target_dtype)

120 + return model

121 +

122 + @classmethod

123 + def from_pretrained(cls, config):

124 + return cls._from_config(config)

125 +

126 +

127 +@model_register.register("wan2_2_container")

128 +class Wan2_2ModelContainer(ModelContainer):

129 + """Wan2.2 multi-model container managing ae / text_encoder / predictor lifecycle; data flow is delegated to :class:`WanT2VPipeline`."""

130 +

131 + # Checkpoint policy is dynamic (trainable sub-models only); flat

132 + # (unprefixed) keys are treated as predictor weights for backward

133 + # compatibility with older checkpoints.

134 + legacy_flat_key_target = "predictor"

135 +

136 + def __init__(self, models: Dict[str, nn.Module], task: str = "t2v"):

137 + super().__init__(models)

138 + self.task = task

139 + # Inter-model data-flow orchestration (condition encoding, noising,

140 + # predictor forward, loss) lives in the pipeline; the container only

141 + # owns component management. The pipeline is a plain object (not an

142 + # nn.Module), so named_modules() paths and FSDP patterns are unaffected.

143 + self.pipeline = WanT2VPipeline(task=task)