Technical Summary: Tevatron 3.0 – Expert-Parallel LLM Reranker Training on an Academic Budget

1. Problem Statement

Modern reranking pipelines increasingly rely on billion-scale cross-encoders, Mixture-of-Experts (MoE) backbones, and knowledge distillation against strong teacher models. However, training these models has outpaced the infrastructure available to most academic research groups. The existing Tevatron toolkit (versions 1.0 and 2.0) relies on the Hugging Face Trainer with PyTorch FSDP (Fully Sharded Data Parallel) and DeepSpeed for memory efficiency. This stack faces three critical limitations when applied to modern reranking workloads:

1. DeepSpeed ZeRO-2 Instability: A known bug involving CPU-offloaded gradient accumulation causes ZeRO-2 sharding to fail in practice, forcing a fallback to PyTorch FSDP.
2. Throughput Inefficiency of FSDP: While PyTorch FSDP (ZeRO-2 and ZeRO-3) is memory-safe, it is throughput-inefficient. For an 8B parameter reranker, full-sharding (ZeRO-3) offers no speed advantage over partial sharding (ZeRO-2) because the optimizer state dominates the memory footprint, and the heavier sharding stages do not hide communication overhead effectively.
3. Lack of Expert Parallelism (EP): PyTorch FSDP does not support Expert Parallelism. Consequently, MoE rerankers (e.g., Qwen3-30B-A3B with 128 experts) cannot be trained efficiently. FSDP can only replicate-then-shard experts as standard parameters, failing to route active subsets to specific devices, making training 30B-parameter MoE models infeasible on standard academic budgets.

2. Methodology: Tevatron 3.0 System Design

The authors present Tevatron 3.0, which integrates a Megatron-Core training backend into the existing Tevatron reranker toolkit. The design philosophy prioritizes backward compatibility: the new backend reuses Tevatron's existing data paths, collators, evaluation pipelines, and Hugging Face-loadable checkpoint formats. Users can switch backends via launch scripts without migrating data or relearning the toolkit.

Core Components

• Megatron Backend Integration: The system introduces a Megatron training backend that supports Tensor Parallelism (TP), Pipeline Parallelism (PP), and Expert Parallelism (EP).
  • Data Parallelism (DP): Uses a ZeRO-1 style distributed optimizer (sharding optimizer states, replicating parameters and gradients).
  • Expert Parallelism: Partitions MoE experts across devices, allowing only the active subset of experts to be processed on specific ranks. This is the critical enabler for training large MoE models.
• Bidirectional Weight Bridge: A mechanism to convert weights between Hugging Face and Megatron parallel layouts.
  • Load: Pretrained HF checkpoints are sharded into the Megatron layout.
  • Save: Trained models are exported back to HF format (supporting both full-parameter and pre-merged LoRA weights), ensuring compatibility with standard HF and vLLM inference paths without conversion steps.
• Unified Evaluation and Serving:
  • Supports two schemas behind a single CLI: Local scoring of existing ranklists and Remote rerank-and-score via HTTP.
  • The remote backend pool supports interchangeable Hugging Face and vLLM servers, allowing users to swap between exact-match training math and high-throughput production inference.
• LoRA for MoE: Introduces a named LoRA target-group registry. This prevents the "foot-gun" where a dense MLP pattern might inadvertently match all 128 expert FFNs. It explicitly maps roles (e.g., moe_experts, moe_shared, moe_router) to correct module patterns, enabling effective low-rank adaptation on MoE layers.

Training Objectives

The system supports both Contrastive Learning (listwise cross-entropy) and Distillation (listwise KL divergence against precomputed teacher scores). The distillation driver is an alternate entry point that consumes offline teacher scores, requiring no additional forward passes during training.

3. Key Contributions

1. C1: Megatron Reranker Training Backend: A drop-in replacement for the HF Trainer that matches FSDP-trained quality and efficiency (at matched topologies) while offering ~22% faster training in recommended single-node configurations. It supports full-parameter and LoRA fine-tuning.
2. C2: Expert-Parallel MoE Training: Enables the training of 30B-parameter MoE rerankers (Qwen3-30B-A3B) on a single-job budget where PyTorch FSDP fails. This is achieved through EP and a specialized LoRA target registry.
3. C3: Controlled Study and Artifacts: A comprehensive evaluation of model architecture (MoE vs. Dense), parameter efficiency (LoRA vs. Full), and loss formulation (Distillation vs. Contrastive) across BEIR-15. All trained checkpoints are released as reusable artifacts.

4. Experimental Results

The authors conducted a controlled study training rerankers on the RLHN-680K dataset and evaluating on BEIR-15 (15 tasks) with three first-stage retrievers (BM25, E5, SPLADE-v3).

Quality and Efficiency

• Quality Parity: The Megatron backend produces rerankers with quality indistinguishable from PyTorch FSDP (agreement within 0.002 NDCG@10) on dense 8B models.
• MoE vs. Dense: The 30B-A3B MoE reranker matches the quality of a dense 8B reranker (within 0.006 NDCG@10) while activating less than half the parameters (~3B active vs. ~8B).
• Training Speed: In recommended configurations, Megatron trains dense 8B models ~22% faster than PyTorch FSDP. However, at matched topologies (DP=8), the speed gap narrows to ~4%, with Megatron trading memory for speed (ZeRO-1 replicates parameters vs. FSDP ZeRO-2 sharding gradients).
• Inference Throughput: The MoE model achieves significantly higher inference throughput than the dense 8B model because it activates fewer parameters per token.
  • Hugging Face: MoE is 1.15× faster.
  • vLLM: MoE is 1.43× faster (2,622 pairs/s vs. 1,834 pairs/s).
  • The combination of vLLM and MoE yields a 4.5× throughput improvement over the most expensive configuration (HF × Dense 8B).

Parameter Efficiency and Distillation

• LoRA: LoRA recovers most of the full-parameter performance (within ~0.01 NDCG@10) with a fraction of trainable parameters. This is particularly effective on MoE models when using the expert-aware target registry.
• Distillation: Distillation acts as a regularizer rather than a uniform improvement. It helps on tasks with noisy/graded relevance but can hurt on tasks requiring sharp answers. The study notes that using a teacher of the same size as the student limits the capacity for effective distillation.

5. Significance and Claims

The paper modestly frames its contribution as a system demonstration rather than a claim of state-of-the-art model performance. The primary significance lies in making advanced reranking recipes (MoE, distillation, billion-scale training) trainable and serveable under an academic compute budget.

• Feasibility: Tevatron 3.0 proves that training 30B-parameter MoE rerankers is feasible on academic hardware (e.g., two 8×H200 nodes) where previous stacks (PyTorch FSDP) fail due to lack of Expert Parallelism.
• Reproducibility: By releasing the framework and trained checkpoints, the authors enable the community to reproduce these results and build upon a standardized, compatible ecosystem.
• Operating Point: The study identifies a favorable operating point where MoE rerankers match dense 8B quality with significantly lower active parameters and higher inference throughput, a capability now accessible to researchers without specialized industrial infrastructure.

The authors explicitly state that the checkpoints are not tuned for leaderboard standing and that the controlled internal comparisons (MoE vs. Dense, LoRA vs. Full) are the primary evidence of the framework's utility, rather than the absolute performance metrics against prior baselines trained on different data.