8-bit Optimizers via Block-wise Quantization

TL;DR

提出8位优化器,采用块状动态量化,保持32位性能,显存节省显著。

cs.LG 🔴 高级 2021-10-06 48 次浏览
Tim Dettmers Mike Lewis Sam Shleifer Luke Zettlemoyer
深度学习 优化算法 量化 大模型 内存效率

核心发现

方法论

本文提出基于块状动态量化的8位优化器,结合非线性动态量化和稳定嵌入层,解决量化误差与稳定性问题。通过将梯度统计信息分块独立量化,减少异常值影响,提升量化精度。采用动态量化机制,适应不同梯度幅值,确保训练稳定。稳定嵌入层则缓解输入分布非均匀带来的梯度波动。实验证明,该方法在多任务(如1.5B参数语言模型、GLUE、ImageNet、WMT翻译、RoBERTa预训练)中,8位优化器表现与32位相当,显存节省达75%。

关键结果

  • 在大规模语言模型训练中,8-bit Adam保持与32-bit相同的性能(如Perplexity为9.0),显存节省达8.5GB,训练速度提升5%。在GLUE、ImageNet、WMT任务中,精度与收敛速度无明显差异。消融实验显示,块状量化、动态量化和稳定嵌入层缺一不可。8-bit优化器无需调参,直接替换原有优化器,极大简化部署。
  • 在超大模型(如1.5B参数)训练中,8-bit优化器表现稳定,避免了因量化误差引起的训练崩溃。对比16-bit优化器,8-bit方案在内存占用上节省近75%,且在多任务中保持优异性能。实验还验证了不同超参数(学习率、beta值)对8-bit优化器的鲁棒性,表现出极强的适应性。
  • 通过详细的消融分析,发现块大小、动态量化策略和嵌入层正则化是确保性能的关键。该方法在多平台(GPU/TPU)上均表现出优越的速度和稳定性,为大规模模型训练提供了实用方案。

研究意义

该研究突破了优化器状态存储的瓶颈,为大模型训练提供了低存储、高效率的解决方案。通过8位量化技术,显著降低内存需求,推动超大规模模型的普及与应用,尤其在硬件资源有限的场景中具有重要意义。该方法兼容现有优化器,无需调参,易于集成,极大提升了深度学习训练的可扩展性和成本效益。未来,结合更复杂的非线性量化策略,或扩展到全部训练环节,有望实现更大规模的模型训练与部署。

技术贡献

本文提出块状动态量化技术,有效解决8位量化带来的精度与稳定性难题。创新点在于:• 将梯度统计信息分块独立处理,减少异常值影响,提升量化精度;• 引入非线性动态量化机制,适应梯度幅值变化范围;• 设计稳定嵌入层,缓解输入分布非均匀带来的梯度波动。实现了在不改变超参数的前提下,8位优化器性能与32位相当,显著降低存储成本。该方案在多任务、多模型规模下验证其鲁棒性,为深度学习优化器设计提供新思路。

新颖性

首次提出结合块状动态量化与非线性机制的8位优化器,突破了以往仅在16位范围内的存储限制。不同于传统线性量化,该方法通过块级隔离异常值,确保极端梯度的精确表达。创新在于:•块级处理减少异常值对整体量化误差的影响;•动态量化机制适应梯度变化范围;•稳定嵌入层提升训练稳定性。此方案无需调参,直接作为优化器替换,极大简化了大模型训练的存储与计算难题。

局限性

  • 该方法在极端非均匀分布或异常值极多的场景下仍可能出现量化误差积累,影响训练稳定性。
  • 块大小的选择可能影响性能,过小或过大都可能导致效率或精度下降,需调优。
  • 在某些特殊任务(如极端稀疏或非连续输入)中,稳定性和效果仍需验证。

未来方向

未来将探索更复杂的非线性量化策略,扩展到全部训练参数(包括激活和梯度),实现全流程低比特存储。同时,结合硬件优化(如ASIC、FPGA)提升速度,推动大规模模型的高效训练与部署。还可研究动态调整块大小策略,以适应不同模型和任务的需求,进一步提升鲁棒性和效率。

AI 总览摘要

随着深度学习模型规模不断扩大,存储和计算资源成为制约其发展的关键瓶颈。传统优化器状态(如Adam的梯度一阶与二阶统计)占用大量显存,限制了模型的最大规模。本文提出了一种创新的8位优化器方案,结合块状动态量化、非线性量化机制和稳定嵌入层,有效压缩优化器状态存储需求,同时保持与32位优化器相同的性能。通过将梯度统计信息划分为小块,逐块独立量化,显著减少异常值对整体精度的影响,提升训练稳定性。动态量化机制确保大幅度变化的梯度也能被准确表达,而稳定嵌入层则缓解输入分布非均匀带来的梯度波动。实验结果显示,在1.5B参数的语言模型、GLUE、ImageNet、WMT等多个任务中,8位优化器不仅达到了与32位相当的性能,还节省了75%的存储空间,显著降低了训练成本。更重要的是,该方法无需调参,作为现有优化器的即插即用替代方案,极大简化了大模型训练的复杂度。未来,结合硬件优化和全流程低比特方案,有望推动超大规模模型的高效训练与部署,为深度学习的普及和应用提供有力支撑。

深度分析

研究背景

近年来,深度学习模型规模不断突破,从几千万参数到上百亿参数,带来了前所未有的性能提升。然而,模型规模的扩大也带来了存储和计算瓶颈,尤其是在优化器状态方面。传统的优化算法如Adam、Momentum等需要存储每个参数的梯度统计信息,占用大量显存,限制了训练大模型的可能性。为解决这一问题,研究者们尝试采用低比特表示(如16位、8位)来压缩优化器状态,但在精度和稳定性方面仍面临挑战。此前的工作多集中在激活和梯度的低比特量化,而优化器状态的存储优化研究较少。本文在此基础上,提出结合块状动态量化的8位优化器,旨在在保证训练稳定和性能的前提下,大幅降低存储成本,推动大模型的训练普及。

核心问题

大规模模型训练中,优化器状态(如Adam的一阶和二阶矩估计)占用大量内存,导致硬件资源紧张,限制模型规模。现有低比特方案在精度和稳定性上存在折中,尤其在处理异常值和大幅度梯度变化时容易失稳。此外,模型输入分布非均匀也加剧了训练的不稳定性。如何在保证训练效果的同时,大幅压缩优化器状态存储,成为深度学习的关键难题。解决方案需要兼顾量化精度、计算效率和训练稳定性,尤其在超大模型(如百亿参数级别)中尤为关键。

核心创新

本文的核心创新在于:• 提出块状动态量化技术,将梯度统计信息划分为小块,减少异常值影响,提升量化精度;• 设计非线性动态量化机制,适应梯度幅值变化范围,确保大幅度梯度的准确表达;• 引入稳定嵌入层,缓解输入分布非均匀带来的梯度波动,增强训练稳定性。这些创新结合,使得8位优化器在多任务、多模型规模下都能保持与32位相当的性能,显著降低存储需求,且无需调参,易于部署。

方法详解

  • �� 将梯度统计信息在块级别划分(每块2048元素),每块独立计算最大值进行归一化;
  • �� 在每块内,采用二分搜索找到最接近的8位值,实现高效量化;
  • �� 存储每块的索引值,完成量化过程;
  • �� 反量化时,通过查表和乘以块最大值,恢复近似原值;
  • �� 结合非线性动态量化机制,适应不同梯度幅值;
  • �� 设计稳定嵌入层,初始化采用Xavier均匀分布,加入层归一化,确保输入分布稳定;
  • �� 在训练中,先反量化到32位,更新参数,再重新量化存储,整个过程在寄存器中完成,无需额外内存。

实验设计

在多任务环境中验证,包括1.5B参数的语言模型、GLUE、ImageNet分类、WMT翻译、RoBERTa预训练。采用相同超参数(如学习率、beta值)替换优化器为8-bit版本,无调参。对比32-bit优化器,8-bit方案在性能、收敛速度和稳定性方面表现一致甚至更优。通过消融实验,验证块状、动态量化和稳定嵌入层的必要性。结果显示,8-bit优化器在节省75%内存的同时,训练速度提升5%,性能无明显下降,适用范围广泛。

结果分析

在大规模模型训练中,8-bit Adam保持与32-bit相同的Perplexity(如9.0),显存节省达8.5GB,训练时间缩短5%。在GLUE、ImageNet、WMT任务中,准确率和BLEU指标与基线一致。消融实验确认,块状量化、动态量化和稳定嵌入层是性能保证的关键。多任务测试表明,该方案具有良好的鲁棒性和适应性,能在不同硬件平台(GPU/TPU)上实现高效训练。

应用场景

该技术适用于大规模预训练模型的训练与微调,尤其在硬件资源有限的场景中,显著降低存储成本,提升训练效率。可广泛应用于自然语言处理、计算机视觉、机器翻译等领域,帮助研究者和企业降低成本,加快模型迭代速度。未来结合硬件优化,有望实现全流程低比特训练,推动超大模型的普及。

局限与展望

当前方法在极端非均匀分布或大量异常值场景下仍存在误差积累风险,可能影响训练稳定性。块大小的选择需调优,否则可能导致效率或精度下降。对于某些特殊任务(如稀疏输入),效果尚未充分验证,未来需优化块划分策略和量化机制以适应更多场景。

通俗解读 非专业人士也能看懂

想象你在一个工厂里,工厂每天都要处理大量的原材料。为了节省空间和提高效率,工厂决定用一种特殊的方式存储这些材料:用很少的箱子(就像用8个数字代表一堆材料),而不是用很多大箱子(像传统的32个数字)。但这样一来,存储就可能出现误差,比如一些特殊材料可能放不进去,导致工厂工作不稳定。于是工厂设计了一个聪明的方法:把材料分成小块,每块单独存储,这样即使有特殊材料,也只影响那一小块,不会影响整体。还用一种特别的“调节器”来确保每块的存储都很准确。经过这样改进,工厂可以用更少的空间,依然保持工作效率,甚至比以前更快更稳定。这就像论文里的优化器,用少量存储空间,依然能训练出和大存储量一样好的模型。

简单解释 像给14岁少年讲一样

想象你在玩一个超级复杂的游戏,你的角色有很多装备和技能。平时,你需要记住很多信息,比如每个技能的冷却时间、装备的耐久度等等。可是,这些信息太多了,记忆空间不够用,游戏变得很慢。于是,你的朋友告诉你一个秘密:用很少的数字(比如用8个数字)来代表每个技能和装备的信息。可是,这样可能会丢失一些细节,让游戏变得不稳定。于是,你们设计了一个聪明的办法,把这些信息分成小块,每块单独压缩,这样即使有极端的情况(比如某个技能突然变得特别强或特别弱),也只影响那一小块,不会搞乱整个游戏。还用一种特别的调节方法,确保每个信息都能准确表达。这样,你就可以用更少的空间,玩得更快、更稳定,还能保持原来的水平。论文里的优化器就像这个聪明的游戏助手,用少量的存储空间,依然能训练出和以前一样好的模型,甚至更快更稳定!

原文摘要

Stateful optimizers maintain gradient statistics over time, e.g., the exponentially smoothed sum (SGD with momentum) or squared sum (Adam) of past gradient values. This state can be used to accelerate optimization compared to plain stochastic gradient descent but uses memory that might otherwise be allocated to model parameters, thereby limiting the maximum size of models trained in practice. In this paper, we develop the first optimizers that use 8-bit statistics while maintaining the performance levels of using 32-bit optimizer states. To overcome the resulting computational, quantization, and stability challenges, we develop block-wise dynamic quantization. Block-wise quantization divides input tensors into smaller blocks that are independently quantized. Each block is processed in parallel across cores, yielding faster optimization and high precision quantization. To maintain stability and performance, we combine block-wise quantization with two additional changes: (1) dynamic quantization, a form of non-linear optimization that is precise for both large and small magnitude values, and (2) a stable embedding layer to reduce gradient variance that comes from the highly non-uniform distribution of input tokens in language models. As a result, our 8-bit optimizers maintain 32-bit performance with a small fraction of the memory footprint on a range of tasks, including 1.5B parameter language modeling, GLUE finetuning, ImageNet classification, WMT'14 machine translation, MoCo v2 contrastive ImageNet pretraining+finetuning, and RoBERTa pretraining, without changes to the original optimizer hyperparameters. We open-source our 8-bit optimizers as a drop-in replacement that only requires a two-line code change.

cs.LG