diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 2d738de25..1d5f87921 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/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.+?)(?P\.(?:fp16|bf16|fp32))(?P-\d+-of-\d+)?(?P\.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, diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index 64ec3b052..5f786eeb8 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/server/consistency_threshold.json b/python/sglang/multimodal_gen/test/server/consistency_threshold.json index 596e98166..45e04b413 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_threshold.json +++ b/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": { - "clip_threshold": 0.92, - "ssim_threshold": 0.91, - "psnr_threshold": 21.5, + "flux_2_image_t2i": { + "clip_threshold": 0.96, + "ssim_threshold": 0.95, + "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, - "ssim_threshold": 0.84, - "psnr_threshold": 19.0, - "mean_abs_diff_threshold": 10.0 + "clip_threshold": 0.94, + "ssim_threshold": 0.88, + "psnr_threshold": 21.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, - "ssim_threshold": 0.40, - "psnr_threshold": 13.2, - "mean_abs_diff_threshold": 32.0 + "clip_threshold": 0.88, + "ssim_threshold": 0.75, + "psnr_threshold": 17.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, - "ssim_threshold": 0.81, - "psnr_threshold": 20.4, - "mean_abs_diff_threshold": 10.0 + "clip_threshold": 0.92, + "ssim_threshold": 0.88, + "psnr_threshold": 22.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, diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py new file mode 100644 index 000000000..71e48e492 --- /dev/null +++ b/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()