阿里云PAI为您提供了部分典型场景下的示例模型,便于您便捷地接入TorchAcc进行训练加速。本文为您介绍如何在ResNet-50分布式训练中接入TorchAcc并实现训练加速。
测试环境配置
测试环境配置方法,请参见配置测试环境。
本案例以DSW环境V100M16卡型为例,例如:节点规格选择ecs.gn6v-c8g1.16xlarge-64c256gNVIDIA V100 * 8
。
接入TorchAcc加速ResNet-50分布式训练
以DSW环境为例:
进入DSW实例页面下载并解压测试代码及脚本文件。
在交互式建模(DSW)页面,单击DSW实例操作列下的打开。
在Notebook页签的Launcher页面,单击快速开始区域Notebook下的Python3。
执行以下命令下载并解压测试代码及脚本文件。
!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
文件。后续,您可以直接在该文件中运行下述步骤中的命令,当成功运行结束一个步骤命令后,再顺次运行下个步骤的命令。
执行以下命令下载测试数据集(默认使用类似imagenet-1k的mock数据集)并安装ResNet-50模型依赖的第三方包。
!bash prepare.sh
分别使用普通训练方法(baseline)和接入TorchAcc进行ResNet-50模型分布式训练,来验证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 --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训练方法的优化配置如下:
baseline:Torch112+DDP+AMPO1
PAI-Opt:Torch112+TorchAcc+AMPO1
执行以下命令,获取性能结果数据。
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进行ResNet-50分布式训练可以明显提升性能。关于接入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
、
- 本页导读 (1)