⚡ Fix Flash Attention x Padding-Free loss

[qgallouedec]

Before

comparing

• BSZ 4, GAS 2
• BSZ 2, GAS 4
• BSZ 1, GAS 8

After

comparing

• BSZ 4, GAS 2
• BSZ 2, GAS 4
• BSZ 1, GAS 8

and comparing

• kernels-community/flash-attn
• kernels-community/flash-attn3
• kernels-community/vllm-flash-attn3

and comparing 1/2 GPUs (slightly different, not sure why)

Bonus: it's waaaay faster

To reproduce

from trl import SFTTrainer, SFTConfig
from datasets import load_dataset
import os

os.environ["TRACKIO_SPACE_ID"] = "qgallouedec/trackio"
os.environ["TRACKIO_PROJECT"] = "fa_loss_bug"

dataset = load_dataset("trl-lib/Capybara", split="train[:5%]")

gas=2
bsz=4

trainer = SFTTrainer(
    model="HuggingFaceTB/SmolLM3-3B",
    args=SFTConfig(
        per_device_train_batch_size=bsz,
        gradient_accumulation_steps=gas,
        logging_steps=1,
        packing=True,
        use_liger_kernel=True,
        max_length=32768,
        model_init_kwargs={"torch_dtype": "bfloat16", "attn_implementation": "kernels-community/vllm-flash-attn3"},
        report_to="trackio",
        run_name=f"after_bsz{bsz}_gas{gas}_vllm_fa3",
    ),
    train_dataset=dataset,
)
trainer.train()

[qgallouedec]

this is the most important line

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[lewtun]

Excellent fix, LGTM