超大模型分布式训练DeepSpeed教程

DeepSpeed教程

• 项目链接

简介

• deep speed是微软的新大规模模型分布式训练的工具。专门为训练超大模型而生。号称可以训练10B参数的模型。比目前最好的模型大10倍，训练速度块10倍。兼容pytorch的模型，可以改动最少代码。下图是展示训练bert需要的时间，基本同gpu的数量成线性相关。
  

安装

下载code(0.3.0)

git clone https://github.com/microsoft/DeepSpeed.git

安装python环境

• 需要注意pytroch cuda的版本，需要彼此对上

pip install torch==1.5.1 torchvision==0.6.1 

pip install cupy_cuda102==7.8.0 

pip install virtualenv==20.0.31 

• 查看cuda版本

import torch
print(torch.version.cuda)

安装deep speed依赖

cd DeepSpeed
pip install -r requirements/requirements-dev.txt 
pip install -r requirements/requirements  
pip install -r requirements/requirements-sparse-attn.txt 
pip install mpi4py
pip install --ignore-installed PyYAML 

分布式ssh

• 多机之间需要通过ssh免密互相登录

git config --global user.name "xxx"
git config --global user.email "[email protected]"
ssh-keygen -t rsa -C "[email protected]"
把生成的公钥拷贝到其他机器.ssh目录下即可

安装DeepSpeed

cd DeepSpeed
./install.sh 
等一会即可完成安装

测试demo

单机测试

• demo是一个简单的分类测试，是单机的

cd DeepSpeed/DeepSpeedExamples/pipeline_parallelism
./run.sh 

多机训练测试demo

• 增加hostfile文件，填写host的相应的gpu数量(slots=4代表有4个gpu)

host1 slots=4
host2 slots=4

• include参数，指定机器和gpu,如下代表使用host1机器的3号和host2的2、3号gpu

--include="host1:3@host2:2,3"

• exclude参数，同include参数，代表不使用相应的gpu
• ds_config.json 文件里面配置训练的参数,如batch_size、优化器参数、log参数度呢

 {
  "train_batch_size" : 256,
  "train_micro_batch_size_per_gpu" : 8,

   "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.001,
      "betas": [
        0.9,
        0.999
      ],
      "eps": 1e-8
    }
  },
  
  "steps_per_print" : 10,
  "wall_clock_breakdown" : false
 }

• 完整run.sh 命令如下,运行即可实验多机、多gpu训练的demo啦。

#!/bin/bash

deepspeed --hostfile=hostfile  --include="host1:3@host2:2,3"  train.py -p 2 --steps=200  --deepspeed_config=ds_config.json