TorchTitan：面向生产环境 LLM 预训练的 PyTorch 原生解决方案
计算与语言
2025-06-10 v3 人工智能
分布式、并行与集群计算
机器学习
摘要
大型语言模型（LLM）的发展对于推动先进自然语言处理应用方面起到了重要作用。拥有数十亿参数和数万亿标记的 LLM 训练需要精细的分布式系统，才能高效地跨越数千个加速器，组合和比较多种最新技术。然而，现有解决方案复杂且分散在多个库/存储库中，缺乏互操作性，使用和维护都相当繁琐。因此，精心挑选和经验性地比较训练配方需要相当大的工程 effort。本文介绍 TorchTitan，一个开源的 PyTorch 原生分布式训练系统，将最新技术统一，简化集成并降低开销。TorchTitan 以模块化方式实现 3D 并行，支持弹性扩展，提供生产就绪训练的完整日志、检查点和调试工具。它还集成了硬件-软件协同设计的解决方案，利用诸如 Float8 训练和 SymmetricMemory 等特性。作为一个灵活的测试平台，TorchTitan 便于定制配方的挑选和比较，我们基于经验开发了针对 Llama 3.1 的优化训练配方，并提供关于如何选择技术以实现最大效率的指导。我们对 TorchTitan 在 Llama 3.1 系列 LLM（从 80 亿到 4050 亿参数）进行了全面评估，展示了其卓越的性能、模块化可组合性和弹性可扩展性。通过叠加训练优化，在 128 GPU 规模下的 1D 并行（Llama 3.1 8B）实现了 65.08% 的加速，在 256 GPU 规模下的 2D 并行（Llama 3.1 70B）实现了额外的 12.59% 加速，在 512 GPU 规模下的 3D 并行（Llama 3.1 405B）在 NVIDIA H100 GPU 上相对于优化基线实现了额外的 30% 加速。
引用
@article{arxiv.2410.06511,
title = {TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training},
author = {Wanchao Liang and Tianyu Liu and Less Wright and Will Constable and Andrew Gu and Chien-Chin Huang and Iris Zhang and Wei Feng and Howard Huang and Junjie Wang and Sanket Purandare and Gokul Nadathur and Stratos Idreos},
journal= {arXiv preprint arXiv:2410.06511},
year = {2025}
}