混合精度訓練 - Amazon SageMaker

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

混合精度訓練

SageMaker 模型平行處理 (SMP) 程式庫 v2 透過與 和 Transformer Engine 等 PyTorch FSDP開放原始碼架構整合,支援立即可用的混合精確度訓練。如需進一步了解,請參閱下列主題。

使用 Transformer Engine 在 P5 執行個體FP8上使用 的混合精確度訓練

從 SageMaker 模型平行處理 (SMP) 程式庫 2.2.0 版開始,SMP程式庫與 Transformer Engine 整合,並支援開箱即用的FP8混合精度訓練,保持與 PyTorch FSDP MixedPrecision的相容性。這表示您可以同時使用 PyTorch FSDP進行混合精度訓練,以及使用 Transformer Engine 進行FP8訓練。對於 Transformer Engine FP8訓練功能不支援的模型層,這些層會回復到 PyTorch FSDP混合精確度。

注意

SMP v2 FP8支援下列 Hugging Face Transformer 模型:

  • GPT-NeoX (可在 SMP v2.2.0 及更新版本中使用)

  • Llama 2 (可在 SMP v2.2.0 及更新版本中使用)

  • Mixtral 8x7b 和 Mixtral 8x22b (在 SMP v2.5.0 和更新版本中提供)

注意

此 P5 功能FP8訓練可在下列 SageMaker 和 PyTorch 程式庫組合中提供:

  • SageMaker Python 2SDK.212.0 版及更新版本

  • PyTorch v2.2.0 及更新版本

FP8 (8 位元浮點精確度) 是一種資料類型,已出現為另一個範例,以加速LLM模型的深度學習訓練。隨著 NVIDIA H100 GPUs支援的FP8資料類型的推出,您可以從配備 H100 的 P5 執行個體的效能改進中受益GPUs,同時透過FP8混合式精準訓練加速分散式訓練。

FP8 資料類型會進一步細分為 E4M3 和 E5M2 格式。E4M3 提供更高的精確度、有限的動態範圍,並且非常適合模型訓練中的向前傳遞。E5M2 具有更寬廣的動態範圍,但精確度降低,更適合向後通過,其中精確度較不重要,且更寬的動態範圍會變得有利。因此,我們建議您使用混合FP8策略配方來有效地利用這些特徵。

對於半精確度資料類型 (FP16 和 BF16),全域損失擴展技術,例如靜態損失擴展或動態損失擴展,可處理由於半精確度四捨五入梯度而導致的資訊損失引起的收斂問題。不過, 的動態範圍FP8甚至更窄,而且全域損失擴展技術還不夠。此時,我們需要更精細的每張張張張量擴展技術。延遲擴展是一種策略,它會根據在許多張量中觀察到的最大絕對值來選取擴展因子,從而形成先前的反覆運算。此策略中存在權衡;它使用FP8運算的完整效能優勢,但需要記憶體來保留張量的最大值歷史記錄。若要進一步了解一般延遲擴展策略,請參閱FP8適用於深度學習的紙本格式

實際上,使用 FP8 有助於 P5 執行個體上的所有訓練案例。我們強烈建議FP8盡可能啟用 ,以提高訓練效能。

SMP v2 支援立即可用的 Transformer Engine。因此,在 SageMaker (ml.p5.48xlarge) 的 P5 執行個體上使用 SMP v2 執行FP8訓練時,您唯一需要做的就是在訓練指令碼torch.sagemaker中匯入,並繼續使用原生 Transformer Engine Python 套件。若要進一步了解如何使用 Transformer Engine 進行一般FP8訓練,請參閱 Transformer Engine 文件中的FP8搭配使用 Transformer Engine。 NVIDIA 下列程式碼程式碼片段顯示匯入SMP程式庫和在訓練指令碼FP8中設定的程式碼行應該看起來如何。

import torch.sagemaker as tsm import transformer_engine.pytorch as te from transformer_engine.common.recipe import DelayedScaling, Format # Initialize the SMP torch.sagemaker API. tsm.init() # Define a transformer model and wrap it with the torch.sagemaker.transform API. from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_config(ModelConfig) model = tsm.transform(model) # Enable E4M3 during forward pass, E5M2 during backward pass. fp8_format = Format.HYBRID # Create an FP8 recipe. fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max") # Enable FP8 autocasting. with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe, fp8_group=tsm.state.world_process_group): out = model(inp) loss = out.sum() loss.backward()

若要尋找在 P5 執行個體上使用 SMP v2 進行FP8訓練的實際範例,請參閱在 P5 執行個體FP8上使用 進行 Llama-v2 (或 GPT-NeoX) 加速 SageMaker PyTorch FSDP訓練的範例筆記本

使用 搭配半精確度資料類型的混合精確度訓練 PyTorch FSDP

SMP v2 PyTorch FSDP MixedPrecision 支援在 P4 和 P5 執行個體上訓練任務。 PyTorch FSDP 提供各種組態,可混合精確度,以提高效能和減少記憶體。

注意

具有 PyTorch FSDP功能的混合精確度訓練可在下列 SageMaker 和 PyTorch 程式庫組合中使用。

  • SMP v2.0.0 及更新版本

  • SageMaker Python 2SDK.200.0 版及更新版本

  • PyTorch v2.0.1 及更新版本

設定混合精度模型的標準方法是在 中建立模型float32,然後允許透過傳遞MixedPrecision政策將參數FSDP轉換為 bfloat16 float16或 快速轉換,如下列程式碼片段所示。如需變更 dtype 中參數、降低或緩衝區混合精確度選項的詳細資訊 PyTorch,請參閱 文件PyTorch FSDPMixedPrecisionAPI中的 。 PyTorch

# Native PyTorch API from torch.distributed.fsdp import MixedPrecision dtype = torch.bfloat16 mixed_precision_policy = MixedPrecision( param_dtype=dtype, reduce_dtype=dtype, buffer_dtype=dtype ) model = FSDP( model, ..., mixed_precision=mixed_precision_policy )

請注意,某些模型 (例如 Hugging Face Transformer Llama 模型) 預期緩衝區為 float32。若要使用 float32,請在定義dtype物件的行torch.float32torch.bfloat16將 取代為 。