fix(diffusion): load serialized H3 INT8 in ComfyUI integrated mode

python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/spec.py

@@ -22,6 +22,12 @@
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
+from sglang.multimodal_gen.runtime.loader.minimax_h3_weights import (
+    comfy_quant_key_filter,
+    inspect_minimax_h3_safetensors,
+    resolve_minimax_h3_checkpoint_quantization,
+)
+from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
 from sglang.multimodal_gen.runtime.loader.weight_utils import (
     safetensors_weights_iterator,
 )
@@ -241,15 +247,45 @@ def load_comfyui_transformer(
                 safetensors_weights_iterator([model_path]), dit_config
             )
 
+        quant_config = None
+        checkpoint_key_filter = None
+        weight_load_plan = None
         # GGUF already sets AdaLN curve from tensor meta. Pruned BF16
         # safetensors keep the same adaln_t_table; without this the DiT is
         # built as the unpruned MLP and load fails on that extra parameter.
         if spec.dit_cls_name == "MiniMaxH3DiTModel":
-            from sglang.multimodal_gen.runtime.loader.minimax_h3_weights import (
-                inspect_minimax_h3_safetensors,
+            adaln_curve_shape, layer_markers = inspect_minimax_h3_safetensors(
+                [model_path]
             )
-
-            adaln_curve_shape, _ = inspect_minimax_h3_safetensors([model_path])
+            if layer_markers:
+                if any(
+                    marker.get("format") != "int8_tensorwise"
+                    for marker in layer_markers.values()
+                ):
+                    raise ValueError(
+                        "ComfyUI H3 integrated mode currently supports serialized INT8 ConvRot quantization"
+                    )
+                if (
+                    server_args.quantization is not None
+                    or server_args.nunchaku_config is not None
+                ):
+                    raise ValueError(
+                        "Checkpoint quantization is encoded in per-layer metadata; do not also set quantization or Nunchaku"
+                    )
+                if server_args.should_use_fsdp_for_component("transformer"):
+                    raise ValueError(
+                        "Comfy quantized checkpoints do not support FSDP inference; use TP and/or sequence parallelism instead"
+                    )
+                quant_config = resolve_minimax_h3_checkpoint_quantization(layer_markers)
+                checkpoint_key_filter = comfy_quant_key_filter
+                checkpoint_device = (
+                    torch.device("cpu")
+                    if server_args.should_start_component_on_cpu("transformer")
+                    else get_local_torch_device()
+                )
+                weight_load_plan = WeightLoadPlan(
+                    checkpoint_load_device=checkpoint_device
+                )
             if adaln_curve_shape is not None:
                 (
                     dit_config.arch_config.adaln_curve_grid,
@@ -269,12 +305,15 @@ def load_comfyui_transformer(
 
         # Weight loading reads param_names_mapping off the model, which inherits it
         # from the class, so the ComfyUI names have to be visible for the whole load.
+        init_params = {"config": dit_config, "hf_config": {}}
+        if quant_config is not None:
+            init_params["quant_config"] = quant_config
         original_mapping = model_cls.param_names_mapping
         model_cls.param_names_mapping = mapping
         try:
             model = maybe_load_fsdp_model(
                 model_cls=model_cls,
-                init_params={"config": dit_config, "hf_config": {}},
+                init_params=init_params,
                 weight_dir_list=[model_path],
                 device=get_local_torch_device(),
                 hsdp_replicate_dim=server_args.hsdp_replicate_dim,
@@ -289,6 +328,8 @@ def load_comfyui_transformer(
                 output_dtype=None,
                 strict=spec.strict,
                 weights_iterator=weights_iterator,
+                checkpoint_key_filter=checkpoint_key_filter,
+                weight_load_plan=weight_load_plan,
             )
         finally:
             model_cls.param_names_mapping = original_mapping

python/sglang/multimodal_gen/test/unit/test_comfyui_h3.py

@@ -1098,3 +1098,116 @@ def test_concat_qkv_attention_skips_grouped_weight_loader() -> None:
     grouped_weight = grouped_attn.qkv_proj.weight
     assert grouped_weight.checkpoint_mapping_unsafe is True
     assert grouped_weight.rank_local_weight_transform is not None
+
+
+@pytest.mark.parametrize("backend,out_features", [("comfy_kitchen", 12), ("jit", 16)])
+def test_integrated_h3_int8_loader_preserves_quantized_weights(
+    tmp_path, monkeypatch, backend, out_features
+):
+    _ensure_single_process_parallel_runtime()
+    import json
+    from types import SimpleNamespace
+
+    from safetensors.torch import save_file
+
+    from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTConfig
+    from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
+    from sglang.multimodal_gen.runtime.loader.comfyui_checkpoints import spec
+    from sglang.multimodal_gen.runtime.loader.comfyui_checkpoints.minimax_h3 import (
+        _build_dit_config,
+    )
+
+    class SmallH3(torch.nn.Module):
+        param_names_mapping = {}
+        _fsdp_forward_methods = ()
+
+        def __init__(self, config, hf_config, quant_config=None):
+            super().__init__()
+            self.proj = ReplicatedLinear(
+                256,
+                out_features,
+                bias=False,
+                params_dtype=torch.bfloat16,
+                quant_config=quant_config,
+                prefix="proj",
+            )
+
+        def post_load_weights(self):
+            pass
+
+    marker = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}
+    weight = torch.arange(out_features * 256).reshape(out_features, 256).to(torch.int8)
+    scales = torch.arange(1, out_features + 1, dtype=torch.float32).reshape(
+        out_features, 1
+    )
+    path = tmp_path / "h3.safetensors"
+    save_file(
+        {
+            "proj.weight": weight,
+            "proj.weight_scale": scales,
+            "proj.comfy_quant": torch.tensor(
+                list(json.dumps(marker).encode()), dtype=torch.uint8
+            ),
+        },
+        path,
+    )
+    monkeypatch.setattr(
+        spec.ModelRegistry, "resolve_model_cls", lambda _: (SmallH3, None)
+    )
+    monkeypatch.setattr(spec, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(
+        spec,
+        "get_comfyui_checkpoint_spec",
+        lambda _: spec.ComfyUICheckpointSpec(
+            dit_cls_name="MiniMaxH3DiTModel", build_dit_config=_build_dit_config
+        ),
+    )
+    monkeypatch.setattr(
+        "sglang.multimodal_gen.runtime.layers.quantization.convrot_int8_comfy_kitchen._load_comfy_kitchen",
+        lambda: None,
+    )
+    monkeypatch.setattr(
+        "sglang.multimodal_gen.runtime.layers.quantization.convrot_int8_jit._load_jit_kernel",
+        lambda: None,
+    )
+    monkeypatch.setenv("SGLANG_DIFFUSION_CONVROT_INT8_BACKEND", backend)
+    arguments = SimpleNamespace(
+        pipeline_config=SimpleNamespace(
+            dit_config=MiniMaxH3DiTConfig(), dit_precision="bf16"
+        ),
+        model_paths={},
+        transformer_weights_path=None,
+        component_weights_paths={},
+        model_path=str(path),
+        dit_precision="bf16",
+        component_precisions={},
+        quantization=None,
+        nunchaku_config=None,
+        hsdp_replicate_dim=1,
+        hsdp_shard_dim=1,
+        pin_cpu_memory=False,
+        should_start_component_on_cpu=lambda _: True,
+        should_use_fsdp_for_component=lambda _: False,
+    )
+    pipeline = SimpleNamespace(
+        pipeline_name="MiniMaxH3Pipeline",
+        model_path=str(path),
+        get_module=lambda _: None,
+    )
+    loaded = spec.load_comfyui_transformer(pipeline, arguments)["transformer"]
+    assert loaded.proj.weight.dtype == torch.int8
+    assert torch.equal(loaded.proj.weight, weight)
+    assert torch.equal(loaded.proj.weight_scale, scales)
+    assert loaded.proj.quant_method.is_checkpoint_serialized
+    assert loaded.proj.quant_method.quant_config.backend == backend
+    assert (
+        arguments.pipeline_config.dit_config.arch_config.qkv_checkpoint_grouped is False
+    )
+    assert all(not parameter.requires_grad for parameter in loaded.parameters())
+    arguments.quantization = "kitchen_int8"
+    with pytest.raises(ValueError, match="per-layer metadata"):
+        spec.load_comfyui_transformer(pipeline, arguments)
+    arguments.quantization = None
+    arguments.should_use_fsdp_for_component = lambda _: True
+    with pytest.raises(ValueError, match="FSDP"):
+        spec.load_comfyui_transformer(pipeline, arguments)