> 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.llama.model

Custom Llama model implementation for NeMo Automodel.

This module provides a self-contained Llama implementation following HuggingFace's
implementation. Uses separate q\_proj/k\_proj/v\_proj and gate\_proj/up\_proj (HF-style).

Example (YAML):

```python
model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.3-70B-Instruct
```

## Module Contents

### Classes

| Name                                                                                         | Description                                                                                                      |
| -------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------- |
| [`LlamaAttention`](#nemo_automodel-components-models-llama-model-LlamaAttention)             | Multi-headed attention from 'Attention Is All You Need' paper.                                                   |
| [`LlamaDecoderLayer`](#nemo_automodel-components-models-llama-model-LlamaDecoderLayer)       | Single Llama decoder layer with RMSNorm, attention, and MLP.                                                     |
| [`LlamaForCausalLM`](#nemo_automodel-components-models-llama-model-LlamaForCausalLM)         | Llama model with causal language modeling head.                                                                  |
| [`LlamaMLP`](#nemo_automodel-components-models-llama-model-LlamaMLP)                         | SwiGLU MLP with separate gate\_proj and up\_proj -- identical to HuggingFace default.                            |
| [`LlamaModel`](#nemo_automodel-components-models-llama-model-LlamaModel)                     | Llama transformer model (embeddings + decoder layers + norm).                                                    |
| [`LlamaPreTrainedModel`](#nemo_automodel-components-models-llama-model-LlamaPreTrainedModel) | An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained |

### Data

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

[`check_model_inputs`](#nemo_automodel-components-models-llama-model-check_model_inputs)

### API

```python
class nemo_automodel.components.models.llama.model.LlamaAttention(
    config: transformers.LlamaConfig,
    layer_idx: int,
    backend: 'BackendConfig' | None = None
)
```

**Bases:** `Module`

Multi-headed attention from 'Attention Is All You Need' paper.

Uses separate q\_proj / k\_proj / v\_proj -- identical to the default
HuggingFace Llama implementation.

**`attention_dropout`** `= config.attention_dropout`

---

**`backend`** `= backend or BackendConfig()`

---

**`head_dim`**

---

**`k_proj`**

---

**`num_key_value_groups`**

---

**`o_proj`**

---

**`q_proj`**

---

**`rope_backend`** `= self.backend.rope`

---

**`rope_fusion`** `= self.backend.rope_fusion`

---

**`scaling`** `= self.head_dim ** -0.5`

---

**`v_proj`**

---

```python
nemo_automodel.components.models.llama.model.LlamaAttention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    cache_position: torch.LongTensor | None = None,
    kwargs: transformers.processing_utils.Unpack[transformers.utils.TransformersKwargs] = {}
) -> tuple[torch.Tensor, torch.Tensor]
```

Run dense attention over padded BSHD or packed THD hidden states.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Hidden states `[B, S, H]` or packed local states
`[T, H]`. `B` is batch, `S` is sequence, `T` is local
total tokens, and `H` is hidden size.

---

**`position_embeddings`** `tuple[torch.Tensor, torch.Tensor]`

RoPE tensors `(cos, sin)` for SDPA/eager or
`(cos, sin, freqs_cis)` for fused TE RoPE. Local cosine/sine
tensors follow `[B, S, D]` or `[T, D]`; `freqs_cis` is
global `[S, 1, 1, D]`.

---

**`attention_mask`** `torch.Tensor | None`

Padded attention mask for BSHD. THD requires
`None` and uses cumulative document lengths from `kwargs`.

---

**`past_key_values`** `Cache | None` — default: None

Optional KV cache for BSHD generation. THD training
does not support a cache.

---

**`cache_position`** `torch.LongTensor | None` — default: None

Optional BSHD cache positions `[S]`.

---

**`**kwargs`** `Unpack[TransformersKwargs]` — default: \{}

THD requires `qkv_format='thd'` and `cu_seqlens`
`[N + 1]`; CP additionally supplies `cp_size` and `cp_rank`.

---

**Returns:** `torch.Tensor`

Attention output shaped like `hidden_states` and optional BSHD

```python
class nemo_automodel.components.models.llama.model.LlamaDecoderLayer(
    config: transformers.LlamaConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `GradientCheckpointingLayer`

Single Llama decoder layer with RMSNorm, attention, and MLP.

Inherits from GradientCheckpointingLayer for efficient activation checkpointing.

**`hidden_size`** `= config.hidden_size`

---

**`input_layernorm`**

---

**`mlp`** `= LlamaMLP(config=config)`

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

```python
nemo_automodel.components.models.llama.model.LlamaDecoderLayer.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    use_cache: bool | None = False,
    cache_position: torch.LongTensor | None = None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    kwargs: transformers.processing_utils.Unpack[transformers.utils.TransformersKwargs] = {}
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.llama.model.LlamaForCausalLM(
    config: transformers.LlamaConfig,
    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), [LlamaPreTrainedModel](#nemo_automodel-components-models-llama-model-LlamaPreTrainedModel)

Llama model with causal language modeling head.

**`_pp_plan`** `= {'lm_head': (['hidden_states'], ['logits'])}`

---

**`_tied_weights_keys`** `= {'lm_head.weight': 'model.embed_tokens.weight'}`

---

**`_tp_plan`** `= {'lm_head': 'colwise_rep'}`

---

**`backend`** `= backend or BackendConfig()`

---

**`lm_head`**

---

**`model`** `= LlamaModel(config=config, backend=(self.backend))`

---

**`state_dict_adapter`** `= LlamaStateDictAdapter(config=(self.config))`

---

**`tie_word_embeddings_support`** `TieSupport = TieSupport.BOTH`

---

**`vocab_size`** `= config.vocab_size`

---

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.forward(
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    labels: torch.LongTensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = None,
    cache_position: torch.LongTensor | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    kwargs: transformers.processing_utils.Unpack[transformers.utils.TransformersKwargs] = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast
```

Forward pass returning CausalLMOutputWithPast.

**Parameters:**

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

Token IDs `[B, S]` or packed local IDs `[T]`.

---

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

Optional padded attention mask. THD uses document
boundaries from `kwargs` instead.

---

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

Position IDs `[B, S]` or packed local IDs `[T]`.

---

**`past_key_values`** `Cache | None` — default: None

Optional BSHD KV cache; unsupported for THD.

---

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

Optional hidden inputs `[B, S, H]` or `[T, H]`.

---

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

Optional labels `[B, S]` or packed `[T]`.

---

**`use_cache`** `bool | None` — default: None

Whether to update the BSHD KV cache.

---

**`output_attentions`** `bool | None` — default: None

Whether to request attention outputs.

---

**`output_hidden_states`** `bool | None` — default: None

Whether to return per-layer hidden states.

---

**`return_dict`** `bool | None` — default: None

Whether to return `CausalLMOutputWithPast`.

---

**`cache_position`** `torch.LongTensor | None` — default: None

Optional BSHD cache positions `[S]`.

---

**`logits_to_keep`** `Union[int, torch.Tensor]` — default: 0

Positions to project from hidden size `H` to
vocabulary size `V`.

---

**`**kwargs`** `Unpack[TransformersKwargs]` — default: \{}

THD metadata. `cu_seqlens` is `[N + 1]` and identifies
packed-document boundaries; CP adds `cp_size` and `cp_rank`.

---

**Returns:** `CausalLMOutputWithPast`

Causal LM output with logits `[B, S, V]`. Packed THD logits are

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.from_config(
    config: transformers.LlamaConfig,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)
```

classmethod

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.get_decoder()
```

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.get_input_embeddings()
```

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.get_output_embeddings()
```

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.set_decoder(
    decoder
)
```

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.set_input_embeddings(
    value
)
```

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.set_output_embeddings(
    new_embeddings
)
```

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.tie_weights(
    _args: object = (),
    _kwargs: object = {}
) -> None
```

```python
class nemo_automodel.components.models.llama.model.LlamaMLP(
    config: transformers.LlamaConfig
)
```

**Bases:** `Module`

SwiGLU MLP with separate gate\_proj and up\_proj -- identical to HuggingFace default.

**`act_fn`** `= ACT2FN[config.hidden_act]`

---

**`down_proj`**

---

**`gate_proj`**

---

**`hidden_size`** `= config.hidden_size`

---

**`intermediate_size`** `= config.intermediate_size`

---

**`up_proj`**

---

```python
nemo_automodel.components.models.llama.model.LlamaMLP.forward(
    x: torch.Tensor
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.llama.model.LlamaModel(
    config: transformers.LlamaConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** [LlamaPreTrainedModel](#nemo_automodel-components-models-llama-model-LlamaPreTrainedModel)

Llama transformer model (embeddings + decoder layers + norm).

**`embed_tokens`**

---

**`layers`**

---

**`norm`**

---

**`padding_idx`** `= config.pad_token_id`

---

**`rotary_emb`**

---

**`vocab_size`** `= config.vocab_size`

---

```python
nemo_automodel.components.models.llama.model.LlamaModel.forward(
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = None,
    cache_position: torch.LongTensor | None = None,
    kwargs: transformers.processing_utils.Unpack[transformers.utils.TransformersKwargs] = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast
```

Run the Llama decoder in padded BSHD or packed THD layout.

**Parameters:**

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

Token IDs `[B, S]` or packed local IDs `[T]`.

---

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

Optional padded mask `[B, S]` or broadcastable
causal mask. THD ignores this mask and uses `cu_seqlens`.

---

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

Position IDs `[B, S]` or packed local IDs `[T]`.

---

**`past_key_values`** `Cache | None` — default: None

Optional BSHD generation cache; unsupported for THD.

---

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

Alternative hidden inputs `[B, S, H]` or `[T, H]`.

---

**`use_cache`** `bool | None` — default: None

Whether to update the BSHD KV cache.

---

**`output_attentions`** `bool | None` — default: None

Whether to request attention outputs.

---

**`output_hidden_states`** `bool | None` — default: None

Whether to retain per-layer hidden states.

---

**`return_dict`** `bool | None` — default: None

Whether to return `BaseModelOutputWithPast`.

---

**`cache_position`** `torch.LongTensor | None` — default: None

Optional BSHD cache positions `[S]`.

---

**`**kwargs`** `Unpack[TransformersKwargs]` — default: \{}

THD metadata including `qkv_format`, `cu_seqlens`,
`max_seqlen`, `cp_size`, and `cp_rank`.

---

**Returns:** `BaseModelOutputWithPast`

Decoder output with final states `[B, S, H]` or packed local

```python
class nemo_automodel.components.models.llama.model.LlamaPreTrainedModel()
```

**Bases:** `PreTrainedModel`

An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
models.

**`_can_record_outputs`**

---

**`_no_split_modules`** `= ['LlamaDecoderLayer']`

---

**`_skip_keys_device_placement`** `= ['past_key_values']`

---

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

---

```python
nemo_automodel.components.models.llama.model.ModelClass = LlamaForCausalLM
```

```python
nemo_automodel.components.models.llama.model.check_model_inputs = get_check_model_inputs_decorator()
```