接入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文件。

   后續，您可以直接在該文件中運行下述步驟中的命令，當成功運行結束一個步驟命令后，再順次運行下個步驟的命令。

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)

   生成如下圖所示結果。

   實驗結果表明，使用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()