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

Abierto Apto para principiantes
  • ¿Algo no está bien?
#5,038 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 3 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
2/5
Tiempo estimado
1-3 horas
Aptitud para principiantes
74/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python
Área
compilers, machine-learning

Línea de trabajo

Empieza en el convertidor grouped_mm de torch_tensorrt/dynamo/conversion/impl/moe.py, donde offs se pasa a cast_trt_tensor en converter_utils.py, y revisa cómo enforce_tensor_types maneja el argumento 2. Ejecuta el repro del issue con un buffer constante offs y confirma el TypeError. Está terminado cuando el caso con offsets constantes compila y coincide con PyTorch dentro de la tolerancia indicada, y las pruebas existentes de grouped_mm siguen pasando. Coordínate con el PR #5021 para que las dos correcciones no entren en conflicto.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Python
Estrellas
3k
Forks
421
Merge medio
4 d 7 h
PR fusionados (30 d)
62

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de pytorch/TensorRT

Todos los issues de pytorch/TensorRT

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.