支持的版本列表

安裝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加速功能已成功啟用。