nemo_automodel.components.models.deepseek_v4.model

DeepSeek V4 Model.

Key architectural points (from official inference/model.py):

HC (Hyper-Connections): Every transformer block maintains hc_mult=4 copies of the hidden state. The embedding output is expanded: [B,S,dim] -> [B,S,hc_mult,dim]. hc_pre reduces [B,S,hc_mult,dim] -> [B,S,dim] before attn/ffn. hc_post expands [B,S,dim] -> [B,S,hc_mult,dim] after attn/ffn. Full HC requires the hc_split_sinkhorn CUDA kernel. Current fallback: mean-pooling for hc_pre, broadcast add for hc_post.

HC parameters (ALL layers, stored in float32): hc_attn_fn : [mix_hc, hc_mult*dim] where mix_hc = (2+hc_mult)hc_mult = 24 hc_attn_base : [mix_hc] hc_attn_scale : [3] hc_ffn_fn : [mix_hc, hc_multdim] hc_ffn_base : [mix_hc] hc_ffn_scale : [3]

Gate hash layers (layer_idx < num_hash_layers): Instead of score-based routing, the gate uses a fixed token-id -> expert-id lookup table (tid2eid: [vocab_size, n_activated_experts]).

All layers use MoE FFN (no dense layers). Compress-ratio sliding-window attention is not yet implemented.

Module Contents

Classes

Name	Description
`DeepseekV4Block`	Single transformer block for DeepSeek V4.
`DeepseekV4CausalLMOutput`	Output of DeepseekV4ForCausalLM.forward.
`DeepseekV4ForCausalLM`	-
`DeepseekV4HashGate`	Hash gate for first num_hash_layers: routes tokens via a fixed lookup table.
`DeepseekV4Model`	-
`DeepseekV4VisionGate`	DSV4 gate with separate visual bias and optional text hash routing.

Functions

Name	Description
`_normalize_thd_packing_metadata`	Accept standard THD offsets at the DSV4 model boundary.
`_seq_lens_from_cu_seqlens`	Convert standard THD cumulative offsets to DSV4’s per-row lengths.
`apply_deepseek_v4_image_visibility`	Make every token inside an image span mutually visible.

Data

ModelClass

API

class nemo_automodel.components.models.deepseek_v4.model.DeepseekV4Block(
    layer_idx: int,
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

Single transformer block for DeepSeek V4.

Uses HuggingFace transformers PR 45616’s HyperConnection decoder-layer pattern: two DeepseekV4HyperConnection modules own the collapse / expand mixer weights at the attention and FFN sites respectively. Checkpoint’s flat hc_attn_* / hc_ffn_* keys are routed into attn_hc.* / ffn_hc.* by the state-dict adapter.

attn_hc

= DeepseekV4HyperConnection(**hc_kwargs)

ffn_hc

= DeepseekV4HyperConnection(**hc_kwargs)

hc_mult

= config.hc_mult

input_layernorm

is_hash_routing_layer

is_vision_routing

mlp

= MoE(moe_config, backend)

post_attention_layernorm

self_attn

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4Block.forward(
    x: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    position_ids: torch.Tensor | None = None,
    position_embeddings_compress: tuple[torch.Tensor, torch.Tensor] | None = None,
    rotary_compress: torch.nn.Module | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    input_ids: torch.Tensor | None = None,
    vision_token_types: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

Transform one HC block.

Parameters:

torch.Tensor

HC streams with layout [batch, sequence, hc_mult, hidden].

position_embeddings

tuple[torch.Tensor, torch.Tensor]

Main RoPE tensors with layout compatible with [batch, sequence, qk_rope_head_dim].

position_ids

torch.Tensor | NoneDefaults to None

Token positions with layout [batch, sequence].

position_embeddings_compress

tuple[torch.Tensor, torch.Tensor] | NoneDefaults to None

Optional compressor RoPE tensors.

rotary_compress

nn.Module | NoneDefaults to None

Optional compressor rotary module.

attention_mask

torch.Tensor | NoneDefaults to None

Additive attention mask with layout [batch, 1, sequence, sequence].

padding_mask

torch.Tensor | NoneDefaults to None

Boolean padding mask with layout [batch, sequence].

input_ids

torch.Tensor | NoneDefaults to None

Token IDs with layout [batch, sequence].

vision_token_types

torch.Tensor | NoneDefaults to None

Visual pseudo-token types with layout [batch, sequence] and -1 at text positions.

Returns: torch.Tensor

HC streams with layout [batch, sequence, hc_mult, hidden].

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4Block.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

class nemo_automodel.components.models.deepseek_v4.model.DeepseekV4CausalLMOutput(
    mtp_per_depth_h: list[torch.Tensor] | None = None,
    mtp_loss_scaling_factor: float | None = None
)

Dataclass

Bases: CausalLMOutputWithPast

Output of DeepseekV4ForCausalLM.forward.

Subclasses transformers.modeling_outputs.CausalLMOutputWithPast so the standard logits / hidden_states fields are present (the recipe’s fused cross-entropy path requires "hidden_states" in out and reads the final hidden states off the output) while the DSV4-specific MTP fields are carried as declared dataclass fields. As required by ModelOutput, every field after the first declares a None default.

mtp_loss_scaling_factor

float | None = None

mtp_per_depth_h

list[Tensor] | None = None

class nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

_keep_in_fp32_modules_strict

backend

= backend or BackendConfig()

lm_head

model

mtp

mtp_config

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM._build_mtp_embed_inputs_for_pp(
    input_ids: torch.Tensor
) -> tuple[torch.Tensor, ...]

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM._is_pipeline_parallel_stage() -> bool

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.customize_pipeline_stage_modules(
    module_names_per_stage: list[list[str]],
    layers_prefix: str,
    text_model: torch.nn.Module | None = None
) -> list[list[str]]

Keep DSV4 non-layer PP dependencies with the stages that need them.

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    mtp_embed_inputs: torch.Tensor = (),
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    pixel_values: torch.Tensor | None = None,
    image_grid_hws: torch.Tensor | None = None,
    n_images_per_sample: torch.Tensor | None = None,
    vision_token_types: torch.Tensor | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: bool | None = None,
    attn_kwargs: typing.Any = {}
) -> 'DeepseekV4CausalLMOutput' | tuple[torch.Tensor, ...] | torch.Tensor

Run causal language modeling with optional DSV4 visual inputs.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Token IDs with layout [batch, sequence].

mtp_embed_inputs

torch.TensorDefaults to ()

Optional PP-propagated MTP embeddings, each with layout [batch, sequence, hidden].

position_ids

torch.Tensor | NoneDefaults to None

Position IDs with layout [batch, sequence].

attention_mask

torch.Tensor | NoneDefaults to None

Valid-token mask with layout [batch, sequence].

padding_mask

torch.Tensor | NoneDefaults to None

Padding mask with layout [batch, sequence].

pixel_values

torch.Tensor | NoneDefaults to None

Concatenated patches with layout [all_patches, 3, patch_size, patch_size].

image_grid_hws

torch.Tensor | NoneDefaults to None

ViT grids with layout [all_images, 2].

n_images_per_sample

torch.Tensor | NoneDefaults to None

Image counts with layout [batch].

vision_token_types

torch.Tensor | NoneDefaults to None

Visual types with layout [batch, sequence] and -1 for text.

logits_to_keep

Union[int, torch.Tensor]Defaults to 0

Number or positions of logits to retain.

output_hidden_states

bool | NoneDefaults to None

Whether to expose final states.

Returns: 'DeepseekV4CausalLMOutput' | tuple[torch.Tensor, ...] | torch.Tensor

DeepseekV4CausalLMOutput outside PP, or the PP stage tensor

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.from_config(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

classmethod

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)

classmethod

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.get_input_embeddings()

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.get_output_embeddings()

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]

Return PP input/output meta tensors for DSV4’s HC and MTP contract.

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.prepare_model_inputs_for_cp(
    batch: dict[str, typing.Any],
    num_chunks: int = 1
) -> dict[str, typing.Any]

Model-owned context-parallel batch prep (Miles-style contiguous shard).

Returns a ContextParallelSharder (under the "cp_sharder" batch key) so the CP dispatch delegates CP sharding back to this model, with the config-derived per-rank shard multiple bound. DSV4 embeds internally, so (unlike VLM models) this does not pre-embed — it leaves input_ids for the sharding callable.

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.set_input_embeddings(
    value
)

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.set_output_embeddings(
    new_embeddings
)

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4ForCausalLM.update_moe_gate_bias() -> None

class nemo_automodel.components.models.deepseek_v4.model.DeepseekV4HashGate(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig
)

Bases: Module

Hash gate for first num_hash_layers: routes tokens via a fixed lookup table.

Instead of computing routing scores, the gate uses tid2eid[token_id] to pre-assign expert indices. The routing weight is still computed from the gate weight but the selection is deterministic per token id.

tid2eid shape: [vocab_size, n_activated_experts] (int64 runtime, non-trainable)

Signature matches components.moe.layers.Gate — forward(x, token_mask, cp_mesh) returning (weights, indices, aux_loss) — so the generic MoE module can call it interchangeably. The per-forward input_ids needed for the tid2eid lookup is stashed on the module by the enclosing Block via :meth:set_input_ids immediately before the MoE call.

_pending_input_ids

Tensor | None = None

n_experts

= moe_config.n_routed_experts

norm_topk_prob

= moe_config.norm_topk_prob

route_scale

= moe_config.route_scale

score_func

= moe_config.score_func

topk

= moe_config.n_activated_experts

weight

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4HashGate.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor | None = None,
    cp_mesh: 'DeviceMesh | None' = None
) -> tuple[torch.Tensor, torch.Tensor, None]

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4HashGate.init_weights(
    init_std: float = 0.02
) -> None

Initialize the trainable gate and a valid deterministic hash table.

Parameters:

init_std

floatDefaults to 0.02

Standard deviation for the routing weight initialization.

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4HashGate.set_input_ids(
    input_ids: torch.Tensor | None
) -> None

Stash the current batch’s input_ids for the next forward call.

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4HashGate.update_bias() -> None

No-op for compat with callers that walk MoE gates and call update_bias.

class nemo_automodel.components.models.deepseek_v4.model.DeepseekV4Model(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    moe_overrides: dict | None = None
)

Bases: Module

aligner

= DeepseekV4VisionAligner(config)

embed_tokens

hc_head

image_end

image_newline

image_pad

image_start

layers

= nn.ModuleDict()

max_seq_len

= config.max_position_embeddings

moe_config

= moe_config or MoEConfig(**moe_defaults)

norm

rotary_emb

rotary_emb_compress

vision

= DeepseekV4VisionTransformer(config)

vision_enabled

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4Model.encode_image(
    patches: torch.Tensor,
    n_vit_h: int,
    n_vit_w: int
) -> torch.Tensor

Encode one image from ViT patches into LLM-width grid features.

Parameters:

patches

torch.Tensor

Image patches with layout [n_vit_h * n_vit_w, 3, patch_size, patch_size].

n_vit_h

int

Number of patch rows.

n_vit_w

int

Number of patch columns.

Returns: torch.Tensor

Aligned features with layout

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4Model.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    pixel_values: torch.Tensor | None = None,
    image_grid_hws: torch.Tensor | None = None,
    n_images_per_sample: torch.Tensor | None = None,
    vision_token_types: torch.Tensor | None = None,
    return_hc_hidden: bool = False,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

Run the DSV4 text backbone with an optional image embedding bridge.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Token IDs with layout [batch, sequence] on the first PP stage, or HC activations on later stages.

inputs_embeds

torch.Tensor | NoneDefaults to None

Optional embeddings with layout [batch, sequence, hidden].

position_ids

torch.Tensor | NoneDefaults to None

Positions with layout [batch, sequence].

attention_mask

torch.Tensor | NoneDefaults to None

Valid-token mask with layout [batch, sequence].

padding_mask

torch.Tensor | NoneDefaults to None

Padding mask with layout [batch, sequence].

pixel_values

torch.Tensor | NoneDefaults to None

Concatenated image patches with layout [all_patches, 3, patch_size, patch_size].

image_grid_hws

torch.Tensor | NoneDefaults to None

Patch grids with layout [all_images, 2].

n_images_per_sample

torch.Tensor | NoneDefaults to None

Image counts with layout [batch].

vision_token_types

torch.Tensor | NoneDefaults to None

Pseudo-token types with layout [batch, sequence] and -1 at text positions.

return_hc_hidden

boolDefaults to False

Whether to also return the uncollapsed HC stream.

Returns: torch.Tensor | tuple[torch.Tensor, torch.Tensor]

Hidden states with layout [batch, sequence, hidden] and,

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4Model.init_weights(
    buffer_device: torch.device | None = None
) -> None

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4Model.merge_image_embeddings(
    inputs_embeds: torch.Tensor,
    pixel_values: torch.Tensor,
    image_grid_hws: torch.Tensor,
    vision_token_types: torch.Tensor,
    n_images_per_sample: torch.Tensor | None
) -> torch.Tensor

Replace pseudo-token embeddings with encoded images and sentinels.

Parameters:

inputs_embeds

torch.Tensor

Text embeddings with layout [batch, sequence, hidden].

pixel_values

torch.Tensor

Concatenated patches with layout [all_patches, 3, patch_size, patch_size].

image_grid_hws

torch.Tensor

ViT grid sizes with layout [all_images, 2].

vision_token_types

torch.Tensor

Pseudo-token types with layout [batch, sequence] and -1 for text.

n_images_per_sample

torch.Tensor | None

Optional counts with layout [batch].

Returns: torch.Tensor

Embeddings with the same layout as inputs_embeds.

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4Model.update_moe_gate_bias() -> None

class nemo_automodel.components.models.deepseek_v4.model.DeepseekV4VisionGate(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    gate_precision: torch.dtype | None,
    hash_routing: bool
)

Bases: Gate

DSV4 gate with separate visual bias and optional text hash routing.

The released vision checkpoint routes visual pseudo tokens by score in all layers. In the first hash layers only text tokens use tid2eid; visual tokens use scores + bias_vl. Later layers select text experts with the normal correction bias and visual experts with bias_vl.

_pending_input_ids

Tensor | None = None

_pending_vision_token_types

Tensor | None = None

vocab_size

= int(config.vocab_size)

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4VisionGate._local_tensor(
    tensor: torch.Tensor
) -> torch.Tensor

staticmethod

Return a local [experts] tensor from a tensor or DTensor.

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4VisionGate._route_scores(
    scores: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Select experts from router logits of layout [tokens, experts].

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4VisionGate.init_dsv4_weights() -> None

Initialize visual bias and a deterministic checkpoint-free hash map.

nemo_automodel.components.models.deepseek_v4.model.DeepseekV4VisionGate.set_routing_context(
    input_ids: torch.Tensor | None,
    vision_token_types: torch.Tensor | None
) -> None

Set token metadata consumed by the next gate call.

Parameters:

input_ids

torch.Tensor | None

Token IDs with layout [batch, sequence].

vision_token_types

torch.Tensor | None

Visual types with layout [batch, sequence] and -1 for text tokens.

nemo_automodel.components.models.deepseek_v4.model._normalize_thd_packing_metadata(
    attn_kwargs: dict[str, typing.Any]
) -> None

Accept standard THD offsets at the DSV4 model boundary.

DSV4 internally uses seq_lens to build document-aware masks. Packed callers commonly provide the equivalent cu_seqlens representation, so normalize it here when context parallelism has not already produced native padded-BSHD lengths.

nemo_automodel.components.models.deepseek_v4.model._seq_lens_from_cu_seqlens(
    cu_seqlens: torch.Tensor,
    name: str
) -> torch.Tensor

Convert standard THD cumulative offsets to DSV4’s per-row lengths.

nemo_automodel.components.models.deepseek_v4.model.apply_deepseek_v4_image_visibility(
    attention_mask: torch.Tensor,
    vision_token_types: torch.Tensor
) -> torch.Tensor

Make every token inside an image span mutually visible.

Parameters:

attention_mask

torch.Tensor

Additive causal/sliding mask with layout [batch, 1, sequence, sequence].

vision_token_types

torch.Tensor

Pseudo-token types with layout [batch, sequence] and -1 at text positions.

Returns: torch.Tensor

Additive mask with the same layout as attention_mask. Text

nemo_automodel.components.models.deepseek_v4.model.ModelClass = DeepseekV4ForCausalLM