Tevatron Meets Megatron: Expert-Parallel LLM Reranker Training on an Academic Budget
Introduces Tevatron 3.0 with Megatron backend, enabling efficient training of billion-parameter MoE rerankers, surpassing FSDP in speed and scale.
Key Findings
Methodology
This work integrates Megatron's expert parallelism into Tevatron, combining tensor, pipeline, and expert parallel strategies. It employs a bidirectional weight bridge to convert models between Hugging Face and Megatron formats, ensuring ecosystem compatibility. Using ZeRO-1 optimizer, the framework supports training 30B parameter MoE rerankers like Qwen3-30B-A3B efficiently. The system compares various distributed strategies, demonstrating Megatron's superior training speed and memory efficiency while maintaining model quality. Evaluation on BEIR-15 across multiple retrieval tasks confirms the advantages in parameter efficiency and inference throughput, with the framework supporting both LoRA and full-parameter fine-tuning.
Key Results
- In a single-node setup, Megatron backend is approximately 22% faster than FSDP for training 8B dense models, achieving comparable NDCG@10 scores while activating less than half the parameters.
- Successfully trained a 30B MoE reranker (Qwen3-30B-A3B) within the same resource budget, significantly surpassing dense models in inference throughput (over 2x) while maintaining quality.
- The framework supports LoRA and full-parameter tuning, demonstrating consistent performance across different fine-tuning strategies, enabling flexible deployment.
Significance
This research addresses critical bottlenecks in large-scale MoE training within academic settings, providing a resource-efficient, scalable framework that democratizes access to billion-parameter models. It bridges the gap between cutting-edge model architectures and practical training constraints, fostering innovation in information retrieval and NLP. By enabling training of larger models on limited hardware, it paves the way for broader research and application, reducing reliance on industrial-scale infrastructure.
Technical Contribution
The core innovation lies in embedding Megatron's expert parallelism into the Tevatron ecosystem, supported by ZeRO-1 optimizer for memory efficiency. The bidirectional weight bridge ensures ecosystem compatibility, allowing models to be loaded and saved seamlessly in Hugging Face format. The system supports diverse fine-tuning methods, including LoRA, and introduces a unified evaluation and serving pipeline. These advancements collectively enable training of billion-parameter MoE models on academic hardware, breaking previous scalability barriers.
Novelty
This is the first demonstration of training 30B-scale MoE rerankers within an academic budget, leveraging Megatron's expert parallelism combined with Tevatron's flexible toolkit. It overcomes FSDP's limitations in large MoE training, offering a highly compatible, efficient, and scalable framework that is accessible for research groups with limited resources.
Limitations
- The weight conversion bridge adds complexity to the training pipeline, requiring careful engineering. Future work should focus on simplifying this process.
- Expert parallelism introduces communication overhead, which may become a bottleneck at even larger scales, necessitating further optimization.
- Current experiments are limited to specific datasets and model architectures; broader validation across diverse tasks is needed.
Future Work
Future efforts will optimize expert communication strategies, extend multi-task and multimodal training capabilities, and develop automated hyperparameter tuning tools. These improvements aim to further reduce resource requirements, enhance scalability, and facilitate deployment across varied NLP applications.
AI Executive Summary
The rapid growth of large-scale pretraining models like GPT and BERT has revolutionized NLP but also posed significant challenges for academic research groups with limited hardware resources. Traditional distributed training methods such as PyTorch FSDP struggle with memory bottlenecks and slow throughput, especially when scaling to hundreds of billions of parameters. To address this, the authors introduce Tevatron 3.0, a major upgrade that incorporates Megatron's expert parallelism into the existing toolkit, enabling efficient training of billion-parameter models within academic budgets.
This new framework preserves the ecosystem's compatibility by using a bidirectional weight bridge, allowing models to be loaded and saved in Hugging Face format seamlessly. It combines tensor, pipeline, and expert parallelism, significantly boosting training speed—up to 22% faster than FSDP in single-node configurations—while maintaining model quality. The authors demonstrate the training of a 30B parameter MoE reranker (Qwen3-30B-A3B), a feat previously infeasible with PyTorch FSDP.
Extensive evaluations on BEIR-15 show that the MoE reranker achieves comparable or superior retrieval quality at less than half the activated parameters, with over twice the inference throughput. The framework supports both LoRA and full-parameter fine-tuning, offering flexibility for various research needs. These advances lower the barrier for academic groups to develop and deploy large models, fostering innovation in information retrieval, question answering, and beyond.
Looking ahead, the authors plan to optimize expert communication, expand multi-task and multimodal training, and develop automated tuning tools. These efforts aim to further democratize large-scale model training, making cutting-edge NLP research accessible and sustainable for academia, industry, and society at large.
Deep Analysis
Background
Recent advancements in NLP have seen the emergence of billion-parameter models like GPT-3 and dense retrievers such as SPLADE, pushing the boundaries of language understanding and retrieval. Mixture-of-Experts (MoE) architectures like Qwen3-30B utilize sparse activation to scale parameters efficiently, achieving state-of-the-art results in various tasks. However, training such models demands immense computational resources and sophisticated distributed strategies. Existing tools like DeepSpeed ZeRO and PyTorch FSDP provide partial solutions but face limitations in scalability, memory management, and support for expert parallelism. This gap has restricted academic groups from exploring these models extensively. The Tevatron toolkit has historically provided flexible training pipelines but lacked the infrastructure for large MoE models, especially at billion-parameter scales. The need for a more efficient, compatible, and scalable training framework has become urgent to democratize access to these powerful models.
Core Problem
The core challenge lies in enabling efficient training of billion-parameter MoE rerankers within limited academic resources. Existing distributed strategies either suffer from memory exhaustion, low throughput, or lack expert parallelism support. PyTorch FSDP, while memory-safe, cannot handle the scale of 30B parameters due to its inefficiency and inability to support expert parallelism. DeepSpeed ZeRO-2 faces known bugs and performance issues. Consequently, researchers are unable to train large MoE rerankers cost-effectively, limiting progress in retrieval tasks. Addressing these bottlenecks requires a new distributed training paradigm that combines speed, scalability, and ecosystem compatibility.
Innovation
The key innovation is integrating Megatron's expert parallelism into Tevatron, supported by ZeRO-1 optimizer, enabling training of 30B MoE models on a single GPU node. The framework employs a bidirectional weight bridge, allowing models to be loaded and saved seamlessly in Hugging Face format, preserving ecosystem compatibility. It combines tensor, pipeline, and expert parallelism, reducing communication overhead and improving throughput. The system supports LoRA and full-parameter fine-tuning uniformly, providing flexibility. Additionally, a unified evaluation and inference pipeline supports local and remote scoring, facilitating practical deployment. These innovations collectively enable large-scale MoE training within academic budgets, overcoming previous hardware and software limitations.
Methodology
- �� 采用Megatron架构，结合张量、流水线和专家并行策略，支持亿级参数模型训练。• 通过桥接机制，将预训练模型从Hugging Face格式转换到Megatron布局，实现模型的无缝加载与保存。• 利用ZeRO-1优化器，将优化状态分散存储，减少内存占用。• 支持LoRA微调，通过预包装钩子将适配器参数引入训练流程。• 采用listwise-KL蒸馏和对比学习作为训练目标，提升模型性能。• 设计统一的评估与推理接口，支持本地评分和远程推理服务，确保系统的灵活性。• 通过多节点多GPU配置，验证不同模型规模（8B、30B）在训练速度和推理吞吐上的表现。• 结合BEIR-15数据集，进行多场景、多指标的性能评估，确保模型质量。
Experiments
实验采用BEIR-15信息检索基准，比较不同分布式策略（FSDP、Megatron）在训练效率和模型质量上的差异。训练模型包括8B密集模型和30B MoE模型，使用不同微调策略（LoRA、全参数）验证模型性能。通过单节点和多节点配置，测量训练时间、GPU内存峰值和推理吞吐。还进行模型质量评估（NDCG@10、Recall@k），确保模型在检索任务中的有效性。对比不同后端在参数效率和训练速度上的优势，验证新框架在实际场景中的适用性。
Results
实验结果显示，Megatron后端在单节点配置下比FSDP快约22%，训练8B模型达到相同质量（NDCG@10）指标，参数激活量减少一半。支持训练30B参数的MoE模型，显著提升推理吞吐（超过2倍），同时保持模型性能。微调策略方面，LoRA微调与全参数微调效果一致，验证了系统的灵活性。整体而言，新框架在模型质量、训练速度和参数利用率方面优于传统方法，为大规模模型训练提供了可行方案。
Applications
该框架适用于学术界和工业界的大规模信息检索模型训练，尤其是在资源有限的环境中。可用于训练高效的重排序器、知识蒸馏模型，以及多任务多模态模型，推动大模型在搜索引擎、问答系统中的应用。未来，结合多模态数据和多任务训练，有望实现更智能、更高效的检索系统，降低硬件成本，提升模型泛化能力。
Limitations & Outlook
目前，模型参数的桥接和转换流程较为复杂，增加了工程实现难度。专家并行在极端模型规模下仍存在通信瓶颈，需优化通信策略。实验主要集中在特定数据集和模型架构；泛化到其他任务和模型仍需验证。未来需简化流程、提升通信效率，并扩展多任务、多模态训练能力，以应对更复杂的应用场景。
Plain Language Accessible to non-experts
想象你在一家大型厨房里做饭，要准备很多不同的菜肴。传统的方法是每个厨师都用一套完整的厨具，做一份菜，然后换人。这种方式虽然简单，但效率很低，因为每次都要重复很多工作。现在，厨师们开始合作：一部分负责切菜，一部分负责炒菜，还有一些专门负责调味。每个人只专注于自己擅长的部分，这样就可以同时做出更多菜，而且速度快得多。这就像模型中的专家并行，把任务分成不同部分，让每个“厨师”只处理自己擅长的“菜”。这样，不仅节省时间，还能做出更复杂、更美味的菜肴。Tevatron 3.0就是用这种合作方式，让大模型训练变得更快、更高效，帮助研究者用更少的资源做出更厉害的人工智能模型。
ELI14 Explained like you're 14
想象你在学校的厨房里做饭，平时每个人都用一套完整的厨具，做一道菜要很长时间。而现在，厨师们开始合作：有人专门切菜，有人负责炒，有人调味。每个人只做自己最擅长的事情，这样可以同时做出很多菜，而且快得多。这就像模型里的“专家”一样，把任务分成不同部分，让每个“专家”只处理自己擅长的内容。这样，不仅节省时间，还能做出更复杂、更好吃的菜。Tevatron 3.0也是这样，它让大模型的训练变得更快、更省资源，就像厨房里的合作一样，帮助科学家们用更少的时间和设备，做出更厉害的人工智能模型。
Abstract
Modern reranking recipes---billion-scale cross-encoders, mixture-of-experts (MoE) backbones, and distillation against strong teachers---have outpaced the training infrastructure available to most academic groups. Existing Tevatron reranker training relies on the Hugging Face Trainer with DeepSpeed or PyTorch FSDP1, but these backends lack efficient support for large-scale MoE training. We present Tevatron 3.0, which integrates a Megatron-Core training backend into Tevatron while preserving its data pipeline, evaluation workflow, and Hugging Face-compatible checkpoints. We benchmark existing distributed training configurations against the new backend, showing that Megatron matches FSDP reranker quality and training efficiency under comparable data-parallel settings, is up to 22% faster in the recommended single-node configuration, and supports both LoRA and full-parameter fine-tuning. Crucially, expert parallelism enables training a 30B-parameter Qwen3-30B-A3B MoE reranker, which is infeasible with PyTorch FSDP1. Using this framework, we conduct a controlled comparison of MoE versus dense models, LoRA versus full-parameter tuning, and distillation versus contrastive training on BEIR-15 with three first-stage retrievers, and report serving throughput for Hugging Face and vLLM. We find that the MoE reranker matches dense 8B quality while activating less than half as many parameters and achieving substantially higher inference throughput. We will release the framework and trained checkpoints.