attention sink not update correctly when fine-tuning gpt-oss with sequence packing
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
Sink related parameters is not stored in optim.state during fine-tuning, making resume training fails. If sink is not stored, training might also be problematic.
After made some initial investigations, I find that the dtensor policy worker will always use flash attn 2 as backend in case of sequence packing. The source implementation does not take into account of the sinks parameters. It is just passed into s_aux args but not used in flash attention forward pass https://github.com/huggingface/transformers/blob/70179949f7899e0ca235210f7188c6c9e0add77b/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L363-L378
**Steps/Code to reproduce bug**
Just fine-tune gpt-oss with sequence packing enabled should reproduce this bug.
Contributor guide
Assessment
This issue has not been assessed yet.