feat: colocated online training with in-process SGLang capture

[maocheng23]

Motivation

Stack 4/5 replacing draft #766 — the colocated online runtime itself. SpecForge online training currently requires a disaggregated producer/consumer topology even when target capture and draft training fit on the same accelerator, adding feature serialization, transport, and a second GPU pool. This PR 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.

Compared with #766, this cut deliberately drops several pieces to keep the change minimal (details in the stack summary): no deployment.colocated config knobs (synchronize_after_capture/zero_copy_features are hardcoded to the validated safe defaults), no global FlexAttention default-backend change (SPECFORGE_FLEX_ATTENTION_BACKEND already exists for opting into FLASH on Blackwell), no CaptureConfig.extra feature-name plumbing on the shared disaggregated path, and the TorchInductor ATen-fallback fix is scoped to the colocated branch instead of every specforge train invocation.

Modifications

• deployment.mode: local_colocated now validates for online runs; colocated checks cover context-length headroom, request/token capacity versus tp_size * batch_size, SGLang DP/EP constraints, and the sp=1 limitation. Disaggregated/offline validation is unchanged.
• build_colocated_online_runtime assembles a synchronous, bounded pull-through LocalRolloutStream (single worker): the trainer requests one batch inside FeatureDataLoader.get(), the local RolloutWorker captures via the offline backend's capture_rows through the new LocalSGLangCaptureAdapter, and the rank-private in-memory LocalFeatureStore is bounded to one training batch. The loader disables prefetch threads (loader_prefetch_safe) and skips the defensive clone; the controller and loader release the previous batch before requesting the next so long-context features never overlap.
• training.tp_size becomes the colocated target-TP island width. Contiguous TP peers capture one aligned TP-wide batch and train disjoint local slices (each peer clones only its slice to release the packed capture allocation); target-DP islands receive deterministic disjoint prompt plans with per-epoch truncation to complete TP-wide batches.
• Resume: checkpointed runs preposition the prompt stream; a checkpoint whose dataset_size/batch_size/tp_size do not match the current plan now raises instead of silently training nothing (a bug in feat: add scalable colocated online training #766).
• Rank zero populates the tokenized prompt cache before peers read it (node-local caches build once per remaining node), preventing a preprocessing process storm and Arrow-cache write races. EAGLE3 colocated runs derive the streaming vocab mapping deterministically from prepared prompts.
• Colocated runs keep TorchInductor's ATen fallback available for variable-length FlexAttention shapes that otherwise raise NoValidChoicesError (torch's own suggested remedy); an explicit TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS stays authoritative.
• perf/* gains capture time/throughput, feature-residency peaks, and accelerator peak-memory at each log interval; perf/global_samples_per_second counts target-TP peers as sample-parallel (unchanged for tp=1 topologies).
• Docs: new colocated_training.md, topology matrix updates, and the reserved examples/configs/online/colocated/ directory gets the Qwen3-8B DSpark 8xH200 recipe.

Related Issues

Splits #766 (supersedes it together with the rest of the stack). Stack: #1 capture-rows ← #2 teardown-abort ← #3 rank0-tracker ← #4 (this) ← #5 hybrid-shard.

Accuracy Test

From the #766 validation run (same runtime semantics; this restructuring hardcodes the two policy defaults that run used):

• Real Qwen3-8B DSpark colocated training: 90 optimizer steps / 360 global samples on 4x H200, finite loss and gradient norms, exit code 0; sequence lengths 43–3911 under a 4096 cap, including shapes that previously raised TorchInductor NoValidChoicesError.
• CPU/unit: full suite on this branch has a failure set identical to main (pre-existing CUDA/sglang-environment failures only); new coverage in test_colocated_online.py (prompt planning disjointness/TP alignment, TP-slice detachment, bounded residency, resume-mismatch rejection, prompt-cache coordination), test_torch_environment.py, launch-plan and schema tests.

Benchmark & Profiling

From the #766 run (Qwen3-8B, DSpark, BF16, SGLang 0.5.14 local capture, SHARD_GRAD_OP, batch 1/rank, 4x H200): 22.93 samples/s steady-state (steps 50–90), 12.94 samples/s including cold start; post-compilation capture ≈11.2% of step wall time; peak 31.94/35.07 GiB allocated/reserved per rank; peak rank-local feature residency 0.179 GiB.

Checklist

• Format your code according to the Code Formatting with Pre-Commit (black --check and isort --check-only).
• Add unit tests.
• Update documentation, architecture notes, and example recipes.
• Provide throughput and memory profiling results.