[Tracing] Support tracing of Gemma3 [#1248]

[kelkelcheng]

SUMMARY:
Add support for tracing of Gemma3: issue#1248.

Steps that I have done:

1. Create gemma3.py from HF and update init.py.
2. Classes and functions that I modified:
   2.1 Gemma3ForConditionalGeneration: _update_causal_mask and forward
   2.2 Gemma3TextModel: _update_causal_mask, forward, and _prepare_4d_causal_attention_mask_with_cache_position

TEST PLAN:
Ran:
llmcompressor.trace --model_id google/gemma-3-4b-it --model_class TraceableGemma3ForConditionalGeneration --ignore "lm_head" "re:vision_tower.*" --modality vision

Output:

This is my first attempt at solving this issue. It is a fun learning experience and please review it carefully.
Gemma3 can go through tracing now, but we might need further tests for the quantization as well.

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[kylesayrs]

I've added some relevant improvements to the tracing system. In the meantime, this is good to land, great job!

[kylesayrs]

@kelkelcheng Would you be so kind as to merge in these changes? #1399. This adds a gemma3 example and fixes style warnings

[kylesayrs]

It looks like there may be some tracing assumptions that are broken, as the example fails.

Line 36

traceback

Traceback (most recent call last): File "/home/kyle/llm-compressor/src/llmcompressor/pipelines/sequential/helpers.py", line 53, in forward outputs = forward_fn(*args, **kwargs) File "", line 36, in forward TypeError: 'NoneType' object is not subscriptable

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/kyle/llm-compressor/src/llmcompressor/modifiers/quantization/gptq/base.py", line 234, in on_initialize
run_sequential(
File "/home/kyle/llm-compressor/src/llmcompressor/pipelines/sequential/pipeline.py", line 67, in run_pipeline
subgraph.forward(model, **inputs)
File "/home/kyle/llm-compressor/src/llmcompressor/pipelines/sequential/helpers.py", line 55, in forward
raise RuntimeError(
RuntimeError: Raised an exception during execution of the following code:

1
2 torch.fx._symbolic_trace.wrap("llmcompressor_transformers_tracing_gemma3_mask_pad_token_id")
3 torch.fx._symbolic_trace.wrap("llmcompressor_transformers_tracing_gemma3_triu_causal_mask")
4 torch.fx._symbolic_trace.wrap("llmcompressor_transformers_tracing_gemma3_bidirectional_mask")
5 torch.fx._symbolic_trace.wrap("llmcompressor_transformers_tracing_gemma3__prepare_4d_causal_attention_mask_with_cache_position")
6
7 def forward(self, input_ids : torch.Tensor, pixel_values : torch.Tensor, attention_mask : torch.Tensor, token_type_ids : torch.Tensor):
8     language_model_model_norm_weight = self.language_model.model.norm.weight
9     language_model_model_rotary_emb_local_inv_freq = self.language_model.model.rotary_emb_local.inv_freq
10     language_model_model_rotary_emb_inv_freq = self.language_model.model.rotary_emb.inv_freq
11     language_model_model_embed_tokens_embed_scale = self.language_model.model.embed_tokens.embed_scale
12     language_model_model_embed_tokens_weight = self.language_model.model.embed_tokens.weight
13     getattr_1 = language_model_model_embed_tokens_weight.dtype
14     gemma3config = transformers_models_gemma3_configuration_gemma3_Gemma3Config(...)
15     getitem_13 = language_model_model_rotary_emb_inv_freq[(None, slice(None, None, None), None)];  language_model_model_rotary_emb_inv_freq = None
16     getitem_16 = language_model_model_rotary_emb_local_inv_freq[(None, slice(None, None, None), None)];  language_model_model_rotary_emb_local_inv_freq = None
17     float_10 = language_model_model_norm_weight.float();  language_model_model_norm_weight = None
18     embedding = torch.nn.functional.embedding(input_ids, language_model_model_embed_tokens_weight, padding_idx = 0, max_norm = None, norm_type = 2.0, scale_grad_by_freq = False, sparse = False);  language_model_model_embed_tokens_weight = None
19     eq_2 = input_ids == 262144   
20     vision_tower = self.vision_tower(pixel_values = pixel_values);  pixel_values = None
21     dim = attention_mask.dim()   
22     dim_1 = attention_mask.dim() 
23     size_2 = attention_mask.size()
24     getitem_8 = attention_mask[(slice(None, None, None), None, None, slice(None, None, None))]
25     to = language_model_model_embed_tokens_embed_scale.to(getattr_1);  language_model_model_embed_tokens_embed_scale = getattr_1 = None
26     mask_pad_token_id = llmcompressor_transformers_tracing_gemma3_mask_pad_token_id(None, input_ids, -1, gemma3config);  input_ids = gemma3config = None
27     float_1 = getitem_13.float();  getitem_13 = None
28     float_5 = getitem_16.float();  getitem_16 = None
29     add_5 = 1.0 + float_10;  float_10 = None
30     unsqueeze = eq_2.unsqueeze(-1);  eq_2 = None
31     getattr_3 = vision_tower.last_hidden_state;  vision_tower = None
32     eq_3 = dim == 4;  dim = eq_3 = None
33     eq_4 = dim_1 == 4;  dim_1 = eq_4 = None
34     getitem_6 = size_2[-1];  size_2 = None
35     mul = embedding * to;  embedding = to = None
36     getitem_55 = mask_pad_token_id[(Ellipsis, slice(1, None, None))];  mask_pad_token_id = None
</details>

[kelkelcheng]

@kelkelcheng Would you be so kind as to merge in these changes? #1399. This adds a gemma3 example and fixes style warnings

Absolutely, and I will try to run the example and see if I can get the same issue

[kylesayrs]

@kelkelcheng It seems that the bug is related to how an if statement was removed when adding mask_pad_token_id, specifically that there used to be an if labels statement before that. With the if statement removed, the model attempts to mask the labels, even if the user doesn't pass any labels. I made that change in #1399.

Once the example runs to competition, this should be good to go

[kylesayrs]

Success!
https://huggingface.co/nm-testing/gemma-3-4b-it-W4A16-G128/tree/main/

[kelkelcheng]

Success! https://huggingface.co/nm-testing/gemma-3-4b-it-W4A16-G128/tree/main/

Awesome—thanks so much for digging in and testing!
I just merged the changes from #1399 and pushed again. Looks like the DCO is acting up, but let me know if there’s anything I need to fix.