From 9cb38a3d5750fa9fd1232ecdedefe714948c8fe7 Mon Sep 17 00:00:00 2001 From: dujifeng <42289967+dujifeng@users.noreply.github.com> Date: Thu, 3 Sep 2026 21:54:11 +0800 Subject: [PATCH] [diffusion] feat: filter duplicate precision variants across custom loaders (#37616) Co-authored-by: mickqian --- .../component_loaders/text_encoder_loader.py | 42 +++----- .../component_loaders/upsampler_loader.py | 6 +- .../loader/component_loaders/vae_loader.py | 12 ++- .../runtime/loader/transformer_load_utils.py | 59 ------------ .../multimodal_gen/runtime/loader/utils.py | 84 ++++++++++------ .../multimodal_gen/runtime/weights/source.py | 39 +++++++- .../test/unit/test_text_encoder_loader.py | 26 +++++ .../test/unit/test_transformer_quant.py | 8 +- .../test/unit/test_weight_source.py | 38 ++++++++ .../test/unit/test_weight_utils.py | 96 +++++++++++++++++++ 10 files changed, 288 insertions(+), 122 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index c0c44e5ae..c7f9ac656 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -67,12 +67,13 @@ from sglang.multimodal_gen.runtime.loader.gguf_weights import ( remap_gguf_tensor_meta, ) from sglang.multimodal_gen.runtime.loader.utils import ( + _list_safetensors_files, + checkpoint_bytes, get_param_names_mapping, set_default_torch_dtype, skip_init_modules, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( - filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, pt_weights_iterator, safetensors_weights_iterator, @@ -462,31 +463,16 @@ def _require_quantized_encoder_layers( ) -def _checkpoint_bytes(model_path: str) -> int: - """On-disk size of a checkpoint, readable before any weight of it is.""" - if os.path.isfile(model_path): - return os.path.getsize(model_path) - total = 0 - for path in glob.glob( - os.path.join(str(model_path), "**", "*.safetensors"), recursive=True - ): - try: - total += os.path.getsize(path) - except OSError: - continue - return total - - def _keep_this_checkpoint_mapped(model_path: str) -> bool: """Whether this encoder's weights should stay on their file mapping.""" - checkpoint_bytes = _checkpoint_bytes(model_path) - if not host_copies_would_not_fit(checkpoint_bytes): + weight_bytes = checkpoint_bytes(model_path) + if not host_copies_would_not_fit(weight_bytes): return False logger.info( "Text encoder checkpoint is %.2f GiB against %.2f GiB of host memory, " "so its compatible weights stay on the checkpoint mapping instead of " "being copied in.", - checkpoint_bytes / 1024**3, + weight_bytes / 1024**3, host_memory_available_bytes() / 1024**3, ) return True @@ -598,20 +584,20 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): hf_weights_files: list[str] = [] for pattern in allow_patterns: - hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) + if pattern == "*.safetensors": + hf_weights_files = _list_safetensors_files( + hf_folder, + index_file=index_file, + key_filter=key_filter, + ) + else: + hf_weights_files = glob.glob(os.path.join(hf_folder, pattern)) if len(hf_weights_files) > 0: if pattern == "*.safetensors": use_safetensors = True break - if use_safetensors: - hf_weights_files = filter_duplicate_safetensors_files( - hf_weights_files, - hf_folder, - index_file, - key_filter=key_filter, - ) - else: + if not use_safetensors: hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) if len(hf_weights_files) == 0: diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py index f6fbbcc12..6a4c9f9c0 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py @@ -1,4 +1,3 @@ -import glob import json import os import re @@ -10,6 +9,7 @@ from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PlainStateDictComponentLoader, ) +from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import ( LatentUpsampler, ) @@ -59,7 +59,7 @@ def _find_safetensors_file(path: str) -> str: return path if os.path.isdir(path): - files = sorted(glob.glob(os.path.join(path, "*.safetensors"))) + files = _list_safetensors_files(path) if len(files) == 1: return files[0] elif len(files) > 1: @@ -75,7 +75,7 @@ def _find_safetensors_file(path: str) -> str: try: maybe_downloaded = maybe_download_model(path) if os.path.isdir(maybe_downloaded): - files = sorted(glob.glob(os.path.join(maybe_downloaded, "*.safetensors"))) + files = _list_safetensors_files(maybe_downloaded) if len(files) == 1: return files[0] elif len(files) > 1: diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 5171b5402..32f754376 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -42,6 +42,9 @@ from sglang.multimodal_gen.runtime.utils.precision import ( resolve_component_precision, resolve_decode_precision, ) +from sglang.multimodal_gen.runtime.weights.source import ( + filter_duplicate_precision_variant_safetensors, +) from sglang.multimodal_gen.utils import PRECISION_TO_TYPE from sglang.srt.model_loader.checkpoint_quantization import ( resolve_checkpoint_quant_spec, @@ -568,7 +571,11 @@ class VAELoader(WeightOverrideComponentLoader): ) safetensors_list = [component_weights_path] else: - safetensors_list = _list_safetensors_files(component_weights_path) + # VAE configs may explicitly choose a precision variant, so their + # selector must run before the canonical fallback. + safetensors_list = _list_safetensors_files( + component_weights_path, raw_candidates=True + ) safetensors_list = self.select_weight_files( safetensors_list, component_weights_path, @@ -576,6 +583,9 @@ class VAELoader(WeightOverrideComponentLoader): component_name, vae_precision, ) + safetensors_list = filter_duplicate_precision_variant_safetensors( + safetensors_list + ) assert len(safetensors_list) >= 1, ( f"Found no safetensors files in {component_weights_path}" 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 5676dd198..78b6df23d 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -14,7 +14,6 @@ from functools import partial from typing import Callable, Optional import torch -from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME from safetensors import safe_open from torch import nn @@ -38,9 +37,6 @@ from sglang.multimodal_gen.runtime.loader.gguf_weights import ( read_gguf_tensor_meta, ) from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files -from sglang.multimodal_gen.runtime.loader.weight_utils import ( - filter_duplicate_safetensors_files, -) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( COMPONENT_OFFLOAD, ComponentResidencyError, @@ -72,9 +68,6 @@ 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)$" -) _MIXED_SAFETENSORS_RE = re.compile(r".*-mixed(?:-\d+-of-\d+)?\.safetensors$") @@ -645,17 +638,7 @@ def resolve_transformer_checkpoint_files( safetensors_list = _list_safetensors_files(component_model_path) if safetensors_list: - # Preserve legacy cleanup for the base component. Explicit overrides - # are resolved above, where an index is already the final authority. - safetensors_list = filter_duplicate_safetensors_files( - safetensors_list, - os.path.dirname(safetensors_list[0]), - SAFE_WEIGHTS_INDEX_NAME, - ) safetensors_list = _prefer_mixed_safetensors_files(safetensors_list) - safetensors_list = _filter_duplicate_precision_variant_safetensors( - safetensors_list - ) if not safetensors_list: raise ValueError(f"no safetensors files found in {component_model_path}") @@ -696,48 +679,6 @@ def _prefer_mixed_safetensors_files(safetensors_list: list[str]) -> list[str]: return mixed_files -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/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index a91623674..7b59cd02a 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -18,9 +18,14 @@ from safetensors.torch import load_file as safetensors_load_file from torch import nn from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.weights.source import ( + filter_duplicate_precision_variant_safetensors, +) logger = init_logger(__name__) +_DEFAULT_SAFETENSORS_INDEX = "diffusion_pytorch_model.safetensors.index.json" + _QUANTIZED_DTYPES = { torch.uint8, torch.float8_e4m3fn, @@ -250,11 +255,15 @@ def _try_redownload_missing_shards(model_path: str, missing: list[str]) -> bool: def checkpoint_bytes(model_path: str) -> int: - """On-disk size of every safetensors under a path, readable before any is.""" + """On-disk size of the selected safetensors checkpoint files.""" + if os.path.isfile(model_path): + return os.path.getsize(model_path) + + paths = sorted( + glob.glob(os.path.join(str(model_path), "**", "*.safetensors"), recursive=True) + ) total = 0 - for path in glob.glob( - os.path.join(str(model_path), "**", "*.safetensors"), recursive=True - ): + for path in filter_duplicate_precision_variant_safetensors(paths): try: total += os.path.getsize(path) except OSError: @@ -289,26 +298,41 @@ def keep_checkpoint_mapped(*, weight_bytes: int, component: str) -> bool: return True -def _list_safetensors_files(model_path: str) -> list[str]: - """List all .safetensors files under a directory. +def _select_safetensors_index_file(model_path: str, preferred_name: str) -> str | None: + preferred_path = os.path.join(str(model_path), preferred_name) + if os.path.exists(preferred_path): + return preferred_path - If a safetensors index file is present, verifies that every shard listed - in the index actually exists on disk. Missing shards are first repaired - automatically via HuggingFace Hub (if the path is an HF cache entry); - if repair fails a clear RuntimeError is raised. + candidates = filter_duplicate_precision_variant_safetensors( + sorted(glob.glob(os.path.join(str(model_path), "*.safetensors.index.json"))) + ) + return candidates[0] if len(candidates) == 1 else None + + +def _list_safetensors_files( + model_path: str, + *, + index_file: str = _DEFAULT_SAFETENSORS_INDEX, + key_filter: Callable[[str], bool] | None = None, + raw_candidates: bool = False, +) -> list[str]: + """Resolve the safetensors files to load from a local component path. + + An index is authoritative when present. Otherwise canonical files are + preferred over precision-suffixed copies. ``raw_candidates`` is reserved + for model-specific selectors that must choose a precision variant first. """ if os.path.isfile(model_path): return [str(model_path)] if str(model_path).endswith(".safetensors") else [] found = sorted(glob.glob(os.path.join(str(model_path), "*.safetensors"))) - index_path = os.path.join( - str(model_path), "diffusion_pytorch_model.safetensors.index.json" - ) - if os.path.exists(index_path): + index_path = _select_safetensors_index_file(model_path, index_file) + if index_path is not None: with open(index_path) as f: index = json.load(f) - expected_shards = sorted(set(index.get("weight_map", {}).values())) + weight_map = index.get("weight_map", {}) + expected_shards = sorted(set(weight_map.values())) found_basenames = {os.path.basename(p) for p in found} missing = [s for s in expected_shards if s not in found_basenames] if missing: @@ -325,24 +349,30 @@ def _list_safetensors_files(model_path: str) -> list[str]: f"`huggingface-cli download {os.path.basename(model_path)}`)." ) - return found + if not raw_candidates: + selected_shards = { + shard + for weight_name, shard in weight_map.items() + if key_filter is None or key_filter(weight_name) + } + return [ + os.path.join(str(model_path), shard) + for shard in sorted(selected_shards) + ] + + if raw_candidates: + return found + return filter_duplicate_precision_variant_safetensors(found) def load_safetensors_state_dict(model_path: str) -> dict[str, torch.Tensor]: """Load one safetensors checkpoint, including an indexed sharded set.""" - index_path = os.path.join( - str(model_path), "diffusion_pytorch_model.safetensors.index.json" - ) + index_path = _select_safetensors_index_file(model_path, _DEFAULT_SAFETENSORS_INDEX) safetensors_files = _list_safetensors_files(model_path) - if os.path.exists(index_path): - with open(index_path) as f: - index = json.load(f) - shard_names = sorted(set(index.get("weight_map", {}).values())) + if index_path is not None: state_dict: dict[str, torch.Tensor] = {} - for shard_name in shard_names: - state_dict.update( - safetensors_load_file(os.path.join(str(model_path), shard_name)) - ) + for path in safetensors_files: + state_dict.update(safetensors_load_file(path)) return state_dict if not safetensors_files: diff --git a/python/sglang/multimodal_gen/runtime/weights/source.py b/python/sglang/multimodal_gen/runtime/weights/source.py index 3b246c615..68a1e49f1 100644 --- a/python/sglang/multimodal_gen/runtime/weights/source.py +++ b/python/sglang/multimodal_gen/runtime/weights/source.py @@ -4,7 +4,8 @@ from __future__ import annotations import json import os -from collections.abc import Callable, Mapping +import re +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, replace from pathlib import Path, PurePosixPath from typing import Literal @@ -18,6 +19,13 @@ WeightSourceKind = Literal["local", "huggingface"] _WEIGHT_SUFFIXES = (".safetensors", ".gguf", ".bin", ".pt", ".pth", ".ckpt") _SAFETENSORS_INDEX_SUFFIX = ".safetensors.index.json" _WEIGHT_REFERENCE_SUFFIXES = _WEIGHT_SUFFIXES + (_SAFETENSORS_INDEX_SUFFIX,) +_PRECISION_VARIANT_SUFFIX_RE = re.compile( + r"^(?P.+?)\.(?:fp16|bf16|fp32)(?:-\d+-of-\d+)?" + r"(?P\.safetensors(?:\.index\.json)?)$" +) +_CANONICAL_SAFETENSORS_SUFFIX_RE = re.compile( + r"^(?P.+?)(?:-\d+-of-\d+)?" r"(?P\.safetensors(?:\.index\.json)?)$" +) @dataclass(frozen=True) @@ -55,6 +63,33 @@ class NoSafetensorsWeightsError(FileNotFoundError): """The source has no safetensors payload to resolve as a weight set.""" +def filter_duplicate_precision_variant_safetensors( + safetensors_files: Sequence[str], +) -> list[str]: + """Prefer canonical files over precision-suffixed copies in each family. + + A precision-only family remains valid. Sharded and unsharded copies are + compared by family so one export cannot be loaded twice merely because its + shard layout differs between variants. + """ + canonical_families: set[tuple[str, str]] = set() + for path in safetensors_files: + if _PRECISION_VARIANT_SUFFIX_RE.match(path) is not None: + continue + if match := _CANONICAL_SAFETENSORS_SUFFIX_RE.match(path): + canonical_families.add((match.group("stem"), match.group("suffix"))) + + selected: list[str] = [] + for path in safetensors_files: + match = _PRECISION_VARIANT_SUFFIX_RE.match(path) + family = ( + (match.group("stem"), match.group("suffix")) if match is not None else None + ) + if family is None or family not in canonical_families: + selected.append(path) + return selected + + def is_explicit_weight_file_reference(source: str) -> bool: """Whether a component override names one weight file, not a component root.""" expanded = os.path.expanduser(source) @@ -442,6 +477,8 @@ def resolve_safetensors_weight_set( _read_safetensors_index(inventory, selected), index_file=selected, ) + weights = tuple(filter_duplicate_precision_variant_safetensors(weights)) + indexes = tuple(filter_duplicate_precision_variant_safetensors(indexes)) if len(indexes) == 1: return ResolvedWeightSet( inventory, diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py index c6fb023ca..b1cd198fb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py @@ -1,6 +1,7 @@ import json import tempfile import unittest +from pathlib import Path from types import SimpleNamespace from unittest import mock @@ -51,6 +52,31 @@ from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLTextMod from sglang.srt.layers.linear import LinearBase as SrtLinearBase +class TestTextEncoderWeightDiscovery(unittest.TestCase): + def test_prepare_weights_prefers_canonical_over_fp16_variant(self): + with tempfile.TemporaryDirectory() as tmpdir: + model_dir = Path(tmpdir) + canonical = model_dir / "model.safetensors" + variant = model_dir / "model.fp16.safetensors" + + canonical.touch() + variant.touch() + + ( + hf_folder, + weight_files, + use_safetensors, + ) = TextEncoderLoader()._prepare_weights( + str(model_dir), + fall_back_to_pt=True, + allow_patterns_overrides=None, + ) + + self.assertEqual(hf_folder, str(model_dir)) + self.assertTrue(use_safetensors) + self.assertEqual(weight_files, [str(canonical)]) + + class TestTextEncoderClassResolution(unittest.TestCase): """load_native must not load encoder-decoder text encoders via AutoModel. diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 3ff46b69f..c1f07cd83 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -96,7 +96,6 @@ from sglang.multimodal_gen.runtime.loader.minimax_h3_weights import ( ) from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( TransformerQuantLoadSpec, - _filter_duplicate_precision_variant_safetensors, _Flux2Nvfp4FallbackAdapter, _needs_device_weight_postprocess, _resolve_quant_config, @@ -124,6 +123,9 @@ from sglang.multimodal_gen.runtime.utils.quantization_utils import ( build_nvfp4_config_from_safetensors_list, get_quant_config, ) +from sglang.multimodal_gen.runtime.weights.source import ( + filter_duplicate_precision_variant_safetensors, +) from sglang.multimodal_gen.tools.build_modelopt_nvfp4_transformer import ( _updated_quant_config, ) @@ -1142,7 +1144,7 @@ class TestTransformerQuantHelpers(unittest.TestCase): "/tmp/transformer/other.safetensors", ] - resolved = _filter_duplicate_precision_variant_safetensors(files) + resolved = filter_duplicate_precision_variant_safetensors(files) self.assertEqual( resolved, @@ -1158,7 +1160,7 @@ class TestTransformerQuantHelpers(unittest.TestCase): "/tmp/transformer/diffusion_pytorch_model.fp16.safetensors", ] - resolved = _filter_duplicate_precision_variant_safetensors(files) + resolved = filter_duplicate_precision_variant_safetensors(files) self.assertEqual(resolved, files) diff --git a/python/sglang/multimodal_gen/test/unit/test_weight_source.py b/python/sglang/multimodal_gen/test/unit/test_weight_source.py index 77ddc2c86..9757b7d6a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_weight_source.py +++ b/python/sglang/multimodal_gen/test/unit/test_weight_source.py @@ -143,6 +143,44 @@ def test_safetensors_weight_set_rejects_unindexed_variants(tmp_path): resolve_safetensors_weight_set(str(tmp_path)) +def test_safetensors_weight_set_prefers_canonical_precision_family(tmp_path): + canonical = tmp_path / "model.safetensors" + canonical.write_bytes(b"canonical") + (tmp_path / "model.fp16.safetensors").write_bytes(b"fp16") + + resolved = resolve_safetensors_weight_set(str(tmp_path)) + + assert resolved.selected_files == (canonical.name,) + + +def test_explicit_precision_variant_overrides_canonical_fallback(tmp_path): + (tmp_path / "model.safetensors").write_bytes(b"canonical") + variant = tmp_path / "model.fp16.safetensors" + variant.write_bytes(b"fp16") + + resolved = resolve_safetensors_weight_set(str(tmp_path), weight_name=variant.name) + + assert resolved.selected_files == (variant.name,) + + +def test_safetensors_weight_set_prefers_canonical_precision_index(tmp_path): + canonical = tmp_path / "model.safetensors" + variant = tmp_path / "model.fp16.safetensors" + canonical.write_bytes(b"canonical") + variant.write_bytes(b"fp16") + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map":{"weight":"model.safetensors"}}' + ) + (tmp_path / "model.fp16.safetensors.index.json").write_text( + '{"weight_map":{"weight":"model.fp16.safetensors"}}' + ) + + resolved = resolve_safetensors_weight_set(str(tmp_path)) + + assert resolved.index_file == "model.safetensors.index.json" + assert resolved.selected_files == (canonical.name,) + + def test_safetensors_index_rejects_non_weight_shard(tmp_path): (tmp_path / "config.json").write_text("{}") (tmp_path / "model.safetensors.index.json").write_text( diff --git a/python/sglang/multimodal_gen/test/unit/test_weight_utils.py b/python/sglang/multimodal_gen/test/unit/test_weight_utils.py index dfe4a769a..85477acd7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_weight_utils.py +++ b/python/sglang/multimodal_gen/test/unit/test_weight_utils.py @@ -4,16 +4,112 @@ import os import tempfile import unittest +from pathlib import Path from unittest.mock import patch +import torch +from safetensors.torch import save_file as safetensors_save_file + +from sglang.multimodal_gen.runtime.loader.utils import ( + _list_safetensors_files, + checkpoint_bytes, + load_safetensors_state_dict, +) from sglang.multimodal_gen.runtime.loader.weight_utils import ( _disable_runai_streamer_rank_discovery_collective, get_lock, ) +from sglang.multimodal_gen.runtime.weights.source import ( + filter_duplicate_precision_variant_safetensors, +) _DIST_STREAMER_MOD = "runai_model_streamer.distributed_streamer.distributed_streamer" +class TestPrecisionVariantSelection(unittest.TestCase): + def test_prefers_canonical_family_across_shard_layouts(self): + files = [ + "/tmp/model.safetensors", + "/tmp/model.fp16-00001-of-00002.safetensors", + "/tmp/model.fp16-00002-of-00002.safetensors", + "/tmp/other.bf16.safetensors", + ] + + self.assertEqual( + filter_duplicate_precision_variant_safetensors(files), + ["/tmp/model.safetensors", "/tmp/other.bf16.safetensors"], + ) + + def test_shared_state_dict_loader_uses_canonical_file(self): + with tempfile.TemporaryDirectory() as tmpdir: + model_dir = Path(tmpdir) + canonical = model_dir / "model.safetensors" + variant = model_dir / "model.fp16.safetensors" + safetensors_save_file({"weight": torch.tensor([1.0])}, canonical) + safetensors_save_file({"weight": torch.tensor([2.0])}, variant) + + state_dict = load_safetensors_state_dict(str(model_dir)) + + self.assertTrue(torch.equal(state_dict["weight"], torch.tensor([1.0]))) + + def test_index_selection_precedes_canonical_fallback(self): + with tempfile.TemporaryDirectory() as tmpdir: + model_dir = Path(tmpdir) + canonical = model_dir / "model.safetensors" + variant = model_dir / "model.fp16.safetensors" + index = model_dir / "model.safetensors.index.json" + canonical.touch() + variant.touch() + index.write_text('{"weight_map":{"weight":"model.fp16.safetensors"}}') + + selected = _list_safetensors_files(str(model_dir), index_file=index.name) + + self.assertEqual(selected, [str(variant)]) + + def test_raw_candidates_preserve_explicit_precision_choice(self): + with tempfile.TemporaryDirectory() as tmpdir: + model_dir = Path(tmpdir) + canonical = model_dir / "model.safetensors" + variant = model_dir / "model.fp16.safetensors" + canonical.touch() + variant.touch() + + selected = _list_safetensors_files(str(model_dir), raw_candidates=True) + + self.assertEqual(selected, [str(variant), str(canonical)]) + + def test_precision_only_index_is_discovered(self): + with tempfile.TemporaryDirectory() as tmpdir: + model_dir = Path(tmpdir) + shard = model_dir / "model.fp16-00001-of-00001.safetensors" + index = model_dir / "model.fp16.safetensors.index.json" + shard.touch() + index.write_text( + '{"weight_map":{"weight":"model.fp16-00001-of-00001.safetensors"}}' + ) + + selected = _list_safetensors_files(str(model_dir)) + + self.assertEqual(selected, [str(shard)]) + + def test_checkpoint_bytes_counts_only_selected_family(self): + with tempfile.TemporaryDirectory() as tmpdir: + model_dir = Path(tmpdir) + canonical = model_dir / "model.safetensors" + variant = model_dir / "model.fp16.safetensors" + canonical.write_bytes(b"a" * 17) + variant.write_bytes(b"b" * 11) + + self.assertEqual(checkpoint_bytes(str(model_dir)), 17) + + def test_checkpoint_bytes_supports_explicit_file(self): + with tempfile.NamedTemporaryFile() as checkpoint: + checkpoint.write(b"checkpoint") + checkpoint.flush() + + self.assertEqual(checkpoint_bytes(checkpoint.name), 10) + + class TestDisableRunaiStreamerRankDiscoveryCollective(unittest.TestCase): def test_never_touches_torch_distributed_even_when_initialized(self): from runai_model_streamer.distributed_streamer.distributed_streamer import (