ChatGPTなどに使われる大規模言語モデルを従来のシステムよりも15倍高速・低コストで学習できる「DeepSpeed-Chat」をMicrosoftが公開

OpenAIが提供する「ChatGPT」などのチャットAIは要約やコーディング、翻訳などを人間の専門家以上の精度で実行できると報告されています。しかしチャットAIの訓練に必要な人間のフィードバックに基づいた強化学習(RLHF)を実行するエンドツーエンドなパイプラインが存在せず、最先端のチャットAIの訓練を行うことは困難でした。しかしMicrosoftが発表した「DeepSpeed-Chat」では誰でもChatGPTのようなモデルを作成可能です。

DeepSpeed/blogs/deepspeed-chat/japanese at master · microsoft/DeepSpeed · GitHub
https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-chat/japanese

これまでChatGPTのようなモデルの訓練で必要となるRLHFを、簡単かつ高い効率で実行できるパイプラインは存在していませんでした。また、ChatGPTのようなAIモデルをトレーニングするには高価なGPUが複数必要になるため、一般の開発者にはこの種のAIモデルを開発することが困難でした。また、GPUを用意したとしても従来のソフトウェアではハードウェアの5％未満の性能しか引き出せず、簡単かつ高速に、かつ低コストで数千億のパラメータを持つモデルの訓練は不可能だったことが報告されています。

そこでMicrosoftは、開発者がより手ごろな価格でチャットAIを開発できるようにすることを目的としたフレームワーク「DeepSpeed-Chat」を発表しました。

DeepSpeed-ChatはChatGPTの元となったInstructGPTで行われた「教師付きファインチューニング」や「報酬モデルのファインチューニング」「RLHFの訓練」の3ステップを実行し、独自のChatGPTライクなモデルを生成できるスクリプトを提供します。また、学習後の会話形式をテストするための推論APIも提供するとのこと。

さらにDeepSpeed-Chatに搭載されている「DeepSpeed-RLHF パイプライン」は、「教師付きファインチューニング」「報酬モデルのファインチューニング」「RLHFの訓練」を行うとともに、研究者や開発者が複数のデータリソースを用いて独自のRLHFモデルを訓練するのを手助けするため、「データの抽象化」や「ブレンド機能」を実行することが可能です。「データの抽象化」では異なるデータセットの形式を統一するために抽象化したデータセットを作成し、「ブレンド機能」は複数のデータセットを適切に融合し、「教師付きファインチューニング」などの3つのトレーニングに分割します。

また、「DeepSpeed-RLHF パイプライン」による学習を幅広いハードウェアで高速かつ低コストで実行するために、これまでDeepSpeedが発表したZeROなどの推論と学習のための全システムを融合した「DeepSpeedハイブリッドエンジン」が構成されています。

DeepSpeedハイブリッドエンジンを搭載したDeepSpeed-Chatを使用し、Microsoft Azure上でデータセンター用のGPU「NVIDIA A100」を64台用いて学習を行った場合、「OPT-13B」モデルは約7.5時間で訓練が完了します。また、その際の費用は1920ドル(約25万円)です。さらに「BLOOM」モデルでは約20時間、5120ドル(約68万円)で訓練が完了するとのこと。これらの数字は既存のRLHFシステムよりもはるかに高速かつ低コストで学習を行うことが可能であることを示しています。

またDeepSpeed-Chatでは、数十億から1兆程度のパラメータを持つ大規模なモデルの訓練と推論が可能で、限られたGPUリソース環境においても訓練と推論を行うことが可能になるとされています。

Hacker Newsでは「DeepSpeed-ChatによってGPT-4の再現が簡単になるというわけではありませんが、再現に向けたいくつかの大きなハードルは間違いなく越えることが可能です」と述べられています。また、MicrosoftはDeepSpeed-Chatを開発するDeepSpeedに無償で100億ドル(約1.3兆円)を出資してChatGPTのような機能をMicrosoftの製品に組み込む研究を支援していることが述べられています。

DeepSpeed-ChatのソースコードなどはGitHub上で公開されています。

GitHub - microsoft/DeepSpeed: DeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.
https://github.com/microsoft/DeepSpeed/