Source-linked AI summary
Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training
Shenggui Li, Hongxin Liu, Zhengda Bian, Jiarui Fang, Haichen Huang, Yuliang Liu, Boxiang Wang, Yang You
TL;DR
Large-scale models create memory bottlenecks on individual GPUs, while selecting and implementing effective distributed parallel strategies requires substantial expertise. Colossal-AI unifies parallelism and heterogeneous training methods behind interfaces designed to preserve sequential coding practices, achieving up to 2.76x speedup over baseline systems.
Problem
Limited GPU memory prevents single devices from efficiently training billion-parameter models, while choosing and adapting parallel strategies remains difficult for users of non-distributed code.
Method
Colossal-AI combines data, pipeline, tensor, sequence, heterogeneous, and automatic parallelization techniques in a unified, modular system with user-friendly APIs.
Results
Up to 2.76x speedup over baseline systems was demonstrated in experiments.
Takeaways & Limitations
Colossal-AI supports flexible combinations of acceleration techniques and provides parallelism strategies intended for different hardware conditions.
Takeaways & Limitations
1D tensor parallelism assumes equal device-interconnect bandwidth and is therefore suited mainly to machines with fully connected NVLinks.
Abstract
from arXiv · showhide
The success of Transformer models has pushed the deep learning model scale to billions of parameters. Due to the limited memory resource of a single GPU, However, the best practice for choosing the optimal parallel strategy is still lacking, since it requires domain expertise in both deep learning and parallel computing. The Colossal-AI system addressed the above challenge by introducing a unified interface to scale your sequential code of model training to distributed environments. It supports parallel training methods such as data, pipeline, tensor, and sequence parallelism, as well as heterogeneous training methods integrated with zero redundancy optimizer. Compared to the baseline system, Colossal-AI can achieve up to 2.76 times training speedup on large-scale models.
1 INTRODUCTION
Colossal-AI addresses the memory and programming challenges of large-scale model training by unifying diverse parallelization and acceleration techniques behind user-friendly interfaces. Its modular system supports flexible combinations of methods and reports improved performance over baseline systems.
- Billion-parameter models exceed typical single-GPU memory because parameters, activations, gradients, and optimizer states jointly consume substantial memory.A 10-billion-parameter FP16 model alone can require 20 GB of model memory, while unoptimized training can exceed 80 GB.
- Existing approaches either cannot handle billion-parameter model data or make distributed programming difficult for users accustomed to sequential code.Pipeline and tensor parallelism address model-data scaling, but adapting ordinary training code to distributed programming remains challenging.
- Colossal-AI unifies multiple training acceleration techniques through a modular system that allows flexible combinations of parallelism methods.The system is designed to preserve users’ coding habits while supporting distributed training.
- Colossal-AI provides optimized parallelism and heterogeneous training methods through friendly APIs requiring minimum code changes.The system also analyzes suitable parallelism strategies under different hardware conditions.
- Figure 1 presents the architecture of Colossal-AI.
2 BACKGROUND
Transformer models and long-sequence applications intensify memory demands, motivating distributed training methods that shard model or activation data. The background surveys data, tensor, pipeline, and sequence parallelism, including their benefits and hardware constraints.
- Transformer Models: Transformer layers combine multi-head attention and feed-forward blocks, and larger Transformer models can improve performance on tasks such as LAMBADA.GPT-3 is reported to outperform smaller models by an 18% absolute increase in prediction accuracy on LAMBADA.
- Data Parallelism: Data parallelism replicates the model across devices, splits the dataset, and synchronizes gradients, but duplicates model data and scales sub-linearly.Zero Redundancy Optimizer partitions parameters, gradients, and optimizer states across devices to reduce this redundancy.
- Tensor Parallelism: Tensor parallelism shards tensors across devices and uses distributed matrix multiplication with collective communication to reduce per-device parameter storage.In Megatron-LM’s 1D approach, each device holds 1/N of the parameters when training on N devices.
- Sequence Parallelism: Tensor parallelism does not eliminate duplicated layer inputs and outputs, while sequence parallelism addresses long-sequence activation memory by partitioning inputs along the sequence dimension.Sequence parallelism replaces self-attention with Ring Self-Attention to exchange partial query, key, and value embeddings across devices.
- Tensor Parallelism: Advanced 2D, 2.5D, and 3D tensor parallelism split tensors across more dimensions, reducing communication volume but imposing GPU-count and topology requirements.1D tensor parallelism works with any number of GPUs, whereas 2D, 2.5D, and 3D methods require n^2, a*n^2, and n^3 GPUs respectively.
- Pipeline Parallelism: Pipeline parallelism assigns consecutive layer chunks to devices and passes intermediate activations and gradients between stages, reducing cross-node communication and increasing throughput.Pipeline execution also introduces bubble time when some devices are idle.
3 DESIGN
Colossal-AI uses a modular unified system to combine model, sequence, pipeline, sharding, offloading, and automatic parallelism techniques for large-scale training. Its advanced tensor-parallel methods reduce communication relative to 1D parallelism, while flexible memory management and experimental search target different hardware and training requirements.
- Unified system: Colossal-AI provides modular acceleration techniques that can be freely combined to support diverse hardware and training settings.The system is designed to achieve maximal performance through unified APIs and modular components.
- Multi-dimensional model parallelism: It supports multiple tensor-parallel methods, enabling users to select strategies according to training requirements and GPU count.Supported methods include 1D, 2D, 2.5D, and 3D tensor parallelism, whereas Megatron-LM supports only 1D tensor splitting.
- Multi-dimensional model parallelism: Advanced tensor parallelism has lower communication volume than 1D parallelism because collective communication is restricted to subgroups of computing nodes.In 2D parallelism, communication involves nodes in one row or column rather than all computing nodes.
- Multi-dimensional model parallelism: Sequence and pipeline parallelism are included so hybrid parallelism is available for large-scale cluster training.These methods complement the supported tensor-parallel strategies.
- Enhanced sharding and offloading: Colossal-AI redesigns tensor sharding and offloading with customizable strategies, lifecycle hooks, and chunked tensor organization.The design supports extensible zero-redundancy data parallelism and improves communication-bandwidth utilization and memory usage.
- Enhanced sharding and offloading: Reusing FP16 storage between forward-pass parameters and backward-pass gradients reduces redundancy and peak memory usage.This reuse allows CPU memory to accommodate larger models.
- Automatic parallelization: Experimental automatic parallelism uses greedy sharding conversion search and integrates activation checkpointing into performance optimization.The feature increases supported sharding dimensions and jointly considers sharding and checkpointing.
4 IMPLEMENTATION
Colossal-AI organizes distributed training through a modular architecture, user-friendly interfaces, and composable acceleration techniques. It minimizes code changes while providing reusable parallelized model components and extension points for customized strategies.
- System architecture: A parallel context manager tracks hybrid distributed-environment metadata and switches automatically to the corresponding parallel mode.The system also includes tensor-parallel model interfaces, activation checkpointing, mixed precision, an execution engine, and a trainer.
- System architecture: Modularity lets users combine different acceleration techniques to pursue maximal performance.The design also supports customized functions, including user-defined sharding strategies and life-cycle hooks for future training methods.
- User interface: User-friendly APIs inject configured acceleration features into the execution engine with minimal changes to sequential training code.Users specify features through a predefined configuration schema, initialize Colossal-AI, and retain familiar data-loading and engine-training operations.
- User interface: Users can directly employ parallelized BERT, GPT, and ViT components without manually designing their parallelism strategy.This reduces the need for domain expertise in distributed training.
- Usage workflow: The usage workflow configures a distributed network, defines training components, initializes Colossal-AI, and executes the training engine.The listing includes a configuration for 1D tensor parallelism with parallel size 4.
5 EVALUATION
Colossal-AI was evaluated across tensor, sequence, and pipeline parallelism under varied hardware conditions, using Megatron-LM and DeepSpeed as baselines. Results show that the best strategy depends on GPU topology, model scale, and memory pressure.
- Evaluation setup: Experiments used multiple hardware systems and compared Colossal-AI methods against Megatron-LM and DeepSpeed baselines.Only a portion of prominent features was tested on each system because of resource constraints.
- Memory Efficiency: 44% and 65% lower memory consumption were achieved by 2.5D and 3D tensor parallelism than 1D tensor parallelism at batch size 512 on 8 GPUs.The comparison reflects maximum allocated CUDA memory during forward and backward passes.
- Hardware Compatibility: 2D, 2.5D, and 3D tensor parallelism could not match 1D throughput on System I with 4 or 8 GPUs.1D benefits from System I’s high communication bandwidth, while advanced methods have higher communication volume with few processors.
- Hardware Compatibility: 2D and 2.5D tensor parallelism delivered 40% higher throughput than 1D with 4 GPUs and 2.5D remained 20.6% faster with 8 GPUs on System II.System II’s low collective-communication bandwidth bottlenecked 1D tensor parallelism, while 3D remained slower because of low scaling.
- Throughput Comparison: 2.76 times speedup was achieved by advanced tensor parallelism over 1D tensor parallelism as the number of GPUs increased from 4 to 64.The paper attributes this scaling advantage to lower communication volume with more processors.
- Sequence Parallelism: Sequence Parallelism supported a maximum batch size 4.44 times larger and training up to 1.43 times faster than 1D tensor parallelism for BERT-Base.It reduces duplicated activation memory by splitting activations along the sequence dimension and requires no activation gather between pipeline stages.
6 FUTURE WORK
Future work focuses on making distributed training more accessible through automated strategy selection and broader model integration.
- Future Work: Future work will design a hardware-aware algorithm to automatically search for an optimal parallelization strategy.The project also plans integration with model zoos such as Hugging Face Transformers.
7 CONCLUSION
The paper presents Colossal-AI as a unified system combining advanced acceleration techniques for large-scale distributed training. Its flexible methods provide robust performance across hardware conditions and achieve up to 2.76x speedup over baseline systems.
- Conclusion: Colossal-AI integrates many advanced acceleration techniques into one unified system with flexible combinations of parallelism methods.The system is designed for large-scale distributed training.
- Conclusion: 2.76x speedup over baseline systems was demonstrated in the experiments.The conclusion reports this as the maximum achieved speedup.