小白入门的DeepSpeed原理和代码实现

这篇文章探讨了在多个GPU上分布式训练神经网络的技术，并测试了微软的DeepSpeed库提供的多种分布式训练优化。

训练的内存需求

本页的大部分内容将集中在减少训练神经网络所需的内存的技术上。因此，为了更好地理解和欣赏这些技术，我将首先分解训练过程中内存消耗的所有方式。

根据Sohoni等人的说法，我将使用以下术语来描述训练期间内存消耗的来源（Sohoni等人只描述了三个来源，但我将梯度内存和优化器内存分开，因为DeepSpeed中也做了同样的区分，稍后将展示）：

• 模型内存：存储模型权重所需的内存。

• 梯度内存：存储每个模型权重梯度所需的内存。

• 优化器内存：存储优化器所需的任何额外状态所需的内存。

• 激活内存：存储在前向传递期间计算的中间值所需的内存。

测量四种内存消耗源

让我们现在看看如何测量这四种来源消耗的内存量，从一个简单的全连接神经网络开始：

model = torch.nn.Sequential(       torch.nn.Linear(512, 1024),       torch.nn.Linear(1024, 1024),       torch.nn.Linear(1024, 1024),       torch.nn.Linear(1024, 512)   )   

要获得总的模型内存，我们计算模型中的参数数量，并乘以每个参数用来表示的字节数（在这种情况下是4个字节，因为每个权重是一个32位浮点数）：

def get_model_memory(model: torch.nn.Module):       """       返回给定模型的内存使用情况       """       total_memory = 0       for param in model.parameters():           total_memory += param.numel() * param.element_size()       return total_memory      print("模型内存：{:,} 字节".format(get_model_memory(model)))      

模型内存：12,597,248 字节   

梯度内存与模型内存相同，因为在后向传递期间我们需要为每个权重存储梯度。

消耗的总优化器内存取决于训练期间使用的优化器类型。下表显示了PyTorch中一些流行的优化器所需的优化器内存，其中M_model代表模型内存[2]：

最后，要计算总的激活内存，我们需要计算在前向传递期间计算的隐藏层输出。有几种不同的方法可以做到这一点；这里，我们将递归地注入PyTorch前向Hook到每个子模块中以捕获该子模块的输出大小。

首先，我们定义了Hook以及一个简单的辅助类ActivationCounter，用于跟踪模型隐藏层输出的总大小：

class ActivationCounter:          def __init__(self):           self.activation_bytes = 0          def add_activations(self, tensor):           self.activation_bytes += tensor.numel() * tensor.element_size()      def activation_counter_hook(counter: ActivationCounter):          def hook(self, input, output):           counter.add_activations(output.data)          return hook   

接下来，我们定义一个函数，它将递归地将Hook注入到模型的每个子模块中：

def register_hooks_recursive(model, counter: ActivationCounter):     for module in model.children():         module.register_forward_hook(activation_counter_hook(counter))         register_hooks_recursive(module, counter)      activation_counter = ActivationCounter()   register_hooks_recursive(model, activation_counter)   

现在要找到激活内存消耗的总量，我们执行模型的一个前向传递，然后打印activation_counter.activation_bytes的值：

inputs = torch.randn(4, 512)   outputs = model(inputs)      # 因为Hook只捕获层输出，我们需要单独添加原始输入张量的大小   activation_counter.add_activations(inputs)      print("激活内存：{:,} 字节".format(     activation_counter.activation_bytes   ))   

激活内存：65,536 字节   

计算密集神经网络模型内存和激活内存的完整脚本可以在GitHub仓库中找到：estimate_nn_memory.py

Transformer模型的内存需求

由于Rajbhandari, Samyam等人。（在他们的论文中介绍了DeepSpeed实现的许多训练优化）将实验集中在Transformer模型上，我想更好地了解这类模型的内存需求。

使用上述相同的方法，我计算了一个基本Transformer模型（详见附录A）的模型和激活内存，作为验证我估计Transformer模型内存的公式，以及Korthikanti, Vijay Anand等人提供的Transformer激活内存公式（都显示在下面）：

给定以下条件：

n：层数
h：隐藏层大小
m：注意力头数
b：批量大小
l：序列长度

分布式训练

DeepSpeed提供的优化主要适用于分布式训练环境（例如，当多个GPU用于训练单个模型时）。有三种常见的多GPU训练并行化范式：数据并行、模型并行和流水线并行。每个都在以下部分中简要解释。

数据并行

数据并行（DP）描述了一个分布式训练过程，整个模型在多个设备（例如GPU）上复制，每个设备并行地在不同的输入数据批次上执行训练。有多种不同的方法可以实现数据并行训练，但这里我只描述PyTorch DistributedDataParallel模块使用的实现[4]。

每个设备接收不同的输入数据批次，并执行前向传递以计算该批次的损失。

当在每个设备上计算梯度时，这些梯度与所有其他设备交换。然后使用梯度的平均值更新每个设备上的模型权重，确保在下一个训练步骤开始时，所有设备都有相同的一组模型权重。

AllReduce

设备之间的梯度交换是使用称为AllReduce的集体通信算法执行的。AllReduce算法在跨多个设备分布的数据上执行归约操作。

下面的可视化展示了在3个设备上执行AllReduce操作以计算分布在3个设备上的值向量的总和：

PyTorch通过其torch.distributed模块提供了对包括AllReduce在内的几种集体通信算法的支持。下面的脚本演示了在3个GPU上进行AllReduce操作：

import os   import torch   import torch.distributed as dist   import torch.multiprocessing as mp      def create_process_group(rank, world_size):       os.environ['MASTER_ADDR'] = 'localhost'       os.environ['MASTER_PORT'] = '12355'          dist.init_process_group(           backend='nccl',           world_size=world_size,           rank=rank       )      def all_reduce_example(rank, world_size):       create_process_group(rank, world_size)          # 在每个设备上创建不同的张量       if rank == 0:           tensor = torch.tensor([1, 2, 3]).to(rank)       elif rank == 1:           tensor = torch.tensor([10, 20, 30]).to(rank)       elif rank == 2:           tensor = torch.tensor([4, 5, 6]).to(rank)          print('Before AllReduce: Rank ', rank, ' has data ', tensor)       dist.all_reduce(tensor, op=dist.ReduceOp.SUM)       print('After AllReduce: Rank ', rank, ' has data ', tensor)      if __name__ == "__main__":       device_count = 3       mp.spawn(all_reduce_example, args=(device_count,), nprocs=device_count)      

Before AllReduce: Rank  0  has data  tensor([1, 2, 3], device='cuda:0')   Before AllReduce: Rank  1  has data  tensor([10, 20, 30], device='cuda:1')   Before AllReduce: Rank  2  has data  tensor([4, 5, 6], device='cuda:2')   After AllReduce:  Rank  2  has data  tensor([15, 27, 39], device='cuda:2')   After AllReduce:  Rank  0  has data  tensor([15, 27, 39], device='cuda:0')   After AllReduce:  Rank  1  has data  tensor([15, 27, 39], device='cuda:1')      

关于数据并行的一些说明

重要的是，数据并行不会减少训练所需的总内存，因为每个设备必须有足够的内存来执行整个模型的前向和后向传递。

因此，数据并行的主要好处是它加快了训练过程，因为每个设备只需要在总训练数据的一小部分上进行训练。然而，加速并不随着设备数量的增加而线性扩展，因为设备之间交换梯度引入了通信开销。例如，Li, Shen等人。2020年观察到，当在256个GPU上训练BERT模型时，扩展因子大约是128。

下图显示了使用1-8个GPU进行数据并行训练，在1000步训练中BERT模型所花费的时间：

模型并行

模型并行（MP）描述了一个分布式训练过程，模型被分割在多个设备上，以至于每个设备只包含模型权重的一部分。前向传递在每个设备上顺序执行，一个设备的输出成为下一个设备的输入。

模型并行在训练非常大的模型时至关重要，这些模型无法适应单个设备。

模型并行训练的一个主要缺点是，前向和后向传递的顺序性质导致每个设备上的空闲时间。例如，在上图中，GPU 0在完成输入批次的前向传递后，它在GPU 1和GPU 2执行它们的前向和后向传递时处于空闲状态。

通过使用PyTorch前向和后向Hook，我们可以测量每个GPU在训练期间的空闲时间：

def idle_time_hook(self, device, forward=True, entering=True):       ""“创建一个PyTorch Hook，记录设备的空闲时间。"""              def hook(*args, **kwargs):           current_timestamp = time.time()           last_timestamp = self.previous_timestamp.get(device, None)                      message = "{} {} pass on device {}".format(               "Entering" if entering else "Finished",               "forward" if forward else "backward",               device           )                      if entering and last_timestamp is not None:               idle_time_ms = (current_timestamp - last_timestamp) * 1000               self.device_idle_time[device] = (                   self.device_idle_time[device][0] + idle_time_ms,                   self.device_idle_time[device][1] + 1               )                              message += f". Idle time: {idle_time_ms:.2f}ms"                      self.previous_timestamp[device] = current_timestamp           self.log(message)       return hook      

2023-06-04 00:20:36.786470 - Entering forward pass on device 0.   2023-06-04 00:20:36.789441 - Finished forward pass on device 0   2023-06-04 00:20:36.790807 - Entering forward pass on device 1.   2023-06-04 00:20:36.793664 - Finished forward pass on device 1   2023-06-04 00:20:36.794982 - Entering forward pass on device 2.   2023-06-04 00:20:36.796839 - Finished forward pass on device 2   2023-06-04 00:20:36.798351 - Entering forward pass on device 3.   2023-06-04 00:20:36.799569 - Finished forward pass on device 3   2023-06-04 00:20:37.911536 - Entering backward pass on device 3. Idle time: 1111.96ms   2023-06-04 00:20:37.913496 - Finished backward pass on device 3   2023-06-04 00:20:37.914812 - Entering backward pass on device 2. Idle time: 1117.97ms   2023-06-04 00:20:37.916677 - Finished backward pass on device 2   2023-06-04 00:20:37.918056 - Entering backward pass on device 1. Idle time: 1124.39ms   2023-06-04 00:20:37.920743 - Finished backward pass on device 1   2023-06-04 00:20:37.922119 - Entering backward pass on device 0. Idle time: 1132.67ms   2023-06-04 00:20:37.923938 - Finished backward pass on device 0      

下面的时间线图示出了在三个GPU上执行模型并行训练时每个GPU的前向和后向执行。每个时间线中的灰色空间代表该GPU空闲的时间。

流水线并行

流水线并行（PP）是模型并行的一种变体，它通过将每个输入数据批次分割成多个较小的“微批次”来减少设备空闲时间[6]。模型参数仅在每个微批次被整个模型处理后更新，这意味着每个设备可以在其他设备仍在处理前一个微批次时开始处理下一个微批次。

支持流水线并行的功能已经内置在PyTorch中，通过torch.distributed.pipeline.sync.Pipe类。然而，这个类有两个主要限制，即（1）它只在模型被实现为torch.nn.Sequential模块时工作，以及（2）它要求每个模块的输入和输出是单个张量或张量元组[7]。

由于这些限制，我不得不修改HuggingFace Transformers库中的BERT模型实现以支持流水线并行。修改后的模型可以在这里找到：distributed-training-and-deepspeed/model/bert_mp.py

使用这个定制的BERT实现，我们可以通过首先将模型转换为torch.nn.Sequential模块，然后将其包装在torch.distributed.pipeline.sync.Pipe对象中来启用流水线并行：

def to_pipeline(self, chunks):       ""“将模型转换为流水线并行。"""              rpc.init_rpc(           name="worker",           rank=0,           world_size=1,           rpc_backend_options=rpc.TensorPipeRpcBackendOptions(               init_method="file://{}".format(tempfile.NamedTemporaryFile().name)           )       )          sequential = torch.nn.Sequential(           self.embeddings,           *self.encoders,           self.head       )       return Pipe(sequential, chunks=chunks)      

比较模型并行和流水线并行

为了比较模型并行和流水线并行的性能，我使用我的定制BERT实现执行了一些训练运行。下图显示了在2、4、6和8个GPU上使用模型并行和流水线并行执行250个训练步骤所花费的时间。对于这些测试，批量大小设置为16，对于流水线并行，微批次的数量设置为4。

DeepSpeed

有了对分布式训练技术的基本情况的了解，我们现在可以开始看看DeepSpeed。

DeepSpeed是微软创建的一个开源深度学习优化工具库，涵盖了模型训练、推理和模型压缩领域。这里，我将只关注与训练相关的优化，并将推理和压缩的主题留给未来的探索。

如前所述，我们可以通过在多个设备上复制整个模型（数据并行）和/或分割模型并将其不同部分存储在不同设备上（模型并行/流水线并行）来进行分布式训练。通常，数据并行在计算效率上比模型并行更高；然而，在模型太大而无法适应单个设备的可用内存的情况下，模型并行是必要的。

DeepSpeed提供的培训优化的核心是Zero Redundancy Optimizer（ZeRO），这是一套减少分布式模型训练所需内存的技术。

ZeRO

数据并行在设备之间引入了显著的内存冗余。例如，考虑在8个GPU上训练一个10亿参数模型。存储模型参数本身所需的内存（假设是32位浮点数）约为3.7GB。当使用数据并行时，每个设备必须存储整个模型的副本，这意味着在所有设备上总共使用了约29.8GB的内存。换句话说，26.1GB的设备内存被冗余的模型参数占据。激活和优化器内存也在设备之间类似地冗余。

ZeRO-DP通过以下方式解决这种冗余：

1. 分割优化器状态

2. 分割梯度

3. 分割模型参数

当使用deepspeed库时，我们通过指定ZeRO Stage 来选择应用哪些内存优化。可用的 Stage 包括：

在接下来的部分中，我们将详细探讨每个 Stage 。所有示例脚本已在装有8个16GB NVIDIA V100 GPU的单台机器上进行了测试。

Stage 0 - 入门

DeepSpeed库像任何其他Python包一样安装：

当使用DeepSpeed执行数据并行训练时，我们不需要像使用PyTorch的DistributedDataParallel那样设置进程组或显式地启动多个进程。相反，我们通过调用deepspeed.initialize将模型包装在DeepSpeedEngine中，它在内部处理所有的分布式训练逻辑。

对于这些示例，我们将使用HuggingFace的预训练的、560M参数变体的BLOOM模型[9]。

# deepspeed_stage_0.py   #   # DeepSpeed自动设置LOCAL_RANK环境变量   # 为当前设备的索引。   rank = int(os.getenv("LOCAL_RANK", "0"))      model = BloomForCausalLM   .from_pretrained("bigscience/bloom-560m")      deepspeed_config = {       "train_micro_batch_size_per_gpu": 1   }      model_engine, _, _, _ = deepspeed.initialize(       model=model,       model_parameters=model.parameters(),       config=deepspeed_config   )      print(f"Device {rank} - ZeRO Stage: {model_engine.zero_optimization_stage()}")      

要启动分布式训练作业，我们使用deepspeed命令行实用程序，它与deepspeed Python包一起安装：

deepspeed deepspeed_stage_0.py      

Device 1 - ZeRO Stage: 0   Device 4 - ZeRO Stage: 0   Device 2 - ZeRO Stage: 0   Device 3 - ZeRO Stage: 0   Device 6 - ZeRO Stage: 0   Device 7 - ZeRO Stage: 0   Device 0 - ZeRO Stage: 0   Device 5 - ZeRO Stage: 0      

在这里我们看到模型复制在所有四个设备上。要允许DeepSpeed创建和管理训练期间使用的优化器，我们向deepspeed_config对象添加适当的配置选项：

deepspeed_config = {       "train_micro_batch_size_per_gpu": 1,       "optimizer": {           "type": "Adam",           "params": {               "lr": 5e-5           }       },   }      model_engine, optimizer, _, _ = deepspeed.initialize(       model=model,       model_parameters=model.parameters(),       config=deepspeed_config   )      # 打印优化器状态信息   optimizer_state = optimizer.param_groups[0]   print(f"Device {rank} - Optimizer: lr={optimizer_state['lr']}; "         f"betas={optimizer_state['betas']}; eps={optimizer_state['eps']}; "         f"parameter count={sum([torch.numel(p) for p in optimizer_state['params']]):,}")      

Device 1 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592   Device 4 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592   Device 2 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592   Device 3 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592   Device 6 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592   Device 7 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592   Device 0 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592   Device 5 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592      

训练DeepSpeedEngine与使用数据并行训练PyTorch模型非常相似。唯一的区别是backward()和step()方法直接在DeepSpeedEngine对象上调用，而不是在损失或优化器上：

# DeepSpeed自动设置WORLD_SIZE环境变量   # 为参与训练作业的设备数量。   world_size = int(os.getenv("WORLD_SIZE", "1"))      tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")   collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)      # 辅助函数加载wikitext数据集   # 实现可以在这里找到：   # https://github.com/gnovack/distributed-training-and-deepspeed/blob/main/util.py   train_dataset = load_wikitext(tokenizer, collator).select(range(64))      train_dataloader = DataLoader(       train_dataset,       batch_size=1,       shuffle=False,       sampler=DistributedSampler(train_dataset, num_replicas=world_size)   )      for batch in train_dataloader:       device = torch.device("cuda", rank)          input_ids = batch['input_ids'].to(device)       labels = batch['labels'].to(device)          outputs = model_engine(input_ids, labels=labels)       model_engine.backward(outputs.loss)       model_engine.step()      

Stage 1 - 优化器状态分割

ZeRO Stage-1将优化器状态（例如，当使用Adam时，一阶和二阶动量）在所有设备上进行分割，使得每个设备只包含这部分状态的一部分。

下图展示了使用传统数据并行（顶部）和ZeRO Stage-1（底部）训练BLOOM模型时每个设备的内存分配情况：

基于之前讨论的Transformer模型内存公式，我们可以估计存储一个560M参数模型的模型参数所需的内存约为2.24GB。当以32位模式使用Adam时，优化器状态将占用每个模型参数8字节的内存，因此优化器将需要4.48GB的内存。

假设我们以批量大小1和输入序列长度512来训练模型，我们可以使用上述公式估计激活所需的内存为1.75GB。

因此，使用传统数据并行训练BLOOM 560M模型所需的总内存大致为：

deepspeed库包括几个实用工具函数，我们可以使用它们来分析训练期间的内存使用情况。例如，要测量当前和峰值内存使用情况，我们在训练循环内调用memory_status()函数：

from deepspeed.runtime.utils import memory_status      for batch in train_dataloader:       input_ids = batch['input_ids'].to(device)       labels = batch['labels'].to(device)          outputs = model_engine(input_ids, labels=labels)          model_engine.backward(outputs.loss)          model_engine.step()          if rank == 0:           memory_status("Memory stats after training step")   

如果我们现在使用deepspeed CLI执行训练脚本，我们将在每个步骤后看到以下行打印出来：

RANK=0 MEMSTATS Memory stats after training step device=cuda:0 current alloc=6.8847GB (delta=0.0000GB max=10.8823GB) current cache=13.4453GB (delta=0.0000GB max=13.4453GB)   

如所示，GPU的峰值内存使用量为10.88GB，非常接近我们估计的10.71GB。接下来，让我们看看使用ZeRO Stage-1可以节省多少内存。要启用ZeRO Stage-1，我们更新deepspeed_config对象以包含zero_optimization配置选项：

deepspeed_config = {       "train_micro_batch_size_per_gpu": 1,       "optimizer": {           "type": "Adam",           "params": {               "lr": 5e-5           }       },       "zero_optimization": {           "stage": 1,       }   }      model_engine, optimizer, _, _ = deepspeed.initialize(       model=model,       model_parameters=model.parameters(),       config=deepspeed_config   )   

现在执行训练脚本时，我们会注意到一些差异。首先，在每个设备上打印优化器参数数量时，显示的参数数量现在是69,901,824，或者是模型中总参数数量的18分之1：

Device 0 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824   Device 6 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824   Device 7 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824   Device 4 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824   Device 5 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824   Device 3 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824   Device 2 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824   Device 1 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824   

其次，当打印每个训练步骤后的内存使用情况时，GPU的峰值内存使用量现在是7.48GB，比使用传统数据并行时测量的10.88GB少了3.4GB：

RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=3.4843GB (delta=0.0000GB max=7.4816GB) current cache=12.7695GB (delta=0.0000GB max=12.7695GB)   

那么这额外的3.4GB来自哪里呢？回想一下，优化器状态占用了大约4.48GB。以前，当使用传统数据并行时，每个设备都持有所有4.48GB的优化器状态。现在，有了ZeRO Stage-1，每个设备只持有优化器状态的18分之1，即0.56GB。

理论上，这意味着使用ZeRO Stage-1时的峰值GPU内存应该是：

2.24 + 2.24 + 0.56 + 1.75 = 6.79GB

我们测量的7.48GB而不是6.79GB可能归因于用于通信更新模型权重的中间缓冲区的使用，但需要进一步调查以确认这一点。

Stage 2 - 梯度分割

ZeRO Stage-2通过在设备间分割优化器状态和梯度，更进一步。下图展示了使用ZeRO Stage-2训练时每个设备的内存分配情况：

Stage-2同样可以通过在deepseed_config中设置stage选项来启用：

deepspeed_config = {       "train_micro_batch_size_per_gpu": 1,       "optimizer": {           "type": "Adam",           "params": {               "lr": 5e-5           }       },       "zero_optimization": {           "stage": 2       }   }   

如果我们启用ZeRO Stage-2并执行训练脚本，我们将发现GPU的峰值内存使用量略有增加，达到8.22GB，相比之下，使用 Stage-1训练时为7.48GB：

RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=3.4843GB (delta=0.0000GB max=8.2187GB) current cache=12.7715GB (delta=0.0000GB max=12.7715GB)   

为了理解GPU峰值内存意外增加的原因，我们需要了解一下ZeRO Stage-2的实现。在后向传递过程中，梯度被平均化并通过一系列Reduce操作放置在一个设备上。

虽然我们可以在计算每个梯度后执行Reduce操作，但通过分块减少梯度可以实现更好的性能。

这些块的大小由deepspeed_config中的reduce_bucket_size选项控制，默认值为5×10^8个元素，或者当每个元素是一个32位浮点数时为2GB。虽然这个默认值适用于更大的模型，但对于我们的560M模型来说太大了，它的梯度总共只占用了大约2.24GB。

如果我们将reduce_bucket_size减少到5×10^6个元素，我们可以将峰值内存使用量降低到6.37GB，与 Stage-1相比额外减少了1.11GB：

RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=3.4843GB (delta=0.0000GB max=6.3750GB) current cache=10.9277GB (delta=0.0000GB max=10.9277GB)   

值得考虑的是6.37GB的峰值内存是否符合我们的预期。鉴于每个设备持有优化器状态的18分之1和梯度的18分之1，我们应该预期峰值内存在：

2.24 + 2.24/8 + 4.48/8 + 1.75 = 4.83GB

这个理论估计假设后向传递中每个操作产生的梯度小于桶大小。 这篇文章探讨了在多个GPU上分布式训练神经网络的技术，并检验了微软的DeepSpeed库提供的多种分布式训练优化。

阶段3 - 参数分割

ZeRO第三阶段在第二阶段的基础上，进一步分割模型参数以及优化器状态和梯度。当我们使用ZeRO第三阶段训练BLOOM 560M模型时，我们会注意到每个训练步骤后分配的内存从第二阶段的3.48GB减少到1.94GB。然而，GPU的峰值增加到7.48GB。

RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=1.9410GB (delta=0.0000GB max=7.4082GB) current cache=10.7129GB (delta=0.0000GB max=10.7129GB)   

目前，GPU峰值内存使用量增加的原因尚不清楚，但我在DeepSpeed GitHub仓库中有一个开放的问题，以更好地理解这种行为：DeepSpeed/issues/3734

通信开销

虽然ZeRO可以减少训练所需的总内存量，但它确实引入了额外的通信开销，因为优化器状态、梯度和模型参数必须在设备之间频繁交换。

DeepSpeed允许我们通过其通信日志设置来测量通信操作所花费的时间。通信日志通过在deepspeed_config中添加comms_logger部分来启用：

deepspeed_config = {       "train_micro_batch_size_per_gpu": 1,       "optimizer": {           "type": "Adam",           "params": {               "lr": 5e-5           }       },       "zero_optimization": {           "stage": 0       },       "comms_logger": {           "enabled": True,           "verbose": False,           "prof_all": True,           "debug": False       }   }   

现在，执行训练脚本时，每个通信操作（AllReduce、AllGather、Reduce Scatter等）将被捕获和测量。要查看这些测量结果，我们调用deepspeed.comm.log_summary()函数：

## 训练循环后   import deepspeed.comm as dist   dist.log_summary()   

Comm. Op            Message Size        Count               Total Latency(ms)   Avg Latency(ms)     tput_avg (Gbps)     busbw_avg (Gbps)   broadcast                       4.0 KB              221                 20.58               0.09                0.36                0.36                       8.0 KB              1                   0.17                0.17                0.38                0.38                       16.0 KB             24                  2.55                0.10                1.25                1.25                       113.27 KB           1                   0.09                0.09                10.70               10.70                       2.0 MB              1                   0.44                0.44                38.01               38.01                       4.0 MB              97                  71.00               0.72                46.41               46.41                       16.0 MB             48                  155.34              3.23                41.67               41.67                       113.27 MB           1                   178.26              178.26              5.33                5.33   all_reduce                       1.24 GB             50                  22385.67            445.17              48.11               36.08   log_summary_barrier                       0B                  1                   69.95               69.95               0.00                0.00      

利用这些日志，我测量了ZeRO每个阶段引入的通信开销。结果如下所示：

多节点训练

到目前为止，我们所看到的所有示例都在单个节点上展示了多个GPU的分布式训练。DeepSpeed也可以应用于多节点训练。要执行多节点训练DeepSpeed，我们可以使用与之前相同的训练脚本，但需要一些额外的设置，以允许多个节点相互通信。本节概述了在多个AWS EC2实例上执行多节点训练的步骤。

假设我们有三个能够相互通信的所有端口的EC2实例。使用AWS CDK创建这些实例的示例包含在GitHub仓库中：distributed-training-and-deepspeed/aws/multi_node_training_stack.py。

为简化此示例，我将使用worker-1、worker-2和worker-3代替三个GPU实例的实际IP地址，并使用localhost来指代我的本地机器。如果你想使用工作器名称进行操作，你可以将它们添加到你的~/.ssh/config文件中，如下所示，将HostNames替换为你的实例的实际IP地址，将IdentityFile替换为你的AWS密钥对的路径：

Host worker-1     HostName ec2-34-222-121-132.us-west-2.compute.amazonaws.com     IdentityFile /home/me/keypair.pem     User ubuntu      Host worker-2     HostName ec2-32-27-113-7.us-west-2.compute.amazonaws.com     IdentityFile /home/me/keypair.pem     User ubuntu      Host worker-3     HostName ec2-52-88-92-190.us-west-2.compute.amazonaws.com     IdentityFile /home/me/keypair.pem     User ubuntu      

要运行多节点训练作业，每个节点都需要能够通过无密码SSH连接到每个其他节点。为此，我们将创建一个新的SSH密钥对并将其分发到每个节点：

# 在localhost上      # 创建一个新的SSH密钥对   mkdir -p multi-node-training-keys   ssh-keygen -t rsa -N "" -f ./multi-node-training-keys/id_rsa      # 将SSH密钥复制到worker-1   scp ./multi-node-training-keys/id_rsa.pub worker-1:/home/ubuntu/.ssh/id_rsa.pub   scp ./multi-node-training-keys/id_rsa worker-1:/home/ubuntu/.ssh/id_rsa      # 在工作器节点上将SSH密钥添加到authorized_keys   ssh worker-1 'cat /home/ubuntu/.ssh/id_rsa.pub >> /home/ubuntu/.ssh/authorized_keys'      # 为worker-2和worker-3重复上述三个命令      

接下来，我们需要在每个工作器上安装DeepSpeed及其依赖项。这个页面的GitHub仓库包括一个requirements.txt文件和一个worker-prereqs.sh脚本，可以用来安装所有先决条件。在每个节点上，我们将克隆GitHub仓库并安装先决条件：

# 在localhost上      # 在worker-1上安装依赖项   ssh worker-1 'rm -rf distributed-training-and-deepspeed && git clone https://github.com/gnovack/distributed-training-and-deepspeed.git && pip install -r distributed-training-and-deepspeed/requirements.txt && distributed-training-and-deepspeed/scripts/worker-prereqs.sh'      # 为worker-2和worker-3重复上述命令      

最后，虽然不是必需的，但在每个工作器上配置SSH主机名别名是很好的。我们可以通过在每个工作器上写入/home/ubuntu/.ssh/config文件来实现：

# 在localhost上；假设你已经在你的~/.ssh/config文件中添加了SSH主机名别名      worker_ip_1=$(ssh -G worker-1 | awk '$1 == "hostname" { print $2 }')   worker_ip_2=$(ssh -G worker-2 | awk '$1 == "hostname" { print $2 }')   worker_ip_3=$(ssh -G worker-3 | awk '$1 == "hostname" { print $2 }')      ssh worker-1 'cat > /home/ubuntu/.ssh/config' << EOF   Host worker-1       HostName $worker_ip_1       StrictHostKeyChecking no      Host worker-2       HostName $worker_ip_2       StrictHostKeyChecking no      Host worker-3       HostName $worker_ip_3       StrictHostKeyChecking no   EOF      # 为worker-2和worker-3重复上述命令      

通过设置StrictHostKeyChecking no，我们绕过了手动接受每个工作器的SSH主机密钥的需要。这简化了这个例子，但不建议在生产环境中这样做。

一个执行所有这些命令的完整脚本适用于所有三个工作器节点，在distributed-training-and-deepspeed/scripts/generate-keys.sh中可用。

现在每个节点都可以通过无密码SSH连接到每个其他节点，下一步是创建一个hostfile，列出每个节点的IP地址，并指定每个节点上可用的GPU数量。假设我们的~/.ssh/config中有每个工作器节点的别名，我们可以使用这些别名而不是IP地址在我们的hostfile中：

worker-1 slots=1   worker-2 slots=1   worker-3 slots=1      

随着hostfile的创建，我们准备 运行一个多节点训练作业。我们可以从任何一个工作器节点启动作业。对于这个例子，我们将从worker-1开始作业，通过将hostfile从localhost复制到worker-1，然后使用--hostfile参数运行训练脚本：

# 在localhost上      # 将hostfile复制到worker-1   scp hostfile worker-1:/home/ubuntu/distributed-training-and-deepspeed/hostfile      # 从SSH配置文件中获取worker-1的IP地址   MASTER_ADDR=$(ssh -G worker-1 | awk '$1 == "hostname" { print $2 }')      # ssh到第一个工作器节点并启动训练   ssh worker-1 'cd distributed-training-and-deepspeed && PATH="/home/ubuntu/.local/bin:$PATH" deepspeed --master_addr=$MASTER_ADDR --hostfile=./hostfile zero_dp_training.py --stage=2 --model_name=facebook/opt-125m'      

总结

这次对分布式训练和DeepSpeed的探索让我学到了很多关于分布式模型训练的细微差别和挑战，以及零冗余优化器如何在多个GPU上训练时减少内存需求。DeepSpeed中还有许多其他优化没有在本文中涵盖，这为我未来的探索留下了几个话题：

1. 卸载：DeepSpeed提供了将优化器状态和模型参数卸载到CPU内存和磁盘的能力，这可以进一步减少内存消耗，以训练时间为代价。

2. 专家混合（MoE）：MoE模型是一类使用稀疏激活层（即在每次前向传递中选择一部分权重使用的层）来增加模型参数总数而不增加计算复杂性的模型。DeepSpeed库包括稀疏激活的PyTorch模型层，可以用来实现MoE模型。

3. 渐进式层丢弃（PLD）：PLD是一种通过允许在训练过程中的不同点切换特定的Transformer层来加速Transformer模型训练的技术。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

第一阶段（10天）：初阶应用

该阶段让大家对大模型 AI有一个最前沿的认识，对大模型 AI 的理解超过 95% 的人，可以在相关讨论时发表高级、不跟风、又接地气的见解，别人只会和 AI 聊天，而你能调教 AI，并能用代码将大模型和业务衔接。

• 大模型 AI 能干什么？
• 大模型是怎样获得「智能」的？
• 用好 AI 的核心心法
• 大模型应用业务架构
• 大模型应用技术架构
• 代码示例：向 GPT-3.5 灌入新知识
• 提示工程的意义和核心思想
• Prompt 典型构成
• 指令调优方法论
• 思维链和思维树
• Prompt 攻击和防范
• …

第二阶段（30天）：高阶应用

该阶段我们正式进入大模型 AI 进阶实战学习，学会构造私有知识库，扩展 AI 的能力。快速开发一个完整的基于 agent 对话机器人。掌握功能最强的大模型开发框架，抓住最新的技术进展，适合 Python 和 JavaScript 程序员。

• 为什么要做 RAG
• 搭建一个简单的 ChatPDF
• 检索的基础概念
• 什么是向量表示（Embeddings）
• 向量数据库与向量检索
• 基于向量检索的 RAG
• 搭建 RAG 系统的扩展知识
• 混合检索与 RAG-Fusion 简介
• 向量模型本地部署
• …

第三阶段（30天）：模型训练

恭喜你，如果学到这里，你基本可以找到一份大模型 AI相关的工作，自己也能训练 GPT 了！通过微调，训练自己的垂直大模型，能独立训练开源多模态大模型，掌握更多技术方案。

到此为止，大概2个月的时间。你已经成为了一名“AI小子”。那么你还想往下探索吗？

• 为什么要做 RAG
• 什么是模型
• 什么是模型训练
• 求解器 & 损失函数简介
• 小实验2：手写一个简单的神经网络并训练它
• 什么是训练/预训练/微调/轻量化微调
• Transformer结构简介
• 轻量化微调
• 实验数据集的构建
• …

第四阶段（20天）：商业闭环

对全球大模型从性能、吞吐量、成本等方面有一定的认知，可以在云端和本地等多种环境下部署大模型，找到适合自己的项目/创业方向，做一名被 AI 武装的产品经理。

• 硬件选型
• 带你了解全球大模型
• 使用国产大模型服务
• 搭建 OpenAI 代理
• 热身：基于阿里云 PAI 部署 Stable Diffusion
• 在本地计算机运行大模型
• 大模型的私有化部署
• 基于 vLLM 部署大模型
• 案例：如何优雅地在阿里云私有部署开源大模型
• 部署一套开源 LLM 项目
• 内容安全
• 互联网信息服务算法备案
• …

学习是一个过程，只要学习就会有挑战。天道酬勤，你越努力，就会成为越优秀的自己。

如果你能在15天内完成所有的任务，那你堪称天才。然而，如果你能完成 60-70% 的内容，你就已经开始具备成为一名大模型 AI 的正确特征了。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】