[feature] add Sequence Parallelism support for offline training

[uygnef]

Motivation

This PR introduces Sequence Parallelism (SP) support for SpecForge training. The primary goal is to enable the training of models with significantly longer sequence lengths, which was previously limited by memory constraints on single GPUs.

Modifications

• Implemented USP Support: Added support for Unified Sequence Parallelism (USP), integrating both Ring Attention and Ulysses attention mechanisms to handle long contexts efficiently.

• Draft Model SP Group: Established a dedicated sequence parallel process group for the draft model to ensure correct communication and synchronization during distributed

Related Issues

#299

Accuracy Test

python tests/test_layers/test_decoder.py

for fp32, output diff less than 1e-5
for bf16, output diff less than 2e-2
grad diff is less than 1e-6

loss

train acc

Benchmark & Profiling

LIMIT

1. Suboptimal Loss Aggregation: Currently uses all_gather within the SP group. Will be optimized to local calculation + reduce_sum to save VRAM.
2. No Online Mode Support: Limited to offline mode. Online training requires resolving TP/SP size mismatches, managing large hidden states for data distribution, and modifying the gradient accumulation loop.
3. Attention Implementation: Currently uses sdpa, which is memory-intensive. Future updates will replace this with flex_attention for better optimization

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]

Summary of Changes

Hello @uygnef, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces Sequence Parallelism (SP) to SpecForge, a critical enhancement for training large language models with extended sequence lengths. By leveraging Unified Sequence Parallelism (USP), which combines Ring Attention and Ulysses attention, the system can now efficiently distribute and process longer sequences across multiple GPUs, effectively mitigating memory constraints. This foundational change paves the way for developing more powerful and context-aware models.

Highlights

• Sequence Parallelism (SP) Support: Introduced comprehensive Sequence Parallelism (SP) support for SpecForge training, enabling the handling of significantly longer sequence lengths than previously possible.
• Unified Sequence Parallelism (USP) Implementation: Implemented Unified Sequence Parallelism (USP), which integrates both Ring Attention and Ulysses attention mechanisms for efficient distributed processing of long contexts.
• Dedicated SP Process Groups: Established dedicated sequence parallel process groups for the draft model to ensure correct communication and synchronization during distributed training.
• New Configuration Arguments: Added new command-line arguments (--sp-ulysses-size, --sp-ring-size) to train_eagle3.py for configuring the dimensions of sequence parallelism.
• Data Collator Adaptation: Modified the DataCollatorWithPadding to adjust padding dynamically, ensuring sequence lengths are multiples of the sequence parallel degree for proper data distribution.
• LlamaUSPAttention Module: Introduced a new LlamaUSPAttention class, extending LlamaAttention, to encapsulate the specific logic for Ulysses and Ring Attention within the draft model.
• Unit Testing for USP: Added a new unit test (tests/test_layers/test_decoder.py) to validate the correctness and numerical stability of the USP implementation against the standard attention mechanism.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces significant new functionality by adding Sequence Parallelism (SP) support, specifically Unified Sequence Parallelism (USP) with Ring and Ulysses attention. This is a valuable addition for training models with long sequences. The implementation is comprehensive, touching distributed setup, data collation, and the core attention mechanism, and includes a new distributed test case.

My review focuses on a few key areas:

• Security and Portability: The example script contains hardcoded, environment-specific, and insecure configurations that should be removed.
• Resource Management: There's a potential resource leak in the distributed setup where a process group is not destroyed.
• Code Clarity and Correctness: I've pointed out areas for simplification, documentation updates, and removal of debugging artifacts.

Overall, this is a strong contribution. Addressing these points will improve the robustness and usability of the new feature.

[xhdidi]

With 8×H800, it seems that TP=8 can support a longer maximum training length than any other USP configuration. Is there any plan to enable flex-attention in the SP?

[uygnef]

With 8×H800, it seems that TP=8 can support a longer maximum training length than any other USP configuration. Is there any plan to enable flex-attention in the SP?

Flex attention does not work well with B/A-series GPUs. I aim to support this PR based on Flash Attention (this PR already exists). If there are precision issues, we will then enable Flex Attention support, which is estimated to take about 2 weeks.

Have you tried sp_ring_size=8? TP=8 maxes out at 16k sequence length, and sp_ring_size=8 should support longer sequences.

[xhdidi]

With 8×H800, it seems that TP=8 can support a longer maximum training length than any other USP configuration. Is there any plan to enable flex-attention in the SP?

Flex attention does not work well with B/A-series GPUs. I aim to support this PR based on Flash Attention (this PR already exists). If there are precision issues, we will then enable Flex Attention support, which is estimated to take about 2 weeks.

Have you tried sp_ring_size=8? TP=8 maxes out at 16k sequence length, and sp_ring_size=8 should support longer sequences.

I tried to train draft model for Qwen3-480B with 8*H800(80G). When tp=8 and using flex-attention backend, the maximum training length could reach 24k, but when ring-size=8 and training with usp backend, the maximum training length is only 12k.
When using the same attention backend, ring-size=8 might support a larger context length than tp=8. However, the memory saved by flex-attention seems to play a greater role while training long texts.

[ggg-s]

hi @uygnef，i want to know Why SP can't use Flex?

[uygnef]

hi @uygnef，i want to know Why SP can't use Flex?

It's easy to implement Ulysses with Flex Attention, but supporting Ring Attention involves complex mask handling. I'm adding USP support on top of Flash Attention, nearly done. PR coming in a few days.

[ggg-s]

@uygnef fine！

[uygnef]

@uygnef fine！

hi, here is the PR #411