Adacc: An Adaptive Framework Unifying Compression and Activation Recomputation for LLM Training

TL;DR

Adacc is an adaptive framework that unifies activation recomputation and data compression, dynamically optimizing memory usage during LLM training to enhance efficiency without sacrificing accuracy.

Contribution

It introduces a fine-grained, tensor-level adaptive strategy that combines multiple memory optimization techniques with global scheduling and dynamic policy updates.

Findings

Improves training throughput by up to 1.37x

Maintains model accuracy comparable to baseline

Effectively balances memory savings and computational overhead

Abstract

Training large language models (LLMs) is often constrained by GPU memory limitations. To alleviate memory pressure, activation recomputation and data compression have been proposed as two major strategies. However, both approaches have limitations: recomputation introduces significant training overhead, while compression can lead to accuracy degradation and computational inefficiency when applied naively. In this paper, we propose Adacc, the first adaptive memory optimization framework that unifies activation recomputation and data compression to improve training efficiency for LLMs while preserving model accuracy. Unlike existing methods that apply static, rule-based strategies or rely solely on one technique, Adacc makes fine-grained, tensor-level decisions, dynamically selecting between recomputation, retention, and compression based on tensor characteristics and runtime hardware constraints. Adacc tackles three key challenges: (1) it introduces layer-specific compression algorithms that mitigate accuracy loss by accounting for outliers in LLM activations; (2) it employs a MILP-based scheduling policy to globally optimize memory strategies across layers; and (3) it integrates an adaptive policy evolution mechanism to update strategies during training in response to changing data distributions. Experimental results show that Adacc improves training throughput by 1.01x to 1.37x compared to state-of-the-art frameworks, while maintaining accuracy comparable to the baseline.