Transformersコミュニティ、FSDP2のロード時メモリ不具合を報告　複数GPUの起動時に完全なモデルが重複して作られる可能性

10月4日、Hugging Face TransformersコミュニティでFSDP2のロードに関する問題が報告された。cpu_ram_efficient_loading=Trueを有効にしても、rank 0以外のプロセスがCPU上に完全なモデルを構築する場合があるという。報告事例では、Transformers 5.12.0、Accelerate 1.15.0、PyTorch 2.10を使用し、単一ノードで16プロセスを起動して、約270億パラメーターのBF16密モデルをロードしたところ、ホストメモリが枯渇した。問題報告

報告者は、問題がfrom_pretrained()の段階で起きると説明している。metaデバイス上のパラメーターを処理する際、該当する分岐がrank 0以外のプロセスでtorch.zeros_like(..., device="cpu")を呼び出し、実体のあるCPUテンソルを作成するという。metaテンソルは通常、形状などの情報だけを保持し、パラメーターの保存領域を確保しない。早い段階で実体化すると、各プロセスが約54 GBを使用する可能性がある。報告された条件に基づく概算では、モデル16個分で約864 GBとなり、ほかのロード時のオーバーヘッドは含まれていない。この数値は報告事例の構成に基づく容量見積もりであり、あらゆる環境に共通するしきい値ではない。コードパスと再現手順

この挙動は、設定の設計目的と食い違っている。Accelerateの公式ドキュメントによると、CPU RAM効率化ロードを有効にした場合、事前学習済みチェックポイントをロードするのは最初のプロセスだけで、ほかのプロセスは空の重みを保持し、その後の同期でパラメーターを取得する想定だ。また、from_pretrained()の前に分散プロセスグループを初期化する必要がある。FSDPのパラメーターシャーディングは学習時のメモリ使用量を抑えられるが、シャーディング前に完全なコピーが作られると、起動時がボトルネックになり得る。FSDP公式ドキュメント

エンジニアリングチームにとって、この報告はモデルのロード中もメモリ容量の検証が必要であることを示している。まず小さめのモデルを使い、from_pretrained()の前後で各rankの常駐メモリを観察し、プロセスグループの初期化順序とパッケージのバージョンを照合するとよい。コストがプロセス数に応じて増えるかも確認できる。現時点で公開されている根拠は単一ユーザーの報告が中心であり、ほかのバージョン、量子化モデル、マルチノード構成への影響は判断できない。今後は、上流で該当分岐のFSDP2の挙動が確認されるか、また修正後も同期が完了するまで空の重みを維持できるかを追う必要がある。