feat(dspark): add Ulysses sequence parallel training

[maocheng23]

Motivation

FSDP shards model parameters, gradients, and optimizer state, but it does not shard sequence activations. Long-context DSpark training therefore remains constrained by replicated 120K-token activations even when model-state memory is manageable.

This PR adds Ulysses sequence parallelism to the offline DSpark path so each SP rank owns only its local sequence slice while the draft remains FSDP-sharded.

Modifications

• Shard and pad offline input_ids, captured target features, final target states, masks, and absolute positions across the SP group.
• Run separate Q/K/V Ulysses all-to-all collectives before local attention.
• Sample anchors consistently across ranks and exchange only the sparse tokens and teacher states required by the DSpark objective.
• Preserve bounded objective-logit memory through chunked anchor blocks.
• Validate that Q and KV head counts divide the Ulysses degree.
• Expose the topology through attention_backend: usp, sp_ulysses_size, and the existing FSDP sharding option.
• Add a GLM-5.2 120K/SP8 recipe, documentation, parity coverage, FSDP integration coverage, and an opt-in 8-GPU 120K smoke test.

The initial implementation is intentionally Ulysses-only. Ring sequence parallelism and online/disaggregated DSpark SP are out of scope.

Related Issues

No linked issue.

Accuracy Test

• 2-GPU Ulysses loss and gradient parity against the full-sequence FlexAttention reference passed.
• 2-GPU DSpark Ulysses + FSDP FULL_SHARD optimizer-step integration passed.
• 8-GPU 120,000-token SP8 + FULL_SHARD smoke test passed.
• Targeted provider/application/loss tests: 43 passed.
• Algorithm suite: 50 passed.
• Model registry/family suite: 18 passed.
• Config suite: 53 passed.
• Black, isort, compileall, and git diff --check passed.

A real Qwen3-8B 120K end-to-end run was also executed on 8x H200:

• real Qwen/Qwen3-8B target, TP8 capture, YaRN context 131,072
• exactly 120,000 input tokens; 119,976 supervised tokens
• auxiliary target state [1, 120000, 20480] BF16
• final target state [1, 120000, 4096] BF16
• 5-layer DSpark draft with 1,126,422,017 parameters
• one Ulysses SP8 + FSDP FULL_SHARD optimizer step completed
• final loss 4.089876; all checkpointed draft parameters finite
• resumable full and per-rank optimizer checkpoints written successfully

The Qwen run used deterministic synthetic text, so it is an integration and capacity check rather than a model-quality or convergence measurement.

Benchmark & Profiling

On the 8x H200 Qwen3-8B run:

• target capture: 2m29s total; 32.1s for the 120K target forward
• training plus checkpoint: 1m23s
• sampled active-step memory: approximately 19.9 GiB per H200
• captured feature artifact: 5.5 GiB
• final resumable checkpoint set: 15 GiB

Checklist

• Format code according to the Code Formatting with Pre-Commit guidance.
• Add unit and distributed integration tests.
• Update documentation and example configuration.
• Provide relevant accuracy and runtime evidence.
• 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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[maocheng23]

Exact online 120K validation completed on 2 nodes / 16x H200.

Configuration:

• Target: GLM-5.2-FP8, SGLang TP8
• Prompt: exactly 120,000 tokens
• Capture layers: 1, 19, 38, 57, 76
• SGLang: max total/prefill 120128, chunked prefill 8192
• Store: Mooncake 128 GiB segment
• Trainer: DSpark, USP Ulysses=8 / Ring=1, FSDP FULL_SHARD, batch=1, accumulation=1, max_steps=1

Observed evidence:

• Every request was delivered to all TP8 scheduler ranks.
• Each 120K prefill completed as 14 x 8192-token chunks + 5312 tokens, HTTP 200, no inference OOM.
• Published tensor specs were hidden_states=[1,120000,30720] bf16, target_last_hidden_states=[1,120000,6144] bf16, input_ids/loss_mask=[1,120000] int64.
• Eight refs (70,794,240,000 resident bytes total) published with zero sink errors.
• The USP8/FSDP consumer completed one optimizer step; consumer=0, inference=0, source_consumed=8.
• A step-1 checkpoint was produced with rank 0 plus all eight FSDP rank state files.

This testing exposed and fixed the forced monolithic-prefill OOM. It also confirmed that the documented 128 GiB Mooncake segment is required; the 32 GiB default fills after three GLM-5.2 samples. The latest lint check passes.