Unable to quantize HF's RTDETRv2

[Rafael-Cast]

I'm trying to quantize RTDETRv2's weights to fp8 to try out torchao with the following script:

import torch
from transformers import RTDetrV2ForObjectDetection
from torchao.quantization import quantize_, Float8WeightOnlyConfig
from torchao.float8.config import e4m3_dtype

resnet_backbone_variant = "r50vd"

pretrained_model = RTDetrV2ForObjectDetection.from_pretrained(
    f"PekingU/rtdetr_v2_{resnet_backbone_variant}",
    num_labels=80,
    ignore_mismatched_sizes=False
).eval().to('cuda')

def _is_linear(m, fqn):
    return isinstance(m, torch.nn.Linear)

cfg = Float8WeightOnlyConfig(weight_dtype=e4m3_dtype)
quantize_(pretrained_model, cfg, _is_linear)

When running the script I get the following error inside torchao:

Traceback (most recent call last):
  File "project_root/object-detect/src/runnables/debug.py", line 18, in <module>
    quantize_(pretrained_model, cfg, _is_linear)
  File "venv_root/lib/python3.12/site-packages/torchao/quantization/quant_api.py", line 544, in quantize_
    _replace_with_custom_fn_if_matches_filter(
  File "venv_root/lib/python3.12/site-packages/torchao/quantization/quant_api.py", line 220, in _replace_with_custom_fn_if_matches_filter
    new_child = _replace_with_custom_fn_if_matches_filter(
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/quantization/quant_api.py", line 220, in _replace_with_custom_fn_if_matches_filter
    new_child = _replace_with_custom_fn_if_matches_filter(
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/quantization/quant_api.py", line 215, in _replace_with_custom_fn_if_matches_filter
    model = replacement_fn(model, *extra_args)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/quantization/quant_api.py", line 1690, in _float8_weight_only_transform
    new_weight = _float8_weight_only_quant_tensor(module.weight, config)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/quantization/quant_api.py", line 1673, in _float8_weight_only_quant_tensor
    new_weight = Float8Tensor.from_hp(
                 ^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/quantization/quantize_/workflows/float8/float8_tensor.py", line 226, in from_hp
    scale = _choose_scale_float8(
            ^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torch/_ops.py", line 1243, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/utils.py", line 658, in _dispatch__torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torch/_ops.py", line 1243, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/quantization/quant_primitives.py", line 2313, in _choose_scale_float8
    tensor_reshaped = tensor.view(shape_for_reduction)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/utils.py", line 674, in _dispatch__torch_dispatch__
    return cls._ATEN_OP_OR_TORCH_FN_TABLE[cls][func](func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/utils.py", line 491, in wrapper
    return func(f, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv_root/lib/python3.12/site-packages/torchao/quantization/quantize_/workflows/float8/float8_tensor.py", line 580, in _
    scale_shape.append(qdata.shape[i] // self.block_size[i])
                       ~~~~~~~~~~~^^^
IndexError: tuple index out of range

The error happend on the class_embed submodule. Interestingly enough, the model's linear layers's weights are already Float8Tensors which are then passed to the from_hp method, so maybe the error happens when traversing the module structure recursively and somehow visiting the same module twice.
The previous root module, that contains the backbone and encoder-decoder, was quantized successfully

My torchao version is 0.14.1.

[jcaip]

IIRC sometimes the first linear layer is tied to the embeddings for transformers models, maybe this is happening for your specific model architecture? Can you try using the config setup mentioned here

quant_config = Float8Config()
quantization_config = TorchAoConfig(quant_type=quant_config, include_embedding=True, untie_embedding_weights=True, modules_to_not_convert=[])

# either use `untied_model_id` or `untied_model_local_path`
quantized_model = AutoModelForCausalLM.from_pretrained(untied_model_id, dtype=torch.float32, device_map="auto", quantization_config=quantization_config)

[Rafael-Cast]

The model does have tied parameters, and they are not only embeddings but also a shared MLP among different decoder layers that maps the output of the decoder layer to bounding box and class predictions.

I've tried your suggestion but I get the following exception:

File "/project_root/src/runnables/debug_2.py", line 12, in <module>
    pretrained_model = RTDetrV2ForObjectDetection.from_pretrained(
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/env_root/lib/python3.12/site-packages/transformers/modeling_utils.py", line 288, in _wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/env_root/lib/python3.12/site-packages/transformers/modeling_utils.py", line 5179, in from_pretrained
    ) = cls._load_pretrained_model(
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/env_root/lib/python3.12/site-packages/transformers/modeling_utils.py", line 5642, in _load_pretrained_model
    _error_msgs, disk_offload_index, cpu_offload_index = load_shard_file(args)
                                                         ^^^^^^^^^^^^^^^^^^^^^
  File "/env_root/lib/python3.12/site-packages/transformers/modeling_utils.py", line 946, in load_shard_file
    disk_offload_index, cpu_offload_index = _load_state_dict_into_meta_model(
                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/env_root/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/env_root/lib/python3.12/site-packages/transformers/modeling_utils.py", line 854, in _load_state_dict_into_meta_model
    hf_quantizer.create_quantized_param(
  File "/env_root/lib/python3.12/site-packages/transformers/quantizers/quantizer_torchao.py", line 258, in create_quantized_param
    input_embed = model.get_input_embeddings()
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/env_root/lib/python3.12/site-packages/transformers/modeling_utils.py", line 1952, in get_input_embeddings
    raise NotImplementedError(
NotImplementedError: `get_input_embeddings` not auto‑handled for RTDetrV2ForObjectDetection; please override in the subclass.

I didn't have had the time to try implementing those methods. I'll try monkey patching the method ASAP, but I'm not quite sure if it makes sense - the model uses dense embeddings (ie. not a nn.Embedding layer).

What did work was quantizing the `inner' model, by swapping this line:
quantize_(pretrained_model, cfg, _is_linear)
To this:
```quantize_(pretrained_model.model, cfg, _is_linear)```

Is this fine? Given what you said previously about tied parameters I guess that my change doesn't really solve the root cause of the problem. Does that make sense?

[Rafael-Cast]

Hi,

The following snippet fixed my issue, thanks!

import torch
from transformers import RTDetrV2ForObjectDetection, TorchAoConfig
from torchao.quantization import Float8WeightOnlyConfig
from torchao.float8.config import e4m3_dtype
from transformers.modeling_utils import find_tied_parameters

resnet_backbone_variant = "r50vd"
tag = f"PekingU/rtdetr_v2_{resnet_backbone_variant}"
cfg_ = Float8WeightOnlyConfig(weight_dtype=e4m3_dtype)
cfg = TorchAoConfig(quant_type=cfg_, include_embedding=True, untie_embedding_weights=True, modules_to_not_convert=[])

def rtdetrv2_get_input_embeddings(self: RTDetrV2ForObjectDetection):
    pass

RTDetrV2ForObjectDetection.get_input_embeddings = rtdetrv2_get_input_embeddings

pretrained_model = RTDetrV2ForObjectDetection.from_pretrained(
    tag,
    num_labels=80,
    ignore_mismatched_sizes=False,
    dtype=torch.float32, device_map="auto", quantization_config=cfg
).eval().to('cuda')


print("Tied params:", find_tied_parameters(pretrained_model))