[论文解读] Sockeye 3: Fast Neural Machine Translation with PyTorch

Felix Hieber, Michael Denkowski+12 more|arXiv (Cornell University)|Jul 12, 2022

Natural Language Processing Techniques被引用 11

一句话总结

Sockeye 3 是一个基于 PyTorch 的神经机器翻译工具包，通过优化的混合精度训练、分布式数据并行和模型量化，加速了训练和推理。与其它 PyTorch 实现相比，其在 GPU 上的推理速度最高提升 126%，在 CPU 上最高提升 292%，同时支持高级功能如源/目标前缀、词汇短名单以及正式性/冗长度控制，适用于生产环境和研究用途。

ABSTRACT

Sockeye 3 is the latest version of the Sockeye toolkit for Neural Machine Translation (NMT). Now based on PyTorch, Sockeye 3 provides faster model implementations and more advanced features with a further streamlined codebase. This enables broader experimentation with faster iteration, efficient training of stronger and faster models, and the flexibility to move new ideas quickly from research to production. When running comparable models, Sockeye 3 is up to 126% faster than other PyTorch implementations on GPUs and up to 292% faster on CPUs. Sockeye 3 is open source software released under the Apache 2.0 license.

研究动机与目标

• 开发一个基于 PyTorch 的更快、更可扩展且适合生产环境的神经机器翻译工具包。
• 通过简化代码库并优化训练与推理流水线，使研究人员能够快速迭代。
• 支持高级神经机器翻译功能，如源和目标前缀、词汇短名单以及混合解码器，以实现细粒度控制。
• 通过高效、量化和分布式训练，弥合研究创新与生产部署之间的差距。
• 提供一个灵活的开源框架，用于在翻译中实验正式性与冗长度控制。

提出的方法

• Sockeye 3 使用 PyTorch 的 JIT 编译器追踪并优化编码器、解码器和束搜索的静态计算图，减少动态控制流开销。
• 通过 PyTorch 的 DDP 和 AMP 实现分布式混合精度训练，支持 FP16/FP32 混合精度以及全 FP16（O2）模式，提升速度与内存效率。
• 数据准备过程实现并行化并分片为二进制格式，支持在任意大规模数据集上训练，且内存使用量固定。
• 通过 PyTorch 的动态量化（INT8）应用于 CPU 推理，以及 GPU 推理的 FP16 转换，实现模型量化，以最小的精度损失加速推理。
• 使用源前缀对模型输出进行条件控制，以实现正式性或冗长度的可控生成，无需重新训练。
• 将贪婪搜索实现为束搜索的轻量级替代方案，在保持翻译质量的同时降低计算开销。

实验结果

研究问题

• RQ1 基于 PyTorch 的神经机器翻译工具包如何显著提升训练与推理吞吐量，相比现有实现？
• RQ2 混合精度与分布式训练在不牺牲模型质量的前提下，能在多大程度上提升神经机器翻译的可扩展性与效率？
• RQ3 源前缀条件控制是否能有效控制神经机器翻译中的输出正式性与冗长度？
• RQ4 量化对推理速度与内存占用有何影响，同时保持翻译质量？
• RQ5 通过优化推理组件（如贪婪搜索与模型追踪）可实现多大的性能提升？

主要发现

• Sockeye 3 在 GPU 上的推理速度相比其他 PyTorch 神经机器翻译实现最高提升 126%，在 CPU 上最高提升 292%。
• 在 8 张 GPU 上使用混合精度训练可实现 7.7 倍的吞吐量加速，GPU 效率达 96.6%，表现出良好的可扩展性。
• 与束搜索（束大小为 1）相比，贪婪搜索在 GPU 上将翻译速度提升 16%，在 CPU 上提升 6%，同时保持翻译质量。
• 在 CPU 上使用 INT8 量化可将推理速度从每秒 2.6 个句子提升至 4.5 个句子，BLEU 分数下降可忽略。
• 正式性控制模型在正式翻译任务中准确率达 97.8%，在 WMT 测试集上 BLEU 得分为 45.0，表明控制精度高。
• 通过长度标记与 N-best 重排序实现的冗长度控制，使英文-德文翻译的最终系统排名指标（BERTScore × LC）最高提升 17.4 分。