feat: support offline training

[Dogacel]

Overview

Primary goal of this PR is to enable offline training by saving and using hidden states on persistent disk. Developers can now train on only 1 GPU. This opens up opportunities for people to develop Torchspec and test their implementations.

Also the test debug and inference flags were not working properly, I think we don't need them since we have offline training now. So they are removed.

Fixes #84.

Design

I've decided to implement the training code as an "inference engine" and a ray actor. So instead of enabling hf, sglang, vllm; user can directly enable the "offline" backend.

Also the generator is implemented as a separate launch script, trying to re-use existing training loop seemed to hacky. Instead of consuming the inference engine generated hidden states directy for training, they are consumed by a offline saving ray actor. This actor persists them on disk, a simple folder format.

Testing

Materialized Qwen3 8B's hidden states on 1000 training samples and trained it for 1 epoch.

Instructions to reproduce is available unde docs/offline_training.md.

---

Disclaimer

This PR is heavily AI assisted (GPT 5.6 Sol XHigh). It consists of two key commits:

1. Part 1: Core pieces to glue offline generation + loading into torchspec. I've inspected all code changes carefully.

2. Part 2: The actual implementation of the generator and the offline state loader. I've inspected all code changes lightly. I am not an expert with Mooncake and ray, so instead I've validated AI's code by running it and inspecting tests mostly.

[Copilot]

Pull request overview

This PR adds an offline replay training mode that materializes target-model outputs (hidden states, etc.) to disk and then replays them through a CPU-only “offline” inference backend, enabling draft-model training with only a single training GPU. It also removes the previously broken debug-only placement modes and updates placement-group creation to support training-only vs inference-only workflows.

Changes:

• Introduces an on-disk offline dataset format plus a materialization CLI (torchspec.offline.generate) and a replay inference engine (offline backend).
• Updates training/controller/inference plumbing to route data_id through queues and to support training-only / inference-only placement groups.
• Updates configs/docs/tests for offline training and removes debug_train_only / debug_inference_only.

Reviewed changes

Copilot reviewed 38 out of 38 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
torchspec/transfer/mooncake/store.py	Avoid CUDA stream creation for CPU-only offline flows (but currently breaks int device inputs).
torchspec/transfer/mooncake/eagle_store.py	Allows CPU tensors to bypass CUDA event/stream sync on put path.
torchspec/training/data_fetcher.py	Adds data_id to TrainSample so offline materialization can persist per-sample identities.
torchspec/train_entry.py	Adds offline mode wiring, placement-role selection, and vocab-mapping load from disk.
torchspec/ray/placement_group.py	Adds roles argument to create training-only or inference-only placement groups; removes debug-only modes.
torchspec/offline/saving_actor.py	New Ray actor that consumes normal training queues and persists Mooncake tensors to disk.
torchspec/offline/generate.py	New CLI entry point to materialize target outputs into an offline dataset.
torchspec/offline/dataset.py	New on-disk dataset implementation (manifest + per-sample .pt files) plus offline arg configuration helper.
torchspec/offline/init.py	Exposes offline dataset APIs.
torchspec/inference/factory.py	Adds offline inference_engine_type that spawns CPU replay actors.
torchspec/inference/engine/offline_replay_engine.py	New inference engine that replays disk-recorded tensors into Mooncake and returns normal engine metadata.
torchspec/inference/engine/init.py	Exports OfflineReplayEngine.
torchspec/controller/training_controller.py	Loads offline dataset IDs, preserves data_id into training queues, and adjusts eval ID behavior for offline mode.
torchspec/controller/eval.py	Makes eval cache key depend on offline dataset path when in offline mode.
torchspec/config/train_config.py	Adds offline config validation, prefixes offline_* flattening, and enforces offline-specific constraints.
torchspec/config/inference_config.py	Adds inference.offline config section (data_path, num_engines).
tests/test_placement_group.py	Adds coverage for training-only and inference-only placement group creation.
tests/test_offline_replay.py	New end-to-end-ish unit coverage for offline dataset, saver, replay engine contract, and config selection.
README.md	Adds a short “Offline Replay Training” section pointing to docs.
docs/ray.md	Updates placement group docs for role-based creation and offline/materialization modes.
docs/offline_training.md	New offline training walkthrough (materialize then train).
configs/vllm_qwen3_8b.yaml	Removes debug-only flags from example config.
configs/vllm_qwen3_8b_dflash.yaml	Removes debug-only flags from example config.
configs/vllm_qwen3_5_35b.yaml	Removes debug-only flags from example config.
configs/trtllm_qwen3_8b.yaml	Removes debug-only flags from example config.
configs/trtllm_qwen3_8b_dflash.yaml	Removes debug-only flags from example config.
configs/train_with_decode/sglang_qwen3_8b.yaml	Removes debug-only flags from example config.
configs/train_with_decode/sglang_kimi_k25_nvfp4.yaml	Removes debug-only flags from example config.
configs/sglang_qwen36_35b_dflash.yaml	Removes debug-only flags from example config.
configs/sglang_qwen3_8b.yaml	Removes debug-only flags from example config.
configs/sglang_qwen3_8b_mla_draft.yaml	Removes debug-only flags from example config.
configs/sglang_qwen3_8b_dspark.yaml	Removes debug-only flags from example config.
configs/sglang_qwen3_8b_dflash.yaml	Removes debug-only flags from example config.
configs/sglang_minimax_m25_5node.yaml	Removes debug-only flags from example config.
configs/sglang_kimi_k25_3node.yaml	Removes debug-only flags from example config.
configs/sglang_kimi_k25_2node.yaml	Removes debug-only flags from example config.
configs/README.md	Documents inference.offline and clarifies custom placement applies to online-only.
configs/hf_qwen3_8b.yaml	Removes debug-only flags from example config.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f27597cae4

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Skip existing rows before dispatching materialization

When --overwrite is not set and output_dir already has a manifest, OfflineDataset.append() later returns False for duplicate data_ids, but these calls still submit every eval/train sample to the live inference manager. Resuming an interrupted materialization therefore reruns target-model inference for all previously saved samples (and can keep stale rows if the source data changed), so large offline jobs waste most of their GPU time instead of only filling missing records.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Avoid Ray metadata in tests after the placement stub

When the full pytest session also collects tests/test_placement_group.py, that module installs a process-wide fake ray.remote in sys.modules; these new tests import Ray actors later and then dereference __ray_metadata__, which the stubbed decorator never creates, so the offline replay tests become order-dependent and fail before their assertions. Instantiate the underlying class without relying on Ray metadata, or scope/restore the Ray stub used by the placement tests.

Useful? React with 👍 / 👎.