nemo_automodel.components.models.deepseek_v41.layers

Single-pass mHC and modality-aware routing for DeepSeek V4.1.

The coefficient handoff follows section 2.4.1 of the official technical report: each sublayer consumes the pre-mix produced by the preceding sublayer.

Module Contents

Classes

Name	Description
`DeepseekV41HyperConnection`	Predict one sublayer’s residual mixing coefficients in FP32.
`DeepseekV41Mix`	FP32 coefficients: pre/post [batch, sequence, streams], comb [batch, sequence, streams, streams].
`DeepseekV41RMSNorm`	Normalize in FP32 and multiply the scale before casting the result.

API

class nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    sinkhorn_backend: typing.Literal['torch', 'tilelang'] = 'torch'
)

Bases: Module

Predict one sublayer’s residual mixing coefficients in FP32.

base

eps

= config.hc_eps

iterations

= config.hc_sinkhorn_iters

norm_eps

= config.rms_norm_eps

scale

= nn.Parameter(torch.ones(3, dtype=(torch.float32)))

streams

= config.hc_mult

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection.collapse(
    hidden_states: torch.Tensor,
    pre_mix: torch.Tensor
) -> torch.Tensor

staticmethod

Collapse streams using the preceding sublayer’s coefficients.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, streams, hidden].

pre_mix

torch.Tensor

FP32 tensor of shape [batch, sequence, streams].

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden], with the input dtype.

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection.expand(
    output: torch.Tensor,
    residual: torch.Tensor,
    mix: nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41Mix
) -> torch.Tensor

staticmethod

Mix the sublayer output and residual in the source’s coefficient orientation.

Parameters:

output

torch.Tensor

Tensor of shape [batch, sequence, hidden].

residual

torch.Tensor

Tensor of shape [batch, sequence, streams, hidden].

mix

DeepseekV41Mix

FP32 pre/post tensors of shape [batch, sequence, streams] and comb of shape [batch, sequence, input_streams, output_streams].

Returns: torch.Tensor

Tensor of shape [batch, sequence, streams, hidden], with output’s dtype.

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection.forward(
    hidden_states: torch.Tensor
) -> nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41Mix

Predict coefficients, preserving projection-before-RMS arithmetic.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, streams, hidden].

Returns: DeepseekV41Mix

Coefficients with pre/post tensors of shape [batch, sequence, streams]

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection.reset_parameters(
    std: float = 0.02
) -> None

Initialize all coefficients before checkpoint-free execution.

class nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41Mix(
    pre: torch.Tensor,
    post: torch.Tensor,
    comb: torch.Tensor
)

Dataclass

FP32 coefficients: pre/post [batch, sequence, streams], comb [batch, sequence, streams, streams].

comb

Tensor

post

Tensor

pre

Tensor

class nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41RMSNorm(
    dim: int,
    eps: float,
    dtype: torch.dtype
)

Bases: Module

Normalize in FP32 and multiply the scale before casting the result.

weight

= nn.Parameter(torch.ones(dim, dtype=dtype))

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41RMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Apply the released reference’s RMS normalization.

Parameters:

hidden_states

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

Returns: torch.Tensor

Tensor of shape […, hidden], with the input dtype.