Transformers、CUDA Graphを妨げるホストテンソルのコピーを除去するGemma 3n向け修正を提案

Transformersコミュニティは11月7日、Gemma 3nテキストモデルのAltUpループを変更し、CUDA Graphのキャプチャを妨げる問題を解消するPR #50457を提出した。11月8日時点で提案はオープンのままであり、正式リリースで修正済みとはまだ言えない。修正提案

問題の報告では、Transformers 5.23.0.dev0、PyTorch 2.11.0、NVIDIA L4が使われた。小型モデルのフォワード演算はeagerモードで実行できる一方、torch.cuda.graph内ではCPUテンソルとCUDAテンソル間のコピーが原因で失敗した。報告者によると、vLLMのTransformersモデルバックエンドでも、コンパイルを有効にするとGemma-3n-E2B-itのグラフキャプチャエラーが発生し、当時はenforce_eager=Trueを指定して実行する必要があった。再現ケース

修正は、振幅の下限を設定する2か所の演算に絞られている。従来はtorch.maximumで計算結果とCPU上のepsilonテンソルを比較し、ループ内で後者を実行先デバイスへ移していた。提案ではこれをnew_magnitude.clamp(min=1e-5)に置き換える。報告者によると、この変更によりテストでフォワード演算のたびに発生していた6回の同期コピーがなくなり、CPUとCUDAの3種類の浮動小数点精度で同じ出力が保たれた。実装と検証の説明

この変更の意義は、推論の実行経路を復旧できる点にある。PyTorchのドキュメントによると、CUDA GraphはGPU処理をあらかじめ記録し、1回の起動で再生することで、Python、C++、ドライバー層のディスパッチコストを減らせる。ただし、キャプチャ中はCPUとGPUの同期を伴う操作が禁止される。そのため、数値安定性のための小さな定数でも、デバイスの扱い方によってはグラフ全体のキャプチャを妨げることがある。CUDA Graphの仕組みと制約

デプロイ環境では、まずTransformersモデルバックエンドを使用しているか確認するとよい。vLLMは--model-impl transformersオプションを提供しており、読み込んだモデルの型も確認できる。今後は提案のマージ、パッケージへの収録、実際のバッチサイズやシーケンス長でのレイテンシ測定を追う必要がある。報告されたテスト結果は修正の方向性を支持するが、スループットのベンチマークは示されていないため、サービスの高速化幅を推定することはできない。vLLMバックエンドのドキュメント