ScratchLLMStepByStep：预训练模型之分布式训练

1.引言

上一节预训练模型之运算加速主要从数据运算层面来优化模型的训练速度，尽管优化效果显著，但始终只能使用单张GPU卡，本节我们将来探讨如何利用多卡来加速模型训练，即分布式训练。

那为什么要使用多张卡来训练呢？原因在于当数据量和运算量很大时，会面临以下问题：

单块显卡内存限制，无法支撑更大的数据量运算。
单块显卡串行计算难以进一步优化，而多块显卡并行运算则能成倍的提升模型训练速度。
单块显卡难以使用更大的batch-size（容易OOM），而大的batch-size比小的batch-size训练过程更稳定，效果也更好。

在pytorch中，支持通过DDP(DistributedDataParallel)数据并行来实现多卡训练。它的基本思想是：将训练数据按照GPU卡的数量划分为多份，并在每个GPU上运行独立的进程来并行训练每份数据，再通过多进程通信来同步梯度和更新参数，从而显著的提高训练速度。

因此，我们需要引入必要的跨进程组件，并对训练过程进行改造，才能支持DDP。

2.基础环节改造

我们会基于上一节构建的训练器Trainer(已经封装到pretrainer_single.py中)进行扩展和改造，以支持并行训练。

首先，引入之前已经封装过的模型结构、数据集和预训练器的实现。

%reload_ext autoreload  
%autoreload 2 
%run transformer.py
%run pretrain_dataset.py
%run pretrainer_single.py

对DDP多卡并行的改造主要包括几方面，分别是：

多进程通信
数据集分片
模型状态同步
验证评估
模型状态保存

下面我们将分别对这几方面进行说明。

2.1 跨进程通信

pytorch中提供了torchrun命令来启动分布式训练，启动命令类似torchrun --nproc_per_node [GPU数量] xxx.py，在这样的命令执行后，torchrun会自动为每个GPU进程设置环境变量，具体包括：

rank：在所有 node 的所有进程中，当前进程全局唯一编号，是从0开始计数的整数；
local rank：当前进程在本地node上的局部编号，也是从0开始计数的整数；
world size：所有并行进程的总数，各个 node 上并行的GPU总数；

有了这些环境变量后，每个进程就可以在启动时知道全局并行的GPU进程总数量和当前的进程编号，这样就能通过一些协议来完成跨进程通信和状态同步。

在分布式训练中，最常用的通信协议是NCCL（NVIDIA Collective Communications Library）, 它提供了以下通信原语来支持跨进程通信。

Broadcast: 将数据从一个 GPU 发送到所有其他 GPU。
All-Reduce: 所有进程的输入数据进行求和、求均值等操作，然后将结果返回给所有进程。
Reduce: 将所有进程的输入数据进行汇总（例如求和），只将结果发送给指定进程。
Send/Receive: 点对点的通信，直接在两个进程间传输数据。

但这些通信原语并不需要我们手动调用，DDP会在每次反向传播后自动执行All-Reduce操作，以保证每个进程的模型参数保持同一状态，我们仅需要确保每个进程都被正确初始化。下面我们为训练器添加一个方法，用于完成分布式环境的初始化。

import os
import torch.distributed as dist
from torch.distributed import init_process_group, destroy_process_group 

def init_distributed_mode(self):
    rank = int(os.environ.get("RANK", -1))
    if rank == -1: 
        self.is_main_process = True
        self.ddp = False
        return
    
    os.environ['NCCL_DEBUG'] = 'WARN'
    world_size = int(os.environ["WORLD_SIZE"])
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
    
    self.ddp = True
    self.rank = rank
    self.is_main_process = self.rank == 0
    self.local_rank = int(os.environ['LOCAL_RANK'])
    self.device = f'cuda:{local_rank}'
    torch.cuda.set_device(device)

setattr(Trainer, "init_distributed_mode", init_distributed_mode)

这个方法主要为训练器初始化了以下几项：

ddp：是否启用分布式训练；
rank: 进程全局编号，用于通信和数据集切分；
local_rank：进程本地编号，用于定位GPU设备；
device：给当前进程分配的GPU设备,如cuda:0；
is_main_process: 用于定位当前进程是否属于主进程，主进程通常需要做一些特殊的操作；

注：init_process_group用于将当前进程加入到nccl进程组中，加入到进程组中的进程可以互相进入跨进程通信。

2.2 数据集分片

上面有提到，数据并行需要将数据集分割成多份，每份分配给一张GPU去训练。在DDP中，这个数据集切割和分配的功能由一个分布式采样器DistributedSampler来完成。具体来讲，我们需要做三点改造：

扩展方法以支持设置原始的训练集和测试集，而不是已经构建好的小批量数据加载器（dataloader)，这样就能对数据集作切分。

def set_dataset(self, train_set, eval_set):
    self.train_set = train_set
    self.eval_set = eval_set
    print(f'set trainset: {len(train_set)}, evalset: {len(eval_set)}') if self.verbose else None

setattr(Trainer, "set_dataset", set_dataset)

创建数据加载器时引入分布式采样器。具体来讲，是要创建一个分布式采样器实例sampler，并用它来构造训练集的加载器，表示训练数据将使用指定的采样器通过部分采样得到，而不是用整个训练数据集。

注：每个GPU进程上都会使用采样器对训练数据集作切分和采样操作，它内部会自动根据world_size来决定将数据集切成多少份，使用local_rank来决定当前进程应该取哪份数据来训练。另外采样器与shuffle操作是互斥的，使用采样器后需要将shuffle关掉。

from torch.utils.data import DataLoader, DistributedSampler, random_split

def init_dataloader(self):
    assert self.train_set and self.eval_set, f"train_set and eval_set can't be empty."
    train_set, eval_set, batch_size = self.train_set, self.eval_set, self.batch_size
    sampler = DistributedSampler(train_set) if self.ddp else None
    self.train_loader = DataLoader(train_set, 
                                   batch_size=batch_size, 
                                   shuffle=(sampler==None), 
                                   num_workers=0, 
                                   drop_last=True, 
                                   sampler=sampler)
    self.eval_loader = DataLoader(eval_set, 
                                  batch_size=batch_size, 
                                  shuffle=True, 
                                  num_workers=0, 
                                  drop_last=False)
    self.steps_per_epoch = len(self.train_loader)
    self.total_steps = self.num_epochs * self.steps_per_epoch
    print(f'init train_loader steps: {len(self.train_loader)}, eval_loader: {len(self.eval_loader)}') if self.verbose else None

setattr(Trainer, "init_dataloader", init_dataloader)

动态分配：一个数据集往往需要训练多轮，DDP支持在每个轮训练开始前，通过set_epoch方法动态计算出数据集索引，这样可以确保每个epoch开始前都可以随机打乱数据顺序，相当于分布式下的shuffle操作。

def train_epoch(self, epoch):
    assert self.train_loader and self.eval_loader, f"train_loader and eval_loader can't be empty."
    self.train_loader.sampler.set_epoch(epoch)
    for i, (X, Y) in enumerate(self.train_loader): 
        ……

2.3 模型复制

在DDP的分布式训练中，每个GPU进程需要复制一份完整的模型参数和优化器状态。因此需要在各张卡上分别加载一份模型，并使用DDP包装，包装的目的是让模型在训练过程中能在各个进程间同步参数状态。

我们在训练器中添加一个wrap_model_with_ddp方法，用于对模型实例进行DDP包装。

@staticmethod
def wrap_model_with_ddp(model, local_rank):
    # 位置编码用的是复数，而nccl不支持复数形式，此变量并不要求在多进程中保持一致，所以暂时屏蔽对此变量的同步
    model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
    model = DistributedDataParallel(model, device_ids=[local_rank])
    print(f"{cur_time()} packaged model with DDP in cuda:{local_rank}")
    return model

setattr(Trainer, "wrap_model_with_ddp", wrap_model_with_ddp)

2.4 模型评估

在分布式训练中，各张卡上的模型状态始终是保持同步的，所以对模型损失的评估只需要在主进程上进行。同时，为了训练流程的简洁，我们将模型评估的检测逻辑独立成一个新方法check_and_evaluate，它的主要逻辑是：在满足设定的eval_steps的前提下，只有主进程才会进行评估验证，并记录训练指标数据。

def check_and_evaluate(self, lr):
    if (self.step + 1) % self.eval_steps != 0:
        return

    if self.is_main_process:
        train_loss = self.train_loss_acc/self.eval_steps
        eval_loss = self.evaluate()
        print(f"{self.cur_time()} lr={lr:.5f}, train_loss: {train_loss:.4f}, "
            + f"eval_loss: {eval_loss:.4f}, steps: {self.step}/{self.total_steps}"
        )
        self.train_loss_acc = 0
        
    dist.barrier() if self.ddp else None

setattr(Trainer, "check_and_evaluate", check_and_evaluate)

注：上面方法中的dist.barrier是为了插入一个同步屏障，起到多进程间训练状态同步的目的。原因在于evaluate方法是一个耗时操作，添加这句代码可以让其它进程等待主进程执行完模型评估再统一进入下一步训练，避免多个进程中的模型状态不同步，触发nccl同步超时。

除此之外在评估阶段，我们需要直接访问模型的原始状态，而不是DDP封装，访问模型的原始状态可以用model.module。

注：这里其实是一个坑，处理不好会造成训练过程中僵死，原因在于DDP封装过的模型在推理时会自动在多卡之间同步运算状态，但由于其它进程并没有运行模型评估，所以这个状态永远同步不到，训练时的表现就是整个训练进度卡住不动。

def evaluate(self):
    # 注意：这里不能多进程同步，必须用原始Model
    model = self.model.module if isinstance(self.model, DistributedDataParallel) else self.model
    model.eval()
    num_batches = len(self.eval_loader)
    total_loss = 0
    
    for (X, Y) in self.eval_loader:
        with torch.no_grad():
            logits = model(X.to(self.device))
        loss = f.cross_entropy(logits.flatten(0, 1), Y.to(self.device).flatten())
        total_loss += loss.item()
    
    model.train()
    return total_loss/num_batches

在分布式训练中，各个进程间的模型参数状态和优化器状态，需要始终保持同步。正常情况下DDP会自动处理各个进程间的训练状态同步，但是像上面这种只需要在主进程上运行的evaluate操作，我们就必须手动加一个barrier操作来处理同步，否则主进程的训练进度就会逐渐落后于其它进程，并且训练的时间越长，模型评估执行的次数越多，落后的进度就越大，最终会超过nccl的多进程同步超时时间而触发NCCL timeout。

2.5 模型保存与恢复

除上面的evaluate操作外，模型的状态保存也只需要在主进程上执行。因此，我们将模型保存的检测逻辑也独立成一个新方法check_and_save_checkpoint，它的主要逻辑是：在满足设定的save_steps的前提下，只有主进程才会保存训练的checkpoint，其它进程需要在dist.barrier()这行代码处等待主进程完成保存操作。

def check_and_save_checkpoint(self, cur_epoch):
    if self.step % self.save_steps != 0:
        return
    
    if self.is_main_process:
        checkpoint_path = f"{self.output_dir}/checkpoint-{self.step}.pth"
        self.save_model(checkpoint_path, cur_epoch)
        print(f"{self.cur_time()} device:{self.device}-save checkpoint: {checkpoint_path}")
        
    # 设置屏障, 让所有进程等待主进程的checkpoint操作
    dist.barrier() if self.ddp else None  

setattr(Trainer, "check_and_save_checkpoint", check_and_save_checkpoint)

3. 训练流程改造

上面增加的环节需要在训练流程中支持才能生效，我们需要分别对单轮训练函数和主训练函数进行改造。

3.1 单轮训练改造

对于单轮训练函数，调整和加入了以下步骤：

分布式模式下，在单轮训练开始前用set_epoch函数打乱顺序；
每次单步训练前，调用adjust_lr动态调整学习率；
评估验证改用封装后的check_and_evaluate方法，兼容单卡训练和多卡训练；
增加保存checkpoint环节的函数调用：check_and_save_checkpoint，也兼容单卡和多卡训练；

def train_epoch(self, cur_epoch):
    assert self.train_loader and self.eval_loader, f"train_loader and eval_loader can't be empty."
    # 每个epoch开始时都重新打乱数据
    self.train_loader.sampler.set_epoch(cur_epoch) if self.ddp else None  
    
    for i, (X, Y) in enumerate(self.train_loader):     
        lr = self.adjust_lr()
        train_loss = self.train_step(X.to(self.device), Y.to(self.device))
        self.train_loss_acc += train_loss.item()
        self.step += 1
        self.check_and_evaluate(lr)
        self.check_and_save_checkpoint(cur_epoch)

setattr(Trainer, "train_epoch", train_epoch)

3.2 主训练改造

对于主训练函数，重构流程如下：

添加对分布式环境的初始化；
添加对恢复训练的支持，能从指定的checkpoint恢复训练；
添加对DDP模型状态同步的支持；
主循环基本不变，只添加了从上次epoch继续训练的支持；
最后训练完后，清理和注销进程资源；

def train(self):
    # 初始化分布式环境
    self.init_distributed_mode()
    # 初始化数据加载器
    self.init_dataloader()
    # 将模型移到指定设备
    self.model.to(self.device)
    # 恢复训练状态
    last_epoch = 0
    if self.last_checkpoint_path:
        last_epoch = self.load_from_checkpoint()
    # 分布式训练需要使用ddp同步模型状态
    if self.ddp:
        self.model = self.wrap_model_with_ddp(self.model, self.local_rank)
    # 打印模型所在的设备  
    model_device = next(self.model.parameters()).device  
    print("Model is on device: ", model_device)  
    # 训练主循环
    for epoch in range(last_epoch, self.num_epochs):
        self.train_epoch(epoch)
    # 注销分布式进程
    dist.destroy_process_group() if self.ddp else None

setattr(Trainer, "train", train)

3.3 main函数改造

编写main函数，用于启动模型训练，主要职责如下：

设置训练参数
配置模型参数的输入/输出路径，用于恢复和保存模型；
创建模型实例，并配置优化器；
加载数据集，并按照比例切分为训练集和测试集；
创建和初始化训练器，并开始训练；

def main():
    epochs = 1
    learning_rate = 1e-3
    batch_size = 16
    train_ratio = 0.98
    weight_decay = 0.01
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    
    last_checkpoint_path = ""
    dataset_path = "/data2/minigpt/dataset/pretrain/mobvoi_seq_monkey_general_open_corpus.bin"
    output_dir = "/data2/minigpt/models/20241210"

    # 模型分布式, autocast会自动将float32绽放为float16（autocast不支持bfloat16），这里不用指定数据类型
    config = GPTConfig(flash_attn=True)
    model = MiniGPT(config)
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

    # 数据加载要设置分布式采样器
    ds = PretrainBinaryDataset(dataset_path, config.context_length)
    train_set, eval_set = split_dataset(ds[:100000], train_ratio)

    train_args = {
        "train_batch_size": batch_size,
        "eval_steps": 500,
        "warmup_steps": 500,
        "save_steps": 12000,
        "num_train_epochs": epochs,
        "output_dir": output_dir,
        "last_checkpoint_path": last_checkpoint_path,
        "use_mixed_precision": True,
    }

    trainer = Trainer(model, optimizer, train_args, device=device, verbose=True)
    trainer.set_seed(123)
    trainer.set_dataset(train_set, eval_set)
    trainer.set_grad_scaler(True)
    trainer.train()

至此，分布式训练的代码部分就基本编写完成，我们将上面编写的整个训练器+main函数代码保存到一个名为pretrainer.py脚本中。

4. 分布式训练

分布式训练需要使用多张卡，我们可以先用nvidia-smi命令看一下机器上有哪些卡可以使用。

!nvidia-smi

    Wed Jan  1 08:04:42 2025       
    +---------------------------------------------------------------------------------------+
    | NVIDIA-SMI 530.30.02              Driver Version: 530.30.02    CUDA Version: 12.1     |
    |-----------------------------------------+----------------------+----------------------+
    | GPU  Name                  Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
    | Fan  Temp  Perf            Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
    |                                         |                      |               MIG M. |
    |=========================================+======================+======================|
    |   0  NVIDIA GeForce RTX 3090         Off| 00000000:3E:00.0 Off |                  N/A |
    | 30%   22C    P8               26W / 350W|      2MiB / 24576MiB |      0%      Default |
    |                                         |                      |                  N/A |
    +-----------------------------------------+----------------------+----------------------+
    |   1  NVIDIA GeForce RTX 3090         Off| 00000000:40:00.0 Off |                  N/A |
    | 30%   22C    P8               18W / 350W|      2MiB / 24576MiB |      0%      Default |
    |                                         |                      |                  N/A |
    +-----------------------------------------+----------------------+----------------------+
    |   2  NVIDIA GeForce RTX 3090         Off| 00000000:41:00.0 Off |                  N/A |
    | 30%   23C    P8               22W / 350W|      2MiB / 24576MiB |      0%      Default |
    |                                         |                      |                  N/A |
    +-----------------------------------------+----------------------+----------------------+
    |   3  NVIDIA GeForce RTX 3090         Off| 00000000:B1:00.0 Off |                  N/A |
    | 30%   20C    P8               20W / 350W|  20498MiB / 24576MiB |      0%      Default |
    |                                         |                      |                  N/A |
    +-----------------------------------------+----------------------+----------------------+
    |   4  NVIDIA GeForce RTX 3090         Off| 00000000:B2:00.0 Off |                  N/A |
    | 30%   19C    P8                6W / 350W|   5554MiB / 24576MiB |      0%      Default |
    |                                         |                      |                  N/A |
    +-----------------------------------------+----------------------+----------------------+
    |   5  NVIDIA GeForce RTX 3090         Off| 00000000:B5:00.0 Off |                  N/A |
    | 30%   23C    P8               14W / 350W|      2MiB / 24576MiB |      0%      Default |
    |                                         |                      |                  N/A |
    +-----------------------------------------+----------------------+----------------------+

我们这里通过环境变量CUDA_VISIBLE_DEVICES来指定要使用的GPU卡号，使用torchrun运行pretrainer.py脚本启动训练。

注意：下面列出的命令只是示意，CUDA_VISIBLE_DEVICES=0,1,2在jupyter notebook中是不起作用的，我们需要从终端运行下面的命令,才能真正启动多个进程。

CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node 3 pretrainer.py

set trainset: 98000, evalset: 2000
init train_loader steps: 2041, eval_loader: 125
2024-12-28 07:33:12 start epoch:0 from step:0
2024-12-28 07:35:16 lr=0.00099, train_loss: 7.3078, eval_loss: 6.4596, grad_norm=0.37541, steps: 199/2041
2024-12-28 07:37:20 lr=0.00097, train_loss: 6.2862, eval_loss: 6.1645, grad_norm=0.24622, steps: 399/2041
2024-12-28 07:39:23 lr=0.00090, train_loss: 6.1241, eval_loss: 6.0611, grad_norm=0.23912, steps: 599/2041
2024-12-28 07:41:27 lr=0.00079, train_loss: 6.0449, eval_loss: 6.0008, grad_norm=0.21517, steps: 799/2041
2024-12-28 07:43:30 lr=0.00064, train_loss: 5.9862, eval_loss: 5.9592, grad_norm=0.20768, steps: 999/2041
2024-12-28 07:45:37 lr=0.00049, train_loss: 5.9465, eval_loss: 5.9235, grad_norm=0.20683, steps: 1199/2041
2024-12-28 07:47:40 lr=0.00034, train_loss: 5.9102, eval_loss: 5.8910, grad_norm=0.19565, steps: 1399/2041
2024-12-28 07:49:44 lr=0.00022, train_loss: 5.8823, eval_loss: 5.8677, grad_norm=0.19772, steps: 1599/2041
2024-12-28 07:51:48 lr=0.00014, train_loss: 5.8657, eval_loss: 5.8471, grad_norm=0.20222, steps: 1799/2041
2024-12-28 07:53:51 lr=0.00010, train_loss: 5.8473, eval_loss: 5.8367, grad_norm=0.19727, steps: 1999/2041
2024-12-28 07:53:54 device:cuda:0-save checkpoint: /data2/minigpt/models/20241015/checkpoint-2000.pth
2024-12-28 07:53:54 barrier wait over of device:cuda:0 at step: 2000.
2024-12-28 07:53:54 barrier wait over of device:cuda:2 at step: 2000.
2024-12-28 07:53:54 barrier wait over of device:cuda:1 at step: 2000.
clean multi process.
train over, steps: 2041
train use time: 21.10min

从上面的训练日志来看，整个训练只用了21分钟，训练数据则与之前相同依然是10万条，相比上一节单卡优化的训练速度50分钟提高了2.38倍，而相比上上一节预训练从零起步的训练速度128分钟则提高了6.1倍。

小结：本文主要介绍了使用DDP进行分布式训练的技术原理和方法，DDP主要采用模型复制、数据并行再结合多进程通信来实现分布式训练。我们先对模型、数据、评估等基础环节进行了改造，之后又对整个训练流程进行分布式改造，最后实际启动脚本进行了分布式训练演示。

总体来讲，多卡并行对训练速度的提升效果非常明显，原理也容易理解，不过分布式训练具体实施过程则相比于单卡训练要复杂很多，特别是涉及到多进程同步的部分，特别容易出问题，需要仔细理解来弄懂其背后的原理。