Gemma2 (quantized) inference is broken - torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment.
Author: royvelich
Pull request: no
Labels: name: bug

### System Info

```
- `transformers` version: 4.49.0
- Platform: Linux-5.4.0-156-generic-x86_64-with-glibc2.39
- Python version: 3.10.14
- Huggingface_hub version: 0.29.1
- Safetensors version: 0.5.3
- Accelerate version: 1.4.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.5.1 (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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-40GB

```

### Who can help?

@SunMarc @MekkCyber

### Information

- [x] 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

Run the following quantized example (taken from the official docs):

```
# pip install bitsandbytes accelerate
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_8bit=True)

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-2b-it",
    quantization_config=quantization_config,
)

input_text = "Write me a poem about Machine Learning."
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=32)
print(tokenizer.decode(outputs[0]))
```

You should get the following error:

```
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 895, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 582, in wrapper
    return inner_fn(self, inst)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1602, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 830, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/variables/misc.py", line 1024, in call_function
    return self.obj.call_method(tx, self.name, args, kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/variables/misc.py", line 774, in call_method
    return self.call_apply(tx, args, kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/variables/misc.py", line 723, in call_apply
    return variables.UserFunctionVariable(fn, source=source).call_function(
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 324, in call_function
    return super().call_function(tx, args, kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 111, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 836, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 3011, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 3139, in inline_call_
    tracer.run()
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 983, in run
    while self.step():
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 895, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 582, in wrapper
    return inner_fn(self, inst)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1692, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 830, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 324, in call_function
    return super().call_function(tx, args, kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/variables/functions.py", line 111, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 836, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 3011, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 3139, in inline_call_
    tracer.run()
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 983, in run
    while self.step():
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 895, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 560, in inner
    raise exc.UserError(
torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond to explicitly capture the control flow. For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#cond-operands

from user code:
   File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
    return func(*args, **kwargs)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 887, in forward
    outputs = self.model(
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 667, in forward
    layer_outputs = decoder_layer(
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 321, in forward
    hidden_states, self_attn_weights = self.self_attn(
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 216, in forward
    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/bitsandbytes/nn/modules.py", line 990, in forward
    out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py", line 509, in matmul
    return MatMul8bitLt.apply(A, B, out, bias, state)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py", line 326, in forward
    CA, SCA, outlier_cols = F.int8_vectorwise_quant(A.to(torch.float16), threshold=state.threshold)
  File "/opt/miniconda3/envs/video-pipe-12-4/lib/python3.10/site-packages/bitsandbytes/functional.py", line 2786, in int8_vectorwise_quant
    if outliers.any():

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

Please note, it does work without problems under `transformers==4.48.3`.


### Expected behavior

Text generation should work correctly.

Comments:
cc @gante as well for generation, although this seems like a tricky bug to diagnose!

Potential solution above !