From 1591dcd91a9e1a5efdfa480fee898a7ba7548efd Mon Sep 17 00:00:00 2001 From: Yiqi Yang <134501931+decajoin@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:33:58 +0800 Subject: [PATCH] [diffusion] fix: fix loading a block-FP8 quantized MiniMax-H3 DiT (#35703) --- .../runtime/models/dits/minimax_h3.py | 60 ++++++-- .../test/unit/test_minimax_h3_dit_contract.py | 136 ++++++++++++++++++ 2 files changed, 188 insertions(+), 8 deletions(-) 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 03ce2b2cc..b6459693d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -71,6 +71,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im is_layerwise_offloaded_module, ) from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT +from sglang.multimodal_gen.runtime.models.parameter import BlockQuantScaleParameter from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, @@ -260,6 +261,30 @@ def _install_qkv_row_reorder( param.rank_local_weight_transform = _maybe_reorder +def _qkv_scale_block_rows(qkv_proj: nn.Module, head_dim: int) -> int: + """Weight rows covered by one row of the qkv projection's scale. + + Per-channel and NVFP4 scales hold one row per weight row and report 1. A + block-FP8 scale holds one row per weight_block_size[0] weight rows, so the + qkv row permutation has to count its rows in blocks instead. Only whole + scale rows can move, so a block spanning two heads' q/k/v rows cannot be + repaired by a permutation and is rejected rather than silently mis-scaled. + """ + quant_config = getattr( + getattr(qkv_proj, "quant_method", None), "quant_config", None + ) + block_size = getattr(quant_config, "weight_block_size", None) + if not block_size: + return 1 + block_rows = block_size[0] + if head_dim % block_rows: + raise ValueError( + "block-quantized qkv needs a block size that divides the head dim: " + f"head_dim={head_dim}, weight_block_size={block_size}." + ) + return block_rows + + def _copy_grouped_qkv_tp_shard( param: torch.Tensor, loaded_weight: torch.Tensor, @@ -755,13 +780,20 @@ class MiniMaxH3Attention(nn.Module): weight.checkpoint_mapping_unsafe = True base_loader = weight.weight_loader - def _reorder_checkpoint_weight(loaded_weight: torch.Tensor) -> torch.Tensor: - return _reorder_grouped_qkv_to_qkv( - loaded_weight, - num_query_groups=arch.num_attention_heads, - heads_per_group=1, - head_dim=arch.attention_head_dim, - ) + def _make_row_reorder( + head_dim: int, + ) -> Callable[[torch.Tensor], torch.Tensor]: + def _reorder(loaded_weight: torch.Tensor) -> torch.Tensor: + return _reorder_grouped_qkv_to_qkv( + loaded_weight, + num_query_groups=arch.num_attention_heads, + heads_per_group=1, + head_dim=head_dim, + ) + + return _reorder + + _reorder_checkpoint_weight = _make_row_reorder(arch.attention_head_dim) def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: # The grouped checkpoint layout is @@ -791,11 +823,22 @@ class MiniMaxH3Attention(nn.Module): # 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. + # A block-FP8 scale is row-indexed too, but in blocks rather than rows: + # it carries one row per block of weight rows, so both its permutation + # and the row count gating it are scaled down by the block height. qkv_rows = 3 * arch.num_attention_heads * arch.attention_head_dim + block_rows = _qkv_scale_block_rows(self.qkv_proj, 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) + rows_per_scale_row = ( + block_rows if isinstance(param, BlockQuantScaleParameter) else 1 + ) + _install_qkv_row_reorder( + param, + _make_row_reorder(arch.attention_head_dim // rows_per_scale_row), + qkv_rows // rows_per_scale_row, + ) def _forward_mps_streamed_attention( self, @@ -2657,5 +2700,6 @@ __all__ = [ "MINIMAX_H3_FP32_BUFFER_NAMES", "MINIMAX_H3_FP32_PARAM_NAMES", "MiniMaxH3DiTModel", + "_qkv_scale_block_rows", "_reorder_grouped_qkv_to_qkv", ] diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py index b4066892a..613d0b518 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py @@ -28,6 +28,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.fp8 import ( Fp8LinearMethod, ) from sglang.multimodal_gen.runtime.layers.usp import _usp_input_all_to_all_packed_qkv +from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( + _needs_device_weight_postprocess, +) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( MINIMAX_H3_FP32_BUFFER_NAMES, @@ -37,6 +40,7 @@ from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( _copy_grouped_qkv_tp_shard, _diffusers_h3_checkpoint, _modulate_gate, + _qkv_scale_block_rows, _reorder_grouped_qkv_to_qkv, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum @@ -491,6 +495,138 @@ def test_online_fp8_keeps_fp32_boundaries_and_ignored_layers_unquantized(): assert isinstance(layer.quant_method, UnquantizedLinearMethod) +def _block_fp8_quant_config(block: int = 128, **kwargs) -> Fp8Config: + return Fp8Config( + is_checkpoint_fp8_serialized=True, + activation_scheme="dynamic", + weight_block_size=[block, block], + **kwargs, + ) + + +def _meta_h3(quant_config) -> MiniMaxH3DiTModel: + _ensure_single_process_parallel_runtime() + with torch.device("meta"): + return MiniMaxH3DiTModel( + config=MiniMaxH3DiTConfig(), hf_config={}, quant_config=quant_config + ) + + +def test_offline_block_fp8_checkpoint_layout_and_cpu_load(): + quant_config = _block_fp8_quant_config( + ignored_layers=[ + "video_patch_proj", + "audio_patch_proj", + "time_embedder.proj_in", + "time_embedder.proj_out", + "final_layer.video_out", + "final_layer.audio_out", + ], + ) + model = _meta_h3(quant_config) + + params = dict(model.named_parameters()) + fp8_weights = {name for name, p in params.items() if p.dtype == torch.float8_e4m3fn} + scales = {name for name in params if name.endswith("weight_scale_inv")} + assert len(fp8_weights) == 260, len(fp8_weights) + assert {f"{n[: -len('weight')]}weight_scale_inv" for n in fp8_weights} == scales + + for scale_name in scales: + weight = params[f"{scale_name[: -len('weight_scale_inv')]}weight"] + n, k = weight.shape + assert params[scale_name].dtype == torch.float32, scale_name + assert tuple(params[scale_name].shape) == ( + -(-n // 128), + -(-k // 128), + ), scale_name + + for layer in ( + model.video_patch_proj, + model.audio_patch_proj, + model.time_embedder.proj_in, + model.time_embedder.proj_out, + model.final_layer.video_out, + model.final_layer.audio_out, + ): + assert isinstance(layer.quant_method, UnquantizedLinearMethod) + assert layer.weight.dtype != torch.float8_e4m3fn + + # False keeps the DiT off the GPU during load, which is the point of + # loading a pre-quantized checkpoint. + assert _needs_device_weight_postprocess(quant_config) is False + assert _needs_device_weight_postprocess(Fp8Config()) is True + + +def test_block_fp8_qkv_scale_follows_its_weight_through_the_grouped_reorder(): + """A block scale is row-indexed in blocks, so it moves with the qkv rows. + + The checkpoint is quantized in the grouped per-head [q, k, v] layout, and + the DiT permutes those rows into [q_all, k_all, v_all] as it loads. Leaving + the scale behind leaves every 128x128 tile scaling rows it no longer covers: + the model loads without error and renders noise. + """ + heads, head_dim, block = 8, 128, 128 + rows, cols = heads * 3 * head_dim, 256 + torch.manual_seed(0) + weight = torch.randn(rows, cols) + + reordered = _reorder_grouped_qkv_to_qkv( + weight, num_query_groups=heads, heads_per_group=1, head_dim=head_dim + ) + # A quantizer sees the grouped layout, so its scales are computed there. + tiles = weight.view(rows // block, block, cols // block, block) + scale = tiles.abs().amax(dim=(1, 3)) / 448.0 + reordered_scale = _reorder_grouped_qkv_to_qkv( + scale, num_query_groups=heads, heads_per_group=1, head_dim=head_dim // block + ) + + def tile_max(w: torch.Tensor) -> torch.Tensor: + return w.view(rows // block, block, cols // block, block).abs().amax(dim=(1, 3)) + + assert torch.equal(tile_max(reordered) / 448.0, reordered_scale) + assert not torch.equal(tile_max(reordered) / 448.0, scale) + + +def test_qkv_block_scale_param_is_reordered_in_blocks(): + model = _meta_h3(_block_fp8_quant_config()) + arch = model.arch + qkv = model.blocks[0].attn.qkv_proj + block = 128 + scale_rows = 3 * arch.num_attention_heads * arch.attention_head_dim // block + assert qkv.weight_scale_inv.shape[0] == scale_rows + + # The installed transform is what rank-local FSDP applies and what the + # wrapped weight_loader runs before handing off, so it is the contract. + loaded = torch.arange(scale_rows * 4, dtype=torch.float32).reshape(scale_rows, 4) + expected = _reorder_grouped_qkv_to_qkv( + loaded, + num_query_groups=arch.num_attention_heads, + heads_per_group=1, + head_dim=arch.attention_head_dim // block, + ) + got = qkv.weight_scale_inv.rank_local_weight_transform(loaded) + assert torch.equal(got, expected) + assert not torch.equal(got, loaded) + + # The weight itself keeps the per-row permutation. + unblocked = torch.zeros(scale_rows * block, 4) + assert qkv.weight.rank_local_weight_transform(unblocked).shape == unblocked.shape + + +def test_unquantized_and_per_tensor_qkv_keep_their_loaders(): + for quant_config in (None, Fp8Config(is_checkpoint_fp8_serialized=True)): + qkv = _meta_h3(quant_config).blocks[0].attn.qkv_proj + assert not hasattr(qkv, "weight_scale_inv") + assert _qkv_scale_block_rows(qkv, 128) == 1 + + +def test_block_that_straddles_heads_is_rejected(): + # head_dim is 128, so a 256-row block covers two heads' rows at once and no + # row permutation can repair it. + with pytest.raises(ValueError, match="divides the head dim"): + _meta_h3(_block_fp8_quant_config(block=256)) + + def test_sdpa_varlen_fallback_matches_naive_packed_reference(): torch.manual_seed(0) heads, dim = 2, 8