接入TorchAcc加速ResNet-50分布式訓練

以DSW環境為例：

1. 進入DSW實例頁面下載并解壓測試代碼及腳本文件。

   1. 在交互式建模（DSW）頁面，單擊DSW實例操作列下的打開。

   2. 在Notebook頁簽的Launcher頁面，單擊快速開始區域Notebook下的Python3。

   3. 執行以下命令下載并解壓測試代碼及腳本文件。

      !wget http://odps-release.cn-hangzhou.oss.aliyun-inc.com/torchacc/accbench/gallery/resnet50.tar.gz && tar -zxvf resnet50.tar.gz

2. 進入ResNet-50目錄，雙擊打開resnet50.ipynb文件。

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

   1. 執行以下命令下載測試數據集（默認使用類似imagenet-1k的mock數據集）并安裝ResNet-50模型依賴的第三方包。

      !bash prepare.sh

   2. 分別使用普通訓練方法（baseline）和接入TorchAcc進行ResNet-50模型分布式訓練，來驗證TorchAcc的性能提升效果。

      普通訓練方法和接入TorchAcc訓練方法的優化配置如下：

      • baseline：Torch112+DDP+AMPO1

      • PAI-Opt：Torch112+TorchAcc+AMPO1

      說明

      • 在測試不同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 --nproc_per_node=1 --amp_level=O1 --batch_size=128
        

      • Pytorch Eager八卡（baseline訓練）

        !#!/bin/bash
        
        !set -ex
        
        !python launch_single_task.py --nproc_per_node=8 --amp_level=O1 --batch_size=128
        

      • TorchAcc單卡（PAI-OPT）

        !#!/bin/bash
        
        !set -ex
        
        !python launch_single_task.py --nproc_per_node=1 --amp_level=O1 --compiler-opt --batch_size=128
        

      • TorchAcc八卡（PAI-OPT）

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

關于接入TorchAcc更詳細的代碼實現原理，請參見代碼實現原理。

代碼實現原理

將上述的ResNet-50模型接入TorchAcc框架進行分布式訓練加速的代碼配置，請參考已下載的代碼文件ResNet-50/main.py。

Import TorchAcc API

在main函數import處添加以下代碼，請參考main.py文件中76-80行代碼：

from logger import create_logger, enable_torchacc_compiler, enable_torchacc_kernel, log_params, log_metrics

if enable_torchacc_compiler():
  import torchacc.torch_xla.core.xla_model as xm
  import torchacc.torch_xla.distributed.xla_backend
  from torchacc.torch_xla.amp import autocast, GradScaler
  from torchacc.torch_xla.amp import syncfree
  dist.get_rank = xm.get_ordinal
  dist.get_world_size = xm.xrt_world_size
else:
  from torch.cuda.amp import autocast, GradScaler

分布式初始化

在調用dist.init_process_group函數時，將backend參數設置為xla：

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

set_replication+封裝dataloader+model placement+optimizer

在模型和dataloader定義完成之后，獲取xla_device并調用set_replication函數，以封裝dataloader并設置模型的設備位置。請參考main.py文件中97-107、124-130、136-140行代碼。

+if enable_torchacc_compiler():
+    dist.init_process_group(backend="xla", init_method="env://")
+    device = xm.xla_device()
+    xm.set_replication(device, [device])
+else:
    args.local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{args.local_rank}")
    dist.init_process_group(backend="nccl", init_method="env://")
    dist.barrier()
args.world_size = dist.get_world_size()
args.rank = dist.get_rank()

+if enable_torchacc_compiler():
+    model.to(device)
+    xm.mark_step()
+else:
    torch.cuda.set_device(device)
    model.cuda(device)
    model = torch.nn.parallel.DistributedDataParallel(model)


+if enable_torchacc_compiler() and args.amp_level != "O0":
+    optimizer_cls = syncfree.AdamW
+else:
    optimizer_cls = torch.optim.AdamW
optimizer = optimizer_cls(model.parameters(), args.lr,
                          weight_decay=args.weight_decay)

梯度allreduce通信

如果啟用了AMP開關，需要在loss backward后對梯度進行allreduce，并在backward和apply計算階段修改代碼。具體請參考main.py文件的240-277行代碼。

    for step, (images, target) in enumerate(train_loader):
      if step > args.max_steps:
        break
+      if not enable_torchacc_compiler():
+          images = images.to(device, non_blocking=True)
+          target = target.to(device, non_blocking=True)

      # compute output
      with autocast_context_manager(args):
        output = model(images)
        loss = criterion(output, target)

      # measure accuracy and record loss
      acc1, acc5 = accuracy(output, target, topk=(1, 5))

      # compute gradient and do optimizer step
      optimizer.zero_grad()
      if args.amp_level != "O0":
        scaler.scale(loss).backward()
+        if enable_torchacc_compiler():
+          gradients = xm._fetch_gradients(optimizer)
+          xm.all_reduce('sum', gradients, scale=1.0/dist.get_world_size())
        scaler.step(optimizer)
        scaler.update()
      else:
        loss.backward()
+        if enable_torchacc_compiler():
+          gradients = xm._fetch_gradients(optimizer)
+          xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())
        optimizer.step()

      if args.rank == 0 and step % args.log_interval == 0:
        # measure elapsed time
        batch_time = (time.time() - end) / args.log_interval
        end = time.time()
        samples_per_step = float(args.batch_size / batch_time) * args.world_size
        peak_mem = torch.cuda.memory_stats()['allocated_bytes.all.peak']/1024.0/1024.0/1024.0
        log_metrics(epoch, step, args.batch_size, loss, batch_time, samples_per_step, peak_mem)

Training Loop封裝

更新代碼邏輯：

• 從dataloader取出樣本（數據）作為后面訓練的輸入，具體請參考main.py文件的243-245行代碼：

  +if not enable_torchacc_compiler():
      images = images.to(device, non_blocking=True)
      target = target.to(device, non_blocking=True)

• 如果啟用了AMP功能，目前TorchAcc只支持使用AMP的autocast功能。因此需要在training loop中添加get_autocast_and_scaler代碼，具體請參考main.py文件的248-250行代碼。

  with autocast_context_manager(args):
      output = model(images)
      loss = criterion(output, target)

  其中autocast_context_manager函數的實現可以參考main.py文件的87-92行代碼。

  def autocast_context_manager(args):
    if args.amp_level != "O0":
      ctx_manager = autocast()
    else:
      ctx_manager = contextlib.nullcontext() if sys.version_info >= (3, 7) else contextlib.suppress()
    return ctx_manager