Llama (Bidirectional) for Embedding#

NeMo AutoModel provides a bidirectional variant of Meta’s Llama for embedding and dense retrieval tasks. Unlike the standard causal (left-to-right) Llama used for text generation, this variant uses bidirectional attention, so each token can attend to both past and future tokens in the sequence, producing richer representations for semantic similarity and dense retrieval.

For the cross-encoder variant, see Llama (Bidirectional) for Reranking.

Tasks

Embedding, Dense Retrieval

Architecture

LlamaBidirectionalModel

Parameters

1B – 8B

HF Org

meta-llama

Available Models#

Any Llama checkpoint can be loaded as a bidirectional backbone. The following configurations are tested:

  • Llama 3.2 1B — fast iteration, fits on a single GPU

  • Llama 3.1 8B — higher-quality embeddings for production use

Embedding Models#

The bidirectional bi-encoder path is used for embedding generation and dense retrieval.

Architecture

Task

Auto Class

Description

LlamaBidirectionalModel

Embedding

NeMoAutoModelBiEncoder

Bidirectional Llama with mean pooling for dense embeddings

Pooling Strategies#

The bi-encoder supports multiple pooling strategies to aggregate token representations into a single embedding vector:

Strategy

Description

avg

Average of all token hidden states (default)

cls

First token hidden state

last

Last non-padding token hidden state

weighted_avg

Weighted average of token hidden states

Example HF Models#

Model

HF ID

Llama 3.2 1B

meta-llama/Llama-3.2-1B

Llama 3.1 8B

meta-llama/Llama-3.1-8B

Example Recipes#

Recipe

Description

llama3_2_1b.yaml

Bi-encoder — Llama 3.2 1B embedding model

llama_embed_nemotron_8b.yaml

Bi-encoder — reproduction recipe for nvidia/llama-embed-nemotron-8b (uses nvidia/embed-nemotron-dataset-v1)

Try with NeMo AutoModel#

1. Install NeMo AutoModel. Refer to the (Installation Guide) for information:

uv pip install nemo-automodel

2. Clone the repo to get the example recipes:

git clone https://github.com/NVIDIA-NeMo/Automodel.git
cd Automodel

3. Run the recipe from inside the repo:

torchrun --nproc-per-node=8 examples/retrieval/bi_encoder/finetune.py --config examples/retrieval/bi_encoder/llama3_2_1b.yaml
Run with Docker

1. Pull the container and mount a checkpoint directory:

docker run --gpus all -it --rm \
  --shm-size=8g \
  -v $(pwd)/checkpoints:/opt/Automodel/checkpoints \
  nvcr.io/nvidia/nemo-automodel:26.04.00

2. Navigate to the AutoModel directory (where the recipes are):

cd /opt/Automodel

3. Run the recipe:

torchrun --nproc-per-node=8 examples/retrieval/bi_encoder/finetune.py --config examples/retrieval/bi_encoder/llama3_2_1b.yaml

See the Installation Guide.

Hugging Face Model Card#

NVIDIA trained and released the Llama Nemotron Embedding 1B model, which leverages a bidirectional attention mechanism for multilingual and cross-lingual question–answer retrieval. The model supports long documents (up to 8,192 tokens) and dynamic embedding sizes via Matryoshka embeddings. For more details, see the model card on Hugging Face.