ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsDeepseek V41

nemo_automodel.components.models.deepseek_v41.fsdp

View as Markdown

FSDP policy for V4.1’s shared V4 vision tower.

Module Contents

Functions

NameDescription
fully_shard_deepseek_v41Reuse V4’s vision norm policy while preserving ordinary language FSDP.

API

nemo_automodel.components.models.deepseek_v41.fsdp.fully_shard_deepseek_v41(
    module: torch.nn.Module,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None = None,
    reshard_after_forward: bool | int | None = None,
    ignored_params: set[torch.nn.Parameter] | None = None
) -> torch.nn.Module

Reuse V4’s vision norm policy while preserving ordinary language FSDP.

Parameters:

module
nn.Module

Module to shard in place; vision towers/blocks retain FP32 norms.

mesh
DeviceMesh

Runtime FSDP device mesh.

mp_policy
MixedPrecisionPolicy

Mixed-precision policy for the module’s non-norm parameters.

offload_policy
OffloadPolicy | NoneDefaults to None

Parameter/gradient offload policy.

reshard_after_forward
bool | int | NoneDefaults to None

PyTorch FSDP resharding setting.

ignored_params
set[nn.Parameter] | NoneDefaults to None

Parameter tensors of arbitrary shapes already managed outside this FSDP unit, retaining their existing DTensor placements.

Returns: nn.Module

The input module with FSDP applied.