nemo_automodel.shared.multimodal_fsdp

Shared helpers for FSDP treatment of multimodal submodules.

Module Contents

Functions

Name	Description
`_container_items`	-
`_is_module_container`	-
`_shard_layer_containers_recursively`	-
`ignored_params_for_root`	Return ignored parameters that are owned by an FSDP root.
`is_multimodal_module_name`	Return True when `name` identifies a known multimodal tower/projector.
`iter_multimodal_modules`	Yield maximal multimodal submodules by qualified name.
`module_is_fully_frozen`	Return whether `module` owns parameters and none require gradients.
`module_parameters`	Return the module’s recursive parameters.
`normalize_frozen_multimodal_sharding`	Validate and normalize the frozen multimodal FSDP policy.
`shard_multimodal_module`	Shard a multimodal module at layer-container granularity when possible.

Data

FrozenMultimodalSharding

MULTIMODAL_MODULE_NAMES

MULTIMODAL_PROJECTOR_NAMES

MULTIMODAL_TOWER_NAMES

VALID_FROZEN_MULTIMODAL_SHARDING

API

nemo_automodel.shared.multimodal_fsdp._container_items(
    module: torch.nn.Module
) -> list[tuple[object, torch.nn.Module]]

nemo_automodel.shared.multimodal_fsdp._is_module_container(
    module: torch.nn.Module
) -> bool

nemo_automodel.shared.multimodal_fsdp._shard_layer_containers_recursively(
    module: torch.nn.Module,
    shard_module: collections.abc.Callable[[nn.Module], object]
) -> bool

nemo_automodel.shared.multimodal_fsdp.ignored_params_for_root(
    root: torch.nn.Module,
    ignored_params: set[torch.nn.Parameter]
) -> set[torch.nn.Parameter] | None

Return ignored parameters that are owned by an FSDP root.

Parameters:

root

nn.Module

Module that will become the FSDP root.

ignored_params

set[nn.Parameter]

Parameters of arbitrary shapes that should remain replicated rather than becoming part of the FSDP root.

Returns: set[nn.Parameter] | None

The subset of ignored_params owned by root, preserving each

nemo_automodel.shared.multimodal_fsdp.is_multimodal_module_name(
    name: str
) -> bool

Return True when name identifies a known multimodal tower/projector.

nemo_automodel.shared.multimodal_fsdp.iter_multimodal_modules(
    model: torch.nn.Module
) -> collections.abc.Iterator[tuple[str, torch.nn.Module]]

Yield maximal multimodal submodules by qualified name.

named_modules is depth-first pre-order, so a parent is always seen before its descendants; already-selected prefixes are skipped to keep each tower/projector maximal (e.g. vision_tower rather than the vision_tower.vision_model nested inside it).

The attribute-scan fallback only serves tests/unit_tests/moe model doubles, which are plain classes rather than nn.Module subclasses. It uses different (non-maximal, top-two-levels-only) selection than the real path, so tests that exercise it are not testing production behavior. Removing it requires converting ~30 doubles across that file to real modules; tracked separately rather than widening this change.

nemo_automodel.shared.multimodal_fsdp.module_is_fully_frozen(
    module: torch.nn.Module
) -> bool

Return whether module owns parameters and none require gradients.

nemo_automodel.shared.multimodal_fsdp.module_parameters(
    module: torch.nn.Module
) -> list[torch.nn.Parameter]

Return the module’s recursive parameters.

nemo_automodel.shared.multimodal_fsdp.normalize_frozen_multimodal_sharding(
    value: str
) -> nemo_automodel.shared.multimodal_fsdp.FrozenMultimodalSharding

Validate and normalize the frozen multimodal FSDP policy.

nemo_automodel.shared.multimodal_fsdp.shard_multimodal_module(
    module: torch.nn.Module,
    shard_module: collections.abc.Callable[[nn.Module], object]
) -> None

Shard a multimodal module at layer-container granularity when possible.

nemo_automodel.shared.multimodal_fsdp.FrozenMultimodalSharding = Literal['root', 'per_layer', 'replicate']

nemo_automodel.shared.multimodal_fsdp.MULTIMODAL_MODULE_NAMES = MULTIMODAL_TOWER_NAMES + MULTIMODAL_PROJECTOR_NAMES

nemo_automodel.shared.multimodal_fsdp.MULTIMODAL_PROJECTOR_NAMES = ('embed_vision', 'embed_audio', 'mm_projector', 'multi_modal_projector', 'multim...

nemo_automodel.shared.multimodal_fsdp.MULTIMODAL_TOWER_NAMES = ('visual', 'vision_tower', 'vision_model', 'vit_model', 'image_encoder', 'vision...

nemo_automodel.shared.multimodal_fsdp.VALID_FROZEN_MULTIMODAL_SHARDING: tuple[FrozenMultimodalSharding, ...] = ('root', 'per_layer', 'replicate')