nnScaler は、単一 GPU での実行用に設計されたディープ ニューラル ネットワーク (DNN) モデルを、複数の GPU で並列実行できるプログラムにコンパイルする並列化エンジンです。
使いやすさ: 自動並列化を有効にするには、数行のコードを変更するだけで済みます。
Pythonic: 並列化出力は PyTorch コードであるため、ユーザーが理解しやすく、さらなる開発やカスタマイズに便利です。
拡張性: nnScaler は、新しいモデルの新しい演算子をサポートする API を公開します。
信頼性: nnScaler は、さまざまなエンドツーエンドのトレーニング セッションを通じて検証され、信頼できるシステムです。
パフォーマンス: 大規模な並列化スペースを探索することで、nnScaler は並列トレーニングのパフォーマンスを大幅に向上させることができます。
DNN 科学者は、並列化の複雑さは nnScaler に任せながら、単一 GPU 上の PyTorch を使用したモデル設計に集中できます。パフォーマンスにおいて既存の手法を上回る革新的な並列処理手法が導入されています。さらに、nnScaler は、新しい構造または実行パターンによる DNN モジュールの拡張をサポートし、ユーザーがカスタム DNN モデルを並列化できるようにします。
DNN システムの専門家は、nnScaler を利用して、新しいモデルの新しい DNN 並列化メカニズムとポリシーを検討できます。 nnScaler で認識されない新しい演算子にユーザー定義関数を提供することで、新しい DNN モデルのシームレスな並列化が保証されます。たとえば、LLM での長いシーケンスのサポートを容易にするためです。
nnScaler をインストールする前に、次のパッケージをインストールします。
Python >= 3.8, < 3.11 (3.10 is recommanded) PyTorch >= 2.0, < 2.4 (2.2.0 is recommanded)
nnScaler ディレクトリで以下のコマンドを実行します。
pip install -r requirements.txt pip install -e .
さらに、 cppimportエラーを回避するには、環境変数PYTHONPATHに nnScaler ディレクトリを含める必要もあります。
export NNSCALER_HOME=$(pwd)
export PYTHONPATH=${NNSCALER_HOME}:$PYTHONPATHLlama-3 の実行に必要なパッケージをインストールします。また、flash-attn のインストール中には、特定のバージョンの CUDA ライブラリが必要です。たとえば、PyTorch 2.20 を使用する場合は、CUDA V11.8 が必要です。
python -m pip install transformers==4.40.0 flash-attn==2.5.5 tensorboard
HuggingFace から Llama-3 モデルへのアクセスを取得します。ここで、環境変数として設定する必要があるアクセス トークンを受け取ります。
export HF_TOKEN=<HUGGINGFACE_ACCESS_TOKEN>
すべてのサンプル コードはexamples/llama3_8B_128Kにあります。以下に示すように、ユーザーは次のことを行う必要があります。
モデルをラップする: 損失計算とその他の必要なコンポーネントを含めます。
コンポーネントの構成: モデル、オプティマイザー、およびデータローダーをセットアップします。
初期化と開始: main 関数で、上記の構成で nnScaler トレーナーを作成し、トレーニング プロセスを開始します。
# nnScaler ビルトインの並列化対応トレーナーをインポートしますfrom nnscaler.cli.trainer import Trainer# 損失計算などを含めるためにモデルをラップします。class WrapperModel(torch.nn.Module):def __init__(self, model_id):super()。 __init__()self.model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation='flash_attention_2')def forward(self, tables):outputs = self.model.model(input_ids=samples['net_input']['src_tokens'],use_cache=False,return_dict=False,
)loss = torch.sum(chunk_linear_cross_entropy(outputs[0], self.model.lm_head.weight, tables['target'], ...))リターンロス、samples['ntokens']、samples['nsentences'] def main(args):# データ configdataloader_config = ... # モデル configmodel_config = ModelConfig(type=WrapperModel,args={'model_id': args.model_id,
}、
)# オプティマイザー ハイパーパラメーター optimizer_config = OptimizerConfig(type=MixedPrecisionAdamW,args={'lr': 2e-5, 'betas': (0.9, 0.95), 'weight_decay': 0.0, 'fused': True},#... )#...# データローダー/モデル/オプティマイザーなどの構成を使用してトレーナーをセットアップします。トレーナー= Trainer(train_args=TrainerArgs(#...model=model_config,optimizer=optimizer_config,dataloader=dataloader_config,#...))trainer.run()その後、サンプルを開始すると、nnScaler によってすべての並列化タスクが自動的に完了します。
cd example/llama3_8B_128K# トレーニング データを準備します:python bookcorpus.py --data_path_or_name bookcorpus/bookcorpus --tokenizer_path_or_name metal-llama/Meta-Llama-3-8B-Instruct --save_path ./bookcorpus_llama3_4K --sequence_length 4096# ミニ モデルをビルドするpython create_mini_model.py --model_id metal-llama/Meta-Llama-3-8B-Instruct --output_id ./llama3_mini#データ並列処理 + zero1torchrun を使用してコンパイルおよび実行 --nproc_per_node=2 train.py --plan_ngpus 1 --runtime_ngpus 2 --name llama3_debug --model_id ./llama3_mini --dataset_path ./bookcorpus_llama3_4K
また、nnScaler の PyTorch Lightning 互換インターフェイスを介してモデルを並列化する方法を示す例も提供します。
nnScaler リポジトリで nanoGPT の例を見つけます。
cd サンプル/nanogpt
nanoGPT の依存関係をインストールします。
pip install -r 要件.txt
データセットを準備します。
python nanoGPT/data/shakespeare_char/prepare.py
単一の GPU でテストする
これで、 torchrun <https://pytorch.org/docs/stable/elastic/run.html>でtrain_nnscaler.py実行できるようになりました。
torchrun --nproc_per_node=1 train_nnscaler.py nanoGPT/config/train_shakespeare_char.py
これにより、単一の GPU でベビー GPT モデルがトレーニングされます。数分かかり、最良の検証損失は約 1.47 になります。
マルチGPUでのテスト
デフォルトでは、nnScaler はデータ並列処理を使用して GPU 上でモデルを並列化します。 1 つのノードに 4 つの GPU がある場合:
torchrun --nproc_per_node=4 train_nnscaler.py nanoGPT/config/train_shakespeare_char.py
または、複数のノードがある場合 (たとえば、それぞれ 4 つの GPU を備えた 2 つのノード):
# on each node torchrun --nnodes=2 --nproc_per_node=4 --rdzv-id=NNSCALER_NANOGPT --rdzv-backend=c10d --rdzv-endpoint=<IP> train_nnscaler.py nanoGPT/config/train_shakespeare_char.py
注: ローカル バッチ サイズはデフォルトで固定されているため、より多くのワーカーを使用すると、グローバル バッチ サイズが大きくなります。
?高度な使用法については、今後のリリースにご期待ください。
nnScaler は、製品と研究の両方を含む複数のプロジェクトで採用されています。
(YOCO) キャッシュは 1 回のみ: 言語モデルのデコーダ-デコーダ アーキテクチャ
LongRoPE: LLM コンテキスト ウィンドウを 200 万トークンを超えて拡張
Phi-3 シリーズのロングコンテキストバージョンのポストトレーニング
OSDI'24 のアーティファクト評価とガイダンスはここにあります。 nnScaler があなたの研究に役立つ場合は、出版物の中でそれを引用してください。
@inproceedings{lin2024nnscaler,
title = {nnScaler: Constraint-Guided Parallelization Plan Generation for Deep Learning Training},
author={Lin, Zhiqi and Miao, Youshan and Zhang, Quanlu and Yang, Fan and Zhu, Yi and Li, Cheng and Maleki, Saeed and Cao, Xu and Shang, Ning and Yang, Yilei and Xu, Weijiang and Yang, Mao and Zhang, Lintao and Zhou, Lidong},
booktitle={18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24)},
pages={347--363},
year={2024}
}このプロジェクトは貢献と提案を歓迎します。ほとんどの投稿では、投稿を使用する権利をお客様が有しており、実際に当社に付与することを宣言する投稿者ライセンス契約 (CLA) に同意する必要があります。詳細については、https://cla.opensource.microsoft.com をご覧ください。
プル リクエストを送信すると、CLA ボットが CLA を提供する必要があるかどうかを自動的に判断し、PR を適切に装飾します (ステータス チェック、コメントなど)。ボットが提供する指示に従ってください。 CLA を使用するすべてのリポジトリでこれを 1 回行うだけで済みます。
このプロジェクトはマイクロソフトのオープンソース行動規範を採用しています。詳細については、行動規範に関する FAQ を参照するか、追加の質問やコメントがあれば [email protected] までお問い合わせください。
このプロジェクトには、プロジェクト、製品、またはサービスの商標またはロゴが含まれている場合があります。 Microsoft の商標またはロゴの許可された使用には、Microsoft の商標およびブランド ガイドラインが適用され、それに従わなければなりません。このプロジェクトの修正バージョンで Microsoft の商標またはロゴを使用することは、混乱を引き起こしたり、Microsoft のスポンサーであることを暗示したりしてはなりません。第三者の商標またはロゴの使用には、それらの第三者のポリシーが適用されます。
弊社のパブリック リポジトリは https://github.com/microsoft/nnscaler または Microsoft 内部リポジトリ https://aka.ms/ms-nnscaler から見つけることができます。ご質問やお問い合わせについては、[email protected] までご連絡ください。
