Pytorch分布式训练快速入门教程（一）：从Accelerate说起

爱被打了一巴掌 2023-09-26 23:30 112阅读 0赞

### 意在从TopDown的模式，从应用出发逐步走向技术的底层。如此不至于在长时间的底层理论学习上花费过多时间，而可以快速上手应用，同时又不会让底层理论缺席。 ###

### **背景知识** ###

为后续专题做知识储备，涵盖数学，DL，和优化的基本理论。

### **DP and DDP** ###

> 参考链接： [pytorch中的分布式训练之DP VS DDP - 知乎][pytorch_DP VS DDP -]

常用的DL框架中，支持DataParallel (DP)和DistributedDataParallel (DDP)两种分布训练方式。主要的区别总结如下：

*  **DP采用的是Parameter Sever (PS)架构，而DDP采用的是Ring-All-Reduce模式**
 *  **DP为单进程多线程实现方式，DDP使用多进程方式**
 *  **DP只支持单机训练，DDP支持单机和分布式训练**
 *  **DDP比DP更快**

首先介绍DP的PS架构：

![format_png][]

DP模型的基本思路是将数据/Batch分发到不同机器，在Forward过程将模型复制，数据分batch后分发给不同的GPU，最后由主GPU gather所有的输出；Backward过程由主GPU计算Loss，再把Loss分发给其他GPUs，其他GPUs计算梯度后传递给主GPU更新模型。

PS模式存在一些弊端，最显著的是**负载不均衡**问题，主GPU承担了额外计算和显存开销，同时卡间大量通信导致的延迟问题也不可忽略，随着数据量和模型大小的增加，**通信成本**逐渐不可忽略。

Ring-All-Reduce模式不同于PS，不存在Sever节点，只有Worker节点之间相互通信构成环：

![format_png 1][]

不再划分主GPU，每个GPU拥有独立进程，每个进程从磁盘加载数据。分布式数据采样器可确保加载的数据在各个进程之间不重叠。

此外介绍一部分需要了解的概念：

1.  **Group**

即进程组，默认情况下一个任务只有一个组，即一个world；在某些需要精细通信的情况下，可以通过创建新的group；

1.  **world size**

全局进程个数，如果是多机多卡就表示机器数量，如果是单机多卡就表示 GPU 数量；

1.  **rank**

表示进程序号，用于进程间通信以表示优先级。rank=0的为master节点。如果是多机多卡就表示对应第几台机器，如果是单机多卡，由于一个进程内就只有一个 GPU，所以 rank 也就表示第几块 GPU。

1.  **local\_rank**

表示进程内GPU编号。例如，多机多卡中 rank = 3，local\_rank = 0 表示第 3 个进程内的第 1 块 GPU。

### **显存机制：以Pytorch框架为例** ###

> 参考链接： [PyTorch显存机制分析 - 知乎][PyTorch_ -]

首先`nvidia-smi`是个不准确的显存计算方式，Pytorch会引入缓存区，即使tensor被释放，空间仍可能被占用。准确的说，此命令计算的是`reserved_memory`和`torch_context`的显存之和。

在分析PyTorch的显存时候，最多的是`torch.cuda.memory_allocated()`和`torch.cuda.max_memory_allocated()`，前者可以精准地反馈当前进程中torch.Tensor所占用的GPU显存(注意是只包括torch.Tensor)，后者则可以告诉我们到调用函数为止所达到的最大的显存占用字节数。还有像`torch.cuda.memory_reserved()`这样的函数则是查看当前进程所分配的显存缓冲区是多少的。

Pytorch context或者说CUDA context是执行第一次CUDA操作时所需要的维护设备间工作相关信息的空间。

Pytorch使用层级模式来分配显存。假如我们需要申请4 Bytes的显存，Pytorch底层会先向CUDA申请比如2MB的显存到缓存区，此时可以通过`torch.cuda.memory_cached()`查询到2MB的空间，再由Pytorch向用户分配如512 Bytes的空间，可以通过`torch.cuda.memory_allocated()`查询。

更详细有关显存分配的机制可见：

[https://github.com/pytorch/pytorch/tree/master/c10/cuda][https_github.com_pytorch_pytorch_tree_master_c10_cuda]

### **实战分析** ###

### **accelerate：轻量化模型加速工具** ###

> 官方文档： [https://huggingface.co/docs/accelerate/index][https_huggingface.co_docs_accelerate_index]

首先我们检查Pytorch基本的tranining loop，如下：

device = "cuda" if torch.cuda.is_availabe() else "cpu"
    model.to(device)
    
    for batch in training_dataloader:
        optimizer.zero_grad()
        inputs, targets = batch
        inputs = inputs.to(device)
        targets = targets.to(device)
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        loss.backward()
        optimizer.step()
        scheduler.step()

如果需要使用Accelerate，需要导入最基础的包：

from accelerate import Accelerator
    
    accelerator = Accelerator()

此外，为了满足对device和dataset等的统一管理，需要使用Accelerate进行统一修饰；

device = accelerator.device
    model.to(device)
    
    # The objects must inherit from Pytorch class, like torch.optim.Adam
    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
        model, optimizer, training_dataloader, scheduler
    )

由于已经对dataloader进行了修饰，后续在trainning loop中需要：

#   inputs = inputs.to(device)
    #   targets = targets.to(device)
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
    #   loss.backward()
    +   accelerator.backward(loss)

对于后续的参数配置，我们需要将整个traning loop封装到callable的函数中：

from accelerate import Accelerator
    
    def main():
      accelerator = Accelerator()
    
      model, optimizer, training_dataloader, scheduler = accelerator.prepare(
          model, optimizer, training_dataloader, scheduler
      )
    
      for batch in training_dataloader:
          optimizer.zero_grad()
          inputs, targets = batch
          outputs = model(inputs)
          loss = loss_function(outputs, targets)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          
    if __name__ == '__main__':
    
      main()

此时，可以通过torchrun或accelerate launch进行调用：

accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} {--arg1} {--arg2} ...

可以通过如下命令查询详细的使用方法：

accelerate launch -h

当然，如果存在较多的参数，每次进行重复调用会显得繁琐，因此可以通过`accelerate config`来设计默认参数。当然，更flexible的方法是传入yaml文件，指定每一次的参数：

compute_environment: LOCAL_MACHINE
    deepspeed_config: {}
    distributed_type: MULTI_GPU
    fsdp_config: {}
    machine_rank: 0
    main_process_ip: null
    main_process_port: null
    main_training_function: main
    mixed_precision: fp16
    num_machines: 1
    num_processes: 2
    use_cpu: false

之后，可以通过如下命令启动训练：

accelerate launch --config_file {path/to/config/my_config_file.yaml} {script_name.py} {--arg1} {--arg2} ...

对于模型的存储，Accelerate提供了一些函数。

from accelerate import Accelerator
    import torch
    
    accelerator = Accelerator()
    
    my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
    
    # Save the starting state
    accelerate.save_state("my/save/path")
    
    device = accelerator.device
    my_model.to(device)
    
    # Perform training
    for epoch in range(num_epochs):
        for batch in my_training_dataloader:
            my_optimizer.zero_grad()
            inputs, targets = batch
            inputs = inputs.to(device)
            targets = targets.to(device)
            outputs = my_model(inputs)
            loss = my_loss_function(outputs, targets)
            accelerator.backward(loss)
            my_optimizer.step()
    
    # Restore previous state
    accelerate.load_state("my/save/path")

对于模型中需要tracking的items，可以通过

*  TensorBoard
 *  WandB
 *  CometML
 *  MLFlow

直接启用Trackers或自定义Trackers，使用代码如下：

from accelerate import Accelerator
    
    accelerator = Accelerator(log_with="wandb")
    
    # add hyperparameters and init trackers
    hps = {"num_iterations": 5, "learning_rate": 1e-2}
    accelerator.init_trackers("my_project", config=hps)
    
    # On training loop
    accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)
    
    # After training
    accelerator.end_training()

Accelerate提供了非常轻量化的APIs用于快速做分布式训练，但是也有些随之而来的问题。对于大型的DL任务，需要添加大量的工程化步骤，如hyperparams的管理，系统状态的监控等。这也是相当繁琐的任务，后续会介绍`Pytorch-Lightning`框架，可以帮助我们更加专注于模型本身而不是工程化考虑上。

--------------------

[Pytorch分布式训练快速入门教程（一）：从Accelerate说起 - 知乎][Pytorch_Accelerate_ -]

[pytorch_DP VS DDP -]: https://zhuanlan.zhihu.com/p/356967195
[format_png]: https://img-blog.csdnimg.cn/img_convert/eac6e78fd826a9476a54cdbb5056ceda.webp?x-oss-process=image/format,png
[format_png 1]: https://img-blog.csdnimg.cn/img_convert/41f8be55173dcd57fddc9caaf2f6c0b0.webp?x-oss-process=image/format,png
[PyTorch_ -]: https://zhuanlan.zhihu.com/p/424512257
[https_github.com_pytorch_pytorch_tree_master_c10_cuda]: https://link.zhihu.com/?target=https%3A//github.com/pytorch/pytorch/tree/master/c10/cuda
[https_huggingface.co_docs_accelerate_index]: https://link.zhihu.com/?target=https%3A//huggingface.co/docs/accelerate/index
[Pytorch_Accelerate_ -]: https://zhuanlan.zhihu.com/p/601522424