Tevatron Meets Megatron: Expert-Parallel LLM Reranker Training on an Academic Budget
Tevatron集成Megatron:学术预算下的专家并行LLM重排序训练
🧭 LLM重排序 · 专家并行高效训练 ★ 深度解读
University of Utah · University of Waterloo · The University of Queensland · Carnegie Mellon University · University of Virginia
为Tevatron重排序工具包引入Megatron训练后端,通过专家并行实现30B MoE重排序器训练,同时保持HF兼容。
📌 工作参考价值(4 点)
- 在资源受限下训练大规模MoE重排序模型:利用Megatron的专家并行(EP)将MoE层的专家分布到多个GPU,仅激活少量专家,可在两台8×H200节点上训练Qwen3-30B-A3B(30B总参,激活约3B),突破PyTorch FSDP1不支持EP的限制。
- 加速训练并保持部署兼容性:Megatron后端在推荐配置(TP=2/DP=4)下比FSDP1全分片快约22%,且训练产出可直接加载为Hugging Face checkpoint,无缝接入vLLM等推理引擎,部署零转换成本。
- LoRA适配MoE时注意目标模块划分:使用命名目标组注册表(如指定moe_experts、moe_shared,冻结router),避免盲目适配全部128个专家FFN或遗漏关键层,这是LoRA在MoE重排序中不严重掉点的关键工程trick。
- 蒸馏训练可离线标注教师分数:将教师分数预计算存入数据集,训练时无需加载教师模型,节省显存和计算量;蒸馏与对比学习在不同数据集上各有优势,可根据业务场景灵活选择,并在重排序任务中简化训练管线。