PAI-TorchAcc (Torch Accelerator) is a training acceleration framework for PyTorch. It uses GraphCapture technology to convert dynamic PyTorch graphs into static execution graphs. TorchAcc then applies distributed and computational optimizations to the computation graph, which improves the training efficiency of PyTorch models and makes the framework easier to use.
Technical overview
TorchAcc is a distributed training acceleration framework that combines the advantages of dynamic and static graphs. It provides the following features:
Converts dynamic graphs to static graphs using GraphCapture technology.
Improves training performance through compilation optimization.
You can reduce resource overhead by optimizing GPU memory.
Improves communication efficiency using optimization techniques such as half-precision communication, communication compression, and communication overlap.
Provides automatic and semi-automatic distributed policies to support efficient training for Large Language Models (LLMs).
Optimizes training data reads:
Prefetcher: Prefetches data to enable concurrent data pre-processing and training. This reduces wait times for data processing and improves training efficiency.
Packed dataset: Uses efficient data packaging to reduce unnecessary computations and improve data read efficiency.
Preprocess Cache: Caches pre-processed data to reduce data pre-processing overhead.
Service architecture
Deep learning frameworks can be divided into two main categories based on their execution mode:
Graph mode: Frameworks such as TensorFlow 1.* run in graph mode. This mode is easy to optimize, production-oriented, and supports integrated training and inference. However, it is not user-friendly, offers limited flexibility for writing code, and is difficult to develop and debug.
Eager mode: Frameworks such as PyTorch run in eager mode. This mode is flexible and easy to develop and debug. However, system-level optimization is difficult.
To address these issues, TorchAcc provides system-level optimizations for model training while maintaining the flexibility of PyTorch. The following figure shows the architecture of TorchAcc.
The core logic of TorchAcc is as follows:
TorchAcc uses LazyTensor and HybridDispatcher to convert eager execution in PyTorch into Intermediate Representation (IR) expressions. It then applies a series of optimizations, such as computation optimization, VARM optimization, and automatic parallelization. The optimized IR is then passed to the backend for further optimization and code generation. TorchAcc also provides multiple kernel implementations for manual operator optimization. These kernels can optimize specific operators to improve their computational efficiency and performance.