sgl-project · maocheng23 · Aug 25, 2026
diff --git a/README.md b/README.md
@@ -19,9 +19,9 @@ We have seen many open-source projects for speculative decoding, but most of the
 
 - regularly maintained by the SpecForge team: the code is runnable out-of-the-box
 - directly compatible with SGLang: no additional porting effort is required
-- able to run online disaggregated training and both colocated and
-  disaggregated offline training through one runtime, including the supported
-  data, tensor, and sequence parallel topologies
+- able to run colocated and disaggregated training, both online and offline,
+  through one runtime, including the supported data, tensor, and sequence
+  parallel topologies
 
 
 Check out [**our documentation**](https://docs.sglang.ai/SpecForge/) to get started.

diff --git a/docs/basic_usage/colocated_training.md b/docs/basic_usage/colocated_training.md
@@ -0,0 +1,67 @@
+# Colocated online training
+
+Colocated mode loads one SGLang target shard and one FSDP draft shard in every
+trainer process. Target capture and draft training alternate on the same GPU,
+so hidden states never cross a process or node. Use it when the combined target,
+KV/cache reservation, draft, optimizer, activations, and collective workspaces
+fit with enough headroom for the longest configured sequence.
+
+Start the Qwen3-8B H200 recipe from the repository root:
+
+```bash
+specforge train --config examples/configs/qwen3-8b-dspark-colocated.yaml
+```
+
+The runtime is pull-through: the trainer requests one batch, SGLang captures
+it synchronously, the rank-local feature slice is trained, and its local store
+entry is released before the next capture. The loader skips its defensive
+clone for the rank-private in-memory store; `CLONE_ON_FETCH=1` has no effect
+here, and the loader stays synchronous regardless of
+`data.dataloader_num_workers` so all CUDA work remains on the trainer thread.
+With target TP greater than one, the capture adapter copies this rank's slice
+once to detach it from the full packed output; otherwise a small rank-local
+view would pin every peer's hidden-state storage.
+
+At each `training.log_interval`, colocated runs add capture time/throughput,
+rank-local feature-residency peaks, and accelerator allocated/reserved peaks to
+the existing `perf/*` metrics. Use `perf/data_wait_time_s` versus
+`perf/train_compute_time_s` to decide whether target TP or draft DP is the next
+scaling bottleneck.
+
+## Memory sizing
+
+`model.sglang_mem_fraction_static` is SGLang's static target budget, not an
+exclusive partition of the GPU. It covers target weights plus SGLang-managed
+KV/cache pools. The remainder is shared by the draft weights/FSDP shards,
+optimizer state, activations, temporary kernels, NCCL buffers, and allocator
+fragmentation.
+
+Measure the post-warm-up peak at the production `max_length`, local batch, and
+accumulation settings. Fixed shapes make the peak repeatable, but lazy kernel
+loading, compilation, communication buffers, and allocator fragmentation still
+make the first steps different from steady state.
+
+## Target islands
+
+`training.tp_size` is the target TP width in colocated mode. The world is split
+into contiguous target-TP islands; every island receives a deterministic,
+disjoint prompt shard. All TP peers capture the same TP-wide batch, and each
+peer trains only its contiguous local slice. The full packed output exists
+transiently during target capture; before draft training begins, each peer
+detaches its slice and releases the other slices' backing allocation.
+
+Set `model.sglang_context_length` to at least `data.max_length + 7`; size
+`sglang_max_running_requests` and `sglang_max_total_tokens` for
+`training.tp_size * training.batch_size` requests per island. Explicit values
+override the safe derived defaults.
+
+Prompt-cache preparation is coordinated: rank zero builds the tokenized Arrow
+cache first (node-local caches are then built once per remaining node) and all
+other ranks take the cache-hit path. The coordination collective runs inside
+`training.dist_timeout`, so raise that timeout or pre-build the cache for very
+large raw datasets.
+
+Colocation maximizes fixed-GPU throughput when it fits because it removes
+feature serialization and transport. Disaggregation remains preferable when a
+target cannot share memory safely, producer/consumer elastic scaling matters,
+or independent fault domains are more important than per-GPU throughput.
diff --git a/docs/basic_usage/disaggregated_training.md b/docs/basic_usage/disaggregated_training.md
@@ -16,8 +16,9 @@ trainer.
 | Online disaggregated, external | Send prompts to existing capture endpoints and publish streaming refs | Mooncake | User or scheduler owns Mooncake and SGLang |
 | Online disaggregated, managed-local | Send prompts to endpoints started from the same run config | Mooncake | SpecForge owns local Mooncake and SGLang |
 
-Online training always uses the producer/consumer topology; there is currently
-no colocated online target-inference path. `external` describes ownership, not
+This page covers the producer/consumer topology; [colocated online
+training](colocated_training.md) runs target capture in the trainer process
+instead. `external` describes ownership, not
 distance—a loopback service started by the user is still external. Conversely,
 `managed-local` still uses distinct producer and consumer roles even though one
 supervisor owns the local process tree.

diff --git a/docs/basic_usage/training.md b/docs/basic_usage/training.md
@@ -34,7 +34,7 @@ examples/configs/
 │   ├── colocated/
 │   └── disaggregated/
 └── online/
-    ├── colocated/                 # reserved; currently unsupported
+    ├── colocated/                 # in-process SGLang capture, role "all"
     └── disaggregated/
         ├── external/
         └── managed-local/
@@ -62,10 +62,12 @@ When it is greater than one, the CLI starts torch distributed itself:
 specforge train -c examples/configs/online/disaggregated/external/qwen3-30b-a3b-eagle3-online.yaml
 ```
 
-Online target inference never runs in the trainer. A patched SGLang server owns
-target parallelism and publishes captured features through Mooncake; every
-consumer rank is data parallel. Offline runs shard fixed feature references
-across trainer ranks and may additionally use EAGLE3 USP. See
+Online target inference either runs locally in each trainer process
+([colocated online training](colocated_training.md)) or on a patched external
+SGLang server that publishes through Mooncake. Colocated target TP is
+configured with `training.tp_size`; disaggregated consumers remain fully data
+parallel. Offline runs shard fixed feature references across trainer ranks and
+may additionally use EAGLE3 USP. See
 [Parallel topologies](#parallel-topologies) for the exact constraints.
 
 Paths in a config are resolved from the current working directory. The example
@@ -253,14 +255,14 @@ launching it.
 
 The unified runtime supports text training in these combinations:
 
-| Strategy | Online disaggregated | Offline colocated | Offline disaggregated |
-| --- | --- | --- | --- |
-| EAGLE3 | Yes, consumer DP | Yes, DP + USP | Yes, consumer DP |
-| DFlash | Yes, consumer DP | Yes, DP | Yes, consumer DP |
-| Domino | Yes, consumer DP | Yes, DP | Yes, consumer DP |
-| DSpark | Yes, consumer DP | Yes, DP | Yes, consumer DP |
-| MTP | Yes, consumer DP | Yes, DP | Yes, consumer DP |
-| P-EAGLE | Yes, consumer DP, batch size 1 | No | No |
+| Strategy | Online colocated | Online disaggregated | Offline colocated | Offline disaggregated |
+| --- | --- | --- | --- | --- |
+| EAGLE3 | Yes, target TP/islands | Yes, consumer DP | Yes, DP + USP | Yes, consumer DP |
+| DFlash | Yes, target TP/islands | Yes, consumer DP | Yes, DP | Yes, consumer DP |
+| Domino | Yes, target TP/islands | Yes, consumer DP | Yes, DP | Yes, consumer DP |
+| DSpark | Yes, target TP/islands | Yes, consumer DP | Yes, DP | Yes, consumer DP |
+| MTP | Yes, target TP/islands | Yes, consumer DP | Yes, DP | Yes, consumer DP |
+| P-EAGLE | Yes, batch size 1 | Yes, consumer DP, batch size 1 | No | No |
 
 Unsupported combinations fail explicitly during config validation or run
 assembly. In particular:
@@ -276,9 +278,10 @@ assembly. In particular:
 - P-EAGLE requires `training.batch_size=1` and reuses EAGLE3's server capture
   schema;
 - offline feature training supports EAGLE3, DFlash, Domino, DSpark, and MTP;
-- every online run is disaggregated and uses `model.target_backend=sglang`;
-  finite runs may omit both step fields so the producer can publish the exact
-  optimizer horizon derived from the prepared prompt plan;
+- every online run uses `model.target_backend=sglang`; disaggregated finite
+  runs may omit both step fields so the producer can publish the exact
+  optimizer horizon, while colocated runs derive it locally from the prompt
+  plan;
 - EAGLE3 offline colocated runs derive and cache a deterministic vocabulary mapping
   from the feature corpus when `model.vocab_mapping_path` is empty. EAGLE3
   disaggregated runs require an explicit shared mapping so producer and
@@ -297,9 +300,12 @@ publishes the exact schedule horizon and the consumer trains to EOF.
 
 The launcher creates every process group from the typed run config:
 
-- Online target TP/EP belongs to each external SGLang capture server, not the
-  trainer. Online consumers keep `training.tp_size` and both SP sizes at 1;
-  every trainer rank receives a disjoint feature stream.
+- Colocated online runs use `training.tp_size` as the target-island width.
+  Target-DP islands receive disjoint prompt streams; TP peers capture the same
+  TP-wide batch and train disjoint local slices.
+- Disaggregated online target TP/EP belongs to each external SGLang capture
+  server. Consumers keep `training.tp_size` and both SP sizes at 1; every
+  trainer rank receives a disjoint feature stream.
 - Offline consumers also keep `training.tp_size` at 1. Without USP, every
   trainer rank receives a disjoint reference shard and participates as data
   parallelism.
@@ -308,16 +314,18 @@ The launcher creates every process group from the typed run config:
   greater than one, USP currently uses `training.batch_size: 1`, and SP peers
   share one sequence while draft-DP groups receive disjoint references.
 
-The world size must be divisible by
-`training.sp_ulysses_size * training.sp_ring_size`. Use a shared `output_dir`
-for multi-rank checkpoints.
+The world size must be divisible by both `training.tp_size` and
+`training.sp_ulysses_size * training.sp_ring_size`. See
+[Colocated online training](colocated_training.md) for memory sizing and
+scaling guidance. Use a shared `output_dir` for multi-rank checkpoints.
 
 ## Loader and profiling controls
 
 `data.dataloader_num_workers` controls ordered background feature
 materialization. If omitted, the former trainer defaults are retained:
 EAGLE3/P-EAGLE use four workers and DFlash-family strategies use eight. Set it
-to zero for fully synchronous loading.
+to zero for fully synchronous loading. Colocated online capture stays
+synchronous regardless, keeping all CUDA work on the trainer thread.
 
 Enable a bounded, per-rank PyTorch trace without a separate profiler entry:
 

diff --git a/docs/get_started/about.md b/docs/get_started/about.md
@@ -6,8 +6,8 @@ Speculative decoding is an important and powerful technique for speeding up infe
 
 - regularly maintained by the SGLang team: the code is runnable out-of-the-box
 - directly compatible with SGLang: no additional porting effort is required
-- able to provide online disaggregated training and both colocated and disaggregated
-  offline training through one runtime, including consumer DP, offline USP,
+- able to provide colocated and disaggregated training, both online and
+  offline, through one runtime, including consumer DP, offline USP,
   evaluation, checkpoint selection, and CUDA/ROCm/Ascend portability
 
 ## ✅ SGLang-ready

diff --git a/docs/index.rst b/docs/index.rst
@@ -25,6 +25,7 @@ SpecForge is an ecosystem project developed by the SGLang team. It is a framewor
 
    basic_usage/data_preparation.md
    basic_usage/training.md
+   basic_usage/colocated_training.md
    basic_usage/disaggregated_training.md
    basic_usage/AMD/amd_rocm.md
    basic_usage/Ascend/ascend_npu.md

@@ -30,7 +30,7 @@ examples/configs/
 | --- | --- | --- | --- |
 | `offline/colocated` | Precomputed `.ckpt` files | Trainer only | Not applicable |
 | `offline/disaggregated` | Precomputed files ingested into `shared_dir` or Mooncake | Producer + consumer | Storage is configured separately; there is no `external`/`managed-local` subdivision |
-| `online/colocated` | — | — | Reserved; the current runtime does not support online colocated training |
+| `online/colocated` | Live in-process SGLang capture | Trainer only (role `all`) | SpecForge owns the in-process SGLang target shard |
 | `online/disaggregated/external` | Live SGLang capture | Producer + consumer | User or scheduler starts Mooncake and SGLang |
 | `online/disaggregated/managed-local` | Live SGLang capture | Producer + consumer | SpecForge starts local Mooncake and SGLang from `managed_local` |
 
@@ -145,6 +145,7 @@ assume the command runs from the repository root.
 
 | Workflow | Canonical starting point |
 | --- | --- |
+| DSpark online, colocated | [`online/colocated/qwen3-8b-dspark-colocated.yaml`](online/colocated/qwen3-8b-dspark-colocated.yaml) |
 | EAGLE3 offline, colocated | [`offline/colocated/qwen3-8b-eagle3-offline.yaml`](offline/colocated/qwen3-8b-eagle3-offline.yaml) |
 | DFlash offline, colocated | [`offline/colocated/qwen3-8b-dflash-offline.yaml`](offline/colocated/qwen3-8b-dflash-offline.yaml) |
 | Domino offline, colocated | [`offline/colocated/qwen3-8b-domino-offline.yaml`](offline/colocated/qwen3-8b-domino-offline.yaml) |

@@ -1,6 +1,13 @@
 # Online colocated recipes
 
-This directory is reserved so the catalog keeps the same mode/topology axes for
-online and offline recipes. All supported online recipes
-use separate producer and consumer roles under `../disaggregated/`, regardless
-of whether one supervisor starts both roles.
+Every trainer rank loads one SGLang target shard next to its FSDP draft shard
+and captures hidden states in process; no producer role, feature transport, or
+second GPU pool is involved. `training.tp_size` is the target-TP island width,
+and the trainer world size must be divisible by it.
+
+| Recipe | Target | Topology |
+| --- | --- | --- |
+| [`qwen3-8b-dspark-colocated.yaml`](qwen3-8b-dspark-colocated.yaml) | Qwen3-8B | 8 islands of TP1 on one 8xH200 node |
+
+See [Colocated online training](../../../../docs/basic_usage/colocated_training.md)
+for memory sizing and scaling guidance.
@@ -0,0 +1,51 @@
+# Qwen3-8B DSpark, online colocated — tuned baseline for one 8xH200 node.
+model:
+  target_model_path: Qwen/Qwen3-8B
+  draft_model_config: configs/qwen3-8b-dspark.json
+  target_backend: sglang
+  embedding_key: model.embed_tokens.weight
+  torch_dtype: bfloat16
+  sglang_attention_backend: fa3
+  sglang_mem_fraction_static: 0.4
+  sglang_context_length: 4103
+  sglang_disable_radix_cache: true
+  sglang_max_running_requests: 1
+  sglang_max_total_tokens: 8192
+
+data:
+  train_data_path: ./cache/dataset/perfectblend_train.jsonl
+  max_length: 4096
+  chat_template: qwen
+  build_dataset_num_proc: 32
+  dataloader_num_workers: 0
+  cache_dir: ./cache
+
+training:
+  strategy: dspark
+  num_epochs: 1
+  max_steps: 10000
+  batch_size: 1
+  accumulation_steps: 5
+  fsdp_sharding: SHARD_GRAD_OP
+  learning_rate: 6.0e-4
+  warmup_ratio: 0.04
+  max_grad_norm: 1.0
+  attention_backend: flex_attention
+  num_anchors: 512
+  loss_decay_gamma: 4.0
+  save_interval: 1000
+  log_interval: 10
+  dist_timeout: 30
+  seed: 42
+
+tracking:
+  report_to: tensorboard
+
+run_id: qwen3-8b-dspark-colocated
+output_dir: ./outputs/qwen3-8b-dspark-colocated
+
+deployment:
+  mode: local_colocated
+  trainer:
+    nnodes: 1
+    nproc_per_node: 8
diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py
@@ -270,7 +270,10 @@ def build_eagle3_model(
         kl_decay=cfg.training.kl_decay,
     ).to(device=_device(), dtype=_torch_dtype(cfg))
     needs_target_head = cfg.mode == "offline" or (
-        cfg.deployment.mode == "disaggregated" and cfg.training.role == "consumer"
+        cfg.mode == "online"
+        and (
+            cfg.deployment.mode == "local_colocated" or cfg.training.role == "consumer"
+        )
     )
     target_head = None
     if needs_target_head:
@@ -310,10 +313,8 @@ def build_peagle_model(
     # same frozen target head used by offline EAGLE3; no target model is loaded
     # in the trainer process.
     target_head = None
-    if (
-        cfg.mode == "online"
-        and cfg.deployment.mode == "disaggregated"
-        and cfg.training.role == "consumer"
+    if cfg.mode == "online" and (
+        cfg.deployment.mode == "local_colocated" or cfg.training.role == "consumer"
     ):
         from specforge.modeling.target.target_head import TargetHead
 

diff --git a/specforge/application/planning.py b/specforge/application/planning.py
@@ -139,37 +139,32 @@ def _validate_training_topology(
             "receives its own data shard"
         )
     if mode is FeatureMode.STREAMING:
-        if deployment_mode != "disaggregated":
-            raise ValueError(
-                "online training requires deployment.mode=disaggregated; "
-                "colocated online training is no longer supported"
-            )
         if cfg.model.target_backend != "sglang":
             raise ValueError(
-                "online training uses an external SGLang capture server and "
+                "online training uses SGLang target capture and "
                 "requires model.target_backend=sglang"
             )
-        deployment = cfg.deployment.disaggregated
-        if deployment is None or deployment.backend != "mooncake":
-            raise ValueError(
-                "online disaggregated training requires "
-                "deployment.disaggregated.backend=mooncake"
-            )
         if cfg.model.shard_target_output:
             raise ValueError(
-                "model.shard_target_output is unavailable with external server "
-                "capture"
-            )
-        if (
-            cfg.training.tp_size != 1
-            or cfg.training.sp_ulysses_size != 1
-            or cfg.training.sp_ring_size != 1
-        ):
-            raise ValueError(
-                "the disaggregated online consumer uses every trainer rank for "
-                "data parallelism; configure target TP on the external server and "
-                "keep training.tp_size/sp sizes at 1"
+                "model.shard_target_output is unavailable with SGLang capture"
             )
+        if deployment_mode == "disaggregated":
+            deployment = cfg.deployment.disaggregated
+            if deployment is None or deployment.backend != "mooncake":
+                raise ValueError(
+                    "online disaggregated training requires "
+                    "deployment.disaggregated.backend=mooncake"
+                )
+            if (
+                cfg.training.tp_size != 1
+                or cfg.training.sp_ulysses_size != 1
+                or cfg.training.sp_ring_size != 1
+            ):
+                raise ValueError(
+                    "the disaggregated online consumer uses every trainer rank "
+                    "for data parallelism; configure target TP on the external "
+                    "server and keep training.tp_size/sp sizes at 1"
+                )
 
     if cfg.training.attention_backend == "usp" and mode is not FeatureMode.OFFLINE:
         raise ValueError("USP attention currently requires offline features")