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

Shared runtime state for Megatron-FSDP modules.

Module Contents#

Classes#

FsdpContext

Runtime stream and prefetch state shared by FSDP roots constructed together.

Functions#

`current_fully_shard_context`	Return the innermost active `fully_shard_context`, or `None`.
`_collect_backward_order`	Collect one root’s static backward prefetch order.
`_collect_fsdp_children`	Collect the nearest registered FSDP descendants of `module`.

Data#

_FSDP_CONTEXT

API#

core.distributed.fsdp.src.megatron_fsdp.experimental.context._FSDP_CONTEXT#: ‘(…)’

class core.distributed.fsdp.src.megatron_fsdp.experimental.context.FsdpContext( device: torch.device | None = None, use_symmetric_memory: bool = False, unify_communication_stream: bool = False, parameter_to_owner: dict[torch.nn.Parameter, int] | None = None, caller_managed_grad_sync: bool = False, )#

Runtime stream and prefetch state shared by FSDP roots constructed together.

Initialization

Create rank-local runtime state for FSDP modules on device.

Parameters:

device – CUDA device on which this context schedules communication. Defaults to the current CUDA device.
use_symmetric_memory – Whether modules constructed in this context allocate communication staging buffers from PyTorch’s NCCL symmetric-memory pool.
unify_communication_stream – Whether all-gathers and reduce-scatters share one communication stream to reduce peak transient memory. See https://github.com/NVIDIA/Megatron-LM/issues/6471.
parameter_to_owner – Construction-time owner assignments for TensorAtomic parameters, keyed by the original parameters before sharding. Owners are ranks in each parameter group’s 1-D data-parallel mesh and must agree across that mesh. Every TensorAtomic parameter needs an entry; other entries are ignored. Tensors are packed by owner without changing logical parameter order.
caller_managed_grad_sync – Disable the automatic autograd completion callback, allowing delayed weight gradients or custom backward schedules. The caller must call finish_grad_sync() after all backward work and before reading or modifying gradients.

allgather_stream: torch.cuda.Stream#: None

reduce_scatter_stream: torch.cuda.Stream#: None

is_last_microbatch: bool#: None

use_symmetric_memory: bool#: None

unify_communication_stream: bool#: None

caller_managed_grad_sync: bool#

None

Leave gradient synchronization to the caller instead of an autograd callback.

The caller must call finish_grad_sync() after all gradient producers, including delayed weight-gradient computation, and before consuming gradients.

forward_order: core.distributed.fsdp.src.megatron_fsdp.experimental.indexed_order.IndexedOrder[torch.nn.Module]#: None

backward_order: core.distributed.fsdp.src.megatron_fsdp.experimental.indexed_order.IndexedOrder[torch.nn.Module]#: None

_roots: core.distributed.fsdp.src.megatron_fsdp.experimental.indexed_order.IndexedOrder[torch.nn.Module]#: None

_module_names: weakref.WeakKeyDictionary[torch.nn.Module, str]#: None

_post_backward_hook_registered: bool#: None

__enter__() → core.distributed.fsdp.src.megatron_fsdp.experimental.context.FsdpContext#: Activate this context for FSDP module construction.

__exit__(exc_type: type[BaseException] | None, *_: object) → None#: Finalize successful construction and always clear the active scope.

register_module(module: torch.nn.Module) → None#: Register a module after its FSDP initialization completes in this context.

finalize() → None#: Finalize roots, names, and cross-root prefetch orders.

is_root(module: torch.nn.Module) → bool#: Return whether module is an outermost FSDP module in this context.

module_name(module: torch.nn.Module) → str#: Return module’s name relative to its FSDP root.

ensure_finalized() → None#: Raise if construction has not completed for this context.

current_stream() → torch.cuda.Stream#: Current stream on this context’s device.

validate_grad_sync() → None#: Require a caller-owned wait or a pending autograd completion callback.

finish_grad_sync() → None#

Order current-stream consumers after all gradient reductions submitted so far.

Call after all backward work, including delayed weight-gradient computation, and before reading or modifying gradients. This enqueues a stream dependency; it does not block the CPU or wait for reductions that have not yet been launched.

post_backward() → None#: Order current-stream consumers after this backward’s gradient reductions.

register_post_backward_hook() → None#

Multiple FSDP roots can share this context. Waiting for the reduce-scatter stream in each root’s post_backward() would prevent one root’s backward compute from overlapping another root’s gradient reductions. Wait once at context-level autograd completion instead.

core.distributed.fsdp.src.megatron_fsdp.experimental.context.current_fully_shard_context() → core.distributed.fsdp.src.megatron_fsdp.experimental.context.FsdpContext | None#

Return the innermost active fully_shard_context, or None.

Read-only counterpart of :func:fully_shard_context: it never creates, joins, or finalizes a context, and returns None whenever no fully_shard_context scope is active. Callers that must share one context – for example per-chunk wrappers built by a single wrap call – use it to join the caller’s ambient context instead of opening a second one.

core.distributed.fsdp.src.megatron_fsdp.experimental.context._collect_backward_order( module: torch.nn.Module, registered_modules: set[torch.nn.Module], order: core.distributed.fsdp.src.megatron_fsdp.experimental.indexed_order.IndexedOrder[torch.nn.Module], ) → None#: Collect one root’s static backward prefetch order.

core.distributed.fsdp.src.megatron_fsdp.experimental.context._collect_fsdp_children( module: torch.nn.Module, registered_modules: set[torch.nn.Module], children: set[torch.nn.Module], ) → None#: Collect the nearest registered FSDP descendants of module.

core.distributed.fsdp.src.megatron_fsdp.experimental.context#