torch.compile fails for gemma-3-1b-it

[InCogNiTo124]

System Info

• transformers version: 4.52.4
• Platform: Linux-6.15.0-1-MANJARO-x86_64-with-glibc2.41
• Python version: 3.12.8
• Huggingface_hub version: 0.32.3
• Safetensors version: 0.5.3
• Accelerate version: 1.7.0
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (GPU?): 2.7.0+cu126 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?: no
• Using GPU in script?: yes
• GPU type: NVIDIA GeForce RTX 3090 Ti

Who can help?

@ArthurZucker @gante

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Running TORCHDYNAMO_VERBOSE=1 TORCH_LOGS="+dynamo" uv run main.py fails:

Minimal reproducible example

import torch
from transformers import GemmaTokenizer, Gemma3ForCausalLM


ckpt = "google/gemma-3-1b-it"
model = Gemma3ForCausalLM.from_pretrained(
    ckpt,
    device_map="cuda:0",
    torch_dtype=torch.bfloat16,
)
processor = GemmaTokenizer.from_pretrained(ckpt)


messages = [{"role": "user", "content": "What is 2^7-2^4??"}]
inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device)


input_len = inputs["input_ids"].shape[-1]


# generate_fn = model.generate

generate_fn = torch.compile(model.generate, fullgraph=True)

generation = generate_fn(**inputs, max_new_tokens=100, do_sample=False)
generation = generation[0][input_len:]


decoded = processor.decode(generation, skip_special_tokens=True)
print(decoded)

Stack trace

Full paste: https://pastebin.com/V103pCWM

  File "/tmp/gemma_torch/.venv/lib/python3.12/site-packages/torch/_dynamo/variables/builtin.py", line 2111, in call_deepcopy
    unimplemented(f"copy.deepcopy {repr(x)}")
  File "/tmp/gemma_torch/.venv/lib/python3.12/site-packages/torch/_dynamo/exc.py", line 439, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: copy.deepcopy UserDefinedObjectVariable(GenerationConfig)

from user code:
   File "/tmp/gemma_torch/.venv/lib/python3.12/site-packages/torch/_dynamo/external_utils.py", line 70, in inner
    return fn(*args, **kwargs)
  File "/tmp/gemma_torch/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/tmp/gemma_torch/.venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 2354, in generate
    generation_config, model_kwargs = self._prepare_generation_config(
  File "/tmp/gemma_torch/.venv/lib/python3.12/site-packages/transformers/generation/utils.py", line 1744, in _prepare_generation_config
    generation_config = copy.deepcopy(generation_config)

Expected behavior

Compilation proceeds

[InCogNiTo124]

Forwarded from pytorch/pytorch#115122 (comment)

[manueldeprada]

cc @zucchini-nlp , can we maybe init a new Gemma3Config for local rope to avoid using deepcopy? whats the state of rope refactor?

transformers/src/transformers/models/gemma3/modular_gemma3.py

Lines 552 to 557 in beaed8c

	# TODO: raushan fix this after RoPE refactor. For now we hack it by reassigning thetas
	# when we want to create a local RoPE layer. Config defaults should hold values for global RoPE
	config = copy.deepcopy(config)
	config.rope_theta = config.rope_local_base_freq
	config.rope_scaling = {"rope_type": "default"}
	self.rotary_emb_local = Gemma3RotaryEmbedding(config=config)

[zucchini-nlp]

Hey @InCogNiTo124 ! This is not an issue for Gemma-3 only and any model would fail with the above script. The reason is that you are comping the model's forward pass before pre-fill which we don't support. Compiling a model before decoding stage in general ins't recommended, as the input shapes in pre-fill are not static and there will be a graph break

Calling a model's generate() with static cache will pre-compile the model automatically for all decoding steps and this is the recommended way currently. We have a small doc section on best practices for compile here as well 🤗

generation = generate_fn(**inputs, max_new_tokens=100, do_sample=False, cache_implementation="static")

[StrongerXi]

@InCogNiTo124 compiling the entire generate function might be difficult, but if you just compile model itself (which effectively compiles the forward function), it should just work out. Can you give it a try?

[InCogNiTo124]

Thanks @StrongerXi @zucchini-nlp it works now! I was so sure I searched all the docs... Closing!

[zucchini-nlp]

Yeah, the docs about compile are hard to discover and we're trying to organize it in a better way