前提条件

• 已经开通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任务 文档进行。过程中，需要将镜像地址设置为 registry.baidubce.com/ai-public/aiak-training，镜像版本设置为 ubuntu20.04-cu11.7-pytorch1.12.0-py38-lastest。

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 执行之前调用（具体见下述示例）。
2. 当使用梯度裁剪时，clip_grad_norm_(model.parameters(), max_norm) 中的 model.parameters() 需要替换为 torch.cuda.amp.aiak_amp_parameters(optimizer)。

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模式自动更新，用户无需手动调整，如手动调整也不会实际生效。

若运行环境未开启 CUDA，或未使用 AIAK-Training 对应运行时，可能出现 Torch not compiled with CUDA enabled，或 torch.cuda.amp 中不存在 aiak_amp_decorate / aiak_amp_parameters 的情况。#### 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 优化器参数详细如下：