(fsdp2 dev. perf) Add sequence packing for EAGLE3, DFlash, and DFlash2

[yushengsu-thu]

Motivation

Variable-length samples waste draft-training work on padding. Add opt-in sequence packing for text EAGLE3, DFlash, and DFlash2 with offline features and online disaggregated consumers.

Modifications

• Enable training.sequence_packing: true with training.attention_backend: flex_attention; the default remains disabled.
• Pack each existing microbatch with document-local positions, isolated attention and boundary-safe labels/teacher shifts. Preserve sample order, logical batch size, loss normalization, accumulation and acknowledgements.
• Preserve DFlash/DFlash2 anchor sampling, use sparse document-aware masks, skip invalid backbone proposals, and restore the original objective layout. Online target capture remains per sample.
• Add focused configuration, collation, model-parity, lifecycle and online-integration tests, plus usage documentation.

USP, multimodal inputs, compact teacher, loss-position trimming and EAGLE3 LK are unsupported. DFlash/DFlash2 LK and D-PACE are supported.

Accuracy Test

Prior Linux/H200 validation (2026-10-02): 204 CPU tests/540 subtests; 57 GPU model/host-sync tests/45 subtests; five live-online/metadata tests/seven subtests. Two-rank FULL_SHARD loss/gradient comparisons passed for all three architectures. Full-model runs verified every trainable tensor received 256 optimizer updates, all five draft layers updated, and weights/losses stayed finite.

Current checks: repository-wide pre-commit and 27 configuration/package-architecture tests (21 subtests) pass. All 17 changed/new production files match the full-model measured-source hashes. The local macOS Torch installation is missing libtorch_cpu.dylib, so model/GPU results above are from the prior completed H200 runs.

BF16 parity is tolerance-based; convergence, final checkpoint quality and speculative-serving acceptance were not evaluated.

Benchmark & Profiling

Complete Qwen3-4B target (36 layers), five-layer drafts, two H200 GPUs, 1,024 ShareGPT conversations, 256 optimizer steps, four runs per mode. Medians:

Model	Training step: padded → packed	Completion: padded → packed	Completion speedup
DFlash	452.79 → 433.64 ms	145.526 → 140.378 s	1.0367×
DFlash2	426.98 → 405.76 ms	137.320 → 132.905 s	1.0332×

Completion includes live target capture, transport, training, acknowledgements and checkpoints at steps 128/256. It excludes preparation, startup and a separate corpus warmup. Training-step diagnostics cover the first 250 steps (1.0442×/1.0523× faster). Gains are workload-specific; the full-vocabulary objective still uses the restored proposal layout.

Related Issues

None.

Checklist

• Format code with repository pre-commit hooks.
• Add unit and integration tests.
• Update usage documentation.
• Provide training performance and numerical-parity results with validation limits.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.