Transformers 社群回報 Gemma 張量平行缺陷，混用 Tensor 與 DTensor 可能中斷推論

Transformers 社群於 10 月 12 日提出 Gemma 多模態模型的張量平行相容性回報，隔日補充修補方向。問題涉及共享嵌入權重與 Gemma 4 的逐層輸入嵌入路徑：部分運算同時收到一般 Tensor 與分散式 DTensor，因而直接拋出例外。回報環境為 Transformers 5.18.0.dev0、PyTorch 2.14.0，不能據此推定所有正式版本都受影響。問題回報

兩種失敗發生在不同位置。Gemma 3、4 的輸出層 lm_head 可能以一般隱藏狀態搭配共享的 DTensor 權重進行矩陣運算；Gemma 4 啟用逐層嵌入時，則可能在 torch.where 中混用 DTensor 填補嵌入與一般張量。附上的縮小模型腳本建立雙 rank 的 CPU 模擬程序群組，套用張量平行並重新綁定權重後，只輸入文字 token 就能觸發錯誤，降低了追查問題的硬體門檻。重現腳本與錯誤堆疊

這類缺陷對多卡部署有實際意義。張量平行會將單層權重切分至多個加速器，讓超出單卡記憶體的模型得以執行；Transformers 透過 DeviceMesh 與 DTensor 表達分片、複製及歸約布局。因此，工程上需要檢查模組邊界的輸入輸出布局，以及共享參數重新綁定後的型別是否仍符合運算要求。這是依官方架構對本案提出的檢查方向。張量平行官方文件

回報作者在 10 月 13 日表示，既有 PR #48938 可處理目前模型錯誤，但方案尚未完全通用，仍需調整。本文未能直接核實該 PR 的完整內容及合併狀態，因此不能將作者的局部結果視為正式修復。公開案例也尚未提供完整檢查點、多模態生成或跨硬體的驗證矩陣。作者後續說明

準備擴展 Gemma 推論的團隊，可先固定套件版本，在實際檢查點上驗證共享嵌入與逐層嵌入設定，再測試前向運算及生成流程。後續值得追蹤的是修補是否涵蓋不同張量布局，以及回歸測試能否同時驗證模型載入、權重共享與多卡執行。