Transformers コミュニティ、ZeRO-3 の読み込みを修正するパッチを提案　全重みの重複読み込みと変換を削減

Hugging Face Transformersコミュニティは10月14日、DeepSpeed ZeRO-3で非量子化モデルを読み込む際に、分散プロセスごとに完全なチェックポイントを重複して読み込み、重み変換を実行する問題に対処するPR #49584を提案した。確認時点で提案はオープンのままで、レビュー承認や正式リリースへの収録は確認されていない。修正提案

最初の報告ではTransformers 5.19.0.dev0を使用し、ランダム初期化した2層のQwen3 MoEを使って最小構成で再現した。2つのrankがそれぞれ3つの重みファイルを読み込み、変換を1回実行する一方、その後チェックポイントの値をモデルにコピーするのはrank 0だけだった。これらの一時テンソルはZeRO-3で分割されたモデルパラメーターには含まれないため、モデル分割後も読み込み段階で重複したメモリ確保が発生し得る。問題報告と再現コード

公式ドキュメントによると、ZeRO-3はパラメーター、勾配、オプティマイザーの状態を分割し、パラメーターをCPUにオフロードすることもできる。この事例は、デプロイ時に、学習中の分割方式と初期化段階で必要となるホストメモリを別々に評価する必要があることを示している。DeepSpeed ZeROドキュメント

提案では、rank 0が重みを読み込んで変換し、キー名、形状、データ型をブロードキャストする。ほかのプロセスはmetaテンソルのプレースホルダーを作成し、既存のパラメーター割り当て機構を引き続き利用する。ZeRO-3で分割されないバッファーは別途ブロードキャストし、読み込みや変換のエラーもほかのrankに伝えることで、プロセスが通信待ちのまま停止するのを防ぐ。修正設計

作者によると、8台のデバイスを使ったGLM-5.3-Flash-BF16のテストでは、修正版で読み込みと15ステップのトレーニングを完了し、パラメーターとバッファーの一致性、およびエラー伝播のテストも添えた。ただし、未修正版の実行はメモリー保護スクリプトによって早期に停止されたため、両者のピークメモリーを完全に比較した結果ではない。今後エンジニアが注目すべき点は、メンテナーによるレビュー、複数モデルでの検証、リリース版への反映であり、読み込みコストは各自の環境でも測定する必要がある。検証結果と制限