fix(models): discover safetensors shards from the index

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
 

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)

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

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))