core.models.common.combined_1f1b_mfsdp_scheduler

Module Contents

Functions

_make_unshard_forward_hook	Create a forward pre-hook that unshards the owning FSDP module before the submodule forward.
_make_unshard_backward_hook	Create a backward pre-hook that unshards the owning FSDP module before the submodule backward.
_module_post_backward_hook
reshard_fsdp_module	Reshard the FSDP module after fine-grained computation.
register_combined_1f1b_hooks	Install the sub-module hooks required by MCore combined 1F1B.

API

core.models.common.combined_1f1b_mfsdp_scheduler._make_unshard_forward_hook(

owner: megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule,

)

Create a forward pre-hook that unshards the owning FSDP module before the submodule forward.

core.models.common.combined_1f1b_mfsdp_scheduler._make_unshard_backward_hook(

owner: megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule,

)

Create a backward pre-hook that unshards the owning FSDP module before the submodule backward.

core.models.common.combined_1f1b_mfsdp_scheduler._module_post_backward_hook(

module: megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule,

) → None

core.models.common.combined_1f1b_mfsdp_scheduler.reshard_fsdp_module(

module: megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule,

) → None

Reshard the FSDP module after fine-grained computation.

core.models.common.combined_1f1b_mfsdp_scheduler.register_combined_1f1b_hooks(

module: megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module.FsdpModule,

) → None

Install the sub-module hooks required by MCore combined 1F1B.