torch_tensorrt.compile() fails with "fake tensor in exported program constant's list" for ModelOpt quantized models unless wrapped in export_torch_mode()

[bradshook]

Bug Description

When compiling a ModelOpt-quantized model with Torch-TensorRT, torch_tensorrt.compile() fails unless it is wrapped in export_torch_mode().

Without export_torch_mode(), FP8 (mtq.FP8_DEFAULT_CFG) ) quantized models fail during export with the following error:

RuntimeError: We found a fake tensor in the exported program constant's list.
This typically means our tracing system encountered an op that we can't trace through.
For the potential source, you can refer to following model attribute:
net.0.input_quantizer.lifted_tensor_0.
Please file an issue on github.

Wrapping the compile call in:

with export_torch_mode():
    trt_model = torch_tensorrt.compile(...)

avoids this error.

To Reproduce

1. Launch the NVIDIA PyTorch container:

   docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.06-py3

2. Quantize a simple model using ModelOpt (mtq.FP8_DEFAULT_CFG).

3. Compile using Torch-TensorRT without export_torch_mode():

import torch
import torch.nn as nn
import torch_tensorrt

import modelopt.torch.quantization as mtq

device = "cuda"

class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(512, 4096),
            nn.GELU(),
            nn.Linear(4096,4096),
            nn.GELU(),
            nn.Linear(4096,512) 
        )

    def forward(self, x):
        return self.net(x)

model = SimpleMLP().to(device).eval()

quant_cfg = mtq.FP8_DEFAULT_CFG  

calib_data = [
    (torch.randn(128, 512, device=device),)
    for _ in range(10)
]

def forward_loop(model):
    with torch.no_grad():
        for (x,) in calib_data:
            model(x)

model = mtq.quantize(
    model,
    quant_cfg,
    forward_loop,
)

inputs = [
    torch_tensorrt.Input(
        shape=(128, 512),
        dtype=torch.float32,
    )
]

trt_model = torch_tensorrt.compile(
    model,
    ir="dynamo",
    inputs=inputs,
)

This produces:

RuntimeError: We found a fake tensor in the exported program constant's list.
This typically means our tracing system encountered an op that we can't trace through.
For the potential source, you can refer to following model attribute:
net.0.input_quantizer.lifted_tensor_0.
Please file an issue on github.

Expected behavior

torch_tensorrt.compile() should successfully export the quantized model, or emit a more descriptive error if export_torch_mode() is required.

Currently, the error does not indicate that wrapping the compilation in export_torch_mode() resolves the issue.

Environment

• Torch-TensorRT Version: 2.13.0a0
• PyTorch Version: 2.13.0a0+8145d630e8.nv26.6.54250401
• CPU Architecture: aarch64
• OS: Ubuntu 24.04.4
• How you installed PyTorch: NVIDIA NGC container (nvcr.io/nvidia/pytorch:26.06-py3)
• Build command you used (if compiling from source): N/A
• Are you using local sources or building from archives: Archive/container image
• Python version: 3.12.3
• CUDA version: 13.3
• GPU models and configuration: NVIDIA Jetson AGX Thor (Blackwell)
• Any other relevant information:
  • TensorRT 11.0.0.114
  • NVIDIA ModelOpt 0.44.0
  • cuDNN 9.23.0

Additional context

This bug showed up when I was trying to figure out this issue related to compiling an FP4 quantized model:
https://forums.developer.nvidia.com/t/fp4-quantization-fails-on-jetson-thor-with-tensorrt-11-0-0-114/377887

[zewenli98]

Hi @bradshook Thanks for bring this to our attention. For now, we support two ways to deal with ModelOpt quantized models: AOT and JIT. torch_tensorrt.compile is not supported yet but it should be. I'll open an PR to do it.

The root cause here is that torch.export.export requires export_torch_mode() wrapper. The error is output by torch:

File "/home/anaconda3/envs/py312/lib/python3.12/site-packages/torch/export/_trace.py", line 2427, in _export_for_training
    raise RuntimeError(error_msg)
RuntimeError: We found a fake tensor in the exported program constant's list. This typically means our tracing system encountered an op that we can't trace through. For the potential source, you can refer to following model attribute: net.0.input_quantizer.lifted_tensor_0. Please file an issue on github. 

[zewenli98]

@bradshook With the PR above, your code should work properly.

[bradshook]

@zewenli98 Yes, that PR fixes the issue, thank you!