Allow compile with bnb
Author: SunMarc
Pull request: yes

# What does this PR do?

This PR enable compilation when generating with bnb models. This is supported with the latest bnb. 

Comments:
When generating by default , is `fullgraph` is set to `False` ? From the codebase it looks like that but I wanted to be sure cc @gante 

any idea why the test is not passing @matthewdouglas ? 

` RUN_SLOW=True pytest tests/quantization/bnb/test_4bit.py::Bnb4bitCompile -s -vvvvv`. I'm using torch 2.7.1 

Getting the following traceback:
```
FAILED tests/quantization/bnb/test_4bit.py::Bnb4bitCompile::test_generate_compile - torch._dynamo.exc.Unsupported: Unsupported method call
  Explanation: Dynamo does not know how to trace method `t` of class `Params4bit`
  Hint: Avoid calling `Params4bit.t` in your code.
  Hint: Please report an issue to PyTorch.

  Developer debug context: call_method UserDefinedObjectVariable(Params4bit) t [] {}


from user code:
   File "/admin/home/marc/transformers/src/transformers/utils/generic.py", line 943, in wrapper
    output = func(self, *args, **kwargs)
  File "/admin/home/marc/transformers/src/transformers/models/llama/modeling_llama.py", line 555, in forward
    outputs: BaseModelOutputWithPast = self.model(
  File "/admin/home/marc/transformers/src/transformers/utils/generic.py", line 943, in wrapper
    output = func(self, *args, **kwargs)
  File "/admin/home/marc/transformers/src/transformers/models/llama/modeling_llama.py", line 443, in forward
    layer_outputs = decoder_layer(
  File "/admin/home/marc/transformers/src/transformers/modeling_layers.py", line 48, in __call__
    return super().__call__(*args, **kwargs)
  File "/admin/home/marc/miniconda3/envs/hf/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/admin/home/marc/miniconda3/envs/hf/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
  File "/admin/home/marc/transformers/src/transformers/models/llama/modeling_llama.py", line 294, in forward
    hidden_states, self_attn_weights = self.self_attn(
  File "/admin/home/marc/transformers/src/transformers/models/llama/modeling_llama.py", line 235, in forward
    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
  File "/admin/home/marc/miniconda3/envs/hf/lib/python3.10/site-packages/bitsandbytes/nn/modules.py", line 496, in forward
    return bnb.matmul_4bit(x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state).to(inp_dtype)

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

```


Same for 8bit 


```

FAILED tests/quantization/bnb/test_mixed_int8.py::Bnb8bitCompile::test_generate_compile - torch._dynamo.exc.Unsupported: Dynamic shape operator
  Explanation: Operator `bitsandbytes.int8_vectorwise_quant.default`'s output shape depends on input Tensor data.
  Hint: Enable tracing of dynamic shape operators with `torch._dynamo.config.capture_dynamic_output_shape_ops = True`

  Developer debug context: bitsandbytes.int8_vectorwise_quant.default


from user code:
   File "/admin/home/marc/transformers/src/transformers/utils/generic.py", line 943, in wrapper
    output = func(self, *args, **kwargs)
  File "/admin/home/marc/transformers/src/transformers/models/llama/modeling_llama.py", line 555, in forward
    outputs: BaseModelOutputWithPast = self.model(
  File "/admin/home/marc/transformers/src/transformers/utils/generic.py", line 943, in wrapper
    output = func(self, *args, **kwargs)
  File "/admin/home/marc/transformers/src/transformers/models/llama/modeling_llama.py", line 443, in forward
    layer_outputs = decoder_layer(
  File "/admin/home/marc/transformers/src/transformers/modeling_layers.py", line 48, in __call__
    return super().__call__(*args, **kwargs)
  File "/admin/home/marc/miniconda3/envs/hf/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/admin/home/marc/miniconda3/envs/hf/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
  File "/admin/home/marc/transformers/src/transformers/models/llama/modeling_llama.py", line 294, in forward
    hidden_states, self_attn_weights = self.self_attn(
  File "/admin/home/marc/transformers/src/transformers/models/llama/modeling_llama.py", line 235, in forward
    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
  File "/admin/home/marc/miniconda3/envs/hf/lib/python3.10/site-packages/bitsandbytes/nn/modules.py", line 1010, in forward
    out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state)
  File "/admin/home/marc/miniconda3/envs/hf/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py", line 369, in matmul
    return MatMul8bitLt.apply(A, B, out, bias, state)
  File "/admin/home/marc/miniconda3/envs/hf/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py", line 196, in forward
    CA, SCA, outlier_cols = F.int8_vectorwise_quant(A.to(torch.float16), threshold=state.threshold)
  File "/admin/home/marc/miniconda3/envs/hf/lib/python3.10/site-packages/bitsandbytes/functional.py", line 2245, in int8_vectorwise_quant
    return torch.ops.bitsandbytes.int8_vectorwise_quant.default(A, threshold)

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

```

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_38886). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Thanks for adding this!

For 4bit this is just a graph break so it does seem to be using `fullgraph=True`. That test would pass on torch >= 2.8. 

I think this default comes from `CompileConfig` here?
https://github.com/huggingface/transformers/blob/c27f628e98744f156e653d7eeab3afe018cd14b0/src/transformers/generation/configuration_utils.py#L1564-L1579

If that's what we expect by default it might be simplest to guard on for torch >= 2.8 for 4bit. Otherwise we might want to find a way to improve UX, like catching this and providing a better trace/error message, or disabling the fullgraph like here:

https://github.com/huggingface/transformers/blob/9cd7570f34fdb833ed874b2eba4d4ea3ae9ccb03/src/transformers/generation/utils.py#L3581-L3586

Regarding 8bit, it's expected that you need `torch._dynamo.config.capture_dynamic_output_shape_ops = True` unless threshold=0.0. Fortunately the error message here isn't too bad and has the right advice, but I think this is also something we should just check for in `is_compileable`.

Thanks for everyone advice ! I will check that later and probably merge it after torch 2.8 it out