Search papers, labs, and topics across Lattice.
The paper introduces Budget-Aware Optimizer Configurator (BAOC), a method for reducing GPU memory consumption during large-scale model training by assigning different optimizer configurations to different network blocks. BAOC samples gradient streams to estimate the performance risk associated with cheaper optimizer configurations (e.g., lower precision or removing momentum) for each block. It then solves a constrained optimization problem to minimize the total performance risk under memory and time constraints, resulting in block-specific optimizer configurations that maintain training quality while reducing memory usage.
Fine-tune optimizer precision block-by-block and slash memory use without sacrificing model quality.
Optimizer states occupy massive GPU memory in large-scale model training. However, gradients in different network blocks exhibit distinct behaviors, such as varying directional stability and scale anisotropy, implying that expensive optimizer states are not universally necessary and using a global optimizer is often memory-inefficient. We propose the Budget-Aware Optimizer Configurator (BAOC) to reduce memory cost by assigning suitable optimizer configurations to individual blocks under given budgets. Specifically, BAOC samples gradient streams to derive statistical metrics that quantify the potential performance risk of applying cheaper configurations (e.g., low precision or removing momentum). It then solves a constrained allocation problem to minimize total risk under memory and time budgets, selecting a budget-feasible configuration for each block. Experiments across vision, language, and diffusion workloads demonstrate that BAOC maintains training quality while significantly reducing the memory usage of optimizer states. The code is available at https://anonymous.4open.science/r/BAOC-45C6.
