PyTorch - Amazon SageMaker AI
配置 Hugging Face 轉換器的 PyTorch 模型

本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。

PyTorch

攜帶您自己的 PyTorch 模型至 SageMaker AI，並使用 SageMaker Training Compiler 執行訓練任務。

配置 Hugging Face 轉換器的 PyTorch 模型

配置 Hugging Face 轉換器的 PyTorch 模型是基於 PyTorch 的 torch.nn.Module API。Hugging Face 轉換器也為 PyTorch 提供訓練器和預先訓練的模型類別，以妥善減少設定自然語言處理 (NLP) 模型的工作量。準備訓練指令碼之後，您可以使用 SageMaker AI PyTorch 或 HuggingFace 估算器與 SageMaker Training Compiler 組態啟動訓練任務，藉此如 啟用 SageMaker Training Compiler 所述繼續進行下一個主題的操作。

提示

當您在訓練指令碼中使用轉換器為 NLP 模型建立字符化工具時，請確保您透過指定 padding='max_length' 來使用靜態輸入張量形狀。請勿使用 padding='longest'，因為填補至批次中最長的序列可能會變更每個訓練批次的張量形狀。動態輸入形狀可觸發模型的重新編譯，並可能會增加總訓練時間。有關轉換器權杖化工具填補選項的詳情，請參閱 Hugging Face 轉換器文件中的填補和截斷。

使用 Hugging Face 轉換器 Trainer 類別的大型語言模型

如果您使用轉換器程式庫的訓練器類別，則不需對訓練指令碼進行任何其他變更。如果您透過估算器類別啟用訓練器模型，SageMaker Training Compiler 會自動編譯訓練器模型。下列程式碼顯示一個 PyTorch 訓練指令碼與 Hugging Face 訓練器 API 的基本形式。

from transformers import Trainer, TrainingArguments

training_args=TrainingArguments(**kwargs)
trainer=Trainer(args=training_args, **kwargs)

適用於單一 GPU 訓練

當您使用 transformers.Trainer 類別時，不需變更程式碼。

適用於分散式訓練

PyTorch v1.11.0 及較新版本

若要使用 SageMaker Training Compiler 執行分散式訓練，您必須在訓練指令碼中新增下列 _mp_fn() 函式並包裝 main() 函式。它會將 _mp_fn(index) 函式呼叫從 PyTorch (pytorchxla) 的 SageMaker AI 分散式執行期重新導向至訓練指令碼的 main() 函式。

def _mp_fn(index):
    main()

此函式會接受 index 引數，以指出叢集中目前 GPU 的等級，藉此進行分散式訓練。若要尋找更多範例指令碼，請參閱 Hugging Face 轉換器語言模型範例指令碼。

適用於轉換器 v4.17 及之前的版本，並搭配 PyTorch v1.10.2 及之前的版本

SageMaker Training Compiler 使用替代機制來啟動分散式訓練任務，您不必在訓練指令碼中進行任何修改。相對而言，SageMaker Training Compiler 會要求您將 SageMaker AI 分散式訓練啟動器指令碼傳遞至 entry_point 引數，並將訓練指令碼傳遞至 SageMaker AI Hugging Face 估算器中的 hyperparameters 引數。

將 SageMaker Training Compiler 搭配使用 Trainer 的最佳實務

  • 設定 transformers.TrainingArgument 時，務必將 optim 引數設定為 adamw_torch_xla，以確保您使用的是 SyncFree 最佳化工具。另請參閱 Hugging Face 轉換器文件中的最佳化工具。

  • 確保資料處理管道的輸送量高於訓練輸送量。您可以調整 transformers.TrainingArgument 類別中的 dataloader_num_workers 和 preprocessing_num_workers 引數，來實現此一操作。通常，這些數量需要大於或等於 GPU 數量，但小於 CPU 數量。

完成訓練指令碼的調整後，請繼續前往使用 SageMaker Training Compiler 執行 PyTorch 訓練任務。

直接使用 PyTorch 的大型語言模型 (無需 Hugging Face 轉換器訓練器 API)

如有直接使用 PyTorch 的訓練指令碼，則需對 PyTorch 訓練指令碼進行其他變更，以實作 PyTorch/XLA。按照指示修改指令碼，以妥善設定 PyTorch/XLA 本原值。

適用於單一 GPU 訓練

  1. 匯入最佳化程式庫。

    import torch_xla
import torch_xla.core.xla_model as xm

  2. 將目標裝置變更為 XLA 而非 torch.device("cuda")

    device=xm.xla_device()

  3. 如果您使用 PyTorch 的 自動混合精確度 (AMP)，請執行下列動作：

    1. 將 torch.cuda.amp 換成下列項目：

      import torch_xla.amp

    2. 將 torch.optim.SGD 和 torch.optim.Adam 取代為下列項目：

      import torch_xla.amp.syncfree.Adam as adam
import torch_xla.amp.syncfree.SGD as SGD

    3. 將 torch.cuda.amp.GradScaler 換成下列項目：

      import torch_xla.amp.GradScaler as grad_scaler

  4. 如果沒有使用 AMP，則請將 optimizer.step() 取代為下列項目：

    xm.optimizer_step(optimizer)

  5. 如果您使用的是分散式資料載入器，請將資料載入器包裝在 PyTorch/XLA 的 ParallelLoader 類別中：

    import torch_xla.distributed.parallel_loader as pl
parallel_loader=pl.ParallelLoader(dataloader, [device]).per_device_loader(device)

  6. 若不使用 parallel_loader 時，請在訓練迴路的結尾新增 mark_step：

    xm.mark_step()

  7. 若要檢查您的訓練，請使用 PyTorch/XLA 的模型檢查點方法：

    xm.save(model.state_dict(), path_to_save)

完成訓練指令碼的調整後，請繼續前往使用 SageMaker Training Compiler 執行 PyTorch 訓練任務。

適用於分散式訓練

除了上個 適用於單一 GPU 訓練 章節所列的變更內容之外，請新增下列變更項目，以便在 GPU 之間妥善分配工作負載。

  1. 如果您使用了 AMP，請在 scaler.scale(loss).backward() 之後新增 all_reduce：

    gradients=xm._fetch_gradients(optimizer)
xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())

  2. 如需為 local_ranks 和 world_size 設定變數，請使用與下列類似的程式碼：

    local_rank=xm.get_local_ordinal()
world_size=xm.xrt_world_size()

  3. 對於任何大於 1 的 world_size (num_gpus_per_node*num_nodes)，您必須定義一個訓練取樣器，該取樣器看起來應該類似於以下內容：

    import torch_xla.core.xla_model as xm

if xm.xrt_world_size() > 1:
    train_sampler=torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=xm.xrt_world_size(),
        rank=xm.get_ordinal(),
        shuffle=True
    )

train_loader=torch.utils.data.DataLoader(
    train_dataset, 
    batch_size=args.batch_size,
    sampler=train_sampler,
    drop_last=args.drop_last,
    shuffle=False if train_sampler else True,
    num_workers=args.num_workers
)

  4. 進行下列變更，以確保您使用 torch_xla distributed 模組提供的 parallel_loader。

    import torch_xla.distributed.parallel_loader as pl
train_device_loader=pl.MpDeviceLoader(train_loader, device)

    類似一般 PyTorch 載入器的 train_device_loader 函式如下所示：

    for step, (data, target) in enumerate(train_device_loader):
    optimizer.zero_grad()
    output=model(data)
    loss=torch.nn.NLLLoss(output, target)
    loss.backward()

    透過所有這些變更，您應該能夠在不使用轉換器訓練器 API 的情況下，使用任何 PyTorch 模型啟動分散式訓練。請注意，這些指示可用於單一節點多重 GPU 和多節點多重 GPU。

  5. 適用於 PyTorch v1.11.0 及較新版本

    若要使用 SageMaker Training Compiler 執行分散式訓練，您必須在訓練指令碼中新增下列 _mp_fn() 函式並包裝 main() 函式。它會將 _mp_fn(index) 函式呼叫從 PyTorch (pytorchxla) 的 SageMaker AI 分散式執行期重新導向至訓練指令碼的 main() 函式。

    def _mp_fn(index):
    main()

    此函式會接受 index 引數，以指出叢集中目前 GPU 的等級，藉此進行分散式訓練。若要尋找更多範例指令碼，請參閱 Hugging Face 轉換器語言模型範例指令碼。

    適用於轉換器 v4.17 及之前的版本，並搭配 PyTorch v1.10.2 及之前的版本

    SageMaker Training Compiler 使用替代機制來啟動分散式訓練任務，並要求您將 SageMaker AI 分散式訓練啟動器指令碼傳遞至 entry_point 引數，並將訓練指令碼傳遞至 SageMaker AI Hugging Face 估算器中的 hyperparameters 引數。

完成訓練指令碼的調整後，請繼續前往使用 SageMaker Training Compiler 執行 PyTorch 訓練任務。

將 SageMaker Training Compiler 與 PyTorch/XLA 搭配使用的最佳實務

如果您想要在原生 PyTorch 訓練指令碼中運用 SageMaker Training Compiler，您可能須先熟悉 XLA 裝置上的 PyTorch。以下各節列出了一些為 PyTorch 啟用 XLA 的最佳實務。

注意

本節的最佳實務假設您會使用下列 PyTorch/XLA 模組：

import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
了解 PyTorch/XLA 中的寬鬆模式

PyTorch/XLA 和原生 PyTorch 之間的一大顯著差異為，PyTorch/XLA 系統在寬鬆模式中運作，而原生 PyTorch 則以嚴格模式執行。寬鬆模式中的張量會用來建置運算圖的預留位置，直到編譯和評估完成後其具體化為止。當您呼叫 PyTorch API 以使用張量和運算符建置運算內容時，PyTorch/XLA 系統將即時建置運算圖。當 pl.MpDeviceLoader/pl.ParallelLoader 明確或隱含地呼叫 xm.mark_step()，或者當您明確請求張量值 (例如透過呼叫 loss.item() 或 print(loss)) 時，運算圖會進行編譯並執行。

使用 pl.MpDeviceLoader/pl.ParallelLoader 和 xm.step_closure 最小化 compilation-and-executions 的數量

為了獲得最佳效能，請記得啟動 compilation-and-executions 的可能方法，如了解 PyTorch/XLA 中的寬鬆模式中所述，並應盡量最小化 compilation-and-executions 的數量。理想情況下，每個訓練重複程序只需一個 compilation-and-execution，並由 pl.MpDeviceLoader/pl.ParallelLoader 自動初始化。MpDeviceLoader 已針對 XLA 進行最佳化，若可能的話，應時時使用以獲得最佳效能。在訓練期間，您可能需要檢查一些中繼結果，例如損失值。在這種情況下，應使用 xm.add_step_closure() 包裝寬鬆張量的列印，以避免不必要的 compilation-and-executions。

使用 AMP 和 syncfree 最佳化工具

在自動混合精確度 (AMP) 模式中進行訓練，並運用 NVIDIA GPU 的 Tensor 核心，大幅加快您的訓練速度。SageMaker Training Compiler 提供針對 XLA 最佳化的 syncfree 最佳化工具，以提升 AMP 效能。目前，以下三個 syncfree 最佳化工具可以使用，如果可能的話，應使用這些工具以獲得最佳效能。

torch_xla.amp.syncfree.SGD
torch_xla.amp.syncfree.Adam
torch_xla.amp.syncfree.AdamW

這些 syncfree 最佳化工具應與漸層擴展/取消擴展搭配 torch_xla.amp.GradScaler 使用。

提示

從 PyTorch 1.13.1 開始，SageMaker Training Compiler 可讓 PyTorch/XLA 在 torch.optim 或 transformers.optimization 當中自動覆寫最佳化工具 (例如 SGD、Adam、AdamW)，同時在 torch_xla.amp.syncfree (例如 torch_xla.amp.syncfree.SGD、torch_xla.amp.syncfree.Adam、torch_xla.amp.syncfree.AdamW) 使用它們本身的 syncfree 版本。您不需要變更您在訓練指令碼中定義最佳化工具的程式碼行。