DeepSpeedがモデル学習の開発者向けに0.20.0を公開した。同じ条件で結果を再現するための実行ポリシーを選べるようになり、演算フローを記録して再実行する機能は、より多様なアクセラレーターに拡張された。学習中の計算値の保持や保存済みの学習状態を復元する処理の不具合も修正した。
1. 再現性を選択機能として導入
今回のバージョンでは、ユーザーが選択して有効にする決定論的実行ポリシーが追加された。決定論的実行は、同じ条件で計算を繰り返す際に結果の再現を図る方式だ。学習結果を比較する開発者にとって、実行方式の新たな選択肢となる。
学習計算の正確性に関する修正も含まれる。複数のサブネットワークをエキスパートとして使い分けるモデルでは、勾配の大きさを制限する計算で、各エキスパートの値を1回ずつだけ数えるよう修正した。このとき、大きさの基準となる値は32ビット浮動小数点で計算する。
エキスパートとテンソルを分割して処理するAutoEPとAutoTPを組み合わせた構成でも、エキスパートごとの勾配調整の不具合を修正した。実行ポリシーの追加とともに、実際の学習計算経路にも手を入れた形だ。
2. グラフキャプチャを非CUDAアクセラレーターに拡張
Hybrid EngineはCUDA以外のアクセラレーターでもグラフキャプチャに対応する。グラフキャプチャは、毎回実行する演算フローをあらかじめ記録し、繰り返し利用する方式だ。今回の変更により、この機能を活用できるハードウェアの範囲が広がる。
グラフキャプチャは、アクセラレーターごとに明示的な有効化が必要となる。Kunlunxin XPUにもグラフキャプチャ対応を追加した。そのため、導入を検討する際は、アクセラレーターごとの対応状況と有効化の要件を併せて確認することが出発点となる。
メモリ管理も実行環境に合わせて変更された。CPUアクセラレーターが認識するメモリ上限は、プロセスに適用されるリソース管理上の上限で制限する。システム全体のメモリ容量と、個々のプロセスに許可された容量を区別して扱う変更だ。
3. ZeROの勾配保持とメモリ処理を修正
学習状態を分割して管理する機能であるZeROのステージ1・2（ZeRO-1/2）は、リクエストの境界を越えても、重みを更新する時点まで勾配を保持する。勾配は、学習過程で重みをどの方向に変えるかを計算した値だ。複数のリクエストの計算結果を集めて1回の更新に使う経路を修正した。
学習状態をCPU側に移す構成では、グループごとに累積器を1つ設け、未使用の勾配を0にする。ZeRO-2は、この構成でリクエスト間の勾配を分割領域内に集めて累積するよう変更された。ここで累積器はグループの勾配を集めるCPUの記憶領域、分割領域は分割管理する勾配のうち該当する部分を保持する領域だ。
ZeRO-3は、デバイス外に置いた更新基準となる重みを、固定メモリバッファーを経由してデバイスにコピーする。重みの分割がすでに基準重みと同じデータ型なら、別のコピーを作らずに再利用する。デバイスへのデータ転送経路とコピーの管理方式を併せて見直した。
4. 学習状態の復元とストレージの不具合も修正
複数の学習構成で共通して読み込める保存形式であるユニバーサルチェックポイントの読み込み時に、演算をまとめて処理する融合型およびグラフキャプチャ対応オプティマイザーの学習ステップ値を、正しい位置に復元するよう修正した。オプティマイザーは、勾配を使って重みを更新する構成要素だ。
CPU上に置くオプティマイザーの状態のうち、元が32ビット浮動小数点の状態は、保存済みの状態を読み込んだ後も同じ精度を維持する。
ストレージに学習状態を置くZeRO-3構成では、パラメーターを先読みする際に、デバイス間の移動に使う一時記憶領域であるスワップバッファーがすでに割り当てられていると発生していた不具合を修正した。DeepNVMeは、アラインメントが揃っていないバッファーで発生していた入出力失敗を修正した。
エンジンとZeROオプティマイザーにあったメモリリークも修正した。学習計算自体だけでなく、状態の復元、データの移動、メモリ管理までが今回の修正範囲に含まれる。
5. 選択設定の柔軟性と検証の負担
複数のリクエストの勾配を集めて1回の更新に使うZeRO-1/2の学習では、リクエスト境界を越えた後の勾配保持が導入判断の基準となる。ZeRO-3では、重みの分割と基準重みのデータ型が同じ構成が、別の基準重みのコピーではなく既存の分割を再利用する変更の検討対象となる。
プロセスにリソース上限が適用されたCPU学習環境では、システム全体の容量ではなく、そのプロセスに許可されたメモリを基準に構成を検討する必要がある。CPUアクセラレーターが認識するメモリ上限にこの制限を反映する変更は、そうした環境に直接関係する。
リリースノートには、決定論的実行ポリシーの具体的な有効化方法と、適用される演算の範囲は説明されていない。
今回のリリースの要点は、再現性の選択肢とアクセラレーターの対応範囲を広げつつ、ZeROの学習・メモリ経路を安定化させたことにある。