TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training

要約

大規模言語モデル (LLM) の開発は、最先端の自然言語処理アプリケーションの進歩に貢献してきました。
数十億のパラメータと数兆のトークンを使用して LLM をトレーニングするには、数千のアクセラレータ間で効率的に拡張するために、いくつかの最先端技術を構成および比較できる高度な分散システムが必要です。
しかし、既存のソリューションは複雑で、複数のライブラリ/リポジトリに分散しており、相互運用性に欠けており、保守が面倒です。
したがって、トレーニング レシピを厳選し、経験的に比較するには、並大抵のエンジニアリング作業が必要です。
このペーパーでは、最先端の技術を統合し、統合を合理化し、オーバーヘッドを削減するオープンソースの PyTorch ネイティブ分散トレーニング システムである、TorchTitan について紹介します。
TorchTitan は、柔軟なスケーリングを備えたモジュール方式で 3D 並列処理を可能にし、実稼働準備が整ったトレーニング用の包括的なロギング、チェックポイント作成、およびデバッグ ツールを提供します。
また、Float8 トレーニングや SymmetricMemory などの機能を活用した、ハードウェアとソフトウェアが共同設計されたソリューションも組み込まれています。
柔軟なテストベッドとして、TorchTitan はカスタム レシピのキュレーションと比較を容易にし、Llama 3.1 用に最適化されたトレーニング レシピを開発し、経験に基づいて効率を最大化するテクニックの選択に関するガイダンスを提供します。
私たちは、80 億から 4,050 億のパラメーターに及ぶ Llama 3.1 ファミリの LLM 上で TorchTitan を徹底的に評価し、その卓越したパフォーマンス、モジュール構成可能性、および柔軟な拡張性を紹介します。
トレーニングの最適化を積み重ねることにより、128 GPU スケール (Llama 3.1 8B) の 1D 並列処理で 65.08%、256 GPU スケール (Llama 3.1 70B) の 2D 並列処理でさらに 12.59%、さらに 30% の高速化を実証しました。
最適化されたベースライン上の NVIDIA H100 GPU 上の 512 GPU スケール (Llama 3.1 405B) での 3D 並列処理。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google