nemo_automodel.components.moe.fsdp_mixin

Module Contents

Classes

Name	Description
`MoEFSDPSyncMixin`	Mixin for managing FSDP synchronization state during MoE model training.

Functions

Name	Description
`_configure_fsdp_module`	-
`_disable_fsdp_for_moe_module`	-
`_iter_fsdp_module_candidates`	-
`_iter_fsdp_modules`	Yield every FSDP unit whose backward state this mixin must drive, deduplicated.
`_reduce_fsdp_for_moe_module`	Let FSDP’s hooks reduce-scatter this micro-batch’s gradients; reshard and finalize only on the last backward.
`_run_post_backward_for_moe_module`	Run FSDP’s post-backward (reduce-scatter) manually for every FSDP unit of a MoE module.
`_run_post_backward_hooks`	Run post-backward for every FSDP state and return the root final callback.
`patched_backward_maybe_with_nosync`	Whether using PP with FSDP or DDP, there are some runtime differences between the last backward step and the

Data

logger

API

class nemo_automodel.components.moe.fsdp_mixin.MoEFSDPSyncMixin()

Mixin for managing FSDP synchronization state during MoE model training.

Controls gradient sync and resharding for FSDP-wrapped modules to optimize performance during gradient accumulation steps.

Usage differs based on parallelism strategy:

Without pipeline parallelism (PP): prepare_for_grad_accumulation() defers sync and resharding at the start of gradient accumulation. prepare_for_final_backward() enables sync and resharding before the last backward pass. FSDP’s autograd hooks automatically handle post-backward synchronization and resharding.
With pipeline parallelism (PP): FSDP state management is handled by patching _PipelineStageBase.backward_maybe_with_nosync (see patched_backward_maybe_with_nosync below). By default (distributed.defer_fsdp_grad_sync: true) the patch disables sync/resharding for all backwards except the last one before the optimizer step, where it manually triggers post-backward hooks and resharding. With defer_fsdp_grad_sync: false every micro-batch reduce-scatters its own gradients (after its full or weight backward); the input backward of a split schedule never reduces or reshards, and resharding/finalization is deferred to the last backward so that pending split weight backwards still find the unsharded parameters.

_fsdp_shard_across_blocks

tuple[str, ...] = ()

nemo_automodel.components.moe.fsdp_mixin.MoEFSDPSyncMixin.prepare_for_final_backward(
    pp_enabled: bool = False
) -> None

Enable gradient sync and resharding for the final backward pass.

Parameters:

pp_enabled

boolDefaults to False

Whether pipeline parallelism is enabled.

When PP is enabled, FSDP state management is handled by the patched _PipelineStageBase.backward_maybe_with_nosync method. This method only applies optimizations for non-PP cases.

nemo_automodel.components.moe.fsdp_mixin.MoEFSDPSyncMixin.prepare_for_grad_accumulation(
    pp_enabled: bool = False
) -> None

Prepare FSDP states before starting gradient accumulation.

Parameters:

pp_enabled

boolDefaults to False

Whether pipeline parallelism is enabled.

When PP is enabled, FSDP state management is handled by the patched _PipelineStageBase.backward_maybe_with_nosync method. This method only applies optimizations for non-PP cases.

nemo_automodel.components.moe.fsdp_mixin._configure_fsdp_module(
    fsdp_module: torch.distributed.fsdp.FSDPModule,
    is_last_backward: bool,
    reshard_after_backward: bool,
    requires_gradient_sync: bool
) -> None

nemo_automodel.components.moe.fsdp_mixin._disable_fsdp_for_moe_module(
    module: torch.nn.Module
) -> None

nemo_automodel.components.moe.fsdp_mixin._iter_fsdp_module_candidates(
    module: torch.nn.Module
) -> typing.Iterator[torch.distributed.fsdp.FSDPModule]

nemo_automodel.components.moe.fsdp_mixin._iter_fsdp_modules(
    module: torch.nn.Module
) -> typing.Iterator[torch.distributed.fsdp.FSDPModule]

Yield every FSDP unit whose backward state this mixin must drive, deduplicated.

A unit that is not yielded here never receives set_is_last_backward / set_reshard_after_backward / set_requires_gradient_sync, so it drops out of the gradient-accumulation state machine and, under pipeline parallelism, out of patched_backward_maybe_with_nosync.

nemo_automodel.components.moe.fsdp_mixin._reduce_fsdp_for_moe_module(
    module: torch.nn.Module,
    last_backward: bool
) -> None

Let FSDP’s hooks reduce-scatter this micro-batch’s gradients; reshard and finalize only on the last backward.

Used for a full backward under distributed.defer_fsdp_grad_sync: false. A split weight backward for an earlier micro-batch may still be pending (mixed schedules such as F0 I0 F1 B1 W0) and looks the parameters up by their unsharded objects, so the parameters must stay unsharded until the actual last backward.

nemo_automodel.components.moe.fsdp_mixin._run_post_backward_for_moe_module(
    module: torch.nn.Module,
    last_backward: bool = True
) -> None

Run FSDP’s post-backward (reduce-scatter) manually for every FSDP unit of a MoE module.

With last_backward (the default) the units are also resharded and the root final callback finalizes the step; otherwise the parameters stay unsharded for the backwards still to come and only the gradients are reduced.

nemo_automodel.components.moe.fsdp_mixin._run_post_backward_hooks(
    fsdp_module: torch.distributed.fsdp.FSDPModule
) -> typing.Callable

Run post-backward for every FSDP state and return the root final callback.

Both the per-state post_backward() calls and the returned callback tolerate FSDP modules that never ran forward in this step. With FusedLinearCrossEntropy under pipeline parallelism the lm_head weight is consumed inside the loss function, so its FSDP module never runs forward and its FSDPCommContext is never lazily initialized; post_backward() then raises AttributeError ('FSDPCommContext' object has no attribute 'post_forward_order'). A never-forwarded group has no gradients to reduce, so it is safe to skip.

nemo_automodel.components.moe.fsdp_mixin.patched_backward_maybe_with_nosync(
    self,
    backward_type,
    bwd_kwargs: dict,
    last_backward: bool = False
) -> tuple[tuple[torch.Tensor | None, ...], list[dict[str, typing.Any]] | None]

Whether using PP with FSDP or DDP, there are some runtime differences between the last backward step and the other steps. Namely, we need to accumulate gradients on previous steps and reduce them on the last step, but there are additional state-variables and performance considerations depending on the data parallelism used. This helper should adapt any pipeline parallel schedule to work with common/supported data parallel libraries.

nemo_automodel.components.moe.fsdp_mixin.logger = logging.getLogger(__name__)