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 one target batch on demand, trains its rank-local slice, and releases it before the next capture.

Design goal for this cut: colocated is a structural peer of disaggregated, and the disaggregated path is not changed. Every shared-path touch is listed under "Shared paths" below.

Structure

Concern	Disaggregated	Colocated (this PR)
Run assembly	training/disaggregated.py::build_disaggregated_run	training/colocated.py::build_colocated_run
Topology dispatch	assembly.build_training_run	same function, topology-only dispatch (_with_run_logger, _build_offline_colocated_run)
Online prompt plan	producer epoch permutation	shared training/prompt_plan.py: same per-epoch permutation and task ids; islands take strided, equal, batch-aligned shards → colocated vs disaggregated runs over the same prompts/seed are comparable at matched samples
model.sglang_* → engine args	launch_plan._sglang_argv	shared launch_plan.resolve_sglang_engine_args (CLI renderer and in-process kwargs are two views of one resolution)
Feature source	SGLangServerCaptureAdapter (Mooncake)	LocalSGLangCaptureAdapter over OfflineSGLangCaptureBackend.capture_rows (#780)
Ref queue	StreamingRefQueue	LocalRolloutStream: pull-through, lazy prompt stream, stages at most one local batch
Trainer	launch._assemble_trainer	same
Topology builder	build_disagg_online_*	build_colocated_online_runtime (now in launch.__all__; five builders as ARCHITECTURE.md states)

Modifications

• deployment.mode: local_colocated 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.
• training.tp_size becomes the target-TP island width. 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); islands are target-DP replicas with disjoint prompt shards.
• LocalRolloutStream ingests exactly one target batch per trainer request from the lazy prompt stream (no eager multi-epoch expansion), fails loud on a short tail or a partial capture, and reports interval capture time/throughput plus peak staged features through perf/*.
• Resume: a checkpoint records the plan (prompt_plan, seed, epochs, island count, dataset size, batch/TP sizes); any recorded value that differs raises instead of silently training the wrong slice (a bug in feat: add scalable colocated online training #766). The consumed prefix is skipped lazily.
• Rank-coordinated prompt-cache preparation (assembly._prepare_prompts_coordinated, generic name and docstring) prevents a tokenization process storm and Arrow-cache write races. The EAGLE3 vocab-map builder (_install_dataset_vocab_mapping) now takes a lazy counter and skips the corpus walk on a cache hit (offline benefits too).
• torch_compat.configure_flex_attention_inductor pins the inductor GEMM candidates to ATEN,TRITON for colocated FlexAttention runs only; the docstring states what was observed and that the mechanism is not isolated (torch 2.11's default already lists ATEN).
• Docs: colocated_training.md (quick-start path fixed, prompt-plan and pull-through sections), topology matrix, DESIGN.md queue contract (loader_prefetch_safe), recipe under examples/configs/online/colocated/.

Shared paths touched (behavior for disaggregated/offline)

• controller.py: del batch after the micro-step (reference release only); perf/global_samples_per_second uses world // sp_size (identical where tp_size=1, which disaggregated/offline enforce); additive peak-memory metrics and an optional runtime_metrics_provider (None for disaggregated).
• feature_dataloader.py: honors a queue's loader_prefetch_safe = False (default True), del batch after yield.
• trainer.py / launch._assemble_trainer: clone_on_fetch threaded through, default True.
• launch.py: the producer's prompt planner moved to training/prompt_plan.py unchanged (import-only change); disaggregated.py uses the shared online_prompt_seed.
• launch_plan._sglang_argv renders from resolve_sglang_engine_args; output identical.
• model_providers.py: needs_target_head includes online local_colocated.

Related Issues

Splits #766 (supersedes it together with the rest of the stack). Stack: #781 ← #782 ← #780 ← #4 (this) ← #784 hybrid-shard.

Accuracy Test

• GPU validation so far is tp_size=1 only, from the feat: add scalable colocated online training #766 run on sglang 0.5.14 (Qwen3-8B DSpark, 90 optimizer steps / 360 samples on 4x H200, finite loss and grad norms, sequence lengths 43–3911 under a 4096 cap). This rebase moves the in-process engine to sglang 0.5.18; the smoke run must be repeated before this leaves draft, and target TP > 1 islands have not been executed on hardware yet (design + CPU tests only). Both are tracked in the checklist.
• CPU: full tests/test_config tests/test_runtime tests/test_optimizer has the same failure set as main (pre-existing CUDA/sglang-environment failures), plus test_sglang_capture_hooks.py which needs sglang (CI). New coverage in test_colocated_online.py: shared planner (island shards partition the producer permutation, epoch task ids match the producer), coordinated cache preparation, run assembly (engine sizing, plan metadata, resume mismatch/prefix/completed), adapter layout and TP-slice detachment, LocalRolloutStream driven through the real RolloutWorker + LocalFeatureStore + adapter at TP=2 in lockstep (disjoint slices, bounded staging, clean status), builder wiring into _assemble_trainer; plus test_sglang_engine_args.py and test_torch_compat.py.

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 staged features 0.179 GiB. To be refreshed on 0.5.18.

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.
• Re-run the Qwen3-8B colocated smoke on sglang 0.5.18 (tp=1).
• Run a target TP ≥ 2 island configuration on GPUs before marking ready for review.

[maocheng23]

GPU validation on sglang 0.5.18 (B300, this branch merged onto main c49bab49 as maocheng/colocate-perf-base, plus #831 for the NVFP4 target head):

run	layout	steps	result
coloc-8gpu-tp1	8 ranks, training.tp_size 1 (8 islands)	300	clean exit; 18.05 samples/s steady state (Qwen3.8-27B NVFP4 target, GQA DFlash2 draft, max_length 8192, batch 2 x accum 4)
coloc-8gpu-tp2	8 ranks, training.tp_size 2 (4 TP islands)	300	clean exit; 19.07 samples/s; every TP peer trained its own slice (8 samples/rank/step), no plan or partition errors
2-GPU smoke, tp1	2 ranks	6	clean exit

Per-rank time split from the observability records: capture 48% / draft compute 52% (tp1) and 54% / 46% (tp2); trainer bubble < 0.4%. So the two open GPU gates listed in the description — 0.5.18 tp=1 smoke and a tp>=2 island run — have now been exercised; the 4x8 HSDP run for #784 is still open.

Two things noticed while running it (not blockers): with a pre-tokenized JSONL prompt source every rank re-parses the file while rank 0 holds the coordination collective (the 40k-row file took ~4 min before step 1), and TP2 islands add ~2 min of extra SGLang/JIT startup versus tp1.