Unable to quantize pytorch model using huggingface export

[ishan-modi]

I am trying to quantize GPT2 (pytorch) model directly without any onnx conversions, following is my code and traceback of the error I am getting.

Code

from transformers import GPT2Tokenizer, GPT2Model
import modelopt.torch.quantization as mtq
from modelopt.torch.export.unified_export_hf import export_hf_checkpoint

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2Model.from_pretrained('gpt2').to('cuda')

config = {
    "quant_cfg": {
        "*weight_quantizer": {"num_bits": 8, "enable": True},
        "*lm_head*": {"enable": False},
        "default": {"enable": False},
    },
    "algorithm": None,
}

mq = mtq.quantize(model, config)

export_hf_checkpoint(mq, save_modelopt_state=True)

Error

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-9-d1815c270de0> in <cell line: 3>()
      1 from modelopt.torch.export.unified_export_hf import export_hf_checkpoint
      2 
----> 3 export_hf_checkpoint(mq, save_modelopt_state=True)

/usr/local/lib/python3.10/dist-packages/modelopt/torch/export/unified_export_hf.py in export_hf_checkpoint(model, dtype, export_dir, save_modelopt_state)
    379             " torch.save for further inspection."
    380         )
--> 381         raise e

/usr/local/lib/python3.10/dist-packages/modelopt/torch/export/unified_export_hf.py in export_hf_checkpoint(model, dtype, export_dir, save_modelopt_state)
    342     export_dir.mkdir(parents=True, exist_ok=True)
    343     try:
--> 344         post_state_dict, hf_quant_config, per_layer_quantization = _export_hf_checkpoint(
    345             model, dtype
    346         )

/usr/local/lib/python3.10/dist-packages/modelopt/torch/export/unified_export_hf.py in _export_hf_checkpoint(model, dtype)
    148     layer_pool = {
    149         f"model.layers.{name}": sub_module
--> 150         for name, sub_module in model.model.layers.named_modules()
    151     }
    152     # NOTE: Speculative decoding models have extra modules that may be quantized

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in __getattr__(self, name)
   1929             if name in modules:
   1930                 return modules[name]
-> 1931         raise AttributeError(
   1932             f"'{type(self).__name__}' object has no attribute '{name}'"
   1933         )

AttributeError: 'GPT2Model' object has no attribute 'model'

Please let me know if I am missing something here, and if at all things can work without ONNX conversion @jingyu-ml @kevalmorabia97

[cjluo-nv]

Thanks for submitting this bug. I guess we didn't focus on GPT2 testing with our latest versions as the model is pretty old. Could you share your context about working on GPT2? CC @Edwardf0t1

[Edwardf0t1]

@ishan-modi Thanks for reporting the issue.

Alternative to the GPT2 model, you can use https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct which has similar size.

Also, our unified HF export API export_hf_checkpoint does not support int8. Could you try fp8? e.g., *weight_quantizer": {"num_bits": (4, 3), "axis": None}

[ishan-modi]

@Edwardf0t1, thanks for the response. I was just running quick experiments on the package so randomly picked model. I tried using the suggested model but ran into the errors while compiling using tensorrt

Code

from transformers import AutoModelForCausalLM, AutoTokenizer
import modelopt.torch.quantization as mtq
import torch
import torch_tensorrt

checkpoint = "HuggingFaceTB/SmolLM2-135M"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint).to("cuda")

config = mtq.FP8_DEFAULT_CFG
mq = mtq.quantize(model, config)

with torch.no_grad():
    trt_model = torch.compile(
        mq.eval().to('cuda'),
        backend="tensorrt"
    )
    trt_model(torch.randint(1,45000, (1,10), device="cuda"))
    print(model.get_memory_footprint()/1e9, trt_model.get_memory_footprint()/1e9)

The error is likely due to input_quantizer not working properly, I am a bit new to tensorrt let me know if I am making a mistake somewhere

err.txt

[meenchen]

Hi @ishan-modi , currently our quantize op is not compatible with torch.compile. To deploy quantized models with torch, you could try our AutoDeploy example: https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/llm_autodeploy. We also have plan to enable torch.compile for the quantize op.

[ishan-modi]

@meenchen, thanks for the response, I will definitely give it a try. Whilst I was going through other examples on llms, I came across LLM_PTQ: https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/llm_ptq.

It seems there is no real quantization happening here, generally fake quantization is used for training and calibration during inference.

• Shouldn't we use real quantization configs in these examples because it is PTQ.
• Also shouldn't we enable lower precision arithmetic (currently even though real_quant is there, every operation is still in float32), due to this activation quantization for inference is not much useful.

Please confirm if it is already there in the roadmap, sorry for the back and forth !

[cjluo-nv]

Thanks @ishan-modi for the good suggestions. It is indeed in our roadmap. And we hope lower precision compute and help us accelerate PTQ, QAT and eval process.

[ishan-modi]

Awesome, thanks ! closing this for now since already in roadmap