透過 Hugging Face 與 Amazon SageMaker 進行 BART 與 T5 的摘要生成分散式訓練

Hugging Face 與 Amazon SageMaker 已完成整合，提供優化的深度學習容器 (DLCs) 以及 SageMaker Python SDK 中專用的 HuggingFace estimator。此合作讓開發者能夠使用分散式訓練策略來訓練最先進的 NLP 模型（例如 BART 與 T5），藉此加速效能並縮短訓練時間。

為 Transformers 提供優化的基礎設施

Amazon SageMaker 現在具備 Hugging Face 優化的深度學習容器，可加速 Transformers 模型之訓練。為了簡化部署流程，SageMaker Python SDK 包含了 HuggingFace estimator，讓使用者能以極少的程式碼啟動訓練任務。

使用 SageMaker Data Parallelism 進行分散式訓練

分散式訓練是透過 SageMaker Data Parallelism 實作的，該功能已整合至 Hugging Face Trainer API 中。藉由在 HuggingFace estimator 中定義 distribution 參數，使用者可以將訓練擴展至多個 GPU 與實例。

配置範例：

distribution = {'smdistributed':{'dataparallel':{ 'enabled': True }}}

在分散式設定中，總批次大小 (total batch size) 是以 per_device_train_batch_size 乘以所有實例所使用的 GPU 數量來計算。

技術實作：針對摘要生成進行 BART 的微調

為了展示此整合，我們在 samsum 資料集上對 BART-large-cnn 模型進行了微調，該資料集包含約 16,000 則類似通訊軟體的對話及其對應的摘要。

硬體與超參數

訓練是使用以下配置進行的：

  • 實例類型： ml.p3dn.24xlarge (包含 8x NVIDIA V100 GPUs)
  • 實例數量： 2 (總計 16 GPUs)
  • 總批次大小： 64 (每個裝置 4 * 16 GPUs)
  • 模型大小： 4 億個參數
  • 訓練輪次 (Epochs)： 3
  • 學習率： 5e-5
  • 精度： 已啟用 FP16

效能與成本

訓練任務在 2,882 個計費秒數內完成。對於使用 16 顆 NVIDIA Tesla V100 GPUs 的此特定配置，成本約為 $28。

模型部署與 Hub 整合

訓練完成後，模型產出物 (artifacts) 會儲存在 Amazon S3 中。分享模型的流程包括：

  1. 下載產出物： 使用 S3Downloader 從 S3 取得訓練好的模型。
  2. 建立模型卡 (Model Card)： 生成包含超參數與評估指標（例如 ROUGE-1、ROUGE-2 與 ROUGE-L）的 README.md。
  3. ** Hub 上傳：** 利用 huggingface_hub SDK 建立儲存庫並將微調後的模型推送到 huggingface.co，以便進行公開存取與託管推論測試。

Sources