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

Optimizer adapter for the minimal Megatron-FSDP path.

Module Contents

Functions

fully_shard_optimizer

Attach FSDP-aware step hooks to an optimizer instance.

API

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.