小白入门的DeepSpeed原理和代码实现（二）分布式训练：数据并行、模型并行、流水线并行-优快云博客

本文链接：https://blog.youkuaiyun.com/Code1994/article/details/145185538

二、分布式训练

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)