TorchTitan统一LLM训练技术,实现高效分布式训练与弹性扩展。
TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training

- 基于PyTorch原生设计,支持3D并行与模块化组合。
- 在512张H100上训练405B模型,性能比基线提升30%。
- 适合需要高效训练大模型的研究者与工业团队。
大语言模型(LLMs)的发展推动了自然语言处理的前沿应用。训练拥有数十亿参数和数万亿标记的模型需依赖复杂的分布式系统,以高效跨数千加速器扩展。然而现有方案分散于多个库中,互操作性差且维护困难,导致训练方案的整合与实证比较耗时费力。本文提出TorchTitan,一个开源的、原生支持PyTorch的分布式训练系统,统一了最先进的训练技术,简化集成并降低开销。TorchTitan以模块化方式支持3D并行与弹性扩展,提供全面的日志、检查点与调试工具,满足生产级训练需求。系统融合软硬件协同设计,利用Float8训练与SymmetricMemory等特性。作为灵活的测试平台,TorchTitan支持自定义训练配方的开发与对比,助力优化Llama 3.1训练方案,并基于经验提供技术选型建议。我们在涵盖80亿至4050亿参数的Llama 3.1系列模型上全面评估了TorchTitan,验证其卓越性能、模块可组合性与弹性扩展能力。通过叠加训练优化,在128卡规模下(Llama 3.1 8B)实现65.08%加速;256卡规模(Llama 3.1 70B)额外提升12.59%;512卡规模(Llama 3.1 405B)再增30%,均基于NVIDIA H100 GPU。
原文摘要 · Abstract (English)
The development of large language models (LLMs) has been instrumental in advancing state-of-the-art natural language processing applications. Training LLMs with billions of parameters and trillions of tokens require sophisticated distributed systems that enable composing and comparing several state-of-the-art techniques in order to efficiently scale across thousands of accelerators. However, existing solutions are complex, scattered across multiple libraries/repositories, lack interoperability, and are cumbersome to maintain. Thus, curating and empirically comparing training recipes require non-trivial engineering effort. This paper introduces TorchTitan, an open-source, PyTorch-native distributed training system that unifies state-of-the-art techniques, streamlining integration and reducing overhead. TorchTitan enables 3D parallelism in a modular manner with elastic scaling, providing comprehensive logging, checkpointing, and debugging tools for production-ready training. It also incorporates hardware-software co-designed solutions, leveraging features like Float8 training and SymmetricMemory. As a flexible test bed, TorchTitan facilitates custom recipe curation and comparison, allowing us to develop optimized training recipes for Llama 3.1 and provide guidance on selecting techniques for maximum efficiency based on our experiences. We thoroughly assess TorchTitan on the Llama 3.1 family of LLMs, spanning 8 billion to 405 billion parameters, and showcase its exceptional performance, modular composability, and elastic scalability. By stacking training optimizations, we demonstrate accelerations of 65.08% with 1D parallelism at the 128-GPU scale (Llama 3.1 8B), an additional 12.59% with 2D parallelism at the 256-GPU scale (Llama 3.1 70B), and an additional 30% with 3D parallelism at the 512-GPU scale (Llama 3.1 405B) on NVIDIA H100 GPUs over optimized baselines.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。