DeepSpeed v0.19.8：MoEの最適化、並列化の拡張

1. 専門家層ごとに設けていた一時記憶領域を共有する

10月20日に公開された今回のバージョンでは、AutoEPが混合エキスパートモデルの複数の層で1つのDeepEPバッファを共有するようになった。層ごとに設けていた一時記憶領域を共用する方式への変更だ。

AutoEPのデフォルトの結合処理では、重複する32ビット浮動小数点数のコピーも削除した。分割計画とトークンのグループ化、入力断片を集める作業を重ねて処理する変更も含まれる。メモリの扱いと処理順序を併せて見直した形だ。

一部の専門家演算をまとめて実行する融合処理と、使用済みメモリを再び確保するメモリ回収ポリシーがオプションとして追加された。開発者が新しい処理経路を選んで適用できる構成だ。

2. 語彙を分割して処理する学習経路が広がった

AutoTPには、語彙並列の言語モデル出力層と埋め込みのサポートが追加された。入力を表す埋め込みから出力層まで語彙を分割して処理する経路を広げ、埋め込みを共有するモデルのサポートも強化した。

語彙を分割して処理しながら学習損失を計算する融合演算も追加された。出力層のサポートと損失計算を併せて拡張した点で、並列学習を適用できる処理範囲が広がる変更だ。

リリース本文は個々の変更の短い見出しを中心に構成されており、性能比較の数値や測定条件は記載されていない。

3. 学習計算の欠落と保存時の不具合を修正した

ZeRO-1・2では、CPU上に非連続に配置されたグラディエント、つまり学習中に計算される勾配値を合算する際、一部のグループが通知なく欠落する問題を修正した。学習に使う計算結果の一部が合算から抜け落ちる不具合を正したものだ。

ZeRO-2では、分割保存の過程でも32ビットの勾配累積精度を維持するようにした。複数のデバイスに分散した学習計算を統合し、保存する過程の正確性に関する修正だ。

学習状態を保存したチェックポイントの変換では、複数の断片が同じ保存断片のアドレスに書き込むよう指定されたマップを拒否する。同じ位置に重複して書き込む配置情報を遮断するようになった。

4. 削除された機能と拒否される設定を先に確認する

参加デバイス数の変化に対応する弾力的学習と、学習中に層を段階的に取り除く段階的層削除のサポートは削除された。古い損失関数2つも削除され、圧縮ライブラリの残存機能と、簡単な課題から学ぶ旧式のカリキュラム学習は非推奨となった。アップグレード時には、依存する機能が既に削除されたのか、非推奨なのかを先に確認するのが実用的だ。

設定ドキュメントでは、sub_group_sizeの説明と、デフォルト値を1e9に訂正する変更が含まれた。既存の構成とドキュメントのデフォルト値を照合する際に確認すべき項目だ。

DeepSpeedCPUAdamは、以前は無視していたamsgrad設定を拒否するようになった。この設定が残っているプロジェクトは、新バージョンを適用する前にオプティマイザーの構成を点検する必要がある。

並列処理機能の拡張と計算不具合の修正が利点である一方、旧機能を使うコードや許可されない設定の整理は、開発者が負担する移行コストとなる。特に機能の削除が含まれるため、パッチリリースという名称よりも、使用中の機能の変更内容が適用判断で重要になる。

5. アクセラレーターのサポートとビルド条件も変わる

NPUには、CPUメモリをデバイスアクセス用に固定するホストメモリ登録と、演算をグラフにまとめて実行するグラフ演算のサポートが追加された。デバイスで実行できる演算とメモリ処理関連のサポートを強化した変更だ。

NPUとXPUの演算ビルダー、つまり演算コードを生成してビルドするツールには、ビルド言語規格の変更が含まれた。

torch 2.12以上ではC++20を使用するようになった。該当する環境で演算をビルドする開発者が確認すべき変更だ。

NPUでは、非推奨のAdam呼び出しを融合演算に置き換えた。リリースに記載された置き換え対象の従来の呼び出しは次のとおりだ。

npu_apply_adam(_w)

今回のリリースは、専門家モデルと並列学習の処理範囲を広げる一方、旧機能の整理とアクセラレーターごとのビルド点検をアップグレードの課題として残す。