> 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.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.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.step5.model

Native Step-5 Preview text and vision-language training models.

## Module Contents

### Classes

| Name                                                                                                           | Description                                                        |
| -------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------ |
| [`Step5ForCausalLM`](#nemo_automodel-components-models-step5-model-Step5ForCausalLM)                           | Step-5 causal LM with native grouped experts and model-owned CP.   |
| [`Step5ForConditionalGeneration`](#nemo_automodel-components-models-step5-model-Step5ForConditionalGeneration) | Step-5 language model with the compatible Step perception encoder. |
| [`Step5TextModel`](#nemo_automodel-components-models-step5-model-Step5TextModel)                               | Embedding, decoder stack and final normalization.                  |

### Data

[`ModelClass`](#nemo_automodel-components-models-step5-model-ModelClass)

### API

```python
class nemo_automodel.components.models.step5.model.Step5ForCausalLM(
    config: nemo_automodel.components.models.step5.configuration.Step5Config | nemo_automodel.components.models.step5.configuration.Step5VLConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)
```

**Bases:** [HFCheckpointingMixin](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `PreTrainedModel`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-MoEFSDPSyncMixin)

Step-5 causal LM with native grouped experts and model-owned CP.

**`_keep_in_fp32_modules`** `= ['rotary_emb']`

---

**`_keep_in_fp32_modules_strict`**

---

**`_no_split_modules`** `= ['Step5DecoderLayer']`

---

**`_pp_keep_self_forward`** `bool = True`

---

**`_pp_return_hidden_states`** `bool = False`

---

**`_pp_return_hidden_states_supported`** `bool = True`

---

**`backend`**

---

**`base_model_prefix`** `= 'model'`

---

**`cp_mesh`** `DeviceMesh | None = None`

---

**`lm_head`**

---

**`model`**

---

**`packed_mask_type`** `= 'document_ids'`

---

**`state_dict_adapter`**

---

**`tie_word_embeddings_support`** `= TieSupport.UNTIED_ONLY`

---

```python
nemo_automodel.components.models.step5.model.Step5ForCausalLM._is_pipeline_parallel_stage() -> bool
```

Identify a stage after AutoPipeline removes modules owned by other stages.

```python
nemo_automodel.components.models.step5.model.Step5ForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    output_hidden_states: bool = False,
    labels: torch.Tensor | None = None,
    _packed_seq_ids: torch.Tensor | None = None,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue | torch.Tensor = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast | torch.Tensor
```

Evaluate all local query positions.

**Parameters:**

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

Int64 tensor of shape \[batch, local_sequence], or the FP32
residual tensor \[batch, local_sequence, hidden] on later pipeline stages.

---

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

Optional tensor of shape \[batch, local_sequence, hidden].

---

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

Optional tensor of shape \[batch, local_sequence], resetting at document starts.

---

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

Optional tensor of shape \[batch, local_sequence], one for valid tokens.

---

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

Optional bool tensor of shape \[batch, local_sequence], true for padding.

---

**`logits_to_keep`** `int | torch.Tensor` — default: 0

Last-token count or tensor of sequence indices to project.

---

**`output_hidden_states`** `bool` — default: False

Return final hidden states for native fused CE.

---

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

Optional tensor of shape \[batch, local_sequence] for serial HF-style shifted CE.

---

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

Optional \[batch, local_sequence] document IDs from native packing.

---

**`**kwargs`** `ConfigValue | torch.Tensor` — default: \{}

Recipe metadata. BSHD inputs are used even for concatenated documents.

---

**Returns:** `CausalLMOutputWithPast | torch.Tensor`

Output with logits \[batch, selected_sequence, vocab], optional final

```python
nemo_automodel.components.models.step5.model.Step5ForCausalLM.from_config(
    config: nemo_automodel.components.models.step5.configuration.Step5Config | nemo_automodel.components.models.step5.configuration.Step5VLConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
) -> nemo_automodel.components.models.step5.model.Step5ForCausalLM
```

classmethod

Construct a model through the AutoModel registry.

```python
nemo_automodel.components.models.step5.model.Step5ForCausalLM.get_input_embeddings() -> torch.nn.Embedding
```

Return the language embedding module.

```python
nemo_automodel.components.models.step5.model.Step5ForCausalLM.get_output_embeddings() -> torch.nn.Module
```

Return the untied language-model head.

```python
nemo_automodel.components.models.step5.model.Step5ForCausalLM.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]
```

Declare stage I/O without rounding the FP32 residual at a pipeline boundary.

**Parameters:**

**`is_first`** `bool`

Whether this stage receives token IDs.

---

**`microbatch_size`** `int`

Number of sequences in each pipeline microbatch.

---

**`seq_len`** `int`

Local sequence length after any context-parallel sharding.

---

**`dtype`** `torch.dtype`

Compute dtype for the final normalized hidden states or logits.

---

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

Input and output metadata tuples. First-stage input IDs have shape

```python
nemo_automodel.components.models.step5.model.Step5ForCausalLM.init_weights(
    buffer_device: torch.device | None = None
) -> None
```

Initialize all trainable tensors for checkpoint-free construction.

```python
nemo_automodel.components.models.step5.model.Step5ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None
```

Initialize after meta allocation and apply precision-aware casting.

```python
nemo_automodel.components.models.step5.model.Step5ForCausalLM.prepare_model_inputs_for_cp(
    batch: dict[str, torch.Tensor],
    num_chunks: int = 1
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder | torch.Tensor]
```

Choose the framework's contiguous sharder.

**Parameters:**

**`batch`** `dict[str, torch.Tensor]`

Mapping with \[batch, global_sequence] token and label tensors,
optional \[batch, global_sequence, hidden] embeddings and image tensors.

---

**`num_chunks`** `int` — default: 1

Number of pipeline microbatches; each uses the same sequence sharder.

---

**Returns:** `dict[str, ContextParallelSharder | torch.Tensor]`

Mapping containing the model's ContextParallelSharder.

```python
class nemo_automodel.components.models.step5.model.Step5ForConditionalGeneration(
    config: nemo_automodel.components.models.step5.configuration.Step5VLConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)
```

**Bases:** [Step5ForCausalLM](#nemo_automodel-components-models-step5-model-Step5ForCausalLM)

Step-5 language model with the compatible Step perception encoder.

**`vision_model`** `= Step5VisionEncoder(config.vision_config)`

---

**`vit_large_projector`**

---

```python
nemo_automodel.components.models.step5.model.Step5ForConditionalGeneration._encode_images(
    pixels: torch.Tensor,
    compute_dtype: torch.dtype
) -> torch.Tensor
```

Encode overview or crop images.

**Parameters:**

**`pixels`** `torch.Tensor`

Tensor of shape \[images, channels, height, width].

---

**`compute_dtype`** `torch.dtype`

Dtype of the completed language embedding projection.

---

**Returns:** `torch.Tensor`

Tensor of shape \[images, image_tokens, hidden].

```python
nemo_automodel.components.models.step5.model.Step5ForConditionalGeneration.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    pixel_values: torch.Tensor | None = None,
    patch_pixel_values: torch.Tensor | None = None,
    num_patches: torch.Tensor | list[int] | None = None,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue | torch.Tensor = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast
```

Inject projected crop/overview tokens and run the text backbone.

**Parameters:**

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

Tensor of shape \[batch, local_sequence].

---

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

Optional tensor of shape \[batch, local_sequence, hidden].

---

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

Optional tensor of shape \[images, channels, height, width].

---

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

Optional tensor of shape \[crops, channels, height, width].

---

**`num_patches`** `torch.Tensor | list[int] | None` — default: None

Crop counts of shape \[images], or equivalent list.

---

**`**kwargs`** `ConfigValue | torch.Tensor` — default: \{}

Tensor and scalar arguments from Step5ForCausalLM.forward.

---

**Returns:** `CausalLMOutputWithPast`

Language output with logits of shape \[batch, local_sequence, vocab].

```python
nemo_automodel.components.models.step5.model.Step5ForConditionalGeneration.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None
```

Initialize both towers after meta allocation, then apply compute dtype.

```python
class nemo_automodel.components.models.step5.model.Step5TextModel(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    dtype: torch.dtype
)
```

**Bases:** `Module`

Embedding, decoder stack and final normalization.

**`csa_backward_full_blocks`**

---

**`csa_forward_full_blocks`**

---

**`embed_tokens`**

---

**`layers`**

---

**`norm`**

---

**`swa_varlen`**

---

```python
nemo_automodel.components.models.step5.model.Step5TextModel.forward(
    embeddings: torch.Tensor,
    positions: torch.Tensor,
    padding_mask: torch.Tensor,
    packed_documents: torch.Tensor | None = None
) -> torch.Tensor
```

Evaluate the decoder stack with FP32 residuals.

**Parameters:**

**`embeddings`** `torch.Tensor`

Tensor of shape \[batch, local_sequence, hidden].

---

**`positions`** `torch.Tensor`

Tensor of shape \[batch, local_sequence].

---

**`padding_mask`** `torch.Tensor`

Boolean tensor of shape \[batch, local_sequence], true for padding.

---

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

Optional independent document IDs \[batch, local_sequence],
positive for real tokens and zero for padding.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, local_sequence, hidden]: FP32 residuals on

```python
nemo_automodel.components.models.step5.model.ModelClass = Step5ForCausalLM
```