[chrisqianz]

Problem

iter_weight_files() only globs *.safetensors at the top level of the model folder. A checkpoint that ships a root model.safetensors.index.json whose weight_map points the shards into a subdirectory (e.g. tensors/model-00000-of-00021.safetensors) resolves to zero weight files: the dense scan silently loads nothing and the failure surfaces much later as an unrelated error.

Fix

When the top-level glob comes back empty, resolve the shard list from the root index's weight_map (sorted for a deterministic load order, deduplicated, joined against the folder). The fast path (flat directory) is unchanged.

Fixes #818

Tested

• Unit tests: tests/models/test_models_loader.py (12 cases, includes the subdirectory layout).
• Checkpoint: nvidia/Qwen3.8-Flash-Next-FP4-Local, ft serve --model ~/models/NVIDIA-Qwen3.8-Flash-Next-FP4-Local/ --moe-strategy offload --moe-cache-auto.
• Hardware: RTX 5090 D 32GB (driver 595.84), Ubuntu 22.04.5, Intel Core Ultra 7 265K, 125 GB RAM, torch 2.14.1+cu130, Python 3.13.

[jason-fxz]

Thanks for the fix. I checked it against the nvidia/Qwen3.8-Flash-Next-FP4-Local index (all 21 shards under tensors/). The unit tests pass, and the subdirectory test fails without the change.

Two requests:

1. Please read the index directly in iter_weight_files instead of calling safetensors_weight_map, and use it whenever it exists:

   def iter_weight_files(model_path: str) -> list[str]:
       model_folder = download_hf_weight(model_path)
       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

   safetensors_weight_map calls iter_weight_files when there is no index, so with the current fallback the two functions call each other. It only stops because both check the same file. Taking the index first keeps the dependency one way. It also matches download_hf_weight, which already downloads only the shards the index names, and it covers checkpoints that split shards between the root and a subdirectory. test_iter_weight_files_prefers_root_shards_over_the_index then needs to flip: a root file that the index does not list is not loaded.

2. _host_ram_fits_parallel in python/freetoken/moe/expert_banks.py also globs only the root. For this layout it finds no files and returns True, so the low-RAM fallback to the serial expert reader never fires. Please use iter_weight_files(model_path) there too.

[chrisqianz]

Both addressed in 4005264:

1. iter_weight_files now reads the index directly and takes it as authoritative whenever it exists, so safetensors_weight_map no longer takes part in shard discovery and the dependency stays one-way. test_iter_weight_files_prefers_root_shards_over_the_index flipped to test_iter_weight_files_lets_the_index_decide_which_shards_load: the index may name both root and tensors/ shards, and a root file it does not list is not loaded.
2. _host_ram_fits_parallel now sizes iter_weight_files(model_path) instead of globbing the root, so the low-RAM serial fallback fires for tensors/ layouts too. Added test_host_ram_fits_parallel_sizes_subdirectory_shards: it fails against the root glob and passes with the index.

[chrisqianz]

Follow-up in ad95ce2, from an independent review pass of my own diff:

• one shared index parser for iter_weight_files / safetensors_weight_map (_index_weight_map), replacing the third copy of the json.load(...)[weight_map] idiom;
• a present-but-unusable index (malformed JSON, missing or empty weight_map) no longer crashes or silently yields zero shards: it warns and falls back to the root glob, so good shards on disk stay loadable (this matches what hf.py already does on the hub path);
• an index key pointing outside the model folder (../x.safetensors) raises instead of reading through it;
• _weight_allow_patterns now fetches model.safetensors.index.json itself, so the loader's index-first premise no longer depends on the probe's cache-dir side effect;
• the checkpoint_diagnose shard_files docstring no longer cites loader semantics this PR changed;
• the flipped test killed the only coverage of the no-index glob + consolidated exclusion, so the fallback paths now have their own tests (no-index, broken/empty index, escaping key). The RAM-probe test gained a positive control (plenty of RAM -> True).

Server: 15 passed in tests/models/test_models_loader.py, RAM test 1 passed; the three new tests fail on 4005264 and pass on ad95ce2.

[jason-fxz]

Thanks, 4005264 is what I was after: the index decides when it exists, and the RAM probe sizes the same files.

Could you revert ad95ce2 and keep only test_iter_weight_files_falls_back_to_root_shards_without_an_index? That test is worth having: once the root-precedence test flipped, nothing else covered the no-index glob and the consolidated.safetensors rule.

The rest is more defensive than this loader needs. A broken index is a broken checkpoint, so a plain error is clearer than a warning and a guess from the root glob, and index paths come from a checkpoint we already trust. The fallback also misses a few shapes (a list at the top level, or a list as weight_map), and the escape check covers iter_weight_files only, not the safetensors_weight_map readers. Rather than complete both, I'd rather keep the PR to the one rule: use the index if there is one, otherwise glob the root.

After that this is ready to merge.

[chrisqianz]

Reverted ad95ce2 in 00b855c, keeping only test_iter_weight_files_falls_back_to_root_shards_without_an_index as requested -- the loader is back to the one rule: use the index if there is one, otherwise glob the root. Your points on the fallback missing shapes and the escape check only guarding iter_weight_files are fair; a full version would need both sides, and a broken index failing loudly is the clearer behavior.

Server re-verified on 00b855c: tests/models/test_models_loader.py 13 passed (the kept test included), RAM probe test 1 passed.