From 5ecd6d794d093c9f6b8f6691a92b1da11df78a72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=A4=E7=94=B7?= Date: Fri, 21 Aug 2026 22:10:50 +0800 Subject: [PATCH] [diffusion] fix: fix quantized qkv scales and missing-param policy for minimax-h3 (#35740) --- .../runtime/models/dits/minimax_h3.py | 50 +++++++++++++++- .../unit/test_minimax_h3_qkv_scale_reorder.py | 60 +++++++++++++++++++ 2 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_minimax_h3_qkv_scale_reorder.py diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 822a80ec8..b84a0655a 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -176,6 +176,39 @@ def _reorder_grouped_qkv_to_qkv( ) +def _install_qkv_row_reorder( + param: torch.Tensor, + reorder: Callable[[torch.Tensor], torch.Tensor], + qkv_rows: int, +) -> None: + """Reorder a per-output-row qkv parameter the same way its rows are reordered. + + Applied to quantization metadata rather than the weight itself. Anything whose + leading dim is not the checkpoint's qkv row count is passed through: per-tensor + scales are scalars, and a swizzled block-scale layout is not row-indexed. + """ + + def _maybe_reorder(loaded_weight: torch.Tensor) -> torch.Tensor: + if loaded_weight.dim() >= 2 and loaded_weight.shape[0] == qkv_rows: + return reorder(loaded_weight) + return loaded_weight + + base_loader = ( + param._weight_loader + if hasattr(param, "_weight_loader") + else param.weight_loader + ) + + def _weight_loader(p: torch.Tensor, loaded_weight: torch.Tensor) -> None: + base_loader(p, _maybe_reorder(loaded_weight)) + + if hasattr(param, "_weight_loader"): + param._weight_loader = _weight_loader + else: + param.weight_loader = _weight_loader + param.rank_local_weight_transform = _maybe_reorder + + def _copy_grouped_qkv_tp_shard( param: torch.Tensor, loaded_weight: torch.Tensor, @@ -662,6 +695,17 @@ class MiniMaxH3Attention(nn.Module): # rank-local FSDP must reorder grouped QKV before selecting each shard weight.rank_local_weight_transform = _reorder_checkpoint_weight + # A quantized checkpoint stores metadata indexed by output row next to the + # rows themselves (NVFP4 block scales, fp8 per-channel scales). Those rows + # are permuted above, so the per-row metadata has to be permuted the same + # way. Row count is the gate: a swizzled scale layout is not row-indexed, + # and per-tensor scales are scalars, so both are passed through untouched. + qkv_rows = 3 * arch.num_attention_heads * arch.attention_head_dim + for name, param in self.qkv_proj.named_parameters(recurse=False): + if name == "weight": + continue + _install_qkv_row_reorder(param, _reorder_checkpoint_weight, qkv_rows) + def _forward_mps_streamed_attention( self, x: torch.Tensor, @@ -1816,7 +1860,11 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): def _mark_missing_params_required(self) -> None: for _, param in self.named_parameters(): - param.missing_param_init = "error" + # A quant method's create_weights() declares its own policy for scales + # it can synthesize (weight-only NVFP4 has no input_scale and marks it + # "ones"); claiming only undeclared params keeps that intact. + if getattr(param, "missing_param_init", None) is None: + param.missing_param_init = "error" def post_load_weights(self) -> None: fp32_param_names = list(_MINIMAX_H3_FP32_PARAM_NAMES_IN_MODEL_ORDER) diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_qkv_scale_reorder.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_qkv_scale_reorder.py new file mode 100644 index 000000000..75285cfeb --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_qkv_scale_reorder.py @@ -0,0 +1,60 @@ +import unittest + +import torch + +from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( + _install_qkv_row_reorder, +) + +QKV_ROWS = 12 + + +def _install(reorder=lambda w: w.flip(0)): + param = torch.zeros(QKV_ROWS) + seen = [] + param.weight_loader = lambda p, loaded_weight: seen.append(loaded_weight) + _install_qkv_row_reorder(param, reorder, QKV_ROWS) + return param, seen + + +class TestMiniMaxH3QkvScaleReorder(unittest.TestCase): + """MiniMax-H3 permutes the fused qkv weight's output rows on load. + + Quantization metadata indexed by output row (NVFP4 block scales, fp8 + per-channel scales) has to move with those rows, or every scale lands on the + wrong row: the model loads and runs, and renders noise. Metadata that is not + row-indexed must be left alone, so the row count is the gate. + """ + + def test_row_indexed_metadata_is_reordered(self): + param, seen = _install() + scales = torch.arange(QKV_ROWS * 2, dtype=torch.float32).reshape(QKV_ROWS, 2) + param.weight_loader(param, scales) + self.assertTrue(torch.equal(seen[0], scales.flip(0))) + + def test_metadata_with_another_row_count_is_passed_through(self): + param, seen = _install() + swizzled = torch.arange(QKV_ROWS, dtype=torch.float32).reshape(QKV_ROWS // 3, 3) + param.weight_loader(param, swizzled) + self.assertTrue(torch.equal(seen[0], swizzled)) + + def test_per_tensor_scale_is_passed_through(self): + param, seen = _install() + scale = torch.tensor(0.5) + param.weight_loader(param, scale) + self.assertTrue(torch.equal(seen[0], scale)) + + def test_rank_local_transform_follows_the_same_gate(self): + # rank-local FSDP reorders before selecting its shard, so it needs the + # transform rather than the wrapped loader. + param, _ = _install() + scales = torch.arange(QKV_ROWS * 2, dtype=torch.float32).reshape(QKV_ROWS, 2) + self.assertTrue( + torch.equal(param.rank_local_weight_transform(scales), scales.flip(0)) + ) + scale = torch.tensor(0.5) + self.assertTrue(torch.equal(param.rank_local_weight_transform(scale), scale)) + + +if __name__ == "__main__": + unittest.main()