TransformersがRoPEの精度修正を提案、ZeRO-3での読み込み時に生じる精度低下に対処

TransformersのメンバーであるCyrilvallezは11月2日、モデルの初期化時にDeepSpeedが回転位置埋め込み（RoPE）の周波数バッファーの精度を下げるのを防ぐ修正を提案しました。これは10月29日の報告を受けた対応です。確認時点でPR #50275はまだオープンで、マージされていません。修正提案

問題はZeRO-3の初期化経路で発生します。報告者によると、deepspeed.zero.Initはtorch.arangeなどのテンソル生成関数をラップし、浮動小数点の結果を学習に使うデータ型へ変換します。TransformersがRoPEモジュールを再構成した後、一部の周波数初期化で浮動小数点数列を使うようになったため、本来float32を保つはずのinv_freqがbf16に変換されていました。その後に値を再計算して既存のバッファーへコピーしても、低精度の結果が残ります。再現コードとバージョン比較

こうした誤差は、特に長いシーケンスで注意が必要です。位置インデックスに周波数を掛けるため、周波数のわずかなずれが大きな回転角の差につながることがあります。著者は、指定されたLlamaの設定で位置65,536をテストし、ZeRO-3経路では最大角度誤差が約117.7ラジアン、比較対象の経路では約0.003ラジアンだったと報告しました。これは特定の構成で行った数値テストであり、モデル品質の低下幅にそのまま換算することはできません。テスト条件

新しい提案は、DeepSpeedがテンソル生成関数をラップする箇所に変更を加えます。報告者はDeepSpeed 0.19.9を使い、CPU上でLlama、Qwen3、Gemma3など6種類のモデルをテストしました。2種類の読み込み方法のいずれでも周波数がfloat32に保たれ、DeepSpeedを使わない場合の結果と一致したとしています。一方、報告者は旧バージョンでは呼び出しインターフェースが異なると指摘しました。著者はその後、互換性対応のコミットを追加しましたが、ページ上では最新バージョンでの完全な再検証は確認できません。テストと互換性に関する議論

公式ドキュメントによると、ZeRO-3は初期化段階からモデルパラメーターを分割するため、読み込み手順もエンジニアリングチームが確認すべき工程です。長コンテキストのファインチューニングでは、修正がマージされるか、旧バージョンのDeepSpeedで検証を通過するか、修正前後で実際の学習結果に差が出るかを引き続き確認する必要があります。現時点の根拠は、主に数値の再現と読み込みテストです。DeepSpeed統合ドキュメント