YaFSDP:高效的分散式資料並行框架,專為大型語言模型設計

連結:https://github.com/yandex/YaFSDP

YaFSDP是一個分散式資料並行框架,專為與transformer類神經網路結構良好協作而設計。它在預訓練大型語言模型(Large Language Models, LLMs)時比傳統的FSDP快20%,並且在高記憶體壓力條件下表現更佳。YaFSDP旨在減少通訊和記憶體操作的開銷。

需求人群:

  • YaFSDP框架適用於需要處理大規模資料和模型的機器學習研究人員和工程師。它特別適合於那些需要在高記憶體壓力環境下進行深度學習模型訓練的場景,例如大型語言模型的預訓練和微調。

使用場景示例:

  • 使用YaFSDP進行7B至70B引數規模的語言模型預訓練。
  • 在64至256個設備上應用YaFSDP進行模型訓練,以提高效率。
  • 利用YaFSDP進行2048至8192個token序列的模型訓練。

產品特色:

  • 支援大規模語言模型的高效預訓練。
  • 最佳化了記憶體和通訊操作,提高了訓練效率。
  • 提供了詳細的使用示例,包括因果預訓練和監督微調。
  • 基於NVIDIA PyTorch影象構建,整合了必要的補丁庫。
  • 支援自訂事件通知,便於開發者根據需要接收更新。
  • 在A100 80G叢集上進行了效能評估,確保了框架的高效能。

使用教學:

1. 克隆YaFSDP的GitHub倉庫到本地環境。

2. 根據示例資料夾中的指導文檔設定Docker環境。

3. 執行docker/build.sh腳本來構建所需的Docker映象。

4. 根據具體的訓練需求,選擇合適的示例腳本進行模型訓練。

5. 監控訓練過程中的記憶體和通訊開銷,確保系統穩定執行。

6. 根據需要調整YaFSDP的配置引數,最佳化模型訓練效能。

返回頂端