下载掌阅APP，畅读海量书库
立即打开
GPU高度并行的设计恰巧迎合了当今人工智能大模型训练的需求，然而随着参数量的不断膨胀，大模型对计算能力的需求也在逐步逼近GPU的上限。2023年初，一场名为ChatGPT的浪潮席卷全球，激发了人们对大型模型潜能的无限想象。其背后起支持作用的模型家族中，GPT-3（Generative Pretrained Transformer）达到了惊人的1750亿参数量。然而，这还仅仅是大型模型的冰山一角。在GPT-3之上，还出现了更为庞大的5000亿级模型，如Megatron Turing NLG等。这里可以简单计算一下，1750亿个参数如果以双精度浮点数存储需要约650GB，然而在实际训练中，参数的梯度、优化器的状态都需要额外占用很大的空间：如果使用Adam优化器，粗略估算梯度和优化器在训练GPT-3时会占用约三倍于模型参数大小的空间。雪上加霜的是，在模型的训练与推理中，往往会产生大量的中间变量，这些中间变量的内存占用取决于模型结构或批次大小，有时甚至会比模型本身的空间占用更大。即使使用现今能够购买到的最大显存的GPU（约80GB），仍然需要至少32个GPU才能开始训练。
如此巨量的参数带来了前所未有的计算挑战，在GPU的计算能力与存储空间都得以充分利用的前提下，借助其他计算设备来辅助计算过程便至关重要。这就引出了异构计算的概念。异构计算下的AI计算旨在统筹协调包括CPU、TPU、GPU、DRAM以及NVMe固态硬盘等在内的多种硬件设备，并充分利用各种设备的独特优势，来实现高效率的模型推理或训练。本节将重点关注通过异构计算实现的高效内存管理。
在GPU的显存非常有限的情况下，一个很自然的想法便是将模型状态挪出GPU来存储。于是，各种数据卸载方案便应运而生。然而对于计算设备，访问主机内存或者硬盘会带来远高于访问显存所需的通信开销，于是开发者们对于大模型中的高效数据存放进行了大胆的尝试。
异构训练是一种利用不同类型的计算设备进行深度学习模型训练的方法。它基于不同设备在计算能力、内存容量和能耗等方面的差异，通过将不同设备的优势结合起来，提高训练效率和性能。
异构训练的基本原理是将计算任务划分为不同的子任务，并将这些子任务分配给适合处理它们的计算设备。一般情况下，异构计算下的AI计算旨在统筹协调包括CPU、TPU、GPU、DRAM以及NVMe固态硬盘等在内的多种硬件设备，并通过充分利用各种设备的独特优势，提高深度学习模型训练的效率和性能。通过合理的任务划分和设备协同工作，异构训练可以加速模型的训练过程，降低训练时间，并提高训练效果。
在实现异构训练时，需要考虑以下几个关键方面。
1. 设备选择
根据任务的计算需求和数据规模，选择合适的设备进行训练。例如，GPU、TPU等硬件非常适合处理如矩阵运算等大规模并行计算任务，而CPU则在高度序列化且读写频繁的特定类型的计算任务上具有优势。
2. 任务划分
将训练过程的不同部分分配到不同的设备上进行计算。例如，将模型的前向和反向传播分配到GPU进行加速，而模型的参数更新则在CPU上进行。
3. 数据通信
在异构训练中，不同设备之间需要进行数据交换和通信。这包括在计算节点之间及在节点内的计算设备之间往复传递中间变量和模型状态等信息。高效的数据通信和同步机制对于异构训练的性能至关重要。
4. 调度和协调
在异构训练中，需要进行设备和通信任务的调度与协调，以确保各个环节间的阻塞最小，同时保证任务顺利执行。这可以通过异步通信、交换机内聚合操作等手段来实现。
异构训练的实现方式可以有多种，取决于具体的任务和计算设备。常见的实现方式包括多GPU训练、CPU-GPU协同训练、GPU-TPU协同训练等。这些方法都旨在充分利用不同设备的计算能力，提高训练速度和效率。
CPU-GPU协同训练是将模型的前向和反向传播部分分配给GPU进行加速，而模型的参数更新部分则在CPU上进行。这样可以利用GPU的并行计算能力加速模型的前向传播过程，并在CPU上进行高效的梯度计算和参数更新。GPU-TPU协同训练是利用GPU和TPU两种不同类型的计算设备进行训练。TPU作为专用加速器，在某些类型的计算任务上具有更高的计算效率。通过将模型的部分计算任务分配给TPU进行加速，可以进一步提高训练效率。
需要注意的是，在异构训练中，设备之间的数据通信和同步是一个关键的挑战。有效的数据传输和同步机制可以减少通信开销，确保设备之间的协同工作。同时，调度和协调任务的方式也会对异构训练的性能产生影响。合理的任务调度和设备协同机制可以提高系统的负载均衡和效率。
异构训练是利用不同计算设备进行深度学习模型训练的重要方法。通过合理的设备选择、数据划分、模型分布、数据通信和调度协调等方式，可以充分利用不同设备的优势，提高训练效率和性能。异构训练为大规模深度学习训练提供了一种有效的解决方案，加速了模型的训练过程，并促进了深度学习技术的发展和应用。下面将重点关注通过异构计算实现的高效内存管理。
在大规模深度学习训练中，异构训练是一种重要的策略，可以利用不同类型的计算设备进行加速，提高训练效率和性能。下面将介绍两种实现异构训练的策略：基于CPU和NVMe的卸载技术。
零冗余优化器（Zero Redundancy Optimizer，简称ZeRO）是2020年Rajbhandari等人提出的一种数据并行方案。ZeRO可以将模型状态均匀地拆分并分散地存储在不同的计算设备中。相较于传统数据并行，这种方案能够极大地减少大规模并行学习训练中的内存使用。ZeRO的具体原理会在第3章中详细说明。
在此基础上，ZeRO-Offload是一种实现CPU卸载的方法，它通过将优化器状态从GPU中转移到主机内存中来减小GPU的存储负担。同时，参数更新等运算也交由CPU同步执行，大大缓解了GPU的计算负担。在ZeRO-Offload中，主机CPU主要负责参数的更新与管理，而高强度的线性代数运算则由外部加速设备负责执行。其优势在于可以充分利用CPU的计算能力，提高训练速度和效率。它可以通过减少加速设备的计算和存储负载，释放更多的计算资源用于训练。
NVMe（Non-Volatile Memory Express）卸载技术是一种通过将数据存储和计算任务分离的方式来实现异构训练的策略。在传统的训练过程中，计算任务和数据存储通常位于同一个设备（如CPU或GPU）中，这会导致数据传输和通信的瓶颈。而NVMe卸载技术则将数据存储从计算设备中解耦，将数据存储在专用的NVMe设备中，并通过高速的PCIe总线与计算设备进行通信。
通过使用NVMe卸载技术，可以实现更高效的数据传输和通信，提高训练效率和性能。计算设备可以通过高速通道直接读取和写入存储设备中的数据，减少数据传输的延迟和开销。这种方式可以充分利用NVMe设备的高速读写能力，加速训练过程。
总结起来，CPU和NVMe卸载技术是实现异构训练的两种重要策略。ZeRO-Offload通过将计算任务从GPU转移到CPU和主机内存，充分利用各种设备的计算能力来提高训练效率。NVMe卸载技术通过将数据存储从计算设备中解耦，并使用高速通道进行数据传输，要使用ZeRO-Offload和NVMe卸载技术，需要进行一些关键的步骤和考虑一些重要因素。以下是一些常见的实现策略。
1）硬件配置和连接：首先，需要合适的硬件配置来支持异构训练。这包括具备较高计算能力的加速设备（如GPU、FPGA或ASIC）和专用的NVMe存储设备。此外，需要确保计算设备和存储设备之间有高速可靠的连接通道，如PCIe Gen4.0或更新的标准。
2）数据分发和同步：在异构训练中，数据的分发和同步是关键问题。需要设计有效的数据分发策略，将训练数据均匀地分配到不同的计算设备上，并确保计算设备之间的数据同步。这可以通过数据并行化和同步机制（如Allreduce）来实现。
3）存储管理和数据预取：对于卸载技术，需要有效管理存储设备中的数据，并实施数据预取机制。这可以通过缓存机制和预取算法来实现，以减少数据访问延迟并提高训练效率。
4）系统优化和性能调优：在实现异构训练时，还需要进行系统优化和性能调优。这包括针对不同硬件设备的优化策略，如GPU的计算图优化、FPGA的并行化设计等。此外，还可以通过调整参数和算法选择等方式来提高训练性能。
1. 早期的数据卸载 ——ZeRO-Offload 简介
数据卸载的历史中，一个颇为成功的方案便是基于ZeRO的ZeRO-Offload。其整个框架的思路非常简单直接：由GPU进行前向传播和反向传播，由CPU进行参数更新等优化器操作，如图2-3所示。这样一来，所有的优化器状态都可以放在主机内存而不是显存中，往返于GPU和内存间的数据只有模型参数和梯度。比起用到哪些参数传递哪些参数的原始想法，这样的卸载方案明显地减少了需要传输的数据量。
图2-3 ZeRO-Offload工作原理示意
从定量的角度审视一下这个方案。每一次训练迭代中，对GPU而言的计算复杂度通常记作O（MB），M为模型大小，B为批次大小。由于CPU的计算能力远低于GPU，便将梯度更新这种每次迭代中仅需一轮计算的操作分配给CPU，复杂度为O（M）。同时，这种方案下的每迭代通信总量也仅为O（MB）。同时为了进一步减小通信量，梯度和参数都以半精度浮点数的形式传递，只有主机内存保留一份单精度浮点数的副本。另一方面为了避免花费大量时间在等待数据传输上，ZeRO-Offload会尽可能重叠数据传输和计算的过程。
然而细心的读者可能很快就会发现ZeRO-Offload的局限性。没错，推理怎么办？由于ZeRO-Offload仅仅将模型训练带来的计算和存储压力用CPU消解掉了，然而计算结果的过程——也就是前向传播的计算量和内存需求并未发生任何变化。这样一来，一个GPU所能承载的模型大小上限仍然没有显著地提升。
2. 更多空间、更加智能 ——NVMe 卸载技术概览
作为后继工作，在将优化器移出GPU的基础之上，ZeRO-Infinity旨在进一步把模型的权重也移出GPU。所有运算需要的参数、变量等都只在运算时被取回GPU，完成后便立即释放。这样一来，理论上运行一个模型所需要的最小显存仅为该模型中最大层及其相关激活层的大小。但是内存是有极限的，越是挑战各种卸载，模型就越可能在某些地方崩溃……除非有其他地方存储。于是引出了ZeRO-Infinity一个最大的亮点：NVMe卸载。NVMe有相较于传统固态硬盘或机械硬盘更低的延迟与更高的带宽，非常适用于模型卸载这种高强度且对延迟敏感的输入输出。由于GPU与NVMe同样使用PCIe总线通讯，与卸载至内存相比，带宽并没有明显受制于通信条件。但仍需注意，NVMe固态硬盘的读写效率远低于DRAM，因此NVMe卸载的性能上限多取决于存储设备本身的读写效率。
为了帮助读者更便捷地评估ZeRO-Infinity的内存/存储需求，在这里定量地讨论一下模型训练中所需要的内存量。假设所有模型权重及优化器状态全部卸载至主机内存，模型参数量为M，所需的内存容量见表2-2。
表2-2 ZeRO-Infinity内存消耗估算
其中缓存的大小往往需要根据模型的具体工作原理和实际需求来确定。更糟的是，ZeRO-Infinity原始的实现并不能有效地卸载中间激活，因此ZeRO-Infinity往往会在训练含有较大中间激活的模型时溢出。激活暂存（Activation Checkpointing）便是一个非常好的解决方案——正常进行前向传播时，所有的中间变量均会被保存下来，以便在反向传播时计算梯度。然而在激活暂存启用时，只挑选前向传播中的某几步保存其输入，并且在反向传播时重算中间的各种变量。这种方法虽然会带来成倍的计算量，然而在数据卸载的背景下，比起传输巨量的激活，重新计算有时会更快一些。
另一方面，为了避免一些经常使用的模型权重和中间变量被往复传输，同时也为了有效地卸载计算中产生的激活，研究者们开发出了更加高效的张量管理方案——PatrickStar（派大星）。通过观察各张量的运行时特征，高效的内存管理器可以按需调度张量的移动与存储。这项功能已经被整合进主流加速框架Colossal-AI的Gemini模块中，这里可以简单体验一下模型卸载的强大力量。此处以GPT模型为例，使用HuggingFace提供的GPT-2模型，参数量约13亿。
通过上述代码，可以轻松地在单个GPU（笔者使用NVIDIA RTX4080，16GB显存）中加载一个10亿级参数的大语言模型。