mixle.utils.parallel.dcp_checkpoint module

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

Note

This is a provisional API with experimental multi-node behavior. Complete-state and failure semantics are covered on CPU; resharding and throughput at multi-node GPU scale still require retained hardware receipts.

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: implemented against the torch 2.4+ DCP and distributed-state-dict APIs, but not verified without retained hardware receipts.

class AsyncTrainingCheckpoint(future, path, generation, rank, world_size, payload, _complete=False)

Bases: object

DCP asynchronous save plus atomic sidecar/manifest finalization.

Parameters:

• future (Any)

• path (Path)

• generation (Path)

• rank (int)

• world_size (int)

• payload (dict[str, Any])

• _complete (bool)

async_save_training_state(module, optimizer, path, *, step, scheduler=None, scaler=None, loader_state=None, parallel_plan=None, typed_scheduler_state=None, extra=None)

Start DCP’s native async save; wait() surfaces write failures.

Parameters:

• module (Any)

• optimizer (Any)

• path (str)

• step (int)

• scheduler (Any)

• scaler (Any)

• loader_state (Any)

• parallel_plan (Any)

• typed_scheduler_state (Any)

• extra (dict[str, Any] | None)

Return type:

AsyncTrainingCheckpoint

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

load_training_metadata(path, *, rank_local_transform=None)

Decode metadata without ever assigning an old rank’s local state to a new rank.

Parameters:

• path (str)

• rank_local_transform (Callable[[tuple[dict[str, Any], ...], int, int, int], dict[str, Any]] | None)

Return type:

dict[str, Any]

load_training_state(module, optimizer, path, *, scheduler=None, scaler=None, restore_rng=True, allow_world_size_change=False, rank_local_transform=None)

Restore complete training state and return its non-tensor metadata.

Parameters:

• module (Any)

• optimizer (Any)

• path (str)

• scheduler (Any)

• scaler (Any)

• restore_rng (bool)

• allow_world_size_change (bool)

• rank_local_transform (Callable[[tuple[dict[str, Any], ...], int, int, int], dict[str, Any]] | None)

Return type:

dict[str, Any]

save_frozen_training_state(model_state, optimizer_state, path, *, step, loader_state=None, parallel_plan=None, typed_scheduler_state=None, extra=None)

Transactionally save already-frozen model and optimizer state dictionaries.

Parameters:

• model_state (dict[str, Any])

• optimizer_state (dict[str, Any])

• path (str)

• step (int)

• loader_state (Any)

• parallel_plan (Any)

• typed_scheduler_state (Any)

• extra (dict[str, Any] | None)

Return type:

None

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

save_training_state(module, optimizer, path, *, step, scheduler=None, scaler=None, loader_state=None, parallel_plan=None, typed_scheduler_state=None, extra=None)

Synchronously save complete, reshardable training state.

Parameters:

• module (Any)

• optimizer (Any)

• path (str)

• step (int)

• scheduler (Any)

• scaler (Any)

• loader_state (Any)

• parallel_plan (Any)

• typed_scheduler_state (Any)

• extra (dict[str, Any] | None)

Return type:

None