feat: support offline training

README.md

@@ -161,6 +161,11 @@ See [examples/README.md](examples/README.md) for more details about each example
 
 ## Training Modes
 
+### Offline Replay Training
+
+Offline replay ([Docs](docs/offline_training.md)) is a training mode that separates target-output generation from draft-model training by reading hidden states from disk. Recommended for testing and development on 1 GPU.
+
+
 ### Resume vs. Continual Training
 
 Both modes use `training.load_path`, but they restore different states:

configs/README.md

@@ -35,14 +35,16 @@ python -m torchspec.train_entry --config configs/sglang_qwen3_8b.yaml training.l
 | `training` | `learning_rate`, `micro_batch_size`, `ttt_length` | Training hyperparameters |
 | `inference` | `inference_engine_type`, `inference_num_gpus` | Inference backend configuration |
 | `inference.sglang` | `tp_size`, `mem_fraction_static`, `extra_args` | SGLang engine settings (nested under inference) |
+| `inference.offline` | `data_path`, `num_engines` | Offline training settings selected by `inference_engine_type: offline` |
 | `mooncake` | `protocol`, `device_name` | Mooncake transfer engine settings |
 
-## Custom Ray placement
+## Custom online Ray placement
 
 Use `training.placement_strategy: custom` when training and inference must run
 on explicitly chosen Ray nodes. This is useful when the default `PACK` placement
 would put actors on nodes with the wrong network locality, cache state, or GPU
-partition.
+partition. Offline training and materialization are single-role workflows and
+use plain `PACK` placement.
 
 IP-based placement uses Ray's built-in `node:<ip>` resource and does not require
 custom Ray labels:

configs/hf_qwen3_8b.yaml

@@ -56,5 +56,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/sglang_kimi_k25_2node.yaml

@@ -80,8 +80,6 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false
 
 
 logging:

configs/sglang_kimi_k25_3node.yaml

@@ -77,8 +77,6 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false
 
 
 logging:

configs/sglang_minimax_m25_5node.yaml

@@ -75,5 +75,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/sglang_qwen36_35b_dflash.yaml

@@ -89,5 +89,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/sglang_qwen3_8b.yaml

@@ -61,5 +61,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/sglang_qwen3_8b_dflash.yaml

@@ -86,5 +86,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/sglang_qwen3_8b_dspark.yaml

@@ -85,5 +85,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/sglang_qwen3_8b_mla_draft.yaml

@@ -57,5 +57,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/train_with_decode/sglang_kimi_k25_nvfp4.yaml

@@ -99,5 +99,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/train_with_decode/sglang_qwen3_8b.yaml

@@ -75,5 +75,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/trtllm_qwen3_8b.yaml

@@ -76,5 +76,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/trtllm_qwen3_8b_dflash.yaml

@@ -89,5 +89,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/vllm_qwen3_5_35b.yaml

@@ -75,5 +75,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/vllm_qwen3_8b.yaml

@@ -67,5 +67,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

configs/vllm_qwen3_8b_dflash.yaml

@@ -83,5 +83,3 @@ model_download_dir: null
 
 debug:
   save_debug_train_data: null
-  debug_train_only: false
-  debug_inference_only: false

docs/offline_training.md

@@ -0,0 +1,69 @@
+# Offline Training
+
+Offline training lets you run target-model inference once, save the hidden
+states, and reuse them for draft-model training. This is useful for development
+because training no longer needs a target inference GPU, thus you can test your
+workflows on only 1 GPU. **Warning:** This workflow is not optimized for
+throughput yet.
+
+## One GPU, 1,000 samples
+
+Run these commands from the TorchSpec repository root.
+
+```bash
+conda activate torchspec
+```
+
+### 1. Save the target hidden states
+
+```bash
+python -m torchspec.offline.generate \
+  --config configs/sglang_qwen3_8b_dflash.yaml \
+  --output outputs/qwen3-8b-dflash-offline-1000 \
+  inference.inference_num_gpus=1 \
+  inference.inference_num_gpus_per_node=1
+```
+
+This runs Qwen3-8B with SGLang on one GPU. The config uses the bundled sample
+dataset and produces 1,000 replay samples. Expect the output to use about 6 GB.
+
+Materialization resumes by default. Add `--overwrite` to replace an existing
+output directory.
+
+### 2. Train from the saved data
+
+```bash
+python -m torchspec.train_entry \
+  --config configs/sglang_qwen3_8b_dflash.yaml \
+  inference.inference_engine_type=offline \
+  inference.offline.data_path=outputs/qwen3-8b-dflash-offline-1000 \
+  inference.offline.num_engines=1 \
+  training.training_num_gpus_per_node=1 \
+  output_dir=outputs/qwen3-8b-dflash-offline-dev
+```
+
+This uses one GPU for training and does not start an inference engine. The
+saved dataset can be reused for as many training runs as needed.
+
+## Use your own data
+
+Override the dataset path while materializing:
+
+```bash
+python -m torchspec.offline.generate \
+  --config configs/sglang_qwen3_8b_dflash.yaml \
+  --output /path/to/offline-data \
+  dataset.train_data_path=/path/to/train.jsonl \
+  inference.inference_num_gpus=1 \
+  inference.inference_num_gpus_per_node=1
+```
+
+Then pass that output directory to
+`inference.offline.data_path` when training.
+
+The materialization and training configs must describe the same target model,
+tokenizer, draft method, and hidden-state layout. Target-model weights must
+remain available during training because TorchSpec still uses the target
+embedding, normalization, and LM-head weights.
+
+Offline training does not currently support USP or `train_with_decode`.

docs/ray.md

@@ -31,17 +31,20 @@ torchspec/controller/
 
 ## Placement Groups
 
-Placement groups reserve GPUs for training and inference as a unit and place them on the correct nodes. `create_placement_groups(args)` is the single entry point.
+Placement groups reserve GPUs and place them on the correct nodes.
+`create_placement_groups(args, roles=...)` is the single entry point. Online
+training requests both roles, offline training defaults to training only, and
+the materialization command requests inference only.
 
 | Mode | Training GPUs | Inference GPUs | Use case |
 |------|--------------|----------------|----------|
 | Default | Sliced from unified PG | Sliced from unified PG | Production: deterministic node-to-role assignment |
-| `custom` | Sliced from custom unified PG | Sliced from custom unified PG | Production: explicit node choice with the same unified reservation semantics |
+| Online `custom` | Sliced from custom unified PG | Sliced from custom unified PG | Production: explicit node choice with the same unified reservation semantics |
 | `colocate` | Shared PG | Shared PG | Dev: share GPUs between train & inference |
-| `debug_train_only` | Dedicated PG | Empty | Debug training without inference |
-| `debug_inference_only` | Empty | Dedicated PG | Debug inference without training |
+| Offline training | Dedicated PG | None | Train from materialized target outputs |
+| Materialization | None | Dedicated PG | Produce target outputs for offline training |
 
-Each placement group probes bundles with a temporary `InfoActor` to discover the actual (node IP, GPU ID) mapping, then sorts by (node, GPU ID) for deterministic ordering. In `custom` mode, TorchSpec sorts by the configured node order first and by physical GPU ID within each selected node.
+Each placement group probes bundles with a temporary `InfoActor` to discover the actual (node IP, GPU ID) mapping, then sorts by (node, GPU ID) for deterministic ordering. `placement_strategy` applies only when online training creates both roles; single-role workflows use `PACK` placement.
 
 ## Ray Cluster Setup
 
@@ -135,7 +138,7 @@ The PACK placement strategy spreads them across nodes automatically.
 | `training.training_num_nodes` | 1 | Number of training nodes |
 | `training.training_num_gpus_per_node` | 1 | GPUs per training node |
 
-### Custom node placement
+### Custom online node placement
 
 By default, TorchSpec creates a unified placement group with Ray's `PACK`
 strategy, probes the resulting bundles, and assigns the ordered bundles to