要在大规模模型上实现数据并行(DP)、张量并行(TP)和流水线并行(PP)等分布式训练策略,需要用到实用的工具。虽然PyTorch等框架提供了基本的分布式数据并行(DistributedDataParallel),但训练拥有数千亿甚至数万亿参数的模型,需要更精密的内存管理和扩展能力。为此,专门的库变得不可或缺。DeepSpeed由微软研究院开发,是一个开源的深度学习优化库,旨在显著提高大型模型训练的速度和规模,同时最大限度地减少所需的硬件资源。它与PyTorch结合使用,提供了一系列优化功能,解决了大型训练中遇到的主要瓶颈,特别是GPU内存限制。标准数据并行中的根本问题是内存冗余。参与DP的每个GPU通常都包含模型权重、梯度和优化器状态(如Adam中的动量和方差缓冲区)的完整副本。对于拥有数十亿参数的模型,仅优化器状态所需的内存(通常以32位精度(FP32)存储)就可以超过即使是高端GPU的可用内存。考虑一个拥有 $1$ 亿参数的模型。权重可能需要 $4$ GB(FP32)或 $2$ GB(FP16)。然而,Adam优化器状态通常每个参数需要额外 $16$ 字节(动量 $4$ 字节,方差 $4$ 字节,FP32主权重 $4$ 字节,如果使用FP16,梯度 $4$ 字节),这意味着 每个GPU 额外增加 $16$ GB。这很快就变得难以维持。DeepSpeed通过几种创新的技术来应对这一问题及其他扩展难题:ZeRO(零冗余优化器): 这可以说是DeepSpeed最受认可的贡献。ZeRO是一系列优化方法,旨在消除数据并行训练中的内存冗余。ZeRO不是在所有数据并行GPU上复制优化器状态、梯度,甚至可能模型权重,而是将这些状态分区到可用设备上。这意味着每个GPU只保留整体状态的一部分,从而大幅减少了每个设备的内存占用。我们将在下一节中查看ZeRO的不同阶段(阶段1、阶段2和阶段3)。内存卸载: DeepSpeed允许将训练状态的部分(优化器状态、激活或参数)从GPU内存卸载到主机CPU的主内存甚至NVMe固态硬盘。虽然访问CPU内存或NVMe比GPU HBM慢,但卸载不常访问的状态可以释放宝贵的GPU内存,从而能够训练比原本能容纳的更大的模型。高效流水线并行: 除了用于数据并行的ZeRO之外,DeepSpeed还包含其自身高度优化的流水线并行实现。这使用户能够将模型的层分区到多个GPU上,减少每个GPU激活所需的内存,并通过分散计算图来支持更大的模型。自定义内核和优化器: DeepSpeed通常包含高度优化的CUDA内核,用于大型模型中常见的操作,例如自定义Transformer层或像FusedAdam这样高效的优化器,它们将多个步骤组合到一个内核启动中以获得更好的性能。简化的混合精度训练: 它提供了易于使用的工具,用于管理FP16或BF16混合精度训练,包括自动损失缩放。DeepSpeed的设计目标之一是易于集成。对于许多常见的使用场景,特别是使用ZeRO时,对现有PyTorch训练脚本只需进行最小的改动。核心修改通常包括使用DeepSpeed的initialize函数封装模型、优化器和数据加载器。以下是DeepSpeed集成的示意图:import torch import deepspeed # 假设模型、优化器、数据加载器已经是已定义的PyTorch对象 # DeepSpeed设置的配置字典(例如,ZeRO阶段,批次大小) config_params = { "train_batch_size": 32, "gradient_accumulation_steps": 1, "optimizer": { "type": "AdamW", "params": { "lr": 1e-5 } }, "fp16": { "enabled": True # 示例:启用混合精度 }, "zero_optimization": { "stage": 2 # 示例:启用ZeRO阶段2 } # ... 其他DeepSpeed配置 } # 初始化DeepSpeed引擎 model_engine, optimizer, _, _ = deepspeed.initialize( model=model, optimizer=optimizer, # 传递原始优化器 config_params=config_params ) # 训练循环修改: # 将 model(inputs) 替换为 model_engine(inputs) # 将 loss.backward() 替换为 model_engine.backward(loss) # 将 optimizer.step() 替换为 model_engine.step() for batch in dataloader: inputs, labels = batch inputs = inputs.to(model_engine.local_rank) # 将数据移动到正确的设备 labels = labels.to(model_engine.local_rank) outputs = model_engine(inputs) loss = calculate_loss(outputs, labels) # 你的损失计算 model_engine.backward(loss) model_engine.step()在此代码片段中,deepspeed.initialize接受标准PyTorch模型和优化器,以及一个配置字典(config_params),并返回一个model_engine。这个引擎在训练循环中取代了原始模型,并且它的方法(backward、step)根据提供的配置处理分布式训练、梯度累积、混合精度和ZeRO优化的复杂性。DeepSpeed提供了一整套工具,旨在使大型语言模型的训练更易于进行且更高效。它通过ZeRO专注于内存优化,结合对卸载和流水线并行的支持,使其成为工程师扩展模型规模的有力选择。我们现在将更仔细地查看ZeRO优化器的不同阶段。