feat: add scalable colocated online training

[maocheng23]

Motivation

SpecForge online training currently requires a disaggregated producer/consumer topology even when target capture and draft training fit on the same accelerator. That adds feature serialization, transport, and a second GPU pool, and makes target-feature locality difficult for K3-class multi-node jobs.

This draft adds a bounded colocated path: every trainer rank owns its local SGLang target shard and FSDP draft shard, captures on demand, trains its rank-local feature slice, and releases it before the next capture.

Modifications

• Add typed deployment.mode: local_colocated validation and launch planning for online SGLang capture.
• Add a synchronous, bounded LocalRolloutStream, local capture adapter, packed variable-length capture, TP-local slicing, and zero-copy feature loading.
• Deterministically shard prompts across target-DP islands while target-TP peers capture one aligned TP-wide batch and train distinct local samples.
• Add HYBRID_SHARD support using target-TP as the FSDP shard group and target-DP as the replica group. Loss/metric reductions remain WORLD-wide and gradient norm counts each shard once.
• Coordinate prompt-cache population so only one rank per shared/node-local cache builds the expensive Arrow map; other ranks take the cache-hit path.
• Make distributed loggers rank-zero-only and add capture, feature residency, step timing, throughput, and accelerator peak-memory metrics.
• Route DSpark capture through native DSpark when available and the SGLang 0.5.14 DFlash capture API otherwise.
• Make rank-local failures abort communicators so torchrun surfaces the originating traceback instead of hanging during NCCL teardown.
• Use stable Triton FlexAttention on Hopper and the FLASH backend on supported Blackwell/PyTorch stacks; retain an ATen GEMM fallback for dynamic shapes.
• Add Qwen3-8B H200 and Kimi K3 4x8 B300/HSDP starting recipes plus topology, memory-sizing, and operations documentation.

Related Issues

No linked issue.

Accuracy Test

• Real Qwen3-8B DSpark colocated training completed 90 optimizer steps / 360 global samples on 4x H200 with finite loss and gradient norms and exit code 0.
• The run covered variable sequence lengths from 43 to 3911 tokens under a configured 4096-token maximum, including shapes that previously triggered a TorchInductor NoValidChoicesError.
• Relevant CPU/unit coverage: 273 passed, 7 skipped, 501 subtests passed.
• Kimi K3/HSDP is configuration- and process-group-tested, but has not yet been run on the intended 4x8 B300 topology; this is why the PR remains draft.

Benchmark & Profiling

Environment: Qwen3-8B, DSpark, BF16, SGLang 0.5.14 local capture, FSDP SHARD_GRAD_OP, batch size 1/rank, 4x H200, 90 optimizer steps, 360 global samples, variable lengths up to 3911 tokens.

Window	Weighted global throughput	Notes
step 50-90	22.93 samples/s	post-compilation steady windows
step 20-90	16.94 samples/s	includes first-time long-shape compilation
step 1-90	12.94 samples/s	includes cold start/compiler warmup

• Post-compilation local target capture consumed about 11.2% of optimizer-step wall time.
• Peak PyTorch allocated/reserved memory per rank: 31.94 / 35.07 GiB.
• Peak rank-local captured-feature residency: 0.179 GiB.
• A separate 512-token two-step smoke exercised capture -> loss -> FSDP backward -> optimizer -> checkpoint end to end.

The historical repository experiment measured 23.85 samples/s for an 8x H200 colocated run, but it used a different GPU count and sample window; that number is directional context, not a direct speedup claim for this PR.

Checklist

• Format your code according to the Code Formatting with Pre-Commit guidance (black --check and isort --check-only).
• Add unit tests.
• Update documentation, architecture notes, and example recipes.
• Provide throughput and memory profiling results.
• Validate Kimi K3 with TP8/HSDP on 4x8 B300 before marking ready for review.

[maocheng23]

This draft has been restructured into a stacked series on latest main (it had grown conflicts with 8 files on main, including the examples/configs re-layout and the trim-loss trainer changes):

1. refactor: extract packed capture_rows from the offline SGLang backend #780 — refactor: packed capture_rows + DSpark dense-target capture hook fallback
2. fix: abort process groups on rank-local training failure #781 — fix: abort process groups on rank-local training failure
3. fix: create the metrics tracker only on trainer rank zero #782 — fix: rank-zero-only metrics tracker
4. feat: colocated online training with in-process SGLang capture #783 — feat: colocated online runtime (core, supersedes the bulk of this PR)
5. feat: HYBRID_SHARD draft training across colocated target islands #784 — feat: HYBRID_SHARD across target islands + K3 recipe

Intentionally dropped from the stack, with rationale in #783's description: the deployment.colocated config knobs (safe defaults hardcoded), the global FlexAttention default-backend change (SPECFORGE_FLEX_ATTENTION_BACKEND=FLASH remains the Blackwell opt-in), the CaptureConfig.extra validation tightening on the shared disaggregated path (worth its own follow-up), and the offline capture_eagle3 attention-mask trimming (kept byte-identical). The silent no-op on mismatched resume was fixed to a hard error, and the mock-driven __dict__ lookup in the metric reduction was replaced with a clean helper.

Keeping this PR open for reference until the stack lands; it can be closed once #783 is in.