nemo_automodel.components.models.deepseek_v4.fsdp

Module Contents

Functions

Name	Description
`_attach_hca_param_sync_group`	-
`_floating_param_dtypes`	-
`_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

API

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_mp_policy(
    mp_policy
)

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,
    mesh,
    mp_policy,
    offload_policy = None,
    fsdp_kwargs = {}
)

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.

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

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....