Gemma 3 AWQ Quantization doesnt work cache_position is NoneType

[Gman3214]

Describe the bug
i am trying to quantize gemma 3 AWQ, i have managed to quantize it using gptq but when trying to do it in AWQ i am getting an error.

Expected behavior
the model to be compressed

Environment
Include all relevant environment information:

1. OS Ubuntu 20.04
2. Python version 3.11
3. LLM Compressor version latest

code:

import torch
from datasets import load_dataset
from transformers import Gemma3ForConditionalGeneration, AutoTokenizer
from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier, AWQMapping

# Select model and load it
MODEL_ID = "google/gemma-3-4b-it-qat-int4-unquantized"
model = Gemma3ForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype="auto", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

# Select calibration dataset
DATASET_ID = "ram3214/calibration_hebrew"
DATASET_SPLIT = "train[:3000]"
NUM_CALIBRATION_SAMPLES = 36
MAX_SEQUENCE_LENGTH = 128

# Load and preprocess dataset
ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
ds = ds.shuffle(seed=42)

# Gemma 3 specific mappings
_gemma3_mappings = [
    AWQMapping(
        smooth_layer="re:.*input_layernorm$",
        balance_layers=[
            "re:.*self_attn\.q_proj$",
            "re:.*self_attn\.k_proj$",
            "re:.*self_attn\.v_proj$",
        ],
    ),
    AWQMapping(
        smooth_layer="re:.*self_attn\.v_proj$",
        balance_layers=["re:.*self_attn\.o_proj$"],
    ),
    AWQMapping(
        smooth_layer="re:.*post_attention_layernorm$",
        balance_layers=[
            "re:.*mlp\.gate_proj$",
            "re:.*mlp\.up_proj$",
        ],
    ),
    AWQMapping(
        smooth_layer="re:.*mlp\.up_proj$",
        balance_layers=["re:.*mlp\.down_proj$"],
    ),
]

# Configure the quantization algorithm
recipe = [
    AWQModifier(
        ignore=[
            "lm_head",
            "re:model\.vision_tower.*",
            "re:model\.multi_modal_projector.*",
        ],
        mappings=_gemma3_mappings,
        scheme="W4A16_ASYM",
        targets=["Linear"],
    ),
]
def data_collator(batch):
    assert len(batch) == 1
    return {key: torch.tensor(value) for key, value in batch[0].items()}


# Apply algorithms
oneshot(
    model=model,
    dataset=ds,
    recipe=recipe,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    data_collator=data_collator,
)

# Save to disk compressed
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-awq-asym"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)

Errors

[/usr/local/lib/python3.11/dist-packages/transformers/models/gemma3/modeling_gemma3.py](https://localhost:8080/#) in prepare_inputs_for_generation(self, input_ids, past_key_values, inputs_embeds, cache_position, position_ids, pixel_values, attention_mask, token_type_ids, use_cache, logits_to_keep, labels, **kwargs)
   1431         # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
   1432         # Otherwise we need pixel values to be passed to model. NOTE: use_cache=False needs pixel_values always
-> 1433         if cache_position[0] == 0:
   1434             model_inputs["pixel_values"] = pixel_values
   1435         is_training = token_type_ids is not None and labels is not None

TypeError: 'NoneType' object is not subscriptable

Additional context
i have ran to this issue before when i quantized the model in gptq all i had to do is change the ignore to this and it worked :

ignore=[
        "lm_head",
        "re:model\.vision_tower.*",
        "re:model\.multi_modal_projector.*",
    ],

[brian-dellabetta]

Hi @Gman3214 , I am able to run AWQ with Gemma3 using this gemma3_example.py script, just replacing GPTQModifier with AWQModifier (import with from llmcompressor.modifiers.awq import AWQModifier). The default mappings should work for Gemma3.

Can you confirm this runs for you? If so, there must be an issue with the custom dataset you are using. You'll have to compare it with the flickr30k dataset to determine what is off

[Gman3214]

@brian-dellabetta hey thanks for the quick reply ill try it now.

[brian-dellabetta]

hi @Gman3214 , I noticed the Gemma architecture is slightly different from llama, I have added mappings for Gemma 3 on this branch:

• AWQ Cohere, Mistral & Gemma mappings #1570

[Gman3214]

@brian-dellabetta ill check that out I tried using the script and switching the for some reason it stalls on this

Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.
Preprocessing: 100%
 512/512 [00:03<00:00, 312.38 examples/s]
Tokenizing: 100%
 512/512 [00:33<00:00,  1.07 examples/s]
2025-06-23T16:36:10.045274+0000 | reset | INFO - Compression lifecycle reset
2025-06-23T16:36:10.046882+0000 | from_modifiers | INFO - Creating recipe from modifiers
2025-06-23T16:36:15.594498+0000 | _set_resolved_mappings | INFO - Excluded 61 from resolved mappings due to shape mismatch

Ill try using your patch and update you, thanks allot for the help !

[brian-dellabetta]

Hi @Gman3214 , try using latest main, there are tons of improvements on it since the 0.5.1 release. We are building a 0.6.0 release today

[brian-dellabetta]

Just FYI, v0.6.0 was just released, feel free to use that

[Gman3214]

@brian-dellabetta Thats awesome thank you so much, ill test it later today, btw just a quick question, because i say some different opinions online, and i was wondering what you think what is generally better awq or gptq?

[Gman3214]

@brian-dellabetta works amazing thank you so much for the help!

for anyone in the future that might want to use his own dataset with AWQ heres my code :

#AWQ test
import requests
import torch
from PIL import Image
from transformers import AutoProcessor, Gemma3ForConditionalGeneration
from datasets import load_dataset
from llmcompressor.transformers import oneshot
from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.modifiers.awq import AWQModifier
import os
from compressed_tensors.quantization import (
    QuantizationArgs,
    QuantizationScheme,
    QuantizationStrategy,
    QuantizationType,
)

# Load model.
model_id = "google/gemma-3-4b-it"
model = Gemma3ForConditionalGeneration.from_pretrained(model_id, device_map="auto", torch_dtype="auto")
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)

# Oneshot arguments
DATASET_ID = "ram3214/HQ-Gemma3-hebrew-dataset-small"
DATASET_SPLIT = "train"
NUM_CALIBRATION_SAMPLES = 512
MAX_SEQUENCE_LENGTH = 2048

# Load dataset and preprocess.
ds = load_dataset(DATASET_ID, split=DATASET_SPLIT,)
ds = ds.shuffle(seed=42)


# Define a oneshot data collator for multimodal inputs.
def data_collator(batch):
    assert len(batch) == 1, "Only batch size of 1 is supported for calibration"
    item = batch[0]
    collated = {}
    import torch


    for key, value in item.items():
        if isinstance(value, torch.Tensor):
            collated[key] = value.unsqueeze(0)
        elif isinstance(value, list) and isinstance(value[0][0], int):
            # Handle tokenized inputs like input_ids, attention_mask
            collated[key] = torch.tensor(value)
        elif isinstance(value, list) and isinstance(value[0][0], float):
            # Handle possible float sequences
            collated[key] = torch.tensor(value)
        elif isinstance(value, list) and isinstance(value[0][0], torch.Tensor):
            # Handle batched image data (e.g., pixel_values as [C, H, W])
            collated[key] = torch.stack(value)  # -> [1, C, H, W]
        elif isinstance(value, torch.Tensor):
            collated[key] = value
        else:
            print(f"[WARN] Unrecognized type in collator for key={key}, type={type(value)}")

    return collated

recipe = [
    AWQModifier(
        targets="Linear",
        scheme="W4A16",
        ignore=[
            "lm_head",
            "re:model\.vision_tower.*",
            "re:model\.multi_modal_projector.*",
        ],
    ),
]

SAVE_DIR=f"{model_id.split('/')[1]}-quantized.w4a16"

# Perform oneshot
oneshot(
    model=model,
    tokenizer=model_id,
    dataset=ds,
    splits=DATASET_SPLIT,
    recipe=recipe,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    trust_remote_code_model=True,
    data_collator=data_collator,
)

SAVE_DIR = model_id.rstrip("/").split("/")[-1] + "-W4A16-G128"
model.save_pretrained(SAVE_DIR, save_compressed=True)
processor.save_pretrained(SAVE_DIR)

[brian-dellabetta]

Hi @Gman3214 , awesome! Glad you were able to resolve. The batch size 1 constraint seems strange, but yeah any training/oneshot is very sensitive to the preprocessing of the vision-language dataset.

I will mark this as resolved for now, feel free to re-open if you hit any other issues