nemo_automodel.components.models.inkling.model

View as Markdown

NeMo AutoModel wrapper for the Inkling multimodal MoE model.

The wrapper reuses the HuggingFace InklingForConditionalGeneration towers, attention, norms, embeddings, and language-model head. Decoder feed-forwards retain the raw checkpoint’s fused interleaved projection layout; sparse layers also use an expert-parallel :class:InklingMoE. This avoids full-size weight conversion copies while preserving Transformers numerics.

Module Contents

Classes

NameDescription
InklingForConditionalGenerationInkling VLM with expert-parallel MoE feed-forwards.
InklingTextModelInkling text backbone that also accepts AutoModel pipeline stages.

API

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

Bases: HFCheckpointingMixin, HFInklingForConditionalGeneration, MoEFSDPSyncMixin

Inkling VLM with expert-parallel MoE feed-forwards.

_keep_in_fp32_modules_strict
= ['_fp32_params']
_pp_keep_self_forward
bool = True
_skip_init_weights_on_load
bool = True
moe_config
state_dict_adapter
tie_word_embeddings_support
TieSupport = TieSupport.UNTIED_ONLY
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: typing.Any | 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 = {}
) -> typing.Any

Run the standard Inkling forward or its pipeline-stage equivalent.

nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.from_config(
config: transformers.models.inkling.configuration_inkling.InklingConfig,
moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
backend: nemo_automodel.components.models.common.BackendConfig | None = None,
kwargs = {}
) -> 'InklingForConditionalGeneration'
classmethod
nemo_automodel.components.models.inkling.model.InklingForConditionalGeneration.from_pretrained(
pretrained_model_name_or_path: str,
model_args = (),
kwargs = {}
) -> 'InklingForConditionalGeneration'
classmethod
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 metadata using Inkling’s unpadded runtime vocabulary.

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

Inkling uses a trained correction bias, so gate-bias updates are a no-op.

class nemo_automodel.components.models.inkling.model.InklingTextModel()

Bases: HFInklingTextModel

Inkling text backbone that also accepts AutoModel pipeline stages.

nemo_automodel.components.models.inkling.model.InklingTextModel.forward(
input_ids: torch.LongTensor | None = None,
attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: typing.Any | None = None,
inputs_embeds: torch.FloatTensor | None = None,
use_cache: bool | None = None,
kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast