split_and_save_layers is AirLLM's checkpoint splitting routine; tests verify it preserves every tensor bit-for-bit while hard-linking single-module shards and materializing real files when multiple modules share a shard. The test suite checks that splitting handles both uniform module-per-shard layouts (hard-link eligible) and complex real-world checkpoints with shared shards, including out-of-order residents and packed 4-bit dtypes that must not be corrupted.
air_llm/tests/test_kimi_k3_split.py bypasses airllm/__init__.py at import time to avoid pulling in the MLX backend and full Transformers stack, making the splitter tests runnable on any platform without those heavy dependencies.[1] The test file pins SafetensorModelPersister explicitly so the splitter tests exercise the Linux/CUDA code path regardless of the platform running the tests.[1]
When a Kimi K3 checkpoint has exactly one module per shard, split_and_save_layers hard-links the shard file instead of copying it, so a 1.56 TB checkpoint does not need double the disk space; this is verified by inode equality in test_kimi_k3_split.py.[1] When multiple modules share a single shard — embed, norm, lm_head, and residual norms in K3 — split_and_save_layers must materialise each module as a separate real file rather than hard-linking; the test asserts the inode differs from the source shard and that no foreign tensors leak into the per-module file.[1] For a standard checkpoint where multiple modules share each shard, split_and_save_layers must likewise write real per-module files; TestStandardCheckpointStillSplits verifies that split-file inodes differ from all source-shard inodes.[1]
split_and_save_layers must be bit-for-bit lossless: every tensor in the original checkpoint must appear in the split output with identical dtype and values; test_kimi_k3_split.py verifies this for both the K3 and standard layouts.[1] Packed 4-bit (uint8) tensor dtypes must be preserved verbatim through the split path; test_kimi_k3_split.py asserts that MXFP4 weight_packed tensors remain torch.uint8 after splitting, since any dtype cast would silently corrupt K3 weights.[1] The K3 fake checkpoint in test_kimi_k3_split.py deliberately places the projector shard before the vision-tower shard — out of the order they appear in the resident list — to ensure split_and_save_layers handles out-of-order resident modules without losing tensors.[1]
Sources