“训练一个万亿参数的大模型，需要180张80GB显存的A100 GPU——但这样的成本和资源，普通团队根本用不起。”这是当前大模型训练领域的真实困境。近日，我国国防科技大学团队在《Frontiers of Information Technology & Electronic Engineering》发表论文，系统梳理了“GPU内存墙”难题的破解之道：从模型参数、训练状态到中间计算结果，逐一“挤”出内存空间，让大模型训练更高效、更普惠。

问题有多难？模型参数涨得比GPU显存快240倍
过去几年，大模型的参数规模像“坐火箭”——从BERT的3.4亿到GPT-3的1750亿，每两年增长约240倍。但GPU显存的增长却慢得多：从2016年的12GB（P100）到2023年的141GB（H200），仅翻了10倍。这就像“给不断膨胀的气球套了个固定大小的套子”——训练大模型时，GPU内存常常“爆仓”。

比如训练GPT-3，仅模型参数就需要650GB显存，再加上训练状态（如梯度、动量）和中间计算结果（如激活值），总需求超过2000GB。而单张A100 GPU只有80GB显存，必须靠多卡并行，但卡越多，卡之间的通信又成了新瓶颈。

怎么“挤”内存？三招破解参数、状态、激活“三大消耗户”
论文指出，GPU内存主要被三部分“吃掉”：模型参数（如权重）、训练状态（如梯度）、中间激活值（前向计算的中间结果）。科学家们针对这三部分，分别研发了“瘦身”策略：

第一招：给模型参数“分房住”
传统方法是让每张GPU复制完整模型（数据并行），但大模型参数太多，单卡存不下。于是科学家用“分块”思路：

模型并行：把模型按层或算子拆分，比如把Transformer的注意力层和全连接层分到不同GPU，每张卡只存一部分参数（如NVIDIA的Megatron-LM）。
流水线并行：把模型像工厂流水线一样分段，前一段计算完的结果传给下一段，减少单卡内存压力（如Google的GPipe）。
混合并行：结合数据并行、模型并行和流水线并行，像“三维空间”一样分配计算任务（如微软的DeepSpeed）。
第二招：给训练状态“减负”
训练时，优化器（如Adam）需要存储梯度、动量等状态，这部分内存消耗是模型参数的3倍。

ZeRO优化器：把梯度、动量等状态按卡数均分，每张卡只存1/N的状态（如用8张卡时，每张卡存1/8），大幅降低单卡内存需求。
PatrickStar：动态管理模型状态，把暂时不用的“搬”到CPU，需要时再“搬”回来，充分利用CPU和GPU的异构内存。
第三招：给中间结果“省着用”
前向计算的中间激活值（如每个神经元的输出）占内存也不小。科学家用两种方法“省内存”：

重计算：把暂时不用的激活值删掉，需要时再重新计算一遍（类似“需要时再算，别提前存着”）。
交换：把用不上的激活值“存”到CPU或硬盘，需要时再“读”回GPU（类似“把东西暂时放仓库，用的时候再拿”）。
未来挑战：平衡速度与内存，让大模型训练更“丝滑”
尽管已有多种方法，但论文也指出，当前技术仍有瓶颈：

并行通信慢：多卡并行时，卡之间传数据的时间可能比计算还长，拖慢整体速度。
优化器不稳定：为了省内存，有些优化器（如Adafactor）可能影响模型收敛，训练效果打折扣。
内存碎片多：动态调整内存时，容易出现“内存碎片”，导致可用空间被分割成小块，用不上。
未来，科学家们计划从三方面突破：一是研发新型并行策略（如结合序列并行和混合精度），二是优化训练状态管理（让优化器既省内存又稳定），三是提升重计算和交换的效率（减少“搬数据”的时间）。

从“内存爆仓”到“精准控存”，这些技术不仅让训练万亿参数模型成为可能，更让大模型研发从“巨头专属”走向“更多人可及”。或许不久的将来，一个小团队用几张GPU，也能训练出自己的“定制大模型”——而这，正是破解“GPU内存墙”的终极意义。

来源: 信息与电子工程前沿FITEE