Description
make_shards() flushes the current shard whenever the next tensor would exceed max_file_size_gb, even when the current shard is still empty. If the first model tensor is larger than the 5 GiB default, the result starts with {}:
from types import SimpleNamespace
from mlx_lm.utils import make_shards
large_weight = SimpleNamespace(nbytes=(5 << 30) + 1)
shards = make_shards({"model.embed_tokens.weight": large_weight})
assert shards == [{"model.embed_tokens.weight": large_weight}]
# Current result: [{}, {"model.embed_tokens.weight": large_weight}]Impact
A single tensor cannot be split by this helper, so an oversized tensor correctly has to live in an oversized shard. The extra empty shard has no useful representation. save_model() nevertheless counts and writes it, producing an unreferenced model-00001-of-00002.safetensors and placing the only real tensor in part 2. mx.save_safetensors() accepts the empty mapping and writes a 41-byte file, so the error is silent: the saved model has a misleading, noncanonical shard set and consumers that enumerate shard files process a shard containing no weights.
This is reachable for large-vocabulary embedding or output matrices that individually exceed the configured shard limit.
Expected behavior
Only flush a shard that already contains at least one weight. For an oversized first tensor, return one oversized, non-empty shard. Normal multi-tensor shard boundaries should remain unchanged.
Proposed fix
Guard the flush with if shard and shard_size + v.nbytes > max_file_size_bytes, plus a regression for an oversized first tensor.
I searched open and closed issues and pull requests for make_shards, empty/oversized shards, and single large tensors. I found no existing report or implementation; #1633 concerns eager evaluation during saving, not shard construction. I am preparing a focused patch.
Description
make_shards()flushes the current shard whenever the next tensor would exceedmax_file_size_gb, even when the current shard is still empty. If the first model tensor is larger than the 5 GiB default, the result starts with{}:Impact
A single tensor cannot be split by this helper, so an oversized tensor correctly has to live in an oversized shard. The extra empty shard has no useful representation.
save_model()nevertheless counts and writes it, producing an unreferencedmodel-00001-of-00002.safetensorsand placing the only real tensor in part 2.mx.save_safetensors()accepts the empty mapping and writes a 41-byte file, so the error is silent: the saved model has a misleading, noncanonical shard set and consumers that enumerate shard files process a shard containing no weights.This is reachable for large-vocabulary embedding or output matrices that individually exceed the configured shard limit.
Expected behavior
Only flush a shard that already contains at least one weight. For an oversized first tensor, return one oversized, non-empty shard. Normal multi-tensor shard boundaries should remain unchanged.
Proposed fix
Guard the flush with
if shard and shard_size + v.nbytes > max_file_size_bytes, plus a regression for an oversized first tensor.I searched open and closed issues and pull requests for
make_shards, empty/oversized shards, and single large tensors. I found no existing report or implementation; #1633 concerns eager evaluation during saving, not shard construction. I am preparing a focused patch.