From 2dc2fee4ea91b8dd83d29724b93572066f2971e5 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Mon, 17 Aug 2026 15:42:07 -0700 Subject: [PATCH 1/2] feat(minimax-h3): support AdaLN curve checkpoints --- .../networks/minimax_h3/infer/post_infer.py | 4 +- .../networks/minimax_h3/infer/pre_infer.py | 19 +++- .../minimax_h3/infer/transformer_infer.py | 4 +- lightx2v/models/networks/minimax_h3/model.py | 8 ++ .../minimax_h3/weights/post_weights.py | 2 +- .../minimax_h3/weights/pre_weights.py | 13 ++- .../minimax_h3/weights/transformer_weights.py | 19 +++- test_cases/test_minimax_h3_adaln_curve.py | 92 +++++++++++++++++++ 8 files changed, 147 insertions(+), 14 deletions(-) create mode 100644 test_cases/test_minimax_h3_adaln_curve.py diff --git a/lightx2v/models/networks/minimax_h3/infer/post_infer.py b/lightx2v/models/networks/minimax_h3/infer/post_infer.py index 5b88dc639..03e8043dd 100644 --- a/lightx2v/models/networks/minimax_h3/infer/post_infer.py +++ b/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)) diff --git a/lightx2v/models/networks/minimax_h3/infer/pre_infer.py b/lightx2v/models/networks/minimax_h3/infer/pre_infer.py index 2ccd86806..94ad9f373 100644 --- a/lightx2v/models/networks/minimax_h3/infer/pre_infer.py +++ b/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) diff --git a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py index f7366425f..e2f17db71 100644 --- a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py +++ b/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) diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 959d71c0e..b4fb213c5 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -47,6 +47,7 @@ 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." @@ -63,6 +64,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 +96,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() diff --git a/lightx2v/models/networks/minimax_h3/weights/post_weights.py b/lightx2v/models/networks/minimax_h3/weights/post_weights.py index 43c7b0f90..f460b9acd 100644 --- a/lightx2v/models/networks/minimax_h3/weights/post_weights.py +++ b/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", diff --git a/lightx2v/models/networks/minimax_h3/weights/pre_weights.py b/lightx2v/models/networks/minimax_h3/weights/pre_weights.py index 974dab411..80d198d12 100644 --- a/lightx2v/models/networks/minimax_h3/weights/pre_weights.py +++ b/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)))]), diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index 9634eb595..c359edcae 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -6,7 +6,8 @@ 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") @@ -14,7 +15,7 @@ def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None): 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): diff --git a/test_cases/test_minimax_h3_adaln_curve.py b/test_cases/test_minimax_h3_adaln_curve.py new file mode 100644 index 000000000..1b4f32811 --- /dev/null +++ b/test_cases/test_minimax_h3_adaln_curve.py @@ -0,0 +1,92 @@ +from types import SimpleNamespace + +import torch +import torch.nn.functional as F + +from lightx2v.models.networks.minimax_h3.infer.pre_infer import interpolate_adaln_curve +from lightx2v.models.networks.minimax_h3.infer.transformer_infer import MiniMaxH3TransformerInfer + + +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) From 105b8313f519967273071201c0694b303f6165b6 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Sat, 22 Aug 2026 22:11:03 -0700 Subject: [PATCH 2/2] fix(minimax-h3): harden AdaLN curve precision paths --- lightx2v/common/ops/mm/mm_weight.py | 39 +++++----------- lightx2v/models/networks/minimax_h3/model.py | 4 ++ test_cases/test_minimax_h3_adaln_curve.py | 49 ++++++++++++++++++++ 3 files changed, 65 insertions(+), 27 deletions(-) diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py index 7cb8f196e..80bbdee8a 100755 --- a/lightx2v/common/ops/mm/mm_weight.py +++ b/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): diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index b4fb213c5..bf776a550 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -54,6 +54,10 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0 ) 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: diff --git a/test_cases/test_minimax_h3_adaln_curve.py b/test_cases/test_minimax_h3_adaln_curve.py index 1b4f32811..7ebf66f53 100644 --- a/test_cases/test_minimax_h3_adaln_curve.py +++ b/test_cases/test_minimax_h3_adaln_curve.py @@ -1,10 +1,13 @@ 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: @@ -90,3 +93,49 @@ def test_full_checkpoint_adaln_preserves_activation_and_dtype(): 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")