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

Shared runtime state for Megatron-FSDP modules.

Module Contents#

Classes#

`PrefetchOrder`	How to choose the module order used for prefetching.
`FsdpContext`	Runtime stream and prefetch state shared by FSDP roots constructed together.

Functions#

`_collect_backward_order`	Collect one root’s static backward prefetch order.
`_collect_fsdp_children`	Collect the nearest registered FSDP descendants of `module`.

API#

class core.distributed.fsdp.src.megatron_fsdp.experimental.context.PrefetchOrder(*args, **kwds)#

Bases: enum.Enum

How to choose the module order used for prefetching.

Initialization

REGISTRATION#: ‘auto(…)’

RECORDED#: ‘auto(…)’

class core.distributed.fsdp.src.megatron_fsdp.experimental.context.FsdpContext( device: torch.device | None = None, runtime_policy: core.distributed.fsdp.src.megatron_fsdp.experimental.runtime_policy.RuntimePolicy = RuntimePolicy(), parameter_to_owner: dict[torch.nn.Parameter, int] | None = None, prefetch_order: core.distributed.fsdp.src.megatron_fsdp.experimental.context.PrefetchOrder = PrefetchOrder.REGISTRATION, )#

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.
runtime_policy – Communication allocation, stream, and gradient-synchronization settings shared by all modules constructed in this context.
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.
prefetch_order – REGISTRATION expects a forward/backward pass matching module registration order. RECORDED gathers parameters on demand and records module calls until the first post_backward(). Subsequent passes must repeat the complete sequence consistently across data-parallel ranks. Recording currently requires the ordinary module hooks; combined 1F1B bypasses them. Discard the context if the recorded pass raises.

allgather_stream: torch.cuda.Stream#: None

inner_reduction_stream: torch.cuda.Stream#: None

outer_reduction_stream: torch.cuda.Stream#: None

is_last_microbatch: bool#: None

use_symmetric_memory: bool#: None

caller_managed_grad_sync: bool#: None

execution_order: core.distributed.fsdp.src.megatron_fsdp.experimental.execution_order.ExecutionOrder#: None

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

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

_post_backward_hook_registered: bool#: None

property is_finalized: bool#: Whether construction has completed for this context.

property is_recording_execution_order: bool#: Whether module calls are being recorded for subsequent prefetching.

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#

Wait on submitted gradient reductions.

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#: Finish backward cleanup and automatically synchronize the last microbatch.

_run_delayed_post_backward_callbacks() → None#: Run deferred module completion callbacks once.

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._collect_backward_order( module: torch.nn.Module, registered_modules: set[torch.nn.Module], order: core.distributed.fsdp.src.megatron_fsdp.experimental.execution_order.ExecutionOrder, ) → 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#

Module Contents#

Classes#

Functions#

API#

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