Speculators Data Generation Design Doc

[markurtz]

Summary

Speculators is a unified library for building, evaluating, and storing speculative decoding algorithms for LLM inference. A smaller, faster speculator drafts tokens, while the larger verifier accepts or rejects them in a single forward pass—delivering lower latency with lossless quality.

Today, Speculators ships configs, model definitions, and converters that standardize outputs from external training flows and in‑repo research prototypes into a format consumable by vLLM. But those flows are brittle, duplicated, and hard to use. This document kicks off the first productization step: data generation. Because it already runs separately and interfaces through datasets, it’s a natural entry point to validate correctness while adding robustness, APIs, CLIs, and features. The goal is to harden the current offline (disk‑based) pathway and expand to online (queue‑based) pipelines for consistent, performant, and easy‑to‑adapt workflows.

Goals

• Goals
  • Productize data generation into robust, reusable APIs + CLI with clear contracts.
  • Support offline (disk-based) and online (queue-based) pipelines with a shared abstraction.
  • Provide a Pydantic-based config for data-specific attributes (datasets, splits, columns, chat template, processors, verifier, states to persist).
  • Ship a Transformers inference engine now; design for vLLM (and others) later without churn.
  • Deliver a dataset reader compatible with PyTorch DataLoader for training pipelines (Eagle3, HASS today).
  • Baseline performance and correctness:
    • Performance/accuracy parity with current research scripts for offline pathways.
    • Throughput uplift for online mode vs. offline: TBD on target.
    • Operational reliability (backpressure, retries, graceful shutdown).
• Non-Goals
  • Re-implement training loops or loss functions (Eagle3/HASS remain in research repos for now).

Requirements

• Functional
  • Provide CLI + Python API to:
    • Launch data generation offline (write serialized shards to disk).
    • Launch data generation online (stream to a shared buffer: multiprocessing queue or ZeroMQ).
  • Config-driven: select dataset(s), split(s), column mapping, chat template, processor settings, verifier model, states to persist (prompt, tokens, logits, hidden states by layer).
  • Inference engine interface: load verifier, run forward pass, return standardized dict per item (or per batch with a single item emission).
  • Storage/sync interface: write/read generated examples; support disk, mp.Queue, ZeroMQ.
  • Training dataset: PyTorch-compatible dataset/iterable that reads from storage/sync.
  • Observability: structured logs, counters (produced, consumed, dropped), basic timing; deterministic run IDs and shard naming.
  • Determinism controls: seeds, reproducible sampling, idempotent writer (no dupes on retries).
• Non-Functional
  • Performance: batching where possible, minimal serialization overhead, streaming-friendly I/O (sharded files, buffered writes).
  • Reliability: backpressure handling, bounded queues, timeout + retry policy, signal-safe shutdown.
  • Compatibility: Python 3.10+, PyTorch DataLoader, Transformers engine; vLLM pluggable later.
  • Security & Safety: safe serialization (no arbitrary code exec), path whitelisting, schema validation at boundaries.
  • Cost efficiency: avoid unnecessary copies; optional compression for disk and wire.
  • Maintainability: small, composable interfaces; Pydantic models as contracts; tests at unit + integration levels.

Design

High-Level Architecture & Main Components

flowchart LR
  A[Source Datasets\nHF IDs, local paths, text files] --> B[Loader + Preprocessor\ncolumns, chat template, processors]
  B --> C[Inference Engine\nTransformers now, vLLM later]
  C --> D[Standardized Example Dict\nprompt, tokens, logits, hidden states]
  D --> E{Storage/Sync}
  E -->|Offline| F[Disk Shards\nmsgpack/zstd, parquet, or arrow]
  E -->|Online\nsingle-node| G[mp.Queue]
  E -->|Online\nmulti-proc/node| H[ZeroMQ Broker]
  F & G & H --> I[Generated Dataset\nPyTorch DataLoader]
  I --> J[Training Pipelines\nEagle3/HASS]

Loading

Key Interfaces (APIs / Contracts)

• Config (DataGenConfig) — data-only knobs:
  • datasets: list of sources (hf_id / local_path / text_file), with split + column mapping.
  • chat_template: hf_id or inline template config (default: verifier’s).
  • processor: overrides for tokenizer/processor settings (default: verifier’s).
  • verifier: model ref (hf_id, local path, pretrained).
  • states: which outputs to persist (prompt, output_tokens, logits, hidden_states={layers:[...], reduce?:mean/none}).
• Inference Pipeline (InferencePipeline)
  • from_config(config: DataGenConfig, **engine_kwargs) -> InferencePipeline
  • forward(batch: list[InputItem]) -> list[ExampleDict]
• Storage/Sync (DataSink & DataSource)
  • sink.write(example: ExampleDict) -> None (or write_batch(list[ExampleDict]))
  • source.read() -> Iterator[ExampleDict]
  • Implementations: DiskStore, MPQueueStore, ZmqStore.
• Generated Dataset (GeneratedDataset)
  • Consumes a DataSource and yields tensors/arrays in training-ready format.
  • Works under PyTorch’s multi-worker DataLoader semantics (init per worker, shard-aware).

Data Model (Example Schema)

• ExampleDict (persisted/generated):

  {
    "prompt": {"text": "string", "ids": [int], "meta": {...}},
    "verifier": {
      "model_id": "hf_or_path",
      "tokens": [int],
      "logits": [[float]],
      "hidden_states": {
        "layer_0": [[float]],
        "layer_31": [[float]]
      }
    },
    "run": {"run_id": "uuid", "engine": "transformers", "timestamp": "iso8601"}
  }
  

Storage Formats

• TBD based on timelines and discussion

Batching Behavior

• Engine may run batched inference but emits 1 example per input (optionally: store batch as N single records to the sink). — Targeted for follow up

Offline vs Online Flows

sequenceDiagram
  participant CLI
  participant Worker as Gen Worker(s)
  participant Sink as Disk/Queue/ZMQ
  participant Train as Training Proc

  CLI->>Worker: generate_data config, engine_kwargs, num_workers
  Worker->>Sink: write example / write_batch
  par Offline
    Train->>Sink: source.read from Disk shards
  and Online - mp.Queue
    Train->>Sink: source.read bounded queue with backpressure
  and Online - ZeroMQ
    Train->>Sink: connect url, source.read brokered
  end

Loading

Implementation

Features to Build

1. Pydantic Configs
   • DataGenConfig (data sources, splits, columns, chat template, processor, verifier, states).
   • EngineConfig (engine-local: batch size, dtype, device map, max seq len).
   • SyncConfig (disk/mp.Queue/ZeroMQ, common args: queue size, URLs, shard sizes, compression).
2. Inference Pipeline Abstraction
   • InferencePipeline interface + TransformersPipeline implementation (initial).
   • Hooks to add vLLMEngine later with the same contract.
3. Storage/Sync
   • DiskStore
   • MPQueueStore
   • ZmqStore
4. Dataset
   • GeneratedDataset (iterable; worker-safe init; shared args propagation; transforms to tensors).
   • Common collation utilities for training pipelines.
5. CLI + API
   • speculators generate-data --config path.yaml [--engine.* --sync.* --run.*]
   • Python: generate_data(config: DataGenConfig, **kwargs)
   • Process helpers: generate_data_process(...), generate_data_worker(...).
6. Observability & Ops
   • Structured logging, progress, counters, metrics stubs.
   • Reproducibility: seeds, run_id, manifest.json per run.
7. Testing
   • Unit tests for config validation, engine I/O, storage round-trip.
   • Integration tests: offline parity, online queue with a toy training loop.

Tasks / Milestones

M0 — Architecture & Contracts (blocking)

• Finalize interfaces: InferenceEngine, DataSink/DataSource, GeneratedDataset.
• Define Pydantic models + JSON schema snapshots.
• Skeleton packages and wiring.

M1 — Offline Path (Disk)

• Implement GeneratedDataset(Disk).
• Checkpoint to validate with save data formats into prototype implementation
• Implement DiskStore.
• Implement TransformersEngine.
• CLI + API happy path; parity tests vs research scripts.

M2 — Online Single-Process/Node (mp.Queue)

• Implement MPQueueStore (bounded, backpressure, graceful stop).
• Launch N generator workers from training proc; ensure no duplication (partitioning / index stride).
• Throughput baseline vs M1.

M3 — Online Multi-Proc/Node (ZeroMQ)

• Implement ZmqStore with a broker process and URL discovery.
• Add simple auth knob (optional), heartbeat, reconnect.
• Validate single-node multi-proc + light multi-node smoke.

M4 — Hardening & Extensibility

• Batched-engine emits single-item records.
• Optional compression modes (fp16/bf16 logits, int8 packs) and payload codecs.
• vLLM engine spike (create scaffolding + small E2E).
• Performance harness + docs.

Appendix: Concrete Interfaces (First Pass)

Packages & Layout

speculators/
  __main__.py                # adds CLI entrypoint for data generation
  train/
    datagen/
      __init__.py
      cli.py                 # argparse/typer glue for generate-data
      api.py                 # generate_data(), process orchestration
      config.py              # DataGenConfig, EngineConfig, SyncConfig
      engines/
        base.py              # InferenceEngine ABC
        transformers.py      # TransformersEngine
        vllm.py              # (scaffold)
      sync/
        base.py              # DataSink/DataSource ABCs
        disk.py              # DiskStore
        mpqueue.py           # MPQueueStore
        zmq.py               # ZmqStore (+ broker proc)
      dataset/
        generated.py         # GeneratedDataset (PyTorch compatible)
      utils/
        io.py, shard.py, ids.py, metrics.py

Pydantic Config Sketch

class DatasetSpec(BaseModel):
    name: str                                   # alias or id
    hf_id: str | None = None
    local_path: str | None = None
    text_file: str | None = None
    split: str = "train"
    columns: dict[str, str] | None = None       # {"text": "prompt", ...}

class StatesSpec(BaseModel):
    save_prompt: bool = True
    save_tokens: bool = True
    save_logits: bool = False
    hidden_state_layers: list[int] = []         # [] => none
    hidden_state_reduce: Literal["none","mean"]="none"

class VerifierSpec(BaseModel):
    id: str                                     # hf id or local path
    dtype: Literal["auto","fp16","bf16","fp32"]="auto"
    max_seq_len: int | None = None

class DataGenConfig(BaseModel):
    datasets: list[DatasetSpec]
    chat_template: str | dict | None = None     # hf id or inline config
    processor_overrides: dict | None = None
    verifier: VerifierSpec
    states: StatesSpec

CLI Examples

# Offline
speculators generate-data \
  --config configs/datagen.eagle3.yaml \
  --engine.batch_size 8 \
  --sync.kind disk --sync.output /data/speculators/runs/2025-09-18

# Online (single node, multiprocessing queue)
speculators generate-data \
  --config configs/datagen.hass.yaml \
  --engine.batch_size 8 \
  --sync.kind mpqueue --sync.maxsize 2048 \
  --run.num_workers 4

# Online (ZeroMQ broker)
speculators generate-data \
  --config configs/datagen.hass.yaml \
  --engine.batch_size 8 \
  --sync.kind zmq --sync.url tcp://0.0.0.0:5555 --sync.bind true