nemo_automodel.components.models.inkling.model

Standalone NeMo AutoModel implementation of the Inkling multimodal MoE.

Module Contents

Classes

Name	Description
`InklingCausalLMOutputWithPast`	Inkling logits, optional loss/cache, and multimodal hidden states.
`InklingForConditionalGeneration`	Native Inkling VLM with expert-parallel feed-forwards.

Data

ModelClass

__all__

API

class nemo_automodel.components.models.inkling.model.InklingCausalLMOutputWithPast(
    loss: torch.FloatTensor | None = None,
    logits: torch.FloatTensor | None = None,
    past_key_values: nemo_automodel.components.models.inkling.text.InklingDynamicCache | None = None,
    hidden_states: tuple[torch.FloatTensor, ...] | None = None,
    attentions: tuple[torch.FloatTensor, ...] | None = None,
    image_hidden_states: torch.FloatTensor | None = None
)

Dataclass

Bases: ModelOutput

Inkling logits, optional loss/cache, and multimodal hidden states.

attentions

tuple[FloatTensor, ...] | None = None

hidden_states

tuple[FloatTensor, ...] | None = None

image_hidden_states

FloatTensor | None = None

logits

FloatTensor | None = None

loss

FloatTensor | None = None

past_key_values

InklingDynamicCache | None = None

class nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration(
    config: nemo_automodel.components.models.inkling.configuration.InklingConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

Native Inkling VLM with expert-parallel feed-forwards.

_keep_in_fp32_modules_strict

= ['_fp32_params']

_pp_keep_self_forward

bool = True

_skip_init_weights_on_load

bool = True

backend

= backend or BackendConfig()

lm_head

model

moe_config

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

vocab_size

= text_config.vocab_size

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.customize_pipeline_stage_modules(
    module_names_per_stage: list[list[str]],
    layers_prefix: str,
    text_model: torch.nn.Module
) -> list[list[str]]

Keep Inkling’s post-embedding norm on the first pipeline stage.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.forward(
    input_ids: torch.LongTensor | None = None,
    pixel_values: torch.FloatTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: nemo_automodel.components.models.inkling.text.InklingDynamicCache | None = None,
    audio_input_ids: torch.LongTensor | None = None,
    audio_input_ids_mask: torch.Tensor | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    labels: torch.LongTensor | None = None,
    use_cache: bool | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.inkling.model.InklingCausalLMOutputWithPast | torch.Tensor

Run native Inkling multimodal conditional generation.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Optional long tensor of shape [batch, sequence]. Non-first pipeline stages may receive hidden states of shape [batch, sequence, hidden] through this argument.

pixel_values

torch.FloatTensor | NoneDefaults to None

Optional tensor of shape [patches, time, height, width, channels].

attention_mask

torch.Tensor | NoneDefaults to None

Optional padding tensor of shape [batch, total_sequence].

position_ids

torch.LongTensor | NoneDefaults to None

Optional long tensor of shape [batch, sequence].

past_key_values

InklingDynamicCache | NoneDefaults to None

Optional model-owned decoding cache.

audio_input_ids

torch.LongTensor | NoneDefaults to None

Optional long tensor of shape [audios, frames, mel_bins].

audio_input_ids_mask

torch.Tensor | NoneDefaults to None

Optional boolean tensor of shape [audios, frames].

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Optional tensor of shape [batch, sequence, hidden].

labels

torch.LongTensor | NoneDefaults to None

Optional long tensor of shape [batch, sequence].

use_cache

bool | NoneDefaults to None

Whether to allocate and return a decoding cache.

logits_to_keep

int | torch.TensorDefaults to 0

Number or indices of trailing logits to compute.

**kwargs

AnyDefaults to {}

Additional text-attention arguments.

Returns: InklingCausalLMOutputWithPast | torch.Tensor

An output whose logits have shape [batch, kept_sequence, vocab] during

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.from_config(
    config: nemo_automodel.components.models.inkling.configuration.InklingConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
) -> 'InklingForConditionalGeneration'

classmethod

Construct an Inkling model from its local AutoModel config.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
) -> 'InklingForConditionalGeneration'

classmethod

Construct the native model tree for checkpoint loading.

The NeMo AutoModel bridge and checkpointer own weight loading. This method resolves only the local Inkling config and native module structure.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.get_image_features(
    pixel_values: torch.Tensor,
    kwargs: typing.Any = {}
) -> typing.Any

Encode image/video patches through the native vision tower.

Parameters:

pixel_values

torch.Tensor

Tensor of shape [patches, time, height, width, channels].

**kwargs

AnyDefaults to {}

Reserved for the common multimodal calling convention.

Returns: Any

An output whose pooler output has shape [patches, text_hidden].

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.get_input_embeddings() -> torch.nn.Module

Return the text token-embedding module.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.get_output_embeddings() -> torch.nn.Module | None

Return the language-model output projection.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.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 metadata using Inkling’s unpadded vocabulary.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype | None = None
) -> None

Initialize every parameter for checkpoint-free construction.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.set_input_embeddings(
    embeddings: torch.nn.Module
) -> None

Replace the text token-embedding module.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.set_output_embeddings(
    embeddings: torch.nn.Module
) -> None

Replace the language-model output projection.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.update_moe_gate_bias() -> None

Keep Inkling’s trained router correction bias unchanged.

nemo_automodel.components.models.inkling.model.ModelClass = InklingForConditionalGeneration

nemo_automodel.components.models.inkling.model.__all__ = ['InklingForConditionalGeneration', 'ModelClass']