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

As large language models (LLMs) pursue higher accuracy, their model sizes have surged, substantially increasing GPU memory consumption. Prior work mitigates this issue by distributing the memory burden across multiple GPUs. However, on clusters interconnected via Ethernet, the resulting computational intensity is insufficient to hide the significant network latency. Achieving a favorable compute-to-communication ratio is further constrained by the memory required to cache the massive activations generated during the forward pass. PyAO, proposed in this paper, effectively offloads activations, selects offloading strategies based on their offloading efficiency, and minimizes data-movement bottlenecks, thereby enabling larger micro-batch sizes. In Ethernet-interconnected cluster environments, experiments on popular models-including OPT-1.3B, GPT-0.8B, and Llama-1.2B-demonstrate that PyAO reduces peak GPU memory by up to 1.94× at the same micro-batch size, enables up to 2.5× larger batch sizes, and accelerates training by up to 3.63× relative to the baseline.