deepspeed分布式训练demo

deepspeed demo

GPT2训练

我们需要修改三个文件：argument.py pretrain_gpt2.py 以及utils.py

这三个分别对应着参数，训练，以及保存加载等过程。

1、对于argument.py我们需要将超参数丢给deepspeed：

    parser = add_data_args(parser)
    # Include DeepSpeed configuration arguments
    parser = deepspeed.add_config_arguments(parser)

2、修改pretrain.py，主要就是，model 、优化器丢到deepspeed

def setup_model_and_optimizer(args):
    """Setup model and optimizer."""
​
    model = get_model(args)
    optimizer = get_optimizer(model, args)
    lr_scheduler = get_learning_rate_scheduler(optimizer, args)
    if args.deepspeed:
        import deepspeed
​
        print_rank_0("DeepSpeed is enabled.")
​
        model, optimizer, _, lr_scheduler = deepspeed.initialize(
            model=model,
            optimizer=optimizer,
            args=args,
            lr_scheduler=lr_scheduler,
            mpu=mpu,
            dist_init_required=False
       )

在设置优化器的时候，我们直接使用Adam系列就行

    optimizer = Adam(param_groups,
                     lr=args.lr, weight_decay=args.weight_decay)

3、设置成deepspeed的反向传播

    #如果使用deepspeed则使用model.backward ,不然是loss.backward()
    if args.deepspeed:
            model.backward(loss)
        else:
            optimizer.zero_grad()
            if args.fp16:
                optimizer.backward(loss, update_master_grads=False)
            else:
                loss.backward()
     #更新参数
    if args.deepspeed:
         model.step()
     else:
         optimizer.step()
​
         # Update learning rate.
         if not (args.fp16 and optimizer.overflow):
             lr_scheduler.step()
         else:
             skipped_iter = 1

所以训练模型的时候都可以这三步走。

底层优化：（可以不看）

反向传播的时候，我们要选择是allreduce还是其他的通信方法

        if args.deepspeed:
            # DeepSpeed backward propagation already addressed all reduce communication.
            # Reset the timer to avoid breaking timer logs below.
            #DeepSpeed反向传播已经处理了所有的reduce通信。
            #重置计时器以避免破坏下面的计时器日志
            timers('allreduce').reset()
        else:
            torch.distributed.all_reduce(reduced_losses.data)
            reduced_losses.data = reduced_losses.data / args.world_size
            if not USE_TORCH_DDP:
                timers('allreduce').start()
                model.allreduce_params(reduce_after=False,
                                       fp32_allreduce=args.fp32_allreduce)
                timers('allreduce').stop()

提过更新主节点的梯度

        # Update master gradients.
        if not args.deepspeed:
            if args.fp16:
                optimizer.update_master_grads()
​
            # Clipping gradients helps prevent the exploding gradient.
            if args.clip_grad > 0:
                if not args.fp16:
                    mpu.clip_grad_norm(model.parameters(), args.clip_grad)
                else:
                    optimizer.clip_master_grads(args.clip_grad)
​
        return lm_loss_reduced
更新模型参数
if args.deepspeed:
         model.step()
     else:
         optimizer.step()
​
         # Update learning rate.
         if not (args.fp16 and optimizer.overflow):
             lr_scheduler.step()
         else:
             skipped_iter = 1

cifar10

将pytorch原始的代码

1配置参数

2训练

3保存节点

环境配置

下载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 "xxx@mobvoi.com"
ssh-keygen -t rsa -C "xxx@mobvoi.com"
#这一步是配置ssh免密连接
把生成的公钥拷贝到其他机器.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