Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsInkling2nemo_automodel.components.models.inkling2.model

nemo_automodel.components.models.inkling2.model

View as MarkdownOpen in Claude

Native FP32/BF16 Inkling2 text reference with explicit accumulation precision.

Module Contents

Classes

NameDescription
Inkling2DecoderLayerOne attention/MoE layer operating on explicit immutable depth state.
Inkling2ForCausalLMText implementation using eager or FlexAttention and common Torch experts.

API

class nemo_automodel.components.models.inkling2.model.Inkling2DecoderLayer(
    config: nemo_automodel.components.models.inkling2.config.Inkling2TextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype
)

Bases: Module

One attention/MoE layer operating on explicit immutable depth state.

attn
attn_norms
attn_query
attn_sconv
is_boundary
= layer_idx % config.n_layers_per_attnres_block == 0
is_local
= layer_idx in config.local_layer_ids
mlp
= Inkling2MoE(config, backend, dtype=dtype)
mlp_norm
mlp_query
mlp_sconv
self_attn
Inkling2Attention

Expose the shared MoE CP setup hook without another parameter owner.

nemo_automodel.components.models.inkling2.model.Inkling2DecoderLayer.forward(
    state: nemo_automodel.components.models.inkling2.layers.ResidualState,
    valid: torch.Tensor,
    positions: torch.Tensor
) -> nemo_automodel.components.models.inkling2.layers.ResidualState

Advance a depth block, with a new block beginning after boundary attention.

Parameters:

state
ResidualState

Completed blocks and partial, each [batch, sequence, hidden].

valid
torch.Tensor

Boolean tensor of shape [batch, sequence].

positions
torch.Tensor

Integer tensor of shape [batch, sequence].

Returns: ResidualState

New state with tensors of shape [batch, sequence, hidden]; previous state is unchanged.

nemo_automodel.components.models.inkling2.model.Inkling2DecoderLayer.reset_parameters(
    init_std: float,
    buffer_device: torch.device
) -> None

Initialize all layer parameters and buffers, including the common experts.

class nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM(
    config: nemo_automodel.components.models.inkling2.config.Inkling2TextConfig | nemo_automodel.components.models.inkling2.config.Inkling2Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    revision: str | None = None
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

Text implementation using eager or FlexAttention and common Torch experts.

This entry point consumes Inkling2TextConfig; the multimodal subclass inserts projected image features before the same decoder.

_packed_cp_attn_backends
= ('eager', 'flex')
_pp_keep_self_forward
bool = True
_pp_return_hidden_states
bool = False
_pp_return_hidden_states_supported
bool = True
backend
= replace(backend, gate_precision=(torch.float32))
cp_mesh
DeviceMesh | None = None
embed_norm
Module | None
embed_tokens
Embedding | None
layers
lm_head
Inkling2LMHead | None
moe_config
name_or_path
= config.name_or_path
norm
Module | None
output_query
ResidualQuery | None
parallelizer
= Inkling2Parallelizer()
state_dict_adapter
= Inkling2StateDictAdapter(config)
tie_word_embeddings_support
TieSupport = TieSupport.UNTIED_ONLY
vocab_size
= config.vocab_size
nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM._pipeline_boundary_metas(
    boundary: int,
    batch: int,
    sequence: int,
    dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]

Describe carried residual history and token layout at a decoder cut.

Returns: torch.Tensor

Compute-dtype history [batch, sequence, sources, hidden] with completed

nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM._pipeline_layer_ids() -> list[int]

Validate contiguous global layer IDs and exclusive parameter ownership.

nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM._project_logits(
    hidden: torch.Tensor
) -> torch.Tensor

Apply the callable head, logit softcap and unused-vocabulary mask.

Parameters:

hidden
torch.Tensor

Compute-dtype tensor of shape [batch, sequence, hidden].

Returns: torch.Tensor

FP32 logits [batch, sequence, vocab]; unused vocabulary rows are

nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM._receive_pipeline_state(
    history: torch.Tensor,
    tensors: tuple[torch.Tensor, ...]
) -> tuple[nemo_automodel.components.models.inkling2.layers.ResidualState, torch.Tensor, torch.Tensor]

Recover differentiable residual views and per-microbatch token metadata.

Parameters:

history
torch.Tensor

Compute-dtype [batch, sequence, sources, hidden] residual history.

tensors
tuple[torch.Tensor, ...]

One int64 [batch, sequence, 2] layout tensor; neither input is mutated.

Returns: ResidualState

Residual state with views [batch, sequence, hidden], boolean validity

nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM._run_decoder(
    state: nemo_automodel.components.models.inkling2.layers.ResidualState,
    valid: torch.Tensor,
    positions: torch.Tensor,
    return_hidden: bool = False
) -> torch.Tensor | tuple[torch.Tensor, ...]

Run local layers; transport immutable depth state or apply final-stage modules.

Parameters:

state
ResidualState

Completed blocks and partial [batch, sequence, hidden].

valid
torch.Tensor

Boolean token validity [batch, sequence].

positions
torch.Tensor

Integer token positions [batch, sequence].

return_hidden
boolDefaults to False

Skip final projection for ordinary selected-logit handling.

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

Non-final stage history/layout as documented by _pipeline_boundary_metas;

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

Assign embedding normalization and the final residual query to one owner.

Parameters:

module_names_per_stage
list[list[str]]

Generated stage FQNs; caller lists are unchanged.

layers_prefix
str

Empty for this native model’s top-level layers.

text_model
nn.Module | NoneDefaults to None

Backbone resolved by the shared splitter.

Returns: list[list[str]]

Copied FQN lists with embed_norm first and output_query last.

nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM.forward(
    input_ids: torch.Tensor,
    pipeline_tensors: torch.Tensor = (),
    inputs_embeds: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    causal_mask_mapping: dict[str, torch.Tensor | None] | None = None,
    seq_lens: torch.Tensor | None = None,
    seq_lens_padded: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    output_hidden_states: bool = False
) -> torch.Tensor | tuple[torch.Tensor, ...] | transformers.modeling_outputs.CausalLMOutputWithPast

Return unshifted text logits; the shared loss path owns the target shift.

Parameters:

input_ids
torch.Tensor

Integer tensor of shape [batch, global_sequence]. Inputs are packed or unpacked rows; the first stage slices them under CP after image splicing. On non-first pipeline stages, this is floating residual history [batch, sequence, sources, hidden] from _pipeline_boundary_metas.

pipeline_tensors
torch.TensorDefaults to ()

Non-first-stage int64 [batch, sequence, 2] token layout. It carries authoritative positions/validity; repeated mask/position kwargs from the shared schedule are only consumed by the first stage.

inputs_embeds
torch.Tensor | NoneDefaults to None

Optional decoder-ready [batch, sequence, hidden] embeddings, with text embedding normalization already applied. Supported on the first stage only; token IDs still supply layout.

attention_mask
torch.Tensor | NoneDefaults to None

Optional tensor of shape [batch, sequence], one/true for real tokens. Document boundaries come from position resets, independently of padding.

padding_mask
torch.Tensor | NoneDefaults to None

Optional boolean tensor of shape [batch, local_sequence], true for padding, as supplied by the shared trainer. Combined with attention_mask on the first stage; later stages use the carried validity tensor.

position_ids
torch.Tensor | NoneDefaults to None

Optional integer tensor of shape [batch, local_sequence], increasing by one within each document, restarting at zero for a new document.

causal_mask_mapping
dict[str, torch.Tensor | None] | NoneDefaults to None

Shared text-pipeline placeholders keyed by attention type. Only None entries are accepted; model-owned masks use validity and positions. Supplied dense masks with arbitrary leading dimensions are rejected.

seq_lens
torch.Tensor | NoneDefaults to None

Optional shared packer real lengths [batch, documents].

seq_lens_padded
torch.Tensor | NoneDefaults to None

Corresponding padded lengths covering each physical row.

logits_to_keep
int | torch.TensorDefaults to 0

Zero for all positions, a positive tail length, or integer indices of shape [kept_sequence]. Used by the shared linear-loss path.

output_hidden_states
boolDefaults to False

Include final hidden states for external loss; targets are already aligned by the caller and are never shifted here.

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

FP32 logits of shape [batch, kept_sequence, vocab], softcapped before

classmethod

Construct the text model, from a text or a released image/text configuration.

The constructor accepts both forms; this classmethod keeps the explicit entry point.

nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM.get_input_embeddings() -> torch.nn.Embedding | None

Return the untied token lookup; embed_norm is applied in forward.

nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM.get_output_embeddings() -> nemo_automodel.components.models.inkling2.layers.Inkling2LMHead | None

Return the callable untied head for shared FSDP and output-loss setup.

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

Supply static shapes/dtypes to Automodel’s shared pipeline runtime.

Parameters:

is_first
bool

Whether the runtime identifies this as logical stage zero.

microbatch_size
int

Number of unpacked sequences per microbatch.

seq_len
int

Global input sequence length, before additional CP padding.

dtype
torch.dtype

Body compute dtype, independent of FP32 master storage.

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

First-stage inputs are global int64 [batch, sequence]. Other token

nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None
) -> None

Initialize a fresh or materialized model without a checkpoint.

nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM.prepare_model_inputs_for_cp(
    batch: nemo_automodel.components.models.inkling2.context_parallel.CPBatch,
    num_chunks: int = 1
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder]

Prepare packed or unpacked CP labels/layout while retaining full inputs for image splicing.