[Feature] USP: Replace SDPA with Flash Attention for memory optimization & Add Online Mode

[uygnef]

waiting for #400 to be merged

Motivation

To accelerate long-context training, this PR integrates Flash Attention into the Unified Sequence Parallelism (USP) framework. By replacing standard PyTorch operations with optimized kernels inside the Ring Attention loop, this implementation significantly improves memory efficiency for Eagle3 draft models.

Additionally, this PR adds support for online mode within the USP framework.

Modifications

• Implemented LlamaUSPFlashAttention: Added a hybrid sequence parallel attention layer that combines Ulysses, Ring Attention, and Flash Attention (flash_attn_func).
• Support online mode for Sequence Parallelism：Enabled online mode support for Sequence Parallelism.

Usage

torchrun \
    ...
    scripts/train_eagle3.py \
    ...
    --attention-backend usp_fa \
    --sp-ulysses-size $ULYSSES_SIZE \
    --sp-ring-size $RING_SIZE

Related Issues

Accuracy Test

python tests/test_layers/test_decoder.py

compare to sdpa implement bf16 diff less than 2e-2, fp16 diff less than 5e-3

Benchmark & Profiling

todo

TODO

1. Optimal Loss Aggregation: Currently uses all_gather within the SP group. Will be optimized to local calculation + reduce_sum to save VRAM.
2. Enhance Online Mode: Set draft micro-batch size.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[jiapingW]

I tried train seqlen=65536 with ring=8, ulysses=2 and it'll cost 94G. The setting with ring=4, ulysses=2 and it'll also cost 94G. Does it normal?