nemo_automodel.components.models.deepseek_v4.vision

Vision encoder and aligner used by DeepSeek-V4-Flash-Vision-Exp.

Module Contents

Classes

Name	Description
`DeepseekV4VisionAligner`	Spatially downsample ViT patches and project them into the LLM width.
`DeepseekV4VisionAttention`	Full bidirectional patch attention with 2D RoPE.
`DeepseekV4VisionBlock`	Pre-norm attention and MLP residual block.
`DeepseekV4VisionMLP`	Bias-free SwiGLU MLP used by each vision block.
`DeepseekV4VisionPatchEmbed`	Linear embedding of flattened RGB patches.
`DeepseekV4VisionRMSNorm`	Reference RMSNorm with an fp32 scale parameter.
`DeepseekV4VisionTransformer`	DeepSeek ViT: full attention over one image with 2D RoPE.

Functions

Name	Description
`apply_vision_rotary`	Apply DeepSeek’s half-split 2D rotary embedding.
`get_vision_cos_sin`	Build the reference 2D rotary table.

Data

__all__

API

class nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionAligner(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config
)

Bases: Module

Spatially downsample ViT patches and project them into the LLM width.

downsample_ratio

= int(config.vision_downsample_ratio)

nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionAligner.forward(
    x: torch.Tensor,
    n_h: int,
    n_w: int
) -> torch.Tensor

Downsample encoded patch states.

Parameters:

torch.Tensor

ViT output with layout [n_h * n_w, vision_dim].

n_h

int

Number of patch rows.

n_w

int

Number of patch columns.

Returns: torch.Tensor

LLM image embeddings with layout

nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionAligner.init_weights(
    init_std: float
) -> None

Initialize all checkpoint-free aligner parameters.

class nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionAttention(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config
)

Bases: Module

Full bidirectional patch attention with 2D RoPE.

head_dim

= vision_dim // self.n_heads

n_heads

= int(config.vision_n_heads)

wqkv

nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionAttention.forward(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor

Attend over one image.

Parameters:

torch.Tensor

Patch states with layout [patches, vision_dim].

cos

torch.Tensor

Cosine table with layout [patches, 1, head_dim / 2].

sin

torch.Tensor

Sine table with layout [patches, 1, head_dim / 2].

Returns: torch.Tensor

Attention output with layout [patches, vision_dim].

class nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionBlock(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config
)

Bases: Module

Pre-norm attention and MLP residual block.

attn

= DeepseekV4VisionAttention(config)

mlp

= DeepseekV4VisionMLP(config)

norm1

= DeepseekV4VisionRMSNorm(int(config.vision_dim))

norm2

= DeepseekV4VisionRMSNorm(int(config.vision_dim))

nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionBlock.forward(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor

Transform patch states with layout [patches, vision_dim].

class nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionMLP(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config
)

Bases: Module

Bias-free SwiGLU MLP used by each vision block.

nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionMLP.forward(
    x: torch.Tensor
) -> torch.Tensor

Map [patches, vision_dim] back to [patches, vision_dim].

class nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionPatchEmbed(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config
)

Bases: Module

Linear embedding of flattened RGB patches.

proj

nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionPatchEmbed.forward(
    patches: torch.Tensor
) -> torch.Tensor

Project [patches, 3, patch_h, patch_w] to [patches, vision_dim].

class nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionRMSNorm(
    dim: int,
    eps: float = 1e-06
)

Bases: Module

Reference RMSNorm with an fp32 scale parameter.

weight

nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionRMSNorm.forward(
    x: torch.Tensor
) -> torch.Tensor

Normalize an input of layout [..., vision_dim].

class nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionTransformer(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config
)

Bases: Module

DeepSeek ViT: full attention over one image with 2D RoPE.

blocks

norm

= DeepseekV4VisionRMSNorm(int(config.vision_dim))

patch_embed

= DeepseekV4VisionPatchEmbed(config)

rope_dim

rope_theta

= float(config.vision_rope_theta)

nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionTransformer.forward(
    patches: torch.Tensor,
    n_h: int,
    n_w: int
) -> torch.Tensor

Encode one image.

Parameters:

patches

torch.Tensor

Normalized RGB patches with layout [n_h * n_w, 3, patch_size, patch_size].

n_h

int

Number of patch rows.

n_w

int

Number of patch columns.

Returns: torch.Tensor

Encoded patches with layout [n_h * n_w, vision_dim].

nemo_automodel.components.models.deepseek_v4.vision.DeepseekV4VisionTransformer.init_weights(
    init_std: float
) -> None

Initialize all checkpoint-free vision parameters.

nemo_automodel.components.models.deepseek_v4.vision.apply_vision_rotary(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor

Apply DeepSeek’s half-split 2D rotary embedding.

Parameters:

torch.Tensor

Query or key tensor with layout [patches, heads, head_dim].

cos

torch.Tensor

Cosine table with layout [patches, 1, head_dim / 2].

sin

torch.Tensor

Sine table with layout [patches, 1, head_dim / 2].

Returns: torch.Tensor

Rotated tensor with the same layout and dtype as x.

nemo_automodel.components.models.deepseek_v4.vision.get_vision_cos_sin(
    n_h: int,
    n_w: int,
    dim: int,
    theta: float,
    device: torch.device
) -> tuple[torch.Tensor, torch.Tensor]

Build the reference 2D rotary table.

Returns: torch.Tensor

Cosine and sine tensors with layout [n_h * n_w, 1, dim]. The

nemo_automodel.components.models.deepseek_v4.vision.__all__ = ['DeepseekV4VisionAligner', 'DeepseekV4VisionTransformer', 'DeepseekV4VisionRMSN...