> 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: typing.Optional['BackendConfig'] = 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.

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

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

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

```python
class nemo_automodel.components.models.llama.model.LlamaForCausalLM(
    config: transformers.LlamaConfig,
    backend: typing.Optional[nemo_automodel.components.models.common.BackendConfig] = 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.

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.forward(
    input_ids: typing.Optional[torch.LongTensor] = None,
    attention_mask: typing.Optional[torch.Tensor] = None,
    position_ids: typing.Optional[torch.LongTensor] = None,
    past_key_values: typing.Optional[transformers.cache_utils.Cache] = None,
    inputs_embeds: typing.Optional[torch.FloatTensor] = None,
    labels: typing.Optional[torch.LongTensor] = None,
    use_cache: typing.Optional[bool] = None,
    output_attentions: typing.Optional[bool] = None,
    output_hidden_states: typing.Optional[bool] = None,
    return_dict: typing.Optional[bool] = None,
    cache_position: typing.Optional[torch.LongTensor] = 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:**

(batch\_size, seq\_len)

Optional attention mask

Optional position indices

Optional cached key/values

Optional pre-computed embeddings

Optional labels for computing loss

Whether to use KV caching

Position in cache

Number of final logits to compute (0=all, N=last N tokens)

**Returns:** `CausalLMOutputWithPast`

CausalLMOutputWithPast with loss, logits, past\_key\_values

```python
nemo_automodel.components.models.llama.model.LlamaForCausalLM.from_config(
    config: transformers.LlamaConfig,
    backend: typing.Optional[nemo_automodel.components.models.common.BackendConfig] = 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.

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

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

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

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

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