Megatron-LLaMA: High-Performance Large Language Model Training Framework

On September 12, 2023, Taotian Group and AiCheng Technology jointly open-sourced the Megatron-LLaMA large model training framework. This framework aims to help developers improve large language model training performance, reduce training costs, and maintain compatibility with the LLaMA community. Testing shows that on 32-card training, Megatron-LLaMA achieves 176% speedup compared to the code version directly obtained from HuggingFace, and demonstrates near-linear scalability in large-scale training with high tolerance to network instability.

Megatron-LLaMA provides a standard Megatron-LM implementation of LLaMA and offers tools for free switching between HuggingFace format, facilitating compatibility with community ecosystem tools. The framework redesigns Megatron-LM's backward process, enabling excellent training performance whether in scenarios requiring large gradient accumulation with fewer nodes or using small GA with more nodes.

The framework includes key technologies such as tensor parallelism (TP), pipeline parallelism (PP), sequence parallelism (SP), and DistributedOptimizer optimization. These technologies significantly reduce memory usage and improve GPU utilization. Megatron-LLaMA's OverlappedDistributedOptimizer addresses the limitation of serial communication in the original Megatron-LM by enabling gradient communication to overlap with computation.

Testing demonstrates that Megatron-LLaMA achieves 176% speedup on 32-card training compared to HuggingFace's direct code version, and even with DeepSpeed and FlashAttention optimizations, it still reduces training time by at least 19%. In large-scale training, Megatron-LLaMA shows near-linear scalability compared to 32 cards. For example, using 512 A100 cards to reproduce LLaMA-13B training, Megatron-LLaMA's backward mechanism saves at least two days compared to Megatron-LM's DistributedOptimizer without any accuracy loss.

The framework demonstrates high tolerance to network instability, achieving 0.85 linear scalability even in cost-effective 4x200Gbps 8xA100-80GB training clusters where network bandwidth is severely limited. This is significantly better than Megatron-LM's performance of less than 0.7 in the same conditions.

Megatron-LLaMA is already widely used within Taotian Group and AiCheng Technology, and the team will continue maintenance and development after open-sourcing. Future plans include adaptive optimal configuration selection, support for more model structures or local design modifications, and extreme performance training solutions in various hardware environments.

Project address: https://github.com/alibaba/Megatron-LLaMA