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

[diffusion] feat: filter duplicate precision variants across custom loaders (#37616)

Browse Source 
Co-authored-by: mickqian <mickqian@users.noreply.github.com>
This commit is contained in:
dujifeng
2026-09-03 21:54:11 +08:00
committed by GitHub
co-authored by mickqian
parent f5bed255c0
commit 9cb38a3d57
10 changed files with 288 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+14 -28
View File
@@ -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:
python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py
+3 -3
View File
@@ -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:
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+11 -1
View File
@@ -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}"
python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
-59
View File
@@ -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<stem>.+?)(?P<precision>\.(?:fp16|bf16|fp32))(?P<shard>-\d+-of-\d+)?(?P<ext>\.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,
python/sglang/multimodal_gen/runtime/loader/utils.py
+57 -27
View File
@@ -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:
python/sglang/multimodal_gen/runtime/weights/source.py
+38 -1
View File
@@ -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<stem>.+?)\.(?:fp16|bf16|fp32)(?:-\d+-of-\d+)?"
r"(?P<suffix>\.safetensors(?:\.index\.json)?)$"
)
_CANONICAL_SAFETENSORS_SUFFIX_RE = re.compile(
r"^(?P<stem>.+?)(?:-\d+-of-\d+)?" r"(?P<suffix>\.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,
python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py
+26
View File
@@ -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.
python/sglang/multimodal_gen/test/unit/test_transformer_quant.py
+5 -3
View File
@@ -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)
python/sglang/multimodal_gen/test/unit/test_weight_source.py
+38
View File
@@ -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(
python/sglang/multimodal_gen/test/unit/test_weight_utils.py
+96
View File
@@ -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 (