> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.

# nemo_automodel.components.models.inkling2.model

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

## Module Contents

### Classes

| Name                                                                                            | Description                                                                |
| ----------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`Inkling2DecoderLayer`](#nemo_automodel-components-models-inkling2-model-Inkling2DecoderLayer) | One attention/MoE layer operating on explicit immutable depth state.       |
| [`Inkling2ForCausalLM`](#nemo_automodel-components-models-inkling2-model-Inkling2ForCausalLM)   | Text implementation using eager or FlexAttention and common Torch experts. |

### API

```python
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.

---

```python
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.

```python
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.

```python
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](/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-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`

---

```python
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

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

Validate contiguous global layer IDs and exclusive parameter ownership.

```python
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

```python
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

```python
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`** `bool` — default: 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;

```python
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 | None` — default: None

Backbone resolved by the shared splitter.

---

**Returns:** `list[list[str]]`

Copied FQN lists with embed_norm first and output_query last.

```python
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.Tensor` — default: ()

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 | None` — default: 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 | None` — default: 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 | None` — default: 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 | None` — default: 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] | None` — default: 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 | None` — default: None

Optional shared packer real lengths \[batch, documents].

---

**`seq_lens_padded`** `torch.Tensor | None` — default: None

Corresponding padded lengths covering each physical row.

---

**`logits_to_keep`** `int | torch.Tensor` — default: 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`** `bool` — default: 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

```python
nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM.from_config(
    config: nemo_automodel.components.models.inkling2.config.Inkling2TextConfig | nemo_automodel.components.models.inkling2.config.Inkling2Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
) -> nemo_automodel.components.models.inkling2.model.Inkling2ForCausalLM
```

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.

```python
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.

```python
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.

```python
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

```python
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.

```python
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.