预算感知自动优化配置器

优化器状态在大规模 模型训练 中占用大量 GPU 内存。然而，不同网络块中的梯度表现出不同的行为，例如不同的方向稳定性和尺度各向异性，这意味着昂贵的优化器状态并不是普遍需要的，并且使用全局优化器通常是内存效率低下的。我们提出预算感知优化器配置器（BAOC），通过在给定预算下将合适的优化器配置分配给各个块来降低内存成本。具体来说，BAOC 对梯度流进行采样，以得出统计指标，量化应用更便宜的配置（例如，低精度或消除动量）的潜在性能风险。然后，它解决了约束分配问题，以最大限度地减少内存和时间预算下的总风险，为每个块选择预算可行的配置。跨视觉、语言和扩散工作负载的实验表明，BAOC 可以保持训练质量，同时显着减少优化器状态的内存使用量。该代码可从 https://anonymous.4open.science/r/BAOC-45C6 获取。