`aten._grouped_mm` converter raises `TypeError` when the offsets are a constant

オープン 初心者向け
  • 何か問題がありますか？
#5,038 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 3 日以内に返信

まだ誰も着手していません。

評価

難易度
2/5
見積もり時間
1〜3時間
初心者へのやさしさ
74/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
python
領域
compilers, machine-learning

調査の方向性

torch_tensorrt/dynamo/conversion/impl/moe.py にある grouped_mm コンバーターから始めてください。ここでは offs が converter_utils.py の cast_trt_tensor に渡されています。そのうえで enforce_tensor_types が引数 2 をどう扱っているかを確認してください。issue の再現手順を定数の offs バッファを使って実行し、TypeError が発生することを確認してください。定数オフセットのケースがコンパイルでき、記載された許容誤差の範囲で PyTorch と一致し、既存の grouped_mm テストがすべて通れば完了です。両方の修正が衝突しないよう、PR #5021 と調整してください。

索引モデルが issue の本文から書いたものです。

説明

What is wrong

In mixture-of-experts models the group offsets of torch._grouped_mm are sometimes fixed, for example a buffer that says how many rows each expert gets. After freezing, the offsets reach the converter as a constant torch.Tensor, not a TensorRT tensor. The converter only forces its first 2 arguments to TensorRT tensors (enforce_tensor_types({0: ..., 1: ...})), then passes offs to cast_trt_tensor, which expects an ITensor. The conversion fails with a TypeError inside that cast, so the grouped matmul does not run in TensorRT.

What this blocks

cosmos3 (pytorch/executorch#24353) hits this: its grouped matmul stays out of TensorRT, so the graph is split around it into more TensorRT engines and CUDA backend programs.
A fix keeps the grouped matmul inside the engine, so mixture-of-experts layers with fixed groups run fully in TensorRT.

Repro
import torch
import torch_tensorrt


class FixedGroups(torch.nn.Module):
    def __init__(self, groups=4, k=64, n=32, rows_per_group=8):
        super().__init__()
        self.w = torch.nn.Parameter(torch.randn(groups, k, n))
        ends = torch.arange(1, groups + 1, dtype=torch.int32) * rows_per_group
        self.register_buffer("offs", ends)

    def forward(self, x):
        return torch._grouped_mm(x, self.w, offs=self.offs)


model = FixedGroups().cuda().bfloat16().eval()
x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
ep = torch.export.export(model, (x,))
trt_gm = torch_tensorrt.dynamo.compile(ep, arg_inputs=[x], min_block_size=1)
torch.testing.assert_close(trt_gm(x), model(x), rtol=2e-2, atol=2e-2)
Expected

The grouped matmul converts and matches PyTorch.

Actual

dynamo.compile raises inside the converter (end of the traceback):

  File ".../torch_tensorrt/dynamo/conversion/impl/moe.py", line 93, in grouped_mm
    offs_i32 = cast_trt_tensor(
  File ".../torch_tensorrt/dynamo/conversion/converter_utils.py", line 276, in cast_trt_tensor
    cast_layer = ctx.net.add_cast(input_val, trt_dtype)
TypeError: add_cast(): incompatible function arguments. The following argument types are supported:
    1. (self: tensorrt_bindings.tensorrt.INetworkDefinition, input: tensorrt_bindings.tensorrt.ITensor, to_type: tensorrt_bindings.tensorrt.DataType) -> tensorrt_bindings.tensorrt.ICastLayer
Invoked with: <tensorrt_bindings.tensorrt.INetworkDefinition object at 0x7f4dc10b3d30>, tensor([ 8, 16, 24, 32], dtype=torch.int32), <DataType.INT32: 3>
While executing %_grouped_mm : [num_users=1] = call_function[target=torch.ops.aten._grouped_mm.default](args = (%x, %w, %offs), kwargs = {})
Versions

torch 2.15.0 (CUDA 13.2), torch_tensorrt 2.15.0, TensorRT 11.4.0.106, H100, driver 580.126.09.

What was tried

Reading the converter. A likely fix is to add 2: (TRTTensor,) to the converter's enforce_tensor_types, or to call get_trt_tensor on offs before the cast. The new grouped GEMM kernel path in the open PR #5021 already does the second; the default dense path does not.

主要言語
Python
スター
3k
フォーク
421
平均マージ
4日 7時間
マージ済み PR（30日）
62

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

pytorch/TensorRT のほかの issue

pytorch/TensorRT の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。