使用AIAK-Training Pytorch版

前提条件

• 已经开通CCE云原生AI服务。

支持的版本列表

AIAK-Training Pytorch版目前支持如下版本列表，用户如有其他版本需求，可提交工单获得支持。

加速特性

• 针对低带宽网络环境的通信优化，新增分层Allreduce算法，支持PowerSGD、FP16等通信压缩算法快速使能。
• 引入nvidia apex amp o2混合精度模式，并提供兼容torch amp原生用法的使用方式，方便更多计算转FP16来加速训练。
• 支持fused optimizer，融合计算Kernel，减少访存、kernel launch等开销，提升参数更新阶段效率。
• 支持lamb优化器算法，可用于解决超大batch训练中的收敛问题。

实践流程

1.获取训练镜像

在CCR公共镜像的“百度智能云AI镜像”中选择“aiak-training”加速镜像，作为训练基础镜像，该镜像内已安装了CUDA、Python、PyTorch 和AIAK-Training加速软件等。

2.任务提交

提交任务，可详细按照新建Pytorch任务文档进行。过程中，需要将镜像地址设置为上述aiak-training的镜像地址。

3.应用加速能力

3.1 加速低配网络条件下的分布式训练

在网络带宽较低的环境下，跨机梯度同步会成为分布式训练速度的主要瓶颈。为了优化该类问题，AIAK-Training新增了分层AllReduce算法实现，来减少跨机通信的影响。同时，我们简化了Pytorch官方提供的PowerSGD、FP16等通信压缩hook的使用方式，用户无需修改代码，仅通过环境变量即可直接开启。注意：在单机多卡或RDMA环境等通信并非瓶颈情况下，上述功能可能不会带来加速收益。

3.2 AMP O2混合精度模式，进一步加速计算

Pytorch官方提供了AMP自动混合精度功能，用户可以使用该技术加速模型的计算效率。AIAK-Training在该功能基础之上，引入了NVIDIA APEX AMP中的O2模式，并在接口层面与Torch AMP做了兼容，用户仅需增加一行代码，即可快速使能更为激进的纯FP16训练方案。具体用法如下：

1. 用户需要在模型和optimizer准备就绪后，增加一行aiak_amp_decorate函数调用的代码，另外当使用DDP时，aiak_amp_decorate需在DDP执行之前调用。（具体见下述示例第6行）
2. 当使用梯度裁剪时，clip_grad_norm(model.parameters(), max_norm)中 model.parameters() 需要替换为 torch.cuda.amp.aiak_amp_parameters(optimizer)。（具体见下述示例第 25行）

1# 定义 model 和优化器
2model = Model().cuda()
3optimizer = optim.SGD(model.parameters(), ...)
4
5# 新增 aiak_amp_decorate 封装调用，用于使能 Apex AMP O2模式
6model, optimizer = torch.cuda.amp.aiak_amp_decorate(model, optimizer) 
7
8# 在O2优化初始化之后构造DDP model
9model = DDP(model)
10
11# torch 原生 amp 用法的相关代码，无需修改
12scaler = GradScaler()
13
14for input, target in data:
15    optimizer.zero_grad()
16    # 在autocast enable 区域运行forward
17    with autocast():
18        output = model(input)
19        loss = loss_fn(output, target)
20    # 用scaler，scale loss(FP16)，backward得到scaled的梯度(FP16)
21    scaler.scale(loss).backward()
22    # 当使用梯度裁剪时，遵守torch原生AMP调用规则， 调用unscale进行梯度恢复
23    scaler.unscale_(optimizer)
24    # 梯度裁剪时，需要替换model.parameters()为amp.amp_parameters(optimizer)
25    torch.nn.utils.clip_grad_norm_(torch.cuda.amp.aiak_amp_parameters(optimizer), max_norm)
26    # scaler 更新参数，会先自动unscale梯度, 如果有nan或inf，自动跳过
27    scaler.step(optimizer)
28    # scaler factor更新
29    scaler.update()

aiak_amp_decorate 接口参数详细如下：

其他注意点：在使用O2模式时，torch原生GradScaler类的backoff_factor、growth_factor、growth_interval等参数，会按照o2模式自动更新，用户无需手动调整，如手动调整也不会实际生效。

3.3 FusedOptimizer 加速参数更新效率

用户使能优化器融合功能，只需设置如下环境变量：

3.4 大 BatchSize 训练，使能 LAMB优化器

LAMB 优化器，和其他优化器使用方式类似，具体举例如下：

1import torch.optim as optim
2...
3optimizer = optim.LAMB(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, eps=args.eps, betas=args.betas)

Lamb 优化器参数详细如下：