From 66cbe325fdec2e8d32bb42cd01536104294183bd Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:31:54 +0100 Subject: [PATCH] MiniMax-H3-Image: fused blocks and tabulated AdaLN modulation (0.18.0) Add MiniMaxH3ImageTransformer2DModel to the minimax_h3 fused family. tabulate_modulation replaces the per-block AdaLN projections, the timestep MLP and the output-norm projection with tables for the model's fixed timesteps (13.0 of 33.1 billion parameters); add_modulation_tables builds the empty tables before a tabulated checkpoint loads. fuse_linears no longer keeps every replaced projection alive until it returns, so a text encoder no longer holds its packed weights twice on the device during fusion. --- README.md | 4 +- pyproject.toml | 2 +- src/orbitquant/__init__.py | 2 +- src/orbitquant/fused/families/minimax_h3.py | 155 +++++++++++++++++++- src/orbitquant/fused/linear.py | 12 +- tests/test_minimax_h3_image.py | 87 +++++++++++ uv.lock | 2 +- 7 files changed, 255 insertions(+), 9 deletions(-) create mode 100644 tests/test_minimax_h3_image.py diff --git a/README.md b/README.md index e7b7e97..b1f680f 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,8 @@ The implementation is clean-room and Apache-2.0 licensed. - Packed-weight CUDA, Triton, and Metal inference paths that avoid a full dequantized weight matrix. - Fused transformer checkpoints for Krea 2, FLUX.2, Ideogram 4, Qwen-Image 2.1, - MiniMax-H3, Boogu-Image, FLUX.3 Klein, MOSS-Image, Fibo 2.0 and Vimoo 1.0 + MiniMax-H3, MiniMax-H3-Image, Boogu-Image, FLUX.3 Klein, MOSS-Image, Fibo 2.0 and + Vimoo 1.0 (`orbitquant.fused`): each block runs grouped INT8 GEMMs with norm/modulation prologues and SwiGLU/gated-residual epilogues plus an INT8 attention kernel, and the checkpoint stores only the @@ -632,6 +633,7 @@ the same for an `orbitquant-v1` component artifact, and | Ideogram 4 | `Ideogram4Transformer2DModel` | sandwich-norm blocks with row-indexed modulation | | Qwen-Image 2.1 | `QwenImage21Transformer2DModel` | block-causal prefill and KV-cached decode | | MiniMax-H3 | `MiniMaxH3Transformer3DModel` | per-row (timestep, modality) AdaLN tables, partial rotate-half RoPE | +| MiniMax-H3-Image | `MiniMaxH3ImageTransformer2DModel` | the MiniMax-H3 blocks; `minimax_h3.tabulate_modulation(model, timesteps)` replaces the per-block AdaLN projections and the timestep MLP with tables for a fixed timestep set (the six-step schedule drops 13.0B of the 33.1B parameters); a timestep outside the set raises `ValueError` | | Boogu-Image | `BooguImageTransformer2DModel` | refiner, single- and double-stream blocks; 8-bit activations | | FLUX.3 Klein | `flux3.dit.DiT` (Black Forest Labs' `flux3` package) | per-mode and single-stream blocks: Q\|K\|V, SwiGLU, gated residuals; one sample at a time | | MOSS-Image | `MOSSImageTransformer` (the model's `pipeline_moss_image.py`) | joint image/text blocks: RMSNorm and shared AdaLN prologues, Q\|K\|V, SiTU gate\|up, gated residuals; one sample at a time | diff --git a/pyproject.toml b/pyproject.toml index c4d8a3e..507058a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.17.1" +version = "0.18.0" description = "Calibration-free OrbitQuant for transformer linear projections" readme = "README.md" requires-python = ">=3.11" diff --git a/src/orbitquant/__init__.py b/src/orbitquant/__init__.py index b73d06c..6c1f932 100644 --- a/src/orbitquant/__init__.py +++ b/src/orbitquant/__init__.py @@ -24,7 +24,7 @@ from orbitquant.quantizer import register_hf_quantizers from orbitquant.recipes import recipe -__version__ = "0.17.1" +__version__ = "0.18.0" register_hf_quantizers() diff --git a/src/orbitquant/fused/families/minimax_h3.py b/src/orbitquant/fused/families/minimax_h3.py index acb250a..49bbe32 100644 --- a/src/orbitquant/fused/families/minimax_h3.py +++ b/src/orbitquant/fused/families/minimax_h3.py @@ -1,4 +1,5 @@ -"""MiniMax-H3 (``MiniMaxH3Transformer3DModel``) main blocks. +"""MiniMax-H3 (``MiniMaxH3Transformer3DModel``) and MiniMax-H3-Image +(``MiniMaxH3ImageTransformer2DModel``) main blocks. One packed video/audio/text sequence; the AdaLN modulation of every row is a row of a (timestep, modality) table. The prologues and residual epilogues read that table through the @@ -6,13 +7,23 @@ RMSNorm and a rotate-half RoPE on the leading 96 channels of each 128-wide head. The token refiner blocks (text only, once per forward) keep their modules. +MiniMax-H3-Image samples a fixed grid of six timesteps (plus one noise level for reference rows), +and its AdaLN projections read nothing but the timestep embedding. ``tabulate_modulation`` stores +their outputs for those timesteps (``modulation.json`` lists them) and removes the projections: +13.0 of the 33.1 billion parameters. ``install`` then runs the model forward with table lookups; +a timestep outside the table raises an error. + Padding rows form their own attention document (the model's mask pairs live rows with live rows and padding rows with padding rows); the two documents run as separate mask-free attentions. """ from __future__ import annotations +import types + import torch +import torch.nn.functional as F +from torch import nn from orbitquant.fused.families import int8_attention_allowed, specs from orbitquant.fused.groups import RowSource @@ -20,7 +31,9 @@ from orbitquant.kernels import triton_int8_gemm as gemm FAMILY = "minimax_h3" -MODEL_CLASSES = ("MiniMaxH3Transformer3DModel",) +MODEL_CLASSES = ("MiniMaxH3Transformer3DModel", "MiniMaxH3ImageTransformer2DModel") +# Rows per timestep in the AdaLN tables: video/image, text, audio. +MODALITIES = 3 def blocks(model): @@ -143,3 +156,141 @@ def _forward(self, hidden_states, temb, adaln_indices, rotary_emb, attention_mas out=h, ) return h[None] + + +class ModulationTable(nn.Module): + """The outputs of one block's ``adaln_proj`` for a fixed set of timesteps. + + ``table`` holds ``adaln_proj.linear(silu(temb))`` reshaped to ``(timesteps * MODALITIES, + 6 * hidden)``, the row layout the block indexes with ``timestep_index * MODALITIES + tag``. + ``forward`` takes the table row of each distinct timestep of the call.""" + + def __init__(self, timesteps: int, width: int, dtype=torch.bfloat16, device=None): + super().__init__() + table = torch.empty(timesteps * MODALITIES, width, dtype=dtype, device=device) + self.register_buffer("table", table) + + def forward(self, rows: torch.Tensor) -> tuple[torch.Tensor, ...]: + tags = torch.arange(MODALITIES, device=rows.device) + index = (rows[:, None] * MODALITIES + tags).reshape(-1) + return self.table.index_select(0, index).chunk(6, dim=-1) + + +class ModulationKeys(nn.Module): + """The tabulated timesteps and the output norm's shift/scale for each of them.""" + + def __init__(self, timesteps: int, hidden: int, dtype=torch.bfloat16, device=None): + super().__init__() + keys = torch.empty(timesteps, dtype=torch.float32, device=device) + self.register_buffer("timesteps", keys) + output = torch.empty(timesteps, 2 * hidden, dtype=dtype, device=device) + self.register_buffer("output", output) + + def rows(self, timestep: torch.Tensor) -> torch.Tensor: + """Table row of every timestep (one host sync).""" + timestep = timestep.to(device=self.timesteps.device, dtype=torch.float32) + match = timestep[:, None] == self.timesteps[None] + if not bool(match.any(dim=1).all()): + raise ValueError( + f"timesteps {timestep.tolist()} are not all in the modulation table " + f"{self.timesteps.tolist()}; this checkpoint supports the six-step sampling grid " + "of MiniMax-H3-Image only" + ) + return match.to(torch.int32).argmax(dim=1) + + +def add_modulation_tables(model, timesteps: int, device=None) -> None: + """Replace the AdaLN projections and the timestep MLP by empty tables (before a tabulated + checkpoint loads).""" + hidden = model.config.hidden_size + dtype = model.transformer_blocks[0].adaln_proj.linear.weight.dtype + if dtype not in (torch.bfloat16, torch.float16, torch.float32): + dtype = torch.bfloat16 + for block in model.transformer_blocks: + block.adaln_proj = ModulationTable(timesteps, 6 * hidden, dtype=dtype, device=device) + model.oq_modulation = ModulationKeys(timesteps, hidden, dtype=dtype, device=device) + del model.time_embedder + del model.time_proj + del model.norm_out.linear + + +@torch.no_grad() +def tabulate_modulation(model, timesteps) -> list[float]: + """Compute the AdaLN tables of ``timesteps`` (unscaled, as the pipeline passes them) from the + loaded weights and remove the projections. Each timestep is projected on its own, as a + text-to-image call projects its one timestep.""" + from diffusers.models.modeling_utils import get_parameter_dtype + + keys = sorted({float(torch.tensor(t, dtype=torch.float32)) for t in timesteps}) + device = model.transformer_blocks[0].adaln_proj.linear.weight.device + embedder_dtype = get_parameter_dtype(model.time_embedder) + tables = [[] for _ in model.transformer_blocks] + output = [] + for t in keys: + timestep = torch.tensor([t], dtype=torch.float32, device=device) + temb = model.time_embedder(model.time_proj(timestep).to(embedder_dtype)) + for i, block in enumerate(model.transformer_blocks): + shift_scale = block.adaln_proj(temb) + tables[i].append(torch.cat(shift_scale, dim=-1)) + linear = model.norm_out.linear + output.append(linear(F.silu(temb).to(linear.weight.dtype))) + dtype = tables[0][0].dtype + add_modulation_tables(model, len(keys), device=device) + for block, rows in zip(model.transformer_blocks, tables, strict=True): + block.adaln_proj.table.copy_(torch.cat(rows, dim=0)) + model.oq_modulation.timesteps.copy_(torch.tensor(keys, dtype=torch.float32)) + model.oq_modulation.output.copy_(torch.cat(output, dim=0).to(dtype)) + install(model) + return keys + + +def install(model): + image_model = type(model).__name__ == "MiniMaxH3ImageTransformer2DModel" + if image_model and hasattr(model, "oq_modulation"): + model.forward = types.MethodType(_image_forward, model) + + +def _image_forward( + self, + hidden_states, + encoder_hidden_states, + timestep, + timestep_indices, + token_tags, + position_ids, + image_indices, + text_indices, + pdd_step_index, + return_dict=True, +): + """``MiniMaxH3ImageTransformer2DModel.forward`` with tabulated modulation.""" + from diffusers.models.modeling_outputs import Transformer2DModelOutput + from diffusers.models.modeling_utils import get_parameter_dtype + + rotary_emb = self.rope(position_ids) + image_embeds = self.proj_in(hidden_states.to(get_parameter_dtype(self.proj_in))) + text_embeds = self.token_refiner( + self.context_embedder(encoder_hidden_states.to(get_parameter_dtype(self.context_embedder))) + ) + sequence_length = position_ids.shape[0] + batch, width = text_embeds.shape[0], text_embeds.shape[-1] + hidden_states = text_embeds.new_zeros((batch, sequence_length, width)) + hidden_states = hidden_states.index_copy(1, text_indices, text_embeds) + hidden_states = hidden_states.index_copy(1, image_indices, image_embeds.to(text_embeds.dtype)) + + rows = self.oq_modulation.rows(timestep) + adaln_indices = timestep_indices * MODALITIES + token_tags + for block in self.transformer_blocks: + hidden_states = block(hidden_states, rows, adaln_indices, rotary_emb) + + proj_out = self.pdd_proj_out[pdd_step_index] + shift, scale = self.oq_modulation.output.index_select(0, rows).chunk(2, dim=-1) + hidden_states = self.norm_out.norm(hidden_states) + scale = scale.index_select(0, timestep_indices) + shift = shift.index_select(0, timestep_indices) + hidden_states = hidden_states * (1.0 + scale) + shift + output = proj_out(hidden_states.to(get_parameter_dtype(proj_out))) + output = output.index_select(1, image_indices) + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) diff --git a/src/orbitquant/fused/linear.py b/src/orbitquant/fused/linear.py index d92fd28..40bf914 100644 --- a/src/orbitquant/fused/linear.py +++ b/src/orbitquant/fused/linear.py @@ -47,9 +47,15 @@ def fuse_linears(model: nn.Module, *, activation_bits: int = 8, config=None) -> runtime = FusedRuntime.from_config(_quant_config(model, config), attention="sdpa") count = 0 - for name, module in list(model.named_modules()): - if not isinstance(module, OrbitQuantLinear) or module.packed_weight_indices is None: - continue + # Names, not modules: a list of the modules would keep every replaced projection's packed + # weights alive until the end, so the device would hold the model twice. + names = [ + name + for name, module in model.named_modules() + if isinstance(module, OrbitQuantLinear) and module.packed_weight_indices is not None + ] + for name in names: + module = model.get_submodule(name) holder = nn.Module() holder.proj = module group = PackedGroup.from_sources( diff --git a/tests/test_minimax_h3_image.py b/tests/test_minimax_h3_image.py new file mode 100644 index 0000000..71c41c6 --- /dev/null +++ b/tests/test_minimax_h3_image.py @@ -0,0 +1,87 @@ +"""MiniMax-H3-Image: tabulated AdaLN modulation reproduces the original transformer forward.""" + +import pytest +import torch + +diffusers = pytest.importorskip("diffusers") +if not hasattr(diffusers, "MiniMaxH3ImageTransformer2DModel"): + pytest.skip("diffusers without MiniMaxH3ImageTransformer2DModel", allow_module_level=True) + +from orbitquant.fused.families import minimax_h3 as family # noqa: E402 + +TIMESTEPS = [0.0, 0.25, 0.5, 0.999] + + +def _model(): + torch.manual_seed(0) + return diffusers.MiniMaxH3ImageTransformer2DModel( + num_attention_heads=2, + attention_head_dim=32, + hidden_size=48, + num_layers=2, + num_refiner_layers=1, + ffn_dim=64, + in_channels=4, + text_dim=16, + freq_dim=32, + time_embed_hidden_dim=48, + time_embed_dim=24, + rope_freq_dim=4, + num_pdd_heads=2, + ).eval() + + +def _inputs(timestep): + image_rows, text_rows = 8, 3 + seq = image_rows + text_rows + grid = torch.meshgrid(torch.arange(2), torch.arange(4), indexing="ij") + grid = torch.stack(grid, -1).reshape(-1, 2) + positions = torch.zeros(seq, 3, dtype=torch.long) + positions[text_rows:, 1:] = grid + positions[:text_rows, 0] = torch.arange(text_rows) + return dict( + hidden_states=torch.randn(1, image_rows, 16), + encoder_hidden_states=torch.randn(1, text_rows, 16), + timestep=torch.tensor(timestep, dtype=torch.float32), + timestep_indices=torch.tensor([0] * text_rows + [len(timestep) - 1] * image_rows), + token_tags=torch.tensor([1] * text_rows + [0] * image_rows), + position_ids=positions, + image_indices=torch.arange(text_rows, seq), + text_indices=torch.arange(text_rows), + pdd_step_index=1, + return_dict=False, + ) + + +@torch.no_grad() +def test_tabulated_forward_matches_original(): + model = _model() + calls = [_inputs([0.25]), _inputs([0.5, 0.999])] + reference = [model(**kw)[0] for kw in calls] + keys = family.tabulate_modulation(model, TIMESTEPS) + assert keys == torch.tensor(sorted(TIMESTEPS), dtype=torch.float32).tolist() + assert not hasattr(model, "time_embedder") + assert all(isinstance(b.adaln_proj, family.ModulationTable) for b in model.transformer_blocks) + for kw, ref in zip(calls, reference, strict=True): + torch.testing.assert_close(model(**kw)[0], ref, rtol=1e-5, atol=1e-5) + + +@torch.no_grad() +def test_tables_load_into_empty_skeleton(): + model = _model() + family.tabulate_modulation(model, TIMESTEPS) + state = model.state_dict() + fresh = _model() + family.add_modulation_tables(fresh, len(TIMESTEPS)) + fresh.load_state_dict(state, strict=True) + family.install(fresh) + kw = _inputs([0.5]) + torch.testing.assert_close(fresh(**kw)[0], model(**kw)[0]) + + +@torch.no_grad() +def test_unknown_timestep_raises(): + model = _model() + family.tabulate_modulation(model, TIMESTEPS) + with pytest.raises(ValueError, match="modulation table"): + model(**_inputs([0.3])) diff --git a/uv.lock b/uv.lock index ea791e1..1afd087 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.17.1" +version = "0.18.0" source = { editable = "." } dependencies = [ { name = "huggingface-hub" },