Prefer safetensors over pickle .bin when splitting model layers (fixes #296)

[Hotragn]

What

Make the layer splitter prefer safetensors over pickle-based pytorch_model.bin, and read any .bin it does load with weights_only=True.

Fixes #296.

Why

split_and_save_layers() resolved the weight map in this order:

if os.path.exists(checkpoint_path / 'pytorch_model.bin.index.json'):
    ...            # <- .bin checked first
elif os.path.exists(checkpoint_path / 'model.safetensors.index.json'):
    ...

Two problems with checking .bin first:

1. Inconsistent with the single-file branch, which already preferred model.safetensors over pytorch_model.bin. So a repo's sharded path and single-file path made opposite choices for the same repo.
2. Unnecessary pickle loading. .bin shards are Python pickles loaded through torch.load(). AirLLM's advertised entry point is an arbitrary Hugging Face repo id (AutoModel.from_pretrained("some/repo")), so when a repo ships both formats — many still do — AirLLM would download and unpickle the .bin shards during the split even though the safe, memory-mapped safetensors were sitting right there. torch.load() documents that unpickling untrusted data can execute arbitrary code.

This aligns the sharded path with the single-file path and with the rest of the ecosystem (transformers itself prefers safetensors when both exist), and it means the common "repo ships both" case no longer touches the pickle path at all.

Changes

• Extract weight-map resolution into a small _resolve_weight_map(checkpoint_path) helper with a clear preference order: sharded safetensors → single-file safetensors → sharded .bin → single-file .bin. .bin is still fully supported as a fallback; nothing is removed.
• Pass weights_only=True to every torch.load() of a checkpoint shard. This restricts the unpickler to tensors and a safe allowlist. setup.py already requires torch>=2.4, where this argument is fully supported (PyTorch made it the default in 2.6), so there's no compatibility cost — and model weight files are plain tensor state dicts, which load fine under it.
• Add air_llm/tests/test_weight_index_resolution.py with offline unit tests.

Testing

New offline unit tests (no network, no GPU) — python -m pytest air_llm/tests/test_weight_index_resolution.py:

• safetensors index preferred when both indexes are present
• falls back to the .bin index when only .bin exists
• single-file safetensors preferred over single-file .bin
• single-file .bin is loaded with weights_only=True
• FileNotFoundError when no weights are present

All 5 pass. I also ran split_and_save_layers() end-to-end on a tiny fake checkpoint that ships both formats and confirmed it now splits from safetensors and produces all per-layer shards, with no change to the resulting layout.