大規模言語モデル (LLM) が簡単になった EasyLM は、JAX/Flax での LLM の事前トレーニング、微調整、評価、提供を行うためのワンストップ ソリューションです。 EasyLM は、JAX の pjit 機能を利用して、LLM トレーニングを数百の TPU/GPU アクセラレータにスケールアップできます。
Hugginface のトランスフォーマーとデータセットの上に構築されたこのリポジトリは、他の多くのフレームワークのような複雑さを伴うことなく、大規模な言語モデルをトレーニングするための使いやすくカスタマイズしやすいコードベースを提供します。
EasyLM は JAX/Flax で構築されています。 JAX の pjit ユーティリティを活用することで、EasyLM はモデルの重みとトレーニング データを複数のアクセラレータにわたってシャーディングすることで、単一のアクセラレータに適合しない大規模なモデルをトレーニングできます。現在、EasyLM は、単一ホストでの複数の TPU/GPU トレーニングと、Google Cloud TPU ポッドでのマルチホスト トレーニングをサポートしています。
現在、次のモデルがサポートされています。
私たちは、JAX での LLM のトレーニングに関するディスカッションのために、非公式の Discord コミュニティ (Google とは無関係) を運営しています。このリンクに従って Discord サーバーに参加します。当社には、EasyLM、JaxSeq、Alpa、Lavanter など、いくつかの JAX ベースの LLM フレームワークの専用チャネルがあります。
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 ホスト VM には、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}
}