minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-04-09 02:44:19 +08:00
committed by GitHub
parent a5ed507a16
commit 4ac6fa0d87
4 changed files with 349 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
+49
View File
@@ -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
+48
View File
@@ -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
+71 -17
View File
@@ -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,
python/sglang/multimodal_gen/test/unit/test_transformer_quant.py
+181
View File
@@ -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()