mixle.utils.parallel.dcp_checkpoint module

Sharded distributed checkpoints via torch.distributed.checkpoint (DCP) – replaces pickle-broadcast at scale.

The gather-to-root + pickle-broadcast that TorchRunEncodedData uses to move a model cannot save a model that does not fit (and is not folded on) one rank. DCP saves each rank’s shard of the (FSDP2-sharded) model + optimizer state in parallel to a checkpoint directory, and loads it back sharded – the standard frontier checkpoint, and the resume hook for StreamingTokenEncodedData.

CUDA / multi-GPU path: correct per the torch 2.4+ DCP + distributed-state-dict APIs, exercised on the cluster.

save_sharded(module, optimizer, path)

Save a sharded (model, optimizer) checkpoint to path – every rank writes its own shard in parallel.

Parameters:

• module (Any)

• optimizer (Any)

• path (str)

Return type:

None

load_sharded(module, optimizer, path)

Load a sharded checkpoint from path into module + optimizer in place (resumable training).

Parameters:

• module (Any)

• optimizer (Any)

• path (str)

Return type:

None