ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsKimi K3

nemo_automodel.components.models.kimi_k3.vision

View as Markdown

Module Contents

Classes

Functions

Data

VL_VISION_ATTENTION_FUNCTIONS

API

class nemo_automodel.components.models.kimi_k3.vision.IdentityMap()

Bases: Module

nemo_automodel.components.models.kimi_k3.vision.IdentityMap.forward(
    x,
    args = (),
    kwargs = {}
)
class nemo_automodel.components.models.kimi_k3.vision.Learnable2DInterpPosEmbDivided_fixed(
    height: int,
    width: int,
    num_frames: int,
    dim: int,
    interpolation_mode: str = 'bicubic'
)

Bases: Module

weight
= nn.Parameter(torch.empty(height, width, dim))
nemo_automodel.components.models.kimi_k3.vision.Learnable2DInterpPosEmbDivided_fixed.forward(
    x: torch.Tensor,
    grid_thws: torch.Tensor
) -> torch.Tensor
nemo_automodel.components.models.kimi_k3.vision.Learnable2DInterpPosEmbDivided_fixed.reset_parameters()
class nemo_automodel.components.models.kimi_k3.vision.MLP(
    config
)

Bases: Module

pre_norm
= nn.LayerNorm(config.mm_hidden_size)
proj
nemo_automodel.components.models.kimi_k3.vision.MLP.forward(
    x,
    args = (),
    kwargs = {}
)
class nemo_automodel.components.models.kimi_k3.vision.MLP2(
    dims: list[int],
    activation,
    bias = True
)

Bases: Module

Parameters:

dims
list[int]

[in_dim, hidden_dim, out_dim]

bias
Defaults to True

whether to use bias in linear layer.

fc0
= nn.Linear(dims[0], dims[1], bias=bias)
fc1
= nn.Linear(dims[1], dims[2], bias=bias)
nemo_automodel.components.models.kimi_k3.vision.MLP2.forward(
    x: torch.Tensor
) -> torch.Tensor
class nemo_automodel.components.models.kimi_k3.vision.MoonViT3dEncoder(
    hidden_dim: int,
    num_layers: int,
    block_cfg: dict,
    use_deterministic_attn: bool = False
)

Bases: Module

blocks
final_layernorm
= nn.LayerNorm(hidden_dim)
rope_2d
nemo_automodel.components.models.kimi_k3.vision.MoonViT3dEncoder.forward(
    hidden_states: torch.Tensor,
    grid_thws: torch.Tensor
) -> torch.Tensor
class nemo_automodel.components.models.kimi_k3.vision.MoonViT3dPretrainedModel(
    config,
    inputs = (),
    kwargs = {}
)

Bases: PreTrainedModel

_no_split_modules
= ['MoonViTEncoderLayer']
encoder
merge_kernel_size
= config.merge_kernel_size
merge_type
= config.merge_type
model_type
= 'moonvit3d'
patch_embed
patch_size
= config.patch_size
nemo_automodel.components.models.kimi_k3.vision.MoonViT3dPretrainedModel.forward(
    pixel_values: torch.Tensor,
    grid_thws: torch.Tensor
) -> torch.Tensor

Parameters:

pixel_values
torch.Tensor

The input pixel values.

grid_thws
torch.Tensor

Temporal, height and width.

Returns: torch.Tensor

torch.Tensor: The output tokens.

class nemo_automodel.components.models.kimi_k3.vision.MoonViTEncoderLayer(
    num_heads: int,
    hidden_dim: int,
    mlp_dim: int,
    qkv_hidden_size: int | None = None,
    norm_type: str = 'layernorm',
    mlp_type: str = 'mlp2',
    attn_implementation: str = 'flash_attention_2',
    activation = F.gelu,
    attn_bias: bool = False,
    linear_bias: bool = True,
    use_deterministic_attn: bool = False
)

Bases: Module

hidden_size_per_attention_head
= self.qkv_hidden_size // self.num_heads
mlp
norm0
= nn.LayerNorm(hidden_dim)
norm1
= nn.LayerNorm(hidden_dim)
qkv_hidden_size
wo
wqkv
nemo_automodel.components.models.kimi_k3.vision.MoonViTEncoderLayer.attention_qkvpacked(
    x: torch.Tensor,
    cu_seqlens: torch.Tensor,
    max_seqlen: torch.Tensor,
    rope_freqs_cis: torch.Tensor | None = None
)

Parameters:

x
torch.Tensor

(batch_size, seqlen, hidden_dim)

cu_seqlens
torch.Tensor
nemo_automodel.components.models.kimi_k3.vision.MoonViTEncoderLayer.forward(
    hidden_states: torch.Tensor,
    cu_seqlens: torch.Tensor,
    max_seqlen: int,
    rope_freqs_cis: torch.Tensor | None = None
)
class nemo_automodel.components.models.kimi_k3.vision.MoonVision3dPatchEmbed(
    out_dim: int,
    in_dim: int = 3,
    patch_size: int | tuple[int, int] = (14, 14),
    pos_emb_height: int = 14,
    pos_emb_width: int = 14,
    pos_emb_time: int = 4,
    pos_emb_type: str = 'divided_fixed',
    patch_embed_proj_bias: bool = True,
    pos_emb_interpolation_mode: str = 'bicubic'
)

Bases: Module

pos_emb
proj
nemo_automodel.components.models.kimi_k3.vision.MoonVision3dPatchEmbed.forward(
    x: torch.Tensor,
    grid_thws: torch.Tensor
) -> torch.Tensor

Parameters:

x
(L, Channels)

input tensor

grid_hws
(N, 3)

temporal, height and width

Returns: torch.Tensor

(L, Cout) tensor

class nemo_automodel.components.models.kimi_k3.vision.PatchMergerMLP(
    config
)

Bases: Module

hidden_size
pre_norm
= nn.LayerNorm(config.mm_hidden_size, eps=eps)
proj
nemo_automodel.components.models.kimi_k3.vision.PatchMergerMLP.forward(
    x,
    args = (),
    kwargs = {}
)
class nemo_automodel.components.models.kimi_k3.vision.PatchMergerMLPV2(
    config
)

Bases: Module

hidden_size
post_norm
= nn.RMSNorm(config.hidden_size, eps=eps)
proj
nemo_automodel.components.models.kimi_k3.vision.PatchMergerMLPV2.forward(
    x,
    args = (),
    kwargs = {}
)
class nemo_automodel.components.models.kimi_k3.vision.Rope2DPosEmbRepeated(
    dim: int,
    max_height: int,
    max_width: int,
    theta_base = 10000
)

Bases: Module

2D rotary position embedding with multi-resolution support.

This class is intended to be used in the following way:

  1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis.
  2. Before each forward pass, call get_freqs_cis_by_* to get the freqs_cis tensor for this iteration.
  3. During the forward pass, pass the freqs_cis tensor to each attention layer, and call apply just before each attention operation. The rope is shared across all attention layers and all heads.

Refs:

Parameters:

dim
int

usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed)

max_height
int

the maximum height of the 2D grid

max_width
int

the maximum width of the 2D grid

theta_base
floatDefaults to 10000

the base of the theta

device
str

the device to store the precomputed cis

nemo_automodel.components.models.kimi_k3.vision.Rope2DPosEmbRepeated._precompute_freqs_cis(
    device: torch.device
) -> torch.Tensor

Calculate the cis(freqs) for each position in the 2D grid.

nemo_automodel.components.models.kimi_k3.vision.Rope2DPosEmbRepeated.extra_repr()
nemo_automodel.components.models.kimi_k3.vision.Rope2DPosEmbRepeated.get_freqs_cis(
    grid_thws: torch.Tensor,
    device: torch.device
) -> torch.Tensor

Parameters:

grid_thws
torch.Tensor

grid time, height and width

Returns: torch.Tensor

tensor of shape (sum(t * height * width), dim//2)

nemo_automodel.components.models.kimi_k3.vision._apply_rope_input_validation(
    x,
    freqs_cis
)
 
nemo_automodel.components.models.kimi_k3.vision.apply_rope(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Returns: xq_out, xk_out: tensors of shape (…, num_heads, head_dim)

Parameters:

xq
torch.Tensor

query, tensor of shape (…, num_heads, head_dim)

xk
torch.Tensor

key, tensor of shape (…, num_heads, head_dim)

freqs_cis
torch.Tensor

tensor of shape (…, head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid.

nemo_automodel.components.models.kimi_k3.vision.eager_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    q_cu_seqlens: typing.Optional[torch.Tensor] = None,
    k_cu_seqlens: typing.Optional[torch.Tensor] = None,
    kwargs = {}
) -> torch.Tensor
 
nemo_automodel.components.models.kimi_k3.vision.get_1d_sincos_pos_embed(
    embed_dim,
    t_size,
    cls_token = False
)

t_size: int of the temporal size return: pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token)

nemo_automodel.components.models.kimi_k3.vision.get_1d_sincos_pos_embed_from_grid(
    embed_dim,
    pos
)

From: https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86 embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

nemo_automodel.components.models.kimi_k3.vision.get_rope_shape(
    org,
    interpolation_mode,
    shape
)
 
nemo_automodel.components.models.kimi_k3.vision.get_rope_shape_decorate(
    func
)
 
nemo_automodel.components.models.kimi_k3.vision.multihead_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    q_cu_seqlens: torch.Tensor | None = None,
    k_cu_seqlens: torch.Tensor | None = None,
    max_seqlen_q: int | None = None,
    max_seqlen_k: int | None = None,
    deterministic: bool = False
)

Multi-head attention using flash attention 2.

Parameters:

q, k, v

tensor of shape (batch_size, seqlen, num_heads, head_dim), or (tot_seqlens, num_heads, head_dim) if packing.

q_cu_seqlens
torch.TensorDefaults to None

cumulative sequence lengths of q. The first element should be 0 and the last element should be q.shape[0].

k_cu_seqlens
torch.TensorDefaults to None

cumulative sequence lengths of k. The first element should be 0 and the last element should be k.shape[0].

Returns:

shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing, where dim = num_heads * head_dim

nemo_automodel.components.models.kimi_k3.vision.tpool_patch_merger(
    x: torch.Tensor,
    grid_thws: torch.Tensor,
    merge_kernel_size: tuple[int, int] = (2, 2)
) -> list[torch.Tensor]
 
nemo_automodel.components.models.kimi_k3.vision.VL_VISION_ATTENTION_FUNCTIONS = {'flash_attention_2': multihead_attention, 'eager': eager_attention}