[diffusion] fix: fix loading multiple ckpts with different precision for a same module (#22360)

python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py

@@ -8,6 +8,7 @@ are handled here behind a small helper/adapter layer.
 
 import json
 import os
+import re
 from dataclasses import dataclass, field
 from functools import partial
 from typing import Callable, Optional
@@ -36,6 +37,10 @@ logger = init_logger(__name__)
 
 PostLoadHook = Callable[[nn.Module], None]
 
+_PRECISION_VARIANT_SUFFIX_RE = re.compile(
+    r"^(?P<stem>.+?)(?P<precision>\.(?:fp16|bf16|fp32))(?P<shard>-\d+-of-\d+)?(?P<ext>\.safetensors)$"
+)
+
 
 @dataclass
 class TransformerQuantLoadSpec:
@@ -172,6 +177,8 @@ def resolve_transformer_safetensors_to_load(
     else:
         safetensors_list = _list_safetensors_files(component_model_path)
 
+    safetensors_list = _filter_duplicate_precision_variant_safetensors(safetensors_list)
+
     if not safetensors_list:
         raise ValueError(
             f"no safetensors files found in {quantized_path or component_model_path}"
@@ -180,6 +187,48 @@ def resolve_transformer_safetensors_to_load(
     return safetensors_list
 
 
+def _filter_duplicate_precision_variant_safetensors(
+    safetensors_list: list[str],
+) -> list[str]:
+    """Drop precision-specific duplicates when a canonical file is present.
+
+    Diffusers checkpoints sometimes ship both `foo.safetensors` and
+    `foo.fp16.safetensors` (and their sharded variants) in the same directory.
+    Loading both is unsafe because duplicate parameter names race and whichever
+    tensor arrives last wins, leading to non-deterministic behavior
+
+    If a canonical unsuffixed (non bf16|fp32) file exists, prefer it and drop the precision
+    variant from the same family. Precision-only families are left untouched.
+    """
+    canonical_paths = set(safetensors_list)
+    filtered: list[str] = []
+    removed: list[str] = []
+
+    for path in safetensors_list:
+        match = _PRECISION_VARIANT_SUFFIX_RE.match(path)
+        if match is None:
+            filtered.append(path)
+            continue
+
+        canonical_path = (
+            f"{match.group('stem')}{match.group('shard') or ''}{match.group('ext')}"
+        )
+        if canonical_path in canonical_paths:
+            removed.append(path)
+            continue
+
+        filtered.append(path)
+
+    if removed:
+        logger.info(
+            "Filtered %d duplicate transformer precision variant file(s): %s",
+            len(removed),
+            removed,
+        )
+
+    return filtered
+
+
 def resolve_transformer_quant_load_spec(
     *,
     hf_config: dict,

python/sglang/multimodal_gen/runtime/loader/weight_utils.py

@@ -8,6 +8,7 @@ import hashlib
 import json
 import os
 import tempfile
+from collections import defaultdict
 from collections.abc import Generator, Iterable
 from pathlib import Path
 
@@ -133,6 +134,51 @@ def _validate_safetensors_file(file_path: str) -> bool:
         return False
 
 
+def _raise_if_duplicate_safetensors_keys(hf_weights_files: list[str]) -> None:
+    """Fail fast when multiple safetensors files define the same tensor name. Make sure runtime behavior is deterministic
+
+    Duplicate keys across files are almost always a packaging error for inference:
+    for example shipping both full and fp16 variants, or mixing consolidated and
+    sharded checkpoints. Continuing would make the final loaded value depend on
+    file iteration or streamer delivery order.
+    """
+    if len(hf_weights_files) <= 1:
+        return
+
+    key_to_file: dict[str, str] = {}
+    duplicate_files_by_key: dict[str, set[str]] = defaultdict(set)
+
+    for st_file in hf_weights_files:
+        with safe_open(st_file, framework="pt", device="cpu") as f:
+            for name in f.keys():  # noqa: SIM118
+                previous_file = key_to_file.get(name)
+                if previous_file is None:
+                    key_to_file[name] = st_file
+                    continue
+                if previous_file == st_file:
+                    continue
+                duplicate_files_by_key[name].update((previous_file, st_file))
+
+    if not duplicate_files_by_key:
+        return
+
+    examples = []
+    for key in sorted(duplicate_files_by_key)[:8]:
+        files = ", ".join(
+            sorted(os.path.basename(p) for p in duplicate_files_by_key[key])
+        )
+        examples.append(f"{key} [{files}]")
+
+    raise ValueError(
+        "Duplicate tensor names detected across safetensors files. Refusing to load "
+        "because final weights would depend on file or streamer ordering. "
+        f"Found {len(duplicate_files_by_key)} duplicate tensor name(s). "
+        f"Examples: {examples}. "
+        "This usually means multiple precision variants or consolidated+sharded "
+        "checkpoints were passed together."
+    )
+
+
 def safetensors_weights_iterator(
     hf_weights_files: list[str],
     to_cpu: bool = True,
@@ -184,6 +230,8 @@ def safetensors_weights_iterator(
             "Please retry - the files will be re-downloaded automatically."
         )
 
+    _raise_if_duplicate_safetensors_keys(hf_weights_files)
+
     if use_runai_model_streamer:
         with SafetensorsStreamer() as streamer:
             streamer.stream_files(hf_weights_files)

python/sglang/multimodal_gen/test/server/consistency_threshold.json

@@ -32,22 +32,28 @@
             "mean_abs_diff_threshold": 8.0
         },
         "qwen_image_t2i_cache_dit_enabled": {
-            "clip_threshold": 0.92,
+            "clip_threshold": 0.99,
             "ssim_threshold": 0.86,
             "psnr_threshold": 17.0,
             "mean_abs_diff_threshold": 10.0
         },
-        "layerwise_offload": {
+        "flux_2_image_t2i": {
-            "clip_threshold": 0.92,
+            "clip_threshold": 0.96,
-            "ssim_threshold": 0.91,
+            "ssim_threshold": 0.95,
-            "psnr_threshold": 21.5,
+            "psnr_threshold": 28.0,
             "mean_abs_diff_threshold": 8.0
         },
+        "layerwise_offload": {
+            "clip_threshold": 0.94,
+            "ssim_threshold": 0.90,
+            "psnr_threshold": 22.0,
+            "mean_abs_diff_threshold": 9.0
+        },
         "zimage_image_t2i_fp8": {
-            "clip_threshold": 0.92,
+            "clip_threshold": 0.94,
-            "ssim_threshold": 0.84,
+            "ssim_threshold": 0.88,
-            "psnr_threshold": 19.0,
+            "psnr_threshold": 21.0,
-            "mean_abs_diff_threshold": 10.0
+            "mean_abs_diff_threshold": 9.5
         },
         "qwen_image_edit_2509_ti2i": {
             "clip_threshold": 0.92,
@@ -55,6 +61,18 @@
             "psnr_threshold": 13.0,
             "mean_abs_diff_threshold": 26.0
         },
+        "qwen_image_edit_ti2i": {
+            "clip_threshold": 0.96,
+            "ssim_threshold": 0.95,
+            "psnr_threshold": 28.0,
+            "mean_abs_diff_threshold": 8.0
+        },
+        "qwen_image_edit_2511_ti2i": {
+            "clip_threshold": 0.96,
+            "ssim_threshold": 0.95,
+            "psnr_threshold": 28.0,
+            "mean_abs_diff_threshold": 8.0
+        },
         "qwen_image_layered_i2i": {
             "clip_threshold": 0.92,
             "ssim_threshold": 0.94,
@@ -68,16 +86,34 @@
             "mean_abs_diff_threshold": 10.0
         },
         "wan2_1_t2v_1_3b_lora_1gpu": {
-            "clip_threshold": 0.54,
+            "clip_threshold": 0.88,
-            "ssim_threshold": 0.40,
+            "ssim_threshold": 0.75,
-            "psnr_threshold": 13.2,
+            "psnr_threshold": 17.0,
-            "mean_abs_diff_threshold": 32.0
+            "mean_abs_diff_threshold": 14.0
+        },
+        "wan2_1_t2v_1.3b": {
+            "clip_threshold": 0.94,
+            "ssim_threshold": 0.94,
+            "psnr_threshold": 26.0,
+            "mean_abs_diff_threshold": 8.0
+        },
+        "wan2_1_t2v_1.3b_teacache_enabled": {
+            "clip_threshold": 0.93,
+            "ssim_threshold": 0.92,
+            "psnr_threshold": 24.0,
+            "mean_abs_diff_threshold": 9.0
+        },
+        "wan2_1_t2v_1.3b_upscaling_4x": {
+            "clip_threshold": 0.94,
+            "ssim_threshold": 0.94,
+            "psnr_threshold": 26.0,
+            "mean_abs_diff_threshold": 8.0
         },
         "wan2_2_ti2v_5b": {
-            "clip_threshold": 0.90,
+            "clip_threshold": 0.92,
-            "ssim_threshold": 0.81,
+            "ssim_threshold": 0.88,
-            "psnr_threshold": 20.4,
+            "psnr_threshold": 22.0,
-            "mean_abs_diff_threshold": 10.0
+            "mean_abs_diff_threshold": 9.0
         },
         "fastwan2_2_ti2v_5b": {
             "clip_threshold": 0.90,
@@ -115,6 +151,24 @@
             "psnr_threshold": 19,
             "mean_abs_diff_threshold": 8.0
         },
+        "flux_2_image_t2i_upscaling_4x": {
+            "clip_threshold": 0.96,
+            "ssim_threshold": 0.95,
+            "psnr_threshold": 28.0,
+            "mean_abs_diff_threshold": 8.0
+        },
+        "flux_2_t2i_customized_vae_path": {
+            "clip_threshold": 0.96,
+            "ssim_threshold": 0.95,
+            "psnr_threshold": 28.0,
+            "mean_abs_diff_threshold": 8.0
+        },
+        "flux_2_ti2i_multi_image_cache_dit": {
+            "clip_threshold": 0.94,
+            "ssim_threshold": 0.90,
+            "psnr_threshold": 22.0,
+            "mean_abs_diff_threshold": 9.0
+        },
         "zimage_image_t2i_2_gpus": {
             "clip_threshold": 0.92,
             "ssim_threshold": 0.90,

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

@@ -0,0 +1,181 @@
+"""
+This unittest is introduced in #22360, preventing duplicate transformer safetensors variants being loaded together
+"""
+
+import json
+import sys
+import tempfile
+import types
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+partial_json_parser = types.ModuleType("partial_json_parser")
+partial_json_parser_core = types.ModuleType("partial_json_parser.core")
+partial_json_parser_exceptions = types.ModuleType("partial_json_parser.core.exceptions")
+partial_json_parser_options = types.ModuleType("partial_json_parser.core.options")
+
+
+class _MalformedJSON(Exception):
+    pass
+
+
+class _Allow:
+    STR = 1
+    OBJ = 2
+    ARR = 4
+    ALL = STR | OBJ | ARR
+
+
+def _loads(input_str, _flags=None):
+    return json.loads(input_str)
+
+
+partial_json_parser_exceptions.MalformedJSON = _MalformedJSON
+partial_json_parser_options.Allow = _Allow
+partial_json_parser.loads = _loads
+sys.modules.setdefault("partial_json_parser", partial_json_parser)
+sys.modules.setdefault("partial_json_parser.core", partial_json_parser_core)
+sys.modules.setdefault(
+    "partial_json_parser.core.exceptions", partial_json_parser_exceptions
+)
+sys.modules.setdefault("partial_json_parser.core.options", partial_json_parser_options)
+
+from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
+    NunchakuConfig,
+)
+from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
+    _filter_duplicate_precision_variant_safetensors,
+    _Flux2Nvfp4FallbackAdapter,
+    resolve_transformer_quant_load_spec,
+    resolve_transformer_safetensors_to_load,
+)
+
+
+class _FakeFluxTransformer:
+    pass
+
+
+class _FakeQuantConfig:
+    @classmethod
+    def get_name(cls):
+        return "modelopt_fp4"
+
+
+class TestTransformerQuantHelpers(unittest.TestCase):
+    def _make_server_args(self, **overrides):
+        defaults = dict(
+            transformer_weights_path=None,
+            pipeline_config=SimpleNamespace(
+                dit_precision="bf16",
+                dit_config=SimpleNamespace(
+                    arch_config=SimpleNamespace(param_names_mapping={})
+                ),
+            ),
+            nunchaku_config=None,
+            tp_size=1,
+            dit_cpu_offload=False,
+            text_encoder_cpu_offload=False,
+        )
+        defaults.update(overrides)
+        return SimpleNamespace(**defaults)
+
+    def test_resolve_transformer_safetensors_to_load_uses_single_override_file(self):
+        with tempfile.NamedTemporaryFile(suffix=".safetensors") as f:
+            server_args = self._make_server_args(transformer_weights_path=f.name)
+            resolved = resolve_transformer_safetensors_to_load(
+                server_args, "/unused/component/path"
+            )
+
+        self.assertEqual(resolved, [f.name])
+
+    def test_filter_transformer_precision_variants_prefers_canonical_file(self):
+        files = [
+            "/tmp/transformer/diffusion_pytorch_model.fp16.safetensors",
+            "/tmp/transformer/diffusion_pytorch_model.safetensors",
+            "/tmp/transformer/other.safetensors",
+        ]
+
+        resolved = _filter_duplicate_precision_variant_safetensors(files)
+
+        self.assertEqual(
+            resolved,
+            [
+                "/tmp/transformer/diffusion_pytorch_model.safetensors",
+                "/tmp/transformer/other.safetensors",
+            ],
+        )
+
+    def test_filter_transformer_precision_variants_keeps_precision_only_family(self):
+        files = [
+            "/tmp/transformer/diffusion_pytorch_model.bf16.safetensors",
+            "/tmp/transformer/diffusion_pytorch_model.fp16.safetensors",
+        ]
+
+        resolved = _filter_duplicate_precision_variant_safetensors(files)
+
+        self.assertEqual(resolved, files)
+
+    @patch(
+        "sglang.multimodal_gen.runtime.loader.transformer_load_utils.build_nvfp4_config_from_safetensors_list",
+        return_value=None,
+    )
+    @patch(
+        "sglang.multimodal_gen.runtime.loader.transformer_load_utils.get_quant_config_from_safetensors_metadata",
+        return_value=None,
+    )
+    @patch(
+        "sglang.multimodal_gen.runtime.loader.transformer_load_utils.get_metadata_from_safetensors_file"
+    )
+    def test_resolve_transformer_quant_load_spec_keeps_nunchaku_hook(
+        self,
+        mock_metadata,
+        _mock_quant_metadata,
+        _mock_nvfp4,
+    ):
+        mock_metadata.return_value = {
+            "config": json.dumps({"_class_name": _FakeFluxTransformer.__name__})
+        }
+        nunchaku_config = NunchakuConfig(
+            transformer_weights_path="/tmp/svdq-int4_r32.safetensors"
+        )
+        server_args = self._make_server_args(
+            transformer_weights_path=nunchaku_config.transformer_weights_path,
+            nunchaku_config=nunchaku_config,
+        )
+
+        spec = resolve_transformer_quant_load_spec(
+            hf_config={},
+            server_args=server_args,
+            safetensors_list=[nunchaku_config.transformer_weights_path],
+            component_model_path="/unused/component/path",
+            model_cls=_FakeFluxTransformer,
+            cls_name=_FakeFluxTransformer.__name__,
+        )
+
+        self.assertIsNone(spec.quant_config)
+        self.assertIs(spec.nunchaku_config, nunchaku_config)
+        self.assertIsNone(spec.param_dtype)
+        self.assertEqual(len(spec.post_load_hooks), 1)
+        self.assertIs(nunchaku_config.model_cls, _FakeFluxTransformer)
+
+    def test_flux2_mixed_nvfp4_fallback_disables_conflicting_offloads(self):
+        server_args = self._make_server_args(
+            transformer_weights_path="/tmp/flux2-dev-nvfp4-mixed.safetensors",
+            tp_size=2,
+            dit_cpu_offload=True,
+            text_encoder_cpu_offload=True,
+        )
+
+        _Flux2Nvfp4FallbackAdapter._maybe_adjust_flux2_nvfp4_fallback_defaults(
+            cls_name="Flux2Transformer2DModel",
+            server_args=server_args,
+            quant_config=_FakeQuantConfig(),
+        )
+
+        self.assertFalse(server_args.dit_cpu_offload)
+        self.assertFalse(server_args.text_encoder_cpu_offload)
+
+
+if __name__ == "__main__":
+    unittest.main()