gpt-oss is not working with flash-attention
Author: avihu111
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 5.0.0.dev0
- Platform: Linux-5.14.0-503.23.1.el9_5.x86_64-x86_64-with-glibc2.34
- Python version: 3.10.15
- Huggingface_hub version: 1.1.5
- Safetensors version: 0.4.5
- Accelerate version: 1.11.0
- Accelerate config:    not found
- DeepSpeed version: 0.17.5
- PyTorch version (accelerator?): 2.7.1+cu126 (CUDA)
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

@ArthurZucker @Cyrilvallez 

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

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

### Reproduction

I ran the following code using the latest main branch:
```
from transformers import AutoTokenizer, AutoModelForCausalLM, Mxfp4Config
import torch
quantization_config = Mxfp4Config(dequantize=True)
tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")
model = AutoModelForCausalLM.from_pretrained("openai/gpt-oss-20b", 
                                             attn_implementation="flash_attention_3",
                                             quantization_config=quantization_config,
                                             dtype=torch.bfloat16,
                                             device_map="cuda",
                                             )
messages = [
    {"role": "user", "content": "Who are you?"},
]
inputs = tokenizer.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=40)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:])) 
```

printed:
```
The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. 
The tokenizer class you load from this checkpoint is 'PreTrainedTokenizerFast'. 
The class this function is called from is 'TokenizersBackend'.
Unrecognized keys in `rope_parameters` for 'rope_type'='yarn': {'truncate'}
Unrecognized keys in `rope_parameters` for 'rope_type'='yarn': {'truncate'}
Loading weights: 100%|███████████████████████████████████████████████| 411/411 [00:05<00:00, 72.97it/s, Materializing param=model.norm.weight]
  medium analysis,  final.      medium.    medium.    medium medium  final medium  medium   ...  ...  ...  ...  ...
```
I didn't trace what's the reason for this issue

### Expected behavior

Less warnings, valid text output from gpt-oss

Thanks a lot!

Comments:
Hey! Did you try with today's main? It should have been fixed a few hours ago!

<img width="1636" height="578" alt="Image" src="https://github.com/user-attachments/assets/50cf4a84-218b-4b53-977a-082d96f23f2c" />I 
I think we just fixed it!

thanks! pulled again and reran and got a slightly different outcome - I'll try to update kernels, remove flash attention and delete gpt-oss cache and update
```Unrecognized keys in `rope_parameters` for 'rope_type'='yarn': {'truncate'}
Unrecognized keys in `rope_parameters` for 'rope_type'='yarn': {'truncate'}
Loading weights: 100%|███████████████████████████████████████████████| 411/411 [00:06<00:00, 66.87it/s, Materializing param=model.norm.weight]

User: 8 medium # inanalysis, # inanalysis, # medium, # inanalysis, # medium, # medium, # medium, # medium, # medium, # inanalysis,```


By the way, what's the issue here: 
```Unrecognized keys in `rope_parameters` for 'rope_type'='yarn':```

IDK! cc @zucchini-nlp I think has a PR to fix it

Hmm, removing flash-attention fixed the issue for me.
I didn't make changes in my flash-attention, it used to work

Update us if it persists! 

Yes, it persists with flash_attention_2/3. without specifying `attn_implementation` it runs OK

Other models that run with the same flash-attention version work well, so I suspect it's specifically related to gpt-oss

Probably due to the sink! If I remember correctly, fa2 cannot work with the sink does it @ArthurZucker? We should probably set `_supports_flash=False` if it's not the case!

GPT Oss needs fa3 or something to have `sinks` in the def

- Fa4 has sinks (not yet integrated into transformers, WIP #42435) 
- Fa2 has an open PR on the flash attn repo side (https://github.com/Dao-AILab/flash-attention/pull/1819)
- Fa3 is not supported at all iirc, but maybe easier to manage through the PR above ^

Only kernels will work with fa for now! I.e. `attn_implementation="kernels-community/vllm-flash-attn3"`

Thanks @vasqu! 
If I recall correctly, `gpt-oss` worked with flash attention in the past. 
What changed?

No no it never did! 

Hey @avihu111 

I can repro your issue. I'm not sure what the RoPE warnings are doing here tho (cc @zucchini-nlp). 
I suspect that you maybe used gpt-oss without having FA3 actually installed? We internally switch to `kernels-community/vllm-flash-attn3` if `flash_attention_3` was requested and it is not found in your env. That's also why @ArthurZucker couldn't repro (no FA3 --> fallback --> kernels work).

Native support as is listed in my comment before

Thanks @vasqu! Maybe that's what happened 🤔 thanks!
Anyway, as @Cyrilvallez said, it would be good to raise an error in case someone (like me) tries to load gpt-oss with flash attention 2/3.  🤕 



Just realized I can't use a padding-free setup with gpt-oss, as it relies on flash attention that knows how to interpret those arguments based on the `position_ids` (e.g. [here](https://github.com/huggingface/transformers/blob/0fa49db1205e0a2745161ccac46184e7e46b6e2b/src/transformers/modeling_flash_attention_utils.py#L354C5-L354C40))


@vasqu I guess the `kernels-community/vllm-flash-attn3` kernel does not support a backward pass. right?

The warning is a false positive, I accidentally deleted `truncate` from allowed keys in yarn config. Will add it back in https://github.com/huggingface/transformers/pull/42513

> I guess the kernels-community/vllm-flash-attn3 kernel does not support a backward pass. right?

Yes, it's only using the fwd pass

> Just realized I can't use a padding-free setup with gpt-oss, as it relies on flash attention that knows how to interpret those arguments based on the position_ids

You can also use paddingfree with all the other attention flavors, see https://github.com/huggingface/transformers/blob/ac0769cd8ca186d4dea6a04cd6e87c38c14a34b1/src/transformers/masking_utils.py#L647-L676

You just have to make sure that you have the correct position ids (just as in flash attention).

> it would be good to raise an error in case someone (like me) tries to load gpt-oss with flash attention 2/3.

Ideally, yes. Gpt-oss is a bit of a special case here and our system usually depends if one works, every other one will work. Open to any PRs here.

> You can also use paddingfree with all the other attention 

Cool! Thanks!
> Open to any PRs here

anything besides what @Cyrilvallez suggested?

>We should probably set _supports_flash=False if it's not the case!

> We should probably set _supports_flash=False if it's not the case!

This suggestion would probably disable even the kernel attn support, so unsure here. I'd rather we have something else tbh

fair enough. keeping this issue open for a little bit, maybe someone would suggest a better solution 🙏 

I mean we could add a warning and then default to `kernels-community/vllm-flash-attn3` . 
I guess that's probably better than setting `_supports_flash=False` . 
I could work on a PR for that if that's a good solution

@i3hz open for a PR here 🤗 I would rather not default but throw an error here. We cannot ensure that people running FA2 for example also have kernels installed (even if they should imo)

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

> > I guess the kernels-community/vllm-flash-attn3 kernel does not support a backward pass. right?
> 
> Yes, it's only using the fwd pass
> 

What happens during training then if only the forward pass is provided? Does it default to some other implementation for the backwards pass?

@ryandeng1 Looking at the source code of the kernel [over here](https://huggingface.co/kernels-community/vllm-flash-attn3/blob/main/build/torch27-cxx11-cu128-x86_64-linux/vllm_flash_attn3/flash_attn_interface.py#L300), it looks like it silently ignores the `sinks` 😢 open to PRs improving the behavior on our side

For training, you could use flex attention instead 

so is there any setting except attn_implementation="eager" (which consumes a lot of VRAM i think) which works for training or not really?

> For training, you could use flex attention instead

`attn_implementation="flex_attention"` should work with gpt oss as well, note that mask creation is a bit expensive so best to create it outside the forward 

Thank you! Indeed seem to be doing something.