本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。
將您自己的 PyTorch 模型帶入 SageMaker AI,並使用 SageMaker Training Compiler 執行訓練任務。
配置 Hugging Face 轉換器的 PyTorch 模型
配置 Hugging Face 轉換器HuggingFace
PyTorch
或估算器搭配 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 訓練
-
匯入最佳化程式庫。
import torch_xla import torch_xla.core.xla_model as xm
-
將目標裝置變更為 XLA 而非
torch.device("cuda")
device=xm.xla_device()
-
如果您使用 PyTorch 的 自動混合精確度
(AMP),請執行下列動作: -
將
torch.cuda.amp
換成下列項目:import torch_xla.amp
-
將
torch.optim.SGD
和torch.optim.Adam
取代為下列項目:import torch_xla.amp.syncfree.Adam as adam import torch_xla.amp.syncfree.SGD as SGD
-
將
torch.cuda.amp.GradScaler
換成下列項目:import torch_xla.amp.GradScaler as grad_scaler
-
-
如果沒有使用 AMP,則請將
optimizer.step()
取代為下列項目:xm.optimizer_step(optimizer)
-
如果您使用的是分散式資料載入器,請將資料載入器包裝在 PyTorch/XLA 的
ParallelLoader
類別中:import torch_xla.distributed.parallel_loader as pl parallel_loader=pl.ParallelLoader(dataloader, [device]).per_device_loader(device)
-
若不使用
parallel_loader
時,請在訓練迴路的結尾新增mark_step
:xm.mark_step()
-
若要檢查您的訓練,請使用 PyTorch/XLA 的模型檢查點方法:
xm.save(model.state_dict(), path_to_save)
完成訓練指令碼的調整後,請繼續前往使用 SageMaker Training Compiler 執行 PyTorch 訓練任務。
適用於分散式訓練
除了上個 適用於單一 GPU 訓練 章節所列的變更內容之外,請新增下列變更項目,以便在 GPU 之間妥善分配工作負載。
-
如果您使用了 AMP,請在
scaler.scale(loss).backward()
之後新增all_reduce
:gradients=xm._fetch_gradients(optimizer) xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())
-
如需為
local_ranks
和world_size
設定變數,請使用與下列類似的程式碼:local_rank=xm.get_local_ordinal() world_size=xm.xrt_world_size()
-
對於任何大於
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 )
-
進行下列變更,以確保您使用
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。
-
適用於 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 模組:
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 版本。您不需要變更您在訓練指令碼中定義最佳化工具的程式碼行。