PyAO: PyTorch-Based Memory-Efficient LLM Training on Ethernet-Interconnected Clusters

Area of Science:

• Artificial Intelligence
• Computer Science
• Machine Learning

Background:

• Large language models (LLMs) are increasing in size, leading to higher GPU memory demands.
• Existing multi-GPU distribution methods struggle with network latency in Ethernet clusters.
• Caching massive activations during forward passes further strains memory resources.

Purpose of the Study:

• To introduce PyAO, a novel system for managing memory consumption in large language models.
• To optimize activation offloading strategies for improved compute-to-communication ratios.
• To enable efficient training of large language models on Ethernet-interconnected clusters.

Main Methods:

• PyAO effectively offloads model activations to manage memory.
• It intelligently selects offloading strategies based on efficiency.
• The system minimizes data-movement bottlenecks to enhance throughput.

Main Results:

• PyAO reduced peak GPU memory by up to 1.94× for models like OPT-1.3B, GPT-0.8B, and Llama-1.2B.
• It enabled batch sizes up to 2.5× larger compared to baseline methods.
• Training acceleration reached up to 3.63× in Ethernet cluster environments.

Conclusions:

• PyAO significantly enhances the feasibility of training large language models on memory-constrained Ethernet clusters.
• The proposed activation offloading and optimization techniques overcome network latency challenges.
• PyAO offers a practical solution for scaling LLM training efficiently.