From f36a3816319d28ebb773f27b4a1f9d52fc205b2d Mon Sep 17 00:00:00 2001 From: chrisqian Date: Sat, 24 Oct 2026 10:14:41 +0800 Subject: [PATCH 1/4] fix(models): find index-only checkpoints that keep shards in a subdirectory --- python/freetoken/models/loader.py | 11 +++++++- tests/models/test_models_loader.py | 45 ++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 4a9ac94c7..05c49b528 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -51,7 +51,16 @@ def shard_tensor( def iter_weight_files(model_path: str) -> list[str]: model_folder = download_hf_weight(model_path) files = glob.glob(f"{model_folder}/*.safetensors") - return [f for f in files if not f.endswith("consolidated.safetensors")] or files + files = [f for f in files if not f.endswith("consolidated.safetensors")] or files + if files: + return files + # Root-glob misses checkpoints that keep their shards in a subdirectory (nvidia's FP4-Local + # releases ship tensors/model-NNNNN-of-NNNNN.safetensors): the index's weight_map holds the + # shard paths relative to the model folder, so resolve them from there. + if not os.path.exists(os.path.join(model_folder, "model.safetensors.index.json")): + return files + shards = sorted({shard for shard in safetensors_weight_map(model_folder).values()}) + return [os.path.join(model_folder, shard) for shard in shards] def safetensors_weight_map(folder: str) -> dict[str, str]: diff --git a/tests/models/test_models_loader.py b/tests/models/test_models_loader.py index 4f634c4c9..debf73a84 100644 --- a/tests/models/test_models_loader.py +++ b/tests/models/test_models_loader.py @@ -31,6 +31,51 @@ def test_shard_tensor_splits_vocab_with_ceil_partition(): assert rank1.tolist() == value[3:5].tolist() +def test_iter_weight_files_reads_shards_from_an_indexed_subdirectory(tmp_path): + """FP4-Local releases keep shards in tensors/; the root glob alone finds nothing.""" + import json + import os + from freetoken.models.loader import iter_weight_files + + sub = tmp_path / "tensors" + sub.mkdir() + for i in (1, 2): + (sub / f"model-{i:05d}-of-00002.safetensors").write_bytes(b"") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "model.embed_tokens.weight": "tensors/model-00001-of-00002.safetensors", + "lm_head.weight": "tensors/model-00002-of-00002.safetensors", + } + } + ) + ) + + files = iter_weight_files(str(tmp_path)) + + assert len(files) == 2 + assert all(os.path.exists(f) for f in files) + + +def test_iter_weight_files_prefers_root_shards_over_the_index(tmp_path): + """Root-level shards keep the historical path; the index fallback is only for subdirectory layouts.""" + import json + from freetoken.models.loader import iter_weight_files + + (tmp_path / "model-00001-of-00001.safetensors").write_bytes(b"") + sub = tmp_path / "tensors" + sub.mkdir() + (sub / "model-00001-of-00002.safetensors").write_bytes(b"") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"model.embed_tokens.weight": "tensors/model-00001-of-00002.safetensors"}}) + ) + + files = iter_weight_files(str(tmp_path)) + + assert files == [str(tmp_path / "model-00001-of-00001.safetensors")] + + def test_iter_root_safetensor_files_from_index_keeps_only_root_index_shards(tmp_path): import json from freetoken.models.loader import iter_root_safetensor_files_from_index From 40052644caac3ddbb75a136ec581af9bf6742fc2 Mon Sep 17 00:00:00 2001 From: chrisqian Date: Sat, 24 Oct 2026 15:15:52 +0800 Subject: [PATCH 2/4] fix(models): make the checkpoint index authoritative for shard discovery The index is read directly so safetensors_weight_map stays the only header-based reader, and _host_ram_fits_parallel sizes the shards the reader will actually open, so the low-RAM serial fallback fires again for subdirectory layouts. --- python/freetoken/models/loader.py | 18 +++++++-------- python/freetoken/moe/expert_banks.py | 8 +++---- tests/models/test_models_loader.py | 21 +++++++++++++----- tests/moe/test_offload.py | 33 ++++++++++++++++++++++++++++ 4 files changed, 59 insertions(+), 21 deletions(-) diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 05c49b528..9ace932f1 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -50,17 +50,15 @@ def shard_tensor( def iter_weight_files(model_path: str) -> list[str]: model_folder = download_hf_weight(model_path) + # The index is authoritative: download_hf_weight only fetches the shards it names, and it + # covers subdirectory layouts (nvidia's FP4-Local releases ship tensors/model-*.safetensors). + index = os.path.join(model_folder, "model.safetensors.index.json") + if os.path.exists(index): + with open(index, encoding="utf-8") as f: + shards = sorted(set(json.load(f)["weight_map"].values())) + return [os.path.join(model_folder, shard) for shard in shards] files = glob.glob(f"{model_folder}/*.safetensors") - files = [f for f in files if not f.endswith("consolidated.safetensors")] or files - if files: - return files - # Root-glob misses checkpoints that keep their shards in a subdirectory (nvidia's FP4-Local - # releases ship tensors/model-NNNNN-of-NNNNN.safetensors): the index's weight_map holds the - # shard paths relative to the model folder, so resolve them from there. - if not os.path.exists(os.path.join(model_folder, "model.safetensors.index.json")): - return files - shards = sorted({shard for shard in safetensors_weight_map(model_folder).values()}) - return [os.path.join(model_folder, shard) for shard in shards] + return [f for f in files if not f.endswith("consolidated.safetensors")] or files def safetensors_weight_map(folder: str) -> dict[str, str]: diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index dcb797970..197dbb9b0 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -8,7 +8,6 @@ from __future__ import annotations -import glob import math import os from dataclasses import dataclass, field @@ -268,13 +267,12 @@ def _host_ram_fits_parallel(model_path: str) -> bool: pass if avail is None: return True - try: # resolve a hub id to its local cache dir (no-op for a local path) so glob sees the shards - from freetoken.utils.hf import download_hf_weight + try: # the loader's index-first shard list also covers subdirectory layouts (a root glob does not) + from freetoken.models.loader import iter_weight_files - model_path = download_hf_weight(model_path) + sizes = [os.path.getsize(p) for p in iter_weight_files(model_path)] except Exception: return True - sizes = [os.path.getsize(p) for p in glob.glob(os.path.join(model_path, "*.safetensors"))] if not sizes: return True return avail > sum(sizes) + max(sizes) diff --git a/tests/models/test_models_loader.py b/tests/models/test_models_loader.py index debf73a84..628124d76 100644 --- a/tests/models/test_models_loader.py +++ b/tests/models/test_models_loader.py @@ -58,22 +58,31 @@ def test_iter_weight_files_reads_shards_from_an_indexed_subdirectory(tmp_path): assert all(os.path.exists(f) for f in files) -def test_iter_weight_files_prefers_root_shards_over_the_index(tmp_path): - """Root-level shards keep the historical path; the index fallback is only for subdirectory layouts.""" +def test_iter_weight_files_lets_the_index_decide_which_shards_load(tmp_path): + """The index is authoritative: it may name root and subdirectory shards; an unlisted root file is not loaded.""" import json from freetoken.models.loader import iter_weight_files - (tmp_path / "model-00001-of-00001.safetensors").write_bytes(b"") + root = tmp_path / "model-00001-of-00002.safetensors" + root.write_bytes(b"") + (tmp_path / "unlisted.safetensors").write_bytes(b"") sub = tmp_path / "tensors" sub.mkdir() - (sub / "model-00001-of-00002.safetensors").write_bytes(b"") + (sub / "model-00002-of-00002.safetensors").write_bytes(b"") (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": {"model.embed_tokens.weight": "tensors/model-00001-of-00002.safetensors"}}) + json.dumps( + { + "weight_map": { + "model.embed_tokens.weight": "model-00001-of-00002.safetensors", + "lm_head.weight": "tensors/model-00002-of-00002.safetensors", + } + } + ) ) files = iter_weight_files(str(tmp_path)) - assert files == [str(tmp_path / "model-00001-of-00001.safetensors")] + assert files == [str(root), str(sub / "model-00002-of-00002.safetensors")] def test_iter_root_safetensor_files_from_index_keeps_only_root_index_shards(tmp_path): diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 8833cdf98..609ef0d41 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -1206,3 +1206,36 @@ def boom(addr, nbytes): with hb.PinPipeline() as pins: pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)}) assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value} + + +def test_host_ram_fits_parallel_sizes_subdirectory_shards(tmp_path, monkeypatch): + """The RAM probe must size the shards the reader will actually open, indexed subdirectory layouts included.""" + import builtins + import io + import json + + from freetoken.moe.expert_banks import _host_ram_fits_parallel + + sub = tmp_path / "tensors" + sub.mkdir() + for i in (1, 2): + (sub / f"model-{i:05d}-of-00002.safetensors").write_bytes(b"x" * (1 << 20)) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "model.layers.0.mlp.experts.0.gate_proj.weight": "tensors/model-00001-of-00002.safetensors", + "model.layers.0.mlp.experts.0.down_proj.weight": "tensors/model-00002-of-00002.safetensors", + } + } + ) + ) + real_open = builtins.open + + def fake_open(file, *args, **kwargs): + if str(file) == "/proc/meminfo": + return io.StringIO("MemAvailable: 512 kB\n") + return real_open(file, *args, **kwargs) + + monkeypatch.setattr(builtins, "open", fake_open) + assert not _host_ram_fits_parallel(str(tmp_path)) From ad95ce2d2f4b4805ececaf87183566c5c5868a8e Mon Sep 17 00:00:00 2001 From: chrisqian Date: Sat, 24 Oct 2026 16:00:21 +0800 Subject: [PATCH 3/4] fix(models): fall back to root shards when the index is unusable Share one index parser between iter_weight_files and safetensors_weight_map instead of a third copy; warn and glob the root when the index is broken, absent or empty; reject index keys that escape the model folder; let the hub download fetch the index itself; tell the truth in the diagnose docstring. --- python/freetoken/models/loader.py | 50 +++++++++++++++++++------ python/freetoken/utils/hf.py | 2 +- scripts/checkpoint_diagnose/diagnose.py | 2 +- tests/models/test_models_loader.py | 50 +++++++++++++++++++++++++ tests/moe/test_offload.py | 15 +++++--- 5 files changed, 100 insertions(+), 19 deletions(-) diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 9ace932f1..64abd1901 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -9,7 +9,11 @@ from typing import Iterable, Iterator import torch + from freetoken.utils import div_ceil, download_hf_weight +from freetoken.utils.logger import init_logger + +logger = init_logger(__name__) SPLIT_DIM_0 = (".q_proj", ".k_proj", ".v_proj", ".gate_proj", ".up_proj") SPLIT_DIM_1 = (".o_proj", ".down_proj") @@ -48,26 +52,48 @@ def shard_tensor( return value +def _index_weight_map(folder: str) -> dict[str, str] | None: + """The index's weight_map, or None when the checkpoint ships no usable one: a bad + index must not hide shards that are on disk.""" + index = os.path.join(folder, "model.safetensors.index.json") + if not os.path.exists(index): + return None + try: + with open(index, encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + except (OSError, ValueError, KeyError) as e: + logger.warning("ignoring unusable %s (%s); falling back to root shards", index, e) + return None + if not weight_map: + logger.warning("ignoring an empty weight_map in %s; falling back to root shards", index) + return None + return weight_map + + def iter_weight_files(model_path: str) -> list[str]: model_folder = download_hf_weight(model_path) - # The index is authoritative: download_hf_weight only fetches the shards it names, and it - # covers subdirectory layouts (nvidia's FP4-Local releases ship tensors/model-*.safetensors). - index = os.path.join(model_folder, "model.safetensors.index.json") - if os.path.exists(index): - with open(index, encoding="utf-8") as f: - shards = sorted(set(json.load(f)["weight_map"].values())) - return [os.path.join(model_folder, shard) for shard in shards] + # The index is authoritative: it also names shards in subdirectories, which the root + # glob misses (nvidia's FP4-Local releases ship tensors/model-*.safetensors). + weight_map = _index_weight_map(model_folder) + if weight_map: + paths = [] + root = os.path.abspath(model_folder) + for shard in sorted(set(weight_map.values())): + path = os.path.abspath(os.path.join(model_folder, shard)) + if not path.startswith(root + os.sep): + raise ValueError(f"index shard {shard!r} escapes the model directory {model_folder}") + paths.append(path) + return paths files = glob.glob(f"{model_folder}/*.safetensors") return [f for f in files if not f.endswith("consolidated.safetensors")] or files def safetensors_weight_map(folder: str) -> dict[str, str]: """Tensor name -> shard basename, from the index or from each shard's header when the checkpoint ships none.""" - index = os.path.join(folder, "model.safetensors.index.json") - if os.path.exists(index): - with open(index, encoding="utf-8") as f: - return json.load(f)["weight_map"] - weight_map: dict[str, str] = {} + weight_map = _index_weight_map(folder) + if weight_map: + return weight_map + weight_map = {} for path in sorted(iter_weight_files(folder)): with open(path, "rb") as fh: n = struct.unpack(" list[str]: SAFE_WEIGHTS_INDEX_NAME, repo_id, e, ) return ["*.safetensors"] - return shards or ["*.safetensors"] + return [*shards, SAFE_WEIGHTS_INDEX_NAME] if shards else ["*.safetensors"] def download_hf_weight(model_path: str) -> str: diff --git a/scripts/checkpoint_diagnose/diagnose.py b/scripts/checkpoint_diagnose/diagnose.py index 0b1b1e446..bd87a8348 100644 --- a/scripts/checkpoint_diagnose/diagnose.py +++ b/scripts/checkpoint_diagnose/diagnose.py @@ -83,7 +83,7 @@ def group_names(names, limit: int = 4) -> str: def shard_files(path: str) -> list[str]: - """The root shards the loaders read: consolidated.safetensors only when it is the only one (loader.iter_weight_files).""" + """Root *.safetensors names, consolidated.safetensors only when it is the only one; subdirectory layouts are reached through safetensors_weight_map below.""" files = sorted(f for f in os.listdir(path) if f.endswith(".safetensors")) return [f for f in files if f != "consolidated.safetensors"] or files diff --git a/tests/models/test_models_loader.py b/tests/models/test_models_loader.py index 628124d76..0c044e066 100644 --- a/tests/models/test_models_loader.py +++ b/tests/models/test_models_loader.py @@ -85,6 +85,56 @@ def test_iter_weight_files_lets_the_index_decide_which_shards_load(tmp_path): assert files == [str(root), str(sub / "model-00002-of-00002.safetensors")] +def test_iter_weight_files_falls_back_to_root_shards_without_an_index(tmp_path): + """No index: the root glob loads, excluding consolidated.safetensors unless it is the only shard.""" + from freetoken.models.loader import iter_weight_files + + a = tmp_path / "model-00001-of-00002.safetensors" + b = tmp_path / "model-00002-of-00002.safetensors" + consolidated = tmp_path / "consolidated.safetensors" + for path in (a, b, consolidated): + path.write_bytes(b"") + + assert sorted(iter_weight_files(str(tmp_path))) == sorted([str(a), str(b)]) + + a.unlink() + b.unlink() + assert iter_weight_files(str(tmp_path)) == [str(consolidated)] + + +def test_iter_weight_files_falls_back_when_the_index_is_unusable(tmp_path): + """A broken or empty index must not hide shards that are on disk.""" + import json + + from freetoken.models.loader import iter_weight_files + + shard = tmp_path / "model.safetensors" + shard.write_bytes(b"") + index = tmp_path / "model.safetensors.index.json" + + index.write_text("{not json") + assert iter_weight_files(str(tmp_path)) == [str(shard)] + + index.write_text(json.dumps({"weight_map": {}})) + assert iter_weight_files(str(tmp_path)) == [str(shard)] + + +def test_iter_weight_files_rejects_an_index_that_escapes_the_folder(tmp_path): + """Index keys are checkpoint-supplied: a ../ path must not make the loader read outside it.""" + import json + + import pytest + + from freetoken.models.loader import iter_weight_files + + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"x": "../outside.safetensors"}}) + ) + + with pytest.raises(ValueError, match="escapes the model directory"): + iter_weight_files(str(tmp_path)) + + def test_iter_root_safetensor_files_from_index_keeps_only_root_index_shards(tmp_path): import json from freetoken.models.loader import iter_root_safetensor_files_from_index diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 609ef0d41..96caba66d 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -1232,10 +1232,15 @@ def test_host_ram_fits_parallel_sizes_subdirectory_shards(tmp_path, monkeypatch) ) real_open = builtins.open - def fake_open(file, *args, **kwargs): - if str(file) == "/proc/meminfo": - return io.StringIO("MemAvailable: 512 kB\n") - return real_open(file, *args, **kwargs) + def fake_open(meminfo): + def open_(file, *args, **kwargs): + if str(file) == "/proc/meminfo": + return io.StringIO(meminfo) + return real_open(file, *args, **kwargs) - monkeypatch.setattr(builtins, "open", fake_open) + return open_ + + monkeypatch.setattr(builtins, "open", fake_open("MemAvailable: 512 kB\n")) assert not _host_ram_fits_parallel(str(tmp_path)) + monkeypatch.setattr(builtins, "open", fake_open("MemAvailable: 999999999 kB\n")) + assert _host_ram_fits_parallel(str(tmp_path)) From 00b855c833bd2886c373dc022417d5942bc2a7da Mon Sep 17 00:00:00 2001 From: chrisqian Date: Sun, 25 Oct 2026 09:42:15 +0800 Subject: [PATCH 4/4] revert(models): back out the defensive index fallback, keep the no-index test Per review: a broken index is a broken checkpoint, so the loader keeps the plain one rule -- use the index if there is one, otherwise glob the root. The no-index fallback test stays as the only coverage of the root glob and the consolidated.safetensors rule. --- python/freetoken/models/loader.py | 50 ++++++------------------- python/freetoken/utils/hf.py | 2 +- scripts/checkpoint_diagnose/diagnose.py | 2 +- tests/models/test_models_loader.py | 33 ---------------- tests/moe/test_offload.py | 15 +++----- 5 files changed, 19 insertions(+), 83 deletions(-) diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 64abd1901..9ace932f1 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -9,11 +9,7 @@ from typing import Iterable, Iterator import torch - from freetoken.utils import div_ceil, download_hf_weight -from freetoken.utils.logger import init_logger - -logger = init_logger(__name__) SPLIT_DIM_0 = (".q_proj", ".k_proj", ".v_proj", ".gate_proj", ".up_proj") SPLIT_DIM_1 = (".o_proj", ".down_proj") @@ -52,48 +48,26 @@ def shard_tensor( return value -def _index_weight_map(folder: str) -> dict[str, str] | None: - """The index's weight_map, or None when the checkpoint ships no usable one: a bad - index must not hide shards that are on disk.""" - index = os.path.join(folder, "model.safetensors.index.json") - if not os.path.exists(index): - return None - try: - with open(index, encoding="utf-8") as f: - weight_map = json.load(f)["weight_map"] - except (OSError, ValueError, KeyError) as e: - logger.warning("ignoring unusable %s (%s); falling back to root shards", index, e) - return None - if not weight_map: - logger.warning("ignoring an empty weight_map in %s; falling back to root shards", index) - return None - return weight_map - - def iter_weight_files(model_path: str) -> list[str]: model_folder = download_hf_weight(model_path) - # The index is authoritative: it also names shards in subdirectories, which the root - # glob misses (nvidia's FP4-Local releases ship tensors/model-*.safetensors). - weight_map = _index_weight_map(model_folder) - if weight_map: - paths = [] - root = os.path.abspath(model_folder) - for shard in sorted(set(weight_map.values())): - path = os.path.abspath(os.path.join(model_folder, shard)) - if not path.startswith(root + os.sep): - raise ValueError(f"index shard {shard!r} escapes the model directory {model_folder}") - paths.append(path) - return paths + # The index is authoritative: download_hf_weight only fetches the shards it names, and it + # covers subdirectory layouts (nvidia's FP4-Local releases ship tensors/model-*.safetensors). + index = os.path.join(model_folder, "model.safetensors.index.json") + if os.path.exists(index): + with open(index, encoding="utf-8") as f: + shards = sorted(set(json.load(f)["weight_map"].values())) + return [os.path.join(model_folder, shard) for shard in shards] files = glob.glob(f"{model_folder}/*.safetensors") return [f for f in files if not f.endswith("consolidated.safetensors")] or files def safetensors_weight_map(folder: str) -> dict[str, str]: """Tensor name -> shard basename, from the index or from each shard's header when the checkpoint ships none.""" - weight_map = _index_weight_map(folder) - if weight_map: - return weight_map - weight_map = {} + index = os.path.join(folder, "model.safetensors.index.json") + if os.path.exists(index): + with open(index, encoding="utf-8") as f: + return json.load(f)["weight_map"] + weight_map: dict[str, str] = {} for path in sorted(iter_weight_files(folder)): with open(path, "rb") as fh: n = struct.unpack(" list[str]: SAFE_WEIGHTS_INDEX_NAME, repo_id, e, ) return ["*.safetensors"] - return [*shards, SAFE_WEIGHTS_INDEX_NAME] if shards else ["*.safetensors"] + return shards or ["*.safetensors"] def download_hf_weight(model_path: str) -> str: diff --git a/scripts/checkpoint_diagnose/diagnose.py b/scripts/checkpoint_diagnose/diagnose.py index bd87a8348..0b1b1e446 100644 --- a/scripts/checkpoint_diagnose/diagnose.py +++ b/scripts/checkpoint_diagnose/diagnose.py @@ -83,7 +83,7 @@ def group_names(names, limit: int = 4) -> str: def shard_files(path: str) -> list[str]: - """Root *.safetensors names, consolidated.safetensors only when it is the only one; subdirectory layouts are reached through safetensors_weight_map below.""" + """The root shards the loaders read: consolidated.safetensors only when it is the only one (loader.iter_weight_files).""" files = sorted(f for f in os.listdir(path) if f.endswith(".safetensors")) return [f for f in files if f != "consolidated.safetensors"] or files diff --git a/tests/models/test_models_loader.py b/tests/models/test_models_loader.py index 0c044e066..30a97b9b4 100644 --- a/tests/models/test_models_loader.py +++ b/tests/models/test_models_loader.py @@ -102,39 +102,6 @@ def test_iter_weight_files_falls_back_to_root_shards_without_an_index(tmp_path): assert iter_weight_files(str(tmp_path)) == [str(consolidated)] -def test_iter_weight_files_falls_back_when_the_index_is_unusable(tmp_path): - """A broken or empty index must not hide shards that are on disk.""" - import json - - from freetoken.models.loader import iter_weight_files - - shard = tmp_path / "model.safetensors" - shard.write_bytes(b"") - index = tmp_path / "model.safetensors.index.json" - - index.write_text("{not json") - assert iter_weight_files(str(tmp_path)) == [str(shard)] - - index.write_text(json.dumps({"weight_map": {}})) - assert iter_weight_files(str(tmp_path)) == [str(shard)] - - -def test_iter_weight_files_rejects_an_index_that_escapes_the_folder(tmp_path): - """Index keys are checkpoint-supplied: a ../ path must not make the loader read outside it.""" - import json - - import pytest - - from freetoken.models.loader import iter_weight_files - - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": {"x": "../outside.safetensors"}}) - ) - - with pytest.raises(ValueError, match="escapes the model directory"): - iter_weight_files(str(tmp_path)) - - def test_iter_root_safetensor_files_from_index_keeps_only_root_index_shards(tmp_path): import json from freetoken.models.loader import iter_root_safetensor_files_from_index diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 96caba66d..609ef0d41 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -1232,15 +1232,10 @@ def test_host_ram_fits_parallel_sizes_subdirectory_shards(tmp_path, monkeypatch) ) real_open = builtins.open - def fake_open(meminfo): - def open_(file, *args, **kwargs): - if str(file) == "/proc/meminfo": - return io.StringIO(meminfo) - return real_open(file, *args, **kwargs) + def fake_open(file, *args, **kwargs): + if str(file) == "/proc/meminfo": + return io.StringIO("MemAvailable: 512 kB\n") + return real_open(file, *args, **kwargs) - return open_ - - monkeypatch.setattr(builtins, "open", fake_open("MemAvailable: 512 kB\n")) + monkeypatch.setattr(builtins, "open", fake_open) assert not _host_ram_fits_parallel(str(tmp_path)) - monkeypatch.setattr(builtins, "open", fake_open("MemAvailable: 999999999 kB\n")) - assert _host_ram_fits_parallel(str(tmp_path))