文档

NLP:TorchAcc提速BERT Base分布式训练

更新时间:

阿里云PAI为您提供了部分典型场景下的示例模型,便于您便捷地接入TorchAcc进行训练加速。本文为您介绍如何在BERT-Base分布式训练中接入TorchAcc并实现训练加速。

测试环境配置

测试环境配置方法,请参见配置测试环境

本案例以DSW环境V100M16卡型为例,例如:节点规格选择ecs.gn6v-c8g1.16xlarge-64c256gNVIDIA V100 * 8

接入TorchAcc加速BERT Base分布式训练

  1. 进入DSW实例页面下载并解压测试代码及脚本文件。

    1. 交互式建模(DSW)页面,单击DSW实例操作列下的打开

    2. Notebook页签的Launcher页面,单击快速开始区域Notebook下的Python3

    3. 执行以下命令下载并解压测试代码及脚本文件。

      !wget http://odps-release.cn-hangzhou.oss.aliyun-inc.com/torchacc/accbench/gallery/bert.tar.gz && tar -zxvf bert.tar.gz
  2. 进入bert目录,双击打开bert.ipynb文件。

    后续,您可以直接在该文件中运行下述步骤中的命令,当成功运行结束一个步骤命令后,再顺次运行下个步骤的命令。image..png

  3. 执行以下命令下载数据集SST并安装BERT-Base模型依赖的第三方包。

    !bash prepare.sh
  4. 分别使用普通训练方法(baseline)和接入TorchAcc进行BERT-Base模型分布式训练,来验证TorchAcc的性能提升效果。

    说明
    • 在测试不同GPU卡型(例如V100、A10等)时,可以通过调整batch_size来适配不同卡型的显存大小。

    • 在测试不同机器实例时,由于单机GPU卡数不同(假设为N),因此可以通过设置nproc_per_node来启动单卡或多卡的任务,其中:1<=nproc_per_node<=N。

    • Pytorch Eager单卡(baseline训练)

      !#!/bin/bash
      
      !set -ex
      
      !python launch_single_task.py --amp_level O1 --batch_size 24 --nproc_per_node 1
    • Pytorch Eager八卡(baseline训练)

      !#!/bin/bash
      
      !set -ex
      
      !python launch_single_task.py --amp_level O1 --batch_size 24 --nproc_per_node 8
    • TorchAcc单卡(PAI-OPT)

      !#!/bin/bash
      
      !set -ex
      
      !python launch_single_task.py --amp_level O1 --batch_size 24 --compiler-opt --nproc_per_node 1
    • TorchAcc八卡(PAI-OPT)

      !#!/bin/bash
      
      !set -ex
      
      !python launch_single_task.py --amp_level O1 --batch_size 24 --compiler-opt --nproc_per_node 8

    其中:普通训练方法和接入TorchAcc训练方法的优化配置如下:

    • baseline:Torch112+DDP+AMPO1

    • PAI-Opt:Torch112+TorchAcc+AMPO1

  5. 执行以下命令,获取性能数据结果。

    import os
    from plot import plot, traverse
    from parser import parse_file
    #import seaborn as sns
    
    if __name__ == '__main__':
        path = "output"
        file_names = {}
        traverse(path, file_names)
    
        for model, tags in file_names.items():
            for tag, suffixes in tags.items():
                title = model + "_" + tag
                label = []
                api_data = []
                for suffix, o_suffixes in suffixes.items():
                    label.append(suffix)
                    for output_suffix, node_ranks in o_suffixes.items():
                        assert "0" in node_ranks
                        assert "log" in node_ranks["0"]
                        parse_data = parse_file(node_ranks["0"]["log"])
                        api_data.append(parse_data)
                
                plot(title, label, api_data)

    生成如下图所示结果。5292d5ac89d7686ca5750cc267966fdb..png

    实验结果表明,使用TorchAcc进行BERT-Base分布式训练可以明显提升性能。接入TorchAcc更详细的代码实现原理,请参见代码实现原理

代码实现原理

将上述的BERT-Base模型接入TorchAcc框架进行分布式训练加速的代码配置,请参考已下载的代码文件bert/bert.py

Import TorchAcc API

TorchAcc在训练时会使用XLA Device,如果使用TorchAcc进行训练,则需要在main函数import处添加以下代码,具体请参考bert.py文件中35-44行代码:

+if enable_torchacc_compiler():
+  import torchacc.torch_xla.core.xla_model as xm
+  import torchacc.torch_xla.distributed.parallel_loader as pl
+  import torchacc.torch_xla.distributed.xla_backend
+  from torchacc.torch_xla.amp import autocast, GradScaler, syncfree
+  dist.init_process_group(backend="xla", init_method="env://")
+else:
  from torch.cuda.amp import autocast, GradScaler
  dist.init_process_group(backend="nccl", init_method="env://")

分布式初始化

在调用dist.init_process_group函数时,将backend参数设置为xla

dist.init_process_group(backend="xla", init_method="env://")

set_replication+封装dataloader+model placement

在模型和dataloader定义完成之后,获取xla_device并调用set_replication函数,以封装dataloader并设置模型的设备位置。

+if args.device == "xla":
+    device = xm.xla_device()
+    xm.set_replication(device, [device])
+    train_device_loader = pl.MpDeviceLoader(train_device_loader, device)
+    model = model.to(device)
+else:
    device = torch.device(f"cuda:{args.local_rank}")
    torch.cuda.set_device(device)
    model = model.cuda()
    model = torch.nn.parallel.DistributedDataParallel(model)

Optimizer封装

若需要启用AMP O1功能,请参考bert.py文件的139-142行进行修改。

if args.device == "xla" and args.amp_level == "O1":
    optimizer = syncfree.Adam(model.parameters(), lr=1e-3)
else:
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

梯度allreduce通信

如果启用了AMP开关,需要在loss backward后对梯度进行allreduce,并在backward和apply计算阶段修改代码。具体请参考bert.py文件的65-67行。

def loop_with_amp(model, inputs, optimizer, autocast, scaler):
  with autocast():
    outputs = model(**inputs)
    loss = outputs["loss"]

  scaler.scale(loss).backward()
+  if args.device == "xla":
+    gradients = xm._fetch_gradients(optimizer)
+    xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())
  scaler.step(optimizer)
  scaler.update()

  return loss, optimizer


def loop_without_amp(model, inputs, optimizer):
  outputs = model(**inputs)
  loss = outputs["loss"]
  loss.backward()
+  if args.device == "xla":
+    xm.optimizer_step(optimizer)
+  else:
    optimizer.step()
  return loss, optimizer

Training Loop封装

更新以下内容:

  • 从dataloader取出样本(数据)作为后面训练的输入,请参考bert.py文件的第93-94行。

    if args.device == "cuda":
        inputs.to(device)
    else:
        # For TorchAcc, wrapper data loader with parallel_loader is enough.
        pass
  • 如果启用了AMP功能,目前TorchAcc只支持使用AMP的autocast功能。因此需要在training loop中添加get_autocast_and_scaler代码。请参考bert.py文件的145-146行。

    if args.amp_level == "O1":
        autocast, scaler = get_autocast_and_scaler()

    其中get_autocast_and_scaler函数的实现可以参考bert.py文件的52-56行。

    # 不考虑AMP O2测试。
    def get_autocast_and_scaler():
      if args.device == "xla":
        return autocast, GradScaler()
    
      return autocast, GradScaler()