安装和使用AIACC-AGSpeed

支持的版本列表

安装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后端。

代码示例

+ 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加速功能已成功启用。