倹約性と精度の融合: AWS Trainium を使用した GPT NeoX および Pythia モデルのコスト効率の高いトレーニング | アマゾン ウェブ サービス

大規模言語モデル (LLM) は、日常の会話の話題になっています。 その急速な普及は、ユーザー数 100 億人に達するまでに要した時間を見れば明らかです。この期間は、「Facebook の場合は 4.5 年」から、過去最低の「ChatGPT の場合はわずか 2 か月」にまで短縮されています。 生成事前トレーニング変換器 (GPT) は、因果的自己回帰更新を使用して予測を行います。 音声認識、テキスト生成、質問応答などのさまざまなタスクが、これらのモデル アーキテクチャによって驚異的なパフォーマンスを発揮することが実証されています。 最近のモデルのいくつか NEOX, ファルコン, ラマ GPT アーキテクチャをバックボーンとして使用します。 LLM のトレーニングには膨大なコンピューティング時間が必要で、それには数百万ドルの費用がかかります。 この記事ではGPTのトレーニング手順をまとめます。 NEOX on AWS トレーニング、ディープ ラーニング トレーニング用に最適化された専用の機械学習 (ML) アクセラレータです。 モデルの品質を損なうことなく、AWS Trainium を使用してそのようなモデルをコスト効率よく (3.2 万トークン/ドル) トレーニングした方法について概要を説明します。

ソリューションの概要

GPT NeoX および Pythia モデル

GPTネオX および ピティア Eleuther-AI によるオープンソースの因果言語モデルであり、NeoX では約 20 億、Pythia では約 6.9 億のパラメーターを備えています。 どちらも、Chat GPT3 と同様のアーキテクチャ設計に従っているデコーダー モデルです。 ただし、いくつかの追加機能もあり、Llama などの最近のモデルにも広く採用されています。 特に、ヘッドの寸法全体にわたって部分的に回転する回転位置埋め込み (ROPE) を備えています。 オリジナルのモデル (NeoX および Pythia 6.9B) は、公開されているモデルでトレーニングされています パイルデータセット 重複排除を使用し、Megatron と Deepspeed バックエンドを使用します。

を使用して、AWS Trainium ベースの Trn1 インスタンス上でこれらのモデルの事前トレーニングと微調整をデモンストレーションします。 ニューロン・ニモ 図書館。 概念実証と迅速な再現を確立するために、GPT2 バイトペア エンコーディング (BPE) トークナイザーを使用してトークン化された、より小さい Wikipedia データセット サブセットを使用します。

チュートリアル

次に示すように、事前にトークン化された Wikipedia データセットをダウンロードします。

export DATA_DIR=~/examples_datasets/gpt2 mkdir -p ${DATA_DIR} && cd ${DATA_DIR} wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.bin . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.idx . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/license.txt . --no-sign-request

NeoX 20B と Pythia 6.9B は両方とも、部分的な回転を伴う ROPE を使用します。たとえば、頭の寸法の 25% を回転し、残りは回転しないままにします。 AWS Trainium アクセラレータで部分回転を効率的に実装するには、回転ディメンションと非回転ディメンションを連結する代わりに、非回転ディメンションにゼロ周波数を追加してから、ヘッドディメンションの完全なセットを回転します。 この簡単なトリックにより、AWS Trainium のスループット (XNUMX 秒あたりに処理されるシーケンス) が向上しました。

トレーニングのステップ

トレーニングを実行するには、SLURM マネージド マルチノード Amazon Elastic Compute Cloud (Amazon EC2) Trn1 クラスター。各ノードには trn1.32xl インスタンスが含まれます。 それぞれ trn1.32xl には 16 個のアクセラレータがあり、アクセラレータごとに XNUMX 人のワーカーがいます。 最新のものをダウンロードしたら、 ニューロン・ニモ パッケージの場合は、提供されているものを使用してください ネオックス および ピュシア 最適化されたハイパーパラメータを使用して事前トレーニングおよび微調整スクリプトを作成し、XNUMX ノードのトレーニングに対して次のコマンドを実行します。

1. コンパイル: XNUMX つのトレーニング反復でモデルをプリコンパイルし、グラフを生成して保存します。

   sbatch --nodes 4 compile.slurm ./neoX_20B_slurm.sh

2. 実行: 最初のステップからキャッシュされたグラフをロードしてトレーニングを実行します。

   sbatch --nodes 4 run.slurm ./neoX_20B_slurm.sh

3. 結果を監視する

   tensorboard --logdir=nemo_experiments/megatron_neox

Pythia 6.9B モデルを置き換えて実行するには、同じ手順に従う必要があります。 neox_20B_slurm.sh by pythia_6.9B_slurm.sh.

事前トレーニングと微調整実験

AWS Trainium で GPT-NeoX および Pythia モデルの事前トレーニングをデモンストレーションします。 ニューロン・ニモ 10k 反復のライブラリと、1k ステップでのこれらのモデルの微調整も示します。 事前トレーニングには、NeMo 内の GPT2 BPE トークナイザーを使用し、同じに従います。 設定 オリジナルモデルで使用されていたものと同じです。 AWS Trainium で微調整するには、いくつかのパラメータの変更が必要です (例: 語彙サイズ分割係数)、これらは、Megatron と NeMo の違い、および GPU と AWS Trainium の変更に対応するために微調整スクリプトで提供されます。 ノード数を変えた場合のマルチノード分散トレーニング スループットを表 1 に示します。

表1。 ノード数を変えて最大 500 ステップをトレーニングする場合の GPT NeoX モデルと Pythia モデルの平均スループットを比較します。 の trn1.32xlの価格 3 年間の予約有効時間当たり料金に基づいています。

次に、AWS Trainium でのモデルトレーニングの損失軌跡も評価し、P4d (Nvidia A100 GPU コア) クラスターでの対応する実行と比較します。 トレーニング損失に加えて、トレーニングの進行状況を監視するために各トレーニング反復で計算されるモデル勾配の 2 ノルムである勾配ノルムなどの有用な指標も比較します。 トレーニング結果を図-1、2に示し、NeoX 20Bの微調整を図-3に示します。

図1。 すべてのワーカーの平均したトレーニング損失 (左) と各ステップのトレーニングにおける勾配ノルム (右)。 NeoX 20B は、同じトレーニング ハイパー パラメーター (グローバル バッチ サイズ = 4) を使用して、GPU および Trainium 上の小さな wiki データセットを使用して 256 つのノードでトレーニングされます。 GPU は BF16 とデフォルトの混合精度を使用していますが、AWS Trainium は確率的丸めを備えた完全な BF16 を使用しています。 損失と勾配ノルムの軌跡は、GPU と AWS Trainium に一致します。

図2。 すべてのワーカーの平均したトレーニング損失 (左) と各ステップのトレーニングにおける勾配ノルム (右)。 図 1 の GPT NeoX と同様に、Pythia 6.9B は、同じトレーニング ハイパー パラメーター (グローバル バッチ サイズ = 4) を使用して、GPU および Trainium 上の小さな wiki データセットを使用して 256 つのノードでトレーニングされます。 損失と勾配ノルムの軌跡は、GPU と Trainium で一致します。

図3。 すべてのワーカーで平均化されたトレーニング損失 (左) と勾配ノルム (右) を使用して、GPU と AWS Trainium で GPT NeoX 20B モデルを微調整します。 微調整のデモンストレーションには、小さな Wiki データセットが使用されます。 損失と勾配ノルムの軌跡は、GPU と AWS Trainium に一致します。

この投稿では、AWS 深層学習ハードウェアでの LLM のコスト効率の高いトレーニングについて説明しました。 Neuron NeMo ライブラリを使用して、AWS Trn20 上で GPT NeoX 6.9B および Pythia 1B モデルをトレーニングしました。 AWS Trainium を使用した 20 億モデルのコスト正規化スループットは、支出額あたり約 3.2 万トークン/ドルとなります。 AWS Trainium でのコスト効率の高いトレーニングに加えて、同様のモデル精度が得られます。これは、トレーニングのステップ損失と勾配ノルム軌道から明らかです。 また、AWS Trainium 上の NeoX 20B モデルで利用可能なチェックポイントを微調整しました。 AWS Trainium での NeMo Megatron による分散トレーニングの詳細については、を参照してください。 NeMo Megatron の AWS Neuron リファレンス。 Llama モデルの微調整を開始するための優れたリソースはここにあります。 Llama2の微調整。 マネージド AWS Trainium の使用を開始するには アマゾンセージメーカーを参照してください AWS Trainium と Amazon SageMaker を使用して ML モデルをトレーニングする.

著者について

ガウラヴ・グプタ 現在、アマゾン ウェブ サービス (AWS) AI ラボの応用科学者です。 グプタ博士は、USC Viterbi で博士号を取得しました。 彼の研究対象は、逐次データ モデリング、偏微分方程式の学習、機械学習のための情報理論、部分力学モデル、複雑なネットワークの領域に及びます。 彼は現在、LLM のトレーニング動作、PDE を使用したビジョン モデル、情報理論のマルチモダリティ モデルに関する応用数学的問題に取り組んでいます。 Gupta 博士は、Neurips、ICLR、ICML、Nature、IEEE Con​​trol Society、ACM Cyber​​-Physical Society などのトップジャーナル/会議で論文を発表しています。

ベン・スナイダー は、AW​​S Deep Learning の応用科学者です。 彼の研究対象には、基礎モデル、強化学習、非同期最適化が含まれます。 仕事以外では、サイクリングとバックカントリーキャンプを楽しんでいます。

アミット(R)マミダラ は、AW​​S Annapurna Labs のシニア機械学習アプリケーション エンジニアリングです。 マミダラ博士は、オハイオ州立大学で高性能コンピューティングと通信の博士号を取得しました。 IBM 研究社に勤務していた間、マミダラ博士は、最も強力で電力効率の高いスーパーコンピューターのトップ 500 ランキングをしばしばリードする BlueGene クラスのコンピューターに貢献しました。 このプロジェクトは、2009 年国家技術革新賞を受賞しました。 金融ヘッジファンドで AI エンジニアとして短期間勤務した後、マミダラ博士はアンナプルナの研究室に加わり、大規模言語モデルのトレーニングに重点を置きました。

ジュン（ルーク）・フアン AWS AI Labs の主席科学者です。 フアン博士は AI とデータ サイエンスに取り組んでいます。 彼は、主要な会議やジャーナルに 180 以上の査読済み論文を発表しています。 彼は 2009 年に NSF 教員早期キャリア開発賞を受賞しました。AWS に入社する前は、Baidu Research で著名な科学者および Baidu ビッグデータ ラボラトリーの所長として働いていました。 AI スタートアップである StylingAI Inc. を設立し、2019 年から 2021 年まで CEO およびチーフサイエンティストを務めました。 業界に入社する前は、カンザス大学 EECS 学部の Charles E. and Mary Jane Spahr 教授を務めていました。

シュルティ・コパーカー AWS のシニア プロダクト マーケティング マネージャーです。 彼女は、顧客が機械学習のニーズに合わせて Amazon EC2 アクセラレーション コンピューティング インフラストラクチャを探索、評価、導入できるよう支援します。

• SEO を活用したコンテンツと PR 配信。 今日増幅されます。
• PlatoData.Network 垂直生成 Ai。 自分自身に力を与えましょう。 こちらからアクセスしてください。
• プラトアイストリーム。 Web3 インテリジェンス。 知識増幅。 こちらからアクセスしてください。
• プラトンESG。 カーボン、 クリーンテック、 エネルギー、 環境、 太陽、 廃棄物管理。 こちらからアクセスしてください。
• プラトンヘルス。 バイオテクノロジーと臨床試験のインテリジェンス。 こちらからアクセスしてください。
• 情報源： https://aws.amazon.com/blogs/machine-learning/frugality-meets-accuracy-cost-efficient-training-of-gpt-neox-and-pythia-models-with-aws-trainium/