feat(minimax-h3): support AdaLN curve checkpoints

lightx2v/common/ops/mm/mm_weight.py

@@ -378,33 +378,18 @@ class MMWeightForceFp32(MMWeight):
     """
 
     def load(self, weight_dict):
-        if not self.create_cuda_buffer and not self.create_cpu_buffer and not self.lazy_load:
-            device_tensors, pin_tensors = create_default_tensors(self.base_attrs, weight_dict)
-            weight = device_tensors.get("weight")
-            bias = device_tensors.get("bias")
-            if weight is not None:
-                weight = weight.to(torch.float32)
-            if bias is not None:
-                bias = bias.to(torch.float32)
-            self.weight = weight
-            self.bias = bias
-            pin_weight = pin_tensors.get("weight")
-            pin_bias = pin_tensors.get("bias")
-            if pin_weight is not None:
-                pin_weight = pin_weight.to(torch.float32)
-            if pin_bias is not None:
-                pin_bias = pin_bias.to(torch.float32)
-            self.pin_weight = pin_weight
-            self.pin_bias = pin_bias
-        else:
-            # Fall back to the Default load path, then force fp32 on the
-            # tensors we expose (buffers for CUDA/CPU streams are left alone
-            # so the copy mechanics keep working; apply() will cast as needed).
-            super().load(weight_dict)
-            if getattr(self, "weight", None) is not None:
-                self.weight = self.weight.to(torch.float32)
-            if getattr(self, "bias", None) is not None:
-                self.bias = self.bias.to(torch.float32)
+        super().load(weight_dict)
+        for attr_name in (
+            "weight",
+            "bias",
+            "pin_weight",
+            "pin_bias",
+            "weight_cuda_buffer",
+            "bias_cuda_buffer",
+        ):
+            tensor = getattr(self, attr_name, None)
+            if tensor is not None:
+                setattr(self, attr_name, tensor.to(torch.float32))
 
 
 class MMWeightQuantTemplate(MMWeightTemplate):

lightx2v/models/networks/minimax_h3/infer/post_infer.py

@@ -7,12 +7,14 @@
 class MiniMaxH3PostInfer:
     def __init__(self, config):
         self.config = config
+        self.h3_adaln_curve = bool(config.get("h3_adaln_curve", False))
 
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
 
     def infer(self, weights, hidden_states, pre_infer_out):
-        shift, scale = weights.norm_out_linear.apply(F.silu(pre_infer_out.temb).to(GET_DTYPE())).chunk(2, dim=-1)
+        adaln_input = pre_infer_out.temb if self.h3_adaln_curve else F.silu(pre_infer_out.temb).to(GET_DTYPE())
+        shift, scale = weights.norm_out_linear.apply(adaln_input).chunk(2, dim=-1)
         indices = pre_infer_out.timestep_indices
         hidden_states = weights.norm_out.apply(hidden_states)
         hidden_states = hidden_states * (1.0 + scale.index_select(0, indices))

lightx2v/models/networks/minimax_h3/infer/pre_infer.py

@@ -22,9 +22,21 @@ def timestep_embedding(timesteps: torch.Tensor, embedding_dim: int = 256) -> tor
     return embedding
 
 
+def interpolate_adaln_curve(table: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
+    """Linearly sample a curve-form H3 AdaLN table at normalized timesteps."""
+    if table.ndim != 2 or table.shape[0] < 2:
+        raise ValueError(f"AdaLN curve table must be [grid>=2, basis], got {tuple(table.shape)}")
+    if timesteps.ndim != 1:
+        raise ValueError(f"timesteps must be one-dimensional, got {tuple(timesteps.shape)}")
+    pos = timesteps.float().clamp(0.0, 1.0) * (table.shape[0] - 1)
+    lower = pos.floor().long().clamp(max=table.shape[0] - 2)
+    return torch.lerp(table[lower], table[lower + 1], (pos - lower).unsqueeze(1))
+
+
 class MiniMaxH3PreInfer:
     def __init__(self, config):
         self.config = config
+        self.h3_adaln_curve = bool(config.get("h3_adaln_curve", False))
         global_num_heads = int(config.get("num_attention_heads", 56))
         if config.get("tensor_parallel", False):
             tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p")
@@ -107,8 +119,11 @@ def infer(self, weights, prompt_embeds):
         hidden_states.index_copy_(0, layout.audio_indices, audio_embeds)
         hidden_states.index_copy_(0, layout.video_indices, video_embeds)
 
-        temb = timestep_embedding(self.scheduler.unique_timesteps, self.freq_dim)
-        temb = weights.time_linear_2.apply(F.silu(weights.time_linear_1.apply(temb.float())))
+        if self.h3_adaln_curve:
+            temb = interpolate_adaln_curve(weights.adaln_t_table.tensor, self.scheduler.unique_timesteps)
+        else:
+            temb = timestep_embedding(self.scheduler.unique_timesteps, self.freq_dim)
+            temb = weights.time_linear_2.apply(F.silu(weights.time_linear_1.apply(temb.float())))
         timestep_indices = self.scheduler.timestep_indices
         adaln_indices = timestep_indices * 3 + layout.token_tags.clamp(min=0)
 

lightx2v/models/networks/minimax_h3/infer/transformer_infer.py

@@ -22,6 +22,7 @@ def __init__(self, config):
         self.num_heads = self.global_num_heads // self.tp_size
         self.head_dim = int(config.get("attention_head_dim", 128))
         self.infer_dtype = GET_DTYPE()
+        self.h3_adaln_curve = bool(config.get("h3_adaln_curve", False))
         if config.get("seq_parallel", False):
             self.seq_p_group = config["device_mesh"].get_group(mesh_dim="seq_p")
             parallel = config.get("parallel", {})
@@ -130,7 +131,8 @@ def infer_block(self, weights, hidden_states, pre_infer_out, modulation=None):
     def _compute_adaln_table(self, weights, pre_infer_out):
         # Activation is evaluated in fp32, then cast to the inference dtype
         # immediately before the (possibly quantized) AdaLN projection.
-        modulation = weights.adaln.apply(F.silu(pre_infer_out.temb).to(self.infer_dtype))
+        adaln_input = pre_infer_out.temb if self.h3_adaln_curve else F.silu(pre_infer_out.temb).to(self.infer_dtype)
+        modulation = weights.adaln.apply(adaln_input)
         modulation = self._gather_tp_last_dim(modulation)
         return modulation.view(-1, 6 * self.hidden_size)
 

lightx2v/models/networks/minimax_h3/model.py

@@ -47,12 +47,17 @@ class MiniMaxH3Model(BaseTransformerModel):
 
     def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0, lora_alpha=None):
         self.lora_alpha = lora_alpha
+        self.h3_adaln_curve = bool(config.get("h3_adaln_curve", False))
         if GET_DTYPE() != torch.bfloat16:
             raise ValueError(
                 "MiniMax-H3 requires DTYPE=BF16. The native loader preserves the released checkpoint's 626 BF16 tensors and 12 FP32 projection/time/head tensors without dtype conversion."
             )
         if config.get("cfg_parallel", False) or config.get("enable_cfg", False):
             raise ValueError("MiniMax-H3 is guidance-distilled and does not have a CFG/unconditional branch")
+        if self.h3_adaln_curve and config.get("dit_quantized", False):
+            raise NotImplementedError(
+                "MiniMax-H3 AdaLN curve checkpoints do not support quantized DiT weights yet: quantized adaln_proj tensors require weight_scale-aware loading, while the curve path requires FP32 AdaLN projections. Use dit_quantized=false."
+            )
         if config.get("dit_quantized", False):
             quant_scheme = config.get("dit_quant_scheme", "Default")
             if quant_scheme not in H3_CHANNEL_QUANT_SCHEMES:
@@ -63,6 +68,11 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0
             raise ValueError("MiniMax-H3 dit_quant_scheme requires a dit_quantized_ckpt")
         if config.get("cpu_offload", False) and config.get("offload_granularity", "model") not in {"model", "block"}:
             raise NotImplementedError("MiniMax-H3 supports model and block CPU offload")
+        if self.h3_adaln_curve:
+            curve_grid = int(config.get("adaln_curve_grid", 0))
+            curve_basis = int(config.get("time_embed_dim", 0))
+            if curve_grid < 2 or curve_basis < 1:
+                raise ValueError(f"MiniMax-H3 curve checkpoints require adaln_curve_grid>=2 and time_embed_dim>=1; got grid={curve_grid}, basis={curve_basis}")
         if config.get("attn_type") == "sol_attn":
             reorder = str(config.get("sol_attn_setting", {}).get("reorder", "none")).lower()
             if reorder != "none":
@@ -90,6 +100,8 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0
             "proj_out",
             "audio_proj_out",
         }
+        if self.h3_adaln_curve:
+            self.sensitive_layer.add("adaln_t_table")
         self._init_infer_class()
         self._init_weights()
         self._init_infer()

lightx2v/models/networks/minimax_h3/weights/post_weights.py

@@ -15,7 +15,7 @@ def __init__(self, config):
         )
         self.add_module(
             "norm_out_linear",
-            MM_WEIGHT_REGISTER["Default"]("norm_out.linear.weight", "norm_out.linear.bias"),
+            MM_WEIGHT_REGISTER["Default-ForceFp32" if config.get("h3_adaln_curve", False) else "Default"]("norm_out.linear.weight", "norm_out.linear.bias"),
         )
         self.add_module(
             "proj_out",

lightx2v/models/networks/minimax_h3/weights/pre_weights.py

@@ -1,7 +1,7 @@
 import torch.distributed as dist
 
 from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList
-from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER
+from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, TENSOR_REGISTER
 
 
 def _linear(name, bias=False, force_fp32=False, config=None, tp_split=None):
@@ -79,13 +79,16 @@ def __init__(self, index, config):
 class MiniMaxH3PreWeights(WeightModule):
     def __init__(self, config):
         super().__init__()
-        # The released checkpoint deliberately keeps the two media projections
-        # and timestep MLP in fp32.  The text projection/refiner stay bf16.
+        # Full checkpoints keep the timestep MLP in fp32. Curve-form
+        # checkpoints replace that MLP with an fp32 interpolation table.
         self.add_module("proj_in", _linear("proj_in", bias=True, force_fp32=True))
         self.add_module("audio_proj_in", _linear("audio_proj_in", bias=True, force_fp32=True))
         self.add_module("context_embedder", _linear("context_embedder", bias=True))
-        self.add_module("time_linear_1", _linear("time_embedder.linear_1", bias=True, force_fp32=True))
-        self.add_module("time_linear_2", _linear("time_embedder.linear_2", bias=True, force_fp32=True))
+        if config.get("h3_adaln_curve", False):
+            self.register_parameter("adaln_t_table", TENSOR_REGISTER["Default"]("adaln_t_table"))
+        else:
+            self.add_module("time_linear_1", _linear("time_embedder.linear_1", bias=True, force_fp32=True))
+            self.add_module("time_linear_2", _linear("time_embedder.linear_2", bias=True, force_fp32=True))
         self.add_module(
             "refiner_blocks",
             WeightModuleList([MiniMaxH3TokenRefinerBlockWeights(i, config) for i in range(int(config.get("num_refiner_layers", 2)))]),

lightx2v/models/networks/minimax_h3/weights/transformer_weights.py

@@ -6,15 +6,16 @@
 from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER
 
 
-def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None):
+def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None, force_fp32=False):
+    mm_type = "Default-ForceFp32" if force_fp32 else config.get("dit_quant_scheme", "Default")
     lora_prefix = "transformer_blocks"
     if config.get("tensor_parallel", False) and tp_split is not None:
         tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p")
         tp_mm_type = config.get("tp_mm_type", "TensorParallel")
         return MM_WEIGHT_REGISTER[tp_mm_type](
             weight_name=f"{name}.weight",
             bias_name=f"{name}.bias" if bias else None,
-            mm_type=config.get("dit_quant_scheme", "Default"),
+            mm_type=mm_type,
             tp_group=tp_group,
             tp_rank=dist.get_rank(tp_group),
             tp_size=dist.get_world_size(tp_group),
@@ -23,7 +24,7 @@ def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None):
             create_cuda_buffer=create_cuda_buffer,
             lora_prefix=lora_prefix,
         )
-    return MM_WEIGHT_REGISTER[config.get("dit_quant_scheme", "Default")](
+    return MM_WEIGHT_REGISTER[mm_type](
         f"{name}.weight",
         f"{name}.bias" if bias else None,
         create_cuda_buffer=create_cuda_buffer,
@@ -123,7 +124,17 @@ def __init__(self, index, config, create_cuda_buffer=False):
         self.add_module("ff", MiniMaxH3FeedForwardWeights(f"{prefix}.ff", config, create_cuda_buffer))
         # AdaLN is the largest per-block projection in H3.  Its output is
         # column-sharded here and gathered once per block before modulation.
-        self.add_module("adaln", _linear(config, f"{prefix}.adaln_proj.linear", bias=True, create_cuda_buffer=create_cuda_buffer, tp_split="col"))
+        self.add_module(
+            "adaln",
+            _linear(
+                config,
+                f"{prefix}.adaln_proj.linear",
+                bias=True,
+                create_cuda_buffer=create_cuda_buffer,
+                tp_split="col",
+                force_fp32=bool(config.get("h3_adaln_curve", False)),
+            ),
+        )
 
 
 class MiniMaxH3TransformerWeights(WeightModule):

test_cases/test_minimax_h3_adaln_curve.py

@@ -0,0 +1,141 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+import torch.nn.functional as F
+
+from lightx2v.common.ops.mm.mm_weight import MMWeightForceFp32
+from lightx2v.models.networks.minimax_h3.infer.pre_infer import interpolate_adaln_curve
+from lightx2v.models.networks.minimax_h3.infer.transformer_infer import MiniMaxH3TransformerInfer
+from lightx2v.models.networks.minimax_h3.model import MiniMaxH3Model
+
+
+class _RecordingProjection:
+    def __init__(self):
+        self.call_count = 0
+        self.last_input = None
+
+    def apply(self, value):
+        self.call_count += 1
+        self.last_input = value
+        return value
+
+
+def _make_transformer_infer(*, curve):
+    infer = MiniMaxH3TransformerInfer.__new__(MiniMaxH3TransformerInfer)
+    infer.h3_adaln_curve = curve
+    infer.infer_dtype = torch.bfloat16
+    infer.hidden_size = 1
+    infer.tp_size = 1
+    return infer
+
+
+def test_interpolate_adaln_curve_clamps_and_interpolates():
+    table = torch.tensor([[0.0, 10.0], [2.0, 12.0], [4.0, 14.0]], dtype=torch.float32)
+    timesteps = torch.tensor([-1.0, 0.0, 0.25, 0.5, 0.75, 1.0, 2.0], dtype=torch.float32)
+    expected = torch.tensor(
+        [
+            [0.0, 10.0],
+            [0.0, 10.0],
+            [1.0, 11.0],
+            [2.0, 12.0],
+            [3.0, 13.0],
+            [4.0, 14.0],
+            [4.0, 14.0],
+        ],
+        dtype=torch.float32,
+    )
+    torch.testing.assert_close(interpolate_adaln_curve(table, timesteps), expected, rtol=0.0, atol=0.0)
+
+
+def test_interpolate_adaln_curve_rejects_invalid_inputs():
+    table = torch.zeros(1, 8)
+    timesteps = torch.zeros(1)
+    try:
+        interpolate_adaln_curve(table, timesteps)
+    except ValueError as error:
+        assert "grid>=2" in str(error)
+    else:
+        raise AssertionError("invalid curve table was accepted")
+
+
+def test_curve_adaln_integrates_with_request_cache():
+    infer = _make_transformer_infer(curve=True)
+    infer._adaln_cache = {}
+    infer._current_adaln_tables = None
+    infer._adaln_cache_hit = False
+    infer.scheduler = SimpleNamespace(unique_timesteps_cpu=torch.tensor([0.0, 0.5], dtype=torch.float32))
+
+    projection = _RecordingProjection()
+    weights = SimpleNamespace(adaln=projection)
+    pre_infer_out = SimpleNamespace(temb=torch.arange(12, dtype=torch.float32).reshape(2, 6))
+
+    infer._prepare_adaln_cache()
+    first = infer._get_or_build_adaln(0, weights, pre_infer_out)
+    assert projection.last_input is pre_infer_out.temb
+    assert first.dtype == torch.float32
+
+    infer._prepare_adaln_cache()
+    second = infer._get_or_build_adaln(0, weights, pre_infer_out)
+    assert infer._adaln_cache_hit
+    assert second is first
+    assert projection.call_count == 1
+
+
+def test_full_checkpoint_adaln_preserves_activation_and_dtype():
+    infer = _make_transformer_infer(curve=False)
+    projection = _RecordingProjection()
+    weights = SimpleNamespace(adaln=projection)
+    temb = torch.arange(12, dtype=torch.float32).reshape(2, 6)
+
+    result = infer._compute_adaln_table(weights, SimpleNamespace(temb=temb))
+    expected = F.silu(temb).to(torch.bfloat16)
+
+    torch.testing.assert_close(projection.last_input, expected, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(result, expected, rtol=0.0, atol=0.0)
+
+
+def test_force_fp32_offload_buffer_preserves_curve_dtype():
+    weight_name = "transformer_blocks.0.adaln_proj.linear.weight"
+    bias_name = "transformer_blocks.0.adaln_proj.linear.bias"
+    checkpoint_weight = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=torch.bfloat16)
+    checkpoint_bias = torch.tensor([0.5, -0.5, 1.0], dtype=torch.bfloat16)
+
+    def _fake_create_cuda_buffers(*args, **kwargs):
+        return {
+            "weight": checkpoint_weight.t().clone(),
+            "bias": checkpoint_bias.clone(),
+        }
+
+    with patch("lightx2v.common.ops.mm.mm_weight.create_cuda_buffers", _fake_create_cuda_buffers):
+        projection = MMWeightForceFp32(weight_name, bias_name, create_cuda_buffer=True)
+        projection.load({})
+
+    assert projection.weight_cuda_buffer.dtype == torch.float32
+    assert projection.bias_cuda_buffer.dtype == torch.float32
+
+    projection.load_state_dict(
+        {
+            weight_name: checkpoint_weight.t().float(),
+            bias_name: checkpoint_bias.float(),
+        },
+        block_index=0,
+    )
+    curve_input = torch.tensor([[2.0, -1.0]], dtype=torch.float32)
+    expected = torch.addmm(checkpoint_bias.float(), curve_input, checkpoint_weight.t().float())
+    torch.testing.assert_close(projection.apply(curve_input), expected, rtol=0.0, atol=0.0)
+
+
+def test_curve_checkpoint_rejects_quantized_dit_before_loading():
+    config = {
+        "h3_adaln_curve": True,
+        "dit_quantized": True,
+    }
+    with patch("lightx2v.models.networks.minimax_h3.model.GET_DTYPE", return_value=torch.bfloat16):
+        try:
+            MiniMaxH3Model("unused", config, torch.device("cpu"))
+        except NotImplementedError as error:
+            assert "weight_scale-aware" in str(error)
+            assert "dit_quantized=false" in str(error)
+        else:
+            raise AssertionError("quantized AdaLN curve checkpoint was accepted")