NVIDIA-NeMo / NVIDIA-NeMo/RL

attention sink not update correctly when fine-tuning gpt-oss with sequence packing

Open
#1,672 3 comments 1 reaction 1 assignee Claimed by @adil-a View on GitHub
bug community-request t-gptoss waiting-on-customer x-mini
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.