下载掌阅APP，畅读海量书库
立即打开
本节将向读者展示如何在真实的硬件环境中进行大型深度学习模型的训练。首先，将详细介绍如何设置和配置Colossal-AI环境，以便为大规模深度学习训练做好准备，目标是帮助读者理解并搭建他们自己的Colossal-AI环境。之后，会指导读者如何在配置好的Colossal-AI环境中训练他们的第一个模型。这里将提供详细的代码和步骤，从数据准备到模型的训练和评估，以帮助读者理解和掌握使用Colossal-AI进行模型训练的全过程。最后，将深入讨论在不同的硬件环境中训练大型AI模型的挑战和策略。这里将介绍如何利用Colossal-AI进行异构训练，比如在多GPU、多节点的设置中，如何优化内存使用，如何平衡计算和通信的需求等，目标是帮助读者理解并实施大型深度学习模型的异构训练。
本节将理论与实践相结合，让读者在了解大型深度学习模型训练的理论基础的同时，也能掌握如何在实际的硬件环境中进行训练的技巧和方法。
可以通过Python的官方索引来安装Colossal-AI软件包。
数据并行是实现加速模型训练的基本方法。通过两步可以实现训练的数据并行：第一步是构建一个配置文件，第二步是在训练脚本中修改很少的几行代码。
（1）Colossal-AI功能配置
为了使用Colossal-AI，在配置好文件后，Colossal-AI提供了一系列的功能来加快训练速度（包括模型并行、混合精度、零冗余优化器等）。每个功能都是由配置文件中的相应字段定义的。如果只用到数据并行，那么只需要具体说明并行模式。本例使用PyTorch最初提出的混合精度训练，只需要定义混合精度配置fp16=dict（mode=AMP_TYPE.TORCH）。
（2）全局超参数
全局超参数包括特定于模型的超参数、训练设置、数据集信息等。
（3）修改训练脚本（/data_parallel/train_with_cifar10.py）
导入Colossal-AI相关模块。
导入其他模块。
（4）启动Colossal-AI
在训练脚本中，在构建好配置文件后，需要为Colossal-AI初始化分布式环境，此过程称为launch。Colossal-AI提供了几种启动方法来初始化分布式后端。在大多数情况下，可以使用colossalai.launch和colossalai.get_default_parser来实现使用命令行传递参数。此外，Colossal-AI可以利用PyTorch提供的现有启动工具，正如许多用户通过使用熟知的colossalai.launch_from_torch那样进行相关操作。更多详细信息，可以参考相关文档。
初始化后，可以使用colossalai.core.global_context访问配置文件中的变量。
学会了Colossal-AI的环境搭建，现在读者可以在Colossal-AI的环境中训练属于自己的第一个模型了。
（1）构建模型
如果只需要数据并行性，则不用对模型代码进行任何更改。这里使用timm中的vit_base_patch16_224。
（2）构建CIFAR-10数据加载器
colossalai.utils.get_dataloader可以轻松构建数据加载器。
（3）定义优化器，损失函数和学习率调度器
Colossal-AI提供了自己的优化器、损失函数和学习率调度器。PyTorch的这些组件也与Colossal-AI兼容。
（4）启动用于训练的Colossal-AI引擎
Engine本质上是对模型、优化器和损失函数的封装类。当使用colossalai.initialize，将返回一个Engine对象，并且它已经按照配置文件中的指定内容，配置了梯度剪裁、梯度累积和零冗余优化器等功能。之后，基于Colossal-AI的Engine可以进行模型训练。
（5）训练：Trainer应用程序编程接口
Trainer是一个更高级的封装类，用户使用更少的代码就可以实现训练。通过传递Engine对象很容易创建Trainer对象。
此外，在Trainer中，用户可以自定义一些挂钩，并将这些挂钩连接到Trainer对象。钩子对象将根据训练方案定期执行生命周期方法。例如，LRSchedulerHook将执行lr_scheduler.step（）在after_train_iter或after_train_epoch阶段更新模型的学习速率。
使用trainer.fit进行训练。
（6）开始训练
DATA是自动下载和存储CIFAR-10数据集的文件路径。<NUM_GPUs> 是要用于使用CIFAR-10 数据集，以数据并行方式训练ViT的GPU数。
零冗余优化器（ZeRO）通过对3个模型状态（优化器状态、梯度和参数）进行划分而不是复制它们，消除了数据并行进程中的内存冗余。该方法与传统的数据并行相比，内存效率得到了极大的提高，而计算粒度和通信效率得到了保留。
● 分片优化器状态：优化器状态（如Adam Optimizer，32位的权重，以及一二阶动量估计）被划分到各个进程中，因此每个进程只更新其分区。
● 分片梯度：在梯度在数据并行进程组内进行Reduction后，梯度张量也被划分，这样每个进程只存储与其划分的优化器状态对应的梯度。注意，Colossal-AI将梯度转换为FP32格式以参与更新参数。
● 分片参数：16位的模型参数被划分到一个数据并行组的进程中。
● Gemini：对于参数、梯度、优化器状态的动态异构内存空间管理器。
下面将介绍基于Chunk内存管理的零冗余优化器。
使用零冗余优化器（ZeRO）时，需要通过切分参数的方式对模型进行分布式存储。这种方法的优点是每个节点的内存负载是完全均衡的。但是这种方式有很多缺点。首先，通信时需要申请一块临时内存用来通信，通信完毕释放，这会导致存在内存碎片化的问题。其次，以Tensor为粒度进行通信，会导致网络带宽无法充分利用。通常来说传输的消息长度越长带宽利用率越高。
利用ColossalAI v0.1.8引入了Chunk机制可以提升ZeRO的性能。将运算顺序上连续的一组参数存入一个Chunk中（Chunk即一段连续的内存空间），每个Chunk的大小相同。Chunk方式组织内存可以保证PCI-e和GPU-GPU之间网络带宽的高效利用，减小了通信次数，同时避免潜在的内存碎片。
在v0.1.8之前，ZeRO在进行参数聚合时通信成本较高，如果一个参数在连续的几次计算中被使用多次，即会发生多次通信，效率较低。这种情况在使用Checkpoint时非常常见，参数在计算Backward时会重计算一遍Forward。这种情况下，ZeRO的效率便不高。
以GPT为例，其Checkpoint会应用在每一个GPT Block上，每一个GPT Block包含一个Self-Attention层和MLP层。在计算Backward时，会依次计算Self-Attention层、MLP层的forward，然后依次计算MLP层、Self-Attention层的Backward。如果使用Chunk机制，将Self-Attention层和MLP层放在同一个Chunk中，在每个GPT Block的Backward的中便不用再通信。
除此之外，由于小Tensor的通信、内存移动没法完全利用NVLINK、PCIE带宽，而且每次通信、内存移动都有Kernel Launch的开销。使用了Chunk之后可以把多次小Tensor的通信、内存移动变为一次大Tensor的通信、内存移动，既提高了带宽利用，也减小了Kernel Launch的开销。
下面提供了轻量级的Chunk搜索机制，帮助用户自动找到内存碎片最小的Chunk尺寸，将运用GeminiDDP的方式来使用基于Chunk内存管理的ZeRO。这是新包装的torch.Module，它使用ZeRO-DP和Gemini，其中ZeRO用于并行，Gemini用于内存管理。同样需要确保模型是在ColoInitContext的上下文中初始化的。
定义模型参数如下：hidden dim是DNN的隐藏维度。用户可以提供这个参数来加快搜索速度。如果用户在训练前不知道这个参数，将使用默认值1024。min_chunk_size_mb是以兆字节为单位的最小块大小。如果参数的总大小仍然小于最小块大小，则所有参数将被压缩为一个小块。
（1）初始化优化器
（2）训练GPT
此例使用Hugging Face Transformers，并以GPT2 Medium为例。必须在允许该例程前安装Transformers。为了简单起见，这里只使用随机生成的数据。只需要引入Huggingface Transformers的GPT2LMHeadModel来定义模型，不需要用户进行模型的定义与修改，方便用户使用。详细代码请访问Colossal-AI教程链接。
1.NVMe Offload
如果模型具有 N 个参数，在使用Adam时，优化器状态具有8 N 个参数。对于10亿规模的模型，优化器状态至少需要32GB内存。GPU显存限制了可以训练的模型规模，这称为GPU显存墙。如果将优化器状态Offload到磁盘，便可以突破GPU内存墙。
编写参与人员实现了一个用户友好且高效的异步Tensor I/O库：TensorNVMe。有了这个库可以简单地实现NVMe Offload。该库与各种磁盘（HDD、SATA SSD和NVMe SSD）兼容。由于HDD或SATA SSD的I/O带宽较低，建议仅在NVMe磁盘上使用此库。
在优化参数时，可以将优化过程分为3个阶段：读取、计算和Offload，并以流水线的方式执行优化过程，这可以重叠计算和I/O，如图2-4所示。
图2-4 优化过程
在开始使用时，请先确保安装了TensorNVMe。
为Adam（CPUAdam和HybridAdam）实现了优化器状态的NVMe Offload。
nvme_offload_fraction是要Offload到NVMe的优化器状态的比例。nvme_offload_dir是保存NVMe Offload文件的目录。如果nvme_offload_dir为None，将使用随机临时目录。它与ColossalAI中的所有并行方法兼容。详细代码请访问Colossal-AI教程链接，相关章节将带读者用不同的方法训练GPT。
结果可以得到，NVMe卸载节省了大约294 MB内存。注意使用Gemini的pin_memory功能可以加速训练，但是会增加内存占用。所以这个结果也是符合预期的。如果关闭pin_memory，仍然可以观察到大约900 MB的内存占用下降。
2. 认识 Gemini：ColossalAI 的异构内存空间管理器
在GPU数量不足情况下，想要增加模型规模，异构训练是最有效的手段。它通过在CPU和GPU中容纳模型数据，并仅在必要时将数据移动到当前设备，可以同时利用GPU内存、CPU内存（由CPU DRAM或NVMe SSD内存组成）来突破单GPU内存墙的限制。在大规模训练下，其他方案如数据并行、模型并行、流水线并行都可以在异构训练基础上进一步扩展GPU规模。下面将描述Colossal-AI的异构内存空间管理模块Gemini的设计细节，它的思想来源于PatrickStar，Colossal-AI根据自身情况进行了重新实现。
（1）解决方案设计
目前的一些解决方案，DeepSpeed采用的Zero-Offload在CPU和GPU内存之间静态划分模型数据，并且它们的内存布局对于不同的训练配置是恒定的。图2-5左边所示，当GPU内存不足以满足其相应的模型数据要求时，即使当时CPU上仍有可用内存，系统也会崩溃。而图2-5右边所示Colossal-AI可以通过将一部分模型数据换出到CPU上来完成训练。
图2-5 优化过程
Colossal-AI设计了Gemini，就像双子星一样，管理CPU和GPU两者内存空间。它可以让张量在训练过程中动态分布在CPU-GPU的存储空间内，从而让模型训练突破GPU的内存墙。内存管理器由两部分组成，分别是MemStatsCollector（MSC）和StatefulTensorMgr（STM）。
这里可以利用深度学习网络训练过程的迭代特性，将迭代分为Warmup和Non-Warmup两个阶段，开始时的一个或若干迭代步属于预热阶段，其余的迭代步属于正式阶段。在Warmup阶段为MSC收集信息，而在Non-Warmup阶段STM利用MSC收集的信息来移动Tensor，以达到最小化CPU-GPU数据移动volume的目的，如图2-6所示。
图2-6 Gemini在不同训练阶段的运行流程
（2）StatefulTensorMgr
STM管理所有Model Data Tensor的信息。在模型的构造过程中，Colossal-AI把所有Model Data张量注册给STM。内存管理器给每个张量标记一个状态信息。状态集合包括HOLD、COMPUTE、FREE三种状态。STM的功能如下。
●查询内存使用：通过遍历所有Tensor的在异构空间的位置，获取模型数据对CPU和GPU的内存占用。
●转换张量状态：它在每个模型数据张量参与算子计算之前，将张量标记为COMPUTE状态，在计算之后标记为HOLD状态。如果张量不再使用则标记的FREE状态。
● 调整张量位置：张量管理器保证COMPUTE状态的张量被放置在计算设备上，如果计算设备的存储空间不足，则需要移动出一些HOLD状态的张量到其他设备上存储。Tensor Eviction Strategy需要MSC的信息将在后面介绍。
（3）MemStatsCollector
在预热阶段，内存信息统计器监测CPU和GPU中模型数据和非模型数据的内存使用情况，供正式训练阶段参考。通过查询STM可以获得模型数据在某个时刻的内存使用。但是非模型的内存使用却难以获取。因为非模型数据的生存周期并不归用户管理，现有的深度学习框架没有暴露非模型数据的追踪接口给用户。MSC通过采样方式在预热阶段获得非模型对CPU和GPU内存的使用情况。具体方法如下。
在算子的开始和结束计算时，触发内存采样操作，称这个时间点为采样时刻（Sampling Moment），两个采样时刻之间的时间称为Period。计算过程是一个黑盒，由于可能分配临时Buffer，内存使用情况很复杂。但是可以较准确地获取Period的系统最大内存使用。非模型数据的使用可以通过两个统计时刻之间系统最大内存使用-模型内存使用获得。
那么如何设计采样时刻呢？选择PreOp的Model Data Layout Adjust之前，如图2-7所示，可以采样获得上一个Period的System Memory Used，和下一个Period的Model Data Memory Used。并行策略会给MSC的工作造成障碍。比如对于ZeRO或者Tensor Parallel，由于Op计算前需要Gather模型数据，会带来额外的内存需求。因此，要求在模型数据变化前进行采样系统内存，这样在一个Period内，MSC会把PreOp的模型变化内存捕捉。比如在Period S2-S3内，考虑的Tensor Gather和Shard带来的内存变化。尽管可以将采样时刻放在其他位置，比如排除Gather Buffer的变动新信息，但是会给造成麻烦。不同并行方式Op的实现有差异，比如对于Linear Op，Tensor Parallel中Gather Buffer的分配在Op中。而对于ZeRO，Gather Buffer的分配是在PreOp中。将放在PreOp开始时采样有利于将两种情况统一。
图2-7 Sampling Based MemStatsCollector
尽管可以将采样时刻放在其他位置，比如排除Gather Buffer的变动新信息，但是会给造成麻烦。不同并行方式Op的实现有差异，比如对于Linear Op，Tensor Parallel中Gather Buffer的分配在Op中。而对于ZeRO，Gather Buffer的分配是在PreOp中。将放在PreOp开始时采样有利于将两种情况统一。
（4）Tensor Eviction Strategy
MSC的重要职责是在调整Tensor Layout位置，比如在图2-7的S2时刻，减少设备上Model Data数据，Period S2-S3计算的峰值内存得到满足。
在Warmup阶段，由于还没执行完毕一个完整的迭代，对内存的真实使用情况尚一无所知。此时限制模型数据的内存使用上限，比如只使用30%的GPU内存。这样保证可以顺利完成预热状态。
在Non-Warmup阶段，需要利用预热阶段采集的非模型数据内存信息，预留出下一个Period在计算设备上需要的峰值内存，这需要移动出一些模型张量。为了避免频繁在CPU-GPU换入换出相同的Tensor，引起类似Cache Thrashing的现象。利用DNN训练迭代特性，设计了OPT Cache换出策略。具体来说，在Warmup阶段，记录每个Tensor被计算设备需要的采样时刻。如果需要驱逐一些HOLD Tensor，那么选择在本设备上最晚被需要的Tensor作为受害者。
为了简化训练逻辑，本例使用合成数据训练BERT。
首先，如下所示，导入需要的依赖。依赖主要来自于PyTorch和Colossal-AI框架本身。这里使用Transformers库来定义BERT模型。注意这里只提及核心依赖，其他必要的依赖请参考完整示例代码。
接下来，传入分布式训练需要的参数和构建模型。Colossal-AI框架提供了十分方便的接口来创建分布式环境。
这里的config可以用来设置精度、并行策略等，本例不需要使用config。通过parser传入的参数（比如placement，dist_plan和vocab_size）可以根据需求手动设置，来确认使用的具体模型设置以及ZeRO策略。完整的示例代码可在随书资源中找到。本例可以拓展兼容PyTorch原生的数据并行进行比较，或者用张量并行来进一步优化，同时可以采用Torch Profiler来对性能进行测量。