`core.distributed.fsdp.src.megatron_fsdp.experimental.checkpoint`#

PyTorch Distributed Checkpoint (DCP) save/load for the experimental Megatron-FSDP path.

After :func:fully_shard, a module’s parameters rest as DTensor views over the optimizer (main_weight) buffers, and the optimizer’s exp_avg/exp_avg_sq states are DTensor s on the same device mesh. The standard DCP state-dict helpers (:func:torch.distributed.checkpoint.state_dict.get_model_state_dict /

func:: ~torch.distributed.checkpoint.state_dict.get_optimizer_state_dict) expose those as FQN-keyed DTensors and initialize the (empty) optimizer state on load, so we do not reimplement that here.

The one Megatron-FSDP-specific step is :func:attach_uneven_dtensor_metadata, which describes each parameter’s true position inside its packed parameter-group buffer. Without it the default planner assumes canonical Shard(0) offsets and silently corrupts the checkpoint.

Module Contents#

Functions#

`save_checkpoint`	Save a `fully_shard`-wrapped model and its optimizer as a DCP checkpoint.
`load_checkpoint`	Load a DCP checkpoint into a `fully_shard`-wrapped model and its optimizer.

Data#

__all__

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.checkpoint.__all__#: [‘save_checkpoint’, ‘load_checkpoint’]

core.distributed.fsdp.src.megatron_fsdp.experimental.checkpoint.save_checkpoint( model: torch.nn.Module, optimizer: torch.optim.Optimizer, checkpoint_dir: str | os.PathLike, ) → None#

Save a fully_shard-wrapped model and its optimizer as a DCP checkpoint.

Parameters:

model – A module tree that has been sharded with :func:fully_shard.
optimizer – Optimizer stepping the sharded parameters.
checkpoint_dir – Destination directory for the DCP checkpoint.

core.distributed.fsdp.src.megatron_fsdp.experimental.checkpoint.load_checkpoint( model: torch.nn.Module, optimizer: torch.optim.Optimizer, checkpoint_dir: str | os.PathLike, *, sync_model_weights: bool = True, ) → None#

Load a DCP checkpoint into a fully_shard-wrapped model and its optimizer.

The model and optimizer must already be sharded with the same layout used at save time (the same module structure and mesh); DCP reshards the on-disk data to this rank’s shards.

Func:

~torch.distributed.checkpoint.state_dict.get_optimizer_state_dict initializes the (empty) optimizer state so DCP has DTensors to load into in place, and the set_* helpers reinstall the loaded state.