From 58229459476d6f8e1af45bf666c0729d0d2f2ab4 Mon Sep 17 00:00:00 2001 From: SHIYANG CHENG Date: Tue, 10 Nov 2026 16:32:35 -0800 Subject: [PATCH 1/8] refactor(minimax-h3): move the DiT TP rewrite to models/ for both neuron backends The TP rewrite (head-sharded attention, split SwiGLU, column-sharded text embedder/norm_out/adaln_proj, the bounded-attention processor and the trace module) only uses difflet.ops, so it moves to difflet/models/minimax_h3/modeling_minimax_h3_tp.py. The Trainium module re-exports every name it had, so its graph is built from the same objects. The fused-SwiGLU key split becomes split_fused_swiglu_keys/ff_prefixes. Co-Authored-By: Claude Opus 5.5 --- .../trainium/minimax_h3/transformer.py | 274 ++-------------- .../minimax_h3/modeling_minimax_h3_tp.py | 294 ++++++++++++++++++ .../models/minimax_h3/test_modeling_tp.py | 28 ++ .../minimax_h3/test_trainium_transformer.py | 14 + 4 files changed, 356 insertions(+), 254 deletions(-) create mode 100644 difflet/models/minimax_h3/modeling_minimax_h3_tp.py create mode 100644 tests/unit/models/minimax_h3/test_modeling_tp.py diff --git a/difflet/backends/trainium/minimax_h3/transformer.py b/difflet/backends/trainium/minimax_h3/transformer.py index 87aba4901..b52d2da8e 100644 --- a/difflet/backends/trainium/minimax_h3/transformer.py +++ b/difflet/backends/trainium/minimax_h3/transformer.py @@ -2,12 +2,10 @@ from __future__ import annotations -import math import os from typing import List import torch -from torch import nn from difflet.backends.trainium.core.application_base import NeuronApplicationBase from difflet.backends.trainium.core.config import InferenceConfig @@ -18,11 +16,22 @@ build_padded_t2va_layout, video_latent_num_frames, ) -from difflet.ops import ( - ColumnParallelLinear, - RowParallelLinear, - attention, - get_tensor_model_parallel_size, +# The TP rewrite is backend-neutral and shared with the neuron backend; re-exported here under +# the names this module has always had. +from difflet.models.minimax_h3.modeling_minimax_h3_tp import ( # noqa: F401 + MiniMaxH3TPAttnProcessor, + _column_parallel_like, + _H3TensorParallelSwiGLU, + _MiniMaxH3TrainiumAttnProcessor, + _MiniMaxH3TransformerTraceModule, + _prefix_bounds, + _replace_attention_for_tp, + _replace_block_for_tp, + _replace_h3_linears_for_tp, + _row_parallel_like, + _safe_tp_size, + ff_prefixes, + split_fused_swiglu_keys, ) @@ -128,235 +137,6 @@ def validate_config(self): ) -class _H3TensorParallelSwiGLU(nn.Module): - """Shard both SwiGLU branches without mixing their checkpoint row layout.""" - - def __init__(self, hidden_size: int, ffn_dim: int) -> None: - super().__init__() - self.up_proj = ColumnParallelLinear(hidden_size, ffn_dim, bias=False, gather_output=False) - self.gate_proj = ColumnParallelLinear(hidden_size, ffn_dim, bias=False, gather_output=False) - self.down_proj = RowParallelLinear(ffn_dim, hidden_size, bias=False, input_is_parallel=True) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.down_proj( - self.up_proj(hidden_states) * torch.nn.functional.silu(self.gate_proj(hidden_states)) - ) - - -def _column_parallel_like(linear: nn.Linear, *, gather_output: bool) -> ColumnParallelLinear: - return ColumnParallelLinear( - linear.in_features, - linear.out_features, - bias=linear.bias is not None, - gather_output=gather_output, - ) - - -def _row_parallel_like(linear: nn.Linear) -> RowParallelLinear: - return RowParallelLinear( - linear.in_features, - linear.out_features, - bias=linear.bias is not None, - input_is_parallel=True, - ) - - -def _safe_tp_size() -> int: - try: - return int(get_tensor_model_parallel_size()) - except AssertionError: - return 1 - - -def _replace_attention_for_tp(attn, tp_degree: int) -> None: - if int(attn.heads) % tp_degree: - raise ValueError(f"MiniMax-H3 attention heads {attn.heads} must divide TP{tp_degree}.") - attn.heads = int(attn.heads) // tp_degree - attn.to_q = _column_parallel_like(attn.to_q, gather_output=False) - attn.to_k = _column_parallel_like(attn.to_k, gather_output=False) - attn.to_v = _column_parallel_like(attn.to_v, gather_output=False) - attn.to_out[0] = _row_parallel_like(attn.to_out[0]) - attn.processor = _MiniMaxH3TrainiumAttnProcessor() - - -def _replace_block_for_tp(block, tp_degree: int, hidden_size: int, ffn_dim: int) -> None: - _replace_attention_for_tp(block.attn, tp_degree) - block.ff = _H3TensorParallelSwiGLU(hidden_size, ffn_dim) - - -def _replace_h3_linears_for_tp(transformer: nn.Module) -> None: - tp_degree = _safe_tp_size() - if tp_degree <= 1: - return - - hidden_size = int(transformer.config.hidden_size) - ffn_dim = int(transformer.config.ffn_dim) - transformer.context_embedder = _column_parallel_like( - transformer.context_embedder, gather_output=True - ) - transformer.norm_out.linear = _column_parallel_like( - transformer.norm_out.linear, gather_output=True - ) - for block in transformer.token_refiner.refiner_blocks: - _replace_block_for_tp(block, tp_degree, hidden_size, ffn_dim) - for block in transformer.transformer_blocks: - _replace_block_for_tp(block, tp_degree, hidden_size, ffn_dim) - # The 13B AdaLN branch is sharded in HBM; its tiny per-step output is - # gathered because the residual stream remains replicated after TP. - # Absent entirely when the modulation table is precomputed on the host. - if block.adaln_proj is not None: - block.adaln_proj.linear = _column_parallel_like( - block.adaln_proj.linear, gather_output=True - ) - - -def _prefix_bounds( - attention_mask: torch.Tensor, - *, - num_heads: int, - query_length: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Compact prefix-validity mask to attention_cte per-query bounds.""" - - valid = attention_mask if attention_mask.dtype == torch.bool else attention_mask != 0 - count = valid.sum(dim=-1, keepdim=True).to(torch.int32) - batch_size = int(valid.shape[0]) - bound_max = ( - count[:, None, None, :] - .expand(batch_size, num_heads, query_length, 1) - .reshape(batch_size * num_heads, query_length, 1) - .contiguous() - ) - return torch.zeros_like(bound_max), bound_max - - -class _MiniMaxH3TrainiumAttnProcessor: - """Full H3 self-attention through attention_cte, including tail bounds.""" - - _attention_backend = None - _parallel_config = None - - def __call__( - self, - attn, - hidden_states: torch.Tensor, - rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, - attention_mask: torch.Tensor | None = None, - ) -> torch.Tensor: - from difflet.models.minimax_h3.modeling_minimax_h3 import _apply_rotary_emb - - query = attn.to_q(hidden_states).unflatten(-1, (attn.heads, -1)) - key = attn.to_k(hidden_states).unflatten(-1, (attn.heads, -1)) - value = attn.to_v(hidden_states).unflatten(-1, (attn.heads, -1)) - query = attn.norm_q(query) - key = attn.norm_k(key) - if rotary_emb is not None: - query = _apply_rotary_emb(query, *rotary_emb) - key = _apply_rotary_emb(key, *rotary_emb) - - q = query.transpose(1, 2) - k = key.transpose(1, 2) - v = value.transpose(1, 2) - batch_size, num_heads, query_length, head_dim = q.shape - key_length = int(k.shape[2]) - kwargs = {} - if attention_mask is not None: - bound_min, bound_max = _prefix_bounds( - attention_mask, - num_heads=num_heads, - query_length=query_length, - ) - kwargs.update(bound_min=bound_min, bound_max=bound_max) - output = attention( - q.reshape(batch_size * num_heads, query_length, head_dim), - k.reshape(batch_size * num_heads, key_length, head_dim), - v.reshape(batch_size * num_heads, key_length, head_dim), - scale=1.0 / math.sqrt(head_dim), - causal=False, - tp_q=True, - tp_k=True, - tp_out=False, - **kwargs, - ) - output = output.reshape(batch_size, num_heads, query_length, head_dim) - output = output.transpose(1, 2).flatten(2, 3).to(query.dtype) - output = attn.to_out[0](output.contiguous()) - return attn.to_out[1](output) - - -class _MiniMaxH3TransformerTraceModule(nn.Module): - def __init__(self, config: MiniMaxH3TransformerInferenceConfig) -> None: - super().__init__() - from difflet.models.minimax_h3.modeling_minimax_h3 import ( - MiniMaxH3Transformer3DModel, - ) - - self.config = config - self.transformer = MiniMaxH3Transformer3DModel( - num_attention_heads=int(config.num_attention_heads), - attention_head_dim=int(config.attention_head_dim), - hidden_size=int(config.hidden_size), - num_layers=int(config.num_layers), - num_refiner_layers=int(config.num_refiner_layers), - ffn_dim=int(config.ffn_dim), - in_channels=int(config.in_channels), - audio_in_channels=int(config.audio_in_channels), - patch_size=tuple(config.patch_size), - text_dim=int(config.text_dim), - freq_dim=int(config.freq_dim), - time_embed_hidden_dim=int(config.time_embed_hidden_dim), - time_embed_dim=int(config.time_embed_dim), - rope_freq_dim=int(config.rope_freq_dim), - rope_theta=float(config.rope_theta), - norm_eps=float(config.norm_eps), - qk_norm_eps=float(config.qk_norm_eps), - final_norm_eps=float(config.final_norm_eps), - precomputed_adaln=bool(config.precomputed_adaln), - ) - _replace_h3_linears_for_tp(self.transformer) - - def restore_checkpoint_fp32_modules(self) -> None: - transformer = self.transformer - transformer.proj_in.float() - transformer.audio_proj_in.float() - transformer.time_embedder.float() - transformer.proj_out.float() - transformer.audio_proj_out.float() - - def forward( - self, - hidden_states, - audio_hidden_states, - encoder_hidden_states, - timestep, - timestep_indices, - token_tags, - position_ids, - video_indices, - audio_indices, - text_indices, - encoder_attention_mask, - attention_mask, - block_modulation=None, - ): - return self.transformer( - hidden_states=hidden_states, - audio_hidden_states=audio_hidden_states, - encoder_hidden_states=encoder_hidden_states, - timestep=timestep, - timestep_indices=timestep_indices, - token_tags=token_tags, - position_ids=position_ids, - video_indices=video_indices, - audio_indices=audio_indices, - text_indices=text_indices, - encoder_attention_mask=encoder_attention_mask, - attention_mask=attention_mask, - block_modulation=block_modulation, - return_dict=False, - ) - - class ModelWrapperMiniMaxH3Transformer(ModelWrapper): def __init__( self, @@ -504,24 +284,10 @@ def convert_hf_to_neuron_state_dict( converted = { key: value for key, value in converted.items() if ".adaln_proj." not in key } - prefixes = [ - *[ - f"transformer.token_refiner.refiner_blocks.{index}.ff" - for index in range(int(config.num_refiner_layers)) - ], - *[ - f"transformer.transformer_blocks.{index}.ff" - for index in range(int(config.num_layers)) - ], - ] - for prefix in prefixes: - fused_key = f"{prefix}.net.0.proj.weight" - if fused_key not in converted: - continue - up_weight, gate_weight = converted.pop(fused_key).chunk(2, dim=0) - converted[f"{prefix}.up_proj.weight"] = up_weight.contiguous() - converted[f"{prefix}.gate_proj.weight"] = gate_weight.contiguous() - converted[f"{prefix}.down_proj.weight"] = converted.pop(f"{prefix}.net.2.weight") + split_fused_swiglu_keys( + converted, + ff_prefixes(int(config.num_refiner_layers), int(config.num_layers)), + ) return converted @staticmethod diff --git a/difflet/models/minimax_h3/modeling_minimax_h3_tp.py b/difflet/models/minimax_h3/modeling_minimax_h3_tp.py new file mode 100644 index 000000000..689367eb1 --- /dev/null +++ b/difflet/models/minimax_h3/modeling_minimax_h3_tp.py @@ -0,0 +1,294 @@ +"""Backend-neutral tensor-parallel rewrite of the MiniMax-H3 Omni Transformer. + +Shared by the Trainium AoT backend (``difflet/backends/trainium/minimax_h3/transformer.py``) and +the PyTorch-native neuron backend (``difflet/backends/neuron/minimax_h3/transformer.py``). Every +layer here comes from ``difflet.ops``, so importing this module binds ``difflet.ops`` to the +selected backend. ``difflet.ops`` is therefore imported inside the functions that build layers +(after the backend's mesh exists), never at module import. + +* ``_replace_h3_linears_for_tp`` head-shards every attention (56 heads / tp), splits each fused + SwiGLU into column-parallel up/gate and a row-parallel down projection, and column-shards the + text embedder, the output norm's modulation and (when present) the 13B ``adaln_proj``. +* ``MiniMaxH3TPAttnProcessor`` runs the self-attention through ``difflet.ops.attention`` with the + packed sequence's padded tail lowered to per-query key bounds. +* ``_MiniMaxH3TransformerTraceModule`` builds the diffusers-layout model from any config object + carrying the transformer's hyper-parameters and applies the TP rewrite. +""" + +from __future__ import annotations + +import math + +import torch +from torch import nn + + +class _H3TensorParallelSwiGLU(nn.Module): + """Shard both SwiGLU branches without mixing their checkpoint row layout.""" + + def __init__(self, hidden_size: int, ffn_dim: int) -> None: + from difflet.ops import ColumnParallelLinear, RowParallelLinear + + super().__init__() + self.up_proj = ColumnParallelLinear(hidden_size, ffn_dim, bias=False, gather_output=False) + self.gate_proj = ColumnParallelLinear(hidden_size, ffn_dim, bias=False, gather_output=False) + self.down_proj = RowParallelLinear(ffn_dim, hidden_size, bias=False, input_is_parallel=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.down_proj( + self.up_proj(hidden_states) * torch.nn.functional.silu(self.gate_proj(hidden_states)) + ) + + +def _column_parallel_like(linear: nn.Linear, *, gather_output: bool) -> nn.Module: + from difflet.ops import ColumnParallelLinear + + return ColumnParallelLinear( + linear.in_features, + linear.out_features, + bias=linear.bias is not None, + gather_output=gather_output, + ) + + +def _row_parallel_like(linear: nn.Linear) -> nn.Module: + from difflet.ops import RowParallelLinear + + return RowParallelLinear( + linear.in_features, + linear.out_features, + bias=linear.bias is not None, + input_is_parallel=True, + ) + + +def _safe_tp_size() -> int: + from difflet.ops import get_tensor_model_parallel_size + + try: + return int(get_tensor_model_parallel_size()) + except AssertionError: + return 1 + + +def _replace_attention_for_tp(attn, tp_degree: int) -> None: + if int(attn.heads) % tp_degree: + raise ValueError(f"MiniMax-H3 attention heads {attn.heads} must divide TP{tp_degree}.") + attn.heads = int(attn.heads) // tp_degree + attn.to_q = _column_parallel_like(attn.to_q, gather_output=False) + attn.to_k = _column_parallel_like(attn.to_k, gather_output=False) + attn.to_v = _column_parallel_like(attn.to_v, gather_output=False) + attn.to_out[0] = _row_parallel_like(attn.to_out[0]) + attn.processor = MiniMaxH3TPAttnProcessor() + + +def _replace_block_for_tp(block, tp_degree: int, hidden_size: int, ffn_dim: int) -> None: + _replace_attention_for_tp(block.attn, tp_degree) + block.ff = _H3TensorParallelSwiGLU(hidden_size, ffn_dim) + + +def _replace_h3_linears_for_tp(transformer: nn.Module) -> None: + tp_degree = _safe_tp_size() + if tp_degree <= 1: + return + + hidden_size = int(transformer.config.hidden_size) + ffn_dim = int(transformer.config.ffn_dim) + transformer.context_embedder = _column_parallel_like( + transformer.context_embedder, gather_output=True + ) + transformer.norm_out.linear = _column_parallel_like( + transformer.norm_out.linear, gather_output=True + ) + for block in transformer.token_refiner.refiner_blocks: + _replace_block_for_tp(block, tp_degree, hidden_size, ffn_dim) + for block in transformer.transformer_blocks: + _replace_block_for_tp(block, tp_degree, hidden_size, ffn_dim) + # The 13B AdaLN branch is sharded in HBM; its tiny per-step output is + # gathered because the residual stream remains replicated after TP. + # Absent entirely when the modulation table is precomputed on the host. + if block.adaln_proj is not None: + block.adaln_proj.linear = _column_parallel_like( + block.adaln_proj.linear, gather_output=True + ) + + +def _prefix_bounds( + attention_mask: torch.Tensor, + *, + num_heads: int, + query_length: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compact prefix-validity mask to attention_cte per-query bounds.""" + + valid = attention_mask if attention_mask.dtype == torch.bool else attention_mask != 0 + count = valid.sum(dim=-1, keepdim=True).to(torch.int32) + batch_size = int(valid.shape[0]) + bound_max = ( + count[:, None, None, :] + .expand(batch_size, num_heads, query_length, 1) + .reshape(batch_size * num_heads, query_length, 1) + .contiguous() + ) + return torch.zeros_like(bound_max), bound_max + + +class MiniMaxH3TPAttnProcessor: + """Full H3 self-attention through attention_cte, including tail bounds.""" + + _attention_backend = None + _parallel_config = None + + def __init__(self) -> None: + # Bound once here, not imported inside the (compiled) call. + from difflet.models.minimax_h3.modeling_minimax_h3 import _apply_rotary_emb + from difflet.ops import attention + + self._attention = attention + self._apply_rotary_emb = _apply_rotary_emb + + def __call__( + self, + attn, + hidden_states: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + _apply_rotary_emb = self._apply_rotary_emb + query = attn.to_q(hidden_states).unflatten(-1, (attn.heads, -1)) + key = attn.to_k(hidden_states).unflatten(-1, (attn.heads, -1)) + value = attn.to_v(hidden_states).unflatten(-1, (attn.heads, -1)) + query = attn.norm_q(query) + key = attn.norm_k(key) + if rotary_emb is not None: + query = _apply_rotary_emb(query, *rotary_emb) + key = _apply_rotary_emb(key, *rotary_emb) + + q = query.transpose(1, 2) + k = key.transpose(1, 2) + v = value.transpose(1, 2) + batch_size, num_heads, query_length, head_dim = q.shape + key_length = int(k.shape[2]) + kwargs = {} + if attention_mask is not None: + bound_min, bound_max = _prefix_bounds( + attention_mask, + num_heads=num_heads, + query_length=query_length, + ) + kwargs.update(bound_min=bound_min, bound_max=bound_max) + output = self._attention( + q.reshape(batch_size * num_heads, query_length, head_dim), + k.reshape(batch_size * num_heads, key_length, head_dim), + v.reshape(batch_size * num_heads, key_length, head_dim), + scale=1.0 / math.sqrt(head_dim), + causal=False, + tp_q=True, + tp_k=True, + tp_out=False, + **kwargs, + ) + output = output.reshape(batch_size, num_heads, query_length, head_dim) + output = output.transpose(1, 2).flatten(2, 3).to(query.dtype) + output = attn.to_out[0](output.contiguous()) + return attn.to_out[1](output) + + +class _MiniMaxH3TransformerTraceModule(nn.Module): + def __init__(self, config) -> None: + super().__init__() + from difflet.models.minimax_h3.modeling_minimax_h3 import ( + MiniMaxH3Transformer3DModel, + ) + + self.config = config + self.transformer = MiniMaxH3Transformer3DModel( + num_attention_heads=int(config.num_attention_heads), + attention_head_dim=int(config.attention_head_dim), + hidden_size=int(config.hidden_size), + num_layers=int(config.num_layers), + num_refiner_layers=int(config.num_refiner_layers), + ffn_dim=int(config.ffn_dim), + in_channels=int(config.in_channels), + audio_in_channels=int(config.audio_in_channels), + patch_size=tuple(config.patch_size), + text_dim=int(config.text_dim), + freq_dim=int(config.freq_dim), + time_embed_hidden_dim=int(config.time_embed_hidden_dim), + time_embed_dim=int(config.time_embed_dim), + rope_freq_dim=int(config.rope_freq_dim), + rope_theta=float(config.rope_theta), + norm_eps=float(config.norm_eps), + qk_norm_eps=float(config.qk_norm_eps), + final_norm_eps=float(config.final_norm_eps), + precomputed_adaln=bool(config.precomputed_adaln), + ) + _replace_h3_linears_for_tp(self.transformer) + + def restore_checkpoint_fp32_modules(self) -> None: + transformer = self.transformer + transformer.proj_in.float() + transformer.audio_proj_in.float() + transformer.time_embedder.float() + transformer.proj_out.float() + transformer.audio_proj_out.float() + + def forward( + self, + hidden_states, + audio_hidden_states, + encoder_hidden_states, + timestep, + timestep_indices, + token_tags, + position_ids, + video_indices, + audio_indices, + text_indices, + encoder_attention_mask, + attention_mask, + block_modulation=None, + ): + return self.transformer( + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + timestep_indices=timestep_indices, + token_tags=token_tags, + position_ids=position_ids, + video_indices=video_indices, + audio_indices=audio_indices, + text_indices=text_indices, + encoder_attention_mask=encoder_attention_mask, + attention_mask=attention_mask, + block_modulation=block_modulation, + return_dict=False, + ) + + +# The Trainium module's historical name for the processor. +_MiniMaxH3TrainiumAttnProcessor = MiniMaxH3TPAttnProcessor + + +def split_fused_swiglu_keys(state_dict: dict, prefixes) -> dict: + """Rename each ``{prefix}.net.0.proj.weight`` (diffusers' fused SwiGLU, rows ``[up; gate]``) + and ``{prefix}.net.2.weight`` to the ``up_proj``/``gate_proj``/``down_proj`` the TP rewrite + uses. In place; returns ``state_dict``.""" + for prefix in prefixes: + fused_key = f"{prefix}.net.0.proj.weight" + if fused_key not in state_dict: + continue + up_weight, gate_weight = state_dict.pop(fused_key).chunk(2, dim=0) + state_dict[f"{prefix}.up_proj.weight"] = up_weight.contiguous() + state_dict[f"{prefix}.gate_proj.weight"] = gate_weight.contiguous() + state_dict[f"{prefix}.down_proj.weight"] = state_dict.pop(f"{prefix}.net.2.weight") + return state_dict + + +def ff_prefixes(num_refiner_layers: int, num_layers: int, root: str = "transformer.") -> list[str]: + """The feed-forward module prefixes of every refiner block and transformer block.""" + return [ + *[f"{root}token_refiner.refiner_blocks.{i}.ff" for i in range(int(num_refiner_layers))], + *[f"{root}transformer_blocks.{i}.ff" for i in range(int(num_layers))], + ] diff --git a/tests/unit/models/minimax_h3/test_modeling_tp.py b/tests/unit/models/minimax_h3/test_modeling_tp.py new file mode 100644 index 000000000..a6c363e5e --- /dev/null +++ b/tests/unit/models/minimax_h3/test_modeling_tp.py @@ -0,0 +1,28 @@ +"""Helpers of the backend-neutral TP rewrite (``difflet/models/minimax_h3/modeling_minimax_h3_tp.py``). +The Trainium re-export check is in ``test_trainium_transformer.py`` (needs the Neuron SDK).""" + +from __future__ import annotations + +import torch + +from difflet.models.minimax_h3 import modeling_minimax_h3_tp as neutral + + +def test_ff_prefixes_cover_refiner_then_blocks(): + assert neutral.ff_prefixes(1, 2) == [ + "transformer.token_refiner.refiner_blocks.0.ff", + "transformer.transformer_blocks.0.ff", + "transformer.transformer_blocks.1.ff", + ] + assert neutral.ff_prefixes(0, 1, root="") == ["transformer_blocks.0.ff"] + + +def test_split_fused_swiglu_keys_is_in_place_and_skips_missing(): + state = { + "a.ff.net.0.proj.weight": torch.arange(8).reshape(4, 2), + "a.ff.net.2.weight": torch.zeros(2, 2), + } + out = neutral.split_fused_swiglu_keys(state, ["a.ff", "missing.ff"]) + assert out is state + assert sorted(state) == ["a.ff.down_proj.weight", "a.ff.gate_proj.weight", "a.ff.up_proj.weight"] + assert state["a.ff.up_proj.weight"].tolist() == [[0, 1], [2, 3]] diff --git a/tests/unit/models/minimax_h3/test_trainium_transformer.py b/tests/unit/models/minimax_h3/test_trainium_transformer.py index 09d5199a1..297e6d60d 100644 --- a/tests/unit/models/minimax_h3/test_trainium_transformer.py +++ b/tests/unit/models/minimax_h3/test_trainium_transformer.py @@ -11,6 +11,8 @@ _prefix_bounds, ) from difflet.backends.trainium.core.config import NeuronConfig +from difflet.backends.trainium.minimax_h3 import transformer as trainium +from difflet.models.minimax_h3 import modeling_minimax_h3_tp as neutral def _config(*, text_seq_len=8): @@ -105,3 +107,15 @@ def test_h3_checkpoint_conversion_splits_swiglu_in_official_order(): [12, 13, 14, 15], ] assert f"{prefix}.net.0.proj.weight" not in converted + + +def test_trainium_reexports_the_neutral_tp_objects(): + for name in ( + "_H3TensorParallelSwiGLU", + "_MiniMaxH3TransformerTraceModule", + "_prefix_bounds", + "_replace_h3_linears_for_tp", + "split_fused_swiglu_keys", + ): + assert getattr(trainium, name) is getattr(neutral, name) + assert trainium._MiniMaxH3TrainiumAttnProcessor is neutral.MiniMaxH3TPAttnProcessor From f95fe649dc8798aa24d0423e5ea8c5a054bbc42c Mon Sep 17 00:00:00 2001 From: SHIYANG CHENG Date: Tue, 10 Nov 2026 16:40:12 -0800 Subject: [PATCH 2/8] feat(neuron, minimax-h3): the H3 DiT as a TorchNeuron application NeuronMiniMaxH3TransformerApplication builds the shared H3 TP trace module on meta, loads each rank's shards (the fused SwiGLU through CheckpointSlices, the fp32 patch projections / timestep MLP / output heads on real fp32 storage) and compiles the 50 blocks and the refiner blocks per block. precomputed_adaln is on by default, so the 13B adaln_proj branch is never loaded. The masks travel as int32; at tp1 the bounded-attention processor is installed too. L1 tests (CPU, gloo): tp1/tp2/tp4 match diffusers' MiniMaxH3Transformer3DModel (fp32, atol 2e-5) with the host modulation table and with the resident branch, eager and per-block compiled. Co-Authored-By: Claude Opus 5.5 --- .../backends/neuron/minimax_h3/__init__.py | 1 + .../backends/neuron/minimax_h3/transformer.py | 316 ++++++++++++++++++ tests/unit/backends/_neuron_h3_toy.py | 159 +++++++++ .../backends/test_neuron_minimax_h3_l1.py | 116 +++++++ 4 files changed, 592 insertions(+) create mode 100644 difflet/backends/neuron/minimax_h3/__init__.py create mode 100644 difflet/backends/neuron/minimax_h3/transformer.py create mode 100644 tests/unit/backends/_neuron_h3_toy.py create mode 100644 tests/unit/backends/test_neuron_minimax_h3_l1.py diff --git a/difflet/backends/neuron/minimax_h3/__init__.py b/difflet/backends/neuron/minimax_h3/__init__.py new file mode 100644 index 000000000..9347e0f4b --- /dev/null +++ b/difflet/backends/neuron/minimax_h3/__init__.py @@ -0,0 +1 @@ +"""MiniMax-H3 on the PyTorch-native neuron backend (TorchNeuron, one process per NeuronCore).""" diff --git a/difflet/backends/neuron/minimax_h3/transformer.py b/difflet/backends/neuron/minimax_h3/transformer.py new file mode 100644 index 000000000..5de729306 --- /dev/null +++ b/difflet/backends/neuron/minimax_h3/transformer.py @@ -0,0 +1,316 @@ +"""Neuron application for the MiniMax-H3 Omni Transformer (the 33B joint video+audio DiT). + +The Trainium counterpart is ``NeuronMiniMaxH3TransformerApplication`` +(``difflet/backends/trainium/minimax_h3/transformer.py``); both build the backend-neutral +``_MiniMaxH3TransformerTraceModule`` (``difflet/models/minimax_h3/modeling_minimax_h3_tp.py``): +the diffusers-layout model under ``self.transformer`` with every attention head-sharded over the +TP group (56 heads / tp), each SwiGLU split into column-parallel ``up_proj``/``gate_proj`` and a +row-parallel ``down_proj``, and the self-attention through ``difflet.ops.attention`` with the +packed sequence's padded tail lowered to per-query key bounds (the bounded flash kernel on the +device). + +Here the lifecycle is ``TorchNeuronApplicationBase``'s: the module is built on ``meta`` after +the mesh exists, each rank reads only its own shards of the diffusers checkpoint, and in compile +mode the 50 identical ``MiniMaxH3TransformerBlock``s are compiled into one shared graph (the two +token-refiner blocks into another) while the top level (patch projections, timestep MLP, RoPE, +packing, the output norm and heads) stays eager. + +H3-specific loading: + +* Mixed precision. The checkpoint keeps the two input patch projections, the timestep MLP and + the two output heads in fp32 (``MiniMaxH3Transformer3DModel._keep_in_fp32_modules``). The meta + loader would give every parameter the application dtype, so ``build_module`` gives those + modules real fp32 host storage first: the loader fills tensors that already have storage in + their own dtype, and moves them to the device with the rest. They run at the top level, eager. +* Fused SwiGLU. The checkpoint's ``ff.net.0.proj.weight`` holds ``[up; gate]`` rows; the TP + module's ``up_proj``/``gate_proj`` are ``CheckpointSlice``s of it (each rank reads its shard of + its half), ``down_proj`` is a rename of ``ff.net.2.weight``. +* ``precomputed_adaln`` (default on). The 13B ``adaln_proj`` branch depends only on the + (timestep, modality) pair, so the host precomputes every block's modulation for the sampling + schedule (``build_adaln_modulation_table``) and passes one step's table per call; the branch is + never built and its checkpoint tensors are reported under ``skipped``. + +The call is the Trainium graph's: ``(hidden_states, audio_hidden_states, encoder_hidden_states, +timestep, timestep_indices, token_tags, position_ids, video_indices, audio_indices, +text_indices, encoder_attention_mask, attention_mask[, block_modulation])``. The two masks are +int32 here (bool tensors cannot be broadcast for the warm-up); the processor reads ``!= 0``. +""" + +from __future__ import annotations + +import dataclasses +import json +from pathlib import Path +from typing import Any + +import torch + +from difflet.backends.neuron.core.application_base import TorchNeuronApplicationBase +from difflet.backends.neuron.core.checkpoint import CheckpointSlice +from difflet.backends.neuron.core.config_checks import ( + check_heads_divide_tp, + check_layer_count, + check_tp_matches, + require_bf16_on_device, +) +from difflet.models.minimax_h3.contracts import ( + PATCH_SIZE, + VAE_SPATIAL_COMPRESSION, + build_padded_t2va_layout, + video_latent_num_frames, +) + +#: The trace module nests the diffusers model under this attribute. +CHECKPOINT_PREFIX = "transformer." +#: Top-level modules the checkpoint stores in fp32 (``_keep_in_fp32_modules`` minus the rope, +#: which is a computed buffer). +FP32_MODULES = ("proj_in", "audio_proj_in", "time_embedder", "proj_out", "audio_proj_out") +#: Distinct timestep slots per call: 0 audio, 1 video/text (the denoise loop's contract). +NUM_TIMESTEP_SLOTS = 2 +#: The model's fixed frame rate: 40 audio latent rows per second per channel. +_FPS = 24 +_AUDIO_LATENTS_PER_SECOND = 40 + +_MODEL_FIELDS = ( + "num_attention_heads", "attention_head_dim", "hidden_size", "num_layers", + "num_refiner_layers", "ffn_dim", "in_channels", "audio_in_channels", "patch_size", + "text_dim", "freq_dim", "time_embed_hidden_dim", "time_embed_dim", "rope_freq_dim", + "rope_theta", "norm_eps", "qk_norm_eps", "final_norm_eps", +) + + +@dataclasses.dataclass +class NeuronMiniMaxH3Config: + """The transformer's hyper-parameters (``transformer/config.json``) plus the one static + request shape and the TP degree. ``num_layers`` below the checkpoint's builds the first N + blocks (an on-device parity check); ``checkpoint_num_layers`` remembers the full count.""" + + num_attention_heads: int = 56 + attention_head_dim: int = 128 + hidden_size: int = 5376 + num_layers: int = 50 + num_refiner_layers: int = 2 + ffn_dim: int = 14336 + in_channels: int = 24 + audio_in_channels: int = 32 + patch_size: tuple[int, int, int] = PATCH_SIZE + text_dim: int = 5120 + freq_dim: int = 256 + time_embed_hidden_dim: int = 5376 + time_embed_dim: int = 2688 + rope_freq_dim: int = 16 + rope_theta: float = 10000.0 + norm_eps: float = 1e-5 + qk_norm_eps: float = 1e-5 + final_norm_eps: float = 1e-5 + # request shape and parallelism + height: int = 256 + width: int = 448 + num_frames: int = 124 + text_seq_len: int = 1024 + tp_degree: int = 4 + precomputed_adaln: bool = True + checkpoint_num_layers: int | None = None + + @classmethod + def from_pretrained(cls, transformer_dir, **overrides: Any) -> "NeuronMiniMaxH3Config": + with open(Path(transformer_dir) / "config.json") as handle: + raw = json.load(handle) + values = {k: raw[k] for k in _MODEL_FIELDS if k in raw} + if "patch_size" in values: + values["patch_size"] = tuple(int(v) for v in values["patch_size"]) + values["checkpoint_num_layers"] = int(raw.get("num_layers", cls.num_layers)) + values.update(overrides) + return cls(**values) + + # -- derived static shapes (the Trainium config's) -------------------------------------- + + @property + def latent_height(self) -> int: + return int(self.height) // VAE_SPATIAL_COMPRESSION + + @property + def latent_width(self) -> int: + return int(self.width) // VAE_SPATIAL_COMPRESSION + + @property + def video_seq_len(self) -> int: + pt, ph, pw = self.patch_size + return (video_latent_num_frames(int(self.num_frames)) // pt + * (self.latent_height // ph) * (self.latent_width // pw)) + + @property + def video_patch_dim(self) -> int: + pt, ph, pw = self.patch_size + return int(self.in_channels) * pt * ph * pw + + @property + def audio_seq_len(self) -> int: + return int(round(int(self.num_frames) / _FPS * _AUDIO_LATENTS_PER_SECOND)) * 2 + + def layout(self, num_text_tokens: int): + return build_padded_t2va_layout( + num_text_tokens=int(num_text_tokens), max_text_tokens=int(self.text_seq_len), + height=int(self.height), width=int(self.width), num_frames=int(self.num_frames), + sequence_alignment=128, + ) + + def validate(self) -> None: + if tuple(self.patch_size) != tuple(PATCH_SIZE): + raise NotImplementedError(f"MiniMax-H3 supports patch_size={PATCH_SIZE}") + if int(self.height) % 32 or int(self.width) % 32: + raise ValueError("MiniMax-H3 height and width must be divisible by 32") + video_latent_num_frames(int(self.num_frames)) # the 17n+5 temporal contract + check_layer_count("MiniMax-H3", self.num_layers, self.checkpoint_num_layers) + check_heads_divide_tp("MiniMax-H3", self.num_attention_heads, self.tp_degree) + if 2 * 3 * int(self.rope_freq_dim) > int(self.attention_head_dim): + raise ValueError("MiniMax-H3 rotary dim exceeds the attention head dim") + + +def checkpoint_key(name: str) -> str: + """Checkpoint key of module parameter ``name``: the wrapper prefix stripped and the split + SwiGLU's ``down_proj`` mapped back to diffusers' ``net.2`` (``up_proj``/``gate_proj`` are + slices, see ``checkpoint_slices``).""" + key = name.removeprefix(CHECKPOINT_PREFIX) + return key.replace(".ff.down_proj.", ".ff.net.2.") + + +def swiglu_slices(config: NeuronMiniMaxH3Config) -> dict[str, CheckpointSlice]: + """``up_proj`` = rows ``[0, ffn)`` and ``gate_proj`` = rows ``[ffn, 2 ffn)`` of every fused + ``ff.net.0.proj.weight`` (diffusers' SwiGLU: ``hidden, gate = proj(x).chunk(2)``).""" + from difflet.models.minimax_h3.modeling_minimax_h3_tp import ff_prefixes + + ffn = int(config.ffn_dim) + slices = {} + for prefix in ff_prefixes(config.num_refiner_layers, config.num_layers): + source = f"{prefix.removeprefix(CHECKPOINT_PREFIX)}.net.0.proj.weight" + slices[f"{prefix}.up_proj.weight"] = CheckpointSlice(source, 0, 0, ffn) + slices[f"{prefix}.gate_proj.weight"] = CheckpointSlice(source, 0, ffn, ffn) + return slices + + +def _give_fp32_host_storage(module: torch.nn.Module) -> None: + """Replace ``module``'s meta parameters with empty fp32 host tensors, which the loader then + fills in fp32 (it keeps the dtype of tensors that already have storage).""" + for sub in module.modules(): + for name, param in list(sub._parameters.items()): + if param is not None and param.is_meta: + sub._parameters[name] = torch.nn.Parameter( + torch.empty(param.shape, dtype=torch.float32), requires_grad=False + ) + + +def _trace_module_class(): + """``_MiniMaxH3TransformerTraceModule`` with the fp32 modules on real host storage. Built + lazily: importing the TP modeling must wait until the backend's mesh exists.""" + from difflet.models.minimax_h3.modeling_minimax_h3_tp import ( + MiniMaxH3TPAttnProcessor, + _MiniMaxH3TransformerTraceModule, + ) + + class NeuronMiniMaxH3TraceModule(_MiniMaxH3TransformerTraceModule): + def __init__(self, config: NeuronMiniMaxH3Config) -> None: + super().__init__(config) + if int(config.tp_degree) == 1: + # The TP rewrite installs the bounded-attention processor only at tp > 1; use it + # at tp1 too, so every degree attends through difflet.ops.attention with the + # int32 masks this application passes. + processor = MiniMaxH3TPAttnProcessor() + for block in (*self.transformer.token_refiner.refiner_blocks, + *self.transformer.transformer_blocks): + block.attn.processor = processor + for name in FP32_MODULES: + _give_fp32_host_storage(getattr(self.transformer, name)) + + return NeuronMiniMaxH3TraceModule + + +class NeuronMiniMaxH3TransformerApplication(TorchNeuronApplicationBase): + """The H3 DiT, sharded across the TP group, at the config's one static shape. + + ``model_path`` is the transformer directory (``/transformer``). On the neuron device + the dtype must be bf16. Only tensor parallelism: no CP, SP, DP or CFG (H3 is + guidance-distilled). + """ + + block_attrs = ("transformer.transformer_blocks", "transformer.token_refiner.refiner_blocks") + + def __init__(self, *, model_path, config: NeuronMiniMaxH3Config, parallel=None, + dtype=torch.bfloat16, exec_mode: str | None = None, device="neuron", + **kwargs: Any) -> None: + super().__init__(model_path=model_path, parallel=parallel, dtype=dtype, + exec_mode=exec_mode, device=device, **kwargs) + name = "NeuronMiniMaxH3Config" + check_tp_matches(name, config, self.parallel) + for field in ("cp_degree", "dp_degree"): + if int(getattr(self.parallel, field, 1) or 1) > 1: + raise NotImplementedError( + f"the MiniMax-H3 DiT on the neuron backend runs tensor parallelism only; " + f"got {field}={getattr(self.parallel, field)}" + ) + for flag in ("cfg_parallel_enabled", "sp_enabled"): + if bool(getattr(self.parallel, flag, False)): + raise NotImplementedError(f"the MiniMax-H3 DiT on the neuron backend has no {flag}") + config.validate() + require_bf16_on_device("MiniMax-H3", self.device.type, self.dtype) + self.config = config + # Checkpoint tensors left out on purpose: the adaln_proj branch (precomputed on the + # host) and the blocks past num_layers. Reported under load_report["skipped"]. + full = int(config.checkpoint_num_layers or config.num_layers) + skipped = [f"transformer_blocks.{i}." for i in range(int(config.num_layers), full)] + if config.precomputed_adaln: + skipped += [f"transformer_blocks.{i}.adaln_proj." for i in range(int(config.num_layers))] + self.skipped_checkpoint_prefixes = tuple(skipped) + + # ------------------------------------------------------------------ hooks + + def build_module(self) -> torch.nn.Module: + return _trace_module_class()(self.config) + + def checkpoint_key(self, name: str) -> str: + return checkpoint_key(name) + + def checkpoint_slices(self) -> dict[str, CheckpointSlice]: + # At tp1 the TP rewrite is not applied and the SwiGLU keeps diffusers' fused layout. + return swiglu_slices(self.config) if int(self.config.tp_degree) > 1 else {} + + def get_example_inputs(self) -> tuple[torch.Tensor, ...]: + """The call at the static shape with 64 live text tokens (as the Trainium trace).""" + cfg = self.config + gen = torch.Generator().manual_seed(0) + live = min(64, int(cfg.text_seq_len)) + layout = cfg.layout(live) + timestep_indices = torch.ones(layout.sequence_length, dtype=torch.int64) + timestep_indices[layout.audio_indices] = 0 + text_mask = torch.zeros(1, int(cfg.text_seq_len), dtype=torch.int32) + text_mask[:, :live] = 1 + inputs = [ + torch.randn(1, cfg.video_seq_len, cfg.video_patch_dim, generator=gen).to(self.dtype), + torch.randn(1, cfg.audio_seq_len, int(cfg.audio_in_channels), generator=gen + ).to(self.dtype), + torch.randn(1, int(cfg.text_seq_len), int(cfg.text_dim), generator=gen + ).to(self.dtype), + torch.tensor([0.3, 0.7], dtype=torch.float32), + timestep_indices, + layout.token_tags.to(torch.int64), + layout.position_ids.to(torch.float32), + layout.video_indices.to(torch.int64), + layout.audio_indices.to(torch.int64), + layout.text_indices.to(torch.int64), + text_mask, + layout.attention_mask.to(torch.int32).unsqueeze(0), + ] + if cfg.precomputed_adaln: + inputs.append(torch.randn(int(cfg.num_layers), 6, NUM_TIMESTEP_SLOTS * 3, + int(cfg.hidden_size), generator=gen).to(self.dtype)) + return tuple(inputs) + + +__all__ = [ + "CHECKPOINT_PREFIX", + "FP32_MODULES", + "NeuronMiniMaxH3Config", + "NeuronMiniMaxH3TransformerApplication", + "checkpoint_key", + "swiglu_slices", +] diff --git a/tests/unit/backends/_neuron_h3_toy.py b/tests/unit/backends/_neuron_h3_toy.py new file mode 100644 index 000000000..1c3e2995d --- /dev/null +++ b/tests/unit/backends/_neuron_h3_toy.py @@ -0,0 +1,159 @@ +"""A tiny diffusers-layout MiniMax-H3 DiT for the neuron H3 tests, and their gloo workers (not a +test module). + +Every run goes through ``_neuron_gloo.run_ranks`` in a fresh process with +``DIFFLET_BACKEND=neuron``, which binds ``difflet.ops`` to the neuron implementations on CPU/gloo. +""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +TINY = dict( + num_attention_heads=4, + attention_head_dim=16, + hidden_size=32, + num_layers=2, + num_refiner_layers=1, + ffn_dim=64, + in_channels=24, + audio_in_channels=32, + patch_size=(1, 2, 2), + text_dim=16, + freq_dim=8, + time_embed_hidden_dim=32, + time_embed_dim=16, + rope_freq_dim=2, + rope_theta=10000.0, +) +TINY_SHAPE = dict(height=32, width=32, num_frames=124) +TINY_TEXT_SEQ_LEN = 8 +LIVE_TEXT = 5 +TIMESTEP = (0.3, 0.7) + + +def _model(precomputed_adaln: bool = False): + from difflet.models.minimax_h3.modeling_minimax_h3 import MiniMaxH3Transformer3DModel + + return MiniMaxH3Transformer3DModel(**TINY, precomputed_adaln=precomputed_adaln) + + +def write_tiny_h3_transformer(root) -> Path: + """Random weights (fp32 everywhere, scaled so the outputs are not dominated by noise), + saved sharded with an index as the real checkpoint is.""" + torch.manual_seed(0) + model = _model() + with torch.no_grad(): + for param in model.parameters(): + param.normal_(0.0, 0.2) + out = Path(root) / "transformer" + model.save_pretrained(out, max_shard_size="200KB", safe_serialization=True) + return out + + +def tiny_inputs(*, scale: float = 1.0): + """The DiT call at the tiny shape (bool masks, as the denoise loop builds them).""" + from difflet.models.minimax_h3.contracts import build_padded_t2va_layout + + layout = build_padded_t2va_layout( + num_text_tokens=LIVE_TEXT, max_text_tokens=TINY_TEXT_SEQ_LEN, sequence_alignment=128, + **TINY_SHAPE, + ) + gen = torch.Generator().manual_seed(1) + video_rows = int(layout.video_indices.numel()) + audio_rows = int(layout.audio_indices.numel()) + timestep_indices = torch.ones(layout.sequence_length, dtype=torch.int64) + timestep_indices[layout.audio_indices] = 0 + text_mask = torch.zeros(1, TINY_TEXT_SEQ_LEN, dtype=torch.bool) + text_mask[:, :LIVE_TEXT] = True + return ( + scale * torch.randn(1, video_rows, 96, generator=gen), + scale * torch.randn(1, audio_rows, 32, generator=gen), + scale * torch.randn(1, TINY_TEXT_SEQ_LEN, 16, generator=gen), + torch.tensor(TIMESTEP, dtype=torch.float32), + timestep_indices, + layout.token_tags.to(torch.int64), + layout.position_ids.to(torch.float32), + layout.video_indices.to(torch.int64), + layout.audio_indices.to(torch.int64), + layout.text_indices.to(torch.int64), + text_mask, + layout.attention_mask.unsqueeze(0), + ) + + +def diffusers_reference(transformer_dir): + """diffusers' model on the same files, fp32, with the adaln branch (no precompute).""" + from difflet.models.minimax_h3.modeling_minimax_h3 import MiniMaxH3Transformer3DModel + + model = MiniMaxH3Transformer3DModel.from_pretrained(transformer_dir, torch_dtype=torch.float32) + model.eval() + with torch.no_grad(): + video, audio = model(*tiny_inputs(), return_dict=False) + return video, audio + + +def modulation_table(transformer_dir) -> torch.Tensor: + from difflet.models.minimax_h3.pipeline import build_adaln_modulation_table + + table = build_adaln_modulation_table( + str(transformer_dir), torch.tensor([TIMESTEP]), num_layers=TINY["num_layers"], + freq_dim=TINY["freq_dim"], time_embed_hidden_dim=TINY["time_embed_hidden_dim"], + time_embed_dim=TINY["time_embed_dim"], hidden_size=TINY["hidden_size"], + dtype=torch.float32, + ) + return table[0] + + +def as_neuron_call(inputs, table=None): + """The neuron application's call: masks as int32, the step's modulation table appended.""" + *head, text_mask, mask = inputs + call = (*head, text_mask.to(torch.int32), mask.to(torch.int32)) + return call if table is None else (*call, table) + + +# --------------------------------------------------------------------------- gloo workers + + +def l1_h3_worker(rank: int, world_size: int, transformer_dir: str, precomputed: bool, + exec_mode: str) -> dict: + from difflet.backends.neuron.minimax_h3.transformer import ( + NeuronMiniMaxH3Config, + NeuronMiniMaxH3TransformerApplication, + ) + from difflet.pipeline.parallel_config import DiffletParallelConfig + + parallel = DiffletParallelConfig(tp_degree=world_size) + config = NeuronMiniMaxH3Config.from_pretrained( + transformer_dir, **TINY_SHAPE, text_seq_len=TINY_TEXT_SEQ_LEN, tp_degree=world_size, + precomputed_adaln=precomputed, + ) + app = NeuronMiniMaxH3TransformerApplication( + model_path=transformer_dir, config=config, parallel=parallel, dtype=torch.float32, + exec_mode=exec_mode, device="cpu", + ) + torch._dynamo.reset() + app.load() + table = modulation_table(transformer_dir) if precomputed else None + video, audio = app(*as_neuron_call(tiny_inputs(), table)) + other, _ = app(*as_neuron_call(tiny_inputs(scale=0.5), table)) + t = app.module.transformer + block0 = t.transformer_blocks[0] + return { + "load_report": app.load_report, + "fp32": {name: str(getattr(t, name).weight.dtype if hasattr(getattr(t, name), "weight") + else next(getattr(t, name).parameters()).dtype) + for name in ("proj_in", "audio_proj_in", "time_embedder", "proj_out", + "audio_proj_out")}, + "block_dtype": str(block0.attn.to_q.weight.dtype), + "local_heads": int(block0.attn.heads), + "processor": type(block0.attn.processor).__name__, + "has_adaln": block0.adaln_proj is not None, + "compiled_blocks": list(app.compiled_blocks), + "params_on_meta": sum(p.is_meta for p in app.module.parameters()), + "video": video.float().tolist(), + "audio": audio.float().tolist(), + "other_differs": not bool(torch.equal(video, other)), + } diff --git a/tests/unit/backends/test_neuron_minimax_h3_l1.py b/tests/unit/backends/test_neuron_minimax_h3_l1.py new file mode 100644 index 000000000..0fc3641a3 --- /dev/null +++ b/tests/unit/backends/test_neuron_minimax_h3_l1.py @@ -0,0 +1,116 @@ +"""L1 unit tests (CPU only): the MiniMax-H3 DiT as a neuron-backend application. + +The model is the tiny diffusers-layout H3 of tests/unit/backends/_neuron_h3_toy.py (2 blocks, 1 +refiner block, 4 heads x 16, a sharded checkpoint with an index). The reference is diffusers' +``MiniMaxH3Transformer3DModel`` itself, reloaded from the same files with its ``adaln_proj`` +branch. Every forward runs in spawned gloo processes: at TP2/TP4 the head shards, the split +SwiGLU (a ``CheckpointSlice`` of the fused weight per rank) and the all-reduces run, and with +``precomputed_adaln`` the host-built modulation table replaces the branch. Nothing here touches +the device. +""" + +import pytest + +torch = pytest.importorskip("torch") + +from types import SimpleNamespace # noqa: E402 + +from difflet.backends.neuron.minimax_h3.transformer import ( # noqa: E402 + NeuronMiniMaxH3Config, + checkpoint_key, + swiglu_slices, +) +from difflet.backends.neuron.ops_impl import parallel_mesh # noqa: E402 +from tests.unit.backends._neuron_gloo import run_ranks # noqa: E402 +from tests.unit.backends._neuron_h3_toy import ( # noqa: E402 + TINY, + TINY_SHAPE, + TINY_TEXT_SEQ_LEN, + diffusers_reference, + l1_h3_worker, + write_tiny_h3_transformer, +) + +ATOL = RTOL = 2e-5 +OFFLINE = {"HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1"} + + +@pytest.fixture(autouse=True) +def _clean_mesh(monkeypatch): + monkeypatch.delenv("DIFFLET_EXEC_MODE", raising=False) + parallel_mesh.destroy_parallel_mesh() + yield + parallel_mesh.destroy_parallel_mesh() + + +@pytest.fixture(scope="module") +def tiny(tmp_path_factory): + pytest.importorskip("safetensors.torch") + pytest.importorskip("accelerate") + pytest.importorskip("diffusers") + transformer = write_tiny_h3_transformer(tmp_path_factory.mktemp("tiny_h3")) + video, audio = diffusers_reference(transformer) + return SimpleNamespace(transformer=str(transformer), video=video, audio=audio) + + +def _run(tiny, *, tp, precomputed=True, exec_mode="eager"): + return run_ranks(l1_h3_worker, tiny.transformer, precomputed, exec_mode, world_size=tp, + env=OFFLINE) + + +def _assert_close(results, tiny): + for r in results: + assert torch.equal(torch.tensor(r["video"]), torch.tensor(results[0]["video"])) + torch.testing.assert_close(torch.tensor(r["video"]), tiny.video, atol=ATOL, rtol=RTOL) + torch.testing.assert_close(torch.tensor(r["audio"]), tiny.audio, atol=ATOL, rtol=RTOL) + assert r["other_differs"] + + +def test_config_reads_checkpoint_and_derives_static_shapes(tiny): + config = NeuronMiniMaxH3Config.from_pretrained( + tiny.transformer, **TINY_SHAPE, text_seq_len=TINY_TEXT_SEQ_LEN, tp_degree=2) + assert config.hidden_size == TINY["hidden_size"] + assert config.checkpoint_num_layers == 2 + assert (config.video_seq_len, config.video_patch_dim, config.audio_seq_len) == (37, 96, 414) + config.validate() + + +def test_checkpoint_keys_and_swiglu_slices(): + assert checkpoint_key("transformer.transformer_blocks.3.ff.down_proj.weight") == ( + "transformer_blocks.3.ff.net.2.weight") + assert checkpoint_key("transformer.proj_in.weight") == "proj_in.weight" + config = NeuronMiniMaxH3Config(num_layers=1, num_refiner_layers=1, ffn_dim=8) + slices = swiglu_slices(config) + up = slices["transformer.transformer_blocks.0.ff.up_proj.weight"] + gate = slices["transformer.transformer_blocks.0.ff.gate_proj.weight"] + assert (up.key, up.dim, up.start, up.length) == ("transformer_blocks.0.ff.net.0.proj.weight", + 0, 0, 8) + assert (gate.start, gate.length) == (8, 8) + assert len(slices) == 4 # one refiner block + one transformer block, two halves each + + +@pytest.mark.parametrize("tp", [1, 2, 4]) +def test_precomputed_adaln_matches_diffusers(tiny, tp): + results = _run(tiny, tp=tp) + _assert_close(results, tiny) + r = results[0] + assert r["params_on_meta"] == 0 + assert set(r["fp32"].values()) == {"torch.float32"} + assert not r["has_adaln"] + assert r["local_heads"] == TINY["num_attention_heads"] // tp + assert r["load_report"]["missing"] == [] and r["load_report"]["unexpected"] == [] + assert len(r["load_report"]["skipped"]) == 2 * 2 # adaln weight + bias per block + if tp > 1: + assert r["processor"] == "MiniMaxH3TPAttnProcessor" + + +def test_resident_adaln_matches_diffusers(tiny): + results = _run(tiny, tp=2, precomputed=False) + _assert_close(results, tiny) + assert results[0]["has_adaln"] and not results[0]["load_report"].get("skipped") + + +def test_compile_mode_matches_eager(tiny): + results = _run(tiny, tp=2, exec_mode="compile") + _assert_close(results, tiny) + assert results[0]["compiled_blocks"] From 5d3ac79385c2f6aa63c11ea9fcd30cc8a7ca76e6 Mon Sep 17 00:00:00 2001 From: SHIYANG CHENG Date: Tue, 10 Nov 2026 17:06:42 -0800 Subject: [PATCH 3/8] feat(minimax-h3): host CPU text encoder and decoders HostQwen3VLTextEncoder builds only the first 50 Qwen3-VL decoder layers (no layers 50-63, final norm, lm_head or vision tower) and captures hidden_states[50] with the Trainium text stage's tokenization and record. load_host_video_decoder / load_host_audio_decoder build the two VAEs in fp32 with the reference attention and the folded weight norm. On trn3 (fox prompt, vs the AOT stage outputs): text live rows rel_l2 9.9e-4 (bf16), mask and token count equal; video decode PSNR 73.4 dB. Co-Authored-By: Claude Opus 5.5 --- difflet/models/minimax_h3/host_components.py | 284 +++++++++++++++++++ 1 file changed, 284 insertions(+) create mode 100644 difflet/models/minimax_h3/host_components.py diff --git a/difflet/models/minimax_h3/host_components.py b/difflet/models/minimax_h3/host_components.py new file mode 100644 index 000000000..86b5d8585 --- /dev/null +++ b/difflet/models/minimax_h3/host_components.py @@ -0,0 +1,284 @@ +"""MiniMax-H3 components that run on the host CPU, for backends without a device port of them. + +Backend-neutral (torch, transformers, safetensors only), used by the neuron backend's staged CLI +(``difflet/cli/orchestrators/minimax_h3.py``): + +* ``HostQwen3VLTextEncoder``: the conditioner, Qwen3-VL's language model run to + ``hidden_states[50]`` (the output of decoder layer 49), as the Trainium text stage captures it + through NxDI. Only the first 50 decoder layers are built and read (the last 14, the final norm, + ``lm_head`` and the vision tower are never loaded). Tokenization is the Trainium stage's: the + raw prompt, no special tokens, right-padded to the 1024-token bucket; the padding rows of the + output are zeroed. +* ``load_host_video_decoder`` / ``load_host_audio_decoder``: the two decoder-only VAEs in fp32, + with the reference attention (no sequence padding) and the audio decoder's weight norm folded + into plain convolution weights. Their callables plug into ``decode_minimax_h3_video`` / + ``decode_minimax_h3_audio`` (``pipeline.py``) unchanged. +""" + +from __future__ import annotations + +import contextlib +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from torch import nn + +from difflet.common.orchestrators import minimax_h3 as h3_common + +_LM_PREFIX = "model.language_model." + + +@contextlib.contextmanager +def intra_op_threads(threads: int | None): + if not threads: + yield + return + before = torch.get_num_threads() + torch.set_num_threads(int(threads)) + try: + yield + finally: + torch.set_num_threads(before) + + +def _safetensors_index(directory: Path) -> dict[str, str]: + """``{key: file}`` for a sharded (index) or single-file safetensors checkpoint.""" + index = sorted(directory.glob("*.safetensors.index.json")) + if index: + with open(index[0]) as handle: + return {k: str(directory / v) for k, v in json.load(handle)["weight_map"].items()} + from safetensors import safe_open + + mapping = {} + for path in sorted(directory.glob("*.safetensors")): + with safe_open(str(path), framework="pt") as handle: + mapping.update({key: str(path) for key in handle.keys()}) + return mapping + + +def _read_tensors(weight_map: dict[str, str], keys) -> dict[str, torch.Tensor]: + from safetensors import safe_open + + by_file: dict[str, list[str]] = {} + for key in keys: + by_file.setdefault(weight_map[key], []).append(key) + out = {} + for path, file_keys in by_file.items(): + with safe_open(path, framework="pt") as handle: + for key in file_keys: + out[key] = handle.get_tensor(key) + return out + + +# ------------------------------------------------------------------------------ text encoder + + +class HostQwen3VLTextEncoder: + """Qwen3-VL's language model to ``hidden_states[layer]`` on the host CPU. + + ``model_dir`` is the H3 snapshot root (``text_encoder/`` and ``tokenizer/`` under it). + ``encode(prompt)`` returns the Trainium text stage's record: ``encoder_hidden_states`` + ``(1, seq_len, 5120)`` bf16 with the padding rows zeroed, ``attention_mask`` ``(1, seq_len)`` + bool, ``num_text_tokens`` and ``text_encoder_layer``. + """ + + def __init__(self, model_dir, *, layer: int = h3_common.TEXT_ENCODER_LAYER, + seq_len: int = h3_common.TEXT_SEQ_LEN, dtype: torch.dtype = torch.bfloat16, + threads: int | None = None) -> None: + self.model_dir = Path(model_dir) + self.layer = int(layer) + self.seq_len = int(seq_len) + self.dtype = dtype + self.threads = threads + self.model = None + self.tokenizer = None + + def load(self) -> "HostQwen3VLTextEncoder": + from accelerate import init_empty_weights + from transformers import AutoConfig, AutoTokenizer + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextModel + + encoder_dir = self.model_dir / "text_encoder" + text_config = AutoConfig.from_pretrained(encoder_dir).text_config + if int(text_config.num_hidden_layers) < self.layer: + raise ValueError( + f"MiniMax-H3 needs Qwen3-VL hidden_states[{self.layer}], but the text encoder " + f"has {text_config.num_hidden_layers} layers" + ) + if getattr(text_config, "pad_token_id", None) is None: + text_config.pad_token_id = 0 + text_config.num_hidden_layers = self.layer # hidden_states[L] is layer L-1's output + text_config._attn_implementation = "sdpa" + with init_empty_weights(include_buffers=False): + model = Qwen3VLTextModel(text_config) + weight_map = _safetensors_index(encoder_dir) + wanted = {} + for key in weight_map: + if not key.startswith(_LM_PREFIX): + continue + name = key[len(_LM_PREFIX):] + if name.startswith("layers.") and int(name.split(".")[1]) >= self.layer: + continue + wanted[key] = name + tensors = _read_tensors(weight_map, wanted) + state = {wanted[k]: v.to(self.dtype) for k, v in tensors.items()} + missing, unexpected = model.load_state_dict(state, strict=False, assign=True) + if missing or unexpected: + raise RuntimeError(f"Qwen3-VL text load: missing {missing[:5]}, unexpected " + f"{unexpected[:5]}") + self.model = model.eval().requires_grad_(False) + tokenizer = AutoTokenizer.from_pretrained(self.model_dir / "tokenizer") + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + self.tokenizer = tokenizer + return self + + def tokenize(self, prompt: str) -> tuple[torch.Tensor, torch.Tensor]: + tokens = self.tokenizer( + prompt, add_special_tokens=False, max_length=self.seq_len, padding="max_length", + truncation=True, return_tensors="pt", return_attention_mask=True, + ) + return tokens.input_ids, tokens.attention_mask + + @torch.no_grad() + def encode(self, prompt: str) -> dict[str, Any]: + if self.model is None: + raise RuntimeError("HostQwen3VLTextEncoder.load() first") + input_ids, attention_mask = self.tokenize(prompt) + token_count = int(attention_mask.sum()) + if token_count < 1: + raise ValueError("MiniMax-H3 prompt must produce at least one token") + captured: list[torch.Tensor] = [] + handle = self.model.layers[self.layer - 1].register_forward_hook( + lambda module, args, output: captured.append( + output[0] if isinstance(output, tuple) else output) + ) + try: + with intra_op_threads(self.threads): + self.model(input_ids=input_ids, attention_mask=attention_mask, + position_ids=torch.arange(self.seq_len).unsqueeze(0), use_cache=False) + finally: + handle.remove() + hidden = captured[-1].to(torch.bfloat16).clone() + hidden[:, token_count:] = 0 # padding rows are never part of the packed document + return { + "encoder_hidden_states": hidden, + "attention_mask": attention_mask.to(torch.bool), + "num_text_tokens": token_count, + "text_encoder_layer": self.layer, + } + + +# ------------------------------------------------------------------------------ decoders + + +@dataclass +class HostDecoder: + """A decoder callable plus the config values the decode helpers take.""" + + module: nn.Module + config: dict[str, Any] + threads: int | None = None + + def __call__(self, latents: torch.Tensor) -> torch.Tensor: + with torch.no_grad(), intra_op_threads(self.threads): + return self.module(latents.float()) + + +class _HostVideoDecoder(nn.Module): + def __init__(self, cfg: dict[str, Any]) -> None: + from difflet.models.minimax_h3.modeling_vae import MiniMaxH3VideoViTDecoder3d + + super().__init__() + spatial = 1 + for value in cfg["spatial_downsample_factors"]: + spatial *= int(value) + temporal = 1 + for value in cfg["temporal_downsample_factors"]: + temporal *= int(value) + latent = int(cfg["latent_channels"]) + self.post_quant_conv = nn.Conv3d(latent, latent, kernel_size=1) + self.decoder = MiniMaxH3VideoViTDecoder3d( + in_channels=latent, out_channels=int(cfg["out_channels"]), patch_size=spatial, + patch_size_t=temporal, num_layers=int(cfg["decoder_num_layers"]), + num_attention_heads=int(cfg["decoder_num_attention_heads"]), + attention_head_dim=int(cfg["decoder_attention_head_dim"]), + num_register_tokens=int(cfg["decoder_num_register_tokens"]), + ffn_mult=int(cfg["decoder_ffn_mult"]), rope_theta=float(cfg["decoder_rope_theta"]), + rope_dim_ratio=float(cfg["decoder_rope_dim_ratio"]), + norm_eps=float(cfg["decoder_norm_eps"]), + ) + + def forward(self, latents: torch.Tensor) -> torch.Tensor: + return self.decoder(self.post_quant_conv(latents)) + + +class _HostAudioDecoder(nn.Module): + def __init__(self, cfg: dict[str, Any]) -> None: + from difflet.models.minimax_h3.modeling_vae import MiniMaxH3AudioBigVGANDecoder + + super().__init__() + self.dec_in_proj = nn.Conv1d(int(cfg["latent_channels"]), int(cfg["latent_dim"]), 1) + self.decoder = MiniMaxH3AudioBigVGANDecoder( + in_channels=int(cfg["latent_dim"]), + upsample_initial_channel=int(cfg["decoder_dim"]), + upsample_rates=tuple(int(v) for v in cfg["decoder_rates"]), + upsample_kernel_sizes=tuple(int(v) for v in cfg["decoder_kernel_sizes"]), + resblock_kernel_sizes=tuple(int(v) for v in cfg["resblock_kernel_sizes"]), + resblock_dilation_sizes=tuple(tuple(int(v) for v in d) + for d in cfg["resblock_dilation_sizes"]), + use_weight_norm=False, + ) + + def forward(self, latents: torch.Tensor) -> torch.Tensor: + return self.decoder(self.dec_in_proj(latents)) + + +def _load_decoder(directory: Path, module: nn.Module, prefixes: tuple[str, ...], + fold_weight_norm: bool) -> None: + weight_map = _safetensors_index(directory) + state = _read_tensors(weight_map, [k for k in weight_map if k.startswith(prefixes)]) + if fold_weight_norm: + # Legacy weight_norm (weight_g, weight_v) -> the identical plain weight. + for key in [k for k in state if k.endswith(".weight_v")]: + prefix = key[: -len(".weight_v")] + state[f"{prefix}.weight"] = torch._weight_norm( + state.pop(key), state.pop(f"{prefix}.weight_g"), 0) + module.load_state_dict({k: v.float() for k, v in state.items()}, strict=True) + + +def load_host_video_decoder(model_dir, *, threads: int | None = None) -> HostDecoder: + directory = Path(model_dir) / "vae" + with open(directory / "config.json") as handle: + cfg = json.load(handle) + module = _HostVideoDecoder(cfg) + _load_decoder(directory, module, ("post_quant_conv.", "decoder."), fold_weight_norm=False) + return HostDecoder(module.eval().requires_grad_(False), cfg, threads) + + +def load_host_audio_decoder(model_dir, *, threads: int | None = None) -> HostDecoder: + directory = Path(model_dir) / "audio_vae" + with open(directory / "config.json") as handle: + cfg = json.load(handle) + module = _HostAudioDecoder(cfg) + _load_decoder(directory, module, ("dec_in_proj.", "decoder."), fold_weight_norm=True) + return HostDecoder(module.eval().requires_grad_(False), cfg, threads) + + +def default_threads() -> int: + """Intra-op threads for the host components: half the CPUs, at most 96 (the host is + shared with the device ranks and, on the trn hosts, with other jobs).""" + return max(1, min(96, (os.cpu_count() or 2) // 2)) + + +__all__ = [ + "HostDecoder", + "HostQwen3VLTextEncoder", + "default_threads", + "load_host_audio_decoder", + "load_host_video_decoder", +] From ff1f46c15a4b134f2fdc0b028c855c5f72ba7803 Mon Sep 17 00:00:00 2001 From: SHIYANG CHENG Date: Tue, 10 Nov 2026 17:06:42 -0800 Subject: [PATCH 4/8] feat(neuron, minimax-h3): staged CLI on the torchrun neuron backend With DIFFLET_BACKEND=neuron the four H3 stages keep their work files: text and both decoders run on the host as one rank, the DiT on --tp-degree torchrun ranks with the AdaLN table built once on rank 0 under /minimax_h3_neuron/. compile only warms the DiT up. The registry lists neuron for minimax_h3; the stage CLI forwards --model-path to every stage on both backends. The c10 refusal tests now restrict H3's entry to trainium themselves. trn3 device 7, 256x448x124, 30 steps, seed 42: difflet generate end to end 242 s (text 23 s, DiT load+warmup 66 s, denoise 21.1 s = 716 ms/step, video VAE 69 s, audio 2.6 s); same scene as the AOT output (video PSNR 22.5 dB, audio mean -58.5 vs -58.0 dB). Co-Authored-By: Claude Opus 5.5 --- difflet/cli/orchestrators/minimax_h3.py | 319 +++++++++++++++++- difflet/registry.py | 4 +- tests/unit/backends/test_neuron_launch_c10.py | 16 +- .../unit/cli/test_orchestrator_minimax_h3.py | 60 ++++ 4 files changed, 390 insertions(+), 9 deletions(-) diff --git a/difflet/cli/orchestrators/minimax_h3.py b/difflet/cli/orchestrators/minimax_h3.py index 390b0a50c..ccf8daf1b 100644 --- a/difflet/cli/orchestrators/minimax_h3.py +++ b/difflet/cli/orchestrators/minimax_h3.py @@ -8,6 +8,18 @@ Only light tokenization, scheduling, tensor layout, normalization and media post-processing remain on the host. Staging is required by HBM capacity, not because the heavy components are intended to execute on the host. + +With DIFFLET_BACKEND naming the non-AoT torchrun backend (neuron, TorchNeuron), the same four +stages and work files run as follows (phase 1 of the port): + +* text: Qwen3-VL to hidden_states[50] on the host CPU (``HostQwen3VLTextEncoder``), one rank; +* generate: the DiT on ``--tp-degree`` torchrun ranks, one NeuronCore each + (``NeuronMiniMaxH3TransformerApplication``, per-block compiled, AdaLN precomputed on the + host); every rank runs the same host denoise loop, rank 0 writes the latents; +* video_vae, audio_vae: the two decoders on the host CPU in fp32, one rank. + +``compile`` then only warms the DiT up (it fills the NEFF cache); the host stages have nothing +to compile. """ from __future__ import annotations @@ -23,6 +35,7 @@ from difflet.cli import runner from difflet.cli.orchestrators.base import ModelOrchestrator +from difflet.cli.orchestrators.torchrun import root_compile_cache_at_cache_dir, torchrun_backend from difflet.common.orchestrators import minimax_h3 as h3_common if TYPE_CHECKING: @@ -95,6 +108,33 @@ def _save_av( return True +class _NeuronH3DiT: + """The denoise loop's DiT callable on the neuron backend: the Trainium graph's positional + call on host tensors in, the host ``(video, audio)`` velocity pair out. The two bool masks + go to the device as int32 (the application's call); every rank must call it.""" + + def __init__(self, app) -> None: + self.app = app + self.seconds: list[float] = [] + + def __call__(self, *inputs): + import time + + import torch + + device = self.app.device + moved = [] + for tensor in inputs: + if tensor.dtype == torch.bool: + tensor = tensor.to(torch.int32) + moved.append(tensor.contiguous().to(device)) + start = time.perf_counter() + video, audio = self.app(*moved) + video, audio = video.cpu(), audio.cpu() + self.seconds.append(time.perf_counter() - start) + return video, audio + + class MiniMaxH3Orchestrator(ModelOrchestrator): def download(self) -> None: from difflet.pipeline.path_resolver import resolve_model_path @@ -114,7 +154,11 @@ def compile(self) -> None: work_dir = Path.home() / ".cache" / "difflet" / "work" / _CLI_NAME work_dir.mkdir(parents=True, exist_ok=True) shared = self._shared_cli_args(stage_mode="compile", work_dir=str(work_dir)) - for stage in _STAGES: + stages = _STAGES + if torchrun_backend() is not None: + root_compile_cache_at_cache_dir(self.args.cache_dir, tag="minimax-h3") + stages = ("generate",) # the host stages have nothing to compile + for stage in stages: runner.run_stage( _HF_MODEL_ID, stage, @@ -130,6 +174,8 @@ def generate(self) -> None: ) work_dir.mkdir(parents=True, exist_ok=True) shared = self._shared_cli_args(stage_mode="generate", work_dir=str(work_dir)) + if torchrun_backend() is not None: + root_compile_cache_at_cache_dir(self.args.cache_dir, tag="minimax-h3") try: for stage in _STAGES: runner.run_stage( @@ -146,12 +192,20 @@ def generate(self) -> None: shutil.rmtree(work_dir, ignore_errors=True) def _run_stage_internal(self, stage: str, args: argparse.Namespace) -> None: - handlers = { - "text": self._stage_text, - "generate": self._stage_generate, - "video_vae": self._stage_video_vae, - "audio_vae": self._stage_audio_vae, - } + if torchrun_backend() is not None: + handlers = { + "text": self._stage_text_host, + "generate": self._stage_generate_neuron, + "video_vae": self._stage_video_vae_host, + "audio_vae": self._stage_audio_vae_host, + } + else: + handlers = { + "text": self._stage_text, + "generate": self._stage_generate, + "video_vae": self._stage_video_vae, + "audio_vae": self._stage_audio_vae, + } try: handler = handlers[stage] except KeyError as exc: @@ -488,6 +542,253 @@ def _stage_audio_vae(self, args: argparse.Namespace) -> None: ) print(f"[audio_vae] AV tensors saved to {tensor_path}") + # ------------------------------------------------- torchrun (neuron) backend stages + + def _stage_text_host(self, args: argparse.Namespace) -> None: + """Qwen3-VL hidden_states[50] on the host CPU; writes the Trainium stage's text.pt.""" + import torch + + from difflet.cli.dp import stage_loop + from difflet.models.minimax_h3.host_components import ( + HostQwen3VLTextEncoder, + default_threads, + ) + from difflet.pipeline.path_resolver import resolve_model_path + + if args.stage_mode == "compile": + return + model_dir = self._model_dir(args, resolve_model_path) + encoder = HostQwen3VLTextEncoder(model_dir, threads=default_threads()).load() + for req in stage_loop.claim_requests(args): + with stage_loop.request_scope(args, req, final=False): + record = encoder.encode(req.prompt) + destination = stage_loop.work_file(args, req, "text.pt") + torch.save(record, destination) + print( + f"[text] Qwen3-VL hidden_states[{record['text_encoder_layer']}] (host) " + f"{tuple(record['encoder_hidden_states'].shape)} " + f"({record['num_text_tokens']} live tokens) -> {destination}", + flush=True, + ) + + def _neuron_adaln_table(self, args, model_dir: str, config, num_steps: int): + """The precomputed AdaLN table for ``num_steps``, built once on rank 0 into + ``/minimax_h3_neuron/`` and read by every rank.""" + import torch + + from difflet.backends.neuron.core.distributed import rank0_call + from difflet.models.minimax_h3.pipeline import ( + build_adaln_modulation_table, + load_minimax_h3_schedulers, + ) + + cache_root = Path(args.cache_dir or Path.home() / ".cache" / "difflet") + path = cache_root / "minimax_h3_neuron" / f"adaln_table_steps{num_steps}.pt" + + def build() -> None: + if path.exists(): + return + video_sched, audio_sched = load_minimax_h3_schedulers(model_dir) + video_sched.set_timesteps(num_steps, device="cpu") + audio_sched.set_timesteps(num_steps, device="cpu") + # Slot order matches the denoise loop: 0 = audio, 1 = video. + timesteps = torch.stack( + [audio_sched.timesteps.to(torch.float32), video_sched.timesteps.to(torch.float32)], + dim=-1, + ) + table = build_adaln_modulation_table( + str(Path(model_dir) / "transformer"), + timesteps, + num_layers=int(config.num_layers), + freq_dim=int(config.freq_dim), + time_embed_hidden_dim=int(config.time_embed_hidden_dim), + time_embed_dim=int(config.time_embed_dim), + hidden_size=int(config.hidden_size), + dtype=torch.bfloat16, + ) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(".tmp") + torch.save(table, tmp) + tmp.replace(path) + print(f"[generate] precomputed AdaLN table for {num_steps} steps " + f"({table.numel() * table.element_size() / 2**20:.0f} MiB) -> {path}", flush=True) + + rank0_call(build, device=self._neuron_device, what="AdaLN table") + return torch.load(path) + + def _stage_generate_neuron(self, args: argparse.Namespace) -> None: + """The DiT on every torchrun rank; each runs the host denoise loop on identical inputs + (one CPU generator per request), rank 0 writes the latents.""" + import time + + import torch + + from difflet.backends import get_backend + from difflet.cli.dp import stage_loop + from difflet.cli.orchestrators.torchrun import replica_request_scope, torchrun_parallel + from difflet.pipeline.path_resolver import resolve_model_path + + parallel = torchrun_parallel(args) + get_backend("neuron").prepare_runtime(parallel) + + from difflet.backends.neuron.core.distributed import world_info + from difflet.backends.neuron.minimax_h3.transformer import ( + NeuronMiniMaxH3Config, + NeuronMiniMaxH3TransformerApplication, + ) + from difflet.models.minimax_h3.pipeline import ( + denoise_minimax_h3_t2va, + load_minimax_h3_schedulers, + ) + + model_dir = self._model_dir(args, resolve_model_path) + transformer_path = str(Path(model_dir) / "transformer") + height = args.height or h3_common.DEFAULT_HEIGHT + width = args.width or h3_common.DEFAULT_WIDTH + num_frames = args.num_frames or h3_common.DEFAULT_NUM_FRAMES + config = NeuronMiniMaxH3Config.from_pretrained( + transformer_path, height=height, width=width, num_frames=num_frames, + text_seq_len=h3_common.TEXT_SEQ_LEN, tp_degree=parallel.tp_degree, + precomputed_adaln=True, + ) + app = NeuronMiniMaxH3TransformerApplication( + model_path=transformer_path, config=config, parallel=parallel, dtype=torch.bfloat16, + ) + self._neuron_device = app.device + rank = world_info()[0] + start = time.perf_counter() + app.load() + if rank == 0: + print(f"[generate] DiT loaded on {app.world_size} ranks in " + f"{time.perf_counter() - start:.1f}s " + f"({ {k: round(v, 1) for k, v in app.phase_seconds.items()} }), " + f"{app.weight_bytes / 2**30:.2f} GiB/rank", flush=True) + if args.stage_mode == "compile": + return + + dit = _NeuronH3DiT(app) + for req in stage_loop.claimed_requests(args): + with replica_request_scope(args, req, final=False): + text = torch.load(stage_loop.work_file(args, req, "text.pt")) + steps = int(stage_loop.effective(req, args, "steps", h3_common.DEFAULT_STEPS)) + table = self._neuron_adaln_table(args, model_dir, config, steps) + video_scheduler, audio_scheduler = load_minimax_h3_schedulers(model_dir) + start = time.perf_counter() + output = denoise_minimax_h3_t2va( + dit, + encoder_hidden_states=text["encoder_hidden_states"], + encoder_attention_mask=text["attention_mask"], + num_text_tokens=int(text["num_text_tokens"]), + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=steps, + seed=req.seed, + video_scheduler=video_scheduler, + audio_scheduler=audio_scheduler, + model_dtype=torch.bfloat16, + adaln_modulation_table=table, + ) + if rank != 0: + continue + elapsed = time.perf_counter() - start + destination = stage_loop.work_file(args, req, "latents.pt") + torch.save( + { + "video_latents": output.video_latents.cpu(), + "audio_latents": output.audio_latents.cpu(), + "height": height, + "width": width, + "num_frames": num_frames, + }, + destination, + ) + print( + f"[generate] {steps} steps in {elapsed:.1f}s " + f"({1000 * sum(dit.seconds) / max(len(dit.seconds), 1):.0f} ms/DiT call); " + f"video latents {tuple(output.video_latents.shape)}, audio latents " + f"{tuple(output.audio_latents.shape)} -> {destination}", + flush=True, + ) + + def _stage_video_vae_host(self, args: argparse.Namespace) -> None: + import time + + import torch + + from difflet.cli.dp import stage_loop + from difflet.models.minimax_h3.host_components import ( + default_threads, + load_host_video_decoder, + ) + from difflet.models.minimax_h3.pipeline import decode_minimax_h3_video + from difflet.pipeline.path_resolver import resolve_model_path + + if args.stage_mode == "compile": + return + model_dir = self._model_dir(args, resolve_model_path) + decoder = load_host_video_decoder(model_dir, threads=default_threads()) + cfg = decoder.config + for req in stage_loop.claimed_requests(args): + with stage_loop.request_scope(args, req, final=False): + latents = torch.load(stage_loop.work_file(args, req, "latents.pt")) + start = time.perf_counter() + video = decode_minimax_h3_video( + decoder, + latents["video_latents"], + latents_mean=cfg["latents_mean"], + latents_std=cfg["latents_std"], + clip_length=int(cfg["clip_length"]), + token_drop=int(cfg["token_drop"]), + ) + destination = stage_loop.work_file(args, req, "video.pt") + torch.save(video.to(torch.float16), destination) + print(f"[video_vae] (host) RGB video {tuple(video.shape)} in " + f"{time.perf_counter() - start:.1f}s -> {destination}", flush=True) + + def _stage_audio_vae_host(self, args: argparse.Namespace) -> None: + import time + + import torch + + from difflet.cli.dp import stage_loop + from difflet.models.minimax_h3.host_components import ( + default_threads, + load_host_audio_decoder, + ) + from difflet.models.minimax_h3.pipeline import decode_minimax_h3_audio + from difflet.pipeline.path_resolver import resolve_model_path + + if args.stage_mode == "compile": + return + model_dir = self._model_dir(args, resolve_model_path) + decoder = load_host_audio_decoder(model_dir, threads=default_threads()) + cfg = decoder.config + for req in stage_loop.claimed_requests(args): + with stage_loop.request_scope(args, req, final=True): + latents = torch.load(stage_loop.work_file(args, req, "latents.pt")) + video = torch.load(stage_loop.work_file(args, req, "video.pt")).float() + start = time.perf_counter() + # The whole latent in one decode (the reference); the Trainium graph decodes + # fixed 48-frame windows only because its shape is static. + audio = decode_minimax_h3_audio( + decoder, + latents["audio_latents"], + latents_mean=cfg["latents_mean"], + latents_std=cfg["latents_std"], + hop_length=800, + ) + print(f"[audio_vae] (host) waveform {tuple(audio.shape)} in " + f"{time.perf_counter() - start:.1f}s", flush=True) + output_path = Path(req.output) + if not _save_av(video, audio, str(output_path), + sample_rate=int(cfg["sampling_rate"])): + tensor_path = output_path.with_suffix(".pt") + tensor_path.parent.mkdir(parents=True, exist_ok=True) + torch.save({"video": video.to(torch.float16), "audio": audio, "fps": 24, + "sample_rate": int(cfg["sampling_rate"])}, tensor_path) + print(f"[audio_vae] AV tensors saved to {tensor_path}", flush=True) + def _validate_topology(self) -> None: if (self.args.tp_degree or 4) != 4: raise NotImplementedError("MiniMax-H3's first qualified graph is fixed at TP4.") @@ -506,6 +807,9 @@ def _validate_topology(self) -> None: ) def _stage_cores(self, stage: str) -> int: + if torchrun_backend() is not None: + # Only the DiT is on the device; the host stages run as one rank. + return (self.args.tp_degree or 4) if stage == "generate" else 1 return (self.args.tp_degree or 4) if stage in {"text", "generate"} else 1 @staticmethod @@ -556,6 +860,7 @@ def _shared_cli_args(self, stage_mode: str, work_dir: str | None = None) -> list ("--output", getattr(args, "output", None)), ("--cache-dir", getattr(args, "cache_dir", None)), ("--revision", getattr(args, "revision", None)), + ("--model-path", getattr(args, "model_path", None)), ("--work-dir", work_dir), ): if value: diff --git a/difflet/registry.py b/difflet/registry.py index 70487a7e5..2a9f5e61c 100644 --- a/difflet/registry.py +++ b/difflet/registry.py @@ -438,7 +438,9 @@ def is_minimax_h3(model_id: str) -> bool: detector=is_minimax_h3, default_parallel=DiffletParallelConfig(tp_degree=4), default_shape={"height": 768, "width": 1344, "num_frames": 124}, - backends=("trainium",), + # neuron: the staged CLI only (text and decoders on the host, the DiT on torchrun + # ranks, difflet/cli/orchestrators/minimax_h3.py); no resident application yet. + backends=("trainium", "neuron"), # The repository also contains the Ref2VA transformer and the original # FL2VA/Ref2VA checkpoints. The first Difflet path is the root-level # diffusers T2VA partition, so avoid downloading those duplicate weights. diff --git a/tests/unit/backends/test_neuron_launch_c10.py b/tests/unit/backends/test_neuron_launch_c10.py index 2daae7e7f..57811c9a9 100644 --- a/tests/unit/backends/test_neuron_launch_c10.py +++ b/tests/unit/backends/test_neuron_launch_c10.py @@ -29,7 +29,7 @@ FP32_TOL = {"atol": 1e-4, "rtol": 1e-4} WAN_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers" # registry backends: trainium, tpu, neuron (W5) HUNYUAN_ID = "hunyuanvideo-community/HunyuanVideo" # registry backends: trainium, neuron -MINIMAX_ID = "MiniMaxAI/MiniMax-H3" # registry backends: trainium +MINIMAX_ID = "MiniMaxAI/MiniMax-H3" # made trainium-only by _trainium_only_minimax below def _isolate_env(monkeypatch, *names): @@ -44,6 +44,18 @@ def _pythonpath() -> str: return os.pathsep.join(filter(None, [str(REPO_ROOT), os.environ.get("PYTHONPATH")])) +def _trainium_only_minimax(monkeypatch) -> None: + """Every registered model now lists the neuron backend; the refusal tests restrict + MiniMax-H3's entry to trainium to keep a model that does not.""" + import dataclasses + + from difflet import registry + + entry = registry._REGISTRY["minimax_h3"] + monkeypatch.setitem(registry._REGISTRY, "minimax_h3", + dataclasses.replace(entry, backends=("trainium",))) + + def _cli_main(): # difflet/cli/__init__.py re-exports the function ``main``, so # ``import difflet.cli.main as m`` would bind the function, not the module. @@ -674,6 +686,7 @@ def test_main_rejects_backend_the_model_does_not_support(monkeypatch, capsys, tm cli_main = _cli_main() _isolate_env(monkeypatch, "DIFFLET_BACKEND", "DIFFLET_EXEC_MODE") + _trainium_only_minimax(monkeypatch) built = [] monkeypatch.setattr(cli_main, "_get_orchestrator", lambda args: built.append(args)) argv = [command, "--model-id", MINIMAX_ID, "--backend", "neuron"] @@ -728,6 +741,7 @@ def test_main_rejects_ambient_non_aot_backend_the_model_does_not_support( cli_main = _cli_main() _isolate_env(monkeypatch, "DIFFLET_BACKEND", "DIFFLET_EXEC_MODE") monkeypatch.setenv("DIFFLET_BACKEND", "neuron") + _trainium_only_minimax(monkeypatch) built = _recording_orchestrator(monkeypatch, cli_main) with pytest.raises(SystemExit) as exc: diff --git a/tests/unit/cli/test_orchestrator_minimax_h3.py b/tests/unit/cli/test_orchestrator_minimax_h3.py index 3ba8a3509..cedc64235 100644 --- a/tests/unit/cli/test_orchestrator_minimax_h3.py +++ b/tests/unit/cli/test_orchestrator_minimax_h3.py @@ -369,3 +369,63 @@ def __call__(self, prompt, **kwargs): assert payload["text_encoder_layer"] == 50 assert payload["encoder_hidden_states"].shape == (1, 1024, 5120) assert torch.count_nonzero(payload["encoder_hidden_states"][:, 3:]) == 0 + + +# ------------------------------------------------------------- torchrun (neuron) backend + + +def test_neuron_backend_compiles_only_the_dit_and_places_host_stages_on_one_rank( + monkeypatch, tmp_path +): + monkeypatch.setattr(h3_mod, "torchrun_backend", lambda: "neuron") + monkeypatch.setattr(h3_mod, "root_compile_cache_at_cache_dir", lambda *a, **k: None) + calls = [] + monkeypatch.setattr( + h3_mod.runner, "run_stage", lambda model, stage, **kwargs: calls.append((stage, kwargs)) + ) + orchestrator = MiniMaxH3Orchestrator(_args(work_dir=str(tmp_path), keep_work_dir=True)) + + orchestrator.compile() + assert [stage for stage, _ in calls] == ["generate"] + + calls.clear() + orchestrator.generate() + assert [(stage, kw["num_cores"]) for stage, kw in calls] == [ + ("text", 1), ("generate", 4), ("video_vae", 1), ("audio_vae", 1) + ] + + +def test_neuron_backend_dispatches_to_host_and_torchrun_stage_bodies(monkeypatch): + monkeypatch.setattr(h3_mod, "torchrun_backend", lambda: "neuron") + orchestrator = MiniMaxH3Orchestrator(_args()) + seen = [] + for name in ("_stage_text_host", "_stage_generate_neuron", "_stage_video_vae_host", + "_stage_audio_vae_host"): + monkeypatch.setattr(orchestrator, name, lambda args, name=name: seen.append(name)) + for stage in ("text", "generate", "video_vae", "audio_vae"): + orchestrator._run_stage_internal(stage, orchestrator.args) + assert seen == ["_stage_text_host", "_stage_generate_neuron", "_stage_video_vae_host", + "_stage_audio_vae_host"] + + +def test_shared_cli_args_forward_model_path(): + args = MiniMaxH3Orchestrator(_args(model_path="/models/h3"))._shared_cli_args("generate") + assert args[args.index("--model-path") + 1] == "/models/h3" + assert "--model-path" not in MiniMaxH3Orchestrator(_args())._shared_cli_args("generate") + + +def test_neuron_dit_adapter_sends_masks_as_int32_and_returns_host_tensors(): + import torch + + class _App: + device = torch.device("cpu") + + def __call__(self, *inputs): + self.inputs = inputs + return torch.ones(1, 2), torch.zeros(1, 3) + + app = _App() + dit = h3_mod._NeuronH3DiT(app) + video, audio = dit(torch.zeros(2), torch.tensor([True, False])) + assert [t.dtype for t in app.inputs] == [torch.float32, torch.int32] + assert video.shape == (1, 2) and audio.shape == (1, 3) and len(dit.seconds) == 1 From 1acfb51cb548e952734773796975b893c6eae6dc Mon Sep 17 00:00:00 2001 From: SHIYANG CHENG Date: Tue, 10 Nov 2026 17:52:50 -0800 Subject: [PATCH 5/8] feat(neuron, minimax-h3): the video VAE decoder on a NeuronCore NeuronMiniMaxH3VideoVaeApplication runs the 36-block ViT decoder on one core, one official 256 px tile per call (sequence padded to 128 with key bounds, blocks compiled per block, block norms in fp32). The neuron video_vae stage uses it by default (--vae-dtype fp32, or bf16); --host-vae keeps the host decode. compile now warms up the DiT and the device decoder. trn3 device 7, 256x448x124 (14 tiles): fp32 775 ms/tile, 11-14 s per clip, PSNR vs the host fp32 decode 80 dB (106.9 dB end to end); bf16 130 ms/tile, 2.3 s, 57 dB. difflet generate end to end 242 s -> 144 s. Co-Authored-By: Claude Opus 5.5 --- difflet/backends/neuron/minimax_h3/vae.py | 141 ++++++++++++++++++ difflet/cli/orchestrators/minimax_h3.py | 82 +++++++++- difflet/models/minimax_h3/host_components.py | 7 +- tests/unit/backends/_neuron_h3_toy.py | 21 +++ .../backends/test_neuron_minimax_h3_vae.py | 59 ++++++++ .../unit/cli/test_orchestrator_minimax_h3.py | 25 +++- 6 files changed, 325 insertions(+), 10 deletions(-) create mode 100644 difflet/backends/neuron/minimax_h3/vae.py create mode 100644 tests/unit/backends/test_neuron_minimax_h3_vae.py diff --git a/difflet/backends/neuron/minimax_h3/vae.py b/difflet/backends/neuron/minimax_h3/vae.py new file mode 100644 index 000000000..8cf44ab53 --- /dev/null +++ b/difflet/backends/neuron/minimax_h3/vae.py @@ -0,0 +1,141 @@ +"""Neuron application for MiniMax-H3's video VAE decoder (a 36-block ViT, about 2.4B params). + +The module is ``_HostVideoDecoder`` (``difflet/models/minimax_h3/host_components.py``): +``post_quant_conv`` then ``MiniMaxH3VideoViTDecoder3d``, with the checkpoint's own keys; the +encoder half is skipped. It decodes one fixed tile per call, the official 256 px tile of 7 +latent frames (``(1, 24, 7, 16, 16)``, 1792 patch tokens + 4 registers + 1 cls), which +``decode_minimax_h3_video`` (``pipeline.py``) stitches into the clip; every request shape uses +that one tile shape, so the warm-up compiles everything a request needs. + +Neuron-specific, as on Trainium: the token sequence is padded to a multiple of 128 with a key +mask (``sequence_alignment=128``), and ``_NeuronH3VaeAttnProcessor`` attends through +``difflet.ops.attention`` with the mask lowered to per-query key bounds (the bounded flash +kernel), after the q/k RMS norms in fp32 as the reference processor computes them. The 36 +identical blocks compile into one shared graph; the patch projection, registers, RoPE, output +norm and unpatchify stay eager. + +Precision: ``dtype`` is the weights' and the blocks' activation dtype (bf16 by default, fp32 +accepted); the block norms run in fp32 as in the reference, and the caller gets fp32 frames. +""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Any + +import torch + +from difflet.backends.neuron.core.application_base import TorchNeuronApplicationBase + +#: The official decode tile: 7 latent frames x 256 px (16 latents at /16). +TILE_LATENT_SHAPE = (1, 24, 7, 16, 16) +SEQUENCE_ALIGNMENT = 128 + + +class _NeuronH3VaeAttnProcessor: + """The reference ``MiniMaxH3VideoAttnProcessor`` with the attention through + ``difflet.ops.attention`` and the padding mask as key bounds.""" + + def __init__(self) -> None: + from difflet.models.minimax_h3.modeling_vae import apply_h3_video_rotary + from difflet.ops import attention + + self._attention = attention + self._rotary = apply_h3_video_rotary + + def __call__(self, attn, hidden_states, rotary_emb=None, attention_mask=None): + dtype = hidden_states.dtype + query = attn.to_q(hidden_states).unflatten(2, (attn.heads, -1)) + key = attn.to_k(hidden_states).unflatten(2, (attn.heads, -1)) + value = attn.to_v(hidden_states).unflatten(2, (attn.heads, -1)) + query = attn.norm_q(query.float()).to(dtype) + key = attn.norm_k(key.float()).to(dtype) + query = self._rotary(query, rotary_emb) + key = self._rotary(key, rotary_emb) + q, k, v = (t.transpose(1, 2) for t in (query, key, value)) + batch, heads, length, head_dim = q.shape + kwargs = {} + if attention_mask is not None: + count = attention_mask.sum(dim=-1, keepdim=True).to(torch.int32) + bound_max = (count[:, None, None, :].expand(batch, heads, length, 1) + .reshape(batch * heads, length, 1).contiguous()) + kwargs = {"bound_min": torch.zeros_like(bound_max), "bound_max": bound_max} + out = self._attention( + q.reshape(batch * heads, length, head_dim), + k.reshape(batch * heads, length, head_dim), + v.reshape(batch * heads, length, head_dim), + scale=1.0 / math.sqrt(head_dim), causal=False, tp_q=True, tp_k=True, tp_out=False, + **kwargs, + ) + out = out.reshape(batch, heads, length, head_dim).transpose(1, 2).flatten(2, 3).to(dtype) + return attn.to_out[0](out.contiguous()) + + +class _NeuronH3VideoDecoder(torch.nn.Module): + def __init__(self, cfg: dict[str, Any]) -> None: + super().__init__() + from difflet.models.minimax_h3.host_components import _HostVideoDecoder + + inner = _HostVideoDecoder(cfg, sequence_alignment=SEQUENCE_ALIGNMENT) + self.post_quant_conv = inner.post_quant_conv + self.decoder = inner.decoder + from difflet.backends.neuron.minimax_h3.transformer import _give_fp32_host_storage + + processor = _NeuronH3VaeAttnProcessor() + for block in self.decoder.transformer_blocks: + block.attn.set_processor(processor) + # The blocks call these norms on fp32 activations (the reference's precision), so + # their weights stay fp32 whatever the application dtype. + _give_fp32_host_storage(block.norm1) + _give_fp32_host_storage(block.norm2) + + def forward(self, latents: torch.Tensor) -> torch.Tensor: + latents = latents.to(self.post_quant_conv.weight.dtype) + return self.decoder(self.post_quant_conv(latents)).float() + + +class NeuronMiniMaxH3VideoVaeApplication(TorchNeuronApplicationBase): + """The H3 video decoder on one rank, one tile per call. ``model_path`` is ``/vae``.""" + + block_attrs = ("decoder.transformer_blocks",) + skipped_checkpoint_prefixes = ("encoder.", "quant_conv.") + + def __init__(self, *, model_path, parallel=None, dtype=torch.bfloat16, + exec_mode: str | None = None, device="neuron", **kwargs: Any) -> None: + super().__init__(model_path=model_path, parallel=parallel, dtype=dtype, + exec_mode=exec_mode, device=device, **kwargs) + if int(self.parallel.world_size) != 1: + raise NotImplementedError("the H3 video VAE runs on one rank") + if self.dtype not in (torch.bfloat16, torch.float32): + raise ValueError(f"the H3 video VAE runs in bf16 or fp32, got {self.dtype}") + with open(Path(str(self.model_path)) / "config.json", encoding="utf-8") as handle: + self.vae_config = json.load(handle) + + def build_module(self) -> torch.nn.Module: + return _NeuronH3VideoDecoder(self.vae_config) + + def get_example_inputs(self) -> tuple[torch.Tensor, ...]: + gen = torch.Generator().manual_seed(0) + return (torch.randn(TILE_LATENT_SHAPE, generator=gen),) + + +class DeviceTileDecoder: + """``decode_minimax_h3_video``'s decoder callable on the device: a host tile in, host fp32 + frames out.""" + + def __init__(self, app: NeuronMiniMaxH3VideoVaeApplication) -> None: + self.app = app + self.config = app.vae_config + + def __call__(self, tile: torch.Tensor) -> torch.Tensor: + out = self.app(tile.float().contiguous().to(self.app.device)) + return out.cpu() + + +__all__ = [ + "DeviceTileDecoder", + "NeuronMiniMaxH3VideoVaeApplication", + "TILE_LATENT_SHAPE", +] diff --git a/difflet/cli/orchestrators/minimax_h3.py b/difflet/cli/orchestrators/minimax_h3.py index ccf8daf1b..af185b3c0 100644 --- a/difflet/cli/orchestrators/minimax_h3.py +++ b/difflet/cli/orchestrators/minimax_h3.py @@ -16,10 +16,13 @@ * generate: the DiT on ``--tp-degree`` torchrun ranks, one NeuronCore each (``NeuronMiniMaxH3TransformerApplication``, per-block compiled, AdaLN precomputed on the host); every rank runs the same host denoise loop, rank 0 writes the latents; -* video_vae, audio_vae: the two decoders on the host CPU in fp32, one rank. +* video_vae: the ViT decoder on one NeuronCore (``NeuronMiniMaxH3VideoVaeApplication``, one + 256 px tile per call, ``--vae-dtype`` fp32 by default or bf16), or on the host in fp32 with + ``--host-vae``; +* audio_vae: the audio decoder on the host CPU in fp32, one rank. -``compile`` then only warms the DiT up (it fills the NEFF cache); the host stages have nothing -to compile. +``compile`` then warms up the device stages only (the DiT and, unless ``--host-vae``, the video +decoder), which fills the NEFF cache; the host stages have nothing to compile. """ from __future__ import annotations @@ -35,7 +38,13 @@ from difflet.cli import runner from difflet.cli.orchestrators.base import ModelOrchestrator -from difflet.cli.orchestrators.torchrun import root_compile_cache_at_cache_dir, torchrun_backend +from difflet.cli.orchestrators.torchrun import ( + decodes_on_device, + refuse_torchrun_conflicts, + root_compile_cache_at_cache_dir, + torchrun_backend, + vae_dtype_conflict, +) from difflet.common.orchestrators import minimax_h3 as h3_common if TYPE_CHECKING: @@ -156,8 +165,10 @@ def compile(self) -> None: shared = self._shared_cli_args(stage_mode="compile", work_dir=str(work_dir)) stages = _STAGES if torchrun_backend() is not None: + refuse_torchrun_conflicts([c for c in (vae_dtype_conflict(self.args),) if c]) root_compile_cache_at_cache_dir(self.args.cache_dir, tag="minimax-h3") - stages = ("generate",) # the host stages have nothing to compile + # The host stages have nothing to compile. + stages = ("generate", "video_vae") if self._device_vae() else ("generate",) for stage in stages: runner.run_stage( _HF_MODEL_ID, @@ -175,6 +186,7 @@ def generate(self) -> None: work_dir.mkdir(parents=True, exist_ok=True) shared = self._shared_cli_args(stage_mode="generate", work_dir=str(work_dir)) if torchrun_backend() is not None: + refuse_torchrun_conflicts([c for c in (vae_dtype_conflict(self.args),) if c]) root_compile_cache_at_cache_dir(self.args.cache_dir, tag="minimax-h3") try: for stage in _STAGES: @@ -196,7 +208,8 @@ def _run_stage_internal(self, stage: str, args: argparse.Namespace) -> None: handlers = { "text": self._stage_text_host, "generate": self._stage_generate_neuron, - "video_vae": self._stage_video_vae_host, + "video_vae": (self._stage_video_vae_neuron if getattr(args, "device_vae", False) + else self._stage_video_vae_host), "audio_vae": self._stage_audio_vae_host, } else: @@ -746,6 +759,59 @@ def _stage_video_vae_host(self, args: argparse.Namespace) -> None: print(f"[video_vae] (host) RGB video {tuple(video.shape)} in " f"{time.perf_counter() - start:.1f}s -> {destination}", flush=True) + def _device_vae(self) -> bool: + """The neuron backend decodes the video on a NeuronCore unless --host-vae.""" + return decodes_on_device(self.args, default=True) + + def _stage_video_vae_neuron(self, args: argparse.Namespace) -> None: + """The video decoder on one NeuronCore, one tile per call; ``compile`` loads it with + the warm-up decode (which compiles the tile), ``generate`` skips the warm-up.""" + import time + + import torch + + from difflet.backends import get_backend + from difflet.cli.dp import stage_loop + from difflet.cli.orchestrators.torchrun import component_dtype + from difflet.pipeline.parallel_config import DiffletParallelConfig + from difflet.pipeline.path_resolver import resolve_model_path + + get_backend("neuron").prepare_runtime(DiffletParallelConfig(tp_degree=1)) + from difflet.backends.neuron.minimax_h3.vae import ( + DeviceTileDecoder, + NeuronMiniMaxH3VideoVaeApplication, + ) + from difflet.models.minimax_h3.pipeline import decode_minimax_h3_video + + model_dir = self._model_dir(args, resolve_model_path) + dtype = getattr(torch, component_dtype(getattr(args, "vae_dtype", None))) + app = NeuronMiniMaxH3VideoVaeApplication(model_path=str(Path(model_dir) / "vae"), + dtype=dtype) + start = time.perf_counter() + app.load(skip_warmup=args.stage_mode != "compile") + print(f"[video_vae] device decoder ({dtype}) loaded in " + f"{time.perf_counter() - start:.1f}s", flush=True) + if args.stage_mode == "compile": + return + decoder = DeviceTileDecoder(app) + cfg = app.vae_config + for req in stage_loop.claimed_requests(args): + with stage_loop.request_scope(args, req, final=False): + latents = torch.load(stage_loop.work_file(args, req, "latents.pt")) + start = time.perf_counter() + video = decode_minimax_h3_video( + decoder, + latents["video_latents"], + latents_mean=cfg["latents_mean"], + latents_std=cfg["latents_std"], + clip_length=int(cfg["clip_length"]), + token_drop=int(cfg["token_drop"]), + ) + destination = stage_loop.work_file(args, req, "video.pt") + torch.save(video.to(torch.float16), destination) + print(f"[video_vae] (device) RGB video {tuple(video.shape)} in " + f"{time.perf_counter() - start:.1f}s -> {destination}", flush=True) + def _stage_audio_vae_host(self, args: argparse.Namespace) -> None: import time @@ -865,6 +931,10 @@ def _shared_cli_args(self, stage_mode: str, work_dir: str | None = None) -> list ): if value: parts.extend([flag, str(value)]) + if torchrun_backend() is not None and self._device_vae(): + parts.append("--device-vae") + if getattr(args, "vae_dtype", None): + parts.extend(["--vae-dtype", str(args.vae_dtype)]) if getattr(args, "requests_dir", None): parts.extend( [ diff --git a/difflet/models/minimax_h3/host_components.py b/difflet/models/minimax_h3/host_components.py index 86b5d8585..ad696857c 100644 --- a/difflet/models/minimax_h3/host_components.py +++ b/difflet/models/minimax_h3/host_components.py @@ -190,7 +190,11 @@ def __call__(self, latents: torch.Tensor) -> torch.Tensor: class _HostVideoDecoder(nn.Module): - def __init__(self, cfg: dict[str, Any]) -> None: + """``post_quant_conv`` then the ViT decoder, keys as the checkpoint's. ``sequence_alignment`` + pads the token sequence (with a key mask) for kernels that need aligned lengths; 1 is the + reference.""" + + def __init__(self, cfg: dict[str, Any], *, sequence_alignment: int = 1) -> None: from difflet.models.minimax_h3.modeling_vae import MiniMaxH3VideoViTDecoder3d super().__init__() @@ -211,6 +215,7 @@ def __init__(self, cfg: dict[str, Any]) -> None: ffn_mult=int(cfg["decoder_ffn_mult"]), rope_theta=float(cfg["decoder_rope_theta"]), rope_dim_ratio=float(cfg["decoder_rope_dim_ratio"]), norm_eps=float(cfg["decoder_norm_eps"]), + sequence_alignment=int(sequence_alignment), ) def forward(self, latents: torch.Tensor) -> torch.Tensor: diff --git a/tests/unit/backends/_neuron_h3_toy.py b/tests/unit/backends/_neuron_h3_toy.py index 1c3e2995d..df8010413 100644 --- a/tests/unit/backends/_neuron_h3_toy.py +++ b/tests/unit/backends/_neuron_h3_toy.py @@ -157,3 +157,24 @@ def l1_h3_worker(rank: int, world_size: int, transformer_dir: str, precomputed: "audio": audio.float().tolist(), "other_differs": not bool(torch.equal(video, other)), } + + +def vae_worker(rank: int, world_size: int, vae_dir: str, exec_mode: str) -> dict: + """The tiny H3 video VAE as a neuron application vs the host reference decoder, one tile.""" + from difflet.backends.neuron.minimax_h3.vae import ( + TILE_LATENT_SHAPE, + DeviceTileDecoder, + NeuronMiniMaxH3VideoVaeApplication, + ) + from difflet.models.minimax_h3.host_components import load_host_video_decoder + + app = NeuronMiniMaxH3VideoVaeApplication(model_path=vae_dir, dtype=torch.float32, + exec_mode=exec_mode, device="cpu") + torch._dynamo.reset() + app.load() + tile = torch.randn(TILE_LATENT_SHAPE, generator=torch.Generator().manual_seed(3)) + out = DeviceTileDecoder(app)(tile) + ref = load_host_video_decoder(str(Path(vae_dir).parent))(tile) + return {"max_abs": float((out - ref).abs().max()), "shape": list(out.shape), + "dtype": str(out.dtype), "report": app.load_report, + "norm_dtype": str(app.module.decoder.transformer_blocks[0].norm1.weight.dtype)} diff --git a/tests/unit/backends/test_neuron_minimax_h3_vae.py b/tests/unit/backends/test_neuron_minimax_h3_vae.py new file mode 100644 index 000000000..b3f3e1784 --- /dev/null +++ b/tests/unit/backends/test_neuron_minimax_h3_vae.py @@ -0,0 +1,59 @@ +"""L1 (CPU): the H3 video VAE decoder as a neuron application, against the reference decoder. + +A tiny decoder (2 blocks, 2 heads x 16) with random weights, saved in the checkpoint layout +(encoder keys included, to be skipped). The application pads the tile's 1797 tokens to 1920 and +attends through ``difflet.ops.attention`` with key bounds; the reference is the unpadded host +decoder with diffusers' attention. One rank, CPU, fp32. +""" + +import pytest + +torch = pytest.importorskip("torch") + +import json # noqa: E402 + +from tests.unit.backends._neuron_gloo import run_ranks # noqa: E402 +from tests.unit.backends._neuron_h3_toy import vae_worker # noqa: E402 + +TINY_VAE = { + "out_channels": 3, "latent_channels": 24, + "spatial_downsample_factors": [2, 2, 2, 2, 1, 1], + "temporal_downsample_factors": [1, 2, 2, 1, 1, 1], + "decoder_num_layers": 2, "decoder_num_attention_heads": 2, + "decoder_attention_head_dim": 16, "decoder_num_register_tokens": 4, "decoder_ffn_mult": 2, + "decoder_rope_theta": 100.0, "decoder_rope_dim_ratio": 0.75, "decoder_norm_eps": 1e-5, + "clip_length": 17, "token_drop": 3, "latents_mean": [0.0] * 24, "latents_std": [1.0] * 24, +} + + +def _write(root): + from safetensors.torch import save_file + + from difflet.models.minimax_h3.host_components import _HostVideoDecoder + + torch.manual_seed(0) + module = _HostVideoDecoder(TINY_VAE) + with torch.no_grad(): + for p in module.parameters(): + p.normal_(0.0, 0.2) + vae = root / "vae" + vae.mkdir() + state = {k: v.contiguous() for k, v in module.state_dict().items()} + state["encoder.dummy.weight"] = torch.zeros(2) + state["quant_conv.weight"] = torch.zeros(2) + save_file(state, str(vae / "diffusion_pytorch_model.safetensors")) + (vae / "config.json").write_text(json.dumps(TINY_VAE)) + return vae + + +@pytest.mark.parametrize("exec_mode", ["eager", "compile"]) +def test_device_vae_matches_reference_decoder(tmp_path, exec_mode): + pytest.importorskip("safetensors.torch") + pytest.importorskip("accelerate") + vae = _write(tmp_path) + (r,) = run_ranks(vae_worker, str(vae), exec_mode, world_size=1) + assert r["shape"] == [1, 3, 28, 256, 256] and r["dtype"] == "torch.float32" + assert r["max_abs"] < 1e-4, r["max_abs"] + assert r["report"]["missing"] == [] and r["report"]["unexpected"] == [] + assert sorted(r["report"]["skipped"]) == ["encoder.dummy.weight", "quant_conv.weight"] + assert r["norm_dtype"] == "torch.float32" diff --git a/tests/unit/cli/test_orchestrator_minimax_h3.py b/tests/unit/cli/test_orchestrator_minimax_h3.py index cedc64235..2dca9ce54 100644 --- a/tests/unit/cli/test_orchestrator_minimax_h3.py +++ b/tests/unit/cli/test_orchestrator_minimax_h3.py @@ -386,7 +386,8 @@ def test_neuron_backend_compiles_only_the_dit_and_places_host_stages_on_one_rank orchestrator = MiniMaxH3Orchestrator(_args(work_dir=str(tmp_path), keep_work_dir=True)) orchestrator.compile() - assert [stage for stage, _ in calls] == ["generate"] + assert [stage for stage, _ in calls] == ["generate", "video_vae"] + assert "--device-vae" in calls[0][1]["cli_args"] calls.clear() orchestrator.generate() @@ -395,17 +396,35 @@ def test_neuron_backend_compiles_only_the_dit_and_places_host_stages_on_one_rank ] +def test_neuron_backend_host_vae_compiles_only_the_dit_and_rejects_vae_dtype( + monkeypatch, tmp_path +): + monkeypatch.setattr(h3_mod, "torchrun_backend", lambda: "neuron") + monkeypatch.setattr(h3_mod, "root_compile_cache_at_cache_dir", lambda *a, **k: None) + calls = [] + monkeypatch.setattr( + h3_mod.runner, "run_stage", lambda model, stage, **kwargs: calls.append((stage, kwargs)) + ) + MiniMaxH3Orchestrator(_args(work_dir=str(tmp_path), host_vae=True)).compile() + assert [stage for stage, _ in calls] == ["generate"] + assert "--device-vae" not in calls[0][1]["cli_args"] + monkeypatch.setattr("difflet.cli.orchestrators.torchrun.torchrun_backend", lambda: "neuron") + with pytest.raises(SystemExit): + MiniMaxH3Orchestrator(_args(host_vae=True, vae_dtype="bf16")).compile() + + def test_neuron_backend_dispatches_to_host_and_torchrun_stage_bodies(monkeypatch): monkeypatch.setattr(h3_mod, "torchrun_backend", lambda: "neuron") orchestrator = MiniMaxH3Orchestrator(_args()) seen = [] for name in ("_stage_text_host", "_stage_generate_neuron", "_stage_video_vae_host", - "_stage_audio_vae_host"): + "_stage_video_vae_neuron", "_stage_audio_vae_host"): monkeypatch.setattr(orchestrator, name, lambda args, name=name: seen.append(name)) for stage in ("text", "generate", "video_vae", "audio_vae"): orchestrator._run_stage_internal(stage, orchestrator.args) + orchestrator._run_stage_internal("video_vae", _args(device_vae=True)) assert seen == ["_stage_text_host", "_stage_generate_neuron", "_stage_video_vae_host", - "_stage_audio_vae_host"] + "_stage_audio_vae_host", "_stage_video_vae_neuron"] def test_shared_cli_args_forward_model_path(): From 61492552d98bbef6442f6246dbe4568d58d18087 Mon Sep 17 00:00:00 2001 From: SHIYANG CHENG Date: Tue, 10 Nov 2026 18:08:27 -0800 Subject: [PATCH 6/8] feat(neuron, minimax-h3): Qwen3-VL conditioner on the device Qwen3VLTextEncoderModel (models/, backend-neutral) is Qwen3-VL's language model to hidden_states[50]: the embedding and the first 50 layers with the difflet.ops TP layers (GQA q/k/v without bias, per-head q/k RMSNorm, head_dim 128, the text case of the interleaved mrope = 1-D RoPE, theta 5e6), causal and unmasked on a right-padded prompt. NeuronMiniMaxH3TextEncoderApp loads model.language_model.* (11.7 GiB/rank at tp4), skips the vision tower, lm_head, the final norm and layers 50-63, and encodes with the Trainium stage's tokenization and text.pt record. L1 (CPU, gloo): tp1/tp2/tp4 eager and tp2 compile match transformers' Qwen3VLForConditionalGeneration hidden_states[3] (fp32, 2e-5) on a tiny model. trn3: hidden_states[50] vs the AOT/NxDI encode rel_l2 9.8e-4, mask equal. Co-Authored-By: Claude Opus 5.5 --- .../neuron/minimax_h3/text_encoder.py | 111 +++++++++++ .../modeling_qwen3vl_text_encoder.py | 172 ++++++++++++++++++ .../test_neuron_minimax_h3_text_encoder.py | 45 +++++ 3 files changed, 328 insertions(+) create mode 100644 difflet/backends/neuron/minimax_h3/text_encoder.py create mode 100644 difflet/models/minimax_h3/modeling_qwen3vl_text_encoder.py create mode 100644 tests/unit/backends/test_neuron_minimax_h3_text_encoder.py diff --git a/difflet/backends/neuron/minimax_h3/text_encoder.py b/difflet/backends/neuron/minimax_h3/text_encoder.py new file mode 100644 index 000000000..203d5a3cc --- /dev/null +++ b/difflet/backends/neuron/minimax_h3/text_encoder.py @@ -0,0 +1,111 @@ +"""Neuron application for MiniMax-H3's conditioner: Qwen3-VL's language model to hidden_states[50]. + +Wraps the backend-neutral ``Qwen3VLTextEncoderModel`` +(``difflet/models/minimax_h3/modeling_qwen3vl_text_encoder.py``): the embedding and the first 50 +decoder layers, TP-sharded over the DiT's TP group (about 12.6 GB per rank in bf16 at tp4). One +static ``(1, 1024)`` token shape (the Trainium text stage's bucket); in compile mode the 50 +identical layers compile into one shared graph. + +Checkpoint: transformers' Qwen3-VL layout, ``model.language_model.*``; the vision tower, +``lm_head``, the final norm and layers 50-63 are skipped on purpose. + +``encode(prompt)`` returns the Trainium text stage's record (``text.pt``): the hidden states with +the padding rows zeroed, the bool mask, the live token count and the layer, identical on every +rank (every rank runs the encode, the TP collectives need them all). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch + +from difflet.backends.neuron.core.application_base import TorchNeuronApplicationBase +from difflet.common.orchestrators import minimax_h3 as h3_common + +_PREFIX = "model.language_model." + + +class NeuronMiniMaxH3TextEncoderApplication(TorchNeuronApplicationBase): + """``model_path`` is ``/text_encoder``; the tokenizer is read from + ``/tokenizer``.""" + + block_attrs = ("layers",) + + def __init__(self, *, model_path, seq_len: int = h3_common.TEXT_SEQ_LEN, + layer: int = h3_common.TEXT_ENCODER_LAYER, parallel=None, + dtype=torch.bfloat16, exec_mode: str | None = None, device="neuron", + **kwargs: Any) -> None: + super().__init__(model_path=model_path, parallel=parallel, dtype=dtype, + exec_mode=exec_mode, device=device, **kwargs) + from difflet.models.minimax_h3.modeling_qwen3vl_text_encoder import ( + Qwen3VLTextEncoderConfig, + ) + + if self.dtype not in (torch.bfloat16, torch.float32): + raise ValueError(f"the Qwen3-VL text encoder runs in bf16 or fp32, got {self.dtype}") + self.seq_len = int(seq_len) + self.layer = int(layer) + self.encoder_config = Qwen3VLTextEncoderConfig.from_pretrained( + self.model_path, num_layers=self.layer) + self.encoder_config.validate(int(self.parallel.tp_degree)) + total = int(self.encoder_config.num_hidden_layers) + self.skipped_checkpoint_prefixes = ( + "model.visual.", "lm_head.", f"{_PREFIX}norm.", + *(f"{_PREFIX}layers.{i}." for i in range(self.layer, total)), + ) + self.tokenizer = None + + def build_module(self) -> torch.nn.Module: + from difflet.models.minimax_h3.modeling_qwen3vl_text_encoder import ( + Qwen3VLTextEncoderModel, + ) + + return Qwen3VLTextEncoderModel(self.encoder_config, max_positions=self.seq_len) + + def checkpoint_key(self, name: str) -> str: + return _PREFIX + name + + def get_example_inputs(self) -> tuple[torch.Tensor, ...]: + return (torch.zeros(1, self.seq_len, dtype=torch.int64),) + + def load(self, *args: Any, **kwargs: Any) -> None: + super().load(*args, **kwargs) + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(Path(str(self.model_path)).parent / "tokenizer") + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + self.tokenizer = tokenizer + + @torch.no_grad() + def encode(self, prompt: str) -> dict[str, Any]: + """The Trainium text stage's tokenization (raw prompt, no special tokens, right-padded to + ``seq_len``) and record.""" + tokens = self.tokenizer(prompt, add_special_tokens=False, max_length=self.seq_len, + padding="max_length", truncation=True, return_tensors="pt", + return_attention_mask=True) + mask = tokens.attention_mask + count = int(mask.sum()) + if count < 1: + raise ValueError("MiniMax-H3 prompt must produce at least one token") + hidden = self(tokens.input_ids.to(torch.int64).to(self.device)) + hidden = hidden.cpu().to(torch.bfloat16) + hidden[:, count:] = 0 # padding rows are never part of the packed document + return {"encoder_hidden_states": hidden, "attention_mask": mask.to(torch.bool), + "num_text_tokens": count, "text_encoder_layer": self.layer} + + def release(self) -> None: + """Drop the weights from the device (the generate stage frees this core's share before + the video decoder loads).""" + self.module = None + self.is_loaded = False + import gc + + gc.collect() + if self.device.type == "neuron": + torch.neuron.empty_cache() + + +__all__ = ["NeuronMiniMaxH3TextEncoderApplication"] diff --git a/difflet/models/minimax_h3/modeling_qwen3vl_text_encoder.py b/difflet/models/minimax_h3/modeling_qwen3vl_text_encoder.py new file mode 100644 index 000000000..921f373d7 --- /dev/null +++ b/difflet/models/minimax_h3/modeling_qwen3vl_text_encoder.py @@ -0,0 +1,172 @@ +"""Qwen3-VL's language model, text-only, as MiniMax-H3's conditioner. Hardware-agnostic. + +H3 conditions its DiT on ``hidden_states[50]`` of Qwen3-VL-32B's language model: the output of +decoder layer 49, before any later layer and without the final norm. This module is therefore +the embedding plus the first ``num_layers`` (50) decoder layers, with the TP layers of +``difflet.ops`` so the neuron backend can shard the ~25 B parameters (about 50 GB in bf16, 12.6 GB +per rank at tp4). Layers 50-63, the final norm, ``lm_head`` and the vision tower are never built. + +It mirrors transformers' ``Qwen3VLTextModel`` (4.57) on the path a text prompt takes, and reuses +the Qwen2.5-VL encoder's pieces (``difflet/models/qwen_image/modeling_qwen_text_encoder.py``): + +* ``embed_tokens``: ``ParallelEmbedding(shard_across_embedding=True)``; +* per layer: RMSNorm -> attention -> residual -> RMSNorm -> gated SiLU MLP -> residual; +* attention (``Qwen3VLTextAttention``): q/k/v without bias, column-parallel (64 / tp query heads, + 8 / tp key/value heads, grouped as in Qwen2.5), a per-head RMSNorm on q and on k + (``q_norm``/``k_norm``, over ``head_dim`` = 128, which is not ``hidden_size / heads``), RoPE, + causal attention, row-parallel ``o_proj``; +* RoPE: Qwen3-VL's interleaved multimodal RoPE gives a text prompt's three axes the same + positions, so every frequency takes the same angle and the result is the ordinary 1-D RoPE with + ``rope_theta`` (5e6) over ``head_dim``; computed once in fp32 for ``0..max_positions-1``. + +No attention mask: the prompt is right-padded and attention is causal, so the valid rows never +see the padding (the padded rows' outputs are dropped by the caller), exactly as the Trainium +stage's NxDI encode and the host encoder treat it. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import torch +import torch.nn as nn + + +@dataclass(frozen=True) +class Qwen3VLTextEncoderConfig: + """``text_config`` of the checkpoint's ``config.json``; ``num_layers`` is how many decoder + layers to build (H3: 50, for ``hidden_states[50]``). Defaults are Qwen3-VL-32B's.""" + + vocab_size: int = 151936 + hidden_size: int = 5120 + intermediate_size: int = 25600 + num_hidden_layers: int = 64 + num_attention_heads: int = 64 + num_key_value_heads: int = 8 + head_dim: int = 128 + rms_norm_eps: float = 1e-6 + rope_theta: float = 5000000.0 + hidden_act: str = "silu" + attention_bias: bool = False + num_layers: int = 50 + + def validate(self, tp_degree: int = 1) -> None: + heads, kv, tp = self.num_attention_heads, self.num_key_value_heads, int(tp_degree) + if heads % kv: + raise ValueError(f"Qwen3-VL has {heads} query heads, not a multiple of {kv} kv heads") + if self.hidden_act != "silu" or self.attention_bias: + raise NotImplementedError("Qwen3-VL text encoder: silu MLP and bias-free q/k/v only") + if not 1 <= int(self.num_layers) <= int(self.num_hidden_layers): + raise ValueError(f"num_layers={self.num_layers} outside 1..{self.num_hidden_layers}") + for name, value in (("num_attention_heads", heads), ("num_key_value_heads", kv), + ("hidden_size", self.hidden_size), + ("intermediate_size", self.intermediate_size)): + if value % tp: + raise ValueError(f"Qwen3-VL {name}={value} does not divide tp={tp}") + + @classmethod + def from_pretrained(cls, text_encoder_dir, **overrides) -> "Qwen3VLTextEncoderConfig": + raw = json.loads((Path(text_encoder_dir) / "config.json").read_text()) + text = dict(raw.get("text_config") or raw) + scaling = text.get("rope_scaling") or {} + if scaling.get("rope_type", scaling.get("type", "default")) != "default": + raise NotImplementedError(f"Qwen3-VL rope_scaling {scaling!r}") + names = set(cls.__dataclass_fields__) + values = {k: v for k, v in text.items() if k in names} + values.update(overrides) + return cls(**values) + + +class Qwen3VLTextAttention(nn.Module): + def __init__(self, config: Qwen3VLTextEncoderConfig, tp: int): + super().__init__() + from difflet.models.qwen_image.modeling_qwen_text_encoder import QwenTextRMSNorm + from difflet.ops import ColumnParallelLinear, RowParallelLinear, attention + + d, hd = config.hidden_size, config.head_dim + self.head_dim = hd + self.heads = config.num_attention_heads // tp + self.kv_heads = config.num_key_value_heads // tp + self.groups = self.heads // self.kv_heads + self.scale = hd**-0.5 + self._attention = attention + self.q_proj = ColumnParallelLinear(d, config.num_attention_heads * hd, bias=False, + gather_output=False) + self.k_proj = ColumnParallelLinear(d, config.num_key_value_heads * hd, bias=False, + gather_output=False) + self.v_proj = ColumnParallelLinear(d, config.num_key_value_heads * hd, bias=False, + gather_output=False) + self.o_proj = RowParallelLinear(config.num_attention_heads * hd, d, bias=False, + input_is_parallel=True, reduce_dtype=torch.bfloat16) + self.q_norm = QwenTextRMSNorm(hd, config.rms_norm_eps) + self.k_norm = QwenTextRMSNorm(hd, config.rms_norm_eps) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + from difflet.models.qwen_image.modeling_qwen_text_encoder import _repeat_kv, _rotate_half + + b, s, _ = x.shape + q = self.q_norm(self.q_proj(x).view(b, s, self.heads, self.head_dim)).transpose(1, 2) + k = self.k_norm(self.k_proj(x).view(b, s, self.kv_heads, self.head_dim)).transpose(1, 2) + v = self.v_proj(x).view(b, s, self.kv_heads, self.head_dim).transpose(1, 2) + q = (q * cos) + (_rotate_half(q) * sin) + k = (k * cos) + (_rotate_half(k) * sin) + k, v = _repeat_kv(k, self.groups), _repeat_kv(v, self.groups) + out = self._attention(q, k, v, scale=self.scale, causal=True, tp_q=True, tp_k=True) + return self.o_proj(out.transpose(1, 2).reshape(b, s, self.heads * self.head_dim)) + + +class Qwen3VLTextDecoderLayer(nn.Module): + """One ``Qwen3VLTextDecoderLayer``; every layer has the same graph (the compile unit).""" + + def __init__(self, config: Qwen3VLTextEncoderConfig, tp: int): + super().__init__() + from difflet.models.qwen_image.modeling_qwen_text_encoder import ( + QwenTextMLP, + QwenTextRMSNorm, + ) + + self.self_attn = Qwen3VLTextAttention(config, tp) + self.mlp = QwenTextMLP(config) + self.input_layernorm = QwenTextRMSNorm(config.hidden_size, config.rms_norm_eps) + self.post_attention_layernorm = QwenTextRMSNorm(config.hidden_size, config.rms_norm_eps) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x = x + self.self_attn(self.input_layernorm(x), cos, sin) + return x + self.mlp(self.post_attention_layernorm(x)) + + +class Qwen3VLTextEncoderModel(nn.Module): + """``input_ids [B, S]`` (right-padded, S <= ``max_positions``) -> ``hidden_states[num_layers]`` + ``[B, S, hidden]`` (no final norm). Parameter names are the checkpoint's without its + ``model.language_model.`` prefix.""" + + def __init__(self, config: Qwen3VLTextEncoderConfig, *, max_positions: int): + super().__init__() + from difflet.models.qwen_image.modeling_qwen_text_encoder import _safe_tp_size, rotary_tables + from difflet.ops import ParallelEmbedding + + tp = _safe_tp_size() + config.validate(tp) + self.config = config + self.embed_tokens = ParallelEmbedding(config.vocab_size, config.hidden_size, + shard_across_embedding=True) + self.layers = nn.ModuleList( + [Qwen3VLTextDecoderLayer(config, tp) for _ in range(int(config.num_layers))] + ) + cos, sin = rotary_tables(config.head_dim, int(max_positions), float(config.rope_theta)) + self.register_buffer("rope_cos", cos, persistent=False) + self.register_buffer("rope_sin", sin, persistent=False) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + x = self.embed_tokens(input_ids) + seq = int(input_ids.shape[1]) + cos = self.rope_cos[:seq].to(x.dtype)[None, None] + sin = self.rope_sin[:seq].to(x.dtype)[None, None] + for layer in self.layers: + x = layer(x, cos, sin) + return x + + +__all__ = ["Qwen3VLTextEncoderConfig", "Qwen3VLTextEncoderModel"] diff --git a/tests/unit/backends/test_neuron_minimax_h3_text_encoder.py b/tests/unit/backends/test_neuron_minimax_h3_text_encoder.py new file mode 100644 index 000000000..d09b0460d --- /dev/null +++ b/tests/unit/backends/test_neuron_minimax_h3_text_encoder.py @@ -0,0 +1,45 @@ +"""L1 (CPU, gloo): H3's Qwen3-VL conditioner as a neuron application vs transformers. + +A tiny Qwen3-VL (4 layers, 8 q heads / 4 kv heads x 16, interleaved mrope config) saved by +transformers; the application builds the first 3 layers and returns their output on the valid +rows of a right-padded prompt, unmasked and causal, against transformers' masked +``hidden_states[3]``. tp1/tp2/tp4 shard the embedding, the GQA heads and the MLP. +""" + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("transformers.models.qwen3_vl") + +from types import SimpleNamespace # noqa: E402 + +from tests.unit.backends._neuron_gloo import run_ranks # noqa: E402 +from tests.unit.backends._neuron_h3_toy import ( # noqa: E402 + TE_LAYER, + te_reference, + te_worker, + write_tiny_qwen3vl, +) + +OFFLINE = {"HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1"} + + +@pytest.fixture(scope="module") +def tiny(tmp_path_factory): + pytest.importorskip("accelerate") + encoder = write_tiny_qwen3vl(tmp_path_factory.mktemp("tiny_qwen3vl")) + return SimpleNamespace(dir=str(encoder), ref=te_reference(encoder)) + + +@pytest.mark.parametrize("tp,exec_mode", [(1, "eager"), (2, "eager"), (4, "eager"), + (2, "compile")]) +def test_text_encoder_matches_transformers(tiny, tp, exec_mode): + results = run_ranks(te_worker, tiny.dir, exec_mode, world_size=tp, env=OFFLINE) + for r in results: + torch.testing.assert_close(torch.tensor(r["hidden"]), tiny.ref, atol=2e-5, rtol=2e-5) + assert r["layers"] == TE_LAYER + assert r["report"]["missing"] == [] and r["report"]["unexpected"] == [] + skipped = r["report"]["skipped"] + assert "lm_head.weight" in skipped + assert any(k.startswith("model.language_model.layers.3.") for k in skipped) + assert any(k.startswith("model.visual.") for k in skipped) From 1ed539d0b86678278e991b94914a9d5bd70ef56b Mon Sep 17 00:00:00 2001 From: SHIYANG CHENG Date: Tue, 10 Nov 2026 18:08:27 -0800 Subject: [PATCH 7/8] feat(neuron, minimax-h3): text, DiT and video decode in one generate launch On the neuron backend the generate stage now encodes the prompt on the device (released before the DiT loads), denoises, and decodes the video with a decoder on every rank, the 256 px tiles spread over the ranks (decode_video_on_ranks: record the tile order, decode i % world, replay the official blending on rank 0; equal to the one-process decode). The separate device video_vae stage is gone; --host-text-encoder / --host-vae bring back the host text and video_vae stages. compile warms up all three components. trn3 device 7, 256x448x124, 30 steps: difflet generate 144 s -> 109 s (text 8.9 s, DiT 21.1 s, video decode 14.2 s -> 6.6 s); compile 70 s warm. Co-Authored-By: Claude Opus 5.5 --- difflet/backends/neuron/minimax_h3/vae.py | 59 ++++- difflet/cli/orchestrators/minimax_h3.py | 246 +++++++++++------- tests/unit/backends/_neuron_h3_toy.py | 110 ++++++++ .../backends/test_neuron_minimax_h3_vae.py | 12 + .../unit/cli/test_orchestrator_minimax_h3.py | 21 +- 5 files changed, 341 insertions(+), 107 deletions(-) diff --git a/difflet/backends/neuron/minimax_h3/vae.py b/difflet/backends/neuron/minimax_h3/vae.py index 8cf44ab53..ce9052a8d 100644 --- a/difflet/backends/neuron/minimax_h3/vae.py +++ b/difflet/backends/neuron/minimax_h3/vae.py @@ -97,7 +97,11 @@ def forward(self, latents: torch.Tensor) -> torch.Tensor: class NeuronMiniMaxH3VideoVaeApplication(TorchNeuronApplicationBase): - """The H3 video decoder on one rank, one tile per call. ``model_path`` is ``/vae``.""" + """The H3 video decoder, one tile per call. ``model_path`` is ``/vae``. + + It has no tensor-parallel layers: under a TP mesh (the DiT's, in the generate launch; a + process holds one mesh) every rank loads the whole decoder and decodes its own tiles + (``decode_video_on_ranks``).""" block_attrs = ("decoder.transformer_blocks",) skipped_checkpoint_prefixes = ("encoder.", "quant_conv.") @@ -106,10 +110,11 @@ def __init__(self, *, model_path, parallel=None, dtype=torch.bfloat16, exec_mode: str | None = None, device="neuron", **kwargs: Any) -> None: super().__init__(model_path=model_path, parallel=parallel, dtype=dtype, exec_mode=exec_mode, device=device, **kwargs) - if int(self.parallel.world_size) != 1: - raise NotImplementedError("the H3 video VAE runs on one rank") if self.dtype not in (torch.bfloat16, torch.float32): raise ValueError(f"the H3 video VAE runs in bf16 or fp32, got {self.dtype}") + for field in ("cp_degree", "dp_degree"): + if int(getattr(self.parallel, field, 1) or 1) > 1: + raise NotImplementedError(f"the H3 video VAE runs under TP only, got {field}") with open(Path(str(self.model_path)) / "config.json", encoding="utf-8") as handle: self.vae_config = json.load(handle) @@ -134,8 +139,56 @@ def __call__(self, tile: torch.Tensor) -> torch.Tensor: return out.cpu() + + +def decode_video_on_ranks(decoder, latents: torch.Tensor, cfg: dict, *, rank: int, world: int, + scratch_dir, device) -> torch.Tensor | None: + """``decode_minimax_h3_video`` with its tiles spread over the ranks, each holding a decoder. + + The decode calls its decoder once per tile in a fixed order, so it runs three times: once + with a recorder that keeps every tile (every rank, the latents are the same on all of them), + once per rank on the tiles ``i % world == rank`` (written to ``scratch_dir``), and once on + rank 0 replaying the decoded tiles in order, which does the official blending and stitching. + A collective after the writes makes rank 0 wait for every rank's file. Returns the video on + rank 0, None elsewhere. + """ + import torch.distributed as dist + + from difflet.models.minimax_h3.pipeline import decode_minimax_h3_video + + kwargs = dict(latents_mean=cfg["latents_mean"], latents_std=cfg["latents_std"], + clip_length=int(cfg["clip_length"]), token_drop=int(cfg["token_drop"])) + tiles: list[torch.Tensor] = [] + + def record(tile: torch.Tensor) -> torch.Tensor: + tiles.append(tile.clone()) + b, _, t, h, w = tile.shape + return torch.zeros(b, 3, t * 4, h * 16, w * 16) + + decode_minimax_h3_video(record, latents, **kwargs) + scratch = Path(scratch_dir) + scratch.mkdir(parents=True, exist_ok=True) + mine = {i: decoder(tiles[i]) for i in range(len(tiles)) if i % world == rank} + torch.save(mine, scratch / f"vae_tiles_rank{rank}.pt") + if world > 1: + done = torch.ones(1, device=device) + dist.all_reduce(done) + if device.type == "neuron": + torch.neuron.synchronize() + if rank != 0: + return None + decoded: dict[int, torch.Tensor] = {} + for r in range(world): + path = scratch / f"vae_tiles_rank{r}.pt" + decoded.update(torch.load(path)) + path.unlink() + order = iter(range(len(tiles))) + return decode_minimax_h3_video(lambda tile: decoded[next(order)], latents, **kwargs) + + __all__ = [ "DeviceTileDecoder", + "decode_video_on_ranks", "NeuronMiniMaxH3VideoVaeApplication", "TILE_LATENT_SHAPE", ] diff --git a/difflet/cli/orchestrators/minimax_h3.py b/difflet/cli/orchestrators/minimax_h3.py index af185b3c0..11855d8e7 100644 --- a/difflet/cli/orchestrators/minimax_h3.py +++ b/difflet/cli/orchestrators/minimax_h3.py @@ -9,20 +9,21 @@ post-processing remain on the host. Staging is required by HBM capacity, not because the heavy components are intended to execute on the host. -With DIFFLET_BACKEND naming the non-AoT torchrun backend (neuron, TorchNeuron), the same four -stages and work files run as follows (phase 1 of the port): - -* text: Qwen3-VL to hidden_states[50] on the host CPU (``HostQwen3VLTextEncoder``), one rank; -* generate: the DiT on ``--tp-degree`` torchrun ranks, one NeuronCore each - (``NeuronMiniMaxH3TransformerApplication``, per-block compiled, AdaLN precomputed on the - host); every rank runs the same host denoise loop, rank 0 writes the latents; -* video_vae: the ViT decoder on one NeuronCore (``NeuronMiniMaxH3VideoVaeApplication``, one - 256 px tile per call, ``--vae-dtype`` fp32 by default or bf16), or on the host in fp32 with - ``--host-vae``; -* audio_vae: the audio decoder on the host CPU in fp32, one rank. - -``compile`` then warms up the device stages only (the DiT and, unless ``--host-vae``, the video -decoder), which fills the NEFF cache; the host stages have nothing to compile. +With DIFFLET_BACKEND naming the non-AoT torchrun backend (neuron, TorchNeuron), the work files +stay the same and the device work runs in one torchrun launch of ``--tp-degree`` ranks, one +NeuronCore each (the generate stage): + +1. text: Qwen3-VL to hidden_states[50] TP-sharded on the ranks + (``NeuronMiniMaxH3TextEncoderApplication``, ``--text-encoder-dtype`` bf16 by default), then + its weights are released; with ``--host-text-encoder`` a host text stage runs it instead; +2. the DiT on every rank (``NeuronMiniMaxH3TransformerApplication``, per-block compiled, AdaLN + precomputed on the host); every rank runs the same host denoise loop; +3. the video decoder on every rank (``NeuronMiniMaxH3VideoVaeApplication``, ``--vae-dtype`` + fp32 by default), the 256 px tiles spread over the ranks; with ``--host-vae`` a host + video_vae stage decodes instead. + +The audio decoder and the AV mux stay on the host (audio_vae, one rank). ``compile`` runs the +generate stage with every device component's warm-up, which fills the NEFF cache. """ from __future__ import annotations @@ -167,8 +168,7 @@ def compile(self) -> None: if torchrun_backend() is not None: refuse_torchrun_conflicts([c for c in (vae_dtype_conflict(self.args),) if c]) root_compile_cache_at_cache_dir(self.args.cache_dir, tag="minimax-h3") - # The host stages have nothing to compile. - stages = ("generate", "video_vae") if self._device_vae() else ("generate",) + stages = ("generate",) # every device component loads in the generate launch for stage in stages: runner.run_stage( _HF_MODEL_ID, @@ -188,8 +188,9 @@ def generate(self) -> None: if torchrun_backend() is not None: refuse_torchrun_conflicts([c for c in (vae_dtype_conflict(self.args),) if c]) root_compile_cache_at_cache_dir(self.args.cache_dir, tag="minimax-h3") + stages = _STAGES if torchrun_backend() is None else self._neuron_stages() try: - for stage in _STAGES: + for stage in stages: runner.run_stage( _HF_MODEL_ID, stage, @@ -208,8 +209,7 @@ def _run_stage_internal(self, stage: str, args: argparse.Namespace) -> None: handlers = { "text": self._stage_text_host, "generate": self._stage_generate_neuron, - "video_vae": (self._stage_video_vae_neuron if getattr(args, "device_vae", False) - else self._stage_video_vae_host), + "video_vae": self._stage_video_vae_host, "audio_vae": self._stage_audio_vae_host, } else: @@ -630,15 +630,22 @@ def build() -> None: return torch.load(path) def _stage_generate_neuron(self, args: argparse.Namespace) -> None: - """The DiT on every torchrun rank; each runs the host denoise loop on identical inputs - (one CPU generator per request), rank 0 writes the latents.""" + """Every device component on every torchrun rank (see the module docstring): the text + encode (``--device-text-encoder``), the DiT denoise and the video decode + (``--device-vae``). Every rank runs the same host loops on identical inputs (one CPU + generator per request); rank 0 writes the work files.""" import time import torch from difflet.backends import get_backend from difflet.cli.dp import stage_loop - from difflet.cli.orchestrators.torchrun import replica_request_scope, torchrun_parallel + from difflet.cli.orchestrators.torchrun import ( + component_dtype, + replica_request_scope, + replica_requests, + torchrun_parallel, + ) from difflet.pipeline.path_resolver import resolve_model_path parallel = torchrun_parallel(args) @@ -654,7 +661,51 @@ def _stage_generate_neuron(self, args: argparse.Namespace) -> None: load_minimax_h3_schedulers, ) + compiling = args.stage_mode == "compile" + device_te = bool(getattr(args, "device_text_encoder", False)) + device_vae = bool(getattr(args, "device_vae", False)) model_dir = self._model_dir(args, resolve_model_path) + rank, world = world_info() + + def log(message: str) -> None: + if rank == 0: + print(message, flush=True) + + # 1. Text: encode every request this replica runs, then free the encoder's HBM. + requests = None + if device_te: + from difflet.backends.neuron.minimax_h3.text_encoder import ( + NeuronMiniMaxH3TextEncoderApplication, + ) + + te_dtype = getattr(torch, component_dtype(getattr(args, "text_encoder_dtype", None), + default="bf16")) + encoder = NeuronMiniMaxH3TextEncoderApplication( + model_path=str(Path(model_dir) / "text_encoder"), parallel=parallel, + dtype=te_dtype) + self._neuron_device = encoder.device + start = time.perf_counter() + encoder.load(skip_warmup=not compiling) + log(f"[text] device Qwen3-VL ({te_dtype}, first {encoder.layer} layers) loaded in " + f"{time.perf_counter() - start:.1f}s, {encoder.weight_bytes / 2**30:.2f} GiB/rank") + if not compiling: + requests = list(replica_requests(args, device=encoder.device)) + for req in requests: + with replica_request_scope(args, req, final=False): + start = time.perf_counter() + record = encoder.encode(req.prompt) + if rank == 0: + destination = stage_loop.work_file(args, req, "text.pt") + torch.save(record, destination) + log(f"[text] (device) hidden_states[{record['text_encoder_layer']}] " + f"({record['num_text_tokens']} live tokens) in " + f"{time.perf_counter() - start:.1f}s") + encoder.release() + del encoder + if requests is None and not compiling: + requests = list(stage_loop.claimed_requests(args)) + + # 2. The DiT. transformer_path = str(Path(model_dir) / "transformer") height = args.height or h3_common.DEFAULT_HEIGHT width = args.width or h3_common.DEFAULT_WIDTH @@ -668,24 +719,42 @@ def _stage_generate_neuron(self, args: argparse.Namespace) -> None: model_path=transformer_path, config=config, parallel=parallel, dtype=torch.bfloat16, ) self._neuron_device = app.device - rank = world_info()[0] start = time.perf_counter() app.load() - if rank == 0: - print(f"[generate] DiT loaded on {app.world_size} ranks in " - f"{time.perf_counter() - start:.1f}s " - f"({ {k: round(v, 1) for k, v in app.phase_seconds.items()} }), " - f"{app.weight_bytes / 2**30:.2f} GiB/rank", flush=True) - if args.stage_mode == "compile": + log(f"[generate] DiT loaded on {world} ranks in {time.perf_counter() - start:.1f}s " + f"({ {k: round(v, 1) for k, v in app.phase_seconds.items()} }), " + f"{app.weight_bytes / 2**30:.2f} GiB/rank") + + # 3. The video decoder, one copy per rank (it is not sharded; the tiles are). + vae = None + if device_vae: + from difflet.backends.neuron.minimax_h3.vae import ( + DeviceTileDecoder, + NeuronMiniMaxH3VideoVaeApplication, + decode_video_on_ranks, + ) + vae_dtype = getattr(torch, component_dtype(getattr(args, "vae_dtype", None))) + # The DiT's mesh (one per process); the decoder has no TP layers, so every rank + # holds all of it. + vae_app = NeuronMiniMaxH3VideoVaeApplication( + model_path=str(Path(model_dir) / "vae"), parallel=parallel, dtype=vae_dtype) + start = time.perf_counter() + # Each rank loads the whole decoder; its warm-up (compile only) runs per rank. + vae_app.load(skip_warmup=not compiling) + vae = DeviceTileDecoder(vae_app) + log(f"[video_vae] device decoder ({vae_dtype}) on every rank, loaded in " + f"{time.perf_counter() - start:.1f}s") + if compiling: return dit = _NeuronH3DiT(app) - for req in stage_loop.claimed_requests(args): + for req in requests: with replica_request_scope(args, req, final=False): text = torch.load(stage_loop.work_file(args, req, "text.pt")) steps = int(stage_loop.effective(req, args, "steps", h3_common.DEFAULT_STEPS)) table = self._neuron_adaln_table(args, model_dir, config, steps) video_scheduler, audio_scheduler = load_minimax_h3_schedulers(model_dir) + dit.seconds.clear() start = time.perf_counter() output = denoise_minimax_h3_t2va( dit, @@ -702,27 +771,36 @@ def _stage_generate_neuron(self, args: argparse.Namespace) -> None: model_dtype=torch.bfloat16, adaln_modulation_table=table, ) - if rank != 0: - continue elapsed = time.perf_counter() - start - destination = stage_loop.work_file(args, req, "latents.pt") - torch.save( - { - "video_latents": output.video_latents.cpu(), - "audio_latents": output.audio_latents.cpu(), - "height": height, - "width": width, - "num_frames": num_frames, - }, - destination, - ) - print( - f"[generate] {steps} steps in {elapsed:.1f}s " + if rank == 0: + destination = stage_loop.work_file(args, req, "latents.pt") + torch.save( + { + "video_latents": output.video_latents.cpu(), + "audio_latents": output.audio_latents.cpu(), + "height": height, + "width": width, + "num_frames": num_frames, + }, + destination, + ) + log(f"[generate] {steps} steps in {elapsed:.1f}s " f"({1000 * sum(dit.seconds) / max(len(dit.seconds), 1):.0f} ms/DiT call); " f"video latents {tuple(output.video_latents.shape)}, audio latents " - f"{tuple(output.audio_latents.shape)} -> {destination}", - flush=True, + f"{tuple(output.audio_latents.shape)}") + if vae is None: + continue + start = time.perf_counter() + video = decode_video_on_ranks( + vae, output.video_latents, vae.config, rank=rank, world=world, + scratch_dir=stage_loop.work_file(args, req, "vae_tiles"), + device=app.device, ) + if rank == 0: + destination = stage_loop.work_file(args, req, "video.pt") + torch.save(video.to(torch.float16), destination) + log(f"[video_vae] (device, tiles over {world} ranks) RGB video " + f"in {time.perf_counter() - start:.1f}s") def _stage_video_vae_host(self, args: argparse.Namespace) -> None: import time @@ -760,57 +838,22 @@ def _stage_video_vae_host(self, args: argparse.Namespace) -> None: f"{time.perf_counter() - start:.1f}s -> {destination}", flush=True) def _device_vae(self) -> bool: - """The neuron backend decodes the video on a NeuronCore unless --host-vae.""" + """The neuron backend decodes the video on the NeuronCores unless --host-vae.""" return decodes_on_device(self.args, default=True) - def _stage_video_vae_neuron(self, args: argparse.Namespace) -> None: - """The video decoder on one NeuronCore, one tile per call; ``compile`` loads it with - the warm-up decode (which compiles the tile), ``generate`` skips the warm-up.""" - import time - - import torch - - from difflet.backends import get_backend - from difflet.cli.dp import stage_loop - from difflet.cli.orchestrators.torchrun import component_dtype - from difflet.pipeline.parallel_config import DiffletParallelConfig - from difflet.pipeline.path_resolver import resolve_model_path - - get_backend("neuron").prepare_runtime(DiffletParallelConfig(tp_degree=1)) - from difflet.backends.neuron.minimax_h3.vae import ( - DeviceTileDecoder, - NeuronMiniMaxH3VideoVaeApplication, + def _device_text_encoder(self) -> bool: + """The neuron backend encodes on the NeuronCores unless --host-text-encoder.""" + return not getattr(self.args, "host_text_encoder", False) + + def _neuron_stages(self) -> tuple[str, ...]: + """The stages a neuron ``generate`` runs: the generate launch, the host stages it does + not absorb, and the audio decode.""" + return ( + *(() if self._device_text_encoder() else ("text",)), + "generate", + *(() if self._device_vae() else ("video_vae",)), + "audio_vae", ) - from difflet.models.minimax_h3.pipeline import decode_minimax_h3_video - - model_dir = self._model_dir(args, resolve_model_path) - dtype = getattr(torch, component_dtype(getattr(args, "vae_dtype", None))) - app = NeuronMiniMaxH3VideoVaeApplication(model_path=str(Path(model_dir) / "vae"), - dtype=dtype) - start = time.perf_counter() - app.load(skip_warmup=args.stage_mode != "compile") - print(f"[video_vae] device decoder ({dtype}) loaded in " - f"{time.perf_counter() - start:.1f}s", flush=True) - if args.stage_mode == "compile": - return - decoder = DeviceTileDecoder(app) - cfg = app.vae_config - for req in stage_loop.claimed_requests(args): - with stage_loop.request_scope(args, req, final=False): - latents = torch.load(stage_loop.work_file(args, req, "latents.pt")) - start = time.perf_counter() - video = decode_minimax_h3_video( - decoder, - latents["video_latents"], - latents_mean=cfg["latents_mean"], - latents_std=cfg["latents_std"], - clip_length=int(cfg["clip_length"]), - token_drop=int(cfg["token_drop"]), - ) - destination = stage_loop.work_file(args, req, "video.pt") - torch.save(video.to(torch.float16), destination) - print(f"[video_vae] (device) RGB video {tuple(video.shape)} in " - f"{time.perf_counter() - start:.1f}s -> {destination}", flush=True) def _stage_audio_vae_host(self, args: argparse.Namespace) -> None: import time @@ -931,10 +974,15 @@ def _shared_cli_args(self, stage_mode: str, work_dir: str | None = None) -> list ): if value: parts.extend([flag, str(value)]) - if torchrun_backend() is not None and self._device_vae(): - parts.append("--device-vae") - if getattr(args, "vae_dtype", None): - parts.extend(["--vae-dtype", str(args.vae_dtype)]) + if torchrun_backend() is not None: + if self._device_vae(): + parts.append("--device-vae") + if getattr(args, "vae_dtype", None): + parts.extend(["--vae-dtype", str(args.vae_dtype)]) + if self._device_text_encoder(): + parts.append("--device-text-encoder") + if getattr(args, "text_encoder_dtype", None): + parts.extend(["--text-encoder-dtype", str(args.text_encoder_dtype)]) if getattr(args, "requests_dir", None): parts.extend( [ diff --git a/tests/unit/backends/_neuron_h3_toy.py b/tests/unit/backends/_neuron_h3_toy.py index df8010413..fbf3623e2 100644 --- a/tests/unit/backends/_neuron_h3_toy.py +++ b/tests/unit/backends/_neuron_h3_toy.py @@ -178,3 +178,113 @@ def vae_worker(rank: int, world_size: int, vae_dir: str, exec_mode: str) -> dict return {"max_abs": float((out - ref).abs().max()), "shape": list(out.shape), "dtype": str(out.dtype), "report": app.load_report, "norm_dtype": str(app.module.decoder.transformer_blocks[0].norm1.weight.dtype)} + + +# --------------------------------------------------------------------------- text encoder + +TINY_QWEN3VL_TEXT = dict( + vocab_size=128, hidden_size=64, intermediate_size=128, num_hidden_layers=4, + num_attention_heads=8, num_key_value_heads=4, head_dim=16, rms_norm_eps=1e-6, + rope_theta=5000000.0, max_position_embeddings=64, + rope_scaling={"rope_type": "default", "mrope_section": [3, 3, 2], "mrope_interleaved": True}, + pad_token_id=0, +) +TE_LAYER = 3 +TE_SEQ = 16 +TE_LIVE = 11 + + +def write_tiny_qwen3vl(root) -> Path: + """A tiny Qwen3VLForConditionalGeneration saved by transformers (``model.language_model.*``, + ``model.visual.*``, ``lm_head``), as H3's ``text_encoder/``.""" + from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration + + torch.manual_seed(0) + config = Qwen3VLConfig( + text_config=TINY_QWEN3VL_TEXT, + vision_config=dict(depth=1, hidden_size=32, intermediate_size=64, num_heads=2, + out_hidden_size=64, patch_size=4, spatial_merge_size=2, + temporal_patch_size=2, deepstack_visual_indexes=[0], + num_position_embeddings=16), + ) + model = Qwen3VLForConditionalGeneration(config).eval() + with torch.no_grad(): + for p in model.parameters(): + p.normal_(0.0, 0.1) + out = Path(root) / "text_encoder" + model.save_pretrained(out, safe_serialization=True) + return out + + +def tiny_te_ids() -> torch.Tensor: + ids = torch.zeros(1, TE_SEQ, dtype=torch.int64) + ids[0, :TE_LIVE] = torch.randint(1, 128, (TE_LIVE,), generator=torch.Generator().manual_seed(2)) + return ids + + +def te_reference(encoder_dir) -> torch.Tensor: + """transformers' model: hidden_states[TE_LAYER] (layer TE_LAYER-1's output), with the mask.""" + from transformers import Qwen3VLForConditionalGeneration + + model = Qwen3VLForConditionalGeneration.from_pretrained(encoder_dir, + torch_dtype=torch.float32).eval() + ids = tiny_te_ids() + mask = (torch.arange(TE_SEQ) < TE_LIVE).to(torch.int64)[None] + with torch.no_grad(): + out = model.model.language_model(input_ids=ids, attention_mask=mask, + output_hidden_states=True, use_cache=False) + return out.hidden_states[TE_LAYER][:, :TE_LIVE] + + +def te_worker(rank: int, world_size: int, encoder_dir: str, exec_mode: str) -> dict: + from difflet.backends.neuron.minimax_h3.text_encoder import ( + NeuronMiniMaxH3TextEncoderApplication, + ) + from difflet.pipeline.parallel_config import DiffletParallelConfig + + app = NeuronMiniMaxH3TextEncoderApplication( + model_path=encoder_dir, seq_len=TE_SEQ, layer=TE_LAYER, + parallel=DiffletParallelConfig(tp_degree=world_size), dtype=torch.float32, + exec_mode=exec_mode, device="cpu", + ) + torch._dynamo.reset() + from difflet.backends.neuron.core.application_base import TorchNeuronApplicationBase + + TorchNeuronApplicationBase.load(app) # no tokenizer in the tiny snapshot + hidden = app(tiny_te_ids()) + return {"hidden": hidden[:, :TE_LIVE].tolist(), "report": app.load_report, + "layers": len(app.module.layers)} + + +def vae_ranks_worker(rank: int, world_size: int, vae_dir: str, scratch: str) -> dict: + """``decode_video_on_ranks`` over the gloo ranks vs one process decoding every tile.""" + from difflet.backends.neuron.minimax_h3.vae import ( + DeviceTileDecoder, + NeuronMiniMaxH3VideoVaeApplication, + decode_video_on_ranks, + ) + from difflet.models.minimax_h3.pipeline import decode_minimax_h3_video + from difflet.pipeline.parallel_config import DiffletParallelConfig + + app = NeuronMiniMaxH3VideoVaeApplication( + model_path=vae_dir, parallel=DiffletParallelConfig(tp_degree=world_size), + dtype=torch.float32, exec_mode="eager", device="cpu") + app.load(skip_warmup=True) + decoder = DeviceTileDecoder(app) + latents = torch.randn(1, 24, 10, 16, 28, generator=torch.Generator().manual_seed(5)) + cfg = app.vae_config + calls = [] + counted = lambda tile: (calls.append(1), decoder(tile))[1] # noqa: E731 + video = decode_video_on_ranks(counted, latents, cfg, rank=rank, world=world_size, + scratch_dir=scratch, device=torch.device("cpu")) + result = {"tiles_decoded_here": len(calls)} + if rank == 0: + ref = decode_minimax_h3_video(decoder, latents, latents_mean=cfg["latents_mean"], + latents_std=cfg["latents_std"], + clip_length=int(cfg["clip_length"]), + token_drop=int(cfg["token_drop"])) + result.update(equal=bool(torch.equal(video, ref)), shape=list(video.shape), + leftovers=sorted(p.name for p in Path(scratch).iterdir())) + else: + result["video_is_none"] = video is None + return result diff --git a/tests/unit/backends/test_neuron_minimax_h3_vae.py b/tests/unit/backends/test_neuron_minimax_h3_vae.py index b3f3e1784..d9421b091 100644 --- a/tests/unit/backends/test_neuron_minimax_h3_vae.py +++ b/tests/unit/backends/test_neuron_minimax_h3_vae.py @@ -57,3 +57,15 @@ def test_device_vae_matches_reference_decoder(tmp_path, exec_mode): assert r["report"]["missing"] == [] and r["report"]["unexpected"] == [] assert sorted(r["report"]["skipped"]) == ["encoder.dummy.weight", "quant_conv.weight"] assert r["norm_dtype"] == "torch.float32" + + +def test_tiles_spread_over_ranks_match_one_process(tmp_path): + pytest.importorskip("safetensors.torch") + pytest.importorskip("accelerate") + from tests.unit.backends._neuron_h3_toy import vae_ranks_worker + + vae = _write(tmp_path) + r0, r1 = run_ranks(vae_ranks_worker, str(vae), str(tmp_path / "tiles"), world_size=2) + assert r0["equal"] and r0["shape"][:2] == [1, 3] and r0["leftovers"] == [] + assert r1["video_is_none"] + assert r0["tiles_decoded_here"] > 0 and r1["tiles_decoded_here"] > 0 diff --git a/tests/unit/cli/test_orchestrator_minimax_h3.py b/tests/unit/cli/test_orchestrator_minimax_h3.py index 2dca9ce54..eb28c54c4 100644 --- a/tests/unit/cli/test_orchestrator_minimax_h3.py +++ b/tests/unit/cli/test_orchestrator_minimax_h3.py @@ -386,14 +386,26 @@ def test_neuron_backend_compiles_only_the_dit_and_places_host_stages_on_one_rank orchestrator = MiniMaxH3Orchestrator(_args(work_dir=str(tmp_path), keep_work_dir=True)) orchestrator.compile() - assert [stage for stage, _ in calls] == ["generate", "video_vae"] - assert "--device-vae" in calls[0][1]["cli_args"] + assert [stage for stage, _ in calls] == ["generate"] + cli = calls[0][1]["cli_args"] + assert "--device-vae" in cli and "--device-text-encoder" in cli calls.clear() orchestrator.generate() + # text and the video decode run inside the generate launch + assert [(stage, kw["num_cores"]) for stage, kw in calls] == [ + ("generate", 4), ("audio_vae", 1) + ] + + calls.clear() + host = MiniMaxH3Orchestrator(_args(work_dir=str(tmp_path), keep_work_dir=True, + host_vae=True, host_text_encoder=True)) + host.generate() assert [(stage, kw["num_cores"]) for stage, kw in calls] == [ ("text", 1), ("generate", 4), ("video_vae", 1), ("audio_vae", 1) ] + cli = calls[0][1]["cli_args"] + assert "--device-vae" not in cli and "--device-text-encoder" not in cli def test_neuron_backend_host_vae_compiles_only_the_dit_and_rejects_vae_dtype( @@ -418,13 +430,12 @@ def test_neuron_backend_dispatches_to_host_and_torchrun_stage_bodies(monkeypatch orchestrator = MiniMaxH3Orchestrator(_args()) seen = [] for name in ("_stage_text_host", "_stage_generate_neuron", "_stage_video_vae_host", - "_stage_video_vae_neuron", "_stage_audio_vae_host"): + "_stage_audio_vae_host"): monkeypatch.setattr(orchestrator, name, lambda args, name=name: seen.append(name)) for stage in ("text", "generate", "video_vae", "audio_vae"): orchestrator._run_stage_internal(stage, orchestrator.args) - orchestrator._run_stage_internal("video_vae", _args(device_vae=True)) assert seen == ["_stage_text_host", "_stage_generate_neuron", "_stage_video_vae_host", - "_stage_audio_vae_host", "_stage_video_vae_neuron"] + "_stage_audio_vae_host"] def test_shared_cli_args_forward_model_path(): From 8995c6cc306857b48e693d729116feb5a1b47977 Mon Sep 17 00:00:00 2001 From: SHIYANG CHENG Date: Tue, 10 Nov 2026 18:22:15 -0800 Subject: [PATCH 8/8] test: allow the moved H3 TP rewrite its fp32 all-reduce default The reduce-dtype coverage test exempted MiniMax-H3 by its Trainium path; the TP rewrite now lives in models/minimax_h3/modeling_minimax_h3_tp.py and keeps the same default on purpose (the Trainium graph is pinned bit-identical). Co-Authored-By: Claude Opus 5.5 --- tests/unit/test_row_parallel_reduce_dtype_coverage.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_row_parallel_reduce_dtype_coverage.py b/tests/unit/test_row_parallel_reduce_dtype_coverage.py index 67798d2a3..d46429c20 100644 --- a/tests/unit/test_row_parallel_reduce_dtype_coverage.py +++ b/tests/unit/test_row_parallel_reduce_dtype_coverage.py @@ -13,7 +13,12 @@ ROOT = Path(__file__).resolve().parents[2] -_ALLOWED_PREFIXES = ("difflet/backends/trainium/minimax_h3/",) +# MiniMax-H3's TP rewrite (shared by the Trainium and neuron backends since it moved to +# models/) keeps the fp32 default; its Trainium graph is pinned bit-identical to it. +_ALLOWED_PREFIXES = ( + "difflet/backends/trainium/minimax_h3/", + "difflet/models/minimax_h3/modeling_minimax_h3_tp.py", +) # Sites still on the NxD fp32 default, by file (MiniMax-H3 aside): none. EXPECTED_DEFAULT_SITES: dict[str, int] = {}