EasyLM 讓大型語言模型 (LLM) 變得簡單,是在 JAX/Flax 中預先訓練、微調、評估和服務 LLM 的一站式解決方案。 EasyLM 可以利用 JAX 的 pjit 功能將 LLM 訓練擴展到數百個 TPU/GPU 加速器。
該儲存庫建立在 Hugginface 的轉換器和資料集之上,提供了一個易於使用且易於自訂的程式碼庫,用於訓練大型語言模型,而無需許多其他框架的複雜性。
EasyLM 使用 JAX/Flax 建構。透過利用 JAX 的 pjit 實用程序,EasyLM 能夠透過在多個加速器上分割模型權重和訓練資料來訓練不適合單一加速器的大型模型。目前,EasyLM 支援在單一主機中進行多個 TPU/GPU 訓練,以及在 Google Cloud TPU Pod 上進行多主機訓練。
目前支援以下型號:
我們正在運行一個非官方的 Discord 社群(與 Google 無關),用於討論與 JAX 培訓法學碩士相關的問題。點擊此連結加入 Discord 伺服器。我們為多個基於 JAX 的 LLM 框架提供專用頻道,包括 EasyLM、JaxSeq、Alpa 和 Levanter。
OpenLLaMA 是我們許可的 LLaMA 複製品,可用於商業目的。在這裡查看項目主頁。 OpenLLaMA 可以取代 EasyLM 中的 LLaMA 權重。請參閱 LLaMA 文件以了解更多詳細資訊。
Koala 是我們在 LLaMA 之上進行微調的新聊天機器人。如果您對我們的 Koala 聊天機器人感興趣,您可以查看有關在本地運行它的部落格文章和文件。
GPU 主機和 Cloud TPU 主機的安裝方法有所不同。第一步是從 GitHub 拉取。
git clone https://github.com/young-geng/EasyLM.git
cd EasyLM
export PYTHONPATH= " ${PWD} : $PYTHONPATH "
GPU環境可以透過Anaconda安裝。
conda env create -f scripts/gpu_environment.yml
conda activate EasyLM
TPU 主機虛擬機器預先安裝了 Python 和 PIP。只需執行以下腳本即可設定 TPU 主機。
./scripts/tpu_vm_setup.sh
EasyLM 文件可以在 docs 目錄中找到。
如果您發現 EasyLM 在您的研究或應用程式中有用,請使用以下 BibTeX 進行引用:
@software{geng2023easylm,
author = {Geng, Xinyang},
title = {EasyLM: A Simple And Scalable Training Framework for Large Language Models},
month = March,
year = 2023,
url = {https://github.com/young-geng/EasyLM}
}