Transformers、テンソル並列処理の修正を提案　勾配フックによるコンパイル推論の遅延を回避

Hugging Face Transformersのテンソル並列処理の実装に、新たな性能改善案が提出された。プロジェクトメンバーは10月26日、PR #49975を提出し、ReplicatedWithGradAllReduceで勾配同期フックを登録するタイミングを変更した。勾配を必要としない推論時に、torch.compileを妨げる可能性のあるモジュールレベルのバックワードフックが残るのを防ぐ狙いだ。確認時点で、ページには提案がマージ待ちのキューに入っていると表示されており、まだマージされていない。修正提案

元の問題は10月23日に報告された。報告には、CPUプロセス2つで実行する小規模なランダムQwen3の再現ケースが2つ添えられていた。tp_plan="auto"を有効にすると、Dynamoは計算グラフの中断を11回記録したが、バックワードフックを外すと0回になった。報告者によると、AWS Trainium上でQwen3-8BをNeuronコア2つに分散した場合、フックを残すとデコード時間は1トークンあたり140ミリ秒、外すと37.6ミリ秒だった。これは特定の環境で報告された数値であり、他のハードウェアやモデルにそのまま当てはめることはできない。問題と再現スクリプト

このフックは、もともと学習用のものだ。公式ドキュメントによると、この戦略ではパラメーターを複製し、勾配に対してall-reduceを実行する。異なる分割層の間にある正規化演算のうち、各プロセスが一部のアテンションヘッドしか参照できない場合に適している。問題は、torch.no_grad()を使っていても、モジュールに登録済みのバックワードフックによってコンパイラーが計算グラフを中断する可能性があることだ。そのため、勾配計算を無効にするだけでは、追加コストを取り除けない場合がある。テンソル並列処理のドキュメント

提案では、フォワード実行時にtorch.is_grad_enabled()を確認し、勾配が必要な場合に限り、実際に読み込まれたパラメーターへ勾配累積後のフックを登録する。登録を遅らせることで、重みの読み込み時にプレースホルダーパラメーターが置き換えられ、フックも失われる問題を避ける。作者はさらに、旧来のフックでは実行時にパラメーターの勾配がまだ生成されておらず、同期がスキップされる可能性にも対処できると説明している。実装の説明

デプロイチームにとって、この変更案は検証可能な性能ボトルネックを示している。テンソル並列処理とコンパイルを併用する場合、デバイス数だけを比較するのではなく、計算グラフが中断される原因も確認するとよい。今後はマージとリリースの状況を追い、実際のモデルでデコード遅延と学習時の勾配同期を検証する必要がある。公開された提案にはテストスクリプトが含まれるが、一般的な高速化の幅を示すのに十分な結果はまだない。