LinkedIn公开多教师蒸馏框架：6亿参数模型让职位搜索NDCG@10提升24.48%

LinkedIn公布了其AI驱动型职位搜索背后的训练基础设施，介绍了一条多教师蒸馏流水线。该流水线将大型教师模型中的知识压缩至一个紧凑的6亿参数排序模型中，主要贡献不只是蒸馏技术本身，而是让蒸馏速度足以支撑快速迭代的系统工程，其中包括一个基于SGLang构建的定制框架，可直接在训练循环中为教师模型提供服务。

训练小语言模型以改善点击和申请等相关性与互动目标，需要针对每个训练样本查询一个或多个大型教师模型，为这些教师模型提供服务可能拖慢整个过程。对于每秒必须处理数十万次查询的LinkedIn排序系统而言，这会成为瓶颈。LinkedIn使用SGLang创建了多教师蒸馏框架，可加载并服务不同规模的教师模型，同时管理张量并行和数据并行配置。训练期间，异步客户端向教师模型发送请求，处理其输出并将结果添加到蒸馏损失中，团队称之为在线多教师蒸馏。通过在多个节点部署本地教师模型副本，该系统将蒸馏过程提速3倍，同时保持较低延迟。

为减少服务开销并避免重复计算，LinkedIn又引入离线多教师蒸馏，预先计算教师模型输出并存储在HDFS或NFS上，训练过程直接使用这些结果而不再实时查询。这种在线与离线相结合的方案还配合了更广泛的训练层优化技术栈：采用LiGer降低内存使用量并将批次大小扩大至原来的2倍，通过多节点训练实现最高3.5倍加速，使用FSDP2再获得20%的性能提升，采用H200多节点集群在此基础上额外提升最高30%。团队曾评估FP8混合精度，但发现它对参数量低于80亿的模型没有帮助，因为类型转换开销超过了计算方面的节省。这些优化叠加带来了约8倍的训练速度提升。

在模型方面，这个6亿参数的学生模型知识来自一个80亿参数的相关性预言模型和一个17亿参数的互动教师模型，将职位搜索的NDCG@10提升了24.48%，从0.7583提高至0.9432。同一项研究还在推理侧采用结构化剪枝和上下文压缩，将每块GPU的排序吞吐量从每秒约290个条目提高到2000多个。该系统已投入生产，为LinkedIn美国用户的自然语言职位搜索提供支持，基于开源大语言模型服务引擎SGLang构建。