TransformersコミュニティがGemmaのテンソル並列処理の不具合を報告、TensorとDTensorの混在で推論が中断する可能性

Transformersコミュニティは10月12日、Gemmaマルチモーダルモデルのテンソル並列処理との互換性に関する問題を報告し、翌日に修正の方向性を補足しました。この問題は、共有埋め込み重みとGemma 4の層ごとの入力埋め込み経路に関係しています。一部の演算に通常のTensorと分散型のDTensorが同時に渡され、例外が発生します。報告された環境はTransformers 5.18.0.dev0、PyTorch 2.14.0であり、これだけを根拠にすべての正式リリースが影響を受けるとは判断できません。問題報告

2種類の失敗は異なる箇所で発生します。Gemma 3および4の出力層lm_headでは、通常の隠れ状態と共有されたDTensorの重みを組み合わせて行列演算する可能性があります。Gemma 4で層ごとの埋め込みを有効にした場合は、torch.where内でDTensorによる埋め込み補完と通常のテンソルが混在する可能性があります。添付された縮小モデルのスクリプトは、2ランクのCPUシミュレーションプロセスグループを作成し、テンソル並列処理を適用して重みを再バインドします。テキストトークンだけを入力すればエラーを再現できるため、問題を追跡する際のハードウェア要件が下がります。再現スクリプトとエラースタック

この種の不具合は、複数GPUへのデプロイに実質的な影響を及ぼします。テンソル並列処理では、各層の重みを複数のアクセラレーターに分割することで、単一GPUのメモリに収まらないモデルを実行できます。TransformersはDeviceMeshとDTensorを使って、シャーディング、複製、リダクションのレイアウトを表現します。そのため、エンジニアリング上はモジュール境界における入出力レイアウトと、共有パラメーターの再バインド後も型が演算要件を満たしているかを確認する必要があります。これは公式のアーキテクチャに基づき、本件について導かれる確認方針です。テンソル並列処理の公式ドキュメント

報告者は10月13日、既存のPR #48938で現在のモデルのエラーに対処できるものの、解決策はまだ完全に汎用的ではなく、調整が必要だと述べました。この記事ではPRの全内容やマージ状況を直接確認できていないため、報告者が示した部分的な結果を正式な修正とみなすことはできません。公開された事例には、完全なチェックポイント、マルチモーダル生成、異なるハードウェアを網羅する検証マトリクスもまだ含まれていません。報告者による続報

Gemmaの推論規模を拡大するチームは、まずパッケージのバージョンを固定し、実際のチェックポイントを使って共有埋め込みと層ごとの埋め込み設定を検証したうえで、前向き計算と生成の流れをテストできます。今後は、修正がさまざまなテンソルレイアウトをカバーするか、また回帰テストでモデルの読み込み、重み共有、マルチGPU実行をまとめて検証できるかを追う必要があります。