core.distributed.fsdp.src.megatron_fsdp.experimental.schedule_policy#
Scheduling configuration for the minimal Megatron-FSDP path.
Module Contents#
Classes#
Control prefetching and execution scheduling for one FSDP module. |
API#
- class core.distributed.fsdp.src.megatron_fsdp.experimental.schedule_policy.SchedulePolicy#
Control prefetching and execution scheduling for one FSDP module.
.. attribute:: forward_prefetch_size
Forward prefetch budget in parameter elements. None prefetches one successor; 0 disables prefetching.
.. attribute:: backward_prefetch_size
Backward prefetch budget in parameter elements. None prefetches one successor; 0 disables prefetching.
.. attribute:: register_hooks
Register automatic forward/backward execution hooks. Disable when an external scheduler drives the FSDP lifecycle explicitly. The state-dict safety hook remains registered independently.
.. attribute:: grad_accumulation_count
Total parameter-gradient callbacks per reduction window, including separate GraphTasks and delayed weight gradients. None defaults to one callback per owned trainable parameter. The scheduler must declare the complete count before sharding the unit.
.. attribute:: unshard_all_planes_in_backward
Whether backward unshards and prefetches of this module gather MXFP8 rowwise planes along with columnwise ones. Enable it when activations recomputed in backward use this module’s MXFP8 weights without re-entering this module’s forward hooks, e.g. a checkpointed submodule. Otherwise, backward unshards gather only columnwise planes, and the recomputation fails in Transformer Engine.
- forward_prefetch_size: int | None#
None
- backward_prefetch_size: int | None#
None
- register_hooks: bool#
True
- grad_accumulation_count: int | None#
None
- unshard_all_planes_in_backward: bool#
False
- __post_init__() None#
Validate non-negative prefetch budgets.