From f851045c948692b44d5998bd337476a3f7709428 Mon Sep 17 00:00:00 2001 From: Tokha233 <61346912+Tokha233@users.noreply.github.com> Date: Fri, 2 Oct 2026 03:58:27 +0800 Subject: [PATCH] fix(diffusion): load serialized H3 INT8 in ComfyUI integrated mode --- .../loader/comfyui_checkpoints/spec.py | 51 ++++++++- .../test/unit/test_comfyui_h3.py | 102 ++++++++++++++++++ 2 files changed, 148 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/spec.py b/python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/spec.py index 34aae2426a83..473b94a2aab7 100644 --- a/python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/spec.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/unit/test_comfyui_h3.py b/python/sglang/multimodal_gen/test/unit/test_comfyui_h3.py index 884b300e5f3d..26b9efc914ca 100644 --- a/python/sglang/multimodal_gen/test/unit/test_comfyui_h3.py +++ b/python/sglang/multimodal_gen/test/unit/test_comfyui_h3.py @@ -1098,3 +1098,105 @@ 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 + + +def test_integrated_h3_int8_loader_preserves_quantized_weights(tmp_path, monkeypatch): + _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, + 12, + 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(12 * 256).reshape(12, 256).to(torch.int8) + scales = torch.arange(1, 13, dtype=torch.float32).reshape(12, 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.kitchen_int8._load_comfy_kitchen", + lambda: None, + ) + 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 ( + 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)