Transformersコミュニティ、ZeRO-3で読み込み処理が重複する問題を報告　各プロセスが完全なチェックポイントを読み込み、変換する可能性

Transformersコミュニティは10月11日、分散モデル読み込みに関する問題を報告した。DeepSpeed ZeRO-3がモデルパラメーターを分割していても、from_pretrained()によって各プロセスが完全なチェックポイントを読み込み、重み変換を重複して実行する可能性があるという。報告では、指定コミットのTransformers 5.19.0.dev0、PyTorch 2.10.0、DeepSpeed 0.18.4を使用している。現在もIssueはオープンで、関連する修正PRは記載されていない。元の報告

ここでは、異なる2つのメモリライフサイクルが関係する。ZeRO-3が分割するのは、学習中のパラメーター、勾配、オプティマイザー状態だ。公式ドキュメントによると、パラメーターは順伝播と逆伝播の際に自動的に集約され、その後再分割される。そのため、学習状態の分割による効果を、チェックポイント読み込み時のメモリ使用量の保証とみなすことはできない。DeepSpeedドキュメント

報告者は、ランダムに初期化した2層のQwen3 MoEから3つのsafetensorsシャードを作成し、2つのプロセスで読み込む構成を用意した。事前学習済みモデルのダウンロードや学習は不要だ。関数の呼び出し回数を数えたところ、両方のrankがそれぞれ3つの重みファイルを読み込み、変換を1回ずつ実行していた。一方、その後のパラメーター書き込みはrank 0に限られていた。報告者は問題箇所を非量子化ZeRO-3の読み込み分岐と特定し、rank 0が読み込みと変換を一括して行い、必要な情報を配布する方法を提案している。再現例とコードパスの分析

大規模モデルのファインチューニングでは、この報告を受け、エンジニアは起動時のホストメモリ使用量と読み込み時間を別途監視する必要がある。これは読み込みの仕組みから導かれる実務上の影響だが、プロセス数に応じてメモリ使用量が倍増すると結論づけることはできない。最小構成のケースは正常に完了しており、読み込み回数は物理ディスクからの読み取り量と同じではない。メモリマッピングやページキャッシュも結果に影響する。

報告者によると、部分的な修正を適用すると、ゼロ以外のrankで読み込みと変換の回数がゼロになり、パラメーターと永続バッファーの整合性も確認できたという。ただし、メンテナーによるレビューと、より広範な検証が必要だ。今後は、マルチノードでの同期、バッファーの配布、正式リリースへの影響を確認する必要がある。調査時にはファイル形式も確認したい。Hugging Faceのドキュメントによると、DeepSpeedネイティブ形式の分割チェックポイントはfrom_pretrained()にそのまま渡せないため、形式の問題と今回の重複読み込みを混同しないよう注意が必要だ。チェックポイントのドキュメント