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):

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`	Multi-headed attention from ‘Attention Is All You Need’ paper.
`LlamaDecoderLayer`	Single Llama decoder layer with RMSNorm, attention, and MLP.
`LlamaForCausalLM`	Llama model with causal language modeling head.
`LlamaMLP`	SwiGLU MLP with separate gate_proj and up_proj — identical to HuggingFace default.
`LlamaModel`	Llama transformer model (embeddings + decoder layers + norm).
`LlamaPreTrainedModel`	An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained

Data

ModelClass

check_model_inputs

API

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.

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

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]

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

Optional[torch.Tensor]

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

past_key_values

Optional[Cache]Defaults to None

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

cache_position

Optional[torch.LongTensor]Defaults to None

Optional BSHD cache positions [S].

**kwargs

Unpack[TransformersKwargs]Defaults to {}

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

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

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

class nemo_automodel.components.models.llama.model.LlamaForCausalLM(
    config: transformers.LlamaConfig,
    backend: typing.Optional[nemo_automodel.components.models.common.BackendConfig] = None
)

Bases: HFCheckpointingMixin, 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

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:

input_ids

Optional[torch.LongTensor]Defaults to None

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

attention_mask

Optional[torch.Tensor]Defaults to None

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

position_ids

Optional[torch.LongTensor]Defaults to None

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

past_key_values

Optional[Cache]Defaults to None

Optional BSHD KV cache; unsupported for THD.

inputs_embeds

Optional[torch.FloatTensor]Defaults to None

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

labels

Optional[torch.LongTensor]Defaults to None

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

use_cache

Optional[bool]Defaults to None

Whether to update the BSHD KV cache.

output_attentions

Optional[bool]Defaults to None

Whether to request attention outputs.

output_hidden_states

Optional[bool]Defaults to None

Whether to return per-layer hidden states.

return_dict

Optional[bool]Defaults to None

Whether to return CausalLMOutputWithPast.

cache_position

Optional[torch.LongTensor]Defaults to None

Optional BSHD cache positions [S].

logits_to_keep

Union[int, torch.Tensor]Defaults to 0

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

**kwargs

Unpack[TransformersKwargs]Defaults to {}

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

nemo_automodel.components.models.llama.model.LlamaForCausalLM.from_config(
    config: transformers.LlamaConfig,
    backend: typing.Optional[nemo_automodel.components.models.common.BackendConfig] = None,
    kwargs = {}
)

classmethod

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

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

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

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

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

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

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

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

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

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

Bases: 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

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

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

Parameters:

input_ids

Optional[torch.LongTensor]Defaults to None

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

attention_mask

Optional[torch.Tensor]Defaults to None

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

position_ids

Optional[torch.LongTensor]Defaults to None

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

past_key_values

Optional[Cache]Defaults to None

Optional BSHD generation cache; unsupported for THD.

inputs_embeds

Optional[torch.FloatTensor]Defaults to None

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

use_cache

Optional[bool]Defaults to None

Whether to update the BSHD KV cache.

output_attentions

Optional[bool]Defaults to None

Whether to request attention outputs.

output_hidden_states

Optional[bool]Defaults to None

Whether to retain per-layer hidden states.

return_dict

Optional[bool]Defaults to None

Whether to return BaseModelOutputWithPast.

cache_position

Optional[torch.LongTensor]Defaults to None

Optional BSHD cache positions [S].

**kwargs

Unpack[TransformersKwargs]Defaults to {}

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

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'

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

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