Update data generation to use vLLM extract_hidden_states

examples/ONLINE_TRAINING.md

@@ -0,0 +1,62 @@
+# Online training
+
+This readme walks through the process of online training an Eagle3 draft model.
+
+## Prepare data
+
+In a python environment with `speculators` installed, prepare the training dataset. Pass in the target model name/path, dataset name/path (you can pass in multiple datasets), and the output directory.
+
+```
+python scripts/prepare_data.py --model Qwen/Qwen3-8B --data sharegpt --output ./output
+```
+
+**Produces:**
+
+```
+./output/
+    data-00000-of-00002.arrow    #  ⎤
+    data-00001-of-00002.arrow    #  | Processed dataset on disk
+    dataset_info.json            #  |
+    state.json                   #  ⎦
+
+    token_freq.pt                # Token frequencies for vocab mapping
+```
+
+## Launch vLLM
+
+In a python environment with `vllm` installed, launch a vllm server configured for hidden states extraction. We provide a wrapper script (`scripts/launch_vllm.py`) to make this easier.
+
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/launch_vllm.py Qwen/Qwen3-8B -- --data-parallel-size 4 --port 8000
+```
+
+Note: anything that comes after the `--` will be passed directly to vllm. The `--data-parallel-size` and `--port` are examples of optional arguments for configuring vLLM. `--tensor-parallel-size` also works as expected.
+
+**Produces:** Model ready to serve requests on on port 8000
+
+## Run training
+
+In a python environment with `speculators` installed, launch the training process. `torchrun` (and the arguments to it) are used to launch a multi-gpu training job. These can be omitted if training on a single gpu.
+
+```
+CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --standalone --nproc_per_node 4 scripts/train.py --verifier-name-or-path Qwen/Qwen3-8B --data-path ./output --vllm-endpoint http://localhost:8000/v1 --save-path ./output/checkpoint --draft-model-size 32000
+```
+
+**Produces:** If `--draft-model-size` is set, vocab mappings will be generated and cached to the `--data-path` directory.
+
+```
+./output/
+    data-00000-of-00002.arrow    #  ⎤
+    data-00001-of-00002.arro     #  | 
+    dataset_info.json            #  | From `scripts/prepare_data.py` step
+    state.json                   #  |
+    token_freq.pt                #  ⎦
+
+    td2.npy                      #  ⎤ Vocab mappings
+    d2t.npy                      #  ⎦
+
+    checkpoints/                 # Training checkpoints (loadable by vLLM)
+        0/
+        1/
+        ...
+```

pyproject.toml

@@ -41,6 +41,7 @@ dependencies = [
     "huggingface-hub",
     "loguru>=0.7.2,<=0.7.3",
     "numpy>=2.0.0,<=2.4.2",
+    "openai>=2.0.0",
     "protobuf",
     "psutil",
     "pydantic>=2.0.0",

scripts/build_vocab_mapping.py

@@ -20,10 +20,10 @@
 
 import numpy as np
 import torch
-from transformers import AutoConfig
 
 from speculators.train.vocab_mapping import (
     build_vocab_mappings_from_distribution,
+    get_target_vocab_size,
 )
 
 logging.basicConfig(
@@ -71,29 +71,6 @@ def parse_args():
     return parser.parse_args()
 
 
-def get_target_vocab_size(args):
-    has_vocab = args.target_vocab_size is not None
-    has_model = args.target_model_path is not None
-
-    if has_vocab and has_model:
-        raise ValueError("Cannot specify both target-vocab-size and target-model-path")
-
-    if not has_vocab and not has_model:
-        raise ValueError("Must specify either target-vocab-size or target-model-path")
-
-    if has_vocab:
-        return args.target_vocab_size
-
-    logger.info(f"Loading target model config from {args.target_model_path}")
-    config = AutoConfig.from_pretrained(args.target_model_path)
-
-    # For multimodal models (Qwen3VL, etc.), extract text_config
-    if hasattr(config, "text_config"):
-        config = config.text_config
-
-    return config.vocab_size
-
-
 def main():
     args = parse_args()
 
@@ -103,7 +80,9 @@ def main():
 
     token_freq_dict = torch.load(token_freq_path, weights_only=True)
 
-    target_vocab_size = get_target_vocab_size(args)
+    target_vocab_size = get_target_vocab_size(
+        args.target_vocab_size, args.target_model_path
+    )
 
     d2t, t2d = build_vocab_mappings_from_distribution(
         token_freq_dict=token_freq_dict,

scripts/data_generation_offline.py

@@ -82,6 +82,7 @@ def parse_args():
     parser.add_argument(
         "--train-data-path",
         type=str,
+        action="append",
         required=True,
         help="Path to training data (same as used in preprocessing)",
     )
@@ -340,13 +341,12 @@ def main():
 
     dataset, _ = load_and_preprocess_dataset(
         target_model_path=args.target_model_path,
-        train_data_path=args.train_data_path,
+        train_data_paths=args.train_data_path,
         seq_length=args.seq_length,
         build_dataset_num_proc=args.num_preprocessing_workers,
         seed=args.seed,
         max_samples=args.max_samples,
         token_freq_path=args.token_freq_path,
-        cache_dir=args.hf_cache_dir,
         assistant_pattern=args.assistant_pattern,
         turn_dropout=args.turn_dropout,
     )