【Hello AI】安装和使用AIACC-AGSpeed（优化PyTorch深度学习模型）

AIACC-AGSpeed（简称AGSpeed）专注于优化PyTorch深度学习模型在阿里云GPU异构计算实例上的计算性能，相比原始的神龙AI加速引擎AIACC，可以实现无感的计算优化性能。本文为您介绍安装和使用AGSpeed的方法。

前提条件

已创建阿里云GPU实例，且GPU实例需满足以下要求：

• 操作系统为Alibaba Cloud Linux、CentOS 7.x或Ubuntu 16.04以上版本。
  
• 已安装NVIDIA Driver和CUDA 10.0或以上版本。
  

支持的版本列表

AGSpeed支持Python、PyTorch以及CUDA版本，版本对应关系及下载地址如下所示。

安装AGSpeed

1. 下载Wheel软件包。
   

根据您环境中的Python、PyTorch以及CUDA版本，在支持的版本列表中选择相应wheel包进行下载。更多信息，请参见支持的版本列表。

2. 执行如下命令，安装AGSpeed。
   

在环境中直接使用pip install命令安装即可。

pip install ${WHEEL_NAME} # 这里的${WHEEL_NAME}替换为您下载的具体wheel软件包名称

使用AGSpeed

建议您尽可能在一切准备工作就绪，即刚好准备执行Train Loop之前，使用agspeed.optimize()对模型进行封装。

例如，模型已经放置到对应的device，并使用DDP优化等操作前，使用agspeed.optimize()对模型进行封装。

1. 执行如下命令，适配代码。
   

import agspeed                  # 导入AGSpeed，用于在PyTorch后端中注册AGSpeed的IR优化Pass和优化后的NvFuser后端。
model = agspeed.optimize(model) # 优化模型，用于模型自动调用计算图抓取的API，并将计算图交由AGSpeed Backend Autotuner优化。

2. 如果您的模型使用的是PyTorch AMP混合精度训练，需要额外在autocast()上下文中增加cache_enabled=False参数，示例代码如下所示。
   

说明：本步骤仅适用于模型使用的是PyTorch AMP混合精度训练场景，使用其他精度（例如FP32）训练场景，请跳过本步骤。

因为TorchDynamo在抓取计算图后，会使用torch.jit.trace进一步将该计算图转换为TorchScript IR，从而调用后端进行优化，在autocast()上下文中直接调用torch.jit.trace会引发冲突，所以需要关闭cache_enabled参数，即您还需要额外在autocast()上下文中增加cache_enabled=False。更多信息，请参见PyTorch commit。 

from torch.cuda.amp.autocast_model import autocast
# ...
# 在autocast上下文参数中添加cache_enabled=False
with autocast(cache_enabled=False):
    loss = model(inputs)
scaler.scale(loss).backward()
scaler.step(optimizer)
# ...

3. 如果您使用的是PyTorch 1.12.x版本，且待训练模型中包含SiLU激活函数，请使用LD_PRELOAD环境变量导入SiLU激活函数符号微分公式。
   

说明：本步骤仅适用于环境为PyTorch 1.12.x版本且待训练模型中包含SiLU激活函数场景，其他场景请跳过本步骤。

PyTorch 1.12.x版本中，TorchScript后端不包含aten::silu的符号微分公式，从而导致aten::silu op不会被纳入可微分子图的范畴，也就无法被后端的NvFuser融合。由于PyTorch的底层实现机制不允许您动态添加符号微分公式，因此AGSpeed将SiLU的符号微分公式集成在另一个动态链接库中（即LD_PRELOAD），将aten::silu的符号微分公式补充到TorchScript后端。在启动训练前，建议您使用LD_PRELOAD环境变量导入SiLU激活函数符号微分公式。

1. 执行以下命令，查看AGSpeed安装路径。
   

python -c "import agspeed; print(agspeed.__path__[0])"

返回页面显示如下，获取AGSpeed安装路径。

2. 执行如下命令，确认上述路径下是否包含libsymbolic_expand.so文件。
   

ls -l ${your_agspeed_install_path} # 将${your_agspeed_install_path}替换为您机器上的AGSpeed安装路径。

返回页面显示如下，表示该路径下已包含libsymbolic_expand.so文件。

3. 执行以下命令，导入LD_PRELOAD环境变量。
   

# 将${your_agspeed_install_path}替换为您机器上的AGSpeed安装路径。
export LD_PRELOAD=${your_agspeed_install_path}/libsymbolic_expand.so
# Start Training...

运行过程中显示如下，表示已将aten::silu的符号微分公式补充到TorchScript后端。

代码示例

在您的训练代码中适配AGSpeed的代码示例如下所示，本示例中，代码前的+号表示本代码行属于新增代码。

+ import agspeed
  # 定义dataloader
  dataloader = ...
  # 定义模型对象
  model = ResNet()
  # 设置模型device
  model.to(device)
  # 定义优化器
  optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
  # 设置DDP
  if distributed:
      model = DDP(model)
+ model = agspeed.optimize(model)
  ############################## Train Loop分为FP32和AMP两种情况演示 ##############################
    ############### FP32 ###############
    # 若是以默认的FP32精度进行训练，不需要修改TrainLoop
  for data, target in dataloader:
      loss = model(data)
      loss.backward()
      optimizer.step()
      optimizer.zero_grad()
    ############### FP32 ###############
    ############### AMP ###############
    # 若是以混合精度进行训练，需要在autocast上下文中增加cache_enabled=False
+ with autocast(cache_enabled=False):
      for data, target in dataloader:
        loss = model(data)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        optimizer.zero_grad()
        scaler.update()
    ############### AMP ###############
  ############################## 通过LD_PRELOAD增加SiLU符号微分公式 ##############################
 # 显示的路径为您机器上的AGSpeed安装路径
  python -c "import agspeed; print(agspeed.__path__[0])"
 # 将${your_agspeed_install_path}替换为您机器上的AGSpeed安装路径
+ export LD_PRELOAD=${your_agspeed_install_path}/libsymbolic_expand.so
 # 执行训练命令
 python train.py

Log示例

Log示例用于帮助您确认是否已成功启用了AGSpeed的加速功能。

• AGSpeed导入成功Log
  

导入AGSpeed时会自动注册AGSpeed的TorchScript IR优化Pass和优化后的NvFuser后端，如果导入操作完成后，出现以下Log表示AGSpeed已成功导入，您可以放心执行下一步操作。

• AGSpeed Autotuning Log
  

AGSpeed会在训练过程中的前几个步骤中进行Autotuning操作，自动选择性能最佳的后端，如果训练过程的Autotuning操作中，出现以下Log表示AGSpeed加速功能已成功启用。

好啦！小弹的分享到此为止。我们更欢迎您分享您对阿里云产品的设想、对功能的建议或者各种吐槽，请扫描提交问卷并获得社区积分或精美礼品一份。https://survey.aliyun.com/apps/zhiliao/P4y44bm_8

【扫码填写上方调研问卷】

欢迎每位来到弹性计算的开发者们来反馈问题哦~