Unable to export quantized models to stablehlo

[Ha-Eyz]

Hello, I am currently trying out torchax, specifically the torchax.export.exported_program_to_stablehlo. While trying out its capabilities, I ran into a problem when trying to export quantized models via the function. As far as I can see it seems that this is simply not supported.

So I wanted to ask if this is not supported yet or a bug. In case of it being a bug, I provide the following code to reproduce my errors:

import torch
import torchax.export
import torch.nn as nn
from torchao.quantization import quantize_
from torchao.quantization.quant_api import Int4WeightOnlyConfig

def llama_mlp_example():
    class LlamaMLP(nn.Module):
        def __init__(self):
            super().__init__()
            self.hidden_size = 16
            self.intermediate_size = 64
            self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
            self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
            self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
            self.act_fn = nn.SiLU()
        def forward(self, x):
            down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
            return down_proj
    m = LlamaMLP().eval()
    example_inputs = (torch.empty(1, 16, dtype=torch.float32),)
    return m, example_inputs

def simple_pt2e_quant_example(m, example_inputs):
    m = torch.export.export(m, example_inputs).module()

    from torchao.quantization.pt2e.quantize_pt2e import (
      prepare_pt2e,
      convert_pt2e,
    )

    from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
      get_symmetric_quantization_config,
      XNNPACKQuantizer,
    )

    quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())
    m = prepare_pt2e(m, quantizer)

    # calibration omitted

    m = convert_pt2e(m)

    exported = torch.export.export(m, example_inputs)

    weights, stablehlo = torchax.export.exported_program_to_stablehlo(exported)

def int4_quant_example(m, example_inputs):
    quant_config = Int4WeightOnlyConfig(
        group_size=32,
    )

    quantize_(m, quant_config)

    exported = torch.export.export(
        m,
        example_inputs,
    )

    weights, stablehlo = torchax.export.exported_program_to_stablehlo(exported)

if __name__ == "__main__":
    llama_mlp_m, llama_mlp_inputs = llama_mlp_example()

    simple_pt2e_quant_example(llama_mlp_m, llama_mlp_inputs)
    # int4_quant_example(llama_mlp_m, llama_mlp_inputs)

The errors that occur are as follows:

AssertionError: quantized_decomposed.dequantize_per_tensor.default

While executing %dequantize_per_tensor : [num_users=1] = call_function[target=torch.ops.quantized_decomposed.dequantize_per_tensor.default](args = (%b__frozen_param0, 1.0, 0, -127, 127, torch.int8), kwargs = {})
Original traceback:
File "/home/Ha-Eyz/testing/issueExample.py", line 19, in forward
    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
  File "/home/Ha-Eyz/roof-mlir/.venv/lib/python3.11/site-packages/torch/nn/modules/linear.py", line 134, in forward
    return F.linear(input, self.weight, self.bias)
Use tlparse to see full graph. (https://github.com/pytorch/tlparse?tab=readme-ov-file#tlparse-parse-structured-pt2-logs)
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

If you run the simple_pt2e_quant_example(...) function

AssertionError: fbgemm.bf16i4bf16_rowwise.default

While executing %bf16i4bf16_rowwise : [num_users=1] = call_function[target=torch.ops.fbgemm.bf16i4bf16_rowwise.default](args = (%view, %p_down_proj_parametrizations_weight_original0, %p_down_proj_parametrizations_weight_original1, %p_down_proj_parametrizations_weight_original2), kwargs = {})
Original traceback:
File "/home/Ha-Eyz/testing/issueExample.py", line 19, in forward
    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
  File "/home/Ha-Eyz/roof-mlir/.venv/lib/python3.11/site-packages/torch/nn/modules/linear.py", line 134, in forward
    return F.linear(input, self.weight, self.bias)
Use tlparse to see full graph. (https://github.com/pytorch/tlparse?tab=readme-ov-file#tlparse-parse-structured-pt2-logs)
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

If you run the int4_quant_example(...) function

For the full error messages you can run the code I provided above.

[qihqi]

hi @Ha-Eyz

So the issue is that this operator (fbgemm.bf16i4bf16_rowwise.default) is not a regular torch op, and torchax currently only implemented mappings from torch (ATen) ops -> jax. i.e. those included in torch.ops.aten.*.

Now, to export a model with these ops, you have provide either the definition of bf16i4bf16_rowwise in terms of JAX ops, OR, in terms of torch Aten ops. (see https://github.com/google/torchax/blob/main/test/test_libraries.py#L54)

Another option is to write your own Quantizer instead of the XNNPackQuantizer, your own Quantizer should swap float Linear with a int Linear written using plain torch aten ops (instead of calling out to fbgemm).

[Ha-Eyz]

Hi @qihqi

thanks for the reply. I would like to ask if the torchax team has any plans of supporting operations used in quantization, specifically "quantized_decomposed.dequantize_per_tensor.default".