Transformers、Whisper と BART のコンパイル修正を提案—LayerDrop 無効時のランダム分岐を除去

Transformers コミュニティは 11 月 8 日、Whisper と BART のエンコーダーおよびデコーダーにおける LayerDrop の判定を調整する PR #50475 を提出しました。問題は、レイヤーをスキップする確率がゼロに設定されていても、訓練経路でランダムテンソルを生成してその値に基づく分岐が実行され、torch.compile(fullgraph=True) でトレースを完了できないことです。確認時点で提案は未完了で、CI の実行にはメンテナーの承認が必要です。修正提案

前日に報告された問題では、Transformers 5.22.2、PyTorch 2.15.0、macOS ARM64 の環境で、小型の Whisper とランダム入力を使い、backend="eager" でエラーを再現しています。この設定は主に Dynamo のフロントエンドにおけるトレースの問題を切り分けるもので、GPU カーネルやコンパイル後の速度を判断することはできません。報告者は BART でも同じエラーを再現し、ローカルで条件チェックを追加すると再現スクリプトが成功したと述べています。問題と再現スクリプト

鍵となるのは、Python の分岐条件がテンソルの値に依存している点です。元のコードはまず torch.rand([]) を呼び出し、その結果がレイヤーをスキップする確率より小さいかを比較していました。確率がゼロでも、トレーサーはデータに依存する制御フローに遭遇します。提案では、先に self.training and self.layerdrop > 0.0 を確認し、LayerDrop が実際に有効な場合にのみサンプリングします。また、両モデルに訓練時の全グラフコンパイルを対象とする回帰テストを追加しています。変更の説明

ファインチューニングでは、この変更によってデフォルト設定でのコンパイル上の障害が一つ取り除かれる可能性があります。PyTorch の公式ドキュメントによると、通常のコンパイルモードでは計算グラフの中断が許容されますが、全グラフモードでは単一の FX グラフの取得が求められ、中断があるとエラーになります。このモードは、ライブラリのコンパイル互換性を確認する用途にも適しています。そのため、全グラフのトレースに成功することにはエンジニアリング上の意義がありますが、それだけでスループットの向上を示すものではありません。PyTorch の全グラフモードのドキュメント

今後は二つの制約に注意が必要です。レイヤーをスキップする確率がゼロより大きい場合、ランダムなスキップによる動的分岐は引き続き発生します。また、不要なサンプリングを省くことで CPU の乱数生成器の状態遷移が変わり、固定シードを使った場合、後続の処理結果が旧バージョンと異なる可能性があります。エンジニアリングチームはマージとリリースの状況を追跡し、実際の訓練構成で勾配、再現性、性能を検証する必要があります。制約と乱数への影響