core.distributed.fsdp.fsdp_unit_collector

Select MCore modules for MFSDP v2 sharding.

Module Contents

Classes

FsdpUnitCollector

Collect selected FSDP units in wrapping order, children before parents.

API

class core.distributed.fsdp.fsdp_unit_collector.FsdpUnitCollector(

fsdp_unit_module_types: tuple[type[torch.nn.Module], ...],

)

Collect selected FSDP units in wrapping order, children before parents.

Initialization

collect(module: torch.nn.Module) → list[torch.nn.Module]

Return selected modules in reverse pre-order, including the root last.

_visit(module: torch.nn.Module, under_shortcut: bool) → None

_should_shard(module: torch.nn.Module, under_shortcut: bool) → bool

Whether this module should be an independent FSDP unit.