按模块分配优化器配置，用更少显存训练大模型
Budget-aware Auto Optimizer Configurator
- 根据梯度行为差异，为各网络块定制优化器配置
- 在显存和时间预算下，降低优化器状态内存占用50%以上
- 适合大模型训练中显存受限的场景
大规模模型训练中，优化器状态占据大量GPU显存。但不同网络模块的梯度行为差异显著，如方向稳定性与尺度各向异性不同，表明并非所有模块都需要高成本优化器状态，全局统一配置效率低下。本文提出预算感知优化器配置器（BAOC），在给定显存和时间预算下，为每个模块分配合适配置以降低内存开销。BAOC通过采样梯度流，计算低精度或移除动量等廉价配置可能带来的性能风险，再求解约束优化问题，在总风险最小的前提下实现各模块配置的可行分配。在视觉、语言和扩散模型任务上的实验表明，BAOC在保持训练质量的同时，显著减少优化器状态显存使用。代码已公开于 https://anonymous.4open.science/r/BAOC-45C6。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。