> This page is for version 0.7.0 · 26.10. > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > 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. # Retrieval Fine-Tuning (Bi-Encoder and Cross-Encoder) > Fine-tune bi-encoder and cross-encoder retrieval models in NeMo AutoModel with distributed training, low-rank adaptation (LoRA), hard-negative mining, and evaluation. ## 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. > **Tip** > > For an end-to-end text bi-encoder workflow starting from your document corpus, follow the > [Nemotron embedding recipe](https://github.com/NVIDIA-NeMo/Nemotron/tree/main/docs/nemotron/embed). > It covers synthetic data generation, data preparation with hard-negative mining, fine-tuning with NeMo AutoModel, > evaluation, model export when required, and deployment. ## Workflow Overview Most retrieval projects move through the same loop: ```text 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 a NeMo AutoModel environment with the full GPU training dependencies installed. For multi-GPU runs, use the NGC container. For source checkouts, refer to [Installation](/nemo/automodel/v0.7/get-started/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: ```bash 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: ```bash 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: ```bash 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: | Need | Use | Why | | ------------------------------------------ | -------------------- | -------------------------------------------------------------------------------------- | | Search across a large corpus | Bi-encoder | Encodes queries and passages independently, so passage embeddings can be indexed once. | | RAG candidate generation | Bi-encoder | Produces dense vectors for approximate nearest-neighbor retrieval. | | Better ranking for a small shortlist | Cross-encoder | Scores each query-passage pair jointly, which is slower but usually more accurate. | | Better negatives for the next training run | Hard-negative mining | Uses a trained bi-encoder to find confusing passages for each query. | The following table compares the recipe components: | Component | Bi-Encoder | Cross-Encoder | | ------------------ | ------------------------------------------------------- | ---------------------------------------------------------- | | Recipe | `TrainBiEncoderRecipe` | `TrainCrossEncoderRecipe` | | Model target | `nemo_automodel.NeMoAutoModelBiEncoder.from_pretrained` | `nemo_automodel.NeMoAutoModelCrossEncoder.from_pretrained` | | Dataset mode | `model_type: bi_encoder` | `model_type: cross_encoder` | | Collator | `BiEncoderCollator` | `CrossEncoderCollator` | | Training objective | Cross entropy over one positive plus negatives | Cross entropy over one positive plus negatives | For Python callers, `NeMoAutoModelCrossEncoder.from_pretrained` follows the Transformers positional argument convention: arguments after the checkpoint path are forwarded to the backbone constructor. Calls that previously passed the attention implementation as the second positional argument must now use a keyword, for example `NeMoAutoModelCrossEncoder.from_pretrained(checkpoint_path, attn_implementation="sdpa")`. 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: | Model `config.model_type` | Bi-Encoder Behavior | Cross-Encoder Behavior | Effective Retrieval Keyword Arguments | | ------------------------------------------------- | ---------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `llama`, `llama_bidirec` | Uses `LlamaBidirectionalModel` with configurable text attention. | Uses `LlamaBidirectionalForSequenceClassification` with configurable text attention. | Bi-encoder: `pooling`, wrapper-level `l2_normalize`, `is_causal`, and top-level recipe `temperature`. Cross-encoder: `pooling`, `num_labels`, `temperature`, and `is_causal` on the Llama scoring config. | | `ministral3` | Uses the stock Hugging Face `Ministral3Model` with configurable text attention. | Uses the stock Hugging Face `Ministral3ForSequenceClassification` with configurable text attention. | Bi-encoder: wrapper-level `pooling`, `l2_normalize`, `is_causal`, plus top-level recipe `temperature`. Cross-encoder: forwards `num_labels` and applies `is_causal`; custom `pooling` and `temperature` are ignored. | | `ministral3_bidirec` | Uses the legacy custom `Ministral3BidirectionalModel` with configurable text attention. | Direct cross-encoder scoring is not supported by the custom retrieval registry. | Bi-encoder: `pooling`, wrapper-level `l2_normalize`, `is_causal`, and top-level recipe `temperature`. | | `llama_nemotron_vl` | Uses `LlamaNemotronVLModel` and applies the attention policy only to its language tower. | Direct cross-encoder scoring is not supported by the custom retrieval registry today. | Bi-encoder: `pooling`, wrapper-level `l2_normalize`, `is_causal`, and top-level recipe `temperature`. | | Any model type without a retrieval registry entry | Falls back to `AutoModel` and applies the resolved attention policy. | Falls back to `AutoModelForSequenceClassification` and applies the resolved attention policy. | Bi-encoder fallback receives standard Hugging Face `from_pretrained` keyword arguments. `pooling`, `l2_normalize`, and `is_causal` apply in the NeMo AutoModel wrapper. Cross-encoder fallback forwards `num_labels` and applies `is_causal`. Custom `pooling` and `temperature` are ignored. | For both retrieval wrappers, `model.is_causal` selects text self-attention. An explicit YAML or API value takes precedence over a saved `is_causal` value in the checkpoint's text config. When neither exists, a bi-encoder defaults to `false` for bidirectional retrieval, while a cross-encoder preserves the scoring backbone's native attention mode. The resolved value is saved with the text config. For composite models such as Llama Nemotron VL, the setting affects only the language tower and does not change vision attention. Changing this setting changes the resulting embeddings. Regenerate stored corpus embeddings before querying an index with a model that uses a different attention policy. 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. The retrieval wrapper applies the resolved `is_causal` value to the text config and attention modules, while registered retrieval backbones such as the Llama path provide retrieval-specific implementations. Confirm that an unregistered architecture honors the selected policy before relying on it for retrieval. ## Prepare Data Use the retrieval dataset format described in [Retrieval Dataset](/nemo/automodel/v0.7/datasets/retrieval-dataset). Choose the data path that matches the workflow you need: | Data Path | Use When | Notes | | --------------------------------------- | ------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------- | | `hf://` NeMo AutoModel retrieval schema | You want a tutorial run or shared public dataset | Requires a NeMo AutoModel-style Hugging Face subset with a companion corpus split. | | Inline JSONL | You want a small custom run without hard-negative mining | Documents are embedded directly in each record. Pure inline records do not provide document IDs for same-document masking. | | Corpus ID-based JSON | You need hard-negative mining, reusable corpora, or same-document masking | Records reference document IDs in a local corpus that can be loaded by Hugging Face `datasets`. | | Normalized Arrow | You are training on a full-scale or image-heavy VL retrieval dataset | Prepare local corpus ID-based JSON sources on CPU, then train from portable Arrow shards with `NormalizedRetrievalDatasetConfig`. | 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](/nemo/automodel/v0.7/datasets/retrieval-dataset#normalized-arrow-for-large-vl-datasets) for the training config and to the [retrieval data preparation tools](https://github.com/NVIDIA-NeMo/Automodel/blob/r0.7.0/tools/retrieval/README.md) for local and Slurm CPU commands. The key field requirements differ by source: | Source | Required Query Field | Required Document Fields | | ----------------------------- | --------------------- | ----------------------------------------------------------------------------------------------------- | | Corpus ID JSON | `question` | `question_id`, `corpus_id`, non-empty `pos_doc`, and present `neg_doc`. | | `hf://` NeMo AutoModel schema | `question` | Non-empty `pos_doc`. `neg_doc` is optional in the source but required before training with negatives. | | Inline JSONL | `query` or `question` | Non-empty `pos_doc`, and present `neg_doc`. | `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. > **Warning** > > `n_passages: 1` is useful for schema checks or custom negative strategies, but it is not a good default training setup. > The standard bi-encoder and cross-encoder recipes need at least one negative candidate for meaningful contrastive or > reranking supervision, unless you add a custom strategy such as qrels-aware in-batch negatives. 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: ```json {"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."]} {"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](/nemo/automodel/v0.7/datasets/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 NeMo AutoModel retrieval schema, such as: ```yaml data_dir_list: - hf://nvidia/embed-nemotron-dataset-v1/FEVER - 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. ```yaml recipe: TrainBiEncoderRecipe seed: 42 temperature: 0.02 step_scheduler: global_batch_size: 4 local_batch_size: 1 max_steps: 10 ckpt_every_steps: 10 val_every_steps: 10 num_epochs: 1 dist_env: backend: nccl timeout_minutes: 30 model: _target_: nemo_automodel.NeMoAutoModelBiEncoder.from_pretrained pretrained_model_name_or_path: meta-llama/Llama-3.2-1B pooling: avg l2_normalize: true torch_dtype: bfloat16 tokenizer: _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained pretrained_model_name_or_path: meta-llama/Llama-3.2-1B add_eos_token: false dataset: _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig model_type: bi_encoder data_dir_list: - /path/to/train.json data_type: train n_passages: 5 seed: 42 do_shuffle: true max_train_samples: 40 dataloader: collate_fn: _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator q_max_len: 512 p_max_len: 512 query_prefix: "query:" passage_prefix: "passage:" pad_to_multiple_of: 8 shuffle: true num_workers: 0 optimizer: _target_: torch.optim.AdamW lr: 5.0e-6 weight_decay: 0.01 lr_scheduler: lr_warmup_steps: 2 lr_decay_style: linear checkpoint: enabled: true checkpoint_dir: ./output/llama3_2_1b_encoder/checkpoints model_save_format: safetensors save_consolidated: true distributed: strategy: fsdp2 dp_size: none tp_size: 1 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`, choose model-side score scaling with `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`. The top-level `temperature` now scales the cross-encoder's listwise training and validation loss, as it does for the bi-encoder. When converting an existing cross-encoder config, set it to `1.0` to preserve the previous loss behavior. For a non-unit top-level temperature, keep `model.temperature: 1.0` so scores are not scaled twice. `model.temperature` also changes exported inference scores; the top-level temperature changes only the recipe loss. ## 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. ```yaml recipe: TrainBiEncoderRecipe temperature: 0.02 model: _target_: nemo_automodel.NeMoAutoModelBiEncoder.from_pretrained pretrained_model_name_or_path: meta-llama/Llama-3.2-1B pooling: avg l2_normalize: true torch_dtype: bfloat16 tokenizer: _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained pretrained_model_name_or_path: meta-llama/Llama-3.2-1B add_eos_token: false dataset: _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig model_type: bi_encoder data_dir_list: - /path/to/train.json data_type: train n_passages: 5 seed: 42 dataloader: collate_fn: _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator q_max_len: 512 p_max_len: 512 query_prefix: "query:" passage_prefix: "passage:" pad_to_multiple_of: 8 shuffle: true 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](https://github.com/NVIDIA-NeMo/Automodel/blob/r0.7.0/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. ```yaml recipe: TrainCrossEncoderRecipe model: _target_: nemo_automodel.NeMoAutoModelCrossEncoder.from_pretrained pretrained_model_name_or_path: meta-llama/Llama-3.2-1B num_labels: 1 pooling: avg temperature: 1.0 torch_dtype: bfloat16 tokenizer: _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained pretrained_model_name_or_path: meta-llama/Llama-3.2-1B add_eos_token: false dataset: _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig model_type: cross_encoder data_dir_list: - /path/to/train.json data_type: train n_passages: 5 seed: 42 dataloader: collate_fn: _target_: nemo_automodel.components.datasets.llm.CrossEncoderCollator rerank_max_length: 512 prompt_template: "question:{query} \n \n passage:{passage}" pad_to_multiple_of: 8 shuffle: true num_workers: 0 ``` The following settings control cross-encoder behavior: * `model.is_causal`: optional text self-attention policy. An explicit value overrides saved metadata; when omitted, the wrapper restores a saved value or preserves the scoring backbone's native mode. * `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](https://github.com/NVIDIA-NeMo/Automodel/blob/r0.7.0/examples/retrieval/cross_encoder/llama3_2_1b.yaml). ## Distributed Launch and Batch Size Launch single-node examples with `automodel --nproc-per-node `. 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: ```bash 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: ```text 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: ```yaml validation_dataset: _target_: nemo_automodel.components.datasets.llm.retrieval_dataset.RetrievalDatasetConfig model_type: bi_encoder data_dir_list: - /path/to/validation.json data_type: eval n_passages: 5 seed: 42 do_shuffle: false validation_dataloader: collate_fn: _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator q_max_len: 512 p_max_len: 512 query_prefix: "query:" passage_prefix: "passage:" pad_to_multiple_of: 8 batch_size: 2 shuffle: false 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. ```bash 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. NeMo 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: ```bash 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 NeMo AutoModel recipes. Uncomment or add `peft` to train low-rank adaptation (LoRA) adapters instead of updating every weight: ```yaml peft: _target_: nemo_automodel.components._peft.lora.PeftConfig target_modules: - q_proj - k_proj - v_proj - o_proj - gate_proj - up_proj - down_proj exclude_modules: [] match_all_linear: false dim: 16 alpha: 32 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](/nemo/automodel/v0.7/datasets/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: ```bash 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 ``, it writes `./embed_nemotron_dataset_v1//.json` and `./embed_nemotron_dataset_v1//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: ```bash 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: ```bash 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 NeMo 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`. NeMo 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: ```yaml checkpoint: enabled: true checkpoint_dir: ./output/llama3_2_1b_encoder/checkpoints model_save_format: safetensors save_consolidated: true ``` Each save creates a versioned directory. The following path is an example: ```text ./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, NeMo AutoModel also writes a Hugging Face-compatible model under the following path: ```text ./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: ```yaml checkpoint: enabled: true checkpoint_dir: ./output/llama3_2_1b_encoder/checkpoints restore_from: LATEST ``` `LATEST` is a resolver keyword: NeMo 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: NeMo AutoModel warns about an incompatibility but still restores the selected checkpoint. If `checkpoint.restore_from` is omitted, NeMo 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: ```python import torch from nemo_automodel import NeMoAutoModelBiEncoder, NeMoAutoTokenizer model_path = "./output/llama3_2_1b_encoder/checkpoints/epoch_0_step_499/model/consolidated" tokenizer = NeMoAutoTokenizer.from_pretrained(model_path, add_eos_token=False) model = NeMoAutoModelBiEncoder.from_pretrained( model_path, pooling="avg", l2_normalize=True, use_liger_kernel=False, ).eval() device = next(model.parameters()).device texts = ["query: what does nvlink do?", "passage: NVLink is a high-bandwidth GPU interconnect."] tokenized = tokenizer(texts, padding=False, truncation=True, max_length=512, return_token_type_ids=False) tokenized = [{key: tokenized[key][idx] for key in tokenized.keys()} for idx in range(len(texts))] tokens = tokenizer.pad(tokenized, padding="longest", return_tensors="pt") tokens = {key: value.to(device) for key, value in tokens.items()} with torch.no_grad(): embeddings = model.encode(tokens) 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. Text reranker training with `CrossEncoderCollator` saves its `prompt_template` as a named `reranking` tokenizer chat template. Sentence Transformers 5.6 or later applies that template automatically when you pass query/document pairs to `CrossEncoder.predict`. Both direct and consolidated saves preserve the template without adding custom Python code to the checkpoint. Any existing model requirement for `trust_remote_code` still applies. This export path supports right truncation and templates ending in `{query}` or `{passage}`. The tokenizer must represent its special tokens as a single-sequence prefix and suffix with zero token type IDs. Setup rejects incompatible tokenizers or templates because their training tokenization cannot be preserved by this export path. The checkpoint's inference limit comes from `tokenizer.model_max_length`; `rerank_max_length` remains the training limit. The following example loads a cross-encoder and ranks two passages: ```python import torch from nemo_automodel import NeMoAutoModelCrossEncoder, NeMoAutoTokenizer model_path = "./output/llama3_2_1b_cross_encoder/checkpoints/epoch_0_step_499/model/consolidated" tokenizer = NeMoAutoTokenizer.from_pretrained(model_path, add_eos_token=False) model = NeMoAutoModelCrossEncoder.from_pretrained( model_path, pooling="avg", num_labels=1, temperature=1.0, use_liger_kernel=False, ).eval() device = next(model.parameters()).device prompt_template = "question:{query} \n \n passage:{passage}" pairs = [ prompt_template.format(query="what does nvlink do?", passage="NVLink is a high-bandwidth GPU interconnect."), prompt_template.format(query="what does nvlink do?", passage="Dropout regularizes neural networks."), ] tokenized = tokenizer(pairs, padding=False, truncation=True, max_length=512, return_token_type_ids=False) tokenized = [{key: tokenized[key][idx] for key in tokenized.keys()} for idx in range(len(pairs))] tokens = tokenizer.pad(tokenized, padding="longest", return_tensors="pt") tokens = {key: value.to(device) for key, value in tokens.items()} with torch.no_grad(): logits = model(tokens).logits.squeeze(-1) 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: | Symptom | Check | | ----------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------- | | Training fails with empty negatives | Ensure every record has `neg_doc` when `n_passages > 1`. | | Dataset records fail to load | Check the supported schemas in [Retrieval Dataset](/nemo/automodel/v0.7/datasets/retrieval-dataset). | | Loss does not move | Verify the positive passage is first and negatives are not duplicates of the positive. | | Poor retrieval quality | Mine harder negatives and align training and inference prefixes. | | Out of memory (OOM) at startup or first batch | Lower `local_batch_size`, `q_max_len`, `p_max_len`, or `rerank_max_length`. Use LoRA for larger backbones. | | Distributed launch times out | Increase `dist_env.timeout_minutes`, especially for first model downloads, slow filesystems, or multi-node runs. | | Batch-size assertion fails | Set `global_batch_size` to a multiple of `local_batch_size * data_parallel_size`. | | `training.jsonl` does not update during a basic verification test | Use standard output and error for live monitoring. JSONL metrics are buffered before flush. | | Run resumes unexpectedly | Use a new or empty `checkpoint_dir`. NeMo AutoModel auto-detects compatible checkpoints when `restore_from` is omitted. | | Different mining and training behavior | Match tokenizer settings, max lengths, and prefix text including trailing spaces across training and mining. | ## Related Files Refer to the following files for implementation and configuration details: * Bi-encoder recipe: [nemo_automodel/recipes/retrieval/train_bi_encoder.py](https://github.com/NVIDIA-NeMo/Automodel/blob/r0.7.0/nemo_automodel/recipes/retrieval/train_bi_encoder.py) * Cross-encoder recipe: [nemo_automodel/recipes/retrieval/train_cross_encoder.py](https://github.com/NVIDIA-NeMo/Automodel/blob/r0.7.0/nemo_automodel/recipes/retrieval/train_cross_encoder.py) * Retrieval dataset guide: [Retrieval Dataset](/nemo/automodel/v0.7/datasets/retrieval-dataset) * Llama-Embed-Nemotron-8B example: [examples/retrieval/bi_encoder/llama_embed_nemotron_8b/llama_embed_nemotron_8b.yaml](https://github.com/NVIDIA-NeMo/Automodel/blob/r0.7.0/examples/retrieval/bi_encoder/llama_embed_nemotron_8b/llama_embed_nemotron_8b.yaml)