feat: add reranker training

[adil-a]

Summary

Adds cross-encoder (reranker) training support alongside the existing bi-encoder pipeline. Cross-encoders jointly attend to query-document pairs, producing relevance scores that are significantly more accurate than bi-encoder similarity — a critical capability for retrieval reranking.

Cross-Encoder Model & Training Recipe

• NeMoAutoModelCrossEncoder — New auto-model class with full infrastructure support (FSDP2, PEFT, kernel patching)
• CrossEncoderModel — Wraps any ForSequenceClassification backbone. Routes through LlamaBidirectionalForSequenceClassification for Llama models, falls back to HF AutoModelForSequenceClassification for all others
• TrainCrossEncoderRecipe — Extends TrainBiEncoderRecipe with cross-entropy loss on reshaped logits, training accuracy tracking, and validation with accuracy@1 and MRR metrics
• accuracy() / batch_mrr() — Pure-function metric utilities for training and validation

Data Pipeline

• CrossEncoderCollator — Concatenates query+passage via configurable prompt template, tokenizes, pads. Compatible with NeMoAutoTokenizer
• flatten_bi_encoder_to_cross_encoder() — Transforms grouped bi-encoder data to flattened cross-encoder format. Same data files work for both model types via model_type config switch
• make_retrieval_dataset(model_type="cross_encoder") — Unified dataset factory supporting both bi-encoder and cross-encoder formats

Architecture Simplification

• Removed dual-encoder (lm_q/lm_p) design with share_encoder — single backbone is simpler and aligns with modern embedding models
• Simplified EncoderStateDictAdapter (was BiencoderStateDictAdapter) — symmetric model. prefix strip/add
• Any HF model as encoder backbone — AutoModel/AutoModelForSequenceClassification fallback means users are not limited to custom bidirectional models
• Unified naming: Biencoder -> BiEncoder, train_n_passages -> n_passages

CI & Testing

• Functional test: 2-GPU FSDP2 training (32 steps) + quality evaluation asserting finetuned model improves positive-pair scoring and achieves >=75% ranking accuracy
• Unit tests: accuracy(), batch_mrr(), CrossEncoderCollator (output shapes, labels, padding), flatten_bi_encoder_to_cross_encoder() (value validation), liger/SDPA retry for cross-encoder
• L2_Retrieval CI job added to cicd-main.yml

Example

model:
  _target_: nemo_automodel.NeMoAutoModelCrossEncoder.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
  num_labels: 1
  pooling: avg
dataloader:
  dataset:
    _target_: nemo_automodel.components.datasets.llm.make_retrieval_dataset
    model_type: cross_encoder
    n_passages: 5
  collate_fn:
    _target_: nemo_automodel.components.datasets.llm.CrossEncoderCollator
    rerank_max_length: 512
    prompt_template: "question:{query} \n \n passage:{passage}"

Test plan

• 2-GPU cross-encoder training completes without errors (32 steps, FSDP2)
• Finetuned model pos-score improvement > 0 (observed: +9.35)
• Finetuned ranking accuracy >= 75% (observed: 92.0%)
• 22 unit tests pass (accuracy, batch_mrr, collator, flatten, retry)
• ruff check clean on nemo_automodel/
• CI L2_Retrieval job passes

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[adil-a]

Thanks for all the useful comments @rnyak . I overlooked these yaml details initially. I've now fixed them. Please take a look when possible.

[adil-a]

/ok to test 3bb2f62

[akoumpa]

/claude review

[claude]

Potential bug: num_labels is read from features[0].get("num_labels"), which reflects the number of unique queries in a single dataset transform call (typically 1 when set_transform processes individual rows). But the labels tensor needs to have size equal to the number of query groups in the entire collated batch (i.e. len(features) // n_passages).

If a DataLoader batch has 8 flat (query, doc) pairs from 4 queries with n_passages=2, num_labels here would be 1 (from a single row's transform), producing labels = torch.zeros(1) instead of torch.zeros(4). This would cause a shape mismatch at F.cross_entropy(logits.view(-1, n_passages), labels) in the training recipe.

This might be masked by how HuggingFace set_transform + __getitems__ batches data in your environment, but it's fragile. Consider computing labels from the actual batch size:

n_queries = len(features) // n_passages  # where n_passages comes from config
labels = torch.zeros(n_queries, dtype=torch.long)