core.distributed.fsdp.src.megatron_fsdp.experimental.optimizer

Optimizer adapter for the minimal Megatron-FSDP path.

Module Contents

Functions

init_optimizer_state	Allocate optimizer state so a DCP load has DTensors to fill.
fully_shard_optimizer	Attach FSDP-aware step hooks to an optimizer instance.

Data

__all__

API

core.distributed.fsdp.src.megatron_fsdp.experimental.optimizer.__all__

[‘fully_shard_optimizer’, ‘init_optimizer_state’]

core.distributed.fsdp.src.megatron_fsdp.experimental.optimizer.init_optimizer_state(optimizer: torch.optim.Optimizer) → None

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

Func:

~torch.distributed.checkpoint.state_dict.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.

Parameters:

optimizer – Optimizer whose (possibly empty) state should be materialized.

core.distributed.fsdp.src.megatron_fsdp.experimental.optimizer.fully_shard_optimizer(

optimizer: torch.optim.Optimizer,

*,

precision_aware: bool = False,

) → None

Attach FSDP-aware step hooks to an optimizer instance.

The adapted optimizer preserves its existing parameter groups, temporarily casts gradients around optimizer steps for FSDP sharded parameters whose data dtype differs from their grad dtype unless the optimizer is precision aware, and refreshes compute weights after each optimizer step.

Alternatives considered: - Monkey-patching optimizer methods directly on the instance. This is more invasive and harder to compose than hooks. - Generating an FSDP-specific subclass per torch.optim.Optimizer. This adds extra class-generation machinery, but would let us instrument zero_grad and __init__ as well as step if needed. - Casting from main_grad.dtype to main_weight.dtype after the last microbatch and casting back before the first microbatch. This should be done from a root post-backward callback if needed later, so users do not need to call fully_shard_optimizer on an existing torch.optim.Optimizer. - Letting the user set main_weight and main_grad to the same dtype. This is enough for an FSDP2 drop-in replacement path and lets optimizers stay unaware of FSDP precision handling.

Parameters:

• optimizer – Optimizer instance to adapt in place.

• precision_aware – Whether the optimizer accepts FSDP’s mixed-precision gradients without temporary casting.