AI加速：使用TorchAcc实现ResNet-50模型分布式训练加速-人工智能平台 PAI(PAI)-阿里云帮助中心

TorchAcc可加速ResNet-50分布式训练。PAI提供示例代码，帮助您在ResNet-50模型上接入TorchAcc，验证训练加速效果。

测试环境配置

测试环境配置方法，请参见配置测试环境。

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

接入TorchAcc加速ResNet-50分布式训练

以DSW环境为例：

进入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
```
进入ResNet-50目录，双击打开resnet50.ipynb文件。
在该文件中按顺序运行各步骤命令。
1. 执行以下命令下载测试数据集（默认使用类似imagenet-1k的mock数据集）并安装ResNet-50模型依赖的第三方包。
```
!bash prepare.sh
```
2. 分别使用普通训练方法（baseline）和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
```

代码实现原理请参见代码实现原理。

代码实现原理

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（参考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)

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