Alpamayo quantized model doesn't run and cannot be converted to TensorRT

Hello,

Description

examples/alpamayo/quantize.py doesn't run the quantized model and doesn't convert the quantized model to TensorRT engine. Once conversion is run:

python quantize.py --ckpt nvidia/Alpamayo-R1-10B --output-dir / --quantize fp8 --limit 10

Traceback (most recent call last):
  File "/Model-Optimizer/examples/alpamayo/quantize.py", line 705, in <module>
    main()
  File "/Model-Optimizer/examples/alpamayo/quantize.py", line 655, in main
    compiled_model = torch_tensorrt.compile(
                     ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/_compile.py", line 286, in compile
    exp_program = dynamo_trace(
                  ^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/_tracer.py", line 83, in trace
    exp_program = export(
                  ^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/__init__.py", line 319, in export
    raise e
  File "/usr/local/lib/python3.12/dist-packages/torch/export/__init__.py", line 286, in export
    return _export(
           ^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1159, in wrapper
    raise e
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1125, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/exported_program.py", line 123, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 2172, in _export
    ep = _export_for_training(
         ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1159, in wrapper
    raise e
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1125, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/exported_program.py", line 123, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 2033, in _export_for_training
    export_artifact = export_func(
                      ^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1975, in _non_strict_export
    aten_export_artifact = _to_aten_func(  # type: ignore[operator]
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1760, in _export_to_aten_ir_make_fx
    gm, graph_signature = transform(_make_fx_helper)(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1901, in _aot_export_non_strict
    gm, sig = aot_export(wrapped_mod, args, kwargs=kwargs, **flags)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1679, in _make_fx_helper
    gm = make_fx(
         ^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 2295, in wrapped
    return make_fx_tracer.trace(f, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 2233, in trace
    return self._trace_inner(f, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 2204, in _trace_inner
    t = dispatch_trace(
        ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 51, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 893, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1221, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1792, in trace
    res = super().trace(root, concrete_args)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 893, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 850, in trace
    (self.create_arg(fn(*args)),),
                     ^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1279, in wrapped
    out = f(*tensors)  # type:ignore[call-arg]
          ^^^^^^^^^^^
  File "<string>", line 1, in <lambda>
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1583, in wrapped_fn
    return tuple(flat_fn(*args))
                 ^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/utils.py", line 184, in flat_fn
    tree_out = fn(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 906, in functional_call
    out = mod(*args[params_len:], **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1862, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 542, in call_module
    ret_val = forward(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 818, in forward
    return _orig_module_call(mod, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1885, in forward
    tree_out = mod(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1862, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 542, in call_module
    ret_val = forward(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 818, in forward
    return _orig_module_call(mod, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1767, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1778, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/alpamayo_r1/src/alpamayo_r1/models/alpamayo_r1.py", line 238, in forward
    vlm_outputs = self.vlm.generate(
                  ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/generation/utils.py", line 2564, in generate
    result = decoding_method(
             ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/generation/utils.py", line 2781, in _sample
    model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/models/qwen3_vl/modeling_qwen3_vl.py", line 1428, in prepare_inputs_for_generation
    if cache_position[0] != 0:
       ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1327, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1374, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_export/non_strict_utils.py", line 976, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py", line 536, in guard_bool
    r = self.evaluate()
        ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py", line 510, in evaluate
    return self.shape_env.evaluate_sym_node(self, size_oblivious)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/symbolic_shapes.py", line 7191, in evaluate_sym_node
    return self.evaluate_expr(
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/symbolic_shapes.py", line 7291, in evaluate_expr
    return self._inner_evaluate_expr(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/recording.py", line 272, in wrapper
    return retlog(fn(*args, **kwargs))
                  ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/symbolic_shapes.py", line 7314, in _inner_evaluate_expr
    return self._evaluate_expr(
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/symbolic_shapes.py", line 7538, in _evaluate_expr
    raise self._make_data_dependent_error(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not guard on data-dependent expression Eq(u3, 1) (unhinted: Eq(u3, 1)).  (Size-like symbols: none)

Caused by: (_export/non_strict_utils.py:976 in __torch_function__)
For more information, run with TORCH_LOGS="dynamic"
For extended logs when we create symbols, also add TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="u3"
If you suspect the guard was triggered from C++, add TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
For more debugging help, see https://docs.google.com/document/d/1HSuTTVvYH1pTew89Rtpeu84Ht3nQEFTYhAX3Ypa_xJs/edit?usp=sharing

For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1

The following call raised this error:
  File "/usr/local/lib/python3.12/dist-packages/transformers/models/qwen3_vl/modeling_qwen3_vl.py", line 1428, in prepare_inputs_for_generation
    if cache_position[0] != 0:

The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.

Process finished with exit code 1

Steps/Code to reproduce bug

• add code for model to have forward function instead, alpamayo_r1/src/alpamayo_r1/models/alpamayo_r1.py line 150

# def sample_trajectories_from_data_with_vlm_rollout(
def forward(
        self,
        # data: dict[str, Any],
        tokenized_data_input_ids: torch.Tensor,
        tokenized_data_attention_mask: torch.Tensor,
        tokenized_data_pixel_values: torch.Tensor,
        tokenized_data_image_grid_thw: torch.Tensor,
        ego_history_xyz: torch.Tensor,
        ego_history_rot: torch.Tensor,
        top_p: float = 0.98,
        top_k: int | None = None,
        temperature: float = 0.6,
        num_traj_samples: int = 6,
        num_traj_sets: int = 1,
        diffusion_kwargs: dict[str, Any] | None = None,
        max_generation_length: int = 256,
        ...
        data = {
            "tokenized_data": {
                "input_ids": tokenized_data_input_ids,
                "attention_mask": tokenized_data_attention_mask,
                "pixel_values": tokenized_data_pixel_values,
                "image_grid_thw": tokenized_data_image_grid_thw,
            },
            "ego_history_xyz": ego_history_xyz,
            "ego_history_rot": ego_history_rot,
        }
        ...
        # max_generation_length = kwargs.get(
        #     "max_generation_length", self.config.tokens_per_future_traj
        # )

• add code for save one input sample for compile, examples/alpamayo/quantize.py line 264

input = None
def make_joint_calibration_forward_loop(
    def _calibration_loop(runtime_model):
               ...
                global input
                input = copy.deepcopy(model_inputs)

                input_ids = input['tokenized_data']['input_ids']
                attention_mask = input['tokenized_data']['attention_mask']
                pixel_values = input['tokenized_data']['pixel_values']
                image_grid_thw = input['tokenized_data']['image_grid_thw']
                ego_history_xyz = input['ego_history_xyz']
                ego_history_rot = input['ego_history_rot']

                with torch.autocast("cuda", dtype=torch.float16):
                    # runtime_model.sample_trajectories_from_data_with_vlm_rollout(
                    runtime_model(
                        # data=model_inputs,
                        input_ids,
                        attention_mask,
                        pixel_values,
                        image_grid_thw,
                        ego_history_xyz,
                        ego_history_rot,

• add code to compile quantized model, examples/alpamayo/quantize.py line 640

    input_ids = input['tokenized_data']['input_ids']
    attention_mask = input['tokenized_data']['attention_mask']
    pixel_values = input['tokenized_data']['pixel_values']
    image_grid_thw = input['tokenized_data']['image_grid_thw']
    ego_history_xyz = input['ego_history_xyz']
    ego_history_rot = input['ego_history_rot']

    with torch.autocast("cuda", dtype=torch.float16):
        # compiled_model = torch.compile(model, backend="torch_tensorrt")
        # output = compiled_model(input_ids, attention_mask, pixel_values, image_grid_thw, ego_history_xyz, ego_history_rot)

        from modelopt.torch.quantization.utils import export_torch_mode
        with export_torch_mode():
            compiled_model = torch_tensorrt.compile(
                model,
                ir="dynamo",
                min_block_size=1,
                truncate_double=True,
                inputs=[
                    torch_tensorrt.Input(min_shape=input_ids.shape, opt_shape=input_ids.shape, max_shape=input_ids.shape, dtype=input_ids.dtype, ),
                    torch_tensorrt.Input(min_shape=attention_mask.shape, opt_shape=attention_mask.shape, max_shape=attention_mask.shape, dtype=attention_mask.dtype, ),
                    torch_tensorrt.Input(min_shape=pixel_values.shape, opt_shape=pixel_values.shape, max_shape=pixel_values.shape, dtype=pixel_values.dtype, ),
                    torch_tensorrt.Input(min_shape=image_grid_thw.shape, opt_shape=image_grid_thw.shape, max_shape=image_grid_thw.shape, dtype=image_grid_thw.dtype, ),
                    torch_tensorrt.Input(min_shape=ego_history_xyz.shape, opt_shape=ego_history_xyz.shape, max_shape=ego_history_xyz.shape, dtype=ego_history_xyz.dtype, ),
                    torch_tensorrt.Input(min_shape=ego_history_rot.shape, opt_shape=ego_history_rot.shape, max_shape=ego_history_rot.shape, dtype=ego_history_rot.dtype, ),
                ]
            )
            output = compiled_model(input_ids, attention_mask, pixel_values, image_grid_thw, ego_history_xyz, ego_history_rot)

Expected behavior

Expected to be able to run quantized model with same inputs as eager.
Expected that quantized model can be converted to TensoRT engine and run.

System information

• OS: Ubuntu 24.04
• CPU architecture: aarch64
• GPU name: RTX PRO 6000
• GPU memory size: 96
• Number of GPUs: 1
• Library versions (if applicable):
  • Python: 3.12
  • ModelOpt version: main
  • CUDA: 13
  • PyTorch: 2.8
  • Transformers: 4.57
  • ONNXRuntime: 1.18/1.29
  • TensorRT: 10.13
• Other details that may help:
  • Eager model runs and calibration runs until conversion.
  • Calibration and conversion does not run out of memory or storage space.

Thanks.