nemo_automodel.components.models.deepseek_v4.fsdp

Module Contents

Classes

Name	Description
`DeepseekV4FSDPKwargs`	Optional `fully_shard` arguments forwarded by `fully_shard_deepseek_v4`.

Functions

Name	Description
`_attach_hca_param_sync_group`	-
`_floating_param_dtypes`	-
`_fp32_module_fsdp_kwargs`	-
`_fp32_mp_policy`	-
`_fsdp_kwargs_for_module`	-
`_fully_shard_once`	-
`_has_fsdp_state`	-
`_hca_param_sync_group_from_1d_mesh`	Return the 1D PyTorch FSDP2 group used for HCA graph alignment.
`_is_deepseek_v4_module`	-
`_iter_dsv4_fp32_modules`	-
`_matches_suffix`	-
`_module_config_model_type`	-
`fully_shard_deepseek_v4`	Apply FSDP2 to DeepSeek-V4 without mixing fp32 and bf16 params in one unit.

Data

_DSV4_CLASS_NAMES

_DSV4_FP32_MODULE_SUFFIXES

_DSV4_RETURNED_FP32_PARAMETER_CLASS_NAMES

_DSV4_SELF_CASTING_FP32_MODULE_CLASS_NAMES

API

class nemo_automodel.components.models.deepseek_v4.fsdp.DeepseekV4FSDPKwargs

Bases: typing.TypedDict

Optional fully_shard arguments forwarded by fully_shard_deepseek_v4.

ignored_params

set[Parameter] | None

reshard_after_forward

bool | int | None

nemo_automodel.components.models.deepseek_v4.fsdp._attach_hca_param_sync_group(
    module: torch.nn.Module,
    mesh
) -> None

nemo_automodel.components.models.deepseek_v4.fsdp._floating_param_dtypes(
    module: torch.nn.Module
) -> set[torch.dtype]

nemo_automodel.components.models.deepseek_v4.fsdp._fp32_module_fsdp_kwargs(
    module: torch.nn.Module,
    fsdp_kwargs: dict
) -> dict

nemo_automodel.components.models.deepseek_v4.fsdp._fp32_mp_policy(
    mp_policy,
    preserve_activation_dtype: bool = False
)

nemo_automodel.components.models.deepseek_v4.fsdp._fsdp_kwargs_for_module(
    module: torch.nn.Module,
    fsdp_kwargs: dict
) -> dict

nemo_automodel.components.models.deepseek_v4.fsdp._fully_shard_once(
    module: torch.nn.Module,
    mesh,
    mp_policy,
    offload_policy,
    fp32_policy: bool = False,
    fsdp_kwargs = {}
)

nemo_automodel.components.models.deepseek_v4.fsdp._has_fsdp_state(
    module: torch.nn.Module
) -> bool

nemo_automodel.components.models.deepseek_v4.fsdp._hca_param_sync_group_from_1d_mesh(
    mesh
)

Return the 1D PyTorch FSDP2 group used for HCA graph alignment.

HCA graph alignment is an FSDP/FSDP2 parameter-sync invariant: ranks that synchronize the same sharded HCA parameters must agree on whether the HCA compressor path participates in backward. This DeepSeek-V4 wrapper gets that domain from its 1D PyTorch FSDP2 mesh. The mesh may be named or unnamed; multi-dimensional meshes need an explicit owner dimension to avoid reducing across unrelated parallel groups. Until that is available, disable HCA graph alignment instead of using a broader or wrong group.

nemo_automodel.components.models.deepseek_v4.fsdp._is_deepseek_v4_module(
    module: torch.nn.Module
) -> bool

nemo_automodel.components.models.deepseek_v4.fsdp._iter_dsv4_fp32_modules(
    module: torch.nn.Module
)

nemo_automodel.components.models.deepseek_v4.fsdp._matches_suffix(
    name: str,
    suffix: str
) -> bool

nemo_automodel.components.models.deepseek_v4.fsdp._module_config_model_type(
    module: torch.nn.Module
) -> str | None

nemo_automodel.components.models.deepseek_v4.fsdp.fully_shard_deepseek_v4(
    module: torch.nn.Module | list[torch.nn.Module],
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None = None,
    fsdp_kwargs: typing.Unpack[nemo_automodel.components.models.deepseek_v4.fsdp.DeepseekV4FSDPKwargs] = {}
) -> torch.distributed.fsdp.FSDPModule | list[torch.distributed.fsdp.FSDPModule]

Apply FSDP2 to DeepSeek-V4 without mixing fp32 and bf16 params in one unit.

This is intentionally model-specific. DeepSeek-V4 keeps a small set of reference-sensitive tensors in fp32, while the existing DeepEP path expects the transformer block itself to remain the main FSDP unit.

Uniform fp32 storage does not necessarily mean fp32 compute: full-parameter training commonly keeps fp32 master weights while the FSDP mixed-precision policy requests bf16 compute. In that case, preserve fp32 compute only for the explicitly listed reference-sensitive islands and let the parent block use the caller’s bf16 policy.

Parameters:

module

nn.Module | list[nn.Module]

One module, or a list of modules sharded together as one cross-block FSDP unit under the caller’s explicit policy.

mesh

DeviceMesh

Data-parallel shard mesh.

mp_policy

MixedPrecisionPolicy

Caller’s mixed-precision policy for the parent unit.

offload_policy

OffloadPolicy | NoneDefaults to None

Optional CPU offload policy.

**fsdp_kwargs

Unpack[DeepseekV4FSDPKwargs]Defaults to {}

Remaining fully_shard options, forwarded unchanged except that ignored_params is filtered per FP32 island.

Returns: FSDPModule | list[FSDPModule]

The sharded module, or the list of sharded modules for a group input.

nemo_automodel.components.models.deepseek_v4.fsdp._DSV4_CLASS_NAMES = {'DeepseekV4ForCausalLM', 'DeepseekV4Model', 'DeepseekV4Block', 'DeepseekV4Visio...

nemo_automodel.components.models.deepseek_v4.fsdp._DSV4_FP32_MODULE_SUFFIXES = ('attn_hc', 'ffn_hc', 'hc_head', 'lm_head', 'self_attn.sinks_param', 'self_attn....

nemo_automodel.components.models.deepseek_v4.fsdp._DSV4_RETURNED_FP32_PARAMETER_CLASS_NAMES = {'DeepseekV4FP32Parameter'}

nemo_automodel.components.models.deepseek_v4.fsdp._DSV4_SELF_CASTING_FP32_MODULE_CLASS_NAMES = {'DeepseekV4VisionRMSNorm'}