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/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/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/difflet/backends/neuron/minimax_h3/vae.py b/difflet/backends/neuron/minimax_h3/vae.py new file mode 100644 index 000000000..ce9052a8d --- /dev/null +++ b/difflet/backends/neuron/minimax_h3/vae.py @@ -0,0 +1,194 @@ +"""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, 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.") + + 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 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) + + 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() + + + + +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/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/cli/orchestrators/minimax_h3.py b/difflet/cli/orchestrators/minimax_h3.py index 390b0a50c..11855d8e7 100644 --- a/difflet/cli/orchestrators/minimax_h3.py +++ b/difflet/cli/orchestrators/minimax_h3.py @@ -8,6 +8,22 @@ 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 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 @@ -23,6 +39,13 @@ from difflet.cli import runner from difflet.cli.orchestrators.base import ModelOrchestrator +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: @@ -95,6 +118,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 +164,12 @@ 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: + 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",) # every device component loads in the generate launch + for stage in stages: runner.run_stage( _HF_MODEL_ID, stage, @@ -130,8 +185,12 @@ 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") + 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, @@ -146,12 +205,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 +555,349 @@ 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: + """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 ( + component_dtype, + replica_request_scope, + replica_requests, + 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, + ) + + 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 + 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 + start = time.perf_counter() + app.load() + 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 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, + 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, + ) + elapsed = time.perf_counter() - start + 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)}") + 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 + + 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 _device_vae(self) -> bool: + """The neuron backend decodes the video on the NeuronCores unless --host-vae.""" + return decodes_on_device(self.args, default=True) + + 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", + ) + + 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 +916,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,10 +969,20 @@ 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: parts.extend([flag, str(value)]) + 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/difflet/models/minimax_h3/host_components.py b/difflet/models/minimax_h3/host_components.py new file mode 100644 index 000000000..ad696857c --- /dev/null +++ b/difflet/models/minimax_h3/host_components.py @@ -0,0 +1,289 @@ +"""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): + """``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__() + 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"]), + sequence_alignment=int(sequence_alignment), + ) + + 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", +] 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/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/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/_neuron_h3_toy.py b/tests/unit/backends/_neuron_h3_toy.py new file mode 100644 index 000000000..fbf3623e2 --- /dev/null +++ b/tests/unit/backends/_neuron_h3_toy.py @@ -0,0 +1,290 @@ +"""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)), + } + + +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)} + + +# --------------------------------------------------------------------------- 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_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/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"] 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) 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..d9421b091 --- /dev/null +++ b/tests/unit/backends/test_neuron_minimax_h3_vae.py @@ -0,0 +1,71 @@ +"""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" + + +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 3ba8a3509..eb28c54c4 100644 --- a/tests/unit/cli/test_orchestrator_minimax_h3.py +++ b/tests/unit/cli/test_orchestrator_minimax_h3.py @@ -369,3 +369,93 @@ 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"] + 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( + 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"): + 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 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 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] = {}