diff --git a/README.md b/README.md index c3e67c99..8329c825 100644 --- a/README.md +++ b/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: diff --git a/configs/README.md b/configs/README.md index 5784404c..627797dd 100644 --- a/configs/README.md +++ b/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:` resource and does not require custom Ray labels: diff --git a/configs/hf_qwen3_8b.yaml b/configs/hf_qwen3_8b.yaml index bf64e68b..bb6768d9 100644 --- a/configs/hf_qwen3_8b.yaml +++ b/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 diff --git a/configs/sglang_kimi_k25_2node.yaml b/configs/sglang_kimi_k25_2node.yaml index b7c5a152..0c954888 100644 --- a/configs/sglang_kimi_k25_2node.yaml +++ b/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: diff --git a/configs/sglang_kimi_k25_3node.yaml b/configs/sglang_kimi_k25_3node.yaml index b7cfc5ba..97cd1f6d 100644 --- a/configs/sglang_kimi_k25_3node.yaml +++ b/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: diff --git a/configs/sglang_minimax_m25_5node.yaml b/configs/sglang_minimax_m25_5node.yaml index 87905418..c198b305 100644 --- a/configs/sglang_minimax_m25_5node.yaml +++ b/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 diff --git a/configs/sglang_qwen36_35b_dflash.yaml b/configs/sglang_qwen36_35b_dflash.yaml index 9e89948a..53c367eb 100644 --- a/configs/sglang_qwen36_35b_dflash.yaml +++ b/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 diff --git a/configs/sglang_qwen3_8b.yaml b/configs/sglang_qwen3_8b.yaml index 8e0975bf..315b8685 100644 --- a/configs/sglang_qwen3_8b.yaml +++ b/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 diff --git a/configs/sglang_qwen3_8b_dflash.yaml b/configs/sglang_qwen3_8b_dflash.yaml index 724ce0af..58b03054 100644 --- a/configs/sglang_qwen3_8b_dflash.yaml +++ b/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 diff --git a/configs/sglang_qwen3_8b_dspark.yaml b/configs/sglang_qwen3_8b_dspark.yaml index 136bc273..a838a517 100644 --- a/configs/sglang_qwen3_8b_dspark.yaml +++ b/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 diff --git a/configs/sglang_qwen3_8b_mla_draft.yaml b/configs/sglang_qwen3_8b_mla_draft.yaml index 023bae46..afce54ce 100644 --- a/configs/sglang_qwen3_8b_mla_draft.yaml +++ b/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 diff --git a/configs/train_with_decode/sglang_kimi_k25_nvfp4.yaml b/configs/train_with_decode/sglang_kimi_k25_nvfp4.yaml index c232fdb8..ca263701 100644 --- a/configs/train_with_decode/sglang_kimi_k25_nvfp4.yaml +++ b/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 diff --git a/configs/train_with_decode/sglang_qwen3_8b.yaml b/configs/train_with_decode/sglang_qwen3_8b.yaml index ecce9049..53525db7 100644 --- a/configs/train_with_decode/sglang_qwen3_8b.yaml +++ b/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 diff --git a/configs/trtllm_qwen3_8b.yaml b/configs/trtllm_qwen3_8b.yaml index b505f76e..c438c564 100644 --- a/configs/trtllm_qwen3_8b.yaml +++ b/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 diff --git a/configs/trtllm_qwen3_8b_dflash.yaml b/configs/trtllm_qwen3_8b_dflash.yaml index 5d829ab4..d5c9e7df 100644 --- a/configs/trtllm_qwen3_8b_dflash.yaml +++ b/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 diff --git a/configs/vllm_qwen3_5_35b.yaml b/configs/vllm_qwen3_5_35b.yaml index 9ca3beab..cc286dd1 100644 --- a/configs/vllm_qwen3_5_35b.yaml +++ b/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 diff --git a/configs/vllm_qwen3_8b.yaml b/configs/vllm_qwen3_8b.yaml index bdb50faa..ff7e73e2 100644 --- a/configs/vllm_qwen3_8b.yaml +++ b/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 diff --git a/configs/vllm_qwen3_8b_dflash.yaml b/configs/vllm_qwen3_8b_dflash.yaml index 5e59ee1c..13d9cf9f 100644 --- a/configs/vllm_qwen3_8b_dflash.yaml +++ b/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 diff --git a/docs/offline_training.md b/docs/offline_training.md new file mode 100644 index 00000000..a36bd753 --- /dev/null +++ b/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`. diff --git a/docs/ray.md b/docs/ray.md index 896b42ab..81444b31 100644 --- a/docs/ray.md +++ b/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 diff --git a/tests/test_offline_replay.py b/tests/test_offline_replay.py new file mode 100644 index 00000000..3decd3ae --- /dev/null +++ b/tests/test_offline_replay.py @@ -0,0 +1,412 @@ +"""Tests for the offline dataset, saver, and replay contract.""" + +from argparse import Namespace +from queue import Queue +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from torchspec.offline.dataset import ( + OFFLINE_SCHEMA_VERSION, + OfflineDataset, +) + + +def _args(tmp_path, **overrides): + values = { + "offline_data_path": str(tmp_path), + "target_model_path": "target/model", + "target_model_backend": "sglang", + "inference_engine_type": "sgl", + "last_hidden_states_prenorm": False, + "chat_template": "llama3", + "max_seq_length": 128, + "aux_hidden_states_layers": [1, 2, 3], + "attention_backend": "sdpa", + "sp_ulysses_size": 1, + "sp_ring_size": 1, + "ttt_length": 7, + } + values.update(overrides) + return Namespace(**values) + + +def _tensors(seq_len=4): + return { + "input_ids": torch.arange(seq_len, dtype=torch.int64).unsqueeze(0), + "hidden_states": torch.randn(1, seq_len, 6, dtype=torch.bfloat16), + "last_hidden_states": torch.randn(1, seq_len, 2, dtype=torch.bfloat16), + "target": None, + } + + +def _write_dataset(tmp_path): + args = _args(tmp_path) + writer = OfflineDataset(tmp_path, create=True, last_hidden_states_prenorm=False) + writer.append( + "train", + data_id="train-1", + tensors=_tensors(), + packed_loss_mask="2,2", + metadata={"source": "test"}, + ) + writer.append( + "eval", + data_id="eval-1", + tensors=_tensors(3), + packed_loss_mask="1,2", + ) + return args, writer + + +def test_offline_dataset_round_trip_and_resume(tmp_path): + _args_obj, writer = _write_dataset(tmp_path) + + assert writer.count("train") == 1 + assert not writer.append( + "train", + data_id="train-1", + tensors=_tensors(), + packed_loss_mask="2,2", + ) + + resumed = OfflineDataset(tmp_path) + assert resumed.count("train") == 1 + assert resumed.metadata["version"] == OFFLINE_SCHEMA_VERSION + assert (tmp_path / "manifest.jsonl").is_file() + assert not (tmp_path / "train" / "manifest.jsonl").exists() + record = resumed.load(resumed.rows("train")[0]) + assert record["data_id"] == "train-1" + assert record["packed_loss_mask"] == "2,2" + assert set(record) >= {"input_ids", "hidden_states", "last_hidden_states"} + assert "input_ids_cpu" not in record + + +def test_offline_dataset_accepts_dflash_hidden_states_only(tmp_path): + writer = OfflineDataset(tmp_path, create=True, last_hidden_states_prenorm=False) + tensors = _tensors() + tensors["last_hidden_states"] = None + + assert writer.append( + "train", + data_id="dflash-1", + tensors=tensors, + packed_loss_mask=None, + ) + + record = OfflineDataset(tmp_path).load("dflash-1") + assert set(record) >= {"input_ids", "hidden_states"} + assert "last_hidden_states" not in record + assert "target" not in record + + +def test_offline_dataset_rejects_cross_split_ids_and_missing_tensors(tmp_path): + _args_obj, writer = _write_dataset(tmp_path) + + with pytest.raises(ValueError, match="already exists in another split"): + writer.append( + "eval", + data_id="train-1", + tensors=_tensors(), + packed_loss_mask="2,2", + ) + + with pytest.raises(ValueError, match="missing tensors"): + writer.append( + "train", + data_id="broken", + tensors={"input_ids": torch.ones(1, 2)}, + packed_loss_mask=None, + ) + + +class _FakeMooncakeStore: + def __init__(self): + self.put_calls = [] + self.flush_count = 0 + + @staticmethod + def _meta(kwargs): + tensors = { + key: value + for key, value in kwargs.items() + if key in ("input_ids", "hidden_states", "target", "last_hidden_states") + and isinstance(value, torch.Tensor) + } + return { + "shapes": {key: tuple(value.shape) for key, value in tensors.items()}, + "dtypes": {key: value.dtype for key, value in tensors.items()}, + } + + def put(self, **kwargs): + self.put_calls.append(kwargs) + return self._meta(kwargs) + + def flush(self): + self.flush_count += 1 + + +def _make_replay_engine(tmp_path, args): + from torchspec.inference.engine.offline_replay_engine import OfflineReplayEngine + + with patch("torchspec.inference.engine.offline_replay_engine.setup_file_logging"): + engine = OfflineReplayEngine(args, rank=0) + dataset = OfflineDataset(tmp_path) + engine._dataset = dataset + rows = dataset.rows("train") + dataset.rows("eval") + engine._rows_by_id = {row["data_id"]: row for row in rows} + engine._mooncake_store = _FakeMooncakeStore() + return engine + + +def test_replay_engine_returns_live_engine_contract(tmp_path): + args, _writer = _write_dataset(tmp_path) + engine = _make_replay_engine(tmp_path, args) + + outputs = engine.generate(data_id=["train-1"]) + + assert len(outputs) == 1 + output = outputs[0] + assert output["mooncake_key"] + assert output["packed_loss_mask"] == "2,2" + assert output["metadata"]["source"] == "test" + assert output["tensor_shapes"]["input_ids"] == (1, 4) + assert len(engine._mooncake_store.put_calls) == 1 + assert engine._mooncake_store.flush_count == 1 + + +def test_replay_engine_rejects_unknown_data_id(tmp_path): + args, _writer = _write_dataset(tmp_path) + engine = _make_replay_engine(tmp_path, args) + + with pytest.raises(KeyError, match="missing"): + engine.generate(data_id="missing") + + +def test_saving_actor_consumes_normal_train_sample(tmp_path): + from torchspec.offline.saving_actor import OfflineSavingActor + from torchspec.training.data_fetcher import TrainSample + + actor_class = OfflineSavingActor.__ray_metadata__.modified_class + actor = object.__new__(actor_class) + actor.dataset = MagicMock() + actor.dataset.append.return_value = True + actor.store = MagicMock() + actor.store.get.return_value.to_tensor_dict.return_value = _tensors() + actor.queues = {"train": Queue()} + actor.queues["train"].put( + TrainSample( + data_id="train-1", + mooncake_key="key", + tensor_shapes={"input_ids": (1, 4), "hidden_states": (1, 4, 6)}, + tensor_dtypes={"input_ids": "int64"}, + packed_loss_mask="2,2", + metadata={"source": "test"}, + ) + ) + + assert actor.save_from_queue("train") == 1 + assert actor.dataset.append.call_args.kwargs["data_id"] == "train-1" + actor.store.remove_eagle3_tensors.assert_called_once() + + +def test_controller_sources_manifest_ids_without_retokenizing(tmp_path): + args, _writer = _write_dataset(tmp_path) + args.inference_engine_type = "offline" + args.per_dp_rank_batch_size = 1 + args.shuffle_dataset = False + args.seed = 0 + + from torchspec.controller.training_controller import AsyncTrainingController + + controller_class = AsyncTrainingController.__ray_metadata__.modified_class + with patch("torchspec.controller.training_controller.Queue", side_effect=lambda: MagicMock()): + controller = controller_class(args, dp_size=1) + + assert controller.load_dataset(args) == 1 + assert controller.load_eval_dataset(args) == 1 + assert controller._stored_dataset == [ + {"data_id": "train-1", "metadata": {"offline_replay": True}} + ] + + controller.submit_eval_chunk(0, 1) + entry = controller.prompt_buffer.popleft() + assert entry.data_id == "eval-1" + assert entry.input_ids is None + + +def test_controller_common_loader_handles_online_train_and_eval(tmp_path): + from torchspec.controller.training_controller import AsyncTrainingController + + args = _args( + tmp_path, + train_data_path="train.jsonl", + eval_data_path="eval.jsonl", + eval_prompt_key="eval_prompt", + per_dp_rank_batch_size=1, + ) + controller_class = AsyncTrainingController.__ray_metadata__.modified_class + with patch("torchspec.controller.training_controller.Queue", side_effect=lambda: MagicMock()): + controller = controller_class(args, dp_size=2) + + train_rows = [{"data_id": "train-1"}] + eval_rows = [{"data_id": f"eval-{i}"} for i in range(3)] + with patch( + "torchspec.data.dataset.load_conversation_dataset", + side_effect=[train_rows, eval_rows], + ) as load_conversation_dataset: + assert controller.load_dataset(args) == 1 + assert controller.load_eval_dataset(args) == 2 + + train_args = load_conversation_dataset.call_args_list[0].args[0] + eval_args = load_conversation_dataset.call_args_list[1].args[0] + assert train_args is args + assert eval_args is not args + assert eval_args.train_data_path == "eval.jsonl" + assert eval_args.prompt_key == "eval_prompt" + + +def test_controller_preserves_data_id_in_train_queue(tmp_path): + from torchspec.controller.training_controller import AsyncTrainingController + from torchspec.utils.types import InferenceOutput + + args, _writer = _write_dataset(tmp_path) + args.per_dp_rank_batch_size = 1 + controller_class = AsyncTrainingController.__ray_metadata__.modified_class + with patch("torchspec.controller.training_controller.Queue", side_effect=lambda: MagicMock()): + controller = controller_class(args, dp_size=1) + + controller._dispatch_to_queues( + [ + InferenceOutput( + data_id="source-id", + mooncake_key="key", + tensor_shapes={}, + tensor_dtypes={}, + ) + ], + controller.train_queues, + ) + + sample = controller.train_queues[0].put.call_args.args[0] + assert sample.data_id == "source-id" + + +def test_factory_creates_cpu_replay_actors(tmp_path): + args, _writer = _write_dataset(tmp_path) + args.offline_num_engines = 2 + actor_class = MagicMock() + actor_class.options.return_value.remote.side_effect = [MagicMock(), MagicMock()] + + with ( + patch("torchspec.inference.factory.ray.remote", return_value=actor_class), + patch("torchspec.inference.factory.get_torchspec_env_vars", return_value={}), + ): + from torchspec.inference.factory import _prepare_offline_replay_engines + + engines, refs = _prepare_offline_replay_engines(args, MagicMock()) + + assert len(engines) == 2 + assert len(refs) == 2 + for call in actor_class.options.call_args_list: + assert call.kwargs["num_gpus"] == 0 + assert call.kwargs["num_cpus"] == 1 + + +def test_offline_training_is_selected_by_engine_type(tmp_path): + from torchspec.config.train_config import config_to_flat_args, load_config + + config = load_config( + cli_args=[ + "inference.inference_engine_type=offline", + f"inference.offline.data_path={tmp_path}", + "inference.offline.num_engines=3", + "dataset.defer_tokenization=true", + ] + ) + args = config_to_flat_args(config) + + assert args.inference_engine_type == "offline" + assert args.offline_data_path == str(tmp_path) + assert args.offline_num_engines == 3 + assert not hasattr(args, "offline_enabled") + assert args.last_hidden_states_prenorm is None + assert args.defer_tokenization is False + assert args.dynamic_loss_mask is False + + +def test_inference_manager_uses_its_normal_input_path_for_offline(tmp_path): + from torchspec.controller.inference_manager import AsyncInferenceManager + from torchspec.utils.types import InferenceInput + + args = _args(tmp_path, inference_engine_type="offline", defer_tokenization=False) + manager_class = AsyncInferenceManager.__ray_metadata__.modified_class + manager = manager_class( + args, + controller=MagicMock(), + inference_engines=[MagicMock()], + ) + entry = InferenceInput(data_id="train-1") + + with patch( + "torchspec.controller.inference_manager.ray.put", return_value="input-ref" + ) as ray_put: + inputs = manager._prepare_engine_inputs([entry]) + + ray_put.assert_called_once_with([None]) + assert inputs["input_ids_ref"] == "input-ref" + + +def test_offline_rejects_usp(tmp_path): + from torchspec.config.train_config import load_config + + with pytest.raises(ValueError, match="usp is not supported offline"): + load_config( + cli_args=[ + "inference.inference_engine_type=offline", + f"inference.offline.data_path={tmp_path}", + "training.attention_backend=usp", + ] + ) + + +def test_mooncake_put_accepts_cpu_replay_tensors_without_cuda_events(): + from torchspec.transfer.mooncake.eagle_store import EagleMooncakeStore + + store = object.__new__(EagleMooncakeStore) + store._gpu_direct_available = False + store._gpu_send_buffer = None + buffer = MagicMock(ptr=123) + store._host_buffer_pool = MagicMock() + store._host_buffer_pool.get_buffer.return_value = buffer + store._async_put_manager = MagicMock() + store._stage_tensors_into_buffer = MagicMock(return_value=([123], [8])) + + store._put_raw_tensors(["key"], [torch.ones(2)]) + + store._async_put_manager.submit.assert_called_once_with(["key"], [123], [8], 123) + + +def test_mooncake_cpu_client_does_not_create_cuda_stream(): + from torchspec.config.mooncake_config import MooncakeConfig + from torchspec.transfer.mooncake.store import MooncakeHiddenStateStore + + client = MagicMock() + client.setup.return_value = 0 + client.batch_remove.__doc__ = "batch_remove(keys, force=False)" + config = MooncakeConfig(async_put_pool_size=0, enable_gpu_direct=False) + store = MooncakeHiddenStateStore(config) + + with ( + patch( + "torchspec.transfer.mooncake.store.MooncakeDistributedStore", + return_value=client, + ), + patch("torchspec.transfer.mooncake.store.torch.cuda.is_available", return_value=True), + patch("torchspec.transfer.mooncake.store.torch.cuda.Stream") as cuda_stream, + ): + store.setup(device=torch.device("cpu")) + + cuda_stream.assert_not_called() + store.close() diff --git a/tests/test_placement_group.py b/tests/test_placement_group.py index 17c52ee6..1bed15f7 100644 --- a/tests/test_placement_group.py +++ b/tests/test_placement_group.py @@ -76,8 +76,6 @@ def __init__(self, **kwargs): def _make_args(**overrides): defaults = dict( placement_strategy="training_first", - debug_train_only=False, - debug_inference_only=False, colocate=False, training_num_nodes=1, training_num_gpus_per_node=2, @@ -280,3 +278,64 @@ def test_custom_colocate_uses_training_topology_for_inference_constraints(): assert kwargs["node_group_indices"] == [0, 0, 0, 0, 1, 1, 1, 1] assert result["training"] == (fake_pg, list(range(8)), list(range(8))) assert result["inference"] == (fake_pg, list(range(8)), list(range(8))) + + +def test_inference_only_placement_reserves_only_inference_gpus(): + args = _make_args(inference_num_gpus=2) + fake_pg = MagicMock(name="pg") + + with ( + patch("torchspec.ray.placement_group._ensure_ray_initialized"), + patch("torchspec.ray.placement_group._wait_for_gpu_resources") as wait_for_gpus, + patch( + "torchspec.ray.placement_group._create_placement_group", + return_value=(fake_pg, [0, 1], [0, 1]), + ) as create_pg, + ): + result = create_placement_groups(args, roles={"inference"}) + + wait_for_gpus.assert_called_once_with(2) + create_pg.assert_called_once_with(2, strategy="PACK", name="inference_pg") + assert result["training"] == (fake_pg, [], []) + assert result["inference"] == (fake_pg, [0, 1], [0, 1]) + + +def test_training_only_placement_ignores_role_ordering_strategy(): + args = _make_args( + inference_engine_type="offline", + placement_strategy="inference_first", + inference_num_gpus=8, + ) + fake_pg = MagicMock(name="pg") + + with ( + patch("torchspec.ray.placement_group._ensure_ray_initialized"), + patch("torchspec.ray.placement_group._wait_for_gpu_resources") as wait_for_gpus, + patch( + "torchspec.ray.placement_group._create_placement_group", + return_value=(fake_pg, [0, 1], [0, 1]), + ) as create_pg, + ): + result = create_placement_groups(args, roles={"training"}) + + wait_for_gpus.assert_called_once_with(2) + create_pg.assert_called_once_with(2, strategy="PACK", name="training_pg") + assert result["training"] == (fake_pg, [0, 1], [0, 1]) + assert result["inference"] == (fake_pg, [], []) + + +def test_single_role_placement_rejects_custom_strategy(): + args = _make_args( + inference_engine_type="offline", + placement_strategy="custom", + training_node_ips=["10.0.0.1"], + ) + + with ( + patch("torchspec.ray.placement_group._ensure_ray_initialized"), + patch("torchspec.ray.placement_group._wait_for_gpu_resources") as wait_for_gpus, + pytest.raises(ValueError, match="only supported when training and inference"), + ): + create_placement_groups(args, roles={"training"}) + + wait_for_gpus.assert_not_called() diff --git a/torchspec/config/inference_config.py b/torchspec/config/inference_config.py index 82218390..0e3a7d32 100644 --- a/torchspec/config/inference_config.py +++ b/torchspec/config/inference_config.py @@ -133,6 +133,14 @@ class TrtllmConfig: extra_args: Dict[str, Any] = field(default_factory=dict) +@dataclass +class OfflineTrainingConfig: + """Configuration for training from materialized target outputs.""" + + data_path: Optional[str] = None + num_engines: int = 1 + + @dataclass class InferenceConfig: aux_hidden_states_layers: Optional[list] = None @@ -148,6 +156,7 @@ class InferenceConfig: last_hidden_states_prenorm: Optional[bool] = None max_sample_pool_size: int = 0 store_last_hidden_states: bool = True + offline: OfflineTrainingConfig = field(default_factory=OfflineTrainingConfig) sglang: SGLangConfig = field(default_factory=SGLangConfig) vllm: VllmConfig = field(default_factory=VllmConfig) trtllm: TrtllmConfig = field(default_factory=TrtllmConfig) diff --git a/torchspec/config/train_config.py b/torchspec/config/train_config.py index ba50bfbc..0995f393 100644 --- a/torchspec/config/train_config.py +++ b/torchspec/config/train_config.py @@ -48,8 +48,6 @@ class DatasetConfig: @dataclass class DebugConfig: - debug_inference_only: bool = False - debug_train_only: bool = False enable_perf_metrics: bool = True max_dump_steps: int = 5 memory_recorder: str = "torch" @@ -202,7 +200,12 @@ class Config: output_dir: str = "" -_ALWAYS_LOCAL_PATH_KEYS = ("output_dir", "cache_dir", "model_download_dir") +_ALWAYS_LOCAL_PATH_KEYS = ( + "output_dir", + "cache_dir", + "model_download_dir", + "inference.offline.data_path", +) _DATA_PATH_KEYS = ("dataset.train_data_path", "dataset.eval_data_path") @@ -251,6 +254,21 @@ def _validate_vllm_config(config: DictConfig) -> None: raise NotImplementedError(f"{label} is not yet supported with the vllm backend!") +def _validate_offline_config(config: DictConfig) -> None: + if config.inference.inference_engine_type == "offline": + if not config.inference.offline.data_path: + raise ValueError( + "inference.offline.data_path is required when " + "inference.inference_engine_type=offline" + ) + if config.training.train_with_decode: + raise ValueError("training.train_with_decode is not supported in offline mode") + if config.training.attention_backend == "usp": + raise ValueError("training.attention_backend=usp is not supported offline") + if config.inference.offline.num_engines <= 0: + raise ValueError("inference.offline.num_engines must be positive") + + def _save_config_snapshot(config: DictConfig) -> None: """Save the resolved config to output_dir/config.yaml if output_dir is set.""" output_dir = OmegaConf.select(config, "output_dir", default=None) @@ -295,6 +313,7 @@ def load_config( _resolve_relative_paths(config, os.getcwd()) _validate_vllm_config(config) + _validate_offline_config(config) if save_snapshot: _save_config_snapshot(config) @@ -306,6 +325,7 @@ def load_config( _PREFIXED_SECTIONS = { "decode": "decode_", "mooncake": "mooncake_", + "offline": "offline_", "sglang": "sglang_", "vllm": "vllm_", "trtllm": "trtllm_", @@ -343,6 +363,9 @@ def _add(key: str, val: Any, origin: str) -> None: # --- Computed / alias fields --- flat["world_size"] = flat["training_num_nodes"] * flat["training_num_gpus_per_node"] flat["rank"] = 0 + if flat.get("inference_engine_type") == "offline": + # Replay records are already tokenized and carry their loss masks. + flat["defer_tokenization"] = False flat["dynamic_loss_mask"] = flat["defer_tokenization"] and not flat["train_with_decode"] flat["use_wandb"] = flat.get("use_wandb", False) or flat.get("report_to") == "wandb" flat["use_tensorboard"] = ( @@ -354,7 +377,9 @@ def _add(key: str, val: Any, origin: str) -> None: if flat.get("continual_training") and not flat.get("load_path"): logger.warning("continual_training=True but no training.load_path was provided") - if "last_hidden_states_prenorm" not in flat or flat["last_hidden_states_prenorm"] is None: + if ( + "last_hidden_states_prenorm" not in flat or flat["last_hidden_states_prenorm"] is None + ) and flat.get("inference_engine_type") != "offline": flat["last_hidden_states_prenorm"] = flat.get("inference_engine_type") == "vllm" return argparse.Namespace(**flat) diff --git a/torchspec/controller/eval.py b/torchspec/controller/eval.py index 24b90eb8..bae435a2 100644 --- a/torchspec/controller/eval.py +++ b/torchspec/controller/eval.py @@ -180,8 +180,13 @@ def setup_eval(controller, train_group, args, eval_dataset_size: int) -> EvalSet if eval_enabled: cache_dir = os.path.abspath(getattr(args, "cache_dir", "./cache")) + data_source = ( + getattr(args, "offline_data_path", "") + if getattr(args, "inference_engine_type", None) == "offline" + else getattr(args, "eval_data_path", "") + ) cache_key = hashlib.md5( - f"{getattr(args, 'eval_data_path', '')}|" + f"{data_source}|" f"{getattr(args, 'target_model_path', '')}|" f"{getattr(args, 'max_seq_length', 0)}".encode() ).hexdigest()[:12] diff --git a/torchspec/controller/training_controller.py b/torchspec/controller/training_controller.py index 7298db91..30b74c18 100644 --- a/torchspec/controller/training_controller.py +++ b/torchspec/controller/training_controller.py @@ -204,12 +204,42 @@ def add_dataset(self, dataset: list) -> int: self.prompt_buffer.append(entry) return len(dataset) - def load_dataset(self, args) -> int: - """Load and store dataset on the controller for later use.""" + def _load_dataset_split(self, args, split: str) -> list: + """Load one split from either replay records or conversation data.""" + if split not in ("train", "eval"): + raise ValueError(f"Unknown dataset split: {split!r}") + + if getattr(args, "inference_engine_type", None) == "offline": + from torchspec.offline.dataset import OfflineDataset + + dataset = OfflineDataset(args.offline_data_path) + return [ + {"data_id": str(row["data_id"]), "metadata": {"offline_replay": True}} + for row in dataset.rows(split) + ] + + data_path = ( + args.train_data_path if split == "train" else getattr(args, "eval_data_path", None) + ) + if not data_path: + return [] + from torchspec.data.dataset import load_conversation_dataset - self._stored_dataset = load_conversation_dataset(args) + dataset_args = args + if split == "eval": + dataset_args = copy.copy(args) + dataset_args.train_data_path = data_path + if getattr(args, "eval_prompt_key", None): + dataset_args.prompt_key = args.eval_prompt_key + return load_conversation_dataset(dataset_args) + + def load_dataset(self, args) -> int: + """Load and store the training dataset for later epochs.""" + self._stored_dataset = self._load_dataset_split(args, "train") if not self._stored_dataset: + if getattr(args, "inference_engine_type", None) == "offline": + raise ValueError("Offline dataset has no train samples") raise ValueError( f"Training dataset is empty after processing. " f"Check train_data_path='{args.train_data_path}', " @@ -269,18 +299,7 @@ def reload_dataset(self) -> int: def load_eval_dataset(self, args) -> int: """Load eval dataset on the controller and store it. Returns size (0 if none).""" - eval_data_path = getattr(args, "eval_data_path", None) - if not eval_data_path: - return 0 - - from torchspec.data.dataset import load_conversation_dataset - - eval_args = copy.copy(args) - eval_args.train_data_path = eval_data_path - eval_prompt_key = getattr(args, "eval_prompt_key", None) - if eval_prompt_key: - eval_args.prompt_key = eval_prompt_key - raw_dataset = load_conversation_dataset(eval_args) + raw_dataset = self._load_dataset_split(args, "eval") raw_count = len(raw_dataset) # Truncate to a multiple of dp_size so every dispatch is a full batch usable = (raw_count // self.dp_size) * self.dp_size @@ -289,9 +308,9 @@ def load_eval_dataset(self, args) -> int: f"Eval dataset truncated from {raw_count} to {usable} samples " f"(dp_size={self.dp_size})" ) - self._stored_eval_dataset = raw_dataset[:usable] if usable > 0 else [] + self._stored_eval_dataset = raw_dataset[:usable] count = len(self._stored_eval_dataset) - logger.info(f"Controller loaded eval dataset: {count} samples from {eval_data_path}") + logger.info(f"Controller loaded eval dataset: {count} samples") return count def get_dataset_size(self) -> int: @@ -524,6 +543,7 @@ def _dispatch_to_queues( packed_loss_mask=result.packed_loss_mask, last_turn_loss_only=last_turn_loss_only, metadata=metadata, + data_id=result.data_id, ) if self.sp_size > 1 and len(queues) == self.queue_count: start = dp_rank * self.sp_size @@ -555,7 +575,11 @@ def _build_eval_entries(self, dataset: list) -> list[InferenceInput]: for sample in dataset: if isinstance(sample, dict): raw_id = sample.get("data_id") or self._generate_data_id() - data_id = f"eval_{raw_id}" + data_id = ( + str(raw_id) + if getattr(self.args, "inference_engine_type", None) == "offline" + else f"eval_{raw_id}" + ) self._eval_data_ids.add(data_id) entry = InferenceInput( data_id=data_id, diff --git a/torchspec/inference/engine/__init__.py b/torchspec/inference/engine/__init__.py index 635b4f6c..ea144c96 100644 --- a/torchspec/inference/engine/__init__.py +++ b/torchspec/inference/engine/__init__.py @@ -21,10 +21,12 @@ from torchspec.inference.engine.base import InferenceEngine from torchspec.inference.engine.hf_engine import HFEngine +from torchspec.inference.engine.offline_replay_engine import OfflineReplayEngine __all__ = [ "InferenceEngine", "HFEngine", + "OfflineReplayEngine", ] try: diff --git a/torchspec/inference/engine/offline_replay_engine.py b/torchspec/inference/engine/offline_replay_engine.py new file mode 100644 index 00000000..86844062 --- /dev/null +++ b/torchspec/inference/engine/offline_replay_engine.py @@ -0,0 +1,140 @@ +"""Inference-engine adapter that replays materialized target tensors from disk.""" + +from __future__ import annotations + +import dataclasses +import uuid +from typing import Any + +import ray +import torch + +from torchspec.inference.engine.base import InferenceEngine +from torchspec.offline.dataset import OfflineDataset, configure_offline_args +from torchspec.ray.ray_actor import RayActor +from torchspec.transfer.mooncake.eagle_store import EagleMooncakeStore +from torchspec.utils.logging import logger, setup_file_logging + + +class OfflineReplayEngine(InferenceEngine, RayActor): + """Serve recorded target-model outputs through the normal Mooncake contract. + + The async inference manager treats this actor like any other inference + engine. Instead of running a target model, ``generate`` looks records up by + data ID, writes their tensors to Mooncake, and returns the same metadata as + a live engine. + """ + + def __init__(self, args, rank: int, engine_group: int = 0, **_kwargs) -> None: + self.args = args + self.rank = rank + self._mooncake_config = None + self._mooncake_store: EagleMooncakeStore | None = None + self._dataset: OfflineDataset | None = None + self._rows_by_id: dict[str, dict[str, Any]] = {} + setup_file_logging("offline_replay", rank, group=engine_group) + + def init(self, mooncake_config=None) -> None: + if mooncake_config is None: + raise ValueError("OfflineReplayEngine requires a Mooncake configuration") + + from torchspec.transfer.mooncake.utils import check_mooncake_master_available + + mooncake_config = dataclasses.replace( + mooncake_config, + local_hostname=self.get_node_ip(), + enable_gpu_direct=False, + ) + check_mooncake_master_available( + mooncake_config.master_server_address, + mooncake_config.metadata_server, + ) + self._mooncake_config = mooncake_config + self._dataset = OfflineDataset(self.args.offline_data_path) + configure_offline_args(self._dataset, self.args) + rows = self._dataset.rows("train") + self._dataset.rows("eval") + self._rows_by_id = {str(row["data_id"]): row for row in rows} + if len(self._rows_by_id) != len(rows): + raise ValueError("Offline dataset contains duplicate data IDs across splits") + + self._mooncake_store = EagleMooncakeStore(mooncake_config) + self._mooncake_store.setup(device=torch.device("cpu")) + logger.info( + "OfflineReplayEngine rank %d initialized with %d records from %s", + self.rank, + len(rows), + self.args.offline_data_path, + ) + + def _replay_one(self, data_id: str) -> dict[str, Any]: + if self._dataset is None or self._mooncake_store is None: + raise RuntimeError("OfflineReplayEngine not initialized. Call init() first.") + row = self._rows_by_id.get(str(data_id)) + if row is None: + raise KeyError(f"Offline dataset has no record for data_id={data_id!r}") + record = self._dataset.load(row) + key = str(uuid.uuid4()) + tensors = { + name: record.get(name) + for name in ("input_ids", "hidden_states", "target", "last_hidden_states") + } + + store_meta = self._mooncake_store.put( + key=key, + hidden_states=tensors["hidden_states"], + input_ids=tensors["input_ids"], + last_hidden_states=tensors["last_hidden_states"], + target=tensors["target"], + ) + + metadata = dict(record.get("metadata") or {}) + return { + "mooncake_key": key, + "tensor_shapes": store_meta["shapes"], + "tensor_dtypes": store_meta["dtypes"], + "packed_loss_mask": record.get("packed_loss_mask"), + "metadata": metadata, + } + + def generate( + self, + data_id: str | list[str], + input_ids_ref: ray.ObjectRef | list[torch.Tensor] | None = None, + packed_loss_mask_list: list[str] | None = None, + formatted_prompts: list[str] | None = None, + return_last_hidden_states: bool = False, + return_logits: bool = True, + multimodal_inputs: list[dict] | None = None, + ) -> list[dict[str, Any]]: + del ( + input_ids_ref, + packed_loss_mask_list, + formatted_prompts, + return_last_hidden_states, + return_logits, + multimodal_inputs, + ) + data_ids = data_id if isinstance(data_id, list) else [data_id] + outputs = [self._replay_one(str(item)) for item in data_ids] + self._mooncake_store.flush() + return outputs + + def health_check(self, timeout: float = 5.0) -> bool: + del timeout + return self._dataset is not None and self._mooncake_store is not None + + def shutdown(self) -> None: + if self._mooncake_store is not None: + self._mooncake_store.close() + self._mooncake_store = None + self._dataset = None + self._rows_by_id.clear() + logger.info("OfflineReplayEngine rank %d shutdown complete", self.rank) + + def get_status(self) -> dict: + return { + "rank": self.rank, + "initialized": self._dataset is not None, + "records": len(self._rows_by_id), + "data_path": getattr(self.args, "offline_data_path", None), + } diff --git a/torchspec/inference/factory.py b/torchspec/inference/factory.py index 1c11859c..dc2f52c9 100644 --- a/torchspec/inference/factory.py +++ b/torchspec/inference/factory.py @@ -34,13 +34,17 @@ def create_inference_engines(args, inference_pg, mooncake_config, engine_group: int = 0): """Create inference engines based on configured engine type (blocking). - Supports "hf", "sgl", and "vllm" engine types via inference_engine_type config. + Supports live inference engines and the CPU-only ``offline`` replay engine. Returns: List of head engines used for dispatching requests. Multi-node TP worker engines (if any) are kept alive internally but not returned. """ engine_type = getattr(args, "inference_engine_type", "hf") + if engine_type == "offline": + engines, init_refs = _prepare_offline_replay_engines(args, mooncake_config, engine_group) + _wait_for_init(init_refs, "OfflineReplay", timeout=300) + return engines if engine_type not in ("hf", "sgl", "vllm", "trtllm"): raise ValueError(f"Unknown inference_engine_type: {engine_type}") @@ -73,6 +77,8 @@ def prepare_inference_engines(args, inference_pg, mooncake_config, engine_group: for dispatching requests, and init_refs are ObjectRefs to wait on. """ engine_type = getattr(args, "inference_engine_type", "hf") + if engine_type == "offline": + return _prepare_offline_replay_engines(args, mooncake_config, engine_group) if engine_type not in ("hf", "sgl", "vllm", "trtllm"): raise ValueError(f"Unknown inference_engine_type: {engine_type}") @@ -95,6 +101,30 @@ def prepare_inference_engines(args, inference_pg, mooncake_config, engine_group: return engines, init_refs +def _prepare_offline_replay_engines( + args, mooncake_config, engine_group: int = 0 +) -> tuple[list, list]: + """Create CPU replay actors; they do not consume inference placement GPUs.""" + from torchspec.inference.engine.offline_replay_engine import OfflineReplayEngine + + num_engines = getattr(args, "offline_num_engines", 1) + if num_engines <= 0: + raise ValueError("inference.offline.num_engines must be positive") + ReplayRayActor = ray.remote(OfflineReplayEngine) + env_vars = get_torchspec_env_vars() + engines = [ + ReplayRayActor.options( + num_cpus=1, + num_gpus=0, + runtime_env={"env_vars": env_vars}, + ).remote(args=args, rank=rank, engine_group=engine_group) + for rank in range(num_engines) + ] + init_refs = [engine.init.remote(mooncake_config=mooncake_config) for engine in engines] + logger.info("Preparing %d CPU offline replay engine(s)", num_engines) + return engines, init_refs + + def init_engines(args, pg, engine_type: str, mooncake_config=None, engine_group: int = 0) -> list: """Initialize inference engines with Ray placement groups. diff --git a/torchspec/offline/__init__.py b/torchspec/offline/__init__.py new file mode 100644 index 00000000..8a703cb5 --- /dev/null +++ b/torchspec/offline/__init__.py @@ -0,0 +1,13 @@ +"""Offline target-data materialization and training support.""" + +from torchspec.offline.dataset import ( + OFFLINE_SCHEMA_VERSION, + OfflineDataset, + configure_offline_args, +) + +__all__ = [ + "OFFLINE_SCHEMA_VERSION", + "OfflineDataset", + "configure_offline_args", +] diff --git a/torchspec/offline/dataset.py b/torchspec/offline/dataset.py new file mode 100644 index 00000000..f025e24d --- /dev/null +++ b/torchspec/offline/dataset.py @@ -0,0 +1,170 @@ +"""Small on-disk dataset used by offline replay.""" + +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import tempfile +from pathlib import Path +from typing import Any + +import torch + +OFFLINE_SCHEMA_VERSION = 1 +_TENSOR_NAMES = ("input_ids", "hidden_states", "target", "last_hidden_states") + + +class OfflineDataset: + """Read and append replay records. + + The format intentionally has only three parts: ``dataset.json``, one + ``manifest.jsonl``, and self-describing ``samples/*.pt`` files. + """ + + def __init__( + self, + root: str | os.PathLike[str], + *, + create: bool = False, + last_hidden_states_prenorm: bool | None = None, + overwrite: bool = False, + ) -> None: + self.root = Path(root).expanduser().resolve() + if overwrite and self.root.exists(): + shutil.rmtree(self.root) + + metadata_path = self.root / "dataset.json" + if create and not metadata_path.exists(): + self.root.mkdir(parents=True, exist_ok=True) + metadata_path.write_text( + json.dumps( + { + "version": OFFLINE_SCHEMA_VERSION, + "last_hidden_states_prenorm": last_hidden_states_prenorm, + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + if not metadata_path.is_file(): + raise FileNotFoundError(f"Offline dataset not found: {metadata_path}") + + self.metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + if self.metadata.get("version") != OFFLINE_SCHEMA_VERSION: + raise ValueError( + f"Unsupported offline dataset version {self.metadata.get('version')!r}" + ) + if ( + create + and last_hidden_states_prenorm is not None + and self.metadata.get("last_hidden_states_prenorm") != last_hidden_states_prenorm + ): + raise ValueError("Offline dataset uses a different hidden-state representation") + + self._rows: dict[str, list[dict[str, str]]] = {"train": [], "eval": []} + self._by_id: dict[str, dict[str, str]] = {} + manifest = self.root / "manifest.jsonl" + if manifest.exists(): + with manifest.open(encoding="utf-8") as stream: + for lineno, line in enumerate(stream, 1): + if not line.strip(): + continue + row = json.loads(line) + split = row.get("split") + data_id = str(row.get("data_id")) + if split not in self._rows: + raise ValueError(f"Invalid split in {manifest}:{lineno}") + if data_id in self._by_id: + raise ValueError(f"Duplicate data_id {data_id!r} in {manifest}") + path = (self.root / row["file"]).resolve() + if self.root not in path.parents or not path.is_file(): + raise FileNotFoundError(f"Offline sample not found: {path}") + item = {"split": split, "data_id": data_id, "file": row["file"]} + self._rows[split].append(item) + self._by_id[data_id] = item + + def rows(self, split: str) -> list[dict[str, str]]: + return list(self._rows[split]) + + def ids(self, split: str) -> set[str]: + return {row["data_id"] for row in self._rows[split]} + + def count(self, split: str) -> int: + return len(self._rows[split]) + + def load(self, row_or_id: dict[str, str] | str) -> dict[str, Any]: + row = self._by_id[str(row_or_id)] if isinstance(row_or_id, str) else row_or_id + path = self.root / row["file"] + record = torch.load(path, map_location="cpu", weights_only=True, mmap=True) + if not isinstance(record, dict): + raise ValueError(f"Offline sample must contain a dict: {path}") + if str(record.get("data_id")) != row["data_id"]: + raise ValueError(f"Offline sample data_id does not match manifest: {path}") + if not all( + isinstance(record.get(name), torch.Tensor) for name in ("input_ids", "hidden_states") + ): + raise ValueError(f"Offline sample has missing tensors: {path}") + return record + + def append( + self, + split: str, + *, + data_id: str, + tensors: dict[str, torch.Tensor | None], + packed_loss_mask: str | None, + metadata: dict[str, Any] | None = None, + ) -> bool: + if split not in self._rows: + raise ValueError(f"Unknown offline split: {split!r}") + data_id = str(data_id) + if data_id in self._by_id: + if self._by_id[data_id]["split"] != split: + raise ValueError(f"data_id {data_id!r} already exists in another split") + return False + + saved = { + name: value.detach().cpu().contiguous() + for name, value in tensors.items() + if name in _TENSOR_NAMES and isinstance(value, torch.Tensor) + } + if not {"input_ids", "hidden_states"}.issubset(saved): + raise ValueError(f"Offline sample {data_id!r} has missing tensors") + + relative = Path("samples") / f"{hashlib.sha1(data_id.encode()).hexdigest()}.pt" + path = self.root / relative + path.parent.mkdir(parents=True, exist_ok=True) + record = { + **saved, + "data_id": data_id, + "packed_loss_mask": packed_loss_mask, + "metadata": dict(metadata or {}), + } + fd, temporary = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.") + os.close(fd) + try: + torch.save(record, temporary) + os.replace(temporary, path) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + row = {"split": split, "data_id": data_id, "file": relative.as_posix()} + with (self.root / "manifest.jsonl").open("a", encoding="utf-8") as stream: + stream.write(json.dumps(row) + "\n") + stream.flush() + os.fsync(stream.fileno()) + self._rows[split].append(row) + self._by_id[data_id] = row + return True + + +def configure_offline_args(dataset: OfflineDataset, args) -> None: + """Use the representation detail recorded with the tensors.""" + value = dataset.metadata.get("last_hidden_states_prenorm") + if value is None: + raise ValueError("Offline dataset does not declare last_hidden_states_prenorm") + args.last_hidden_states_prenorm = bool(value) diff --git a/torchspec/offline/generate.py b/torchspec/offline/generate.py new file mode 100644 index 00000000..ac9e035e --- /dev/null +++ b/torchspec/offline/generate.py @@ -0,0 +1,177 @@ +"""Materialize live target-model outputs for offline training.""" + +from __future__ import annotations + +import argparse +import os +import time +from pathlib import Path + +import ray +import torch +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + +from torchspec.config.train_config import config_to_flat_args, load_config +from torchspec.controller.inference_manager import AsyncInferenceManager +from torchspec.controller.loop import _safe_training_cleanup +from torchspec.controller.setup import build_mooncake_config +from torchspec.controller.training_controller import AsyncTrainingController +from torchspec.inference.factory import create_inference_engines +from torchspec.offline.dataset import OfflineDataset +from torchspec.offline.saving_actor import OfflineSavingActor +from torchspec.ray.placement_group import _ensure_ray_initialized, create_placement_groups +from torchspec.train_entry import ( + _get_draft_model_config, + _resolve_batch_size, + _validate_and_configure_dflash, +) +from torchspec.transfer.mooncake.utils import launch_mooncake_master +from torchspec.utils.env import get_torchspec_env_vars +from torchspec.utils.logging import logger + + +def _parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", "-c", required=True) + parser.add_argument("--output", "-o", required=True) + parser.add_argument("--overwrite", action="store_true") + cli, unknown = parser.parse_known_args() + args = config_to_flat_args(load_config(cli.config, cli_args=unknown or None)) + args.rank = 0 + args.world_size = args.training_num_nodes * args.training_num_gpus_per_node + if getattr(args, "attention_backend", None) == "usp": + raise ValueError("Offline materialization does not support USP") + _resolve_batch_size(args) + return cli, args + + +def _save_vocab_mapping(controller, output_dir: str, draft_config) -> None: + draft_vocab_size = getattr(draft_config, "draft_vocab_size", None) + if draft_vocab_size is None or draft_vocab_size == draft_config.vocab_size: + return + path = Path(output_dir) / "vocab_mapping.pt" + d2t, t2d = ray.get( + controller.compute_vocab_mapping.remote(draft_config.vocab_size, draft_vocab_size) + ) + temporary = path.with_suffix(".pt.tmp") + torch.save({"d2t": d2t.cpu(), "t2d": t2d.cpu()}, temporary) + os.replace(temporary, path) + + +def materialize(args, output_dir: str, *, overwrite: bool = False) -> dict[str, int]: + if getattr(args, "inference_engine_type", None) == "offline": + raise ValueError("Materialization requires a live inference engine") + + draft_config = _get_draft_model_config(args) + args.draft_model_config_obj = draft_config + _validate_and_configure_dflash(args, draft_config) + # A one-item dispatch lets the saving actor persist every source sample, + # including a final partial inference batch. + args.per_dp_rank_batch_size = 1 + OfflineDataset( + output_dir, + create=True, + last_hidden_states_prenorm=args.last_hidden_states_prenorm, + overwrite=overwrite, + ) + + _ensure_ray_initialized() + driver_node_id = ray.get_runtime_context().get_node_id() + controller = AsyncTrainingController.options( + runtime_env={"env_vars": get_torchspec_env_vars()}, + scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=driver_node_id, soft=False), + ).remote(args, 1) + train_size, eval_size = ray.get( + [controller.load_dataset.remote(args), controller.load_eval_dataset.remote(args)] + ) + _save_vocab_mapping(controller, output_dir, draft_config) + + inference_pg = create_placement_groups(args, roles={"inference"})["inference"] + launch_mooncake_master(args) + mooncake_config = build_mooncake_config(args) + engines = create_inference_engines(args, inference_pg, mooncake_config) + + saver = OfflineSavingActor.remote(output_dir, mooncake_config) + train_queues, eval_queues = ray.get( + [controller.get_train_queues.remote(), controller.get_eval_queues.remote()] + ) + ray.get(saver.set_queues.remote(train_queues[0], eval_queues[0])) + + if getattr(args, "max_sample_pool_size", 0) <= 0: + args.max_sample_pool_size = max( + 8, + getattr(args, "inference_batch_size", 1) + * getattr(args, "max_concurrent_batches", 1) + * 4, + ) + manager = AsyncInferenceManager.remote( + args, + controller, + inference_engines=engines, + max_concurrent_batches=getattr(args, "max_concurrent_batches", 1), + ) + if eval_size: + ray.get(controller.submit_eval_chunk.remote(0, eval_size)) + ray.get(controller.submit_training_dataset.remote()) + + manager_future = manager.run.remote() + expected = {"train": train_size, "eval": eval_size} + processed = {"train": 0, "eval": 0} + counts = None + try: + while processed != expected: + made_progress = False + for split, method in ( + ("eval", controller.try_dispatch_eval_batch), + ("train", controller.try_dispatch_batch), + ): + if processed[split] < expected[split] and ray.get(method.remote()): + ray.get(saver.save_from_queue.remote(split)) + processed[split] += 1 + made_progress = True + + if made_progress: + logger.info("Materialization progress: %s / %s", processed, expected) + continue + + manager_status, controller_status = ray.get( + [manager.get_status.remote(), controller.get_full_status.remote()] + ) + drained = ( + manager_status["prompt_buffer_size"] == 0 + and manager_status["pending_tasks"] == 0 + and controller_status["prompt_buffer_size"] == 0 + and controller_status["sample_pool_size"] == 0 + and ray.get(controller.get_eval_pool_size.remote()) == 0 + ) + if drained: + raise RuntimeError( + f"Inference drained early: processed={processed}, expected={expected}" + ) + time.sleep(0.05) + + if eval_size: + ray.get(controller.finalize_eval_dispatch.remote()) + counts = ray.get(saver.counts.remote()) + finally: + try: + ray.get(saver.close.remote()) + except Exception as exc: + logger.warning("Failed to close offline saver: %s", exc) + _safe_training_cleanup(args, manager, manager_future, engines) + try: + ray.get(controller.shutdown.remote()) + except Exception as exc: + logger.warning("Failed to stop materialization controller: %s", exc) + + logger.info("Offline dataset ready at %s: %s", output_dir, counts) + return counts + + +def main() -> None: + cli, args = _parse_args() + materialize(args, cli.output, overwrite=cli.overwrite) + + +if __name__ == "__main__": + main() diff --git a/torchspec/offline/saving_actor.py b/torchspec/offline/saving_actor.py new file mode 100644 index 00000000..910ab9e9 --- /dev/null +++ b/torchspec/offline/saving_actor.py @@ -0,0 +1,74 @@ +"""Queue consumer that saves inference outputs instead of training on them.""" + +from __future__ import annotations + +import dataclasses + +import ray +import torch + +from torchspec.offline.dataset import OfflineDataset +from torchspec.ray.ray_actor import RayActor +from torchspec.transfer.mooncake.eagle_store import EagleMooncakeStore + + +@ray.remote(num_cpus=1, num_gpus=0) +class OfflineSavingActor: + """Consume the normal training queues and persist their tensors.""" + + def __init__(self, output_dir, mooncake_config) -> None: + self.dataset = OfflineDataset(output_dir) + config = dataclasses.replace( + mooncake_config, + local_hostname=RayActor.get_node_ip(), + global_segment_size=0, + async_put_pool_size=0, + enable_gpu_direct=False, + ) + self.store = EagleMooncakeStore(config) + self.store.setup(device=torch.device("cpu")) + self.queues = {} + + def set_queues(self, train_queue, eval_queue) -> None: + self.queues = {"train": train_queue, "eval": eval_queue} + + def save_from_queue(self, split: str, count: int = 1) -> int: + """Save exactly ``count`` items and return the number newly written.""" + written = 0 + for _ in range(count): + sample = self.queues[split].get() + if not sample.data_id: + raise ValueError("Offline saving requires TrainSample.data_id") + dtypes = { + key: getattr(torch, value.replace("torch.", "")) + if isinstance(value, str) + else value + for key, value in (sample.tensor_dtypes or {}).items() + } + output = self.store.get( + key=sample.mooncake_key, + shapes=sample.tensor_shapes, + dtypes=dtypes, + device=torch.device("cpu"), + ) + try: + written += self.dataset.append( + split, + data_id=sample.data_id, + tensors=output.to_tensor_dict(), + packed_loss_mask=sample.packed_loss_mask, + metadata=sample.metadata, + ) + finally: + self.store.remove_eagle3_tensors( + sample.mooncake_key, + has_last_hidden_states="last_hidden_states" in sample.tensor_shapes, + has_target="target" in sample.tensor_shapes, + ) + return written + + def counts(self) -> dict[str, int]: + return {split: self.dataset.count(split) for split in ("train", "eval")} + + def close(self) -> None: + self.store.close() diff --git a/torchspec/ray/placement_group.py b/torchspec/ray/placement_group.py index 0bff3791..068120ca 100644 --- a/torchspec/ray/placement_group.py +++ b/torchspec/ray/placement_group.py @@ -322,11 +322,7 @@ def _ensure_ray_initialized(): def _get_expected_gpu_count(args) -> int: training_gpus = args.training_num_nodes * args.training_num_gpus_per_node inference_gpus = getattr(args, "inference_num_gpus", 0) - if ( - getattr(args, "colocate", False) - or getattr(args, "debug_train_only", False) - or getattr(args, "debug_inference_only", False) - ): + if getattr(args, "colocate", False): return max(training_gpus, inference_gpus) return training_gpus + inference_gpus @@ -354,55 +350,6 @@ def _wait_for_gpu_resources(expected_gpus: int, timeout: int = 300, poll_interva ) -def _create_custom_role_placement_group( - args, - role: str, - *, - total_gpus: int, - gpus_per_node: int, - name: str, -): - constraints = _normalize_node_constraints(args, role, required=True) - bundles, bundle_label_selectors, node_group_indices = _build_custom_bundles( - role, - constraints, - total_gpus, - gpus_per_node, - ) - logger.info( - f"Creating custom {role} placement group with {total_gpus} GPU(s) on " - f"{[constraint.selector_for_log for constraint in constraints]}" - ) - return _create_placement_group( - total_gpus, - strategy="PACK", - name=name, - bundles=bundles, - bundle_label_selector=_merge_bundle_label_selectors(bundle_label_selectors), - node_group_indices=node_group_indices, - ) - - -def _create_role_placement_group( - args, - role: str, - *, - total_gpus: int, - gpus_per_node: int, - name: str, - custom: bool, -): - if custom: - return _create_custom_role_placement_group( - args, - role, - total_gpus=total_gpus, - gpus_per_node=gpus_per_node, - name=f"custom_{name}", - ) - return _create_placement_group(total_gpus, strategy="PACK", name=name) - - def _create_custom_unified_placement_group(args, num_training_gpus: int, num_inference_gpus: int): training_constraints = _normalize_node_constraints( args, "training", required=num_training_gpus > 0 @@ -497,32 +444,6 @@ def _validate_custom_placement_constraints(args) -> None: if getattr(args, "placement_strategy", "training_first") != "custom": return - if args.debug_train_only: - num_training_gpus = args.training_num_nodes * args.training_num_gpus_per_node - training_constraints = _normalize_node_constraints( - args, "training", required=num_training_gpus > 0 - ) - _build_custom_bundles( - "training", - training_constraints, - num_training_gpus, - args.training_num_gpus_per_node, - ) - return - - if args.debug_inference_only: - num_inference_gpus = args.inference_num_gpus - inference_constraints = _normalize_node_constraints( - args, "inference", required=num_inference_gpus > 0 - ) - _build_custom_bundles( - "inference", - inference_constraints, - num_inference_gpus, - args.inference_num_gpus_per_node, - ) - return - if args.colocate: num_gpus = args.training_num_nodes * args.training_num_gpus_per_node _role, constraints = _get_custom_colocated_constraints(args) @@ -584,49 +505,49 @@ def _create_custom_colocated_placement_group(args, num_gpus: int): ) -def create_placement_groups(args): - """Initialize Ray, wait for GPU resources, and create placement groups. +def create_placement_groups(args, roles: set[str] | None = None): + """Create placement groups for the requested training/inference roles.""" + if roles is None: + roles = {"training", "inference"} + roles = frozenset(roles) + unknown = roles - {"training", "inference"} + if not roles or unknown: + raise ValueError(f"Invalid placement roles: {sorted(roles)}") - This is the single entry point for all GPU placement setup. - """ _ensure_ray_initialized() + + if len(roles) == 1: + if getattr( + args, "placement_strategy", "training_first" + ) == "custom" or _has_custom_placement_fields(args): + raise ValueError( + "Custom placement is only supported when training and inference " + "roles are created together" + ) + + role = next(iter(roles)) + if role == "training": + num_gpus = args.training_num_nodes * args.training_num_gpus_per_node + else: + num_gpus = args.inference_num_gpus + if not isinstance(num_gpus, int) or num_gpus <= 0: + raise ValueError(f"{role.capitalize()} placement requires a positive GPU count") + + _wait_for_gpu_resources(num_gpus) + logger.info("Creating %s-only placement with %d GPUs...", role, num_gpus) + pg, bundle_indices, gpu_ids = _create_placement_group( + num_gpus, strategy="PACK", name=f"{role}_pg" + ) + empty = (pg, [], []) + result = {"training": empty, "inference": empty} + result[role] = (pg, bundle_indices, gpu_ids) + return result + _validate_custom_strategy_usage(args) _validate_custom_placement_constraints(args) _wait_for_gpu_resources(_get_expected_gpu_count(args)) placement_strategy = getattr(args, "placement_strategy", "training_first") - if args.debug_train_only: - num_training_gpus = args.training_num_nodes * args.training_num_gpus_per_node - logger.info(f"Creating training placement group with {num_training_gpus} GPUs...") - training_pg, training_bundle_indices, training_gpu_ids = _create_role_placement_group( - args, - "training", - total_gpus=num_training_gpus, - gpus_per_node=args.training_num_gpus_per_node, - name="training_pg", - custom=placement_strategy == "custom", - ) - return { - "training": (training_pg, training_bundle_indices, training_gpu_ids), - "inference": (training_pg, [], []), - } - - if args.debug_inference_only: - num_inference_gpus = args.inference_num_gpus - logger.info(f"Creating inference placement group with {num_inference_gpus} GPUs...") - inference_pg, inference_bundle_indices, inference_gpu_ids = _create_role_placement_group( - args, - "inference", - total_gpus=num_inference_gpus, - gpus_per_node=args.inference_num_gpus_per_node, - name="inference_pg", - custom=placement_strategy == "custom", - ) - return { - "training": (inference_pg, [], []), - "inference": (inference_pg, inference_bundle_indices, inference_gpu_ids), - } - if args.colocate: num_gpus = args.training_num_nodes * args.training_num_gpus_per_node logger.info(f"Creating colocated placement group with {num_gpus} GPUs...") diff --git a/torchspec/train_entry.py b/torchspec/train_entry.py index 12415fd5..3455250e 100644 --- a/torchspec/train_entry.py +++ b/torchspec/train_entry.py @@ -34,6 +34,7 @@ from typing import Any, Generator import ray +import torch from omegaconf import OmegaConf from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy @@ -110,7 +111,7 @@ def parse_config(): model, dataset, training, debug, inference, logging, mooncake, decode. The config is flattened via config_to_flat_args(), with prefixed sections: - mooncake_*, sglang_*, decode_*. + mooncake_*, sglang_*, offline_*, decode_*. """ parser = argparse.ArgumentParser(description="Eagle3 speculative decoding training") @@ -137,8 +138,6 @@ def parse_config(): defaults = { "colocate": False, - "debug_train_only": False, - "debug_inference_only": False, "dp_size": None, "save_debug_train_data": None, } @@ -149,6 +148,21 @@ def parse_config(): _resolve_batch_size(flat_args) _validate_usp_args(flat_args) + if ( + getattr(flat_args, "inference_engine_type", None) == "offline" + and getattr(flat_args, "max_sample_pool_size", 0) <= 0 + ): + flat_args.max_sample_pool_size = max( + flat_args.global_batch_size, + getattr(flat_args, "inference_batch_size", 1) + * getattr(flat_args, "offline_num_engines", 1) + * 2, + ) + logger.info( + "Offline replay set max_sample_pool_size=%d for bounded Mooncake staging", + flat_args.max_sample_pool_size, + ) + return flat_args @@ -261,9 +275,10 @@ def _validate_and_configure_dflash(args, draft_model_config) -> None: algo = "DSpark" if is_dspark else "DFlash" engine_type = getattr(args, "inference_engine_type", "hf") - if engine_type not in ("vllm", "sgl", "trtllm"): + if engine_type not in ("vllm", "sgl", "trtllm", "offline"): raise NotImplementedError( - f"{algo} supports inference_engine_type in ('vllm', 'sgl', 'trtllm'), got '{engine_type}'." + f"{algo} supports inference_engine_type in " + f"('vllm', 'sgl', 'trtllm', 'offline'), got '{engine_type}'." ) if getattr(args, "defer_tokenization", False): raise NotImplementedError("DFlash does not support defer_tokenization=True.") @@ -289,7 +304,7 @@ def _validate_and_configure_dflash(args, draft_model_config) -> None: def train_async_no_generation(args): - """Entry point for Eagle3 online training. + """Entry point for Eagle3 asynchronous training. Supports prefill-only mode (default) and decode mode (train_with_decode=True) with speculative decoding. Uses distributed Ray actors with placement groups. @@ -301,6 +316,15 @@ def train_async_no_generation(args): ): raise ValueError("train_with_decode=True requires inference_engine_type=sgl") + if getattr(args, "inference_engine_type", None) == "offline": + from torchspec.offline.dataset import ( + OfflineDataset, + configure_offline_args, + ) + + offline_dataset = OfflineDataset(args.offline_data_path) + configure_offline_args(offline_dataset, args) + init_tracking(args) timer = _InitTimer() @@ -326,7 +350,12 @@ def train_async_no_generation(args): # [3] Do initialization that doesn't depend on dataset in parallel with timer.phase("Driver-side init"): - pgs = create_placement_groups(args) + roles = ( + {"training"} + if getattr(args, "inference_engine_type", None) == "offline" + else {"training", "inference"} + ) + pgs = create_placement_groups(args, roles=roles) launch_mooncake_master(args) mooncake_config = build_mooncake_config(args) @@ -346,13 +375,23 @@ def train_async_no_generation(args): vocab_size = draft_model_config.vocab_size if draft_vocab_size is not None and draft_vocab_size != vocab_size: with timer.phase("Vocab mapping"): - logger.info( - f"Computing vocab mapping on controller " - f"(target={vocab_size}, draft={draft_vocab_size})..." - ) - vocab_mapping = ray.get( - controller.compute_vocab_mapping.remote(vocab_size, draft_vocab_size) - ) + if getattr(args, "inference_engine_type", None) == "offline": + mapping_path = os.path.join(args.offline_data_path, "vocab_mapping.pt") + if not os.path.isfile(mapping_path): + raise FileNotFoundError( + "Offline replay requires the vocabulary mapping produced during " + f"materialization: {mapping_path}" + ) + saved_mapping = torch.load(mapping_path, map_location="cpu", weights_only=True) + vocab_mapping = (saved_mapping["d2t"], saved_mapping["t2d"]) + else: + logger.info( + f"Computing vocab mapping on controller " + f"(target={vocab_size}, draft={draft_vocab_size})..." + ) + vocab_mapping = ray.get( + controller.compute_vocab_mapping.remote(vocab_size, draft_vocab_size) + ) logger.info( f"Generated vocab mapping: " f"d2t={vocab_mapping[0].shape}, t2d={vocab_mapping[1].shape}" diff --git a/torchspec/training/data_fetcher.py b/torchspec/training/data_fetcher.py index 27ebf75d..fa98f5c5 100644 --- a/torchspec/training/data_fetcher.py +++ b/torchspec/training/data_fetcher.py @@ -53,6 +53,7 @@ class TrainSample: packed_loss_mask: Optional[str] = None last_turn_loss_only: Optional[bool] = None metadata: Optional[Dict[str, Any]] = None + data_id: Optional[str] = None class MooncakeDataset(IterableDataset): diff --git a/torchspec/transfer/mooncake/eagle_store.py b/torchspec/transfer/mooncake/eagle_store.py index 718b9360..42502d01 100644 --- a/torchspec/transfer/mooncake/eagle_store.py +++ b/torchspec/transfer/mooncake/eagle_store.py @@ -86,6 +86,13 @@ def _put_raw_tensors(self, keys: List[str], tensors: List[torch.Tensor]) -> None self._async_put_manager.check_last_error() self._async_put_manager.wait_for_buffer(buf.ptr) + # Offline replay tensors are already on the host, so CUDA event and + # stream synchronization is both unnecessary and unavailable. + if not any(tensor.is_cuda for tensor in tensors): + buffer_ptrs, sizes = self._stage_tensors_into_buffer(buf, tensors) + self._async_put_manager.submit(keys, buffer_ptrs, sizes, buf.ptr) + return + compute_event = torch.cuda.Event() compute_event.record() diff --git a/torchspec/transfer/mooncake/store.py b/torchspec/transfer/mooncake/store.py index 363c2d32..820f31c0 100644 --- a/torchspec/transfer/mooncake/store.py +++ b/torchspec/transfer/mooncake/store.py @@ -118,7 +118,8 @@ def setup(self, device: torch.device | int | None = None) -> None: if self.config.enable_gpu_direct and torch.cuda.is_available(): self._setup_gpu_direct(device) - if torch.cuda.is_available(): + # Can't create copy stream for "cpu" device (offline replay engine). + if torch.cuda.is_available() and (device is None or device.type == "cuda"): cuda_device = device if device is not None else torch.device("cuda") self._copy_stream = torch.cuda.Stream(device=cuda_device) logger.info("DtoH copy stream created on %s", cuda_device)