`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#

`_init_optimizer_state`	Allocate optimizer state so a DCP load has DTensors to fill.
`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._init_optimizer_state(optimizer: torch.optim.Optimizer) → None#

Allocate optimizer state so a DCP load has DTensors to fill.

Func:: get_optimizer_state_dict initializes empty optimizer state via torch’s _init_optim_state, but that assigns a parameter-dtype gradient. A Megatron-FSDP sharded parameter advertises the FSDP gradient dtype through grad_dtype, which differs from the (main-weight) parameter dtype under mixed precision, and rejects a mismatched gradient. So initialize the state here with a grad_dtype-matched zero gradient; the subsequent load overwrites it. This is a no-op once the state exists (for example after a training step).

TODO: this function becomes unnecessary once torch’s _init_optim_state honors a parameter’s grad_dtype when it allocates the placeholder gradient (torch.zeros_like(param) in torch/distributed/checkpoint/state_dict.py); an upstream issue is being filed.

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.