core.distributed.fsdp.src.megatron_fsdp.experimental.runtime_policy#
Shared runtime configuration for Megatron-FSDP contexts.
Module Contents#
Classes#
Configure communication and gradient synchronization for an FSDP context. |
API#
- class core.distributed.fsdp.src.megatron_fsdp.experimental.runtime_policy.RuntimePolicy#
Configure communication and gradient synchronization for an FSDP context.
.. attribute:: use_symmetric_memory
Allocate communication staging buffers from PyTorch’s NCCL symmetric-memory pool.
.. attribute:: overlap_allgather_reduction
Use separate all-gather and reduction streams. Disable to reduce peak transient memory. See https://github.com/NVIDIA/Megatron-LM/issues/6471.
.. attribute:: overlap_outer_communication
Run the last-microbatch DP-outer reduction on a separate stream, independently of overlap_allgather_reduction.
.. attribute:: caller_managed_grad_sync
Leave gradient synchronization to the caller while retaining automatic weight cleanup. The caller must call finish_grad_sync() after all backward work, including delayed weight gradients, and before reading or modifying gradients.
- use_symmetric_memory: bool#
False
- overlap_allgather_reduction: bool#
True
- overlap_outer_communication: bool#
False
- caller_managed_grad_sync: bool#
False