Olmo-3 flash-attention-2 support
Author: Zephyr271828
Pull request: no

## Problem Description
Hi! I've trying to post-train the Olmo-3 model recently, but it seems Olmo-3 does not support flash-attention-2 at this moment. Though switching to sdpa is a workaround, it consumes more memory and time compare to flash-attn-2. Therefore, I wonder if you have plans to support Olmo-3 at this moment:)

## Error Message
```
  File "/work/nvme/bdhh/yxu21/offrl/math/training/verl/verl/workers/fsdp_workers.py", line 823, in init_model
    self.ref_module_fsdp = self._build_model_optimizer(
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nvme/bdhh/yxu21/offrl/math/training/verl/verl/workers/fsdp_workers.py", line 377, in _build_model_optimizer
    actor_module = actor_module_class.from_pretrained(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nvme/bdhh/yxu21/venvs/synlogic/lib/python3.12/site-packages/transformers/models/auto/auto_factory.py", line 571, in from_pretrained
    return model_class.from_pretrained(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nvme/bdhh/yxu21/venvs/synlogic/lib/python3.12/site-packages/transformers/modeling_utils.py", line 279, in _wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/work/nvme/bdhh/yxu21/venvs/synlogic/lib/python3.12/site-packages/transformers/modeling_utils.py", line 4336, in from_pretrained
    config = cls._autoset_attn_implementation(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nvme/bdhh/yxu21/venvs/synlogic/lib/python3.12/site-packages/transformers/modeling_utils.py", line 2109, in _autoset_attn_implementation
    cls._check_and_enable_flash_attn_2(
  File "/work/nvme/bdhh/yxu21/venvs/synlogic/lib/python3.12/site-packages/transformers/modeling_utils.py", line 2233, in _check_and_enable_flash_attn_2
    raise ValueError(
ValueError: Olmo3ForCausalLM does not support Flash Attention 2.0 yet. Please request to add support where the model is hosted, on its model hub page: https://huggingface.co//work/nvme/bdhh/yxu21/offrl/olmo/checkpoints/generic/Olmo-3-1025-7B/discussions/new or in the Transformers GitHub repo: https://github.com/huggingface/transformers/issues/new
```

Comments:
Olmo3 should support FA https://github.com/huggingface/transformers/blob/e8e142de378e15626ace04b02256a9d8ee158550/src/transformers/models/olmo3/modeling_olmo3.py#L351

I also tried
```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "allenai/Olmo-3-7B-Instruct",
    device_map="auto",
    torch_dtype="auto",
    attn_implementation="flash_attention_2",
).eval()
tokenizer = AutoTokenizer.from_pretrained("allenai/Olmo-3-7B-Instruct")
tokenizer.pad_token = tokenizer.eos_token
inputs = tokenizer(
    ["Hello, how are you?", "is this life?"],
    padding=True,
    padding_side="left",
    return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=50, do_sample=False)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
```
as sanity check and it worked. What transformers version are you running? I tested on v5 main

@vasqu Hi! Sorry for the late reply. I tried again and FA2 worked for OLMo-3. Thank you for your replyl!!