save_sharded

→None

save_sharded(obj: object, path: str | bytes, shard_size_mb: float = 1024.0, pickle_protocol: int = 4)

source

Save obj as a sharded, self-describing checkpoint directory.

Splits a state dict across multiple shard files capped at shard_size_mb MiB each and emits a JSON index listing which keys live in which shard. Useful for very large models where a single monolithic file is impractical to host, partial-load, or download-resume.

Parameters

objobject

Object to save. dict / OrderedDict state dicts are sharded; other types are written as a single shard for compatibility with load_sharded.

pathstr or bytes

Destination directory. Created if it does not exist.

shard_size_mbfloat= 1024.0

Soft cap on per-shard size in MiB. A tensor exceeding the cap gets its own shard. Default 1024.0.

pickle_protocolint= 4

Pickle protocol forwarded to every per-shard save. Default 4.

Notes

The packing algorithm is a first-fit pass that never splits a single tensor across shards — preserving the property that any given key resolves to exactly one file. The index records the optional _metadata attribute when present on a state dict.

Examples

>>> import lucid
>>> sd = {"w": lucid.randn(4096, 4096)}
>>> lucid.serialization.save_sharded(sd, "ckpt_dir", shard_size_mb=64)