Retrieval Fine-Tuning (Bi-Encoder and Cross-Encoder)

Introduction

Retrieval fine-tuning adapts a model for search, retrieval-augmented generation (RAG), semantic similarity, and reranking. NeMo AutoModel provides two retrieval fine-tuning recipes:

• Bi-encoder fine-tuning trains one encoder to produce query and passage embeddings. Use it when you need fast nearest-neighbor search over a document index.
• Cross-encoder fine-tuning trains a reranker that scores a query and passage together. Use it after a retriever has produced a shortlist and you want stronger ranking quality.

Both recipes use retrieval examples where the first passage is positive and the remaining passages are negatives. A common workflow is to train a bi-encoder, use it to mine harder negatives, then train either a stronger bi-encoder or a cross-encoder reranker.

Workflow Overview

Most retrieval projects move through the same loop:

Prepare retrieval data
  -> Train a bi-encoder
  -> Validate candidate-group ranking quality
  -> Mine hard negatives
  -> Retrain the bi-encoder or train a cross-encoder reranker
  -> Use the consolidated checkpoint for indexing, mining, reranking, or serving

Start with a bi-encoder when you need embeddings for approximate nearest-neighbor search. Add hard-negative mining after the first pass if the model mostly sees easy negatives. Train a cross-encoder when a separate retriever already produces a small candidate set and you want a stronger reranking stage.

Quickstart

Before running the examples:

• Use an AutoModel environment with the full GPU training dependencies installed. For multi-GPU runs, use the NGC container. For source checkouts, refer to Installation and run uv sync --locked --all-groups --extra all.
• Run the commands from a source checkout or an NGC container workspace that contains the repository examples/ tree. The YAML configs and mining helpers use repository-relative paths. For an installed package without the repository, place the referenced files in your project and update the paths.
• From a source checkout, use uv run automodel .... From an installed environment that has local copies of the configs, use automodel ....
• Accept the access terms for the configured Hugging Face model and set HF_TOKEN, or select a model that your environment can download. Refer to the support matrix before changing model families.
• Make sure every rank can read the dataset paths or hf:// sources.

The examples use automodel. From a source checkout, prefix these commands with uv run. For direct torchrun commands, use uv run torchrun ... from a source checkout or activate an installed environment first.

Run a single-GPU basic verification test first. The timestamped checkpoint directory keeps this command from silently resuming or appending metrics to an older run:

$
RUN_ID=$(date +%Y%m%d-%H%M%S)
$
automodel examples/retrieval/bi_encoder/llama3_2_1b.yaml --nproc-per-node 1 \
>
  --checkpoint.checkpoint_dir ./output/retrieval_smoke_${RUN_ID}/checkpoints \
>
  --dist_env.timeout_minutes 30 \
>
  --step_scheduler.global_batch_size 4 \
>
  --step_scheduler.local_batch_size 1 \
>
  --step_scheduler.max_steps 10 \
>
  --dataset.max_train_samples 40

max_train_samples shortens the training rows after the configured hf:// split and corpus are loaded. This is a training-step verification, not a lightweight data-loading check. First-run downloads and corpus loading can still take time and disk. For a small data-loading check, point the config at a local retrieval JSON or JSONL sample.

The first artifact to check is training.jsonl under checkpoint.checkpoint_dir. JSONL metrics are buffered, so standard output and error are still the best live signals during a very short run.

Scale the Llama 3.2 1B bi-encoder example to the GPUs on your machine:

$
RUN_ID=$(date +%Y%m%d-%H%M%S)
$
automodel examples/retrieval/bi_encoder/llama3_2_1b.yaml --nproc-per-node 8 \
>
  --checkpoint.checkpoint_dir ./output/llama3_2_1b_encoder_${RUN_ID}/checkpoints

Run the matching cross-encoder example:

$
RUN_ID=$(date +%Y%m%d-%H%M%S)
$
automodel examples/retrieval/cross_encoder/llama3_2_1b.yaml --nproc-per-node 8 \
>
  --checkpoint.checkpoint_dir ./output/llama3_2_1b_cross_encoder_${RUN_ID}/checkpoints

Adjust --nproc-per-node to the number of GPUs on your machine. The examples use Fully Sharded Data Parallel version 2 (FSDP2) and bfloat16 by default. The example scheduler uses global_batch_size: 128 and local_batch_size: 4, so GPU counts that do not divide 32 need an explicit --step_scheduler.global_batch_size override. For example, a 6-GPU run can use --step_scheduler.global_batch_size 120 or another multiple of 4 * 6.

Choose a Recipe

Choose the recipe that matches your retrieval stage:

The following table compares the recipe components:

The bi-encoder computes a query embedding and passage embeddings independently. The cross-encoder formats each query-passage pair into one sequence and predicts a score for each candidate passage.

The following table summarizes the supported model families and their effective retrieval keyword arguments:

Known model types with a registry entry fail fast when the requested retrieval task is unsupported rather than falling back silently. For example, the legacy ministral3_bidirec type supports bi-encoder embeddings but not cross-encoder scoring. The stock ministral3 type is not a custom registry entry, so it follows the Hugging Face fallback and supports both Ministral3Model embeddings and Ministral3ForSequenceClassification scoring. If you are extracting a text tower from a parent checkpoint, set model.extract_submodel: language_model. Extracted text backbones use the extraction path, where supported extracted types use registered retrieval classes and other extracted types can fall back to Hugging Face sequence classification for cross-encoder scoring.

Treat unregistered decoder-only fallback models as an architecture experiment, not just a drop-in model swap. AutoModel sets config.is_causal: false for embedding fallbacks, while registered retrieval backbones such as the Llama bidirectional path use retrieval-specific implementations. Confirm that an unregistered architecture honors the is_causal flag and produces the expected bidirectional attention before relying on it for symmetric retrieval.

Prepare Data

Use the retrieval dataset format described in Retrieval Dataset. Choose the data path that matches the workflow you need:

For typical text-only retrieval, load the source directly with RetrievalDatasetConfig. For full-scale or image-heavy VL retrieval, prepare normalized Arrow before requesting GPUs. Loading a large image corpus and building its dataset cache during distributed startup can leave every allocated GPU waiting and waste GPU-hours before the first training step.

The normalized preparation tool currently accepts local corpus ID-based JSON sources. Convert hf:// or inline JSONL data to that layout first, or keep loading those sources directly. Refer to Retrieval Dataset for the training config and to the retrieval data preparation tools for local and Slurm CPU commands.

The key field requirements differ by source:

neg_doc must be present for local JSON and JSONL sources. It can be [] only when n_passages: 1. When n_passages > 1, provide at least one negative.

For quick custom experiments, inline JSONL is the simplest format. RetrievalDatasetConfig selects the inline loading path from the .jsonl extension. Switch to corpus ID-based JSON before hard-negative mining or full-corpus evaluation:

1
{"query":"What does NVLink do?","pos_doc":"NVLink is a high-bandwidth GPU interconnect.","neg_doc":["CUDA is a programming model.","Tensor Cores accelerate matrix math."]}
2
{"query":"What is retrieval augmented generation?","pos_doc":"RAG grounds generation by retrieving relevant context.","neg_doc":["Beam search expands candidate tokens.","Dropout regularizes training."]}

To migrate inline data to corpus ID JSON, assign a stable document ID to each unique passage, write those passages into a corpus split with id and text columns, then replace inline pos_doc and neg_doc strings with those IDs. Keep all known positives for each query in your query relevance judgments (qrels) or source metadata, even if each training row uses only the first positive. Otherwise, in-batch negatives and mined hard negatives can accidentally treat another relevant passage as a negative. The detailed source schemas and conversion rules are in Retrieval Dataset.

For larger corpora, use the corpus ID-based JSON format from the dataset guide. The same nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig also loads hf:// sources that already follow the AutoModel retrieval schema, such as:

1
data_dir_list:
2
  - hf://nvidia/embed-nemotron-dataset-v1/FEVER
3
  - hf://nvidia/embed-nemotron-dataset-v1/SyntheticClassificationData

n_passages controls the size of each query group. For example, n_passages: 5 means one positive and four negatives. All supported sources use pos_doc[0] by default. Set cycle_positive_docs: true to rotate the supervised positive across epochs when a record has multiple positives. Negatives are taken from neg_doc in order. If a record has fewer negatives than requested, they are repeated cyclically to fill the group. Treat repetition as a fallback for shape compatibility. Prefer enough distinct negatives.

The training recipe does not load a separate qrels file. Materialize qrels into retrieval records before fine-tuning. For mining, keep every known positive for the query in pos_doc so the miner can exclude those IDs. It does not read an external qrels file. If you expand multi-positive queries into one row per positive, make sure sibling positives are removed from neg_doc and audited out of mined negatives before training. For a multi-positive input row, examples/retrieval/data_utils/unroll_pos_docs.py emits only a suffixed question_id plus question, corpus_id, pos_doc, and neg_doc. It does not add original_question_id. Keep a separate mapping from each suffixed ID to its source question if you need that lineage. The current miner also normalizes rows to the required retrieval fields and does not preserve extra row metadata. Keep sibling-positive rows out of the same in-batch-negative training batch, disable distributed in-batch negatives, or add qrels-aware sampling or masking. Keep the original qrels for offline recall at K (Recall@K), mean reciprocal rank at K (MRR@K), and normalized discounted cumulative gain at K (nDCG@K) evaluation.

Minimal Config Anatomy

This minimal bi-encoder config shows the pieces that must be present in a runnable retrieval fine-tuning job. The sections below explain the model-specific parts in more detail.

1
recipe: TrainBiEncoderRecipe
2
seed: 42
3
temperature: 0.02
4
5
step_scheduler:
6
  global_batch_size: 4
7
  local_batch_size: 1
8
  max_steps: 10
9
  ckpt_every_steps: 10
10
  val_every_steps: 10
11
  num_epochs: 1
12
13
dist_env:
14
  backend: nccl
15
  timeout_minutes: 30
16
17
model:
18
  _target_: nemo_automodel.NeMoAutoModelBiEncoder.from_pretrained
19
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
20
  pooling: avg
21
  l2_normalize: true
22
  torch_dtype: bfloat16
23
24
tokenizer:
25
  _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained
26
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
27
  add_eos_token: false
28
29
dataset:
30
  _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig
31
  model_type: bi_encoder
32
  data_dir_list:
33
    - /path/to/train.json
34
  data_type: train
35
  n_passages: 5
36
  seed: 42
37
  do_shuffle: true
38
  max_train_samples: 40
39
40
dataloader:
41
  collate_fn:
42
    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
43
    q_max_len: 512
44
    p_max_len: 512
45
    query_prefix: "query:"
46
    passage_prefix: "passage:"
47
    pad_to_multiple_of: 8
48
  shuffle: true
49
  num_workers: 0
50
51
optimizer:
52
  _target_: torch.optim.AdamW
53
  lr: 5.0e-6
54
  weight_decay: 0.01
55
56
lr_scheduler:
57
  lr_warmup_steps: 2
58
  lr_decay_style: linear
59
60
checkpoint:
61
  enabled: true
62
  checkpoint_dir: ./output/llama3_2_1b_encoder/checkpoints
63
  model_save_format: safetensors
64
  save_consolidated: true
65
66
distributed:
67
  strategy: fsdp2
68
  dp_size: none
69
  tp_size: 1
70
  cp_size: 1

For a cross-encoder, change recipe, model._target_, dataset.model_type, and dataloader.collate_fn to the cross-encoder values in the following example. Also set model.num_labels: 1, set the loss temperature under model.temperature, replace q_max_len and p_max_len with rerank_max_length in the collator, and use a separate checkpoint.checkpoint_dir such as ./output/llama3_2_1b_cross_encoder/checkpoints.

Configure a Bi-Encoder

A bi-encoder config has four important parts: the model, tokenizer, retrieval dataset, and BiEncoderCollator. This example is an excerpt. Keep the scheduler, optimizer, checkpoint, and distributed sections from the full config or one of the examples.

1
recipe: TrainBiEncoderRecipe
2
3
temperature: 0.02
4
5
model:
6
  _target_: nemo_automodel.NeMoAutoModelBiEncoder.from_pretrained
7
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
8
  pooling: avg
9
  l2_normalize: true
10
  torch_dtype: bfloat16
11
12
tokenizer:
13
  _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained
14
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
15
  add_eos_token: false
16
17
dataset:
18
  _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig
19
  model_type: bi_encoder
20
  data_dir_list:
21
    - /path/to/train.json
22
  data_type: train
23
  n_passages: 5
24
  seed: 42
25
26
dataloader:
27
  collate_fn:
28
    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
29
    q_max_len: 512
30
    p_max_len: 512
31
    query_prefix: "query:"
32
    passage_prefix: "passage:"
33
    pad_to_multiple_of: 8
34
  shuffle: true
35
  num_workers: 0

The following settings control bi-encoder behavior:

• pooling: controls how token hidden states become one embedding. Common single-vector choices are avg, cls, last, and weighted_avg. The shipped hard-negative miner does not expose a pooling override. When a checkpoint contains supported Sentence Transformers metadata, the miner restores the saved pooling mode. Otherwise, the loader uses the backbone configuration or avg. Do not use the miner with colbert pooling, which returns token-level embeddings.
• l2_normalize: normalizes embeddings before scoring. When enabled, the recipe divides scores by temperature. The shipped miner does not expose a normalization override. When a checkpoint contains supported Sentence Transformers metadata, the miner restores the saved normalization setting. Otherwise, the loader enables L2 normalization.
• q_max_len and p_max_len: set separate truncation lengths for queries and passages.
• query_prefix and passage_prefix: add task-specific text before tokenization. Keep these aligned between training, hard-negative mining, and inference.
• do_distributed_inbatch_negative: optional model setting that treats passages from other data-parallel ranks as additional negatives. Enable it with model.do_distributed_inbatch_negative: true or the CLI override --model.do_distributed_inbatch_negative true. Today it all-gathers over the default process group, so use it only for data-parallel FSDP2 retrieval runs (tp_size: 1, cp_size: 1). Reliable same-document masking requires stable doc_id fields from corpus-backed or custom datasets. Pure inline JSONL produces empty IDs, and the collator omits same-document masking if any document ID in the batch is missing or incomplete. The collator hashes raw document IDs across the full batch, so make IDs globally unique when a run combines multiple corpora. For multi-positive queries expanded into separate rows, keep distributed in-batch negatives disabled unless your sampler or masking prevents sibling positives from becoming negatives.

The complete example is examples/retrieval/bi_encoder/llama3_2_1b.yaml.

Configure a Cross-Encoder

A cross-encoder config uses the same retrieval dataset config, but sets model_type: cross_encoder and uses CrossEncoderCollator. The dataset transform flattens each query with its positive and negative passages so the model scores each query-passage pair. This example is an excerpt. Keep the same scheduler, optimizer, checkpoint, and distributed structure as the bi-encoder config.

1
recipe: TrainCrossEncoderRecipe
2
3
model:
4
  _target_: nemo_automodel.NeMoAutoModelCrossEncoder.from_pretrained
5
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
6
  num_labels: 1
7
  pooling: avg
8
  temperature: 1.0
9
  torch_dtype: bfloat16
10
11
tokenizer:
12
  _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained
13
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
14
  add_eos_token: false
15
16
dataset:
17
  _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig
18
  model_type: cross_encoder
19
  data_dir_list:
20
    - /path/to/train.json
21
  data_type: train
22
  n_passages: 5
23
  seed: 42
24
25
dataloader:
26
  collate_fn:
27
    _target_: nemo_automodel.components.datasets.llm.CrossEncoderCollator
28
    rerank_max_length: 512
29
    prompt_template: "question:{query} \n \n passage:{passage}"
30
    pad_to_multiple_of: 8
31
  shuffle: true
32
  num_workers: 0

The following settings control cross-encoder behavior:

• rerank_max_length: maximum combined query-passage sequence length.
• prompt_template: controls how the pair is serialized before tokenization. It must include {query} and {passage}.
• n_passages: number of candidates per query. The positive passage must remain first in each group because labels point to index 0.

The complete example is examples/retrieval/cross_encoder/llama3_2_1b.yaml.

Distributed Launch and Batch Size

Launch single-node examples with automodel <config.yaml> --nproc-per-node <gpus>. The documented retrieval configuration uses data-parallel FSDP2. Keep tp_size: 1 and cp_size: 1. The recipes raise an error when pipeline parallelism is enabled.

For multi-node runs, launch with your cluster launcher or an external torchrun command so every node has an explicit rank and rendezvous endpoint:

$
uv run torchrun \
>
  --nnodes 2 \
>
  --nproc-per-node 8 \
>
  --node-rank ${NODE_RANK} \
>
  --rdzv-backend c10d \
>
  --rdzv-endpoint ${MASTER_ADDR}:${MASTER_PORT} \
>
  -m nemo_automodel.cli.app examples/retrieval/bi_encoder/llama3_2_1b.yaml

Mount dataset and Hugging Face cache paths at the same location on every node. If first-time data loading is expensive, populate the cache before requesting GPUs. Use a unique checkpoint_dir for each experiment. For multi-node training, checkpoint_dir must be on a persistent filesystem mounted at the same path from every node. Relative ./output/... paths are appropriate only when they resolve to that storage. Increase dist_env.timeout_minutes for first model downloads, slow filesystems, multi-node collectives, or large checkpoint writes.

The step scheduler computes gradient accumulation from:

gradient_accumulation_steps = global_batch_size / (local_batch_size * data_parallel_size)

global_batch_size must be divisible by local_batch_size * data_parallel_size, and the result must be at least 1. For the supported data-parallel FSDP2 configuration, data_parallel_size is the total GPU count. For example, two 8-GPU nodes have data_parallel_size: 16.

For Distributed Data Parallel (DDP) recipes, keep distributed.static_graph: false whenever gradient_accumulation_steps > 1. If the model contains parameters that do not contribute to the loss, also set distributed.find_unused_parameters: true when the static graph is disabled. Recalculate the accumulation steps after changing the GPU count, global_batch_size, or local_batch_size. When the result is 1 and the model follows the same execution path on every iteration, you can set distributed.static_graph: true to reduce DDP overhead.

local_batch_size is the number of query groups per rank. For memory pressure, reduce step_scheduler.local_batch_size first, then sequence lengths (q_max_len, p_max_len, or rerank_max_length), then n_passages. Bi-encoders scale memory with query length plus local_batch_size * n_passages passage sequences. Cross-encoders scale with local_batch_size * n_passages combined query-passage sequences.

Current retrieval datasets are map-style datasets loaded in each process, not streaming distributed inputs. Pre-cache Hugging Face data on each node or use a shared cache. Budget CPU RAM and local disk per rank for corpus-backed datasets.

Add Validation

Both examples include commented validation_dataset and validation_dataloader blocks. Enable them when you have a held-out retrieval file. Use RetrievalDatasetConfig for every supported source. It selects the loading path from the source URI and file extension. This corpus-backed example mirrors the shipped configs:

1
validation_dataset:
2
  _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig
3
  model_type: bi_encoder
4
  data_dir_list:
5
    - /path/to/validation.json
6
  data_type: eval
7
  n_passages: 5
8
  seed: 42
9
  do_shuffle: false
10
11
validation_dataloader:
12
  collate_fn:
13
    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
14
    q_max_len: 512
15
    p_max_len: 512
16
    query_prefix: "query:"
17
    passage_prefix: "passage:"
18
    pad_to_multiple_of: 8
19
  batch_size: 2
20
  shuffle: false
21
  num_workers: 0

Validation logs val_loss, val_acc1, and val_mrr to validation.jsonl under checkpoint.checkpoint_dir. These metrics measure ranking within each candidate group in the validation file. They are not full-corpus Recall@K or nDCG metrics. For cross-encoder validation, use model_type: cross_encoder and CrossEncoderCollator instead. In multi-rank runs, validation uses the same distributed sampler path as training and can drop tail examples to keep rank shapes even. Make the validation set divisible by data_parallel_size * validation_dataloader.batch_size for comparable per-batch val_loss, or run validation on one GPU when you need every candidate group included.

$
tail -n 5 ./output/llama3_2_1b_encoder/checkpoints/validation.jsonl

Evaluate Retrieval Quality

Candidate-group validation is a basic verification of the training objective. To decide whether a bi-encoder is useful for RAG candidate generation, evaluate against a fixed held-out corpus and qrels:

1. Encode corpus passages with the same tokenizer, pooling, normalization, passage prefix, and p_max_len used in training.
2. Build an approximate nearest-neighbor (ANN) or exact top-k index. With l2_normalize: true, use inner product or cosine similarity.
3. Encode held-out queries with the matching query prefix and q_max_len.
4. Report full-corpus Recall@K, MRR@K, and nDCG@K for the K values your application uses.

AutoModel does not currently provide a one-command full-corpus retrieval evaluator in this guide. Use your existing information retrieval evaluation stack or a small script around the consolidated checkpoint. Report enough run details to make the result repeatable. Include the query count, corpus size, qrels source, judged and unjudged handling, exact or ANN search settings, K values, and baseline checkpoint. State whether you used confidence intervals or significance tests. At minimum, specify a consolidated bi-encoder checkpoint, corpus and query tables with stable IDs, and qrels keyed by those IDs. Also specify the query and passage prefixes, maximum lengths, and K values to report.

For cross-encoders, freeze a first-stage retriever, rerank its top-K candidates, and report reranking metrics on that same candidate set. Also report first-stage candidate recall or coverage: if a query’s positive document is missing from the retriever top-K, count that query as a miss rather than dropping it from reranker evaluation. Do not compare cross-encoder candidate-group validation directly to full-corpus bi-encoder metrics.

Monitor Training

Training metrics are written to training.jsonl under checkpoint.checkpoint_dir. The file logger buffers records in chunks before writing and flushes the remaining records on close, so tail -f is useful for completed or longer runs but might not update during a short basic verification test:

$
tail -f ./output/llama3_2_1b_encoder/checkpoints/training.jsonl

Use standard output and error as the live per-step signals today. Watch loss, grad_norm, learning rate, GPU memory, and step time before scaling to a longer run. On preempted or timed-out jobs, recent buffered JSONL metrics might be missing even when standard output and error showed them.

The examples include a commented wandb block. Enable it when you want remote tracking, and tune step_scheduler.log_remote_every_steps to control remote logging cadence.

Enable Low-Rank Adaptation (LoRA)

Retrieval recipes support the same parameter-efficient fine-tuning (PEFT) block used by other AutoModel recipes. Uncomment or add peft to train low-rank adaptation (LoRA) adapters instead of updating every weight:

1
peft:
2
  _target_: nemo_automodel.components._peft.lora.PeftConfig
3
  target_modules:
4
    - q_proj
5
    - k_proj
6
    - v_proj
7
    - o_proj
8
    - gate_proj
9
    - up_proj
10
    - down_proj
11
  exclude_modules: []
12
  match_all_linear: false
13
  dim: 16
14
  alpha: 32
15
  use_triton: true

Use either match_all_linear: true or an explicit target_modules list. Setting both selection modes raises a configuration error.

Use LoRA when you need lower memory use or want to ship a small adapter. Use full fine-tuning when you can afford the memory and want maximum adaptation.

Mine Hard Negatives

After an initial bi-encoder run, use examples/retrieval/data_utils/mine_hard_negatives.py to retrieve confusing passages from a corpus. The miner accepts the corpus ID-based local JSON format described in Retrieval Dataset. It does not accept inline JSONL or hf:// sources directly. Mine one corpus-backed JSON file at a time.

The miner enforces that the input resolves to exactly one corpus, but it does not enforce every required data invariant. Validate these conditions before launch:

• Every row’s corpus_id must match the declared corpus.
• Every question_id must be unique. Duplicate IDs overwrite entries in the miner’s result lookup.
• Use an absolute corpus path, or keep the output beside the input. Relative corpus paths are copied without being rewritten when the output moves.
• Make train_file_output_path different from train_qa_file_path, and make sure the output does not already exist. The miner opens the output in write mode and overwrites an existing file without confirmation.
• Keep lineage or custom metadata in a separate table. Mined rows are rebuilt from normalized retrieval fields, so arbitrary row-level fields are not preserved.

The bi-encoder and cross-encoder quickstarts use hf:// sources. For nvidia/embed-nemotron-dataset-v1, restore the maintained local layout with:

$
uv run python examples/retrieval/bi_encoder/llama_embed_nemotron_8b/data_preparation.py \
>
  --download-path ./embed_nemotron_dataset_v1

The script has no subset selector. It scans the repository and restores every discovered directory containing a corpus.parquet. For each subset <name>, it writes ./embed_nemotron_dataset_v1/<name>/<name>.json and ./embed_nemotron_dataset_v1/<name>/corpus/train.parquet. The JSON stores the corpus directory as an absolute path. For example, use ./embed_nemotron_dataset_v1/FEVER/FEVER.json as the FEVER mining input. The maintained examples/retrieval/bi_encoder/llama_embed_nemotron_8b/llama_embed_nemotron_8b.yaml recipe consumes this same layout.

Before mining, load the resulting JSON through the retrieval dataset path or run a basic verification test to confirm the schema and document IDs. The Llama 3.2 quickstarts train from both FEVER and SyntheticClassificationData. For a complete train, mine, and retrain workflow, mine each desired subset separately. Give each output and embedding cache a run-specific path, then list all mined JSON files in the next config’s data_dir_list.

The mining script can load the consolidated checkpoint from the Llama 3.2 quickstart:

$
uv run torchrun --standalone --nproc-per-node 8 examples/retrieval/data_utils/mine_hard_negatives.py \
>
  --config examples/retrieval/data_utils/mining_config.yaml \
>
  --mining.model_name_or_path ./output/llama3_2_1b_encoder/checkpoints/epoch_0_step_499/model/consolidated/ \
>
  --mining.train_qa_file_path ./embed_nemotron_dataset_v1/FEVER/FEVER.json \
>
  --mining.train_file_output_path /path/to/retrieval-data/mined/FEVER/train.json \
>
  --mining.cache_embeddings_dir /shared/path/to/empty-cache/fever-run-001 \
>
  --mining.query_prefix "query: " \
>
  --mining.passage_prefix "passage: " \
>
  --mining.query_max_length 512 \
>
  --mining.passage_max_length 512

For multi-node mining, replace --standalone with explicit rendezvous flags. Every rank must see the model, tokenizer, input JSON, referenced corpus, and shared cache at the same paths:

$
uv run torchrun \
>
  --nnodes 2 \
>
  --nproc-per-node 8 \
>
  --node-rank ${NODE_RANK} \
>
  --rdzv-backend c10d \
>
  --rdzv-endpoint ${MASTER_ADDR}:${MASTER_PORT} \
>
  examples/retrieval/data_utils/mine_hard_negatives.py \
>
  --config examples/retrieval/data_utils/mining_config.yaml \
>
  --mining.model_name_or_path /path/to/compatible-hf-bi-encoder \
>
  --mining.train_qa_file_path ./embed_nemotron_dataset_v1/FEVER/FEVER.json \
>
  --mining.train_file_output_path /path/to/retrieval-data/mined/FEVER/train.json \
>
  --mining.cache_embeddings_dir /shared/path/to/empty-cache/fever-run-001 \
>
  --mining.query_prefix "query: " \
>
  --mining.passage_prefix "passage: " \
>
  --mining.query_max_length 512 \
>
  --mining.passage_max_length 512

The mining script loads a concrete Hugging Face model path. It does not resolve AutoModel’s LATEST keyword or pointer files. It creates the output parent directory and overwrites train_file_output_path if that file already exists. The miner preserves the input’s top-level corpus value without rewriting relative paths.

Match these mining settings to training:

• Set query_prefix and passage_prefix to the same text used for training. Include the trailing space in mining prefixes because the miner concatenates prefixes directly, while BiEncoderCollator inserts the space.
• Match query_max_length, passage_max_length, add_bos_token, and add_eos_token to the training tokenizer and truncation behavior.
• Set use_negatives_from_file when you want to keep existing negatives before appending mined negatives. Deduplicate and audit the combined list.

The miner does not accept pooling or l2_normalize settings. It calls NeMoAutoModelBiEncoder.from_pretrained without those arguments. The loader restores pooling and normalization from supported Sentence Transformers checkpoint metadata. If the metadata is unavailable, pooling uses the backbone configuration or avg, and L2 normalization is enabled. Confirm that these fallback settings match training before mining from a legacy or external checkpoint.

Use a new, empty cache_embeddings_dir for every model, input, prefix, sequence length, corpus_chunk_size, and world size. The miner reuses expected shard files even when load_embeddings_from_cache is false, and it does not record a fingerprint for the model, data, or embedding settings. Set load_embeddings_from_cache: true only when rerunning an unchanged job with both consolidated embedding archives present.

hard_negatives_to_mine sets the maximum number of mined negatives per query. The miner can return fewer when the corpus has too few non-positive candidates. Margin filtering can leave -inf entries in the top results when too few finite candidates remain. Inspect the scores and per-query counts before training. The checked-in examples/retrieval/data_utils/mining_config.yaml documents the remaining mining defaults.

Hard-negative mining distributes embedding generation, but rank 0 assembles the cache, performs final exact scoring, and materializes the full document embedding matrix. Use a smaller mining slice or a custom ANN or blockwise workflow for very large corpora.

Use the mined output as a corpus-backed data_dir_list source for another bi-encoder pass or for cross-encoder training. If the previous run used multiple sources, list each mined source file. The miner excludes IDs in pos_doc, but it cannot read external qrels or identify semantic duplicates. It also drops extra row metadata such as original_question_id.

AutoModel does not provide a mined-negative audit or cleanup utility. Before training, verify the output:

1. Every question ID is unique, every row references the declared corpus, and every document ID exists.
2. No known positive or sibling positive appears in neg_doc.
3. Negative IDs are unique within each row.
4. Positive and negative scores are present and finite. Remove entries with scores such as -inf.
5. Each row has enough distinct negatives for the next run’s n_passages.
6. The corpus path resolves from the mined file, and required lineage metadata is joined back explicitly.

Do not mine from validation or test corpora. Retain the original qrels for full-corpus evaluation.

Save and Resume from a Checkpoint

Set checkpointing in the config:

1
checkpoint:
2
  enabled: true
3
  checkpoint_dir: ./output/llama3_2_1b_encoder/checkpoints
4
  model_save_format: safetensors
5
  save_consolidated: true

Each save creates a versioned directory. The following path is an example:

./output/llama3_2_1b_encoder/checkpoints/epoch_0_step_499/

Checkpoint directory names use the scheduler step at save time. The saved scheduler state advances to the next step, so for exact paths prefer the Saving checkpoint to ... log line or the LATEST pointer over hand-constructing a step number.

With save_consolidated: true and full fine-tuning, AutoModel also writes a Hugging Face-compatible model under the following path:

./output/llama3_2_1b_encoder/checkpoints/epoch_0_step_499/model/consolidated/

Use the concrete epoch_*_step_* directory printed in your logs. Some workflows also create a LATEST symlink, but direct Hugging Face and mining loads expect a real exported model path. If your run produced LATEST.txt instead of a symlink, read that file and substitute the resolved checkpoint directory before calling from_pretrained() or mine_hard_negatives.py.

PEFT runs using LoRA save adapter artifacts under the checkpoint model/ directory instead of the previously described full consolidated export path. Resume LoRA training from the AutoModel checkpoint directory. If you need a standalone checkpoint for inference or mining, first produce a Hugging Face-loadable merged encoder with your adapter workflow.

The LATEST symlink points to the most recent checkpoint when it is valid. To resume from the latest resolved checkpoint, set:

1
checkpoint:
2
  enabled: true
3
  checkpoint_dir: ./output/llama3_2_1b_encoder/checkpoints
4
  restore_from: LATEST

LATEST is a resolver keyword: AutoModel follows the symlink or pointer file and can fall back to the highest epoch_*_step_* checkpoint directory if the pointer is not usable. An explicit restore target, including LATEST, is authoritative: AutoModel warns about an incompatibility but still restores the selected checkpoint. If checkpoint.restore_from is omitted, AutoModel selects the latest candidate and resumes only when that candidate is compatible. It does not continue scanning older checkpoints when the newest candidate is incompatible. Use a new or empty checkpoint_dir for fresh experiments, and rotate or clear training.jsonl and validation.jsonl if you do not want logs from multiple runs appended together.

When checkpoint.is_async: true, the LATEST symlink can lag the most recent write at job end. For final mining, export, or evaluation workflows, prefer the explicit epoch_*_step_* checkpoint directory or keep async checkpointing disabled for the final save.

Use the Model

Use a bi-encoder checkpoint to encode passages, build an approximate nearest-neighbor index, encode queries, and search the index. Keep the same tokenizer, pooling, normalization, prefixes, and max lengths that you used for training. These retrieval wrappers instantiate the model on the current CUDA device, so the examples require CUDA.

Pass wrapper-only settings explicitly when they differ from the defaults because the current export does not preserve all of them.

The following example loads a bi-encoder and computes one similarity score:

1
import torch
2
3
from nemo_automodel import NeMoAutoModelBiEncoder, NeMoAutoTokenizer
4
5
model_path = "./output/llama3_2_1b_encoder/checkpoints/epoch_0_step_499/model/consolidated"
6
tokenizer = NeMoAutoTokenizer.from_pretrained(model_path, add_eos_token=False)
7
model = NeMoAutoModelBiEncoder.from_pretrained(
8
    model_path,
9
    pooling="avg",
10
    l2_normalize=True,
11
    use_liger_kernel=False,
12
).eval()
13
device = next(model.parameters()).device
14
15
texts = ["query: what does nvlink do?", "passage: NVLink is a high-bandwidth GPU interconnect."]
16
tokenized = tokenizer(texts, padding=False, truncation=True, max_length=512, return_token_type_ids=False)
17
tokenized = [{key: tokenized[key][idx] for key in tokenized.keys()} for idx in range(len(texts))]
18
tokens = tokenizer.pad(tokenized, padding="longest", return_tensors="pt")
19
tokens = {key: value.to(device) for key, value in tokens.items()}
20
with torch.no_grad():
21
    embeddings = model.encode(tokens)
22
score = embeddings[0] @ embeddings[1]

The example produces one scalar tensor named score. Its value depends on the fine-tuned checkpoint.

Use a cross-encoder checkpoint to rerank a shortlist from a retriever. Cross-encoders score each query-passage pair jointly, so they are usually too expensive for first-stage full-corpus search.

The following example loads a cross-encoder and ranks two passages:

1
import torch
2
3
from nemo_automodel import NeMoAutoModelCrossEncoder, NeMoAutoTokenizer
4
5
model_path = "./output/llama3_2_1b_cross_encoder/checkpoints/epoch_0_step_499/model/consolidated"
6
tokenizer = NeMoAutoTokenizer.from_pretrained(model_path, add_eos_token=False)
7
model = NeMoAutoModelCrossEncoder.from_pretrained(
8
    model_path,
9
    pooling="avg",
10
    num_labels=1,
11
    temperature=1.0,
12
    use_liger_kernel=False,
13
).eval()
14
device = next(model.parameters()).device
15
16
prompt_template = "question:{query} \n \n passage:{passage}"
17
pairs = [
18
    prompt_template.format(query="what does nvlink do?", passage="NVLink is a high-bandwidth GPU interconnect."),
19
    prompt_template.format(query="what does nvlink do?", passage="Dropout regularizes neural networks."),
20
]
21
tokenized = tokenizer(pairs, padding=False, truncation=True, max_length=512, return_token_type_ids=False)
22
tokenized = [{key: tokenized[key][idx] for key in tokenized.keys()} for idx in range(len(pairs))]
23
tokens = tokenizer.pad(tokenized, padding="longest", return_tensors="pt")
24
tokens = {key: value.to(device) for key, value in tokens.items()}
25
with torch.no_grad():
26
    logits = model(tokens).logits.squeeze(-1)
27
ranking = torch.argsort(logits, descending=True)

For a checkpoint that learned the reranking task, ranking[0] is expected to be 0.

Bi-encoder scores are comparable only within the same model, tokenizer, prefix, max-length, pooling, normalization, and indexing setup. Mining scores are raw embedding similarities from that exact setup. Cross-encoder logits are uncalibrated reranking signals. Do not mix them with bi-encoder scores or use one global threshold across model versions without calibration.

Troubleshooting

Use the following checks to diagnose common retrieval training problems:

Related Files

Refer to the following files for implementation and configuration details:

• Bi-encoder recipe: nemo_automodel/recipes/retrieval/train_bi_encoder.py
• Cross-encoder recipe: nemo_automodel/recipes/retrieval/train_cross_encoder.py
• Retrieval dataset guide: Retrieval Dataset
• Llama-Embed-Nemotron-8B example: examples/retrieval/bi_encoder/llama_embed_nemotron_8b/llama_embed_nemotron_8b.yaml