diff --git a/docs/cli/train.md b/docs/cli/train.md index d6bdd2bb3..57d42d245 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -71,10 +71,7 @@ torchrun --standalone --nproc_per_node=4 -m speculators.train \ - `warn`: Skip the sample with a warning, pads to fill batch. - `raise`: Raise an error -- **`--on-generate`** (choice: `cache`|`delete`, default: `"delete"`) Behavior after generating new hidden states (only applies if `--on-missing=generate`): - - - `delete`: Delete hidden states after loading (pure online training) - - `cache`: Store hidden states for reuse in future epochs (hybrid training) +- **`--on-generate`** (choice: `delete`, default: `"delete"`) Behavior after generating new hidden states (only applies if `--on-missing=generate`). Generated hidden states are always deleted after loading. Retained for backwards compatibility with existing command lines. - **`--hidden-states-path`** (str, default: `{data-path}/hidden_states`) Path where cached hidden states files are stored (or will be stored if generating). @@ -290,22 +287,6 @@ speculators train \ --lr 3e-5 ``` -### Hybrid Training (Cache on First Epoch) - -```bash -speculators train \ - --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \ - --data-path ./training_data \ - --hidden-states-path ./hidden_states \ - --vllm-endpoint http://localhost:8000/v1 \ - --on-missing generate \ - --on-generate cache \ - --save-path ./checkpoints \ - --draft-vocab-size 32000 \ - --epochs 10 \ - --lr 3e-5 -``` - ### Multi-GPU Training with WandB Logging ```bash diff --git a/docs/user_guide/algorithms/dflash.md b/docs/user_guide/algorithms/dflash.md index b8ea95760..dc77d8dbd 100644 --- a/docs/user_guide/algorithms/dflash.md +++ b/docs/user_guide/algorithms/dflash.md @@ -56,4 +56,4 @@ DFlash is based on research from Z Lab: [DFlash Project Page](https://z-lab.ai/p - [DFlash2](dflash2.md) -- Adds local convolution and candidate selection - [DSpark](dspark.md) -- Builds on DFlash with a sequential Markov head and a confidence head -- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select DFlash or DSpark, then online, offline, or hybrid) +- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select DFlash, DFlash2, or DSpark, then online or offline) diff --git a/docs/user_guide/algorithms/dspark.md b/docs/user_guide/algorithms/dspark.md index d5ad81cd5..e842ec85e 100644 --- a/docs/user_guide/algorithms/dspark.md +++ b/docs/user_guide/algorithms/dspark.md @@ -60,4 +60,4 @@ DSpark is based on research from DeepSeek: [arXiv Paper](https://arxiv.org/abs/2 ## See Also - [DFlash](dflash.md) -- The base algorithm DSpark extends -- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select DSpark, then online, offline, or hybrid) +- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select DSpark, then online or offline) diff --git a/docs/user_guide/algorithms/eagle3.md b/docs/user_guide/algorithms/eagle3.md index d6a69e5ee..f99508a7c 100644 --- a/docs/user_guide/algorithms/eagle3.md +++ b/docs/user_guide/algorithms/eagle3.md @@ -43,4 +43,4 @@ Eagle-3 is based on research from SafeAI Lab: [EAGLE Repository](https://github. ## See Also -- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select Eagle-3, then online, offline, or hybrid) +- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select Eagle-3, then online or offline) diff --git a/docs/user_guide/algorithms/mtp.md b/docs/user_guide/algorithms/mtp.md index 86927294b..6471be5e8 100644 --- a/docs/user_guide/algorithms/mtp.md +++ b/docs/user_guide/algorithms/mtp.md @@ -41,5 +41,5 @@ MTP finetuning is based on the FastMTP method from Tencent: [FastMTP Repository] ## See Also -- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select MTP, then online, offline, or hybrid) +- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select MTP, then online or offline) - [vLLM Recipes](https://recipes.vllm.ai/) -- Deployment commands for serving MTP models diff --git a/docs/user_guide/algorithms/peagle.md b/docs/user_guide/algorithms/peagle.md index 855def811..a9d98bb20 100644 --- a/docs/user_guide/algorithms/peagle.md +++ b/docs/user_guide/algorithms/peagle.md @@ -56,4 +56,4 @@ P-EAGLE is based on research from AWS AI Labs: [arXiv Paper](https://arxiv.org/a ## See Also -- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select P-EAGLE, then online, offline, or hybrid) +- [Train a Speculator](../tutorials/train.md) -- Step-by-step training guide (select P-EAGLE, then online or offline) diff --git a/docs/user_guide/getting_started.md b/docs/user_guide/getting_started.md index 32547dbfa..1b9703d1b 100644 --- a/docs/user_guide/getting_started.md +++ b/docs/user_guide/getting_started.md @@ -41,11 +41,10 @@ If a pre-trained speculator isn't available for your target model, you can train - **Online training** -- Hidden states are generated on-the-fly during training. Easier to get started, lower disk usage. - **Offline training** -- Hidden states are pre-generated and cached. -- **Hybrid training** -- Hidden states are generated on-the-fly during the first epoch, cached, and reused after. ### Tutorials -- [Train a Speculator](tutorials/train.md) -- Recommended starting point. Covers Eagle-3, P-EAGLE, DFlash, DSpark, and MTP in all three modes +- [Train a Speculator](tutorials/train.md) -- Recommended starting point. Covers Eagle-3, P-EAGLE, DFlash, DFlash2, DSpark, and MTP in online and offline modes - [Evaluating Model Performance](tutorials/evaluating_performance.md) -- Benchmark your trained speculator - [Response Regeneration](tutorials/response_regeneration.md) -- Improve training data quality diff --git a/docs/user_guide/tutorials/index.md b/docs/user_guide/tutorials/index.md index 1a0594af2..e895e7365 100644 --- a/docs/user_guide/tutorials/index.md +++ b/docs/user_guide/tutorials/index.md @@ -4,7 +4,7 @@ Step-by-step tutorials to guide you through complete workflows, from data prepar ## [Train a Speculator](train.md) -The main end-to-end walkthrough: prepare data, generate hidden states, train, and serve. Covers Eagle-3, P-EAGLE, DFlash, DSpark, and MTP, in online, offline, or hybrid mode -- pick your algorithm and mode at the top of the page. +The main end-to-end walkthrough: prepare data, generate hidden states, train, and serve. Covers Eagle-3, P-EAGLE, DFlash, DFlash2, DSpark, and MTP, in online or offline mode -- pick your algorithm and mode at the top of the page. ## [Response Regeneration](response_regeneration.md) diff --git a/docs/user_guide/tutorials/train.md b/docs/user_guide/tutorials/train.md index 968582fab..67101e625 100644 --- a/docs/user_guide/tutorials/train.md +++ b/docs/user_guide/tutorials/train.md @@ -1,6 +1,6 @@ # Train a Speculator -This tutorial walks you through training a speculator model end to end, from raw data to a checkpoint served in vLLM. It covers **Eagle-3**, **P-EAGLE**, **DFlash**, **DFlash2**, **DSpark**, and **MTP** in all three training modes. +This tutorial walks you through training a speculator model end to end, from raw data to a checkpoint served in vLLM. It covers **Eagle-3**, **P-EAGLE**, **DFlash**, **DFlash2**, **DSpark**, and **MTP**, in both online and offline training modes. Pick an algorithm and a training mode below; the rest of the walkthrough is the same for every combination. The examples use `Qwen/Qwen3-8B` as the target model -- except for MTP, which needs a verifier with native MTP layers and so uses `Qwen/Qwen3.5-9B`. The process is the same for other models. @@ -58,13 +58,7 @@ Hidden states are generated on demand from a live vLLM server, then discarded. U /// tab | Offline -Hidden states are pre-generated to disk, then read back. Use when GPU resources are limited, so you can give them all to generation and then all to training. Needs substantial disk space. - -/// - -/// tab | Hybrid - -Hidden states are generated on demand during epoch 0, cached, then reused. Use when you want to pay the generation cost once and reuse it across epochs. +Hidden states are pre-generated to disk, then read back. Use when GPU resources are limited, so you can give them all to generation and then all to training. Needs substantial disk space. The cache depends only on the verifier, so one generation pass can feed several training runs. /// @@ -99,26 +93,45 @@ Note: if you are using an experiment tracker (e.g. trackio, wandb, tensorboard, - Python 3.10+ - One or more accelerators (NVIDIA GPU, AMD GPU, or Ascend NPU) -- For offline and hybrid modes, disk space for the cached hidden states -- see [Estimating Disk Space](#estimating-disk-space-requirements) +- For offline mode, disk space for the cached hidden states -- see [Estimating Disk Space](#estimating-disk-space-requirements) - For MTP, a verifier with native MTP layers (e.g. `Qwen/Qwen3.5-9B`, `Qwen/Qwen3.5-0.8B`) ## Step 1: Prepare Your Data -Speculator training data must contain responses produced by the target model. You can create it with [Response Regeneration](response_regeneration.md) or supply on-policy data from your own generation pipeline. +Speculator training data must contain responses produced by the target model. What you run depends on the format that data is already in. + +### Data already in speculator format + +Rows carrying `input_ids` and `loss_mask` are read directly -- no tokenization and no vLLM server, so this step can run before Step 2. -Response Regeneration writes speculator-format rows containing `input_ids` and `loss_mask`, which `speculators prepare-data` can use directly: +This walkthrough uses the `tutorial_regen` split of [`inference-optimization/speculators-ci-datasets`](https://huggingface.co/datasets/inference-optimization/speculators-ci-datasets), 5K UltraChat conversations whose assistant turns were regenerated by `Qwen/Qwen3-8B` itself, so it is already on-policy for this tutorial's target model: ```bash # in speculators venv speculators prepare-data \ --model Qwen/Qwen3-8B \ + --data hf:inference-optimization/speculators-ci-datasets:tutorial_regen \ + --output ./output \ + --max-samples 5000 \ + --seq-length 8192 +``` + +For a different target model, generate your own on-policy data with [Response Regeneration](response_regeneration.md) -- it writes this format -- and pass the file instead: + +```bash +speculators prepare-data \ + --model \ --data ./target_responses.jsonl \ --output ./output \ --max-samples 5000 \ --seq-length 8192 ``` -If your generation pipeline saves natural-language conversations instead, start the target model's vLLM server as described in Step 2, then use its render endpoint to convert those responses into speculator format: +### Data in natural-language conversations + +These still need tokens and a loss mask, and both come from the target model's own vLLM server via `--render-endpoint`. That keeps a single tokenizer behind the mask, the hidden states, and serving; a local tokenizer can drift from the deployed chat template and mask the wrong spans. + +Start the server as described in Step 2 -- you need it there anyway, and any vLLM server for the same model will do. Rendering only applies the template and tokenizes, so it writes no hidden states: ```bash speculators prepare-data \ @@ -130,12 +143,17 @@ speculators prepare-data \ --seq-length 8192 ``` -The render endpoint applies the serving chat template, tokenizes each turn, and derives its loss mask. It does not generate responses or make a dataset on-policy, so the assistant responses must already come from the same target model and generation configuration used for training. +What this costs: + +- Two render calls per assistant turn -- three when the template pre-fills a reasoning scaffold -- so time scales with turns, not just samples. +- Each assistant turn becomes its own row carrying the history before it, so `--max-samples` counts rows, not conversations. + +Rendering does not generate responses or make a dataset on-policy: the assistant responses must already come from the same target model and generation configuration used for training. **Parameters explained:** - `--model` - The target model you want to accelerate -- `--data` - On-policy target-model data, either natural-language `conversations` or speculator-format `input_ids` and `loss_mask`. Can be supplied multiple times to combine datasets. +- `--data` - On-policy target-model data, either natural-language `conversations` or speculator-format `input_ids` and `loss_mask`. Accepts a local `.json`/`.jsonl` file, a directory of them, a built-in preset, or `hf:[::]` for any HuggingFace dataset. Can be supplied multiple times to combine datasets. - `--render-endpoint` - Target model's vLLM base URL; required only for natural-language conversations. - `--output` - Where to save preprocessed data - `--max-samples` - Limit samples (optional, good for testing/getting started) @@ -247,6 +265,10 @@ output/hidden_states/ └── hs_4999.safetensors # ⎦ ``` +**Time:** ~56 minutes for 5K samples on four H100s with `--data-parallel-size 4`, producing ~460 GB. + +The cache is keyed on the verifier, not on the speculator algorithm, so you can generate once and then train Eagle-3, P-EAGLE, DFlash, DFlash2, and DSpark from the same directory. + **Optimizing generation speed:** ```bash @@ -286,12 +308,6 @@ speculators generate-offline-data \ /// -/// tab | Hybrid - -Nothing to do up front. The first epoch generates hidden states from the live vLLM server and caches them to `--hidden-states-path`; subsequent epochs read from that cache. Leave vLLM running and continue to Step 4. - -/// - ## Step 4: Train //// tab | Online @@ -593,165 +609,6 @@ speculators stitch-mtp \ //// -//// tab | Hybrid - -Wait for vLLM to finish launching. In a **separate terminal** on the same node, start training. The first epoch generates hidden states from the live vLLM server and caches them; later epochs read the cache. vLLM can be stopped after the first epoch. - -The commands below assume a four-GPU node: vLLM holds GPUs 0-1 from Step 2, so training takes 2-3. Adjust `CUDA_VISIBLE_DEVICES` and `--nproc_per_node` to your machine - -/// tab | Eagle-3 - -```bash -# in speculators venv -CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 \ - -m speculators.train \ - --verifier-name-or-path Qwen/Qwen3-8B \ - --data-path ./output \ - --save-path ./output/checkpoints \ - --draft-vocab-size 32000 \ - --epochs 5 \ - --total-seq-len 8192 \ - --hidden-states-path ./output/hidden_states \ - --vllm-endpoint http://localhost:8000/v1 \ - --on-missing generate \ - --on-generate cache -``` - -/// - -/// tab | P-EAGLE - -```bash -# in speculators venv -CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 \ - -m speculators.train \ - --verifier-name-or-path Qwen/Qwen3-8B \ - --data-path ./output \ - --save-path ./output/checkpoints \ - --draft-vocab-size 32000 \ - --epochs 5 \ - --total-seq-len 8192 \ - --speculator-type peagle \ - --num-layers 4 \ - --num-depths 4 \ - --no-norm-before-residual \ - --scheduler-type cosine \ - --lr 6e-4 \ - --hidden-states-path ./output/hidden_states \ - --vllm-endpoint http://localhost:8000/v1 \ - --on-missing generate \ - --on-generate cache -``` - -/// - -/// tab | DFlash - -```bash -# in speculators venv -CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 \ - -m speculators.train \ - --verifier-name-or-path Qwen/Qwen3-8B \ - --data-path ./output \ - --save-path ./output/checkpoints \ - --draft-vocab-size 32000 \ - --epochs 5 \ - --total-seq-len 8192 \ - --speculator-type dflash \ - --num-layers 5 \ - --lr 3e-4 \ - --hidden-states-path ./output/hidden_states \ - --vllm-endpoint http://localhost:8000/v1 \ - --on-missing generate \ - --on-generate cache -``` - -/// - -/// tab | DFlash2 - -```bash -# in speculators venv -CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 \ - -m speculators.train \ - --verifier-name-or-path Qwen/Qwen3-8B \ - --data-path ./output \ - --save-path ./output/checkpoints \ - --epochs 5 \ - --total-seq-len 8192 \ - --speculator-type dflash2 \ - --num-layers 5 \ - --lr 3e-4 \ - --hidden-states-path ./output/hidden_states \ - --vllm-endpoint http://localhost:8000/v1 \ - --on-missing generate \ - --on-generate cache -``` - -/// - -/// tab | DSpark - -```bash -# in speculators venv -CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 \ - -m speculators.train \ - --verifier-name-or-path Qwen/Qwen3-8B \ - --data-path ./output \ - --save-path ./output/checkpoints \ - --draft-vocab-size 32000 \ - --epochs 5 \ - --total-seq-len 8192 \ - --speculator-type dspark \ - --num-layers 5 \ - --lr 3e-4 \ - --loss-fn '{"ce": 0.1, "tv": 0.9}' \ - --hidden-states-path ./output/hidden_states \ - --vllm-endpoint http://localhost:8000/v1 \ - --on-missing generate \ - --on-generate cache -``` - -/// - -/// tab | MTP - -```bash -# in speculators venv -CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node 2 \ - -m speculators.train \ - --verifier-name-or-path Qwen/Qwen3.5-9B \ - --data-path ./output \ - --save-path ./output/checkpoints \ - --epochs 3 \ - --total-seq-len 8192 \ - --speculator-type mtp \ - --target-layer-ids 32 \ - --hidden-states-path ./output/hidden_states \ - --vllm-endpoint http://localhost:8000/v1 \ - --on-missing generate \ - --on-generate cache -``` - -Then stitch the finetuned MTP weights back into the verifier checkpoint. This produces a self-contained checkpoint deployable on vLLM with native MTP speculative decoding: - -```bash -speculators stitch-mtp \ - ./output/checkpoints/checkpoint_best \ - Qwen/Qwen3.5-9B \ - --output-path ./output/stitched -``` - -/// - -**Flags specific to hybrid mode:** - -- `--hidden-states-path` - Where the first epoch writes its cache -- `--on-missing generate` - Generate hidden states on-the-fly when not already cached -- `--on-generate cache` - Keep generated hidden states for reuse in later epochs - -//// - **Shared parameters:** - `--draft-vocab-size 32000` - Reduced vocabulary size. MTP omits it and uses the full verifier vocabulary. @@ -794,6 +651,15 @@ vllm serve ./output/checkpoints/checkpoint_best --port 8000 That single argument is enough because the checkpoint is self-describing: vLLM reads the `speculators_config` from its `config.json`, loads the verifier named there, and enables speculative decoding. To override the defaults -- a different `num_speculative_tokens`, or pairing the speculator with a quantized verifier -- use the long form in [Serve in vLLM](serve_vllm.md). +DFlash drafts a whole 16-token block per step, which needs more scheduling headroom than the default: + +```bash +# in vllm venv +vllm serve ./output/checkpoints/checkpoint_best \ + --max-num-batched-tokens 32768 \ + --port 8000 +``` + MTP is served differently. Its weights live inside the verifier rather than in a standalone draft model, so serve the stitched checkpoint from Step 4 and enable its native MTP head: ```bash @@ -819,52 +685,38 @@ Check vLLM logs for speculative decoding metrics. ## Expected Results -These are sanity-check runs, not production numbers. With just 5K samples, model performance will be limited, so the point is to verify that the pipeline works and the model is learning. For production quality, train on significantly more data. Reference numbers are published for DFlash and P-EAGLE so far; as a rough guide, most methods reach around 40% first-token acceptance rate on similar data. +These runs exist to confirm the pipeline works and the model is learning, not to produce a good speculator. 5K samples is far too small to train a drafter you would deploy -- the models below start overfitting within a few epochs -- so treat these as a smoke test and train on substantially more data for anything real. -### DFlash +Measured on four H100s with the offline commands above -- `Qwen/Qwen3-8B`, the `tutorial_regen` split, 5 epochs, `checkpoint_best` -- and evaluated with `evaluate.py throughput` across all nine subsets of `RedHatAI/speculator_benchmarks`: -5K ShareGPT samples, 5 epochs, Qwen3-8B, measured on MT-Bench (80 prompts, 2048 max output tokens): +| Algorithm | Draft tokens per step | Acceptance length | Position-0 acceptance | Training time | +| --------- | --------------------- | ----------------- | --------------------- | ------------- | +| Eagle-3 | 3 | 1.89 | 52.4% | 15 min | +| DFlash | 15 | 1.99 | 57.9% | 19 min | +| DFlash2 | 7 | 1.90 | 52.1% | 20 min | +| DSpark | 8 | 2.02 | 54.2% | 17 min | -| Metric | Value | -| ----------------- | ------------ | -| Acceptance rate | 5.90% | -| Acceptance length | 1.47 | -| Output throughput | 129.41 tok/s | +DFlash2 differs from the other three rows in two ways, both forced rather than chosen. It requires the full verifier vocabulary, so it trains a 151936-token output layer where the others use the pruned 32000-token one -- which also means it needs a `--data-path` without the `t2d.npy` / `d2t.npy` written by Step 1, since those files set the draft vocabulary and override the CLI. And it was served on vLLM 0.28.1, the first version with DFlash2 support; the other three rows were served on 0.27.1. -### P-EAGLE +Acceptance length is the number of tokens accepted per verifier step, so it is the figure that tracks end-to-end speedup. Compare algorithms on it rather than on the raw accepted/drafted ratio, which falls as a drafter proposes more tokens per step. -5K ShareGPT samples, 5 epochs, Qwen3-8B, measured on SpecBench (80 prompts, 256 output tokens): - -| Metric | Value | -| ----------------- | ------ | -| Acceptance rate | 13.35% | -| Acceptance length | 1.53 | - -Per-position acceptance: - -| Position | Acceptance | -| -------- | ---------- | -| 0 | 40.84% | -| 1 | 10.84% | -| 2 | 1.58% | -| 3 | 0.15% | - -> **Note:** these numbers were measured with `--total-seq-len 4096`, not the 8192 used throughout this tutorial. +End to end on a four-GPU node: about 17 seconds to prepare the data, 56 minutes to generate hidden states, and 15-20 minutes to train each speculator. Because all four read the same hidden-state cache, the four checkpoints together take under two hours. ## Estimating Disk Space Requirements -Only relevant for offline and hybrid modes. +Only relevant for offline mode. Cost scales with the *actual* token count of each sample, not with `--total-seq-len`: ```python -# For Llama-3.1-8B: # avg_seq_len × num_layers × hidden_size × dtype_bytes -# 8192 × 4 × 4096 × 2 = ~268 MB per sample +# 3000 × 4 × 4096 × 2 = ~98 MB per sample (Qwen3-8B, 4 captured layers) -# For 50K samples: ~13 TB -# For 10K samples: ~2.6 TB -# For 1K samples: ~260 GB +# For 5K samples: ~460 GB +# For 10K samples: ~960 GB +# For 50K samples: ~4.8 TB ``` +The 5K figure is what this tutorial's dataset actually produced: **463 GB** for 4993 samples averaging 3025 tokens. Using `--total-seq-len` (8192) in place of the average would overestimate it by about 2.7x. + ## Common Issues & Solutions ### Issue: Out of Memory (Training) @@ -935,7 +787,7 @@ python scripts/launch_vllm.py model -- --tensor-parallel-size 2 ### Issue: Inconsistent training utilization -**Symptoms:** Training logs are bursty, GPU utilization/power draw is inconsistent for the training process. Applies to online and hybrid modes. +**Symptoms:** Training logs are bursty, GPU utilization/power draw is inconsistent for the training process. Applies to online mode. **Solutions:** diff --git a/hs_connectors/src/hs_connectors/transfer.py b/hs_connectors/src/hs_connectors/transfer.py index cf3b69aea..122d736c5 100644 --- a/hs_connectors/src/hs_connectors/transfer.py +++ b/hs_connectors/src/hs_connectors/transfer.py @@ -5,7 +5,6 @@ import dataclasses import fcntl import os -import shutil import socket import time from abc import ABC, abstractmethod @@ -57,9 +56,6 @@ def get_cached(self, file_idx: int) -> dict[str, torch.Tensor] | None: def get_generated(self, handle: str) -> dict[str, torch.Tensor] | None: """Retrieve a freshly generated sample by its vLLM-returned handle.""" - def cache(self, handle: str, file_idx: int) -> None: # noqa: B027 - """Persist a generated sample to the cache location.""" - def delete(self, handle: str) -> None: # noqa: B027 """Clean up a generated sample (e.g. delete a temp file).""" @@ -146,11 +142,6 @@ def get_cached(self, file_idx: int) -> dict[str, torch.Tensor] | None: def get_generated(self, handle: str) -> dict[str, torch.Tensor] | None: return _load_hs_file(Path(handle)) - def cache(self, handle: str, file_idx: int) -> None: - self.hidden_states_path.mkdir(parents=True, exist_ok=True) - target = self.hidden_states_path / f"hs_{file_idx}.safetensors" - shutil.move(handle, target) - def delete(self, handle: str) -> None: Path(handle).unlink() diff --git a/scripts/benchmark.py b/scripts/benchmark.py index 326d92f7c..3206ed995 100644 --- a/scripts/benchmark.py +++ b/scripts/benchmark.py @@ -310,7 +310,6 @@ def _build_train_loader( transfer=transfer, vllm_endpoint=train_args.vllm_endpoint, on_missing=train_args.on_missing, - on_generate=train_args.on_generate, verifier_name_or_path=train_args.verifier_name_or_path, request_timeout=train_args.request_timeout, max_retries=train_args.max_retries, diff --git a/src/speculators/train/cli.py b/src/speculators/train/cli.py index 2f20306e5..6cdbfe637 100644 --- a/src/speculators/train/cli.py +++ b/src/speculators/train/cli.py @@ -665,7 +665,6 @@ def main(cfg: TrainConfig): # noqa: C901 transfer=transfer, vllm_endpoint=args.vllm_endpoint, on_missing=args.on_missing, - on_generate=args.on_generate, verifier_name_or_path=args.verifier_name_or_path, request_timeout=args.request_timeout, max_retries=args.max_retries, diff --git a/src/speculators/train/config/schema.py b/src/speculators/train/config/schema.py index 3dc678c5b..de69752e7 100644 --- a/src/speculators/train/config/schema.py +++ b/src/speculators/train/config/schema.py @@ -270,11 +270,11 @@ class GenerationArgs(_Group): "Default 'generate' generates them on demand via the vLLM endpoint; the others " "skip the sample, skip with a warning, or raise.", ) - on_generate: Literal["cache", "delete"] = Field( + on_generate: Literal["delete"] = Field( default="delete", description="Behaviour after generating a hidden state (only if " - "--on-missing=generate). 'delete' discards it once loaded; 'cache' stores it " - "in the hidden states path, enabling hybrid online/offline training.", + "--on-missing=generate). Generated hidden states are always discarded once " + "loaded. Retained for backwards compatibility with existing command lines.", ) request_timeout: float = Field( default=DEFAULT_REQUEST_TIMEOUT, diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index 3fb6b4b2a..8bd3a1f64 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -159,7 +159,6 @@ def __init__( transfer: HiddenStatesTransfer | None = None, vllm_endpoint: str = "http://localhost:8000/v1", on_missing: Literal["generate", "skip", "warn", "raise"] = "generate", - on_generate: Literal["cache", "delete"] = "delete", train_ratio: float = 1.0, split: Literal["train", "val"] = "train", transform: TransformTensors | None = None, @@ -193,7 +192,6 @@ def __init__( self.transfer = transfer or FileTransfer(Path(datapath) / "hidden_states") self.vllm_endpoint = vllm_endpoint self.on_missing = on_missing - self.on_generate = on_generate self.client: openai.OpenAI | None = None self.model = model self.request_timeout = request_timeout @@ -236,7 +234,6 @@ def _compute_approx_lengths(self) -> list[int]: def _generate_hidden_states_once( self, - index: int, dataset_item: dict, client_item: ClientItem, ) -> dict[str, torch.Tensor]: @@ -260,19 +257,15 @@ def _generate_hidden_states_once( # transfer performs manifest/checksum validation first. check_hidden_states(loaded_hs, dataset_item["input_ids"].tolist()) - file_idx = self._map_to_file_idx(index) - if self.on_generate == "cache": - self.transfer.cache(handle, file_idx) - else: - try: - self.transfer.delete(handle) - except Exception as cleanup_error: # noqa: BLE001 - logger.warning( - "Loaded a valid hidden-state sample but failed to delete " - "handle %s: %s", - handle, - cleanup_error, - ) + try: + self.transfer.delete(handle) + except Exception as cleanup_error: # noqa: BLE001 + logger.warning( + "Loaded a valid hidden-state sample but failed to delete " + "handle %s: %s", + handle, + cleanup_error, + ) return loaded_hs except Exception: if handle is not None: @@ -295,7 +288,6 @@ def _get_raw_data(self, index: int) -> BatchType | SampleUnavailable: client_item = build_client_item(dataset_item) loaded_hs = self.generation_recovery.run( lambda: self._generate_hidden_states_once( - index, dataset_item, client_item, ), diff --git a/src/speculators/train/dataloader.py b/src/speculators/train/dataloader.py index b6d859412..3cae3037c 100644 --- a/src/speculators/train/dataloader.py +++ b/src/speculators/train/dataloader.py @@ -95,7 +95,6 @@ def create_train_val_loaders( transfer: HiddenStatesTransfer | None = None, vllm_endpoint: str, on_missing: Literal["generate", "skip", "warn", "raise"], - on_generate: Literal["cache", "delete"], verifier_name_or_path: str, request_timeout: float | None, max_retries: int, @@ -127,7 +126,6 @@ def create_train_val_loaders( transfer=transfer, vllm_endpoint=vllm_endpoint, on_missing=on_missing, - on_generate=on_generate, transform=noise_transform, train_ratio=train_data_ratio, split="train", @@ -144,7 +142,6 @@ def create_train_val_loaders( transfer=transfer, vllm_endpoint=vllm_endpoint, on_missing=on_missing, - on_generate=on_generate, train_ratio=train_data_ratio, split="val", model=verifier_name_or_path, diff --git a/tests/unit/train/test_data.py b/tests/unit/train/test_data.py index f0ea37d38..741a87199 100644 --- a/tests/unit/train/test_data.py +++ b/tests/unit/train/test_data.py @@ -5,7 +5,6 @@ import torch from datasets import Dataset -from safetensors.torch import save_file import speculators.train.data as data_module from speculators.models.eagle3.data import shift_batch @@ -235,43 +234,6 @@ def test_arrow_dataset_default_train_ratio_does_not_crash(tmp_path: Path): assert arrow_ds._map_to_file_idx(5) == 5 -def test_arrow_dataset_on_generate_cache_creates_hidden_states_dir(tmp_path: Path): - """on_generate="cache" must create the cache dir when cache() is called — - otherwise shutil.move into it raises FileNotFoundError, which generation recovery - downgrades to a warning, so caching silently fails for every sample.""" - ds = Dataset.from_dict( - { - "input_ids": [[1, 2, 3]], - "loss_mask": [[1, 1, 1]], - "seq_len": [3], - } - ) - ds.save_to_disk(str(tmp_path / "data")) - - arrow_ds = ArrowDataset( - max_len=128, - datapath=str(tmp_path / "data"), - on_missing="generate", - on_generate="cache", - ) - - assert hasattr(arrow_ds.transfer, "hidden_states_path") - # Directory is created lazily when cache() is called - assert not arrow_ds.transfer.hidden_states_path.exists() - - # Simulate caching a generated sample - - temp_file = tmp_path / "temp_hs.safetensors" - save_file({"hidden_states": torch.zeros(1, 1)}, temp_file) - - arrow_ds.transfer.cache(str(temp_file), file_idx=0) - - # Now the directory should exist - assert arrow_ds.transfer.hidden_states_path.is_dir() - # And the cached file should exist - assert (arrow_ds.transfer.hidden_states_path / "hs_0.safetensors").exists() - - class _SequenceTransfer: """Minimal transfer fake which returns or raises queued generated results.""" @@ -294,9 +256,6 @@ def get_generated(self, _handle): def delete(self, handle): self.deleted.append(handle) - def cache(self, _handle, _file_idx): - return None - def _make_generation_dataset( tmp_path: Path,