Transformersコミュニティ、ZeRO-3初期化時の精度低下を報告　RoPE周波数がBF16に変換される可能性

Hugging Face Transformersコミュニティは10月29日、長系列学習における精度上の問題を報告した。DeepSpeed ZeRO-3とBF16を使ってモデルを読み込むと、回転位置エンコーディング（RoPE）の逆周波数バッファーinv_freqがBF16で作成され、本来のFP32を維持する動作から逸脱する可能性がある。報告者は5.17.0と開発ブランチで再現し、単一のH200でも発生するとしている。問題報告

問題はモデル初期化と混合精度処理の境界にある。報告者の分析によると、deepspeed.zero.Initはテンソルの構築を捕捉し、浮動小数点の結果を学習精度に変換する。以前のRoPEのリファクタリングでは浮動小数点のtorch.arangeを使うようになったため、torch.floatを指定しても、先に精度が低下する可能性がある。その後、新しい値を既存のBF16バッファーにコピーしても、元の精度は復元できない。修正の説明

報告では、head dimensionを256、RoPE thetaを1,000万に設定し、FP64の参照値と比較した。位置65,536におけるLlamaの最大未剰余角度誤差は、ZeRO-3で117.7ラジアンだったのに対し、比較対象の経路では0.003ラジアンだった。これは特定の設定における数値誤差であり、精度（accuracy）の低下にそのまま換算できるものではない。報告者はまた、テストで使ったYaRNの経路には影響がなかったと述べている。再現手順と測定結果

DeepSpeedの公式ドキュメントによると、ZeRO-3はモデルパラメーターを分割してメモリーを節約し、大規模モデルの構築にはzero.Initを利用できる。そのため、初期化段階の精度処理は確認すべきポイントとなる。エンジニアリング上の推論として、重みの精度を確認するだけでは、位置エンコーディングのバッファーに生じる数値差を排除するには不十分だ。DeepSpeedドキュメント

修正案では、int64の系列を先に作成してから浮動小数点数に変換する方式に戻しており、203ファイルが対象となっている。報告者は、修正後にFP32へ戻ったことと、ローカルでのテスト結果を示した。確認時点ではPRはオープンのままで、レビュー承認は得られていない。長文脈のファインチューニングを行うチームは、メンテナーによる再現確認と正式リリースへの取り込み状況に加え、自分たちのモデル読み込み後のinv_freq精度や学習結果の比較を確認すべきだ。この測定結果をすべてのRoPEモデルに一般化しないよう注意が必要である。修正案