diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 4a9ac94c7..9ace932f1 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -50,6 +50,13 @@ 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") return [f for f in files if not f.endswith("consolidated.safetensors")] or files 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 4f634c4c9..30a97b9b4 100644 --- a/tests/models/test_models_loader.py +++ b/tests/models/test_models_loader.py @@ -31,6 +31,77 @@ 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_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 + + 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-00002-of-00002.safetensors").write_bytes(b"") + (tmp_path / "model.safetensors.index.json").write_text( + 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(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_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 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))