diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 000000000..f5b21dcb4 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,41 @@ +# See https://help.github.com/articles/about-codeowners/ +# for more info about CODEOWNERS file + +/src/speculators/ @fynnsu @shanjiaz @rahul-tuli @orestis-z +/src/speculators/convert/ @shanjiaz @fynnsu @rahul-tuli +/src/speculators/convert/mtp/ @rahul-tuli @fynnsu +/src/speculators/convert/eagle/ @shanjiaz @fynnsu +/src/speculators/data_generation/ @shanjiaz @fynnsu @rahul-tuli +/src/speculators/models/ @fynnsu @shanjiaz @rahul-tuli @orestis-z +/src/speculators/models/mtp/ @rahul-tuli @fynnsu +/src/speculators/models/peagle/ @orestis-z @shanjiaz +/src/speculators/models/eagle3/ @fynnsu @shanjiaz @orestis-z +/src/speculators/models/dflash/ @fynnsu @shanjiaz @orestis-z +/src/speculators/train/ @fynnsu @shanjiaz + +/tests/ @fynnsu @shanjiaz @rahul-tuli @orestis-z +/tests/unit/convert/ @shanjiaz @rahul-tuli @fynnsu +/tests/unit/models/ @fynnsu @rahul-tuli @shanjiaz +/tests/unit/train/ @fynnsu @shanjiaz +/tests/integration/ @fynnsu @shanjiaz @orestis-z +/tests/e2e/ @shanjiaz @fynnsu @rahul-tuli +/tests/e2e/regression/ @shanjiaz @fynnsu @rahul-tuli @dsikka +/tests/e2e/smoke/ @shanjiaz @rahul-tuli @dsikka + +/scripts/ @fynnsu @shanjiaz @rahul-tuli @orestis-z + +/hs_connectors/ @fynnsu @shanjiaz + +/.buildkite/ @dsikka @dhuangnm @deepak-kumar-neu + +/.github/ @dsikka @fynnsu @shanjiaz + +/README.md @dsikka @fynnsu @shanjiaz @rahul-tuli +/CONTRIBUTING.md @dsikka @fynnsu @shanjiaz @rahul-tuli +/CODE_OF_CONDUCT.md @dsikka @fynnsu @shanjiaz @rahul-tuli +/Makefile @dsikka @fynnsu @shanjiaz @rahul-tuli +/mkdocs.yml @dsikka @fynnsu @shanjiaz @rahul-tuli +/pyproject.toml @dsikka @fynnsu @shanjiaz @rahul-tuli +/setup.py @dsikka @fynnsu @shanjiaz @rahul-tuli +/.coderabbit.yaml @dsikka @fynnsu @shanjiaz @rahul-tuli + diff --git a/.github/mergify.yml b/.github/mergify.yml index 83f2062e4..ae92a2ae8 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -79,13 +79,28 @@ pull_request_rules: - quality-failed merge_protections: + - name: Require approval from approved reviewers list + description: >- + All pull requests must have at least one approving review from a + member of the approved reviewers list before merging. + if: [] + success_conditions: + - &approved_reviewer_check + or: + - approved-reviews-by = fynnsu + - approved-reviews-by = shanjiaz + - approved-reviews-by = dsikka + - approved-reviews-by = rahul-tuli + - approved-reviews-by = orestis-z + - name: Require two reviews description: >- - PRs labelled "two-reviews" must have at least two approving reviews - before merging. + PRs labelled "two-reviews" must have at least two approving reviews, + with at least one from the approved reviewers list, before merging. if: - label = two-reviews success_conditions: - "#approved-reviews-by >= 2" + - *approved_reviewer_check merge_protections_settings: reporting_method: check-runs diff --git a/.gitignore b/.gitignore index cb421e12c..6b972c21a 100644 --- a/.gitignore +++ b/.gitignore @@ -173,4 +173,7 @@ cython_debug/ .DS_Store # Visual Studio Code -.vscode/ \ No newline at end of file +.vscode/ + +output +output/* diff --git a/README.md b/README.md index 5bfd906df..c5d190940 100644 --- a/README.md +++ b/README.md @@ -41,11 +41,10 @@ Big updates have landed in Speculators! To get a more in-depth look, check out t Some of the exciting new features include: +- **DSpark Training Algorithm**: Added support for the DSpark training algorithm, which extends DFlash's anchored-block drafting with a Markov head that conditions each draft position on the previous token within the block, plus a confidence head that predicts per-position acceptance probability. DSpark checkpoints can warm-start from existing DFlash checkpoints. - **P-EAGLE Training Support**: Added support for the [P-EAGLE training algorithm](https://docs.vllm.ai/projects/speculators/en/latest/user_guide/algorithms/peagle), which extends EAGLE-3's architecture with parallel multi-token prediction via Conditional-On-Distribution (COD) sampling. Rather than generating draft tokens sequentially, P-EAGLE predicts multiple tokens in a single forward pass, reducing drafting latency. The Red Hat team published a [P-EAGLE speculator for Qwen3-8B](https://huggingface.co/RedHatAI/Qwen3-8B-speculator.peagle). - **MTP Finetuning Support**: Added support for finetuning the native Multi-Token Prediction (MTP) heads of models like Qwen3-Next on domain-specific data, following the [FastMTP](https://arxiv.org/abs/2509.18362) approach. Because the MTP head is small (~100M–400M params), it can be trained on pre-extracted hidden states without loading the full verifier - **Sliding Window Attention for DFlash and DSpark**: DFlash and DSpark speculators use sliding window attention on all draft layers by default. Use `--sliding-window` to set the window size and `--full-attention-indices` to opt specific layers into full attention. Sliding window attention reduces KV cache allocation for long-context sequences and can improve per-position acceptance rates compared to full attention. -- **Qwen3-8B DFlash Speculator**: The RedHat team published a [DFlash speculator for Qwen3-8B](https://huggingface.co/RedHatAI/Qwen3-8B-speculator.dflash), achieving average speculative token acceptance lengths of up to 3.74 on `math_reasoning`. -- **Gemma 4 Speculators**: The RedHat team published speculators for Gemma 4 31B-it, including both [DFlash](https://huggingface.co/RedHatAI/gemma-4-31B-it-speculator.dflash) and [EAGLE-3](https://huggingface.co/RedHatAI/gemma-4-31B-it-speculator.eagle3) checkpoints, enabling production-grade speculative decoding for Gemma 4 models. - **DFlash Training Algorithm**: Added support for the DFlash training algorithm with anchored-block drafting, using auxiliary hidden states from multiple verifier layers. Includes CLI options for block size and max anchors, plus DFlash metrics, utilities, and draft model. DFlash models trained through Speculators can now run seamlessly in vLLM as of [vLLM PR #38300](https://github.com/vllm-project/vllm/pull/38300). - **Online Training Support**: Added support for online training using the new [vLLM hidden extraction system](https://github.com/vllm-project/vllm/pull/33736), enabling real-time hidden state generation during training without requiring separate offline data generation steps. @@ -118,11 +117,16 @@ The following table summarizes the models that have been trained end-to-end by o ✅ - Qwen3 MoE + Qwen3 MoE 30B-Instruct EAGLE-3 - ✅ + ✅
DFlash ✅ + ✅ + + + 30B + DFlash ✅ ✅ @@ -149,10 +153,10 @@ The following table summarizes the models that have been trained end-to-end by o ✅ -Mistral 3 Large -675B-Instruct -EAGLE-3 ⏳ -⏳ +Mistral Small 4 +119B +DFlash ✅ +✅ Gemma 4 @@ -166,6 +170,18 @@ The following table summarizes the models that have been trained end-to-end by o EAGLE-3 ✅ ✅ + +NVIDIA Nemotron 3 Ultra +550B-A55B +DFlash ✅ +✅ + + +NVIDIA Nemotron 3 Super +120B-A12B +DFlash ✅ +✅ + diff --git a/docs/cli/launch_vllm.md b/docs/cli/launch_vllm.md index a14e907b4..ba96dfecb 100644 --- a/docs/cli/launch_vllm.md +++ b/docs/cli/launch_vllm.md @@ -20,9 +20,9 @@ python scripts/launch_vllm.py meta-llama/Llama-3.1-8B-Instruct - **`--target-layer-ids`** (int list, default: auto-select) Space-separated list of integer layer IDs from which to capture hidden states. Note: if `--include-last-layer` is enabled (default), the model's last layer will be appended to this list. Default: `[2, num_layers//2, num_layers-3]` - **Important:** If set, you must also pass the same layer ids to the training script using `--target-layer-ids`. + **Important:** If set, you must also pass the same layer ids to the training script using `--target-layer-ids`, **excluding** the final layer — training takes the auxiliary layers only. For the [full example](#full-example) below, that is `--target-layer-ids 5 20 40`. -- **`--include-last-layer` / `--no-include-last-layer`** (flag, default: `True`) For DFlash models, append the last layer (`num_hidden_layers`) to `target_layer_ids` for verifier hidden states extraction. +- **`--include-last-layer` / `--no-include-last-layer`** (flag, default: `True`) Append the last layer (`num_hidden_layers`) to `target_layer_ids` for verifier hidden states extraction. - **`--dry-run`** (flag) Print the command that would be executed without running it. diff --git a/docs/cli/response_regeneration.md b/docs/cli/response_regeneration.md index 35e1c07ad..cfbf6deb0 100644 --- a/docs/cli/response_regeneration.md +++ b/docs/cli/response_regeneration.md @@ -1,6 +1,6 @@ # response_regeneration -Regenerates assistant responses in existing datasets using a vLLM-served model. Given a dataset containing user prompts (e.g., Magpie, UltraChat), this pipeline extracts the prompts, sends them to a vLLM server, and produces a new dataset with the original prompts paired with freshly generated responses from the target model. This is useful for creating training data where you want a specific model's outputs in place of the original assistant responses. +Regenerates assistant responses in existing datasets using a vLLM-served model. Given a dataset containing conversations (e.g., Magpie, UltraChat, GSM8K), this pipeline extracts conversation turns, regenerates each assistant response turn-by-turn against the model's own prior outputs, and produces pre-tokenized training samples. For multi-turn conversations, each turn conditions on the regenerated history, producing on-policy training data. The pipeline consists of two scripts: @@ -33,6 +33,10 @@ Orchestrates the entire pipeline: starts a vLLM server (with optional data/tenso - **`--tp-size`** (int) Tensor parallel size per replica (maps to vLLM's `--tensor-parallel-size`). +- **`--max-model-len`** (int) Maximum model context length (passed to `vllm serve --max-model-len`). + +- **`--reasoning-parser`** (str) Reasoning parser for the vLLM server (passed to `vllm serve --reasoning-parser`). + - **`--keep-server`** (flag) Don't stop the vLLM server after processing completes. - **`--tool-call-parser`** (str) vLLM tool-call parser (e.g. `hermes`, `llama3_json`). Adds `--enable-auto-tool-choice --tool-call-parser` to the server; required for tool-call regeneration, otherwise tool calls arrive as raw text and are not regenerated as tools. @@ -53,13 +57,15 @@ All other arguments are passed through to `script.py`. ## script.py -Extracts user prompts from a dataset, sends them to a vLLM chat completion endpoint, and writes out new prompt-response pairs with the target model's generated responses. +Extracts conversation turns from a dataset, regenerates each assistant response turn-by-turn via a vLLM chat completion endpoint, and writes out pre-tokenized training samples with generation boundaries marked in the loss mask. ### Features +- **Multi-turn support** — detects `messages`/`conversations` fields and regenerates each assistant turn against the model's own prior responses - **Auto-detects model** from vLLM server (no need to specify `--model`) -- **Resume capability** to skip already-processed rows +- **Resume capability** to skip already-processed conversations - **Async processing** with configurable concurrency +- **Automatic retries** with exponential backoff on transient failures ### Basic Usage @@ -95,6 +101,8 @@ python scripts/response_regeneration/script.py --dataset magpie - **`--sampling-params`** (str, default: `None`) JSON object merged into each chat-completion request, e.g. `'{"temperature": 0.6, "top_p": 0.95, "seed": 0}'`. Unset keys use the server defaults. +- **`--max-retries`** (int, default: `3`) Max retry attempts per request on transient HTTP failures (408, 409, 425, 429, 5xx) with exponential backoff. Permanent errors (e.g., 400, 404) fail immediately. + #### Output Arguments - **`--outfile`** (str, default: auto-generated) Output JSONL path. If not specified, auto-generated as `{dataset}_{model}.jsonl`. diff --git a/docs/cli/train.md b/docs/cli/train.md index bb0c6f9f2..4c9b41e35 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -100,7 +100,7 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--mask-token-id`** (int, default: auto-detect) Token ID to use as mask token (for DFlash). Auto-detected if not provided. -- **`--target-layer-ids`** (int list, default: auto-select) Space-separated list of layer IDs used for hidden states. Default: `[2, num_layers//2, num_layers-3]` **Must match the values used when launching vLLM if custom layers were specified.** +- **`--target-layer-ids`** (int list, default: auto-select) Space-separated list of layer IDs for the auxiliary hidden states. Default: `[2, num_layers//2, num_layers-3]` **If custom layers were specified when launching vLLM, pass the same ids here, excluding the final layer `launch_vllm.py` appends** — that one reaches training separately as the verifier's last hidden states. ### Distributed Training Arguments diff --git a/docs/index.md b/docs/index.md index 112554304..cff0781c9 100644 --- a/docs/index.md +++ b/docs/index.md @@ -81,7 +81,7 @@ The following table summarizes the models that have been trained end-to-end by o Qwen3 8B -Eagle-3 ✅ +Eagle-3 ✅
DFlash ✅
P-EAGLE ✅ ✅ @@ -108,11 +108,16 @@ The following table summarizes the models that have been trained end-to-end by o ✅ - Qwen3 MoE + Qwen3 MoE 30B-Instruct Eagle-3 - ✅ + ✅
DFlash ✅ + ✅ + + + 30B + DFlash ✅ ✅ @@ -139,15 +144,9 @@ The following table summarizes the models that have been trained end-to-end by o ✅ -Mistral 3 Large -675B-Instruct -Eagle-3 ⏳ -⏳ - - -Qwen3.5 -9B -MTP ✅ +Mistral Small 4 +119B +DFlash ✅ ✅ @@ -162,6 +161,18 @@ The following table summarizes the models that have been trained end-to-end by o Eagle-3 ✅ ✅ + +NVIDIA Nemotron 3 Ultra +550B-A55B +DFlash ✅ +✅ + + +NVIDIA Nemotron 3 Super +120B-A12B +DFlash ✅ +✅ + diff --git a/docs/user_guide/tutorials/response_regeneration.md b/docs/user_guide/tutorials/response_regeneration.md index c12a55f78..47488f0eb 100644 --- a/docs/user_guide/tutorials/response_regeneration.md +++ b/docs/user_guide/tutorials/response_regeneration.md @@ -1,6 +1,6 @@ # Response Regeneration -This tutorial walks you through regenerating assistant responses in an existing dataset using a target model served by vLLM. The resulting dataset pairs the original user prompts with freshly generated responses (on-policy data), and is the recommended starting point for speculator training: the drafter learns to predict what the target model actually generates, not what the dataset's original authors wrote. Training directly on the dataset's original responses (off-policy) is a cheaper fallback, since it skips a full target-model pass over the data, but costs acceptance length at inference time. +This tutorial walks you through regenerating assistant responses in an existing dataset using a target model served by vLLM. The resulting dataset pairs the original user prompts with freshly generated responses (on-policy data), and is the recommended starting point for speculator training: the drafter learns to predict what the target model actually generates, not what the dataset's original authors wrote. For multi-turn conversations, each assistant turn is regenerated sequentially against the model's own prior responses, keeping the entire history on-policy. Training directly on the dataset's original responses (off-policy) is a cheaper fallback, since it skips a full target-model pass over the data, but costs acceptance length at inference time. ## Overview @@ -26,8 +26,8 @@ The simplest way to regenerate responses is using the `run_all.sh` script, which This will: 1. Start a vLLM server with the specified model -2. Extract prompts from the dataset and generate new responses -3. Save results to a JSONL file (e.g., `magpie_Llama-3.3-70B-Instruct.jsonl`) +2. Extract conversation turns from the dataset and regenerate assistant responses turn-by-turn +3. Save pre-tokenized results to a JSONL file (e.g., `magpie_Llama-3.3-70B-Instruct.jsonl`) 4. Stop the server ### Multi-GPU Configurations @@ -48,6 +48,34 @@ For larger models, use data parallelism and/or tensor parallelism: --dataset magpie ``` +### Tool-Call Regeneration + +For tool-calling datasets (e.g. `hermes-fc`), pass the model's `--tool-call-parser` (and `--reasoning-parser`, for thinking models) so the server returns structured `tool_calls` and separated reasoning instead of plain text. Both are model-specific, so look them up in the model's [vLLM recipe](https://recipes.vllm.ai/): + +```bash +# Qwen3 +./scripts/response_regeneration/run_all.sh \ + --model "Qwen/Qwen3-8B" \ + --tool-call-parser hermes --reasoning-parser qwen3 \ + --dataset hermes-fc + +# Gemma 4 +./scripts/response_regeneration/run_all.sh \ + --model "google/gemma-4-E2B-it" \ + --tool-call-parser gemma4 --reasoning-parser gemma4 \ + --dataset hermes-fc + +# gpt-oss (reasoning is parsed automatically) +./scripts/response_regeneration/run_all.sh \ + --model "openai/gpt-oss-20b" \ + --tool-call-parser openai \ + --dataset hermes-fc +``` + +This is **semi-on-policy** tool-call regeneration: the target regenerates the tool-call tokens on-policy, but tools are not executed. The *i*-th cached tool result from the source data is spliced positionally after the target's *i*-th regenerated call. + +**Limitation:** parallel tool calls are under development; the turn is currently truncated to the first call. + ## Step 2: Verify the Output The output is a JSONL file with one pre-tokenized row per target generation. `loss_mask` is `0` over the prompt the target conditioned on and `1` over the tokens it generated, so training needs no further masking: @@ -74,6 +102,30 @@ The output is a JSONL file with one pre-tokenized row per target generation. `lo Each assistant turn produces at least one row — and more when the target calls a tool, since every call is its own generation — so expect more lines than input conversations. `conversations` is a review-only twin of `input_ids`; training drops it. +For multi-turn datasets, later turns include the regenerated history as context. For example, the second turn of the same conversation would be: + +```json +{ + "id": "conv-abc_gen1", + "primary_id": "conv-abc", + "input_ids": [151644, 872, ...], + "loss_mask": [0, 0, ..., 1, 1], + "conversations": [ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "The capital of France is Paris."}, + {"role": "user", "content": "What about Germany?"}, + {"role": "assistant", "content": "The capital of Germany is Berlin."} + ], + "metadata": { + "idx": 0, + "finish_reason": "stop", + "is_tool_call": false, + "usage": {...}, + "endpoint": "http://127.0.0.1:8000/v1/chat/completions" + } +} +``` + Check that the output looks correct: ```bash diff --git a/scripts/benchmark.py b/scripts/benchmark.py new file mode 100644 index 000000000..309a24867 --- /dev/null +++ b/scripts/benchmark.py @@ -0,0 +1,661 @@ +#!/usr/bin/env python3 +"""Training benchmark harness for speculators. + +Measures training loop throughput, timing breakdown, and GPU memory usage +with full provenance tracking. Results are stored as JSON for comparison +across code changes. + +Uses the real Trainer class so measurements stay in sync with the actual +training code path. + +Subcommands: + run Run a training benchmark + compare Compare two benchmark result files + +Examples: + # Synthetic benchmark (no dataset / vLLM needed) + python scripts/benchmark.py run --synthetic \\ + -- --verifier-name-or-path Qwen/Qwen3-8B --total-seq-len 4096 + + # Real data benchmark + python scripts/benchmark.py run \\ + -- --verifier-name-or-path Qwen/Qwen3-8B --data-path ./output \\ + --on-missing skip + + # Multi-GPU + torchrun --standalone --nproc_per_node 2 scripts/benchmark.py run \\ + --synthetic -- --verifier-name-or-path Qwen/Qwen3-8B + + # Compare two runs + python scripts/benchmark.py compare baseline.json candidate.json +""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import logging +import socket +import statistics +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +import torch +import torch.distributed as dist + +# Allow importing from the scripts/ directory (peer imports like train.py). +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from train import ( + build_draft_model, + parse_vocab_mappings, + set_seed, +) + +from hs_connectors import HiddenStatesBackend +from speculators.model import SpeculatorModel +from speculators.models.eagle3.data import shift_batch +from speculators.models.mtp.data import shift_batch_mtp +from speculators.train.config import TrainConfig +from speculators.train.dataloader import create_train_val_loaders +from speculators.train.distributed import ( + get_rank, + maybe_destroy_distributed, + maybe_setup_distributed, +) +from speculators.train.logger import setup_root_logger +from speculators.train.trainer import Trainer, TrainerConfig + +BENCHMARK_VERSION = "1.0" + +TIMING_KEYS = ( + "step_ms", + "fwd_ms", + "bwd_ms", + "opt_ms", + "fetch_ms", + "tokens_per_s", +) + + +# --------------------------------------------------------------------------- +# Metric capture +# --------------------------------------------------------------------------- + + +class _MetricCapture(logging.Handler): + """Captures profile dicts emitted by the Trainer via metric_logger.""" + + def __init__(self): + super().__init__() + self.profiles: list[dict] = [] + + def emit(self, record): + msg = record.msg + if isinstance(msg, dict) and "profile" in msg and msg["profile"] is not None: + self.profiles.append(msg["profile"]) + + +# --------------------------------------------------------------------------- +# Synthetic data loader +# --------------------------------------------------------------------------- + + +class _SyntheticLoader: + """DataLoader-like wrapper that yields the same batch repeatedly. + + Satisfies the Trainer's expectations: ``__len__``, ``__iter__``, and a + ``batch_sampler`` with ``set_epoch``. + """ + + class _BatchSampler: + def set_epoch(self, _epoch): + pass + + def __init__(self, batch: dict[str, torch.Tensor], num_steps: int): + self._batch = batch + self._num_steps = num_steps + self.batch_sampler = self._BatchSampler() + + def __len__(self): + return self._num_steps + + def __iter__(self): + for _ in range(self._num_steps): + yield self._batch + + +# --------------------------------------------------------------------------- +# Provenance +# --------------------------------------------------------------------------- + + +def collect_provenance() -> dict: + """Gather system and version metadata for reproducibility.""" + try: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], # noqa: S607 + capture_output=True, + text=True, + timeout=5, + check=False, + ) + git_sha = result.stdout.strip() if result.returncode == 0 else "unknown" + except OSError: + git_sha = "unknown" + + gpu_info = [] + if torch.cuda.is_available(): + for i in range(torch.cuda.device_count()): + props = torch.cuda.get_device_properties(i) + gpu_info.append( + { + "name": props.name, + "total_memory_gb": round(props.total_memory / 2**30, 2), + "compute_capability": [props.major, props.minor], + } + ) + + def _version(pkg: str) -> str: + try: + return importlib.metadata.version(pkg) + except importlib.metadata.PackageNotFoundError: + return "unknown" + + return { + "git_sha": git_sha, + "timestamp": datetime.now(timezone.utc).isoformat(), + "hostname": socket.gethostname(), + "python_version": sys.version.split()[0], + "pytorch_version": torch.__version__, + "cuda_version": torch.version.cuda or "none", + "speculators_version": _version("speculators"), + "transformers_version": _version("transformers"), + "gpu_info": gpu_info, + "num_gpus": (torch.cuda.device_count() if torch.cuda.is_available() else 0), + } + + +# --------------------------------------------------------------------------- +# Synthetic data +# --------------------------------------------------------------------------- + + +def create_synthetic_batch( + total_seq_len: int, + hidden_size: int, + num_target_layers: int, + vocab_size: int = 32000, + dtype: torch.dtype = torch.bfloat16, + device: torch.device | int = 0, +) -> dict[str, torch.Tensor]: + """Create a random batch matching the post-collation training shape.""" + hs_dim = num_target_layers * hidden_size + return { + "hidden_states": torch.randn( + 1, total_seq_len, hs_dim, dtype=dtype, device=device + ), + "input_ids": torch.randint(0, vocab_size, (1, total_seq_len), device=device), + "verifier_last_hidden_states": torch.randn( + 1, total_seq_len, hidden_size, dtype=dtype, device=device + ), + "loss_mask": torch.ones(1, total_seq_len, dtype=torch.bool, device=device), + "position_ids": torch.arange( + 1, total_seq_len + 1, device=device, dtype=torch.long + ).unsqueeze(0), + "document_ids": torch.zeros(1, total_seq_len, dtype=torch.long, device=device), + } + + +# --------------------------------------------------------------------------- +# Statistics +# --------------------------------------------------------------------------- + + +def compute_statistics(values: list[float]) -> dict[str, float]: + """Compute summary statistics for a list of measurements.""" + return { + "mean": statistics.mean(values), + "std": statistics.stdev(values) if len(values) > 1 else 0.0, + "min": min(values), + "max": max(values), + "median": statistics.median(values), + "count": len(values), + } + + +# --------------------------------------------------------------------------- +# Benchmark runner +# --------------------------------------------------------------------------- + + +def _build_train_loader( + bench_args, + train_args, + hidden_size, + num_target_layers, + vocab_size, + hidden_states_dtype, + total_steps, +): + """Build either a synthetic or real data loader for benchmarking.""" + if bench_args.synthetic: + synth_batch = create_synthetic_batch( + total_seq_len=train_args.total_seq_len, + hidden_size=hidden_size, + num_target_layers=num_target_layers, + vocab_size=vocab_size, + dtype=hidden_states_dtype, + device="cpu", + ) + return _SyntheticLoader(synth_batch, total_steps), True + + preprocess_fns = { + "eagle3": shift_batch, + "peagle": shift_batch, + "mtp": shift_batch_mtp, + } + preprocess = preprocess_fns.get(train_args.speculator_type) + + backend_registry = HiddenStatesBackend.registry + backend_cls = backend_registry[train_args.hidden_states_backend] + transfer = backend_cls.from_train_args(train_args, train_args.data_path) + + train_loader, _ = create_train_val_loaders( + data_path=train_args.data_path, + total_seq_len=train_args.total_seq_len, + hidden_states_dtype=hidden_states_dtype, + noise_std=train_args.noise_std, + legacy_data=train_args.legacy_data, + 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, + hidden_size=hidden_size, + num_target_layers=num_target_layers, + num_workers=train_args.num_workers, + prefetch_factor=train_args.prefetch_factor, + preprocess=preprocess, + train_data_ratio=train_args.train_data_ratio, + ) + return train_loader, False + + +def run_benchmark(bench_args, train_args) -> dict: + """Execute the benchmark using the real Trainer and return results.""" + set_seed( + train_args.seed, + getattr(train_args, "deterministic_cuda", False), + ) + setup_root_logger() + maybe_setup_distributed() + + rank = get_rank() + total_steps = bench_args.warmup_steps + bench_args.measured_steps + hidden_states_dtype = getattr(torch, train_args.hidden_states_dtype) + + # --- Build model --- + if train_args.speculator_type == "mtp": + d2t, t2d, draft_vocab_size = None, None, None + train_args.mask_token_id = None + else: + d2t, t2d, draft_vocab_size = parse_vocab_mappings(train_args) + + model_class = SpeculatorModel.registry[train_args.speculator_type] + draft_model = build_draft_model(train_args, model_class, t2d, d2t, draft_vocab_size) + + num_target_layers = len(draft_model.target_layer_ids) + hidden_size = draft_model.config.transformer_layer_config.hidden_size + vocab_size = draft_model.config.transformer_layer_config.vocab_size + + # --- Build data loader --- + train_loader, is_synthetic = _build_train_loader( + bench_args, + train_args, + hidden_size, + num_target_layers, + vocab_size, + hidden_states_dtype, + total_steps, + ) + + # --- Get forward kwargs --- + train_call_kwargs, _ = model_class.get_trainer_kwargs(**vars(train_args)) + + # --- Build TrainerConfig --- + trainer_config = TrainerConfig( + lr=train_args.lr, + num_epochs=1, + save_path="benchmark_unused", + resume_from_checkpoint=False, + train_call_kwargs=train_call_kwargs, + optimizer=train_args.optimizer, + weight_decay=train_args.weight_decay, + muon_lr=train_args.muon_lr, + muon_momentum=train_args.muon_momentum, + muon_weight_decay=train_args.muon_weight_decay, + muon_ns_steps=train_args.muon_ns_steps, + muon_adjust_lr_fn=train_args.muon_adjust_lr_fn, + scheduler_type="none", + hidden_states_dtype=hidden_states_dtype, + log_freq=1, + fsdp_shard=train_args.fsdp_shard, + max_steps=total_steps, + ) + + # --- Construct Trainer (handles GPU placement, DDP, optimizer) --- + trainer = Trainer(draft_model, trainer_config, train_loader) + + if rank == 0: + print( + f"Benchmarking: {bench_args.warmup_steps} warmup + " + f"{bench_args.measured_steps} measured steps" + ) + + # --- Attach metric capture --- + metric_logger = logging.getLogger("speculators.metrics") + capture = _MetricCapture() + metric_logger.addHandler(capture) + + # --- Reset memory tracking before the run --- + local_rank = trainer.local_rank + torch.cuda.reset_peak_memory_stats(local_rank) + + # --- Run the real training loop --- + trainer.train_epoch(0) + + # --- Remove capture handler --- + metric_logger.removeHandler(capture) + + # --- Collect memory --- + peak_allocated_mb = torch.cuda.max_memory_allocated(local_rank) / (1024**2) + peak_reserved_mb = torch.cuda.max_memory_reserved(local_rank) / (1024**2) + + # --- Split warmup / measured profiles --- + all_profiles = capture.profiles + measured_profiles = all_profiles[bench_args.warmup_steps :] + + if not measured_profiles: + if rank == 0: + print( + f"WARNING: Got {len(all_profiles)} profiles but expected " + f"{total_steps}. Check warmup_steps={bench_args.warmup_steps}." + ) + measured_profiles = all_profiles + + # --- Aggregate --- + timing_agg = {} + for key in TIMING_KEYS: + values = [s[key] for s in measured_profiles] + timing_agg[key] = compute_statistics(values) + + num_gpus_used = dist.get_world_size() if dist.is_initialized() else 1 + + results = { + "benchmark_version": BENCHMARK_VERSION, + "provenance": collect_provenance(), + "config": { + "speculator_type": train_args.speculator_type, + "verifier_name_or_path": train_args.verifier_name_or_path, + "total_seq_len": train_args.total_seq_len, + "hidden_size": hidden_size, + "num_target_layers": num_target_layers, + "optimizer": train_args.optimizer, + "lr": train_args.lr, + "hidden_states_dtype": train_args.hidden_states_dtype, + "synthetic_data": is_synthetic, + "fsdp_shard": train_args.fsdp_shard, + "num_gpus_used": num_gpus_used, + "warmup_steps": bench_args.warmup_steps, + "measured_steps": bench_args.measured_steps, + "seed": train_args.seed, + }, + "memory": { + "peak_allocated_mb": round(peak_allocated_mb, 2), + "peak_reserved_mb": round(peak_reserved_mb, 2), + }, + "timing": timing_agg, + } + + if not bench_args.no_per_step: + results["per_step"] = [ + {"step": i, **p} for i, p in enumerate(measured_profiles) + ] + + # --- Write results (rank 0 only) --- + if rank == 0: + output_path = Path(bench_args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump(results, f, indent=2) + print(f"\nResults written to {output_path}") + _print_summary(results) + + # --- Cleanup --- + del trainer, draft_model + if torch.cuda.is_available(): + torch.cuda.empty_cache() + maybe_destroy_distributed() + + return results + + +def _print_summary(results: dict) -> None: + """Print a compact summary of benchmark results to stdout.""" + timing = results["timing"] + memory = results["memory"] + + print(f"\n{'Metric':<16} {'Mean':>10} {'Std':>10} {'Min':>10} {'Max':>10}") + print("-" * 58) + for key in TIMING_KEYS: + stats = timing[key] + print( + f"{key:<16} {stats['mean']:>10.2f} " + f"{stats['std']:>10.2f} " + f"{stats['min']:>10.2f} {stats['max']:>10.2f}" + ) + print( + f"\nPeak memory: {memory['peak_allocated_mb']:.1f} MB " + f"allocated, {memory['peak_reserved_mb']:.1f} MB reserved" + ) + + +# --------------------------------------------------------------------------- +# Compare +# --------------------------------------------------------------------------- + + +def _nested_get(d: dict, key_path: str): + """Traverse a dotted key path into a nested dict.""" + for part in key_path.split("."): + if not isinstance(d, dict): + return None + d = d.get(part) # type: ignore[assignment] + return d + + +def _get_gpu_name(result: dict) -> str: + """Extract the first GPU name from a result dict.""" + info = result.get("provenance", {}).get("gpu_info") + if info: + return info[0].get("name", "unknown") + return "unknown" + + +def compare_benchmarks(baseline_path: str, candidate_path: str) -> None: + """Load two result files and print a comparison table.""" + with open(baseline_path) as f: + baseline = json.load(f) + with open(candidate_path) as f: + candidate = json.load(f) + + # --- Comparability warnings --- + comparability_checks = [ + ("config.speculator_type", "Speculator type"), + ("config.hidden_size", "Hidden size"), + ("config.total_seq_len", "Sequence length"), + ("config.num_gpus_used", "GPU count"), + ("config.fsdp_shard", "FSDP shard"), + ("config.optimizer", "Optimizer"), + ("config.hidden_states_dtype", "Dtype"), + ] + warnings_found = False + for key_path, label in comparability_checks: + val_a = _nested_get(baseline, key_path) + val_b = _nested_get(candidate, key_path) + if val_a != val_b: + if not warnings_found: + print("COMPARABILITY WARNINGS:") + warnings_found = True + print(f" {label}: {val_a} vs {val_b}") + + gpu_a = _get_gpu_name(baseline) + gpu_b = _get_gpu_name(candidate) + if gpu_a != gpu_b: + if not warnings_found: + print("COMPARABILITY WARNINGS:") + print(f" GPU: {gpu_a} vs {gpu_b}") + + # --- Header --- + sha_a = baseline.get("provenance", {}).get("git_sha", "unknown")[:12] + sha_b = candidate.get("provenance", {}).get("git_sha", "unknown")[:12] + print(f"\nBaseline: {baseline_path}") + print(f" Git SHA: {sha_a}") + print(f"Candidate: {candidate_path}") + print(f" Git SHA: {sha_b}") + + # --- Timing comparison --- + col_w = 26 + print( + f"\n{'Metric':<16} " + f"{'Baseline (mean +/- std)':<{col_w}} " + f"{'Candidate (mean +/- std)':<{col_w}} " + f"{'Delta':>10} {'Delta %':>10}" + ) + print("-" * (16 + col_w * 2 + 22)) + + for key in TIMING_KEYS: + ba = baseline.get("timing", {}).get(key, {}) + ca = candidate.get("timing", {}).get(key, {}) + ba_mean = ba.get("mean", 0) + ba_std = ba.get("std", 0) + ca_mean = ca.get("mean", 0) + ca_std = ca.get("std", 0) + delta = ca_mean - ba_mean + pct = (delta / ba_mean * 100) if ba_mean != 0 else 0 + + ba_str = f"{ba_mean:>8.2f} +/- {ba_std:<6.2f}" + ca_str = f"{ca_mean:>8.2f} +/- {ca_std:<6.2f}" + print( + f"{key:<16} {ba_str:<{col_w}} {ca_str:<{col_w}} " + f"{delta:>+10.2f} {pct:>+9.1f}%" + ) + + # --- Memory comparison --- + print(f"\n{'Memory':<24} {'Baseline':>12} {'Candidate':>12} {'Delta':>12}") + print("-" * 62) + for key in ("peak_allocated_mb", "peak_reserved_mb"): + ba_val = baseline.get("memory", {}).get(key, 0) + ca_val = candidate.get("memory", {}).get(key, 0) + delta = ca_val - ba_val + print(f"{key:<24} {ba_val:>9.1f} MB {ca_val:>9.1f} MB {delta:>+9.1f} MB") + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def build_parser(): + """Build the top-level argument parser.""" + parser = argparse.ArgumentParser( + description="Training benchmark harness for speculators.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "Pass train.py flags after '--'. Example:\n" + " python scripts/benchmark.py run --synthetic " + "-- --verifier-name-or-path Qwen/Qwen3-8B" + ), + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + # --- run --- + run_parser = subparsers.add_parser("run", help="Run a training benchmark") + run_parser.add_argument( + "--synthetic", + action="store_true", + help="Use synthetic random data instead of a real dataset.", + ) + run_parser.add_argument( + "--warmup-steps", + type=int, + default=10, + help="Warmup steps (not measured). Default: 10.", + ) + run_parser.add_argument( + "--measured-steps", + type=int, + default=50, + help="Number of measured steps. Default: 50.", + ) + run_parser.add_argument( + "--output", + type=str, + default=None, + help="Output JSON path. Default: benchmark_.json.", + ) + run_parser.add_argument( + "--no-per-step", + action="store_true", + help="Omit per-step timing data from the output JSON.", + ) + + # --- compare --- + cmp_parser = subparsers.add_parser( + "compare", help="Compare two benchmark result files" + ) + cmp_parser.add_argument("baseline", help="Path to baseline result JSON.") + cmp_parser.add_argument("candidate", help="Path to candidate result JSON.") + + return parser + + +def main(): + parser = build_parser() + + # Split argv on '--' to separate benchmark / train.py args. + argv = sys.argv[1:] + if "--" in argv: + sep_idx = argv.index("--") + bench_argv = argv[:sep_idx] + train_argv = argv[sep_idx + 1 :] + else: + bench_argv = argv + train_argv = [] + + bench_args = parser.parse_args(bench_argv) + + if bench_args.command == "compare": + compare_benchmarks(bench_args.baseline, bench_args.candidate) + return + + # --- run command --- + if bench_args.output is None: + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + bench_args.output = f"benchmark_{ts}.json" + + # Resolve train.py args through the shared TrainConfig, flattened to the + # same argparse.Namespace the model layer consumes in train.main(). + train_args = argparse.Namespace(**TrainConfig.resolve(train_argv).flatten()) + + run_benchmark(bench_args, train_args) + + +if __name__ == "__main__": + main() diff --git a/scripts/train.py b/scripts/train.py index db0462e7c..515c2328e 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -15,19 +15,15 @@ from transformers.models.qwen3.configuration_qwen3 import Qwen3Config from hs_connectors import HiddenStatesBackend -from speculators.data_generation.vllm_client import ( - DEFAULT_MAX_RETRIES, - DEFAULT_REQUEST_TIMEOUT, -) from speculators.model import SpeculatorModel from speculators.models.eagle3.data import shift_batch from speculators.models.eagle3.rotary_partial import install_partial_neox_rotary -from speculators.models.metrics import resolve_loss_config from speculators.models.mtp.data import shift_batch_mtp from speculators.models.utils import ( get_verifier_config, resolve_draft_intermediate_size, ) +from speculators.train.config import TrainConfig from speculators.train.dataloader import create_train_val_loaders from speculators.train.distributed import ( get_rank, @@ -37,12 +33,11 @@ ) from speculators.train.logger import setup_metric_logger, setup_root_logger from speculators.train.trainer import Trainer, TrainerConfig -from speculators.train.utils import resolve_mask_token_id, save_train_command +from speculators.train.utils import resolve_mask_token_id from speculators.train.vocab_mapping import ( build_vocab_mappings_from_distribution, get_target_vocab_size, ) -from speculators.utils.argparse_utils import explicitly_provided_dests from speculators.utils.loading import is_config_only_dir logger = logging.getLogger(__name__) @@ -512,7 +507,13 @@ def build_draft_model( ) -def main(args: argparse.Namespace): # noqa: C901 +def main(cfg: TrainConfig): # noqa: C901 + # Phase-1 adapter: the model layer still consumes a flat vars(args)-shaped + # dict via **kwargs, so flatten the typed config back into a namespace here. + # New code should read cfg.. directly and must NOT add new + # args.* accesses below this line. + args = argparse.Namespace(**cfg.flatten()) + # Set random seed for reproducibility set_seed(args.seed, args.deterministic_cuda) @@ -523,7 +524,36 @@ def main(args: argparse.Namespace): # noqa: C901 ) # Setup distributed training - maybe_setup_distributed() + maybe_setup_distributed(sp_size=args.sp_size) + + if args.sp_size > 1: + if not is_distributed(): + raise ValueError( + "--sp-size > 1 requires launching with torchrun/distributed " + "training; otherwise sequence parallelism has no effect." + ) + if args.total_seq_len % args.sp_size != 0: + raise ValueError( + f"--total-seq-len ({args.total_seq_len}) must be divisible " + f"by --sp-size ({args.sp_size})" + ) + if args.speculator_type not in ("eagle3", "dflash", "dspark", "peagle"): + raise ValueError( + f"Sequence parallelism (--sp-size > 1) is currently only " + f"supported for eagle3, dflash, dspark, and peagle, " + f"got --speculator-type={args.speculator_type}" + ) + if not args.fsdp_shard: + logger.warning( + "--sp-size > 1 without --fsdp-shard uses DDP, which is " + "supported but less memory-efficient for long-context SP " + "training. Consider adding --fsdp-shard." + ) + if args.speculator_type == "dflash" and args.max_anchors % args.sp_size != 0: + raise ValueError( + f"--max-anchors ({args.max_anchors}) must be divisible " + f"by --sp-size ({args.sp_size}) for DFlash SP" + ) if args.fsdp_shard and not is_distributed(): raise ValueError( @@ -538,24 +568,14 @@ def main(args: argparse.Namespace): # noqa: C901 "Installed partial-neox rotary patch for HF/vLLM RoPE alignment " "(draft_mrope_full_head_hack=False)" ) + # Write the reproducibility artifacts (run.yaml + train_command.txt) next to + # the checkpoints at rank 0 only, so every checkpoint carries the resolved + # config that produced it. if get_rank() == 0: - save_train_command(args.save_path) + cfg.save(args.save_path) - if not hasattr(torch, args.hidden_states_dtype): - raise ValueError( - "--hidden-states-dtype must be a dtype attribute of torch. e.g. `bfloat16`" - ) hidden_states_dtype = getattr(torch, args.hidden_states_dtype) - if hidden_states_dtype == torch.float16: - raise NotImplementedError( - "--hidden-states-dtype=float16 is not supported. " - "float16 with torch.autocast requires gradient scaling (GradScaler) to " - "prevent gradient underflow, which is not implemented. " - "Use bfloat16 instead, which provides the same memory savings with " - "better numerical stability and no gradient scaling required." - ) - if args.speculator_type == "mtp": if args.draft_attn_impl != "simple_flex_attention": raise ValueError( @@ -628,6 +648,12 @@ def main(args: argparse.Namespace): # noqa: C901 backend_registry = HiddenStatesBackend.registry backend_cls = backend_registry[args.hidden_states_backend] + # from_train_args is the live runtime consumer of the backend's (mirrored) + # train-args, read off the flattened namespace. hs_connectors stays + # argparse-based and standalone so vLLM can use it without speculators; that + # is why the backend's train-args are mirrored into the pydantic schema rather + # than the plugin depending on pydantic. test_backend_reconciliation.py keeps + # the mirror complete so nothing read here was dropped during resolution. transfer = backend_cls.from_train_args(args, args.data_path) train_loader, val_loader = create_train_val_loaders( @@ -678,6 +704,7 @@ def main(args: argparse.Namespace): # noqa: C901 hidden_states_dtype=hidden_states_dtype, log_freq=args.log_freq, fsdp_shard=args.fsdp_shard, + max_steps=args.max_steps, ) trainer = Trainer(draft_model, trainer_config, train_loader, val_loader) @@ -692,638 +719,8 @@ def main(args: argparse.Namespace): # noqa: C901 maybe_destroy_distributed() -def _checkpoint_freq(value: str) -> float: - fvalue = float(value) - if fvalue <= 0: - raise argparse.ArgumentTypeError("--checkpoint-freq must be > 0") - if fvalue > 1 and not fvalue.is_integer(): - raise argparse.ArgumentTypeError( - f"--checkpoint-freq={fvalue} is not an integer. Values > 1 are treated " - "as epoch counts and must be whole numbers." - ) - return fvalue - - -# CLI flags that synthesize the draft decoder shape. They conflict with both -# --from-pretrained and --draft-config, each of which fully defines the draft. -DECODER_SHAPING_FLAGS: dict[str, str] = { - "num_layers": "--num-layers", - "draft_arch": "--draft-arch", - "draft_hidden_act": "--draft-hidden-act", - "sliding_window": "--sliding-window", - "full_attention_indices": "--full-attention-indices", -} - - -def validate_draft_init_args( - parser: argparse.ArgumentParser, - args: argparse.Namespace, - provided: set[str], -) -> None: - """Enforce the draft-init contract. - - The draft model may be defined in exactly one way: - - * ``--from-pretrained`` -- load a complete speculator checkpoint (or a - config-only directory); or - * ``--draft-config`` -- load just the decoder config and build the rest of - the speculator from the other CLI args; or - * the decoder-shaping flags (``--num-layers`` etc.) -- synthesize everything. - - ``--from-pretrained`` takes precedence over all other model-definition - options: it is mutually exclusive with ``--draft-config`` and with the - decoder-shaping flags, since those values come from the checkpoint. - ``--draft-config`` is likewise incompatible with the decoder-shaping flags. - MTP from scratch (``--speculator-type mtp`` without ``--from-pretrained``) - reuses the verifier's own decoder config, so ``--draft-config`` and the - decoder-shaping flags do not apply and are rejected. - - ``provided`` is the set of decoder-shaping dests the user explicitly passed - (see :func:`speculators.utils.argparse_utils.explicitly_provided_dests`); a flag - passed at its default value still counts as a conflict. - """ - shaping = [flag for dest, flag in DECODER_SHAPING_FLAGS.items() if dest in provided] - if args.from_pretrained: - conflicting = shaping + (["--draft-config"] if args.draft_config else []) - if conflicting: - parser.error( - "--from-pretrained loads a complete draft model and takes precedence " - "over all other model-definition options, so these conflict with it " - f"(remove them): {', '.join(conflicting)}" - ) - return - if args.speculator_type == "mtp": - # MTP-from-scratch reuses the verifier's own decoder config and extracts the - # native MTP head weights; --draft-config and the decoder-shaping flags do not - # apply, so reject them rather than silently ignoring them. - conflicting = shaping + (["--draft-config"] if args.draft_config else []) - if conflicting: - parser.error( - "--speculator-type mtp reuses the verifier's decoder config, so these " - f"options do not apply (remove them): {', '.join(conflicting)}" - ) - return - if args.draft_config and shaping: - parser.error( - "--draft-config defines the draft decoder, so these flags conflict with " - f"it (remove them): {', '.join(shaping)}" - ) - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--verifier-name-or-path", type=str, required=True) - parser.add_argument( - "--trust-remote-code", - action="store_true", - help="Allow executing code from HF Hub when loading the verifier's tokenizer.", - ) - parser.add_argument( - "--speculator-type", - type=str, - default="eagle3", - help="Type of speculator model to train (eagle3, dflash, dspark, peagle, mtp)", - ) - parser.add_argument( - "--from-pretrained", - type=str, - default="", - help="Path or HF id of a pretrained draft. May also point to a " - "local directory containing only a config.json, in which case a " - "fresh draft is initialized from that full speculator config. Takes precedence " - "over and is mutually exclusive with --draft-config and the decoder-shaping " - "flags (--num-layers, --draft-arch, --draft-hidden-act, --sliding-window, " - "--full-attention-indices).", - ) - parser.add_argument( - "--draft-config", - type=str, - default="", - help="HF id, directory, or JSON path of a decoder config (LlamaConfig for " - "eagle3/peagle, Qwen3Config for dflash) to use as the draft " - "transformer_layer_config; the rest of the speculator is built from the other " - "CLI args. Mutually exclusive with --from-pretrained and with the " - "decoder-shaping flags (--num-layers, --draft-arch, --draft-hidden-act, " - "--sliding-window, --full-attention-indices).", - ) - parser.add_argument( - "--dry-run", - action="store_true", - default=False, - help="Build the speculator, initialize weights, save a checkpoint to " - "--save-path, then exit before training. Useful to validate the config and " - "weights (e.g. in vLLM) before launching a full run. Can be combined with " - "--draft-config or --from-pretrained.", - ) - parser.add_argument( - "--data-path", - type=str, - default="./output", - help=( - "Root data directory containing the preprocessed dataset, " - "vocab mappings (d2t.npy, t2d.npy), token frequencies " - "(token_freq.pt), and hidden states (default: ./output)" - ), - ) - backend_registry = HiddenStatesBackend.registry - parser.add_argument( - "--hidden-states-backend", - choices=list(backend_registry.keys()), - default="file", - help=( - "Hidden states transfer backend. Each backend may add its own " - "CLI arguments (see below). Default: 'file'." - ), - ) - for backend_cls in backend_registry.values(): - backend_cls.add_train_args(parser) - - parser.add_argument( - "--vllm-endpoint", - type=str, - default="http://localhost:8000/v1", - help=( - "vLLM endpoint address to use if generating hidden states on-demand." - " Only required if `--on-missing=generate` and samples are missing." - " Note: the vLLM instance must be configured to cache hidden states" - " to a location that is accessible from the training instance. i.e." - " on the same node, or a shared network drive. (Default: 'http://localhost:8000/v1')" - ), - ) - parser.add_argument( - "--on-missing", - choices=["generate", "skip", "warn", "raise"], - default="generate", - help=( - "Dataloader behaviour when there are no cached hidden states for a sample." - "Default: 'generate', which attempts to generate the hidden states on-" - "demand using the provided vLLM endpoint. The other options skip the sample" - ", skip and warn, or raise an error respectively." - ), - ) - parser.add_argument( - "--on-generate", - choices=["cache", "delete"], - default="delete", - help=( - "Dataloader behaviour when a new hidden state has been generated" - " (only applies if args.on_missing=='generate'). Default: 'delete', " - "deletes hidden states once they are loaded. 'cache' will instead store" - "the hidden states in the args.hidden_states_path. This can be used to " - "enable hybrid online/offline training, with hidden states generated on the" - "first epoch, and reused on subsequent epochs." - ), - ) - parser.add_argument( - "--request-timeout", - type=float, - default=DEFAULT_REQUEST_TIMEOUT, - help=( - "Timeout in seconds for each individual vLLM request " - f"(default: {DEFAULT_REQUEST_TIMEOUT}). " - "Only applies if --on-missing=generate." - ), - ) - parser.add_argument( - "--max-retries", - type=int, - default=DEFAULT_MAX_RETRIES, - help=( - "Maximum number of retry attempts per vLLM request on failure " - f"(default: {DEFAULT_MAX_RETRIES}). " - "Only applies if --on-missing=generate." - ), - ) - parser.add_argument( - "--legacy-data", - action="store_true", - help=( - "DEPRECATED. Use the old data format which stores hidden states alongside " - "token_ids and assistant_masks, in data_i.pt files. This option will be " - "removed soon." - ), - ) - parser.add_argument("--save-path", type=str, default="./output/checkpoints") - parser.add_argument("--epochs", type=int, default=20) - parser.add_argument("--lr", type=float, default=1e-4) - parser.add_argument("--train-data-ratio", type=float, default=0.9) - parser.add_argument("--no-resume-from-checkpoint", action="store_true") - parser.add_argument( - "--logger", - type=str, - default="", - help=( - "One of 'trackio', 'wandb', 'tensorboard', 'mlflow' or " - "comma separated list." - ), - ) - parser.add_argument("--total-seq-len", type=int, default=8192) - parser.add_argument( - "--log-freq", - type=int, - default=1, - help="Log training metrics every N steps (default: 1)", - ) - parser.add_argument("--log-dir", type=str, default="./logs") - parser.add_argument("--run-name", type=str, default=None) - parser.add_argument("--num-layers", type=int, default=1) - parser.add_argument( - "--draft-arch", - type=str, - default=None, - choices=list(DRAFT_ARCH_CONFIGS.keys()), - help="Architecture for draft decoder layers " - "(default: 'llama' for eagle3, 'qwen3' otherwise).", - ) - parser.add_argument( - "--draft-hidden-act", - type=str, - default="silu", - help="Activation function for draft decoder layers. Defaults to 'silu' for " - "sigmoid linear unit. Qwen3 layers of dflash expect 'silu' activation for " - "vLLM deployment. If another function is desired, set as a string or leave " - "as None to automatically fall back to the verifier's activation function.", - ) - parser.add_argument( - "--draft-mrope-full-head-hack", - action=argparse.BooleanOptionalAction, - default=True, - help=( - "For MRoPE configs with partial_rotary_factor < 1, rescale " - "mrope_section and set partial_rotary_factor=1.0 so HF training " - "and vLLM inference use equivalent full-head rotary semantics." - ), - ) - parser.add_argument( - "--target-layer-ids", - type=int, - nargs="+", - help=( - "(Optional) A (space separated) list of integer layer ids. Defaults to" - "[2, num_hidden_layers // 2, num_hidden_layers - 3, num_hidden_layers]. " - "Note: must be set explicitly if custom values were used to launch vllm" - ), - ) - parser.add_argument( - "--token-freq-path", - type=str, - default=None, - help=( - "Path to token frequency distribution file (.pt). Used together with " - "--draft-vocab-size to build vocab mappings at training time. Falls back " - "to '/token_freq.pt' if not provided. If neither that file " - "exists nor --draft-vocab-size is set, vocab mapping is skipped and the " - "full verifier vocab is used." - ), - ) - parser.add_argument( - "--draft-vocab-size", - type=int, - default=None, - help=( - "Vocabulary size for the draft model. Must be provided together with a " - "token frequency file (--token-freq-path or '/token_freq.pt') " - "to generate vocab mappings. If either is absent, vocab mapping is skipped " - "and the full verifier vocab is used, making this argument a no-op." - ), - ) - parser.add_argument("--d2t-path", type=str, default=None) - parser.add_argument("--t2d-path", type=str, default=None) - parser.add_argument("--mask-token-id", type=int, default=None) - parser.add_argument("--ttt-steps", type=int, default=3) - parser.add_argument( - "--num-speculative-steps", - type=int, - default=3, - help="Number of MTP prediction steps (default: 3). Only used with MTP.", - ) - parser.add_argument("--ttt-step-loss-decay", type=float, default=1.0) - parser.add_argument( - "--loss-fn", - type=str, - default="kl_div", - help=( - "Loss function specification. Pass a name for a single loss " - "(kl_div, rkl, jsd, ce, tv, nla, lk_hybrid) or a JSON dict for a weighted " - 'combination, e.g. \'{"ce": 0.1, "tv": 0.9}\'.' - ), - ) - parser.add_argument( - "--step-weight-beta", - type=float, - default=0.6, - help=( - "Exponential decay factor for MTP step weights. " - "Higher values weight earlier prediction steps more heavily. " - "Only used with MTP algorithm." - ), - ) - parser.add_argument( - "--seed", type=int, default=42, help="Random seed for reproducibility" - ) - parser.add_argument( - "--hidden-states-dtype", - type=str, - default="bfloat16", - help="Data type for dataloader hidden states and autocast compute. " - "Model master weights are always kept in fp32. " - "Options: float32 (full precision), bfloat16 (recommended). " - "Note: float16 is not supported (requires gradient scaling).", - ) - parser.add_argument( - "--deterministic-cuda", - action="store_true", - default=False, - help="Sets cuda to deterministic mode. This may impact performance.", - ) - # Model hyperparameters - parser.add_argument( - "--norm-before-residual", - action=argparse.BooleanOptionalAction, - default=True, - help="Toggle normalization before residual connections (default: True)", - ) - parser.add_argument( - "--embed-requires-grad", - action=argparse.BooleanOptionalAction, - default=False, - help="Whether to train embedding layer weights (default: False)", - ) - parser.add_argument( - "--norm-before-fc", - action=argparse.BooleanOptionalAction, - default=None, - help="Apply a single RMSNorm to the concatenated auxiliary hidden states " - "before the FC projection (gpt-oss style). See --fc-norm for the " - "per-layer alternative from the Eagle 3.1 paper. " - "(default: True for eagle3, False otherwise). " - "Disable with --no-norm-before-fc.", - ) - parser.add_argument( - "--fc-norm", - action="store_true", - default=False, - help="Apply per-layer RMSNorm to each auxiliary hidden state before " - "concatenation and FC projection (Eagle 3.1 paper approach).", - ) - parser.add_argument( - "--norm-output", - action=argparse.BooleanOptionalAction, - default=None, - help="Feed post-norm hidden states back across TTT steps to stabilize " - "magnitude drift across speculation depths " - "(default: True for eagle3, False otherwise). " - "Disable with --no-norm-output.", - ) - # D-Flash specific parameters - parser.add_argument( - "--block-size", - type=int, - default=8, - help="Block size for DFlash model (default: 8)", - ) - parser.add_argument( - "--sample-from-anchor", - action=argparse.BooleanOptionalAction, - default=None, - help="Sample from the anchor position (all positions predict). " - "Default: False for dflash, True for dspark. ", - ) - parser.add_argument( - "--max-anchors", - type=int, - default=3072, - help="Maximum anchor positions for DFlash, DSpark, " - "and P-EAGLE training (default: 3072).", - ) - parser.add_argument( - "--dflash-decay-gamma", - type=float, - default=4.0, - help="Decay gamma for DFlash/DSpark loss weighting (default: 4.0)", - ) - # D-Pace specific arguments (loss weight option + smoothing) - parser.add_argument( - "--per-position-loss-weight", - choices=["fixed-exp-decay", "dpace"], - default="fixed-exp-decay", - help="Per-position loss weight option for D-PACE support" - "default: fixed-exp-decay", - ) - parser.add_argument( - "--dpace-alpha", - type=float, - default=0.5, - help="Smoothing constant for D-PACE loss (default: 0.5)", - ) - # DSpark-specific arguments (sequential Markov head + confidence head). - parser.add_argument( - "--markov-rank", - type=int, - default=256, - help="DSpark: low-rank dim of the Markov logit-bias head (0 disables it).", - ) - parser.add_argument( - "--markov-head-type", - type=str, - default="vanilla", - choices=["vanilla", "gated", "rnn"], - help="DSpark: sequential head variant (default: vanilla).", - ) - parser.add_argument( - "--enable-confidence-head", - action=argparse.BooleanOptionalAction, - default=True, - help="DSpark: attach the per-position acceptance confidence head.", - ) - parser.add_argument( - "--confidence-head-with-markov", - action=argparse.BooleanOptionalAction, - default=True, - help="DSpark: feed the Markov previous-token embedding into the " - "confidence head alongside the backbone hidden state.", - ) - parser.add_argument( - "--confidence-head-alpha", - type=float, - default=1.0, - help="DSpark: weight of the confidence-head BCE term (default: 1.0).", - ) - parser.add_argument( - "--draft-attn-impl", - type=str, - default="simple_flex_attention", - choices=["simple_flex_attention", "sdpa", "eager"], - help="Attention implementation for draft layers. " - "Use 'sdpa' or 'eager' for hardware that doesn't support flex attention." - "Not supported for MTP.", - ) - # P-EAGLE specific parameters - parser.add_argument( - "--num-depths", - type=int, - default=8, - help="Number of parallel prediction depths for P-EAGLE (default: 8)", - ) - parser.add_argument( - "--down-sample-ratio", - type=float, - default=0.7, - help="Geometric decay ratio for COD sampling in P-EAGLE (default: 0.7)", - ) - parser.add_argument( - "--down-sample-ratio-min", - type=float, - default=0.2, - help="Minimum retention ratio for COD sampling in P-EAGLE (default: 0.2)", - ) - parser.add_argument( - "--sliding-window", - type=int, - default=2048, - help="Sliding window size for sliding window attention layers (default: 2048). " - "All draft layers use sliding window by default (except mtp).", - ) - parser.add_argument( - "--full-attention-indices", - type=int, - nargs="+", - default=[], - help="(Optional) Space-separated draft layer indices that should use full " - "attention instead of sliding window. All draft layers use sliding window " - "by default (except mtp). " - "(e.g. '--full-attention-indices 0 2' makes layers 0 and 2 use full " - "attention; the rest use sliding window).", - ) - parser.add_argument( - "--sliding-window-non-causal", - action="store_true", - default=False, - help="Use non-causal (bidirectional) masking within draft blocks for sliding " - "window attention layers. Full attention layers are always bidirectional. " - "Note: vLLM currently doesn't support these models.", - ) - # Dataloader parameters - parser.add_argument( - "--num-workers", type=int, default=12, help="Number of dataloader workers" - ) - parser.add_argument( - "--prefetch-factor", type=int, default=4, help="Dataloader prefetch factor" - ) - parser.add_argument( - "--noise-std", - type=float, - default=0.05, - help="Standard deviation for noise augmentation", - ) - # Checkpoint Parameters - parser.add_argument( - "--checkpoint-freq", - type=_checkpoint_freq, - default=1.0, - help="Save a checkpoint every N epochs. Values < 1 enable sub-epoch " - "checkpointing (e.g. 0.5 = every half epoch).", - ) - parser.add_argument( - "--save-best", - action="store_true", - default=False, - help="Pointing to checkpoint with lowest validation loss.", - ) - - # distributed strategy - parser.add_argument( - "--fsdp-shard", - action="store_true", - default=False, - help="Shard model parameters across GPUs with FSDP. By default, " - "parameters are fully replicated (DDP-like). Enable this when the " - "model does not fit in a single GPU's memory.", - ) - - # lr scheduler - parser.add_argument( - "--scheduler-type", - type=str, - default="linear", - choices=["linear", "cosine", "none"], - ) - parser.add_argument("--scheduler-warmup-steps", type=int, default=None) - parser.add_argument( - "--scheduler-warmup-ratio", - type=float, - default=None, - help=( - "Warmup as a fraction of total scheduler steps, in [0, 1]. Ignored " - "(with a warning) when --scheduler-warmup-steps is also set." - ), - ) - parser.add_argument("--scheduler-total-steps", type=int, default=None) - parser.add_argument("--scheduler-num-cosine-cycles", type=float, default=0.5) - - # optimizer - parser.add_argument( - "--optimizer", - type=str, - default="muon", - choices=["adamw", "muon"], - help=( - "Optimizer to use. 'muon' applies Muon to 2D weight matrices and AdamW to " - "the remaining params (norms, biases, embeddings, lm_head)." - ), - ) - parser.add_argument( - "--weight-decay", - type=float, - default=0.01, - help="Weight decay for the AdamW optimizer (and the AdamW group in muon mode).", - ) - parser.add_argument( - "--muon-lr", - type=float, - default=None, - help="LR for the Muon (2D weights) group. Only used with --optimizer muon. " - "Defaults to 10*lr (and --lr defaults to 1e-4)", - ) - parser.add_argument("--muon-momentum", type=float, default=0.95) - parser.add_argument("--muon-weight-decay", type=float, default=0.1) - parser.add_argument("--muon-ns-steps", type=int, default=5) - parser.add_argument( - "--muon-adjust-lr-fn", - type=str, - default="match_rms_adamw", - choices=["original", "match_rms_adamw"], - help="Muon LR adjustment. 'match_rms_adamw' matches AdamW's update RMS.", - ) - - args = parser.parse_args() - - is_eagle3 = args.speculator_type == "eagle3" - if args.draft_arch is None: - args.draft_arch = "llama" if is_eagle3 else "qwen3" - if args.norm_before_fc is None: - args.norm_before_fc = is_eagle3 - if args.norm_output is None: - args.norm_output = is_eagle3 - if args.muon_lr is None: - args.muon_lr = 10 * args.lr - - provided = explicitly_provided_dests(parser, DECODER_SHAPING_FLAGS) - validate_draft_init_args(parser, args, provided) - resolve_loss_config(args.loss_fn) - - if args.per_position_loss_weight == "dpace": - if args.loss_fn != "ce": - parser.error("--per-position-loss-weight=dpace requires --loss-fn=ce") - if not 0.0 < args.dpace_alpha <= 1.0: - raise ValueError(f"alpha must be in (0, 1], got {args.dpace_alpha}") - - return args - - if __name__ == "__main__": - args = parse_args() - main(args) + main(TrainConfig.resolve()) # RUN WITH: diff --git a/src/speculators/data_generation/vllm_client.py b/src/speculators/data_generation/vllm_client.py index cbf0fd634..f784240b8 100644 --- a/src/speculators/data_generation/vllm_client.py +++ b/src/speculators/data_generation/vllm_client.py @@ -96,6 +96,8 @@ def extract_output( token_ids: list[int], ) -> str: if isinstance(response, Completion): + if not response.choices: + raise InvalidResponseError("Response has no choices") prompt_token_ids = getattr(response.choices[0], "prompt_token_ids", None) else: prompt_token_ids = getattr(response, "prompt_token_ids", None) @@ -188,8 +190,8 @@ async def generate_hidden_states_async( timeout: float | None = DEFAULT_REQUEST_TIMEOUT, ) -> str: """ - Runs decode w/ max_tokens 1 to generate hidden states and returns path to - hidden states file. + Runs prefill to extract hidden states and returns path to hidden states + file. Args: client: The async OpenAI client. @@ -241,8 +243,8 @@ def generate_hidden_states( timeout: float | None = DEFAULT_REQUEST_TIMEOUT, ) -> str: """ - Runs decode w/ max_tokens 1 to generate hidden states and returns path to - hidden states file. + Runs prefill to extract hidden states and returns path to hidden states + file. """ token_ids = client_item["input_ids"] messages = client_item.get("messages") diff --git a/src/speculators/models/attention.py b/src/speculators/models/attention.py index 4467700a4..3e6eee19d 100644 --- a/src/speculators/models/attention.py +++ b/src/speculators/models/attention.py @@ -26,10 +26,12 @@ def flex_attention_forward( scaling: float | None = None, **_kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Shared flex attention forward implementation. + """Shared flex attention forward with optional Ulysses sequence parallelism. - This function is used by both EAGLE3 and DFlash attention mechanisms to avoid - code duplication and ensure consistent behavior. + When ``sp_size > 1``, Q/K/V are transposed via all-to-all from + sequence-parallel layout ``(B, H, S_local, D)`` to head-parallel + layout ``(B, H/sp, S_full, D)`` before attention, and the output + is transposed back afterwards. Args: module: The attention module (unused but required for interface compatibility). @@ -44,6 +46,47 @@ def flex_attention_forward( Tuple of (attention_output, None) where attention_output has shape (batch, seq_len, num_heads, head_dim) and None represents no attention weights. """ + from speculators.train.distributed import get_sp_group, get_sp_size # noqa: PLC0415 + from speculators.train.sequence_parallel import ( # noqa: PLC0415 + maybe_replicate_kv_heads, + ulysses_gather, + ulysses_scatter, + ) + + sp_size = get_sp_size() + use_sp = sp_size > 1 + + if use_sp: + sp_group = get_sp_group() + if sp_group is None: + raise RuntimeError( + "SP group did not initialize for sequence parallelism, \ + something went wrong!" + ) + + key, value = maybe_replicate_kv_heads(key, value, sp_size) + + q_len = query.shape[2] + kv_len = key.shape[2] + + query = ulysses_scatter(query, sp_group, sp_size) + + if kv_len > q_len: + # KV cache from TTT: K/V = [step0_local | step1_local | ...]. + # Scatter each step separately so the result is + # [step0_full | step1_full | ...] matching the extended mask. + k_parts = key.split(q_len, dim=2) + v_parts = value.split(q_len, dim=2) + key = torch.cat( + [ulysses_scatter(k, sp_group, sp_size) for k in k_parts], dim=2 + ) + value = torch.cat( + [ulysses_scatter(v, sp_group, sp_size) for v in v_parts], dim=2 + ) + else: + key = ulysses_scatter(key, sp_group, sp_size) + value = ulysses_scatter(value, sp_group, sp_size) + num_query_heads = query.shape[1] num_key_value_heads = key.shape[1] enable_gqa = num_query_heads != num_key_value_heads @@ -61,7 +104,11 @@ def flex_attention_forward( enable_gqa=enable_gqa, scale=scaling, ) - attention_output: torch.Tensor = flex_attention_output + attention_output: torch.Tensor = flex_attention_output # type: ignore[assignment] + + if use_sp: + attention_output = ulysses_gather(attention_output, sp_group, sp_size) # type: ignore[arg-type] + attention_output = attention_output.transpose(1, 2).contiguous() return attention_output, None @@ -106,3 +153,16 @@ def block_mask_to_dense_attention_mask( # Singleton registry for attention functions (shared across all models) ALL_ATTENTION_FUNCTIONS = AttentionInterface() ALL_ATTENTION_FUNCTIONS.register("simple_flex_attention", flex_attention_forward) + + +def _dflash_flex_attention_forward(*args, **kwargs): + from speculators.train.sequence_parallel import ( # noqa: PLC0415 + dflash_flex_attention_forward, + ) + + return dflash_flex_attention_forward(*args, **kwargs) + + +ALL_ATTENTION_FUNCTIONS.register( + "dflash_flex_attention", _dflash_flex_attention_forward +) diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index 0bad5418e..590fa6b50 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -2,6 +2,7 @@ from typing import ClassVar import torch +import torch.distributed as dist from torch import nn from torch.nn.attention.flex_attention import create_block_mask, create_mask from transformers import PretrainedConfig @@ -51,6 +52,8 @@ class DFlashDraftModel(DraftVocabMixin, SpeculatorModel): t2d: torch.Tensor | None d2t: torch.Tensor | None + _sp_splits_batch: ClassVar[bool] = False + def __init__( self, config: DFlashSpeculatorConfig, @@ -68,6 +71,14 @@ def __init__( if self._attn_impl == "eager" else create_mask ) + + from speculators.train.distributed import get_sp_size # noqa: PLC0415 + + if get_sp_size() > 1: + config.transformer_layer_config._attn_implementation = ( # noqa: SLF001 + "dflash_flex_attention" + ) + super().__init__(config=config) self._init_vocab(config) @@ -289,12 +300,25 @@ def _create_attention_mask( @torch.compiler.disable def _build_attention_mask(self, loss_mask, max_anchors, document_ids, device): + from speculators.train.distributed import ( # noqa: PLC0415 + get_sp_group, + get_sp_rank, + get_sp_size, + ) + total_seq_len = loss_mask.shape[1] + sp_size = get_sp_size() anchor_positions, anchor_valid = select_anchors( loss_mask, max_anchors, self.block_size ) + if sp_size > 1: + sp_group = get_sp_group() + src = dist.get_process_group_ranks(sp_group)[0] + dist.broadcast(anchor_positions, src=src, group=sp_group) + dist.broadcast(anchor_valid, src=src, group=sp_group) + full_attn_mask = None if self.uses_full_attn: full_attn_mask = self._create_attention_mask( @@ -316,6 +340,13 @@ def _build_attention_mask(self, loss_mask, max_anchors, document_ids, device): sliding_window_non_causal=self.sliding_window_non_causal, ) + if sp_size > 1: + sp_rank = get_sp_rank() + n_per_rank = max_anchors // sp_size + local_start = sp_rank * n_per_rank + anchor_positions = anchor_positions[local_start : local_start + n_per_rank] + anchor_valid = anchor_valid[local_start : local_start + n_per_rank] + return full_attn_mask, sliding_window_attn_mask, anchor_positions, anchor_valid def _backbone_forward( @@ -333,11 +364,26 @@ def _backbone_forward( Returns ``(hidden, logits, targets, aligned_loss_mask, anchored_block_indices)``. DSpark reuses this and adds its Markov and confidence heads before computing its own loss. + + When ``sp_size > 1``, each SP rank processes a local slice of the + context and a partition of the anchors. The attention function + handles the all-to-all communication to reconstruct the global + context and noise inside each layer. """ + from speculators.train.distributed import ( # noqa: PLC0415 + get_sp_rank, + get_sp_size, + ) + device = hidden_states.device total_seq_len = hidden_states.shape[1] num_anchors = kwargs.pop("max_anchors", 3072) + sp_size = get_sp_size() + sp_rank = get_sp_rank() + local_seq_len = total_seq_len // sp_size if sp_size > 1 else total_seq_len + sp_start = sp_rank * local_seq_len if sp_size > 1 else 0 + if position_ids is None: position_ids = torch.arange( total_seq_len, dtype=torch.long, device=device @@ -347,7 +393,8 @@ def _backbone_forward( self._build_attention_mask(loss_mask, num_anchors, document_ids, device) ) - mask_tokens_size = num_anchors * self.block_size + local_num_anchors = len(anchor_positions) + mask_tokens_size = local_num_anchors * self.block_size mask_token_ids = torch.full( (1, mask_tokens_size), @@ -359,33 +406,50 @@ def _backbone_forward( noise_embedding = self.embed_tokens(mask_token_ids) # shape: [1, num_anchors*block_size, hidden_size] - fc_output = self.fc(hidden_states) + fc_output = self.fc(hidden_states[:, sp_start : sp_start + local_seq_len]) fc_output = self.hidden_norm(fc_output) - # shape: [1, total_seq_len, hidden_size] + # shape: [1, local_seq_len, hidden_size] mask_position_ids = get_base_indices_for_anchored_blocks( position_ids[0, anchor_positions], self.block_size ) - position_ids = torch.cat([position_ids, mask_position_ids.unsqueeze(0)], dim=1) + local_position_ids = torch.cat( + [ + position_ids[:, sp_start : sp_start + local_seq_len], + mask_position_ids.unsqueeze(0), + ], + dim=1, + ) # shape: [1, total_seq_len + num_anchors*block_size] # the hidden_states shape doesn't match position_ids but doesn't need # to, as hidden_states is only used to set dtype and device in rotary_emb - position_embeddings = self.rotary_emb(hidden_states, position_ids) + position_embeddings = self.rotary_emb( + hidden_states[:, sp_start : sp_start + local_seq_len], + local_position_ids, + ) anchored_block_indices = get_base_indices_for_anchored_blocks( anchor_positions, self.block_size ) # shape: [num_anchors*block_size] with torch.no_grad(): - verifier_logits = self.verifier_lm_head( - self.verifier_norm(verifier_last_hidden_states) - ) - if not self.config.sample_from_anchor: - # False: shift right by 1 so slot j predicts token at position j - verifier_logits = torch.roll(verifier_logits, 1, dims=1) - # else: True, slot k predicts token at position k+1 (next), no shift - targets = verifier_logits[:, anchored_block_indices] + if anchored_block_indices.numel() < total_seq_len: + target_indices = ( + anchored_block_indices + if self.config.sample_from_anchor + else (anchored_block_indices - 1) % total_seq_len + ) + targets = self.verifier_lm_head( + self.verifier_norm(verifier_last_hidden_states[:, target_indices]) + ) + else: + verifier_logits = self.verifier_lm_head( + self.verifier_norm(verifier_last_hidden_states) + ) + if not self.config.sample_from_anchor: + verifier_logits = torch.roll(verifier_logits, 1, dims=1) + targets = verifier_logits[:, anchored_block_indices] # shape: [1, num_anchors*block_size, draft_vocab_size] for layer_idx, layer in enumerate(self.layers): @@ -395,7 +459,7 @@ def _backbone_forward( attention_mask=sliding_window_attn_mask if layer_idx in self.sliding_window_indices else full_attn_mask, - position_ids=position_ids, + position_ids=local_position_ids, use_cache=False, position_embeddings=position_embeddings, **kwargs, diff --git a/src/speculators/models/dflash/model_definitions.py b/src/speculators/models/dflash/model_definitions.py index 0391e2a17..8ac22b6cb 100644 --- a/src/speculators/models/dflash/model_definitions.py +++ b/src/speculators/models/dflash/model_definitions.py @@ -107,23 +107,16 @@ def forward( # Instead of computing the k and v matricies from the hidden states, # the target_hidden is injected into the kv cache, (shape is context # length + block size) - bsz, q_len = hidden_states.shape[:-1] - ctx_len = target_hidden.shape[1] q = self.q_proj(hidden_states) - q = q.view(bsz, q_len, -1, self.head_dim) + q = q.unflatten(-1, (-1, self.head_dim)) q = self.q_norm(q).transpose(1, 2) # This is the main difference from the usual attention mechanism. k_ctx = self.k_proj(target_hidden) k_noise = self.k_proj(hidden_states) v_ctx = self.v_proj(target_hidden) v_noise = self.v_proj(hidden_states) - k = torch.cat([k_ctx, k_noise], dim=1).view( - bsz, ctx_len + q_len, -1, self.head_dim - ) - # note the length becomes context length + block size - v = torch.cat([v_ctx, v_noise], dim=1).view( - bsz, ctx_len + q_len, -1, self.head_dim - ) + k = torch.cat([k_ctx, k_noise], dim=1).unflatten(-1, (-1, self.head_dim)) + v = torch.cat([v_ctx, v_noise], dim=1).unflatten(-1, (-1, self.head_dim)) k = self.k_norm(k).transpose(1, 2) v = v.transpose(1, 2) cos, sin = position_embeddings @@ -150,7 +143,7 @@ def forward( sliding_window=self.sliding_window, **kwargs, ) - attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = attn_output.flatten(2) attn_output = self.o_proj(attn_output) return attn_output, attn_weights diff --git a/src/speculators/models/dspark/core.py b/src/speculators/models/dspark/core.py index a2382f717..6b42c338d 100644 --- a/src/speculators/models/dspark/core.py +++ b/src/speculators/models/dspark/core.py @@ -138,9 +138,8 @@ def forward( ) ) - # DSpark: add the Markov logit bias and predict per-position confidence. - num_blocks = max_anchors block = self.block_size + num_blocks = anchored_block_indices.shape[0] // block mask_tokens_size = num_blocks * block # Ground-truth block tokens (verifier vocab); position 0 is the anchor. block_tokens = input_ids[0, anchored_block_indices].view(num_blocks, block) diff --git a/src/speculators/models/dspark/metrics.py b/src/speculators/models/dspark/metrics.py index 4b5d54e5f..391f226e9 100644 --- a/src/speculators/models/dspark/metrics.py +++ b/src/speculators/models/dspark/metrics.py @@ -35,14 +35,25 @@ def _masked_decayed_mean( decay_fn: Callable[..., torch.Tensor] | None, ) -> torch.Tensor: """Masked, optionally position-decayed mean of a precomputed per-position term.""" + import torch.distributed as dist # noqa: PLC0415 + + from speculators.train.distributed import get_sp_group, get_sp_size # noqa: PLC0415 + loss_mask = loss_mask.to(elementwise.dtype) weighted = elementwise * loss_mask if decay_fn is not None: weighted = weighted * decay_fn( pos_idx.to(weighted.dtype), elementwise_loss=elementwise ) + numerator = weighted.sum(dim=1) denominator = loss_mask.sum(dim=1) + _EPS - return (weighted.sum(dim=1) / denominator).mean() + + if get_sp_size() > 1: + sp_group = get_sp_group() + dist.all_reduce(numerator, group=sp_group) + dist.all_reduce(denominator, group=sp_group) + + return (numerator / denominator).mean() def compute_metrics( diff --git a/src/speculators/models/eagle3/core.py b/src/speculators/models/eagle3/core.py index e4e006bbd..01158e451 100644 --- a/src/speculators/models/eagle3/core.py +++ b/src/speculators/models/eagle3/core.py @@ -200,13 +200,15 @@ def forward( # noqa: C901 **kwargs, ): device = hidden_states.device - total_seq_len = hidden_states.shape[1] + # With Ulysses SP, hidden_states is split but document_ids is full. + local_seq_len = hidden_states.shape[1] + total_seq_len = document_ids.shape[1] if position_ids is None: position_ids = 1 + torch.arange( - total_seq_len, dtype=torch.long, device=device + local_seq_len, dtype=torch.long, device=device ).unsqueeze(0) - # shape: [1, total_seq_len] + # shape: [1, local_seq_len] past_key_values = DynamicCache() @@ -237,7 +239,7 @@ def forward( # noqa: C901 dim=-1, ) hidden_states = self.fc(hidden_states) - # shape: [1, total_seq_len, hidden_size] + # shape: [1, local_seq_len, hidden_size] original_input_ids = input_ids.detach().clone() return_loss = verifier_last_hidden_states is not None @@ -246,7 +248,7 @@ def forward( # noqa: C901 targets = self.verifier_lm_head( self.verifier_norm(verifier_last_hidden_states) ) - # shape: [1, total_seq_len, draft_vocab_size] + # shape: [1, local_seq_len, draft_vocab_size] loss = torch.tensor(0.0, device=device) # prev_correct is a boolean tensor that is True for tokens that have been @@ -256,7 +258,7 @@ def forward( # noqa: C901 prev_correct = ( loss_mask.clone() if loss_mask is not None - else torch.ones(1, total_seq_len, device=device, dtype=torch.bool) + else torch.ones(1, local_seq_len, device=device, dtype=torch.bool) ) metrics = {} @@ -264,14 +266,14 @@ def forward( # noqa: C901 for ttt_step in range(ttt_steps): with torch.no_grad(): input_embeds = self.embed_tokens(input_ids) - # shape: [1, total_seq_len, hidden_size] + # shape: [1, local_seq_len, hidden_size] cache_position = torch.arange( - ttt_step * total_seq_len, - (ttt_step + 1) * total_seq_len, + ttt_step * local_seq_len, + (ttt_step + 1) * local_seq_len, dtype=torch.long, device=device, ) - # shape: [total_seq_len] + # shape: [local_seq_len] hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) # shape: [1, total_seq_len, 2 * hidden_size] diff --git a/src/speculators/models/metrics.py b/src/speculators/models/metrics.py index 77105369d..101f22e30 100644 --- a/src/speculators/models/metrics.py +++ b/src/speculators/models/metrics.py @@ -514,6 +514,12 @@ def loss_function( ): """Compute masked, optionally position-decayed training loss. + When Ulysses SP is active (``sp_size > 1``), each rank holds a + local chunk of the sequence. The numerator and denominator are + all-reduced across the SP group so that every rank computes the + same globally-averaged loss, producing correct gradients after + the SP gradient all-reduce. + Args: logits: Draft model logits. targets: Target model logits. @@ -525,6 +531,10 @@ def loss_function( Returns: Scalar mean loss across the batch. """ + import torch.distributed as dist # noqa: PLC0415 + + from speculators.train.distributed import get_sp_group, get_sp_size # noqa: PLC0415 + elementwise_loss = loss_fn(logits, targets) # shape: [1, seq_len] loss_mask = loss_mask.to(elementwise_loss.dtype) @@ -536,7 +546,15 @@ def loss_function( ) elementwise_loss = elementwise_loss * decay_mult + numerator = torch.sum(elementwise_loss, dim=1) denominator = loss_mask.sum(dim=1) + _EPS - batch_loss = torch.sum(elementwise_loss, dim=1) / denominator # shape: [1] + # will need to manually detach graphs if all_reduce + # ever starts getting tracked by autograd + if get_sp_size() > 1: + sp_group = get_sp_group() + dist.all_reduce(numerator, group=sp_group) + dist.all_reduce(denominator, group=sp_group) + + batch_loss = numerator / denominator # shape: [1] return batch_loss.mean() # shape: [] diff --git a/src/speculators/models/mtp/core.py b/src/speculators/models/mtp/core.py index a3d3c67e6..a74f37388 100644 --- a/src/speculators/models/mtp/core.py +++ b/src/speculators/models/mtp/core.py @@ -143,6 +143,13 @@ def forward( (lengths, verifier_last_hidden_states) :return: Tuple of (logits_list, loss, metrics) """ + from speculators.train.distributed import get_sp_size # noqa: PLC0415 + + if get_sp_size() > 1: + raise NotImplementedError( + "MTPDraftModel does not yet support sequence parallelism (sp_size > 1)" + ) + input_ids = input_ids.long() device = input_ids.device batch_size, seq_len = input_ids.shape diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index b8d11edde..2e3022367 100644 --- a/src/speculators/models/peagle/core.py +++ b/src/speculators/models/peagle/core.py @@ -3,6 +3,7 @@ from typing import ClassVar import torch +import torch.distributed as dist from transformers import PretrainedConfig from speculators.config import SpeculatorsConfig, VerifierConfig @@ -31,6 +32,7 @@ class PEagleDraftModel(Eagle3DraftModel): *Eagle3DraftModel._keys_to_ignore_on_load_missing, # noqa: SLF001 "mask_hidden", ] + _sp_splits_batch: ClassVar[bool] = False def __init__( self, @@ -50,6 +52,79 @@ def __init__( torch.randn(1, 1, num_aux * self.hidden_size) ) + @torch.compiler.disable + def _sample_and_partition( + self, + seq_length: int, + loss_mask: torch.Tensor, + num_depths: int, + down_sample_ratio: float, + down_sample_ratio_min: float, + max_anchors: int | None, + device: torch.device, + ): + """Generate COD samples and partition across SP ranks. + + Returns (full_anchor_pos, full_depth, local_anchor_pos, local_depth, + local_sampled, pad_mask). + """ + from speculators.train.distributed import ( # noqa: PLC0415 + get_sp_group, + get_sp_rank, + get_sp_size, + ) + + sp_size = get_sp_size() + + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_length, + loss_mask=loss_mask, + num_depths=num_depths, + down_sample_ratio=down_sample_ratio, + down_sample_ratio_min=down_sample_ratio_min, + max_anchors=max_anchors, + ) + + pad_mask: torch.Tensor | None = None + + if sp_size > 1: + sp_group = get_sp_group() + sp_rank = get_sp_rank() + src = dist.get_process_group_ranks(sp_group)[0] + + size_t = torch.tensor(anchor_pos.shape[0], device=device) + dist.broadcast(size_t, src=src, group=sp_group) + n_real = int(size_t.item()) + if anchor_pos.shape[0] != n_real: + anchor_pos = torch.empty(n_real, device=device, dtype=torch.long) + depth = torch.empty(n_real, device=device, dtype=torch.long) + dist.broadcast(anchor_pos, src=src, group=sp_group) + dist.broadcast(depth, src=src, group=sp_group) + + remainder = n_real % sp_size + if remainder != 0: + pad_len = sp_size - remainder + anchor_pos = torch.nn.functional.pad(anchor_pos, (0, pad_len)) + depth = torch.nn.functional.pad(depth, (0, pad_len)) + + total_sampled = anchor_pos.shape[0] + n_per_rank = total_sampled // sp_size + local_start = sp_rank * n_per_rank + local_anchor_pos = anchor_pos[local_start : local_start + n_per_rank] + local_depth = depth[local_start : local_start + n_per_rank] + local_sampled = n_per_rank + + if n_real < total_sampled: + full_valid = torch.ones(total_sampled, device=device) + full_valid[n_real:] = 0 + pad_mask = full_valid[local_start : local_start + n_per_rank] + else: + local_anchor_pos = anchor_pos + local_depth = depth + local_sampled = anchor_pos.shape[0] + + return anchor_pos, depth, local_anchor_pos, local_depth, local_sampled, pad_mask + @conditional_torch_compile def forward( self, @@ -91,24 +166,34 @@ def forward( if loss_mask is None: loss_mask = torch.ones_like(input_ids, dtype=torch.float32) - # Generate COD sampling indices - anchor_pos, depth = generate_cod_sample_indices( - seq_length=seq_length, - loss_mask=loss_mask, - num_depths=num_depths, - down_sample_ratio=down_sample_ratio, - down_sample_ratio_min=down_sample_ratio_min, - max_anchors=max_anchors, + # COD sampling + SP broadcast/partition (compiler-disabled to avoid + # graph breaks from collectives and .item()) + ( + full_anchor_pos, + full_depth, + local_anchor_pos, + local_depth, + local_sampled, + pad_mask, + ) = self._sample_and_partition( + seq_length, + loss_mask, + num_depths, + down_sample_ratio, + down_sample_ratio_min, + max_anchors, + device, ) - total_sampled = anchor_pos.shape[0] - orig_positions = anchor_pos + depth - is_depth_0 = depth == 0 # [total_sampled] + total_sampled = full_anchor_pos.shape[0] + local_orig_positions = local_anchor_pos + local_depth + + is_depth_0 = local_depth == 0 # Build sampled input_ids: real tokens for depth 0, mask for others sampled_ids = torch.where( is_depth_0, - input_ids[0, orig_positions], + input_ids[0, local_orig_positions], torch.tensor(self.mask_token_id, dtype=input_ids.dtype, device=device), ).unsqueeze(0) # [1, total_sampled] inputs_embeds = self.embed_tokens(sampled_ids).to( @@ -119,8 +204,8 @@ def forward( mask_hidden = self.mask_hidden.to(device=device, dtype=hidden_states.dtype) sampled_hidden = torch.where( is_depth_0.unsqueeze(-1), - hidden_states[0, orig_positions], - mask_hidden.squeeze(0).expand(orig_positions.shape[0], -1), + hidden_states[0, local_orig_positions], + mask_hidden.squeeze(0).expand(local_sampled, -1), ).unsqueeze(0) # [1, total_sampled, 3*hidden_size] # Project concatenated hidden states (3*hidden_size) -> hidden_size @@ -138,16 +223,17 @@ def forward( [inputs_embeds, sampled_hidden], dim=-1 ) # [1, total_sampled, 2*hidden_size] - position_ids = orig_positions.unsqueeze(0) # [1, total_sampled] + position_ids = local_orig_positions.unsqueeze(0) # [1, total_sampled] position_embeddings = self.rotary_emb(layer_input, position_ids) doc_ids_1d = document_ids.squeeze(0).to(device) + # Build masks at FULL total_sampled scale (matches post-all-to-all layout) def _build_attn_mask(sliding_window=None): mask_mod = create_peagle_mask_mod( - anchor_pos=anchor_pos, - depth=depth, + anchor_pos=full_anchor_pos, + depth=full_depth, document_ids=doc_ids_1d, sliding_window=sliding_window, ) @@ -191,16 +277,17 @@ def _build_attn_mask(sliding_window=None): self.verifier_norm(verifier_last_hidden_states) ) - targets = targets[:, orig_positions, :] # [1, total_sampled, vocab_size] + targets = targets[:, local_orig_positions, :] # [1, total_sampled, vocab_size] loss, metrics = compute_metrics( logits=logits, targets=targets, loss_mask=loss_mask, - anchor_pos=anchor_pos, - depth=depth, + anchor_pos=local_anchor_pos, + depth=local_depth, num_depths=num_depths, loss_config=loss_config, + pad_mask=pad_mask, ) return None, loss, metrics diff --git a/src/speculators/models/peagle/metrics.py b/src/speculators/models/peagle/metrics.py index a658af36f..9d3c6c3b2 100644 --- a/src/speculators/models/peagle/metrics.py +++ b/src/speculators/models/peagle/metrics.py @@ -22,6 +22,7 @@ def compute_metrics( depth: torch.Tensor, num_depths: int, loss_config: LossConfig | None = None, + pad_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, Any]]: """Compute loss and accuracy metrics for P-EAGLE predictions. @@ -46,6 +47,8 @@ def compute_metrics( # shared loss_function/compute_accuracy_multi_step shape contract orig_positions = anchor_pos + depth # [total_sampled] sampled_loss_mask = loss_mask[:, orig_positions] # [1, total_sampled] + if pad_mask is not None: + sampled_loss_mask = sampled_loss_mask * pad_mask.unsqueeze(0) loss, term_losses = compound_loss( logits, targets, sampled_loss_mask, depth.unsqueeze(0), loss_config=loss_config diff --git a/src/speculators/train/config/__init__.py b/src/speculators/train/config/__init__.py new file mode 100644 index 000000000..1011fec2c --- /dev/null +++ b/src/speculators/train/config/__init__.py @@ -0,0 +1,9 @@ +"""Config-file-first configuration for ``scripts/train.py``. + +One public type, :class:`~.schema.TrainConfig`, split across ``schema`` (fields), +``resolution`` (CLI + precedence), and ``artifacts`` (reproducibility I/O). +""" + +from .schema import TrainConfig + +__all__ = ["TrainConfig"] diff --git a/src/speculators/train/config/artifacts.py b/src/speculators/train/config/artifacts.py new file mode 100644 index 000000000..1c0d8aaaf --- /dev/null +++ b/src/speculators/train/config/artifacts.py @@ -0,0 +1,76 @@ +"""Reproducibility artifacts: the ``run.yaml`` + ``train_command.txt`` writers. + +``run.yaml`` is the resolved, provenance-free config that round-trips through +``--config``; ``train_command.txt`` records argv, git SHA, world size, versions. +""" + +from pathlib import Path +from typing import TYPE_CHECKING + +import yaml + +from .schema import nest_flat + +if TYPE_CHECKING: + from .schema import TrainConfig + + +def dump_yaml(cfg: "TrainConfig") -> str: + """Serialize a resolved config to the stage-shaped, round-trippable YAML. + + Everything nests under a top-level ``train:`` key, so the file re-loads + cleanly via ``--config`` and leaves room for a future ``prepare_data:`` / + ``launch_vllm:`` stage. + + Only values a non-default layer supplied (flag or yaml) are emitted; defaults + are omitted, which is what makes the file round-trip: the draft-init conflict + check treats any YAML-supplied decoder-shaping key as explicit, so persisting + every materialized default would make a ``--from-pretrained`` run.yaml reject + itself on reload. Re-resolving the emitted subset re-derives the omitted + defaults, yielding an identical config. The provenance record only selects + which keys to emit; it is never inlined. + """ + # Iterate the pinned flatten() order (not the provenance dict, whose key + # order varies with how the config was populated) so the emitted YAML is + # byte-stable across reloads regardless of flag-vs-yaml source. + resolved = cfg.flatten() + if cfg.provenance: + provided = { + dest: value + for dest, value in resolved.items() + if cfg.provenance[dest] != "default" + } + else: + # A from_flat config carries no layer provenance, so "customized" is + # defined as "differs from the default-constructed config". + baseline = type(cfg)().flatten() + provided = { + dest: value + for dest, value in resolved.items() + if value != baseline.get(dest) + } + return yaml.safe_dump( + {"train": nest_flat(provided)}, + sort_keys=False, + default_flow_style=False, + ) + + +def save(cfg: "TrainConfig", save_dir: str) -> None: + """Write ``run.yaml`` + ``train_command.txt`` next to the checkpoints. + + Called at rank 0 by ``scripts/train.py`` so every checkpoint carries the + config that produced it. ``run.yaml`` is the clean resolved config (re-run via + ``--config run.yaml``); ``train_command.txt`` records the exact argv this + config was resolved from plus the environment manifest. + """ + # Local import: utils pulls in the (heavier) preprocessing stack, so keeping + # it out of config import time avoids a needless cost and any import cycle. + from speculators.train.utils import save_train_command # noqa: PLC0415 + + out_dir = Path(save_dir) + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "run.yaml").write_text(dump_yaml(cfg)) + # Record the argv this config was resolved from (falls back to the live + # sys.argv when the config was built off-argv, e.g. via from_flat). + save_train_command(save_dir, argv=list(cfg.argv)) diff --git a/src/speculators/train/config/resolution.py b/src/speculators/train/config/resolution.py new file mode 100644 index 000000000..cc03b1faf --- /dev/null +++ b/src/speculators/train/config/resolution.py @@ -0,0 +1,455 @@ +"""Turns the schema fields into an argparse CLI and applies precedence. + +Precedence is ``flag > yaml > default`` via pydantic-settings source ordering; +``resolve`` turns any config failure into a clean ``SystemExit(2)``. +""" + +import argparse +import sys +import types +import warnings +from pathlib import Path +from typing import Any, Literal, Union, get_args, get_origin + +import yaml +from pydantic import ValidationError +from pydantic.fields import FieldInfo + +from hs_connectors import HiddenStatesBackend +from speculators.train.config.schema import ( + _GROUPS, + _ROOT_FIELDS, + CONFIG_DESTS, + TrainConfig, + nest_flat, +) + +__all__ = [ + "DECODER_SHAPING_FLAGS", + "ConfigError", + "add_config_cli_arguments", + "build_from_sources", + "build_parser", + "resolve", +] + + +class ConfigError(ValueError): + """A user-facing configuration error (e.g. a draft-init conflict). + + Subclasses ``ValueError`` so the pure core can raise it for a unit test to + catch, while the :func:`resolve` boundary renders it as a clean + ``SystemExit(2)`` instead of a traceback. + """ + + +def _dest_to_flag(dest: str) -> str: + """The argparse flag string for a config dest (``muon_lr`` -> ``--muon-lr``).""" + return "--" + dest.replace("_", "-") + + +# Decoder-shaping dests -- the flags that synthesize the draft decoder and so +# conflict with --from-pretrained / --draft-config (each of which fully defines +# the draft). Derived from the schema so the set stays in one place. +# dest -> CLI flag; the flag strings are used only for human-readable ConfigError +# messages. +DECODER_SHAPING_FLAGS: dict[str, str] = { + dest: _dest_to_flag(dest) + for dest in ( + "num_layers", + "draft_arch", + "draft_hidden_act", + "sliding_window", + "full_attention_indices", + ) +} + + +def _required_flags() -> dict[str, str]: + """The dests marked ``cli_required`` in the schema, mapped to their flags. + + Scans the root scalars and every group for the ``cli_required`` marker so the + required contract stays schema-driven -- adding another required field needs no + change here. Enforced post-build (see :func:`build_from_sources`), not by + argparse, so a value supplied only in ``--config`` still satisfies it. + """ + flags: dict[str, str] = {} + fields = dict(TrainConfig.model_fields) + for gmodel in _GROUPS.values(): + fields.update(gmodel.model_fields) + for dest, field in fields.items(): + extra = field.json_schema_extra + if isinstance(extra, dict) and extra.get("cli_required"): + flags[dest] = _dest_to_flag(dest) + return flags + + +# Flat dests the schema marks required (``cli_required``), mapped to their flags. +REQUIRED_FLAGS: dict[str, str] = _required_flags() + + +# Algorithm group -> the speculator types that consume it. A group set to a +# non-default value under a speculator_type absent from its set is ignored, and +# warns. DSpark is-a DFlash, so a dspark run reads the dflash group too; the +# dspark-exclusive heads belong only to dspark. eagle3 uses no group. +_ALGORITHM_GROUP_USERS: dict[str, frozenset[str]] = { + "dflash": frozenset({"dflash", "dspark"}), + "dspark": frozenset({"dspark"}), + "peagle": frozenset({"peagle"}), + "mtp": frozenset({"mtp"}), +} + + +# Marker key (schema ``_CLI_CHOICES``) -> the live registry supplying a field's +# argparse choices. Binding the registry here -- not in the schema -- keeps the schema +# free of runtime backend objects. The schema, not each backend's add_train_args, now +# generates the flag surface: every backend's train-args are mirrored as schema fields +# (e.g. the 'file' backend's --hidden-states-path is DataArgs' hidden_states_path) so +# they survive resolution's CONFIG_DESTS filter; test_backend_reconciliation.py guards +# against a backend adding a train-arg with no matching schema field (which would +# otherwise be silently dropped). +_CLI_CHOICE_REGISTRIES: dict[str, Any] = { + "hidden_states_backends": HiddenStatesBackend.registry, +} + + +def _annotation_spec(annotation: Any) -> tuple[Any, bool, list[Any] | None]: + """Reduce a field annotation to ``(base_type, is_list, choices)``. + + Strips an optional ``| None``; maps ``list[T]`` to a ``nargs`` list and + ``Literal[...]`` to ``choices``. ``base_type`` is the argparse ``type=`` + callable; for a ``Literal`` it is the element type so parsed values and + ``choices`` compare like-for-like. A genuine multi-type union (e.g. + ``int | str``) is rejected rather than silently truncated. + """ + origin = get_origin(annotation) + if origin in (Union, types.UnionType): + non_none = [arg for arg in get_args(annotation) if arg is not type(None)] + if len(non_none) != 1: + raise TypeError( + f"CLI generation supports only single-type optionals; got " + f"{annotation!r} with {len(non_none)} non-None arms." + ) + annotation = non_none[0] + origin = get_origin(annotation) + if origin is list: + (inner,) = get_args(annotation) + return inner, True, None + if origin is Literal: + choices = list(get_args(annotation)) + elem_types = {type(choice) for choice in choices} + base = elem_types.pop() if len(elem_types) == 1 else str + return base, False, choices + return annotation, False, None + + +def _add_field_argument( + parser: argparse._ArgumentGroup, name: str, field: FieldInfo +) -> None: + """Add one argparse flag derived from a pydantic field. + + Defaults are ``SUPPRESS``ed so the parsed namespace holds only the flags the + user actually passed -- that set is the ``flag`` layer of the precedence walk. + """ + base, is_list, choices = _annotation_spec(field.annotation) + flag = _dest_to_flag(name) + kwargs: dict[str, Any] = {"dest": name, "default": argparse.SUPPRESS} + if field.description: + kwargs["help"] = field.description + extra = field.json_schema_extra if isinstance(field.json_schema_extra, dict) else {} + + # A str field may draw its choices from a dynamic registry named indirectly by a + # _CLI_CHOICES marker, so the choice set is resolved here (where the registry is + # known) with no field-name literal. + choices_key = extra.get("cli_choices") + if isinstance(choices_key, str): + choices = list(_CLI_CHOICE_REGISTRIES[choices_key]) + + if base is bool and not is_list: + # True/None-default bools, or those explicitly tagged, render as + # --x/--no-x; a plain False-default bool is a simple store_true flag. + # store_true can only set true, so it can't flip a YAML-set true back to + # false -- a bool needing CLI false-override must be tagged _CLI_BOOL_OPTIONAL. + optional = ( + extra.get("cli_bool") == "optional" + or field.default is None + or field.default is True + ) + action = argparse.BooleanOptionalAction if optional else "store_true" + parser.add_argument(flag, action=action, **kwargs) + return + + if is_list: + kwargs["nargs"] = "+" + kwargs["type"] = base + if choices: + kwargs["choices"] = choices + parser.add_argument(flag, **kwargs) + + +def add_config_cli_arguments(parser: argparse.ArgumentParser) -> None: + """Register every config field as a flag, grouped by concern for ``--help``. + + This is the whole tunable surface: generated from the schema, so a new field + becomes a new flag with the right type, choices, bool style, and help. + """ + general = parser.add_argument_group("general") + for name in _ROOT_FIELDS: + _add_field_argument(general, name, TrainConfig.model_fields[name]) + for gname, gmodel in _GROUPS.items(): + group = parser.add_argument_group(gname) + for name, field in gmodel.model_fields.items(): + _add_field_argument(group, name, field) + + +def build_parser() -> argparse.ArgumentParser: + """The ``train.py`` parser: ``--config`` plus the schema-generated flags.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--config", + dest="config", + default=None, + metavar="PATH", + help="YAML config file (stage-shaped: trainer keys under a top-level " + "'train:' key; a bare mapping is also accepted). CLI flags override it.", + ) + parser.add_argument( + "--dump-config", + dest="dump_config", + action="store_true", + help="Print the fully-resolved config as a valid run.yaml to stdout and " + "exit, so a working invocation can be scaffolded into a shareable config.", + ) + add_config_cli_arguments(parser) + return parser + + +def _validate_draft_init(cfg: TrainConfig, provided: set[str]) -> None: + """Enforce the draft-init contract; raise :class:`ConfigError` on conflict. + + The draft model may be defined in exactly one way: ``--from-pretrained`` (a + complete checkpoint), ``--draft-config`` (a decoder config, rest built from + the other flags), or the decoder-shaping flags (synthesize everything). + ``--from-pretrained`` takes precedence over all others; ``--draft-config`` is + incompatible with the decoder-shaping flags; and MTP-from-scratch reuses the + verifier's own decoder, so ``--draft-config`` and the shaping flags do not + apply to it. + + ``provided`` is the set of dests won by a non-default layer, so a shaping flag + passed at its default value still counts as a conflict -- preserving the + pre-refactor behavior. + """ + shaping = [flag for dest, flag in DECODER_SHAPING_FLAGS.items() if dest in provided] + with_draft_config = shaping + (["--draft-config"] if cfg.draft.draft_config else []) + if cfg.draft.from_pretrained: + _reject_conflicts( + with_draft_config, + "--from-pretrained loads a complete draft model and takes precedence " + "over all other model-definition options, so these conflict with it", + ) + return + if cfg.speculator_type == "mtp": + _reject_conflicts( + with_draft_config, + "--speculator-type mtp reuses the verifier's decoder config, so these " + "options do not apply", + ) + return + if cfg.draft.draft_config: + _reject_conflicts( + shaping, + "--draft-config defines the draft decoder, so these flags conflict with it", + ) + + +def _reject_conflicts(conflicting: list[str], reason: str) -> None: + if conflicting: + raise ConfigError(f"{reason} (remove them): {', '.join(conflicting)}") + + +def _warn_mismatched_algorithm_blocks(cfg: TrainConfig, provided: set[str]) -> None: + """Warn when an algorithm group is set under a speculator type that ignores it. + + Every algorithm group and its flags stay present in the schema for every + speculator type (back-compat requires the flags to always exist), so a + mismatched block is not a parse error -- it is silently ignored by the model + layer. Surface that so a misplaced recipe is visible, without rejecting it: + the pre-refactor lenient behavior. + + Whether a group was "set" is read from the same in-memory winning-layer record + (``provided`` = won by flag or yaml) as the draft-init check, so a mismatched + block supplied entirely via YAML warns identically to one on the CLI. + """ + for group, users in _ALGORITHM_GROUP_USERS.items(): + if cfg.speculator_type in users: + continue + ignored = sorted( + _dest_to_flag(dest) + for dest in _GROUPS[group].model_fields + if dest in provided + ) + if ignored: + warnings.warn( + f"--speculator-type {cfg.speculator_type} does not use the '{group}' " + f"algorithm group; these settings are ignored: {', '.join(ignored)}", + stacklevel=2, + ) + + +def build_from_sources( + cls: type[TrainConfig], + *, + cli: dict[str, Any], + config_path: str | None, + argv: list[str], +) -> TrainConfig: + """Layer the sources into a validated config, recording each value's origin. + + The pure core behind :meth:`TrainConfig.from_sources`: it loads the optional + stage-shaped ``config_path`` YAML, feeds the ``flag`` and ``yaml`` layers to + pydantic-settings in precedence order (so the merge yields + ``flag > yaml > default``), records per dest which layer won, and enforces the + cross-field draft-init contract. Raises :class:`ConfigError` / + ``ValidationError`` on bad input; never exits the process. + """ + yaml_nested: dict[str, Any] = {} + yaml_dests: set[str] = set() + if config_path is not None: + yaml_nested = _unwrap_stage(_load_config_file(config_path)) + yaml_dests, unknown = _partition_yaml_keys(yaml_nested) + if unknown: + warnings.warn( + f"--config '{config_path}' has unrecognised keys (ignored): " + f"{', '.join(sorted(unknown))}", + stacklevel=2, + ) + + cfg = cls(_layers={"flag": nest_flat(cli), "yaml": yaml_nested}) # type: ignore[call-arg] + cfg._provenance = { + dest: ("flag" if dest in cli else "yaml" if dest in yaml_dests else "default") + for dest in CONFIG_DESTS + } + cfg._argv = list(argv) + + provided = {dest for dest, layer in cfg._provenance.items() if layer != "default"} + _warn_mismatched_algorithm_blocks(cfg, provided) + _validate_draft_init(cfg, provided) + _validate_required(cfg) + return cfg + + +def _validate_required(cfg: TrainConfig) -> None: + """Enforce the ``cli_required`` contract once every layer has been merged. + + argparse can't do this: the value may arrive from ``--config`` rather than a + flag, and requiring it at parse time would reject a ``run.yaml`` that supplies + it. So the schema keeps an empty placeholder default and we check the resolved + value here, naming the flag the user can set (or provide in ``--config``). + """ + for dest, flag in REQUIRED_FLAGS.items(): + if cfg.provenance[dest] == "default": + raise ConfigError( + f"missing required value: set {flag} or provide it in --config" + ) + + +def _load_config_file(path: str) -> dict[str, Any]: + """Parse a ``--config`` YAML file into a top-level mapping. + + Raises :class:`ConfigError` -- not a traceback -- for the mistakes a user + actually makes: an unreadable path, unparseable YAML, or a top-level document + that is not a mapping. The :func:`resolve` boundary renders that as exit 2. + """ + try: + text = Path(path).read_text() + data = yaml.safe_load(text) + except OSError as exc: + raise ConfigError( + f"--config '{path}' could not be read: {exc.strerror or exc}" + ) from exc + except yaml.YAMLError as exc: + raise ConfigError(f"--config '{path}' is not valid YAML: {exc}") from exc + if data is None: + return {} + if not isinstance(data, dict): + raise ConfigError(f"--config '{path}' must contain a top-level mapping.") + return data + + +def _unwrap_stage(data: dict[str, Any]) -> dict[str, Any]: + """Return the trainer config, unwrapping the ``train:`` stage block if present. + + The canonical file is stage-shaped -- trainer keys nest under ``train:`` so a + future ``prepare_data:`` / ``launch_vllm:`` stage extends the file rather than + replacing it. Sibling stage keys are ignored: a file authored for + the whole pipeline still trains today using only its ``train:`` block. Loading + stays lenient -- a bare top-level mapping (no ``train:`` key, since no config + group is named ``train``) is accepted unchanged for back-compat. + """ + stage = data.get("train") + return stage if isinstance(stage, dict) else data + + +def _partition_yaml_keys(yaml_nested: dict[str, Any]) -> tuple[set[str], set[str]]: + """Split a parsed YAML mapping's leaves into ``(known, unknown)`` flat dests. + + A group block contributes its leaf keys; a root scalar contributes itself. + Each group leaf is checked against *its own* group's fields, not the global + dest set, so a real field placed under the wrong block reads as unknown (and + stays out of provenance) rather than being accepted because the name exists in + some other group. Unknown keys are warned about and ignored. + """ + known: set[str] = set() + unknown: set[str] = set() + for key, value in yaml_nested.items(): + if key in _GROUPS and isinstance(value, dict): + group_fields = _GROUPS[key].model_fields + for leaf in value: + (known if leaf in group_fields else unknown).add(leaf) + elif key in _ROOT_FIELDS: + known.add(key) + else: + unknown.add(key) + return known, unknown + + +def _format_config_error(exc: ValidationError) -> str: + """Render a pydantic ``ValidationError`` as a concise, flag-oriented message.""" + lines = [] + for err in exc.errors(): + label = ".".join(str(part) for part in err["loc"]) or "" + message = err["msg"].removeprefix("Value error, ") + lines.append(f"{label}: {message}") + return "invalid configuration:\n " + "\n ".join(lines) + + +def resolve(cls: type[TrainConfig], argv: list[str] | None) -> TrainConfig: + """Parse argv, layer the sources, validate; exit(2) cleanly on any error. + + The only function that touches ``sys.argv`` or raises ``SystemExit``. Any + configuration error -- a missing required value (which names its flag, e.g. + ``--verifier-name-or-path``), a draft-init conflict, a bad value, or a broken + config file -- surfaces through ``parser.error`` as ``SystemExit(2)`` with no + traceback. + """ + parser = build_parser() + namespace = parser.parse_args(argv) + config_path = namespace.config + cli = { + dest: value for dest, value in vars(namespace).items() if dest in CONFIG_DESTS + } + # The command recorded for provenance: the live argv on the real launch path, + # or a reconstruction when resolve() is driven with an explicit argv. + full_argv = list(sys.argv) if argv is None else [sys.argv[0], *argv] + try: + cfg = cls.from_sources(cli=cli, config_path=config_path, argv=full_argv) + except ValidationError as exc: + parser.error(_format_config_error(exc)) + except ConfigError as exc: + parser.error(str(exc)) + # --dump-config turns this working invocation into a shareable run.yaml: print + # the resolved config and exit cleanly (exit 0) before any training happens. + if namespace.dump_config: + sys.stdout.write(cfg.dump_yaml()) + raise SystemExit(0) + return cfg diff --git a/src/speculators/train/config/schema.py b/src/speculators/train/config/schema.py new file mode 100644 index 000000000..4f5fd1710 --- /dev/null +++ b/src/speculators/train/config/schema.py @@ -0,0 +1,787 @@ +"""The trainer's tunable fields: the single source of truth for every flag. + +Each tunable is one typed ``pydantic`` field; adding one is a one-field edit. +:class:`TrainConfig` is the sole public type (``resolution`` builds the CLI). +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Any, Literal, cast + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + JsonValue, + PrivateAttr, + field_validator, + model_validator, +) +from pydantic_settings import ( + BaseSettings, + InitSettingsSource, + PydanticBaseSettingsSource, + SettingsConfigDict, +) + +from speculators.data_generation.vllm_client import ( + DEFAULT_MAX_RETRIES, + DEFAULT_REQUEST_TIMEOUT, +) +from speculators.models.metrics import resolve_loss_config + +# A bool that must render as an argparse ``--x/--no-x`` (``BooleanOptionalAction``) +# even though it defaults to False. Bools defaulting to True or None get that form +# automatically; a plain False-default bool renders as ``store_true`` unless tagged +# here. The explicit ``dict[str, JsonValue]`` annotation is load-bearing: a bare +# literal infers ``dict[str, str]``, which pydantic's ``json_schema_extra`` rejects. +_CLI_BOOL_OPTIONAL: dict[str, JsonValue] = {"cli_bool": "optional"} + +# A field the CLI must require (argparse ``required=True``). The schema gives it a +# placeholder default so a config is still constructible from defaults; the +# resolution layer turns the marker back into a required flag / missing-value error. +_CLI_REQUIRED: dict[str, JsonValue] = {"cli_required": True} + +# A str field whose valid values are a dynamic registry (argparse ``choices=``) +# resolved at CLI-generation time. The marker names the registry indirectly -- a +# string key the resolution layer maps to the live registry -- so the schema keeps no +# runtime dependency on the backends and stays purely declarative. +_CLI_CHOICES: dict[str, JsonValue] = {"cli_choices": "hidden_states_backends"} + + +class _Group(BaseModel): + """Base for a config group: mutable (post-resolution) and lenient on extras.""" + + model_config = ConfigDict(extra="ignore", validate_assignment=False) + + +class VerifierArgs(_Group): + verifier_name_or_path: str = Field( + default="", + description="Path or HF id of the verifier/target model.", + json_schema_extra=_CLI_REQUIRED, + ) + trust_remote_code: bool = Field( + default=False, + description="Allow executing code from HF Hub when loading the verifier's " + "tokenizer.", + ) + + +class DraftArgs(_Group): + """The draft model definition: init source, decoder shape, shared + normalization hyperparameters, and vocabulary mapping.""" + + from_pretrained: str = Field( + default="", + description="Path or HF id of a pretrained draft. May also point to a local " + "directory containing only a config.json, in which case a fresh draft is " + "initialized from that full speculator config. Takes precedence over and is " + "mutually exclusive with --draft-config and the decoder-shaping flags " + "(--num-layers, --draft-arch, --draft-hidden-act, --sliding-window, " + "--full-attention-indices).", + ) + draft_config: str = Field( + default="", + description="HF id, directory, or JSON path of a decoder config (LlamaConfig " + "for eagle3/peagle, Qwen3Config for dflash) to use as the draft " + "transformer_layer_config; the rest of the speculator is built from the other " + "CLI args. Mutually exclusive with --from-pretrained and the decoder-shaping " + "flags (--num-layers, --draft-arch, --draft-hidden-act, --sliding-window, " + "--full-attention-indices).", + ) + num_layers: int = Field( + default=1, description="Number of draft decoder layers to synthesize." + ) + draft_arch: Literal["llama", "qwen3"] | None = Field( + default=None, + description="Architecture for draft decoder layers " + "(default: 'llama' for eagle3, 'qwen3' otherwise).", + ) + draft_hidden_act: str = Field( + default="silu", + description="Activation function for draft decoder layers. Defaults to 'silu'. " + "Qwen3 layers of dflash expect 'silu' for vLLM deployment. Leave as None to " + "fall back to the verifier's activation function.", + ) + draft_mrope_full_head_hack: bool = Field( + default=True, + description="For MRoPE configs with partial_rotary_factor < 1, rescale " + "mrope_section and set partial_rotary_factor=1.0 so HF training and vLLM " + "inference use equivalent full-head rotary semantics.", + ) + target_layer_ids: list[int] | None = Field( + default=None, + description="(Optional) space-separated list of integer layer ids for the " + "auxiliary hidden states. Defaults to [2, n//2, n-3]. If custom values were " + "used to launch vllm, pass the same ids here, excluding the final layer " + "launch_vllm.py appends.", + ) + token_freq_path: str | None = Field( + default=None, + description="Path to token frequency distribution file (.pt). Used with " + "--draft-vocab-size to build vocab mappings at training time. Falls back to " + "'/token_freq.pt'. If neither exists nor --draft-vocab-size is set, " + "vocab mapping is skipped and the full verifier vocab is used.", + ) + draft_vocab_size: int | None = Field( + default=None, + description="Vocabulary size for the draft model. Must be provided together " + "with a token frequency file to generate vocab mappings; otherwise a no-op.", + ) + d2t_path: str | None = Field( + default=None, description="Draft-to-target vocab mapping file (.npy)." + ) + t2d_path: str | None = Field( + default=None, description="Target-to-draft vocab mapping file (.npy)." + ) + mask_token_id: int | None = Field( + default=None, + description="Token id used to mask positions during training. Resolved from " + "the verifier tokenizer when unset.", + ) + norm_before_residual: bool = Field( + default=True, + description="Toggle normalization before residual connections (default: True).", + ) + embed_requires_grad: bool = Field( + default=False, + description="Whether to train embedding layer weights (default: False).", + json_schema_extra=_CLI_BOOL_OPTIONAL, + ) + norm_before_fc: bool | None = Field( + default=None, + description="Apply a single RMSNorm to the concatenated auxiliary hidden " + "states before the FC projection (gpt-oss style). See --fc-norm for the " + "per-layer alternative from the Eagle 3.1 paper. " + "(default: True for eagle3, False otherwise).", + ) + fc_norm: bool = Field( + default=False, + description="Apply per-layer RMSNorm to each auxiliary hidden state before " + "concatenation and FC projection (Eagle 3.1 paper approach).", + ) + norm_output: bool | None = Field( + default=None, + description="Feed post-norm hidden states back across TTT steps to stabilize " + "magnitude drift across speculation depths " + "(default: True for eagle3, False otherwise).", + ) + sliding_window: int = Field( + default=2048, + description="Sliding window size for sliding window attention layers " + "(default: 2048). All draft layers use sliding window by default (except mtp).", + ) + full_attention_indices: list[int] = Field( + default_factory=list, + description="(Optional) space-separated draft layer indices that should use " + "full attention instead of sliding window. All draft layers use sliding window " + "by default (except mtp). (e.g. '--full-attention-indices 0 2' makes layers 0 " + "and 2 use full attention; the rest use sliding window).", + ) + sliding_window_non_causal: bool = Field( + default=False, + description="Use non-causal (bidirectional) masking within draft blocks for " + "sliding window attention layers. Full attention layers are always " + "bidirectional. Note: vLLM currently doesn't support these models.", + ) + draft_attn_impl: Literal["simple_flex_attention", "sdpa", "eager"] = Field( + default="simple_flex_attention", + description="Attention implementation for draft layers. Use 'sdpa' or 'eager' " + "for hardware that doesn't support flex attention. Not supported for MTP.", + ) + + +class DataArgs(_Group): + data_path: str = Field( + default="./output", + description="Root data directory with the preprocessed dataset, vocab " + "mappings (d2t.npy, t2d.npy), token frequencies (token_freq.pt), and hidden " + "states.", + ) + hidden_states_backend: str = Field( + default="file", + description="Hidden states transfer backend. Each backend may add its own CLI " + "arguments. Default: 'file'.", + json_schema_extra=_CLI_CHOICES, + ) + hidden_states_path: str | None = Field( + default=None, + description="Path where cached hidden states are stored " + "(default: /hidden_states). Contributed by the 'file' backend.", + ) + total_seq_len: int = Field( + default=8192, description="Maximum training sequence length, in tokens." + ) + train_data_ratio: float = Field( + default=0.9, + description="Fraction of the dataset used for training; the remainder is held " + "out for validation.", + ) + noise_std: float = Field( + default=0.05, description="Standard deviation for noise augmentation." + ) + legacy_data: bool = Field( + default=False, + description="DEPRECATED. Use the old data format which stores hidden states " + "alongside token_ids and assistant_masks in data_i.pt files. Will be removed " + "soon.", + ) + hidden_states_dtype: str = Field( + default="bfloat16", + description="Data type for dataloader hidden states and autocast compute. " + "Model master weights are always kept in fp32. Options: float32, bfloat16 " + "(recommended). float16 is not supported (requires gradient scaling).", + ) + num_workers: int = Field(default=12, description="Number of dataloader workers.") + prefetch_factor: int = Field(default=4, description="Dataloader prefetch factor.") + max_anchors: int = Field( + default=3072, + description="Maximum anchor positions for DFlash, DSpark, and P-EAGLE training " + "(default: 3072).", + ) + + @field_validator("hidden_states_dtype") + @classmethod + def _dtype_is_torch_attr(cls, v: str) -> str: + # Local import keeps this schema module purely declarative and torch-free + # at import time; torch is only needed to validate the dtype string. + import torch # noqa: PLC0415 + + # hasattr alone is too weak: torch.nn / torch.cuda / torch.Tensor all + # exist but are not dtypes and would only fail later as an opaque + # autocast error. Require the resolved attribute to be a torch.dtype. + resolved = getattr(torch, v, None) + if not isinstance(resolved, torch.dtype): + raise ValueError( + "hidden_states_dtype must name a torch dtype, e.g. `bfloat16`." + ) + if resolved is torch.float16: + raise ValueError( + "hidden_states_dtype float16 is not supported (requires gradient " + "scaling); use bfloat16." + ) + return v + + +class GenerationArgs(_Group): + """On-demand hidden-state generation against a vLLM endpoint.""" + + vllm_endpoint: str = Field( + default="http://localhost:8000/v1", + description="vLLM endpoint used to generate hidden states on demand. Only " + "needed if --on-missing=generate and samples are missing. The vLLM instance " + "must cache hidden states to a location reachable from the training instance.", + ) + on_missing: Literal["generate", "skip", "warn", "raise"] = Field( + default="generate", + description="Dataloader behaviour when a sample has no cached hidden states. " + "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( + 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.", + ) + request_timeout: float = Field( + default=DEFAULT_REQUEST_TIMEOUT, + description="Timeout in seconds for each individual vLLM request. Only applies " + "if --on-missing=generate.", + ) + max_retries: int = Field( + default=DEFAULT_MAX_RETRIES, + description="Maximum retry attempts per vLLM request on failure. Only applies " + "if --on-missing=generate.", + ) + + +class LossArgs(_Group): + loss_fn: str = Field( + default="kl_div", + description="Loss function specification. A name (kl_div, rkl, jsd, ce, tv, " + 'nla, lk_hybrid) or a JSON dict for a weighted combination, e.g. \'{"ce": 0.1, ' + '"tv": 0.9}\'.', + ) + ttt_steps: int = Field( + default=3, + description="Number of test-time-training (TTT) steps the draft is unrolled " + "over during training.", + ) + ttt_step_loss_decay: float = Field( + default=1.0, + description="Multiplicative loss weight decay applied per TTT step " + "(1.0 = weight every step equally).", + ) + + @field_validator("loss_fn") + @classmethod + def _loss_parseable(cls, v: str) -> str: + resolve_loss_config(v) + return v + + +class OptimizerArgs(_Group): + optimizer: Literal["adamw", "muon"] = Field( + default="muon", + description="Optimizer to use. 'muon' applies Muon to 2D weight matrices and " + "AdamW to the remaining params (norms, biases, embeddings, lm_head).", + ) + lr: float = Field(default=1e-4, description="Learning rate (AdamW / base group).") + weight_decay: float = Field( + default=0.01, + description="Weight decay for the AdamW optimizer (and the AdamW group in muon " + "mode).", + ) + muon_lr: float | None = Field( + default=None, + description="LR for the Muon (2D weights) group. Only used with --optimizer " + "muon. Defaults to 10*lr (and --lr defaults to 1e-4).", + ) + muon_momentum: float = Field( + default=0.95, description="Momentum for the Muon group." + ) + muon_weight_decay: float = Field( + default=0.1, description="Weight decay for the Muon group." + ) + muon_ns_steps: int = Field( + default=5, description="Newton-Schulz iteration steps for Muon." + ) + muon_adjust_lr_fn: Literal["original", "match_rms_adamw"] = Field( + default="match_rms_adamw", + description="Muon LR adjustment. 'match_rms_adamw' matches AdamW's update RMS.", + ) + + +class SchedulerArgs(_Group): + scheduler_type: Literal["linear", "cosine", "none"] = Field( + default="linear", description="LR scheduler type." + ) + scheduler_warmup_steps: int | None = Field( + default=None, description="Warmup steps (default: scheduler-dependent)." + ) + scheduler_warmup_ratio: float | None = Field( + default=None, + description="Warmup as a fraction of total scheduler steps, in [0, 1]. Ignored " + "(with a warning) when --scheduler-warmup-steps is also set.", + ) + scheduler_total_steps: int | None = Field( + default=None, description="Total scheduler steps (default: inferred)." + ) + scheduler_num_cosine_cycles: float = Field( + default=0.5, description="Number of cosine cycles for the cosine scheduler." + ) + + +class TrainerArgs(_Group): + epochs: int = Field(default=20, description="Number of training epochs.") + checkpoint_freq: float = Field( + default=1.0, + description="Save a checkpoint every N epochs. Values < 1 enable sub-epoch " + "checkpointing (e.g. 0.5 = every half epoch).", + ) + save_best: bool = Field( + default=False, + description="Also point a checkpoint at the lowest validation loss.", + ) + no_resume_from_checkpoint: bool = Field( + default=False, description="Do not resume training from an existing checkpoint." + ) + log_freq: int = Field( + default=1, description="Log training metrics every N steps (default: 1)." + ) + save_path: str = Field( + default="./output/checkpoints", + description="Directory to write checkpoints and the resolved run.yaml.", + ) + fsdp_shard: bool = Field( + default=False, + description="Shard model parameters across GPUs with FSDP. By default " + "parameters are fully replicated (DDP-like). Enable when the model does not " + "fit in a single GPU's memory.", + ) + sp_size: int = Field( + default=1, + ge=1, + description="Ulysses sequence parallelism degree. Shards the sequence " + "dimension across this many GPUs within each node, reducing per-GPU " + "memory for long-context training. Must divide --total-seq-len and " + "world_size. Currently supported for eagle3, dflash, dspark, and peagle.", + ) + max_steps: int | None = Field( + default=None, + ge=1, + description="Stop training after this many optimizer steps (counted across " + "epochs). Useful for quick smoke runs. Default: run all epochs to completion.", + ) + + @field_validator("checkpoint_freq") + @classmethod + def _validate_checkpoint_freq(cls, v: float) -> float: + if v <= 0: + raise ValueError("--checkpoint-freq must be > 0") + if v > 1 and not float(v).is_integer(): + raise ValueError( + f"--checkpoint-freq={v} is not an integer. Values > 1 are treated as " + "epoch counts and must be whole numbers." + ) + return v + + +class LoggingArgs(_Group): + logger: str = Field( + default="", + description="One of 'trackio', 'wandb', 'tensorboard', 'mlflow' or a comma " + "separated list.", + ) + log_dir: str = Field(default="./logs", description="Directory for training logs.") + run_name: str | None = Field( + default=None, + description="Name for this run (used by metric loggers). Auto-generated when " + "unset.", + ) + + +class DFlashArgs(_Group): + """DFlash-family backbone knobs (also used by DSpark, which is-a DFlash).""" + + block_size: int = Field( + default=8, description="Block size for DFlash model (default: 8)." + ) + sample_from_anchor: bool | None = Field( + default=None, + description="Sample from the anchor position (all positions predict). " + "Default: False for dflash, True for dspark.", + ) + dflash_decay_gamma: float = Field( + default=4.0, description="Decay gamma for DFlash/DSpark loss weighting." + ) + per_position_loss_weight: Literal["fixed-exp-decay", "dpace"] = Field( + default="fixed-exp-decay", + description="Per-position loss weight option for D-PACE support " + "(default: fixed-exp-decay).", + ) + dpace_alpha: float = Field( + default=0.5, + description="Smoothing constant for the D-PACE loss (default: 0.5). Must be in " + "(0, 1] when --per-position-loss-weight=dpace.", + ) + + +class DSparkArgs(_Group): + """DSpark-exclusive heads (sequential Markov head + confidence head).""" + + markov_rank: int = Field( + default=256, + description="DSpark: low-rank dim of the Markov logit-bias head " + "(0 disables it).", + ) + markov_head_type: Literal["vanilla", "gated", "rnn"] = Field( + default="vanilla", description="DSpark: sequential head variant." + ) + enable_confidence_head: bool = Field( + default=True, + description="DSpark: attach the per-position acceptance confidence head.", + ) + confidence_head_with_markov: bool = Field( + default=True, + description="DSpark: feed the Markov previous-token embedding into the " + "confidence head alongside the backbone hidden state.", + ) + confidence_head_alpha: float = Field( + default=1.0, description="DSpark: weight of the confidence-head BCE term." + ) + + +class PEagleArgs(_Group): + num_depths: int = Field( + default=8, + description="Number of parallel prediction depths for P-EAGLE (default: 8).", + ) + down_sample_ratio: float = Field( + default=0.7, description="Geometric decay ratio for COD sampling in P-EAGLE." + ) + down_sample_ratio_min: float = Field( + default=0.2, description="Minimum retention ratio for COD sampling in P-EAGLE." + ) + + +class MTPArgs(_Group): + num_speculative_steps: int = Field( + default=3, description="Number of MTP prediction steps (default: 3)." + ) + step_weight_beta: float = Field( + default=0.6, + description="Exponential decay factor for MTP step weights; higher weights " + "earlier prediction steps more. Only used with MTP.", + ) + + +# Group attribute name -> group model. Order defines both the flatten() key order +# (after the root scalars) and a dumped run.yaml's group layout. +_GROUPS: dict[str, type[_Group]] = { + "verifier": VerifierArgs, + "draft": DraftArgs, + "data": DataArgs, + "generation": GenerationArgs, + "loss": LossArgs, + "optimizer": OptimizerArgs, + "scheduler": SchedulerArgs, + "trainer": TrainerArgs, + "logging": LoggingArgs, + "dflash": DFlashArgs, + "dspark": DSparkArgs, + "peagle": PEagleArgs, + "mtp": MTPArgs, +} + + +class TrainConfig(BaseSettings): + """Top-level trainer configuration -- the package's one public type. + + Composed from the typed groups (:data:`_GROUPS`) plus the root-level scalars + declared here. Constructible from defaults with no arguments: the verifier + path is not required at the schema level (a config may supply it later via + YAML), and the resolution layer restores the required-flag contract. + + Subclasses ``BaseSettings`` to reuse pydantic-settings' source ordering as + the precedence engine: the resolution layer passes the ``flag`` and ``yaml`` + layers as a private ``_layers`` kwarg and lets pydantic deep-merge them over + the field defaults, giving ``flag > yaml > default`` without a hand-rolled + merge (see :meth:`settings_customise_sources`). Plain keyword construction + (:meth:`from_flat`) still works and bypasses that machinery. + """ + + model_config = SettingsConfigDict( + extra="ignore", + validate_assignment=False, + nested_model_default_partial_update=True, + ) + + @classmethod + def settings_customise_sources( + cls, + settings_cls: type[BaseSettings], + init_settings: PydanticBaseSettingsSource, + env_settings: PydanticBaseSettingsSource, # noqa: ARG003 + dotenv_settings: PydanticBaseSettingsSource, # noqa: ARG003 + file_secret_settings: PydanticBaseSettingsSource, # noqa: ARG003 + ) -> tuple[PydanticBaseSettingsSource, ...]: + """Turn the precedence rule ``flag > yaml > default`` into source ordering. + + pydantic-settings calls this hook by keyword, so all five parameters are + part of the contract even though env/dotenv/secret sources are excluded: + the trainer is configured only via CLI + YAML. The resolution layer + passes the two explicit layers as a private ``_layers`` kwarg (popped here + so it never becomes model data); each becomes its own + :class:`InitSettingsSource` in precedence order, above the field defaults. + A plain keyword construction carries no ``_layers`` and behaves like a + normal ``BaseModel``. + """ + init_kwargs = dict(cast("InitSettingsSource", init_settings).init_kwargs) + layers = init_kwargs.pop("_layers", None) + if layers is None: + return (init_settings,) + return tuple( + InitSettingsSource(settings_cls, init_kwargs=layers.get(name, {})) + for name in ("flag", "yaml") + ) + + # Resolution bookkeeping: the per-dest winning layer + # (``flag | yaml | default``) and the argv the run was resolved from. In + # memory only -- consumed by the draft-init conflict check and the + # reproducibility manifest, never persisted and not part of the config shape. + # Exposed read-only via the ``provenance`` / ``argv`` properties below. + _provenance: dict[str, str] = PrivateAttr(default_factory=dict) + _argv: list[str] = PrivateAttr(default_factory=list) + + @property + def provenance(self) -> Mapping[str, str]: + """The per-field winning layer (``'flag'`` | ``'yaml'`` | ``'default'``).""" + return MappingProxyType(self._provenance) + + @property + def argv(self) -> Sequence[str]: + """The argv the run was resolved from.""" + return tuple(self._argv) + + speculator_type: str = Field( + default="eagle3", + description="Type of speculator model to train " + "(eagle3, dflash, dspark, peagle, mtp).", + ) + dry_run: bool = Field( + default=False, + description="Build the speculator, initialize weights, save a checkpoint to " + "--save-path, then exit before training. Useful to validate the config and " + "weights (e.g. in vLLM) before a full run. Can be combined with --draft-config " + "or --from-pretrained.", + ) + seed: int = Field(default=42, description="Random seed for reproducibility.") + deterministic_cuda: bool = Field( + default=False, + description="Set cuda to deterministic mode. This may impact performance.", + ) + + verifier: VerifierArgs = Field(default_factory=VerifierArgs) + draft: DraftArgs = Field(default_factory=DraftArgs) + data: DataArgs = Field(default_factory=DataArgs) + generation: GenerationArgs = Field(default_factory=GenerationArgs) + loss: LossArgs = Field(default_factory=LossArgs) + optimizer: OptimizerArgs = Field(default_factory=OptimizerArgs) + scheduler: SchedulerArgs = Field(default_factory=SchedulerArgs) + trainer: TrainerArgs = Field(default_factory=TrainerArgs) + logging: LoggingArgs = Field(default_factory=LoggingArgs) + dflash: DFlashArgs = Field(default_factory=DFlashArgs) + dspark: DSparkArgs = Field(default_factory=DSparkArgs) + peagle: PEagleArgs = Field(default_factory=PEagleArgs) + mtp: MTPArgs = Field(default_factory=MTPArgs) + + @model_validator(mode="after") + def _resolve_derived_defaults(self) -> "TrainConfig": + """Fill defaults that derive from other fields, mirroring the tail of the + pre-refactor ``parse_args``: unset ``draft_arch`` -> ``llama`` for eagle3 else + ``qwen3``; unset ``norm_before_fc`` / ``norm_output`` -> ``True`` for eagle3 + else ``False``; unset ``muon_lr`` -> ``10 * lr``. + + Idempotent: a concrete value (as produced by :meth:`flatten`) is left + untouched, so :meth:`from_flat` round-trips. + """ + is_eagle3 = self.speculator_type == "eagle3" + if self.draft.draft_arch is None: + self.draft.draft_arch = "llama" if is_eagle3 else "qwen3" + if self.draft.norm_before_fc is None: + self.draft.norm_before_fc = is_eagle3 + if self.draft.norm_output is None: + self.draft.norm_output = is_eagle3 + if self.optimizer.muon_lr is None: + self.optimizer.muon_lr = 10 * self.optimizer.lr + return self + + @model_validator(mode="after") + def _validate_dpace(self) -> "TrainConfig": + """D-PACE per-position loss weighting requires CE loss and a smoothing constant + in ``(0, 1]``. Only enforced when actually selected, so it never fires on a + default (``fixed-exp-decay``) run.""" + if self.dflash.per_position_loss_weight == "dpace": + if self.loss.loss_fn != "ce": + raise ValueError( + "--per-position-loss-weight=dpace requires --loss-fn=ce" + ) + if not 0.0 < self.dflash.dpace_alpha <= 1.0: + raise ValueError( + f"--dpace-alpha must be in (0, 1], got {self.dflash.dpace_alpha}" + ) + return self + + def flatten(self) -> dict[str, Any]: + """The flat ``vars(args)``-shaped dict the ``SpeculatorModel`` classes + consume via ``**kwargs``. + + Keys are emitted in schema-declaration order (root scalars, then each + group's fields); consumers bind by name, so only the key set and values + matter, but the order is deterministic so the ``run.yaml`` dump is stable. + """ + values = {name: getattr(self, name) for name in _ROOT_FIELDS} + for gname in _GROUPS: + group = getattr(self, gname) + for fname in type(group).model_fields: + values[fname] = getattr(group, fname) + return values + + @classmethod + def from_flat(cls, flat: dict[str, Any]) -> "TrainConfig": + """Inverse of :meth:`flatten`: recover the typed config from a flat + working-dict. Non-config keys are dropped; the value validators are + idempotent on already-resolved values.""" + known = {dest: value for dest, value in flat.items() if dest in CONFIG_DESTS} + return cls(**nest_flat(known)) + + @classmethod + def from_sources( + cls, + *, + cli: dict[str, Any], + config_path: str | None = None, + argv: list[str], + ) -> "TrainConfig": + """The pure, argv-free core: layer the sources into a validated config and + record each value's origin. Raises on bad input; never exits. + + ``cli`` is the flat ``{dest: value}`` map of flags the user passed; + ``config_path`` names an optional stage-shaped YAML file layered beneath + the flags; ``argv`` is the command recorded for the reproducibility + manifest. This is the primary test seam -- it can be exercised without + ``sys.argv``. + """ + from speculators.train.config import resolution # noqa: PLC0415 + + return resolution.build_from_sources( + cls, cli=cli, config_path=config_path, argv=argv + ) + + @classmethod + def resolve(cls, argv: list[str] | None = None) -> "TrainConfig": + """The impure CLI boundary: parse argv, layer, validate. + + Turns any configuration error into a clean ``SystemExit(2)``. This is what + ``scripts/train.py`` calls. + """ + from speculators.train.config import resolution # noqa: PLC0415 + + return resolution.resolve(cls, argv) + + def dump_yaml(self) -> str: + """Serialize this config to the stage-shaped, round-trippable YAML. + + The ``train:``-wrapped form, clean of provenance annotations, so it + re-loads identically via ``--config`` (see ``artifacts`` for the writer). + """ + from speculators.train.config import artifacts # noqa: PLC0415 + + return artifacts.dump_yaml(self) + + def save(self, save_dir: str) -> None: + """Write the reproducibility artifacts (``run.yaml`` + ``train_command.txt``) + next to the checkpoints. Called at rank 0.""" + from speculators.train.config import artifacts # noqa: PLC0415 + + artifacts.save(self, save_dir) + + +# Root-level scalar fields (not inside any group), in declaration order. +_ROOT_FIELDS: tuple[str, ...] = tuple( + name for name in TrainConfig.model_fields if name not in _GROUPS +) + + +def _build_dest_to_group() -> dict[str, str | None]: + """Map each flat dest to its owning group (``None`` for a root scalar).""" + mapping: dict[str, str | None] = dict.fromkeys(_ROOT_FIELDS) + for gname, gmodel in _GROUPS.items(): + for fname in gmodel.model_fields: + if fname in mapping: + raise RuntimeError(f"duplicate config field '{fname}' across groups") + mapping[fname] = gname + return mapping + + +# Flat dest -> owning group attribute (or None for a root scalar). +_DEST_TO_GROUP: dict[str, str | None] = _build_dest_to_group() + +# Every flat dest owned by the config. +CONFIG_DESTS: frozenset[str] = frozenset(_DEST_TO_GROUP) + + +def nest_flat(flat: dict[str, Any]) -> dict[str, Any]: + """Turn a flat ``{dest: value}`` dict into the nested group structure that + :class:`TrainConfig` is constructed from.""" + nested: dict[str, Any] = {} + for dest, value in flat.items(): + group = _DEST_TO_GROUP.get(dest) + if group is None: + nested[dest] = value + else: + nested.setdefault(group, {})[dest] = value + return nested diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index 4ccd62587..a692bfaaa 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -136,12 +136,18 @@ def build_client_item(dataset_item: dict) -> ClientItem: Text-only EAGLE-3 models (e.g. Llama) use a plain tokenizer, so ``messages`` is never created and this guard is a no-op. """ - out_dict: dict = {"input_ids": dataset_item["input_ids"].tolist()} + ids = dataset_item["input_ids"].tolist() if "messages" in dataset_item and _has_multimodal_content(dataset_item["messages"]): - out_dict["messages"] = dataset_item["messages"] + return cast( + "ClientItem", {"input_ids": ids, "messages": dataset_item["messages"]} + ) - return cast("ClientItem", out_dict) + # Text-only / Completions API path: drop the last token so that + # len(prompt) + max_tokens(=1) <= max_model_len even when the sample + # fills the full context window. The hidden state at the final + # position is never used for loss (no next-token target). + return cast("ClientItem", {"input_ids": ids[:-1]}) class BaseDataset(Dataset): @@ -208,7 +214,8 @@ def __init__( vllm_endpoint: str = "http://localhost:8000/v1", on_missing: Literal["generate", "skip", "warn", "raise"] = "generate", on_generate: Literal["cache", "delete"] = "delete", - split_ratio: float = 1.0, + train_ratio: float = 1.0, + split: Literal["train", "val"] = "train", transform: TransformTensors | None = None, hidden_states_dtype=torch.bfloat16, model: str | None = None, @@ -216,19 +223,24 @@ def __init__( max_retries: int = DEFAULT_MAX_RETRIES, ): self.data = load_from_disk(datapath) - self.start_file_idx = 0 - if split_ratio == 1.0: - pass - elif 1.0 > split_ratio > 0: - self.start_file_idx = 0 - split_idx = int(len(self.data) * split_ratio) - self.data = self.data.select(range(split_idx)) - elif -1.0 < split_ratio < 0: - split_idx = int(len(self.data) * (1.0 + split_ratio)) - self.start_file_idx = split_idx - self.data = self.data.select(range(split_idx, len(self.data))) - else: - raise ValueError("split_ratio must be in range (-1.0, 1.0] excluding 0.0.") + if not 0.0 < train_ratio <= 1.0: + raise ValueError(f"train_ratio must be in (0.0, 1.0], got {train_ratio}") + if split == "val" and train_ratio == 1.0: + raise ValueError("train_ratio=1.0 leaves no validation split") + + # Both splits derive their boundary from this one expression, + # so they are exactly complementary. + split_idx = int(len(self.data) * train_ratio) + start, stop = ( + (0, split_idx) if split == "train" else (split_idx, len(self.data)) + ) + if start >= stop: + raise ValueError( + f"{split} split is empty (dataset has {len(self.data)} rows, " + f"train_ratio={train_ratio} gives split_idx={split_idx})" + ) + self.start_file_idx = start + self.data = self.data.select(range(start, stop)) self.transfer = transfer or FileTransfer(Path(datapath) / "hidden_states") self.vllm_endpoint = vllm_endpoint @@ -287,7 +299,7 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None: if loaded_hs is None: raise ValueError(f"Failed to load hidden states for handle {handle}") - check_hidden_states(loaded_hs, dataset_item["input_ids"].tolist()) + check_hidden_states(loaded_hs, client_item["input_ids"]) file_idx = self._map_to_file_idx(index) match self.on_generate: @@ -335,23 +347,28 @@ def _get_raw_data(self, index): # "token_ids": [seq_len] # } - if not torch.equal(loaded_hs["token_ids"], self.data[index]["input_ids"]): + cached_ids = loaded_hs["token_ids"] + full_ids = self.data[index]["input_ids"] + # Cached hidden states may cover N-1 tokens (last token trimmed + # before sending to vLLM); accept any valid prefix of the full ids. + if not torch.equal(cached_ids, full_ids[: len(cached_ids)]): warnings.warn( - f"Loaded token ids {loaded_hs['token_ids']} for index {index} don't" - f"match input ids {self.data[index]['input_ids']}", + f"Loaded token ids {cached_ids} for index {index} don't" + f"match input ids {full_ids}", stacklevel=1, ) return None + hs_len = cached_ids.shape[0] return { "hidden_states": loaded_hs["hidden_states"][:, :-1].flatten( 1 - ), # [seq_len, 3 * hidden_size] - "input_ids": loaded_hs["token_ids"], # [seq_len] + ), # [hs_len, 3 * hidden_size] + "input_ids": cached_ids, # [hs_len] "verifier_last_hidden_states": loaded_hs["hidden_states"][ :, -1 - ], # [seq_len, hidden_size] - "loss_mask": self.data[index]["loss_mask"], # [seq_len] + ], # [hs_len, hidden_size] + "loss_mask": self.data[index]["loss_mask"][:hs_len], # [hs_len] } diff --git a/src/speculators/train/dataloader.py b/src/speculators/train/dataloader.py index 754e932db..c7a006f97 100644 --- a/src/speculators/train/dataloader.py +++ b/src/speculators/train/dataloader.py @@ -122,7 +122,8 @@ def create_train_val_loaders( on_missing=on_missing, on_generate=on_generate, transform=noise_transform, - split_ratio=train_data_ratio, + train_ratio=train_data_ratio, + split="train", model=verifier_name_or_path, hidden_states_dtype=hidden_states_dtype, request_timeout=request_timeout, @@ -135,7 +136,8 @@ def create_train_val_loaders( vllm_endpoint=vllm_endpoint, on_missing=on_missing, on_generate=on_generate, - split_ratio=train_data_ratio - 1.0, + train_ratio=train_data_ratio, + split="val", model=verifier_name_or_path, hidden_states_dtype=hidden_states_dtype, request_timeout=request_timeout, diff --git a/src/speculators/train/distributed.py b/src/speculators/train/distributed.py index 7291befb1..2e0a5893f 100644 --- a/src/speculators/train/distributed.py +++ b/src/speculators/train/distributed.py @@ -13,6 +13,7 @@ import torch import torch.distributed as dist from torch.distributed import ProcessGroup +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard logger = logging.getLogger("speculators") @@ -33,6 +34,7 @@ _sp_group: ProcessGroup | None = None _dp_group: ProcessGroup | None = None +_device_mesh: DeviceMesh | None = None # --------------------------------------------------------------------------- @@ -80,6 +82,10 @@ def get_dp_rank() -> int: return _dp_rank +def get_device_mesh() -> DeviceMesh | None: + return _device_mesh + + # --------------------------------------------------------------------------- # Initialization # --------------------------------------------------------------------------- @@ -92,6 +98,7 @@ def _init_sp_process_groups(rank: int, world_size: int, sp_size: int) -> None: DP groups use strided ranks (e.g. sp_size=2, world_size=4: {0,2}, {1,3}). """ global _sp_group, _dp_group, _sp_size, _sp_rank, _dp_size, _dp_rank # noqa: PLW0603 + global _device_mesh # noqa: PLW0603 if sp_size <= 0: raise ValueError(f"sp_size must be positive, got {sp_size}") @@ -103,6 +110,7 @@ def _init_sp_process_groups(rank: int, world_size: int, sp_size: int) -> None: dp_size = world_size // sp_size + # get the SP group for this rank sp_group = None for i in range(dp_size): sp_ranks = list(range(i * sp_size, (i + 1) * sp_size)) @@ -110,6 +118,7 @@ def _init_sp_process_groups(rank: int, world_size: int, sp_size: int) -> None: if rank in sp_ranks: sp_group = pg + # get the DP group for this rank dp_group = None for i in range(sp_size): dp_ranks = list(range(i, world_size, sp_size)) @@ -127,6 +136,13 @@ def _init_sp_process_groups(rank: int, world_size: int, sp_size: int) -> None: _dp_size = dp_size _dp_rank = rank // sp_size + if sp_size > 1: + _device_mesh = init_device_mesh( + "cuda", + (dp_size, sp_size), + mesh_dim_names=("dp", "sp"), + ) + def maybe_setup_distributed(sp_size: int = 1) -> None: """Set up distributed training if launched with ``torchrun``. @@ -171,7 +187,7 @@ def maybe_destroy_distributed() -> None: """Destroy the distributed process group if using distributed training.""" global _is_distributed, _local_rank, _rank, _world_size # noqa: PLW0603 global _sp_size, _sp_rank, _dp_size, _dp_rank # noqa: PLW0603 - global _sp_group, _dp_group # noqa: PLW0603 + global _sp_group, _dp_group, _device_mesh # noqa: PLW0603 if not _is_distributed: return @@ -192,6 +208,7 @@ def maybe_destroy_distributed() -> None: _dp_rank = 0 _sp_group = None _dp_group = None + _device_mesh = None def apply_fully_sharded( @@ -202,15 +219,54 @@ def apply_fully_sharded( Assumes the model has a `layers` attribute containing the decoder layers. Model should be validated with SpeculatorModel.verify_training_compatible() before calling this function. + + When ``sp_size > 1``, FSDP shards only across the DP sub-mesh so that + SP ranks hold identical parameters. """ mp_policy = MixedPrecisionPolicy( param_dtype=param_dtype, reduce_dtype=torch.float32, ) + mesh = _device_mesh + fsdp_kwargs: dict = {"mp_policy": mp_policy} + if mesh is not None and _sp_size > 1: + fsdp_kwargs["mesh"] = mesh["dp"] + for layer in model.layers: # type: ignore[union-attr] - fully_shard(layer, mp_policy=mp_policy) + fully_shard(layer, **fsdp_kwargs) - fully_shard(model, mp_policy=mp_policy) + fully_shard(model, **fsdp_kwargs) return model + + +def register_sp_gradient_hooks(model: torch.nn.Module) -> list: + """Register hooks to all-reduce gradients across the SP group. + + Each SP rank computes gradients from its own sequence chunk. + These hooks ensure the partial gradients are summed so that + the optimizer sees the correct total gradient. + + Returns the hook handles for cleanup. + """ + if _sp_size <= 1 or _sp_group is None: + return [] + + from torch.distributed.tensor import DTensor # noqa: PLC0415 + + sp_group = _sp_group + + def _sp_all_reduce(p: torch.nn.Parameter, _group: ProcessGroup = sp_group) -> None: + grad = p.grad + if isinstance(grad, DTensor): + dist.all_reduce(grad._local_tensor, group=_group) # noqa: SLF001 + else: + dist.all_reduce(grad, group=_group) + + hooks = [] + for param in model.parameters(): + if param.requires_grad: + hook = param.register_post_accumulate_grad_hook(_sp_all_reduce) + hooks.append(hook) + return hooks diff --git a/src/speculators/train/distributed_batch_sampler.py b/src/speculators/train/distributed_batch_sampler.py index e1b089cbc..d785744ab 100644 --- a/src/speculators/train/distributed_batch_sampler.py +++ b/src/speculators/train/distributed_batch_sampler.py @@ -41,11 +41,16 @@ class _Bin(NamedTuple): """Helper named tuple for `lpt_packed_batch`""" fill: int # sum of items in _Bin - rank: int # global rank _Bin is associated with + slot: int # heap slot id (0..num_replicas-1) def _lpt_packed_batch( - lengths: np.ndarray, max_len: int, num_replicas: int, start_index: int, rank: int + lengths: np.ndarray, + max_len: int, + num_replicas: int, + start_index: int, + rank: int, + rotation: int, ) -> None | list: """ Check if lengths can be distributed into `num_replicas` machines with at most @@ -54,16 +59,23 @@ def _lpt_packed_batch( Uses the LPT (Longest processing time first scheduling) algorithm Time: O(|lengths| log |lengths| + |lengths| log replicas) + Bins are packed to a token budget, which leaves sample counts skewed: with every + bin empty the heap tie-breaks on slot id, so slot 0 always takes the largest sample + and the highest slot absorbs the small ones. `rotation` re-maps slots to ranks so + that role cycles across batches. + Returns: `None` if unable to find a valid packing. Otherwise, return the batch indices that correspond to `rank`. """ - # Greedily assign lengths (in decreasing order) to the least full rank until they + # Greedily assign lengths (in decreasing order) to the least full slot until they # are all assigned or we run out of space. local_batch = [] heap = [_Bin(0, i) for i in range(num_replicas)] + target_slot = (rank - rotation) % num_replicas + # sort in descending order indices = np.argsort(lengths)[::-1] @@ -73,11 +85,11 @@ def _lpt_packed_batch( # Size doesn't fit in least full batch (or any others), report failure. return None - if heap[0].rank == rank: + if heap[0].slot == target_slot: # minimum bucket corresponds to this rank -> add idx to local batch local_batch.append(start_index + idx) - _ = heapreplace(heap, _Bin(new_fill, heap[0].rank)) + _ = heapreplace(heap, _Bin(new_fill, heap[0].slot)) return local_batch @@ -104,7 +116,7 @@ def _assign_to_packed_batches( lengths_so_far = 0 ind = 0 - result = [] + result: list = [] lengths_cumsum = np.cumsum(lengths) # binary search for max integer x such that the next x elements in shuffled lengths @@ -122,11 +134,14 @@ def _assign_to_packed_batches( lengths_cumsum[ind:], lengths_so_far + max_len * replicas, "right" ) + # Cycle the slot->rank mapping so no rank is permanently the many-samples one. + rotation = len(result) + batch = None while right - left > 1 and right > replicas: mid = (left + right) // 2 batch = _lpt_packed_batch( - lengths[ind : ind + mid], max_len, replicas, ind, rank + lengths[ind : ind + mid], max_len, replicas, ind, rank, rotation ) if batch is None: right = mid @@ -135,7 +150,7 @@ def _assign_to_packed_batches( if batch is None: batch = _lpt_packed_batch( - lengths[ind : ind + left], max_len, replicas, ind, rank + lengths[ind : ind + left], max_len, replicas, ind, rank, rotation ) ind += left diff --git a/src/speculators/train/sequence_parallel.py b/src/speculators/train/sequence_parallel.py new file mode 100644 index 000000000..76ccf293a --- /dev/null +++ b/src/speculators/train/sequence_parallel.py @@ -0,0 +1,297 @@ +"""Ulysses sequence parallelism utilities. + +Provides batch splitting and all-to-all communication primitives +for distributing packed sequences across SP ranks. +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +_SPLIT_KEYS = frozenset( + { + "hidden_states", + "verifier_last_hidden_states", + "input_ids", + "loss_mask", + "position_ids", + } +) + +_KEEP_FULL_KEYS = frozenset({"document_ids"}) + +_MIN_SPLIT_NDIM = 2 + + +# --------------------------------------------------------------------------- +# Differentiable all-to-all primitive (torch.compile-friendly custom op) +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("speculators::all_to_all_sp", mutates_args=()) +def all_to_all_sp( + input_tensor: torch.Tensor, + scatter_dim: int, + gather_dim: int, +) -> torch.Tensor: + from speculators.train.distributed import get_sp_group, get_sp_size # noqa: PLC0415 + + sp_size = get_sp_size() + if sp_size == 1: + return input_tensor + + sp_group = get_sp_group() + input_chunks = [ + c.contiguous() for c in input_tensor.chunk(sp_size, dim=scatter_dim) + ] + output_chunks = [torch.empty_like(c) for c in input_chunks] + dist.all_to_all(output_chunks, input_chunks, group=sp_group) + return torch.cat(output_chunks, dim=gather_dim).contiguous() + + +@all_to_all_sp.register_fake +def _all_to_all_sp_fake( + input_tensor: torch.Tensor, + scatter_dim: int, + gather_dim: int, +) -> torch.Tensor: + from speculators.train.distributed import get_sp_size # noqa: PLC0415 + + sp_size = get_sp_size() + if sp_size == 1: + return input_tensor.new_empty(input_tensor.shape) + + shape = list(input_tensor.shape) + shape[scatter_dim] //= sp_size + shape[gather_dim] *= sp_size + return input_tensor.new_empty(shape) + + +def _all_to_all_sp_setup(ctx, inputs, output): # noqa: ARG001 + _, scatter_dim, gather_dim = inputs + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + + +def _all_to_all_sp_backward(ctx, grad_output): + return ( + torch.ops.speculators.all_to_all_sp( + grad_output, ctx.gather_dim, ctx.scatter_dim + ), + None, + None, + ) + + +torch.library.register_autograd( + "speculators::all_to_all_sp", + _all_to_all_sp_backward, + setup_context=_all_to_all_sp_setup, +) + + +# --------------------------------------------------------------------------- +# Batch splitting +# --------------------------------------------------------------------------- + + +def split_batch_for_sp( + batch: dict[str, torch.Tensor], + sp_rank: int, + sp_size: int, +) -> dict[str, torch.Tensor]: + """Split batch tensors along the sequence dimension for SP. + + Tensors in ``_SPLIT_KEYS`` are chunked along dim 1 and the chunk + for *sp_rank* is kept. ``document_ids`` is kept full (needed for + attention mask construction over the global sequence). + + Adds ``full_seq_len`` to the returned dict so downstream code + can recover the original sequence length. + """ + if sp_size <= 1: + return batch + + out: dict[str, torch.Tensor] = {} + full_seq_len: int | None = None + device: torch.device | None = None + + for key, tensor in batch.items(): + if key in _SPLIT_KEYS and tensor.dim() >= _MIN_SPLIT_NDIM: + seq_len = tensor.shape[1] + if full_seq_len is None: + full_seq_len = seq_len + device = tensor.device + chunks = tensor.chunk(sp_size, dim=1) + out[key] = chunks[sp_rank].contiguous() + elif key in _KEEP_FULL_KEYS: + if full_seq_len is None and tensor.dim() >= _MIN_SPLIT_NDIM: + full_seq_len = tensor.shape[1] + device = tensor.device + out[key] = tensor + else: + out[key] = tensor + + if full_seq_len is not None: + out["full_seq_len"] = torch.tensor(full_seq_len, device=device) + + return out + + +# --------------------------------------------------------------------------- +# GQA head replication +# --------------------------------------------------------------------------- + + +def maybe_replicate_kv_heads( + key: torch.Tensor, + value: torch.Tensor, + sp_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Replicate KV heads when ``sp_size > num_kv_heads`` for GQA models. + + When ``sp_size <= num_kv_heads``, validates divisibility and returns + K/V unchanged. When ``sp_size > num_kv_heads``, replicates each KV + head so the total becomes divisible by ``sp_size``, allowing the + all-to-all scatter to split heads evenly across SP ranks. + + Args: + key: ``(B, num_kv_heads, S, D)`` + value: ``(B, num_kv_heads, S, D)`` + sp_size: Sequence-parallel world size. + """ + num_kv_heads = key.shape[1] + + if sp_size <= num_kv_heads: + if num_kv_heads % sp_size != 0: + raise ValueError( + f"num_kv_heads ({num_kv_heads}) must be divisible by " + f"sp_size ({sp_size})" + ) + return key, value + + if sp_size % num_kv_heads != 0: + raise ValueError( + f"sp_size ({sp_size}) must be divisible by " + f"num_kv_heads ({num_kv_heads}) when sp_size > num_kv_heads" + ) + + replication_factor = sp_size // num_kv_heads + key = key.repeat_interleave(replication_factor, dim=1) + value = value.repeat_interleave(replication_factor, dim=1) + return key, value + + +# --------------------------------------------------------------------------- +# Ulysses scatter / gather for attention +# --------------------------------------------------------------------------- + + +def ulysses_scatter( + x: torch.Tensor, + sp_group: ProcessGroup, # noqa: ARG001 + sp_size: int, # noqa: ARG001 +) -> torch.Tensor: + """All-to-all: sequence-parallel to head-parallel layout. + + Input: ``(B, H, S_local, D)`` + Output: ``(B, H // sp_size, S_full, D)`` + + Scatters heads (dim 1) and gathers sequence chunks (dim 2). + Gradients flow back through the reverse all-to-all automatically. + """ + return torch.ops.speculators.all_to_all_sp(x, 1, 2) + + +def ulysses_gather( + x: torch.Tensor, + sp_group: ProcessGroup, # noqa: ARG001 + sp_size: int, # noqa: ARG001 +) -> torch.Tensor: + """All-to-all: head-parallel to sequence-parallel layout. + + Input: ``(B, H_local, S_full, D)`` + Output: ``(B, H_local * sp_size, S_full // sp_size, D)`` + + Inverse of :func:`ulysses_scatter`. + """ + return torch.ops.speculators.all_to_all_sp(x, 2, 1) + + +# --------------------------------------------------------------------------- +# DFlash-specific Ulysses attention +# --------------------------------------------------------------------------- + + +def dflash_flex_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask, + scaling: float | None = None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Ulysses SP attention for DFlash with asymmetric Q/KV. + + DFlash K/V are ``[context | noise]`` concatenated along the sequence + dim. A naive all-to-all would interleave context and noise from + different ranks, breaking the mask. This function splits K/V at + ``ctx_len``, does separate all-to-alls, and concatenates + ``[global_ctx | global_noise]`` to preserve the layout that + ``create_anchor_block_mask_mod`` expects. + + Falls back to the base ``flex_attention_forward`` when SP is + inactive (``sp_size <= 1``). + """ + from speculators.train.distributed import get_sp_size # noqa: PLC0415 + + sp_size = get_sp_size() + + if sp_size <= 1: + from speculators.models.attention import ( # noqa: PLC0415 + flex_attention_forward, + ) + + return flex_attention_forward( + module, query, key, value, attention_mask, scaling=scaling, **kwargs + ) + + from torch.nn.attention.flex_attention import flex_attention # noqa: PLC0415 + + key, value = maybe_replicate_kv_heads(key, value, sp_size) + + # DFlash K/V = [context | noise]. Q is noise-only, + # so ctx_len = kv_seq_len - q_seq_len. + ctx_len = key.shape[2] - query.shape[2] + + k_ctx, k_noise = key.split([ctx_len, key.shape[2] - ctx_len], dim=2) + v_ctx, v_noise = value.split([ctx_len, value.shape[2] - ctx_len], dim=2) + + query = torch.ops.speculators.all_to_all_sp(query, 1, 2) + k_ctx = torch.ops.speculators.all_to_all_sp(k_ctx, 1, 2) + k_noise = torch.ops.speculators.all_to_all_sp(k_noise, 1, 2) + v_ctx = torch.ops.speculators.all_to_all_sp(v_ctx, 1, 2) + v_noise = torch.ops.speculators.all_to_all_sp(v_noise, 1, 2) + + key = torch.cat([k_ctx, k_noise], dim=2) + value = torch.cat([v_ctx, v_noise], dim=2) + + num_q_local = query.shape[1] + num_kv_local = key.shape[1] + + attn_output = flex_attention( + query.contiguous(), + key.contiguous(), + value.contiguous(), + score_mod=None, + block_mask=attention_mask, + enable_gqa=num_q_local != num_kv_local, + scale=scaling, + ) + + attn_output = torch.ops.speculators.all_to_all_sp(attn_output, 2, 1) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, None diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py index 2c63fe312..caa0f253c 100644 --- a/src/speculators/train/trainer.py +++ b/src/speculators/train/trainer.py @@ -28,12 +28,17 @@ ) from speculators.train.distributed import ( apply_fully_sharded, + get_dp_group, get_local_rank, get_rank, + get_sp_rank, + get_sp_size, is_distributed, + register_sp_gradient_hooks, ) from speculators.train.graceful_shutdown import with_graceful_shutdown from speculators.train.optimizers import build_optimizers +from speculators.train.sequence_parallel import split_batch_for_sp from speculators.train.utils import normalize_counted_metrics root_logger = logging.getLogger("speculators") @@ -93,6 +98,10 @@ def profile(self, num_tokens: int) -> dict[str, float] | None: warnings.filterwarnings("ignore", category=TqdmExperimentalWarning) MIN_STEP_PCT = 0.25 +# Re-synchronise ranks every N validation batches to bound cross-rank skew, which would +# otherwise blow the NCCL watchdog at the end-of-epoch metrics all-reduce. 0 disables. +_VAL_SYNC_INTERVAL = 50 + class TrainerConfig(NamedTuple): lr: float @@ -118,6 +127,7 @@ class TrainerConfig(NamedTuple): hidden_states_dtype: torch.dtype = torch.bfloat16 log_freq: int = 1 fsdp_shard: bool = False + max_steps: int | None = None def _resolve_scheduler_steps( @@ -298,6 +308,7 @@ def _setup_model_fsdp(self, load_checkpoint: bool): full_state_dict = self.model.state_dict() apply_fully_sharded(self.model, param_dtype=self.config.hidden_states_dtype) + self._sp_grad_hooks = register_sp_gradient_hooks(self.model) if load_checkpoint: self.checkpointer.load_model_state_dict(self.model) @@ -327,8 +338,12 @@ def _setup_model_ddp(self, load_checkpoint: bool): dist.broadcast(param.data, src=0) dist.barrier() - # DDP constructor broadcasts rank 0's params to all ranks - self.model = DistributedDataParallel(self.model) # type: ignore[assignment] + self._sp_grad_hooks = register_sp_gradient_hooks(self.model) + + # Scope DDP to the DP sub-group so its all-reduce averages across DP + # peers only; the SP gradient hooks handle the cross-SP sum separately. + dp_group = get_dp_group() + self.model = DistributedDataParallel(self.model, process_group=dp_group) # type: ignore[assignment] def setup_optimizer(self): # Setup optimizer(s). The "muon" option returns two optimizers (Muon for the @@ -450,6 +465,8 @@ def train_epoch(self, epoch: int): else v for k, v in batch.items() } + if getattr(self.model, "_sp_splits_batch", True): + gpu_batch = split_batch_for_sp(gpu_batch, get_sp_rank(), get_sp_size()) with torch.autocast( self.device_type, dtype=self.config.hidden_states_dtype @@ -502,6 +519,12 @@ def train_epoch(self, epoch: int): ) self.global_step += 1 + if ( + self.config.max_steps is not None + and self.global_step >= self.config.max_steps + ): + break + if ( step_interval is not None and not self.config.save_best @@ -511,6 +534,12 @@ def train_epoch(self, epoch: int): ): self.maybe_save_checkpoint(epoch, local_step=local_step) + def _maybe_val_sync(self, batch_index: int) -> None: + if not self.is_distributed or _VAL_SYNC_INTERVAL <= 0: + return + if batch_index > 0 and batch_index % _VAL_SYNC_INTERVAL == 0: + dist.barrier() + @torch.no_grad() def val_epoch(self, epoch: int) -> dict[str, float] | None: if self.val_loader is None: @@ -524,13 +553,16 @@ def val_epoch(self, epoch: int) -> dict[str, float] | None: accumulated: dict[str, torch.Tensor] = {} num_batches = len(val_loader) - for batch in val_loader: + for i, batch in enumerate(val_loader): + self._maybe_val_sync(i) gpu_batch = { k: v.to(self.local_rank, non_blocking=True) if isinstance(v, torch.Tensor) else v for k, v in batch.items() } + if getattr(self.model, "_sp_splits_batch", True): + gpu_batch = split_batch_for_sp(gpu_batch, get_sp_rank(), get_sp_size()) with torch.autocast( self.device_type, dtype=self.config.hidden_states_dtype diff --git a/src/speculators/train/utils.py b/src/speculators/train/utils.py index da85c308a..cce04d695 100644 --- a/src/speculators/train/utils.py +++ b/src/speculators/train/utils.py @@ -100,8 +100,13 @@ def normalize_counted_metrics( return metrics -def save_train_command(save_path: str) -> None: - """Write the launch command and provenance header to save_path/train_command.txt.""" +def save_train_command(save_path: str, argv: list[str] | None = None) -> None: + """Write the launch command and provenance header to save_path/train_command.txt. + + ``argv`` is the exact command the run was resolved from (``TrainConfig`` records + it during resolution); it falls back to the live ``sys.argv`` when a caller has + no recorded argv, so a direct call is unchanged. + """ try: sha = subprocess.run( ["git", "rev-parse", "HEAD"], # noqa: S607 @@ -129,7 +134,7 @@ def save_train_command(save_path: str) -> None: ] ) - command = shlex.join(sys.argv) + command = shlex.join(argv or sys.argv) content = f"{header}\n{command}\n" path = Path(save_path) diff --git a/src/speculators/utils/argparse_utils.py b/src/speculators/utils/argparse_utils.py deleted file mode 100644 index de637bcd0..000000000 --- a/src/speculators/utils/argparse_utils.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Small, reusable argparse helpers.""" - -import argparse -from collections.abc import Iterable - - -def explicitly_provided_dests( - parser: argparse.ArgumentParser, - dests: Iterable[str], - argv: list[str] | None = None, -) -> set[str]: - """Return the subset of ``dests`` whose option actually appeared on the CLI. - - Re-parses ``argv`` into a throwaway namespace pre-seeded with a unique sentinel - per dest. argparse skips applying an action's default when the attribute is already - present on the namespace, and only overwrites it when the option is provided -- so a - value that merely *equals* the default is still detected as explicitly provided - (unlike comparing the parsed value against ``parser.get_default(dest)``). - - :param parser: The parser to re-run. - :param dests: Argparse destination names to check (e.g. ``"num_layers"``). - :param argv: Arguments to re-parse; defaults to ``None`` (i.e. ``sys.argv``). Pass - the same ``argv`` the parser was originally invoked with so the detected - destinations match the arguments actually parsed. - :return: The subset of ``dests`` that were explicitly provided. - """ - sentinels = {dest: object() for dest in dests} - namespace = argparse.Namespace(**sentinels) - parser.parse_args(args=argv, namespace=namespace) - return { - dest - for dest, sentinel in sentinels.items() - if getattr(namespace, dest) is not sentinel - } diff --git a/tests/unit/models/test_dflash_targets.py b/tests/unit/models/test_dflash_targets.py new file mode 100644 index 000000000..9dbcb365f --- /dev/null +++ b/tests/unit/models/test_dflash_targets.py @@ -0,0 +1,63 @@ +import pytest +import torch +from transformers.models.qwen3.modeling_qwen3 import Qwen3Config + +from speculators.models.dflash import DFlashSpeculatorConfig +from speculators.models.dflash.core import DFlashDraftModel + + +def _tiny_model(sample_from_anchor: bool) -> DFlashDraftModel: + tl_config = Qwen3Config( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=8, + vocab_size=64, + _attn_implementation="eager", # type: ignore[call-arg] + ) + config = DFlashSpeculatorConfig( + transformer_layer_config=tl_config, + draft_vocab_size=64, + block_size=4, + aux_hidden_state_layer_ids=[0, 1], + mask_token_id=0, + sample_from_anchor=sample_from_anchor, + ) + model = DFlashDraftModel(config) + torch.nn.init.normal_(model.verifier_lm_head.weight) + torch.nn.init.ones_(model.verifier_norm.weight) + return model.eval() + + +@pytest.mark.parametrize("max_anchors", [5, 16]) +@pytest.mark.parametrize("sample_from_anchor", [False, True]) +def test_targets_match_full_sequence_roll(sample_from_anchor, max_anchors): + torch.manual_seed(0) + model = _tiny_model(sample_from_anchor) + seq_len = 32 + hidden_states = torch.randn(1, seq_len, 2 * 16) + verifier_last_hidden_states = torch.randn(1, seq_len, 16) + input_ids = torch.randint(0, 64, (1, seq_len)) + loss_mask = torch.ones(1, seq_len) + document_ids = torch.zeros(1, seq_len, dtype=torch.long) + + with torch.no_grad(): + _, _, targets, _, anchored_block_indices = model._backbone_forward( + hidden_states, + input_ids, + loss_mask, + verifier_last_hidden_states, + document_ids, + max_anchors=max_anchors, + ) + + full_logits = model.verifier_lm_head( + model.verifier_norm(verifier_last_hidden_states) + ) + if not sample_from_anchor: + full_logits = torch.roll(full_logits, 1, dims=1) + expected = full_logits[:, anchored_block_indices] + + torch.testing.assert_close(targets, expected, atol=1e-5, rtol=0) diff --git a/tests/unit/scripts/test_benchmark.py b/tests/unit/scripts/test_benchmark.py new file mode 100644 index 000000000..fc29c0936 --- /dev/null +++ b/tests/unit/scripts/test_benchmark.py @@ -0,0 +1,453 @@ +"""Unit tests for the benchmark harness.""" + +from __future__ import annotations + +import json +import logging +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import torch + +# Add scripts/ to the import path the same way the benchmark script does. +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "scripts")) + +from benchmark import ( # type: ignore[import-not-found] + _MetricCapture, + _SyntheticLoader, + collect_provenance, + compare_benchmarks, + compute_statistics, + create_synthetic_batch, +) + +# --------------------------------------------------------------------------- +# compute_statistics +# --------------------------------------------------------------------------- + + +class TestComputeStatistics: + def test_basic(self): + values = [10.0, 20.0, 30.0, 40.0, 50.0] + result = compute_statistics(values) + assert result["mean"] == pytest.approx(30.0) + assert result["min"] == 10.0 + assert result["max"] == 50.0 + assert result["median"] == 30.0 + assert result["count"] == 5 + assert result["std"] > 0 + + def test_single_value(self): + result = compute_statistics([42.0]) + assert result["mean"] == 42.0 + assert result["std"] == 0.0 + assert result["min"] == 42.0 + assert result["max"] == 42.0 + assert result["median"] == 42.0 + assert result["count"] == 1 + + def test_identical_values(self): + result = compute_statistics([5.0, 5.0, 5.0]) + assert result["mean"] == 5.0 + assert result["std"] == 0.0 + + +# --------------------------------------------------------------------------- +# collect_provenance +# --------------------------------------------------------------------------- + + +class TestCollectProvenance: + @patch("benchmark.torch") + @patch("benchmark.subprocess.run") + def test_keys_present(self, mock_run, mock_torch): + mock_run.return_value = MagicMock(returncode=0, stdout="abc123\n") + mock_torch.cuda.is_available.return_value = False + mock_torch.cuda.device_count.return_value = 0 + mock_torch.__version__ = "2.9.0" + mock_torch.version.cuda = "12.4" + + result = collect_provenance() + + expected_keys = { + "git_sha", + "timestamp", + "hostname", + "python_version", + "pytorch_version", + "cuda_version", + "speculators_version", + "transformers_version", + "gpu_info", + "num_gpus", + } + assert set(result.keys()) == expected_keys + assert result["git_sha"] == "abc123" + assert result["num_gpus"] == 0 + + @patch("benchmark.torch") + @patch("benchmark.subprocess.run") + def test_git_failure(self, mock_run, mock_torch): + mock_run.return_value = MagicMock(returncode=128, stdout="") + mock_torch.cuda.is_available.return_value = False + mock_torch.cuda.device_count.return_value = 0 + mock_torch.__version__ = "2.9.0" + mock_torch.version.cuda = None + + result = collect_provenance() + assert result["git_sha"] == "unknown" + assert result["cuda_version"] == "none" + + +# --------------------------------------------------------------------------- +# create_synthetic_batch +# --------------------------------------------------------------------------- + + +class TestCreateSyntheticBatch: + def test_shapes(self): + seq_len = 128 + hidden_size = 64 + num_layers = 3 + batch = create_synthetic_batch( + total_seq_len=seq_len, + hidden_size=hidden_size, + num_target_layers=num_layers, + device="cpu", + ) + + assert batch["hidden_states"].shape == ( + 1, + seq_len, + num_layers * hidden_size, + ) + assert batch["input_ids"].shape == (1, seq_len) + assert batch["verifier_last_hidden_states"].shape == ( + 1, + seq_len, + hidden_size, + ) + assert batch["loss_mask"].shape == (1, seq_len) + assert batch["position_ids"].shape == (1, seq_len) + assert batch["document_ids"].shape == (1, seq_len) + + def test_dtypes(self): + batch = create_synthetic_batch( + total_seq_len=64, + hidden_size=32, + num_target_layers=2, + dtype=torch.bfloat16, + device="cpu", + ) + + assert batch["hidden_states"].dtype == torch.bfloat16 + assert batch["verifier_last_hidden_states"].dtype == torch.bfloat16 + assert batch["input_ids"].dtype == torch.long + assert batch["loss_mask"].dtype == torch.bool + assert batch["position_ids"].dtype == torch.long + assert batch["document_ids"].dtype == torch.long + + def test_position_ids_start_at_one(self): + batch = create_synthetic_batch( + total_seq_len=10, + hidden_size=16, + num_target_layers=1, + device="cpu", + ) + assert batch["position_ids"][0, 0].item() == 1 + assert batch["position_ids"][0, -1].item() == 10 + + def test_document_ids_all_zero(self): + batch = create_synthetic_batch( + total_seq_len=10, + hidden_size=16, + num_target_layers=1, + device="cpu", + ) + assert (batch["document_ids"] == 0).all() + + def test_all_keys_present(self): + batch = create_synthetic_batch( + total_seq_len=8, + hidden_size=16, + num_target_layers=1, + device="cpu", + ) + expected_keys = { + "hidden_states", + "input_ids", + "verifier_last_hidden_states", + "loss_mask", + "position_ids", + "document_ids", + } + assert set(batch.keys()) == expected_keys + + +# --------------------------------------------------------------------------- +# _MetricCapture +# --------------------------------------------------------------------------- + + +class TestMetricCapture: + def test_captures_profile_dicts(self): + capture = _MetricCapture() + profile = {"step_ms": 45.0, "fwd_ms": 20.0} + record = logging.LogRecord( + name="speculators.metrics", + level=logging.INFO, + pathname="", + lineno=0, + msg={"train": {}, "profile": profile, "epoch": 0}, + args=None, + exc_info=None, + ) + capture.emit(record) + assert len(capture.profiles) == 1 + assert capture.profiles[0] is profile + + def test_ignores_records_without_profile(self): + capture = _MetricCapture() + record = logging.LogRecord( + name="speculators.metrics", + level=logging.INFO, + pathname="", + lineno=0, + msg={"train": {}, "epoch": 0}, + args=None, + exc_info=None, + ) + capture.emit(record) + assert len(capture.profiles) == 0 + + def test_ignores_none_profile(self): + capture = _MetricCapture() + record = logging.LogRecord( + name="speculators.metrics", + level=logging.INFO, + pathname="", + lineno=0, + msg={"train": {}, "profile": None, "epoch": 0}, + args=None, + exc_info=None, + ) + capture.emit(record) + assert len(capture.profiles) == 0 + + def test_ignores_non_dict_messages(self): + capture = _MetricCapture() + record = logging.LogRecord( + name="speculators.metrics", + level=logging.INFO, + pathname="", + lineno=0, + msg="some string message", + args=None, + exc_info=None, + ) + capture.emit(record) + assert len(capture.profiles) == 0 + + def test_captures_multiple(self): + capture = _MetricCapture() + for i in range(5): + record = logging.LogRecord( + name="speculators.metrics", + level=logging.INFO, + pathname="", + lineno=0, + msg={"profile": {"step_ms": float(i)}, "train": {}}, + args=None, + exc_info=None, + ) + capture.emit(record) + assert len(capture.profiles) == 5 + assert capture.profiles[3]["step_ms"] == 3.0 + + +# --------------------------------------------------------------------------- +# _SyntheticLoader +# --------------------------------------------------------------------------- + + +class TestSyntheticLoader: + def test_len(self): + batch = {"x": torch.zeros(1)} + loader = _SyntheticLoader(batch, num_steps=7) + assert len(loader) == 7 + + def test_iter_yields_correct_count(self): + batch = {"x": torch.zeros(1)} + loader = _SyntheticLoader(batch, num_steps=3) + batches = list(loader) + assert len(batches) == 3 + + def test_iter_yields_same_batch(self): + batch = {"x": torch.tensor([1.0, 2.0])} + loader = _SyntheticLoader(batch, num_steps=3) + for b in loader: + assert b is batch + + def test_batch_sampler_has_set_epoch(self): + batch = {"x": torch.zeros(1)} + loader = _SyntheticLoader(batch, num_steps=1) + assert hasattr(loader.batch_sampler, "set_epoch") + loader.batch_sampler.set_epoch(5) + + +# --------------------------------------------------------------------------- +# Warmup / measured split +# --------------------------------------------------------------------------- + + +class TestWarmupMeasuredSplit: + """Tests for the profile slicing logic used in run_benchmark.""" + + def test_discard_warmup(self): + warmup_steps = 3 + all_profiles = [{"step_ms": float(i)} for i in range(13)] + measured = all_profiles[warmup_steps:] + assert len(measured) == 10 + assert measured[0]["step_ms"] == 3.0 + + def test_exact_boundary(self): + warmup_steps = 5 + all_profiles = [{"step_ms": float(i)} for i in range(5)] + measured = all_profiles[warmup_steps:] + assert len(measured) == 0 + + def test_zero_warmup(self): + warmup_steps = 0 + all_profiles = [{"step_ms": float(i)} for i in range(10)] + measured = all_profiles[warmup_steps:] + assert len(measured) == 10 + assert measured[0]["step_ms"] == 0.0 + + def test_fallback_on_insufficient_profiles(self): + """When fewer profiles than expected, fallback uses all of them.""" + warmup_steps = 10 + all_profiles = [{"step_ms": float(i)} for i in range(5)] + measured = all_profiles[warmup_steps:] + if not measured: + measured = all_profiles + assert len(measured) == 5 + + +# --------------------------------------------------------------------------- +# compare_benchmarks +# --------------------------------------------------------------------------- + + +def _make_result( + step_ms_mean=45.0, + step_ms_std=1.0, + peak_alloc=2048.0, + git_sha="aaa", + gpu_name="H100", + speculator_type="eagle3", +): + """Create a minimal benchmark result dict for testing.""" + timing = {} + for key in ( + "step_ms", + "fwd_ms", + "bwd_ms", + "opt_ms", + "fetch_ms", + "tokens_per_s", + ): + timing[key] = { + "mean": step_ms_mean, + "std": step_ms_std, + "min": step_ms_mean - 2, + "max": step_ms_mean + 2, + "median": step_ms_mean, + "count": 50, + } + return { + "benchmark_version": "1.0", + "provenance": { + "git_sha": git_sha, + "gpu_info": [{"name": gpu_name, "total_memory_gb": 80.0}], + }, + "config": { + "speculator_type": speculator_type, + "hidden_size": 4096, + "total_seq_len": 8192, + "num_gpus_used": 1, + "fsdp_shard": False, + "optimizer": "muon", + "hidden_states_dtype": "bfloat16", + }, + "memory": { + "peak_allocated_mb": peak_alloc, + "peak_reserved_mb": peak_alloc + 1024, + }, + "timing": timing, + } + + +class TestCompareBenchmarks: + def test_basic_compare(self, tmp_path, capsys): + baseline = _make_result(step_ms_mean=50.0, git_sha="aaa111") + candidate = _make_result(step_ms_mean=45.0, git_sha="bbb222") + + baseline_path = tmp_path / "baseline.json" + candidate_path = tmp_path / "candidate.json" + baseline_path.write_text(json.dumps(baseline)) + candidate_path.write_text(json.dumps(candidate)) + + compare_benchmarks(str(baseline_path), str(candidate_path)) + + output = capsys.readouterr().out + assert "aaa111" in output + assert "bbb222" in output + assert "step_ms" in output + assert "-5.00" in output or "-10.0%" in output + + def test_comparability_warning_gpu(self, tmp_path, capsys): + baseline = _make_result(gpu_name="H100") + candidate = _make_result(gpu_name="A100") + + baseline_path = tmp_path / "b.json" + candidate_path = tmp_path / "c.json" + baseline_path.write_text(json.dumps(baseline)) + candidate_path.write_text(json.dumps(candidate)) + + compare_benchmarks(str(baseline_path), str(candidate_path)) + + output = capsys.readouterr().out + assert "GPU" in output + assert "H100" in output + assert "A100" in output + + def test_comparability_warning_config(self, tmp_path, capsys): + baseline = _make_result(speculator_type="eagle3") + candidate = _make_result(speculator_type="dflash") + + baseline_path = tmp_path / "b.json" + candidate_path = tmp_path / "c.json" + baseline_path.write_text(json.dumps(baseline)) + candidate_path.write_text(json.dumps(candidate)) + + compare_benchmarks(str(baseline_path), str(candidate_path)) + + output = capsys.readouterr().out + assert "Speculator type" in output + + def test_memory_delta(self, tmp_path, capsys): + baseline = _make_result(peak_alloc=2000.0) + candidate = _make_result(peak_alloc=1800.0) + + baseline_path = tmp_path / "b.json" + candidate_path = tmp_path / "c.json" + baseline_path.write_text(json.dumps(baseline)) + candidate_path.write_text(json.dumps(candidate)) + + compare_benchmarks(str(baseline_path), str(candidate_path)) + + output = capsys.readouterr().out + assert "peak_allocated_mb" in output + assert "-200.0" in output diff --git a/tests/unit/train/config/__init__.py b/tests/unit/train/config/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/train/config/test_artifacts.py b/tests/unit/train/config/test_artifacts.py new file mode 100644 index 000000000..48f7e3557 --- /dev/null +++ b/tests/unit/train/config/test_artifacts.py @@ -0,0 +1,191 @@ +"""Serialization / reproducibility seam tests. + +The reproducibility contract is the round-trip at the ``from_sources`` seam: +``dump_yaml()`` -> re-load via ``config_path`` -> an identical resolved config. +Also asserted here: the emitted YAML is ``train:``-wrapped and clean (no +provenance inlined, re-loadable), ``cfg.save(dir)`` writes both artifacts, and +``--dump-config`` prints a valid config and exits cleanly at the ``resolve`` +boundary. No golden ``vars(args)`` snapshot. +""" + +import warnings + +import pytest +import yaml + +from speculators.train.config import TrainConfig + + +def _resolved(**cli): + # Silence the mismatched-algorithm-block warnings that fire when a cross-type + # recipe is resolved; they are exercised in test_resolution. + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=".*does not use the.*algorithm group.*", + category=UserWarning, + ) + return TrainConfig.from_sources(cli=cli, config_path=None, argv=["train.py"]) + + +def _reload(tmp_path, cfg): + path = tmp_path / "run.yaml" + path.write_text(cfg.dump_yaml()) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=".*does not use the.*algorithm group.*", + category=UserWarning, + ) + return TrainConfig.from_sources( + cli={}, config_path=str(path), argv=["train.py"] + ) + + +# --------------------------------------------------------------------------- # +# dump_yaml: the round-trip reproducibility contract +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "cli", + [ + {"verifier_name_or_path": "m"}, + {"verifier_name_or_path": "m", "lr": 0.9, "num_layers": 4}, + # --from-pretrained is the case a full materialized dump could not survive: + # persisting the derived decoder-shaping defaults would self-reject on reload. + {"verifier_name_or_path": "m", "from_pretrained": "ckpt"}, + {"verifier_name_or_path": "m", "speculator_type": "dspark", "markov_rank": 128}, + # Non-scalar recipe values: a nargs list and a JSON-dict loss spec (both + # appear verbatim in examples/train/*.sh) must survive the YAML round-trip. + {"verifier_name_or_path": "m", "target_layer_ids": [2, 5, 9]}, + {"verifier_name_or_path": "m", "loss_fn": '{"ce": 0.1, "tv": 0.9}'}, + ], +) +def test_dump_yaml_round_trips_to_identical_resolved_config(tmp_path, cli): + cfg = _resolved(**cli) + assert cfg.flatten() == _reload(tmp_path, cfg).flatten() + + +def test_dump_yaml_is_byte_stable_across_reload(tmp_path): + # Emission follows the pinned flatten() order, not the provenance dict order + # (which varies with flag-vs-yaml source), so dumping a reloaded config yields + # byte-identical YAML -- the run.yaml artifact is stable, not just equivalent. + cfg = _resolved(verifier_name_or_path="m", lr=0.9, num_layers=4, epochs=7) + first = cfg.dump_yaml() + assert _reload(tmp_path, cfg).dump_yaml() == first + + +def test_dump_yaml_is_stage_wrapped_and_clean(tmp_path): + cfg = _resolved(verifier_name_or_path="m", lr=0.9) + doc = yaml.safe_load(cfg.dump_yaml()) + # Canonical stage shape: a single top-level ``train:`` block. + assert set(doc) == {"train"} + # Clean: only supplied values, no provenance annotations, no default noise. + assert doc["train"] == { + "verifier": {"verifier_name_or_path": "m"}, + "optimizer": {"lr": 0.9}, + } + + +def test_saved_run_yaml_reloads_to_identical_config(tmp_path): + # The file written next to the checkpoints re-loads (via the ``config_path`` + # seam) to the same resolved config. + cfg = _resolved(verifier_name_or_path="m", num_layers=6, lr=0.3) + cfg.save(str(tmp_path)) + reran = TrainConfig.from_sources( + cli={}, config_path=str(tmp_path / "run.yaml"), argv=["train.py"] + ) + assert reran.flatten() == cfg.flatten() + + +# --------------------------------------------------------------------------- # +# save: the reproducibility artifacts next to the checkpoints +# --------------------------------------------------------------------------- # + + +def test_save_writes_run_yaml_and_train_command(tmp_path): + cfg = _resolved(verifier_name_or_path="m") + cfg.save(str(tmp_path)) + assert (tmp_path / "run.yaml").exists() + manifest = (tmp_path / "train_command.txt").read_text() + # The manifest records the argv this config was resolved from + the environment. + assert "train.py" in manifest + assert "# Git SHA:" in manifest + assert "# World size:" in manifest + for pkg in ("speculators", "transformers", "torch"): + assert f"# {pkg}:" in manifest + + +def test_no_provenance_sidecar_is_written(tmp_path): + _resolved(verifier_name_or_path="m").save(str(tmp_path)) + assert not (tmp_path / "run.provenance.yaml").exists() + assert {p.name for p in tmp_path.iterdir()} == {"run.yaml", "train_command.txt"} + + +def test_dump_yaml_handles_from_flat_config_without_provenance(tmp_path): + # A config rebuilt via from_flat carries no layer provenance; dump_yaml/save + # must still work (the old code KeyError'd on the empty provenance dict), + # falling back to "differs from defaults" to pick the emitted keys. + resolved = _resolved( + verifier_name_or_path="m", + speculator_type="eagle3", + save_path="/tmp/ckpt", + epochs=7, + lr=0.3, + ) + rebuilt = TrainConfig.from_flat(resolved.flatten()) + # Precondition that triggered the old KeyError. + assert rebuilt.provenance == {} + + doc = rebuilt.dump_yaml() + assert doc + assert "train:" in doc + # The customized (differs-from-default) values survive into the emitted YAML. + assert "save_path" in doc + assert "/tmp/ckpt" in doc + + rebuilt.save(str(tmp_path)) + assert (tmp_path / "run.yaml").exists() + + +# --------------------------------------------------------------------------- # +# --dump-config: the config-out surface at the resolve boundary +# --------------------------------------------------------------------------- # + + +def test_dump_config_prints_valid_config_and_exits(capsys): + with pytest.raises(SystemExit) as exc: + TrainConfig.resolve( + ["--verifier-name-or-path", "m", "--lr", "0.5", "--dump-config"] + ) + assert exc.value.code == 0 + printed = capsys.readouterr().out + doc = yaml.safe_load(printed) + assert doc["train"]["optimizer"]["lr"] == 0.5 + assert doc["train"]["verifier"]["verifier_name_or_path"] == "m" + + +def test_dump_config_output_is_a_reloadable_config(tmp_path, capsys): + with pytest.raises(SystemExit): + TrainConfig.resolve( + ["--verifier-name-or-path", "m", "--num-layers", "7", "--dump-config"] + ) + path = tmp_path / "run.yaml" + path.write_text(capsys.readouterr().out) + reran = TrainConfig.from_sources(cli={}, config_path=str(path), argv=["train.py"]) + assert reran.draft.num_layers == 7 + assert reran.verifier.verifier_name_or_path == "m" + + +def test_dump_config_includes_config_file_values(tmp_path, capsys): + # A value supplied only in --config surfaces in the --dump-config output. + config = tmp_path / "run.yaml" + config.write_text("train:\n optimizer:\n lr: 0.2\n") + with pytest.raises(SystemExit) as exc: + TrainConfig.resolve( + ["--verifier-name-or-path", "m", "--config", str(config), "--dump-config"] + ) + assert exc.value.code == 0 + doc = yaml.safe_load(capsys.readouterr().out) + assert doc["train"]["optimizer"]["lr"] == 0.2 diff --git a/tests/unit/train/config/test_backend_reconciliation.py b/tests/unit/train/config/test_backend_reconciliation.py new file mode 100644 index 000000000..0544c2aa2 --- /dev/null +++ b/tests/unit/train/config/test_backend_reconciliation.py @@ -0,0 +1,70 @@ +"""Drift-guard: every backend train-arg must be mirrored as a schema field. + +``hs_connectors`` is speculators-agnostic (argparse-based, no pydantic) because +vLLM consumes it standalone. So speculators no longer builds its parser by +calling each backend's ``add_train_args``; the schema generates the whole flag +surface, and each backend's train-args are *mirrored* as schema fields (e.g. +``FileBackend``'s ``--hidden-states-path`` is mirrored as ``DataArgs``' +``hidden_states_path``). + +That mirroring is the risk this test guards. :func:`resolution.resolve` filters +the parsed namespace down to :data:`~schema.CONFIG_DESTS`, so a backend train-arg +with no matching schema field would be *silently dropped* during resolution. A +new backend (NIXL/RDMA is on the roadmap) could introduce that gap without any +error. This test makes the gap loud: for every registered backend it collects the +argparse dests ``add_train_args`` registers and asserts each one is a config dest. +""" + +import argparse + +import pytest + +from hs_connectors import HiddenStatesBackend +from speculators.train.config.schema import CONFIG_DESTS + + +def _backend_train_arg_dests(backend_cls: type[HiddenStatesBackend]) -> set[str]: + """The argparse dests a backend registers via ``add_train_args``. + + Introspected from a throwaway parser's ``_actions``, skipping argparse's + auto-added ``help`` action (and any suppressed dest) so only genuine + backend train-args remain. + """ + scratch_parser = argparse.ArgumentParser() + backend_cls.add_train_args(scratch_parser) + return { + action.dest + for action in scratch_parser._actions + if action.dest not in ("help", argparse.SUPPRESS) + } + + +def test_registry_is_populated_with_file_backend(): + # Prove the parametrized test below is exercising something real: the registry + # is non-empty and the one backend that exists today mirrors its train-arg. + assert HiddenStatesBackend.registry, "no backends registered" + assert "file" in HiddenStatesBackend.registry + assert "hidden_states_path" in _backend_train_arg_dests( + HiddenStatesBackend.registry["file"] + ) + assert "hidden_states_path" in CONFIG_DESTS + + +@pytest.mark.parametrize( + ("name", "backend_cls"), + sorted(HiddenStatesBackend.registry.items()), + ids=sorted(HiddenStatesBackend.registry), +) +def test_backend_train_args_are_mirrored_in_schema( + name: str, backend_cls: type[HiddenStatesBackend] +): + # Every dest a backend adds via add_train_args must have a matching schema + # field, or resolve() (which filters to CONFIG_DESTS) drops the value. + dests = _backend_train_arg_dests(backend_cls) + missing = sorted(dests - CONFIG_DESTS) + assert not missing, ( + f"Backend '{name}' registers train-arg dest(s) {missing} via add_train_args " + f"that have no matching schema field. speculators mirrors backend train-args " + f"as schema fields; add them (e.g. to DataArgs) or resolve() will silently " + f"drop them (filtered by CONFIG_DESTS)." + ) diff --git a/tests/unit/train/config/test_resolution.py b/tests/unit/train/config/test_resolution.py new file mode 100644 index 000000000..8358976bb --- /dev/null +++ b/tests/unit/train/config/test_resolution.py @@ -0,0 +1,547 @@ +"""Resolution seam tests plus the example-recipe back-compat suite. + +The recipe suite is the central backward-compatibility guard: each real +``examples/train/*.sh`` recipe's ``train.py`` invocation is extracted (bash does +the variable expansion and word-splitting, exactly as a user's shell would) and +run through :meth:`TrainConfig.resolve`; the flags it sets must resolve to their +expected values in the flat dict. A renamed, dropped, or retyped flag breaks a +real recipe. There is deliberately no golden ``vars(args)`` snapshot. +""" + +import re +import shutil +import subprocess +import warnings +from pathlib import Path + +import pytest +import yaml +from pydantic import ValidationError + +from speculators.train.config import TrainConfig +from speculators.train.config.resolution import ( + ConfigError, + build_parser, +) +from speculators.train.config.schema import nest_flat + +EXAMPLES = Path(__file__).resolve().parents[4] / "examples" / "train" + + +# --------------------------------------------------------------------------- # +# from_sources: the pure, argv-free core +# --------------------------------------------------------------------------- # + + +def test_from_sources_is_pure_and_flag_beats_default(): + cfg = TrainConfig.from_sources( + cli={"verifier_name_or_path": "some-model", "lr": 0.5}, argv=["train.py"] + ) + flat = cfg.flatten() + assert flat["verifier_name_or_path"] == "some-model" + assert flat["lr"] == 0.5 # flag wins + assert flat["epochs"] == 20 # untouched -> schema default + + +def test_from_sources_records_winning_layer_in_memory(): + cfg = TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", "lr": 0.5}, argv=["train.py"] + ) + assert cfg._provenance["lr"] == "flag" + assert cfg._provenance["epochs"] == "default" + + +def test_from_sources_raises_instead_of_exiting(): + # checkpoint_freq=1.5 fails the schema validator (>1 must be a whole number). + with pytest.raises(ValidationError): + TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", "checkpoint_freq": 1.5}, + argv=["train.py"], + ) + + +def test_from_sources_rejects_float16_dtype(): + # float16 needs gradient scaling, which the trainer does not implement; the + # dtype validator rejects it at config-construction time. + with pytest.raises(ValidationError): + TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", "hidden_states_dtype": "float16"}, + argv=["train.py"], + ) + + +def test_from_sources_rejects_non_dtype_torch_attr(): + # hasattr(torch, "nn") is True but it is not a dtype; the validator must + # reject it at config time rather than let it fail later as an opaque + # autocast error. + with pytest.raises(ValidationError): + TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", "hidden_states_dtype": "nn"}, + argv=["train.py"], + ) + + +def test_from_sources_missing_required_names_the_flag(): + # The verifier is required, but the requirement is enforced post-build (so a + # --config can supply it), naming the flag when nothing does. + with pytest.raises(ConfigError, match="--verifier-name-or-path"): + TrainConfig.from_sources(cli={}, argv=["train.py"]) + + +def test_from_sources_rejects_draft_init_conflict(): + # The conflict is decided from the in-memory winning-layer record: num_layers + # won by the flag layer, so it conflicts with --from-pretrained. + with pytest.raises(ConfigError, match="from-pretrained"): + TrainConfig.from_sources( + cli={ + "verifier_name_or_path": "m", + "from_pretrained": "ckpt", + "num_layers": 2, + }, + argv=["train.py"], + ) + + +# --------------------------------------------------------------------------- # +# from_sources: the YAML layer and full flag > yaml > default precedence +# --------------------------------------------------------------------------- # + + +def _write(tmp_path: Path, text: str) -> str: + path = tmp_path / "run.yaml" + path.write_text(text) + return str(path) + + +def test_yaml_beats_default_and_flag_beats_yaml(tmp_path): + # lr: set in all three layers -> flag wins; epochs: only YAML -> YAML beats + # default; weight_decay: nowhere -> schema default. + config_path = _write( + tmp_path, + "train:\n optimizer:\n lr: 0.2\n trainer:\n epochs: 7\n", + ) + cfg = TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", "lr": 0.9}, + config_path=config_path, + argv=["train.py"], + ) + flat = cfg.flatten() + assert flat["lr"] == 0.9 + assert cfg._provenance["lr"] == "flag" + assert flat["epochs"] == 7 + assert cfg._provenance["epochs"] == "yaml" + assert flat["weight_decay"] == 0.01 + assert cfg._provenance["weight_decay"] == "default" + + +def test_flag_and_yaml_merge_within_the_same_group(tmp_path): + # The crux of the precedence engine: two fields of the SAME group arrive from + # different layers -- lr from the flag, weight_decay from YAML -- and both must + # survive the partial-update merge (neither layer's optimizer blob clobbers the + # other's field). + config_path = _write(tmp_path, "train:\n optimizer:\n weight_decay: 0.5\n") + cfg = TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", "lr": 0.9}, + config_path=config_path, + argv=["train.py"], + ) + flat = cfg.flatten() + assert flat["lr"] == 0.9 + assert cfg._provenance["lr"] == "flag" + assert flat["weight_decay"] == 0.5 + assert cfg._provenance["weight_decay"] == "yaml" + + +def test_bare_top_level_mapping_loads(tmp_path): + config_path = _write(tmp_path, "verifier:\n verifier_name_or_path: from-bare\n") + cfg = TrainConfig.from_sources(cli={}, config_path=config_path, argv=["train.py"]) + assert cfg.flatten()["verifier_name_or_path"] == "from-bare" + assert cfg._provenance["verifier_name_or_path"] == "yaml" + + +@pytest.mark.parametrize("contents", ["", "---\n", "# comment\n"]) +def test_empty_or_comment_only_config_is_a_noop(tmp_path, contents): + # An empty / comment-only --config contributes nothing; the flag still applies. + config_path = _write(tmp_path, contents) + cfg = TrainConfig.from_sources( + cli={"verifier_name_or_path": "m"}, + config_path=config_path, + argv=["train.py"], + ) + assert cfg.flatten()["verifier_name_or_path"] == "m" + + +def test_train_block_wins_over_sibling_stage_keys(tmp_path): + # A file authored for the future pipeline trains today using only `train:`. + config_path = _write( + tmp_path, + "prepare_data:\n data_path: ignored\n" + "launch_vllm:\n port: 8200\n" + "train:\n data:\n data_path: from-train\n", + ) + cfg = TrainConfig.from_sources( + cli={"verifier_name_or_path": "m"}, config_path=config_path, argv=["train.py"] + ) + assert cfg.flatten()["data_path"] == "from-train" + + +def test_unknown_key_warns_and_is_ignored(tmp_path): + config_path = _write( + tmp_path, + "train:\n optimizer:\n lr: 0.3\n nonsense_knob: 1\n bogus_group: {}\n", + ) + with pytest.warns(UserWarning, match="unrecognised keys"): + cfg = TrainConfig.from_sources( + cli={"verifier_name_or_path": "m"}, + config_path=config_path, + argv=["train.py"], + ) + assert cfg.flatten()["lr"] == 0.3 # the valid sibling key still applies + + +def test_yaml_only_draft_init_conflict_is_rejected(tmp_path): + # The conflict lives entirely in the file: from_pretrained plus a shaping flag. + config_path = _write( + tmp_path, + "train:\n draft:\n from_pretrained: ckpt\n num_layers: 4\n", + ) + with pytest.raises(ConfigError, match="from-pretrained"): + TrainConfig.from_sources(cli={}, config_path=config_path, argv=["train.py"]) + + +def test_split_draft_init_conflict_is_rejected(tmp_path): + # The conflict straddles both layers: the init source in YAML, the shaping flag + # on the CLI. It must be caught identically to a same-layer conflict, because + # "explicitly provided" is decided from the layer-agnostic winning-layer record. + config_path = _write(tmp_path, "train:\n draft:\n from_pretrained: ckpt\n") + with pytest.raises(ConfigError, match="from-pretrained"): + TrainConfig.from_sources( + cli={"num_layers": 4}, config_path=config_path, argv=["train.py"] + ) + + +def test_split_draft_init_conflict_reversed_is_rejected(tmp_path): + # Same conflict, layers swapped: shaping flag in YAML, init source on the CLI. + config_path = _write(tmp_path, "train:\n draft:\n num_layers: 4\n") + with pytest.raises(ConfigError, match="from-pretrained"): + TrainConfig.from_sources( + cli={"from_pretrained": "ckpt"}, config_path=config_path, argv=["train.py"] + ) + + +@pytest.mark.parametrize( + "cli", + [ + {"from_pretrained": "ckpt"}, # one init source alone + {"draft_config": "decoder"}, # the other init source alone + {"num_layers": 4, "draft_arch": "qwen3"}, # shaping flags alone + ], +) +def test_single_draft_init_source_is_accepted(cli): + # Exactly one of the three draft-init mechanisms is not a conflict. + cfg = TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", **cli}, argv=["train.py"] + ) + assert cfg.speculator_type == "eagle3" + + +def test_every_flag_is_settable_via_yaml(tmp_path): + # Each schema dest round-trips through nest_flat -> YAML -> resolved config. + probes = { + dest: value + for dest, value in TrainConfig().flatten().items() + if isinstance(value, (int, float, str)) and not isinstance(value, bool) + } + probes["verifier_name_or_path"] = "m" # required: give it a non-empty value + yaml_text = yaml.safe_dump({"train": nest_flat(probes)}) + config_path = _write(tmp_path, yaml_text) + cfg = TrainConfig.from_sources(cli={}, config_path=config_path, argv=["train.py"]) + for dest in probes: + assert cfg._provenance[dest] == "yaml", dest + + +# --------------------------------------------------------------------------- # +# from_sources: algorithm-block mismatch warning +# --------------------------------------------------------------------------- # + + +def test_mismatched_algorithm_block_warns_via_flag(): + # markov_rank belongs to the dspark group; an eagle3 run ignores it, so warn. + with pytest.warns(UserWarning, match="does not use the 'dspark'"): + TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", "markov_rank": 128}, argv=["train.py"] + ) + + +def test_mismatched_algorithm_block_warns_via_yaml(tmp_path): + # The mismatched block is set entirely in YAML: warns identically to the flag. + config_path = _write(tmp_path, "train:\n dflash:\n block_size: 16\n") + with pytest.warns(UserWarning, match="does not use the 'dflash'"): + TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", "speculator_type": "peagle"}, + config_path=config_path, + argv=["train.py"], + ) + + +def test_matching_algorithm_block_does_not_warn(): + # dspark consumes both the dflash and dspark groups: setting either is fine. + with warnings.catch_warnings(): + warnings.simplefilter("error") + TrainConfig.from_sources( + cli={ + "verifier_name_or_path": "m", + "speculator_type": "dspark", + "block_size": 16, + "markov_rank": 128, + }, + argv=["train.py"], + ) + + +def test_default_algorithm_block_does_not_warn(): + # An untouched (all-default) mismatched group is not "set", so it stays silent. + with warnings.catch_warnings(): + warnings.simplefilter("error") + TrainConfig.from_sources( + cli={"verifier_name_or_path": "m", "speculator_type": "eagle3"}, + argv=["train.py"], + ) + + +# --------------------------------------------------------------------------- # +# resolve: the impure CLI boundary +# --------------------------------------------------------------------------- # + + +def test_resolve_missing_required_names_the_flag(capsys): + # No verifier anywhere: the post-build required check surfaces through + # parser.error as a clean exit(2) whose message still names the flag. + with pytest.raises(SystemExit) as exc: + TrainConfig.resolve([]) + assert exc.value.code == 2 + assert "--verifier-name-or-path" in capsys.readouterr().err + + +def test_resolve_verifier_from_config_only_succeeds(tmp_path): + # The reproducibility contract: reloading a run.yaml that supplies the verifier + # only inside the file must succeed even with no verifier flag on the CLI -- + # the required check runs post-build, after the YAML layer is merged. + config = tmp_path / "run.yaml" + config.write_text("train:\n verifier:\n verifier_name_or_path: from-yaml\n") + cfg = TrainConfig.resolve(["--config", str(config)]) + assert cfg.flatten()["verifier_name_or_path"] == "from-yaml" + + +def test_resolve_config_error_exits_cleanly(capsys): + with pytest.raises(SystemExit) as exc: + TrainConfig.resolve( + ["--verifier-name-or-path", "m", "--checkpoint-freq", "1.5"] + ) + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "invalid configuration" in err + assert "Traceback" not in err + + +def test_resolve_float16_dtype_exits_cleanly(capsys): + # The dtype validator's float16 rejection surfaces as a clean exit(2). + with pytest.raises(SystemExit) as exc: + TrainConfig.resolve( + ["--verifier-name-or-path", "m", "--hidden-states-dtype", "float16"] + ) + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "invalid configuration" in err + assert "Traceback" not in err + + +def test_resolve_malformed_config_exits_cleanly(tmp_path, capsys): + bad = tmp_path / "run.yaml" + bad.write_text("train:\n optimizer:\n lr: [unclosed\n") + with pytest.raises(SystemExit) as exc: + TrainConfig.resolve(["--verifier-name-or-path", "m", "--config", str(bad)]) + assert exc.value.code == 2 + err = capsys.readouterr().err + assert str(bad) in err + assert "Traceback" not in err + + +def test_resolve_config_not_a_mapping_exits_cleanly(tmp_path, capsys): + bad = tmp_path / "run.yaml" + bad.write_text("- just\n- a\n- list\n") + with pytest.raises(SystemExit) as exc: + TrainConfig.resolve(["--verifier-name-or-path", "m", "--config", str(bad)]) + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "top-level mapping" in err + assert "Traceback" not in err + + +def test_resolve_flag_overrides_config_file(tmp_path): + config = tmp_path / "run.yaml" + config.write_text("train:\n optimizer:\n lr: 0.2\n") + cfg = TrainConfig.resolve( + ["--verifier-name-or-path", "m", "--config", str(config), "--lr", "0.5"] + ) + assert cfg.flatten()["lr"] == 0.5 + + +def test_help_is_grouped_by_concern(): + titles = {group.title for group in build_parser()._action_groups} + assert {"general", "verifier", "draft", "data", "optimizer", "mtp"} <= titles + + +# --------------------------------------------------------------------------- # +# Example-recipe back-compat suite +# --------------------------------------------------------------------------- # + + +def _config_block(lines: list[str]) -> list[str]: + """The recipe's ``# ==== Configuration ====`` variable-assignment block.""" + start = next(i for i, line in enumerate(lines) if "Configuration" in line) + block: list[str] = [] + for line in lines[start + 1 :]: + if re.match(r"^#\s*=+\s*$", line.strip()): + break + block.append(line) + return block + + +def _train_invocation_args(lines: list[str]) -> str: + """The argument text after ``scripts/train.py``, with continuations joined.""" + start = next(i for i, line in enumerate(lines) if "scripts/train.py" in line) + block: list[str] = [] + for line in lines[start:]: + block.append(line) + if not line.rstrip().endswith("\\"): + break + joined = " ".join(line.rstrip().rstrip("\\").strip() for line in block) + return joined.split("scripts/train.py", 1)[1] + + +def _recipe_argv(path: Path) -> list[str]: + """Extract the argv a recipe threads into ``train.py`` via a real shell.""" + lines = path.read_text().splitlines() + script = ( + "\n".join(_config_block(lines)) + + '\n__emit() { printf "%s\\n" "$@"; }\n__emit ' + + _train_invocation_args(lines) + + "\n" + ) + out = subprocess.run( # noqa: S603 + [shutil.which("bash") or "bash", "-c", script], + capture_output=True, + text=True, + check=True, + ) + return out.stdout.splitlines() + + +RECIPES: dict[str, dict] = { + "dflash_qwen3_8b_sharegpt_online_5k.sh": { + "verifier_name_or_path": "Qwen/Qwen3-8B", + "data_path": "./output/dflash_qwen3_8b_sharegpt", + "vllm_endpoint": "http://localhost:8000/v1", + "save_path": "./output/dflash_qwen3_8b_sharegpt/checkpoints", + "draft_vocab_size": 32000, + "epochs": 5, + "lr": 3e-4, + "total_seq_len": 8192, + "speculator_type": "dflash", + "block_size": 8, + "max_anchors": 3072, + "num_layers": 5, + "target_layer_ids": [2, 18, 33], + "on_missing": "generate", + "on_generate": "delete", + }, + "dspark_qwen3_0_6b_sharegpt_online.sh": { + "verifier_name_or_path": "Qwen/Qwen3-0.6B", + "data_path": "./output/dspark_qwen3_0_6b_sharegpt", + "vllm_endpoint": "http://localhost:8000/v1", + "save_path": "./output/dspark_qwen3_0_6b_sharegpt/checkpoints", + "draft_vocab_size": 32000, + "epochs": 5, + "lr": 3e-4, + "total_seq_len": 4096, + "speculator_type": "dspark", + "block_size": 8, + "max_anchors": 3072, + "num_layers": 3, + "target_layer_ids": [2, 14, 25], + "markov_rank": 256, + "markov_head_type": "vanilla", + "enable_confidence_head": True, + "confidence_head_with_markov": True, + "loss_fn": '{"ce": 0.1, "tv": 0.9}', + "confidence_head_alpha": 1.0, + "on_missing": "generate", + "on_generate": "delete", + }, + "eagle3_llama3_8b_ultrachat_offline_5k.sh": { + "verifier_name_or_path": "meta-llama/Llama-3.1-8B-Instruct", + "data_path": "./output", + "hidden_states_path": "./output/hidden_states", + "save_path": "./output/checkpoints", + "draft_vocab_size": 32000, + "epochs": 5, + "lr": 1e-4, + "total_seq_len": 8192, + "on_missing": "raise", + }, + "eagle3_qwen3_8b_sharegpt_online_5k.sh": { + "verifier_name_or_path": "Qwen/Qwen3-8B", + "data_path": "./output", + "vllm_endpoint": "http://localhost:8000/v1", + "save_path": "./output/checkpoints", + "draft_vocab_size": 32000, + "epochs": 5, + "lr": 1e-4, + "total_seq_len": 8192, + "on_missing": "generate", + "on_generate": "delete", + }, + "mtp_qwen3_5_9b_gsm8k_online.sh": { + "verifier_name_or_path": "Qwen/Qwen3.5-9B", + "data_path": "./output", + "vllm_endpoint": "http://localhost:8000/v1", + "save_path": "./output/checkpoints", + "speculator_type": "mtp", + "num_speculative_steps": 3, + "target_layer_ids": [32], + "step_weight_beta": 0.6, + "epochs": 3, + "lr": 1e-4, + "total_seq_len": 8192, + "on_missing": "generate", + "on_generate": "delete", + }, + "peagle_qwen3_8b_sharegpt_online_5k.sh": { + "verifier_name_or_path": "Qwen/Qwen3-8B", + "data_path": "./output/peagle_qwen3_8b_sharegpt", + "vllm_endpoint": "http://localhost:8108/v1", + "hidden_states_path": "./output/peagle_qwen3_8b_sharegpt/hidden_states", + "save_path": "./output/peagle_qwen3_8b_sharegpt/checkpoints", + "epochs": 5, + "lr": 6e-4, + "total_seq_len": 4096, + "speculator_type": "peagle", + "num_layers": 4, + "num_depths": 4, + "down_sample_ratio": 0.7, + "down_sample_ratio_min": 0.2, + "norm_before_residual": False, + "scheduler_type": "cosine", + "on_missing": "generate", + "on_generate": "delete", + }, +} + + +@pytest.mark.skipif(shutil.which("bash") is None, reason="requires bash") +@pytest.mark.parametrize(("recipe", "expected"), RECIPES.items(), ids=list(RECIPES)) +def test_recipe_flags_resolve_unchanged(recipe: str, expected: dict): + flat = TrainConfig.resolve(_recipe_argv(EXAMPLES / recipe)).flatten() + for dest, value in expected.items(): + assert flat[dest] == value, dest diff --git a/tests/unit/train/config/test_schema.py b/tests/unit/train/config/test_schema.py new file mode 100644 index 000000000..7d733cd93 --- /dev/null +++ b/tests/unit/train/config/test_schema.py @@ -0,0 +1,76 @@ +"""Data-model seam tests for the train config schema. + +These exercise the schema at its two pure seams -- ``flatten()`` and +``from_flat()`` -- not CLI parsing or YAML. Backward compatibility against the +real parser is proven separately by the example-recipe tests, not a golden +``vars(args)`` snapshot here. +""" + +import pytest + +from speculators.train.config import TrainConfig +from speculators.train.config.schema import ( + CONFIG_DESTS, + DraftArgs, + OptimizerArgs, +) + + +def test_constructs_from_defaults(): + # The whole point of the schema seam: a config exists with no inputs. + TrainConfig() + + +def test_flatten_covers_exactly_the_schema_fields(): + # flatten() emits every schema dest and nothing else; consumers bind the flat + # dict by name (**kwargs / args.), so the key set is the contract. + flat = TrainConfig().flatten() + assert set(flat) == CONFIG_DESTS + # Order is deterministic (declaration order) so the run.yaml dump stays stable. + assert list(flat) == list(TrainConfig(speculator_type="dflash").flatten()) + + +def test_flatten_resolves_eagle3_derived_defaults(): + # Mirrors the tail of the pre-refactor parse_args for the default (eagle3) run. + flat = TrainConfig().flatten() + assert flat["draft_arch"] == "llama" + assert flat["norm_before_fc"] is True + assert flat["norm_output"] is True + assert flat["muon_lr"] == pytest.approx(10 * flat["lr"]) + + +def test_flatten_resolves_non_eagle3_derived_defaults(): + flat = TrainConfig(speculator_type="dflash").flatten() + assert flat["draft_arch"] == "qwen3" + assert flat["norm_before_fc"] is False + assert flat["norm_output"] is False + + +def test_from_flat_inverts_flatten(): + cfg = TrainConfig( + speculator_type="dspark", + draft=DraftArgs(num_layers=4, full_attention_indices=[2, 18, 33]), + optimizer=OptimizerArgs(lr=3e-4), + ) + assert TrainConfig.from_flat(cfg.flatten()) == cfg + + +def test_from_flat_default_round_trip(): + cfg = TrainConfig() + assert TrainConfig.from_flat(cfg.flatten()) == cfg + + +def test_from_flat_ignores_non_config_keys(): + flat = TrainConfig().flatten() + flat["config"] = "run.yaml" + flat["dump_config"] = True + recovered = TrainConfig.from_flat(flat) + assert recovered == TrainConfig() + + +def test_from_flat_accepts_partial_working_dict(): + recovered = TrainConfig.from_flat({"lr": 5e-4, "num_layers": 6}) + assert recovered.optimizer.lr == pytest.approx(5e-4) + assert recovered.draft.num_layers == 6 + # Untouched fields fall back to their schema defaults. + assert recovered.trainer.epochs == 20 diff --git a/tests/unit/train/test_batch_sampler_balance.py b/tests/unit/train/test_batch_sampler_balance.py new file mode 100644 index 000000000..8161a2b8e --- /dev/null +++ b/tests/unit/train/test_batch_sampler_balance.py @@ -0,0 +1,101 @@ +"""Balance and partition guarantees for MultipackDistributedBatchSamplerV2.""" + +import numpy as np +import pytest + +from speculators.train.distributed_batch_sampler import ( + MultipackDistributedBatchSamplerV2, +) + +MAX_LEN = 8192 + + +def _skewed_lengths(n: int = 4000, seed: int = 0) -> np.ndarray: + rng = np.random.default_rng(seed) + return np.clip(rng.lognormal(mean=7.6, sigma=0.9, size=n).astype(int), 64, MAX_LEN) + + +def _shards(lengths: np.ndarray, replicas: int) -> list[list[np.ndarray]]: + return [ + list( + iter( + MultipackDistributedBatchSamplerV2( + batch_max_length=MAX_LEN, + lengths=lengths, + num_replicas=replicas, + rank=rank, + ) + ) + ) + for rank in range(replicas) + ] + + +@pytest.mark.parametrize("replicas", [2, 3, 4, 8]) +def test_sample_counts_are_balanced_across_ranks(replicas): + lengths = _skewed_lengths() + counts = [sum(len(b) for b in batches) for batches in _shards(lengths, replicas)] + + assert min(counts) > 0 + assert max(counts) / min(counts) < 1.15, f"sample counts unbalanced: {counts}" + + +@pytest.mark.parametrize("replicas", [2, 3, 4, 8]) +def test_token_counts_stay_balanced_across_ranks(replicas): + lengths = _skewed_lengths() + tokens = [ + sum(int(lengths[b].sum()) for b in batches) + for batches in _shards(lengths, replicas) + ] + + assert max(tokens) / min(tokens) < 1.15, f"token counts unbalanced: {tokens}" + + +@pytest.mark.parametrize("replicas", [2, 3, 4, 8]) +def test_ranks_form_a_disjoint_partition(replicas): + lengths = _skewed_lengths() + shards = _shards(lengths, replicas) + + seen: list[set[int]] = [] + for batches in shards: + idxs = [int(i) for b in batches for i in b] + assert len(idxs) == len(set(idxs)), "duplicate index within a single rank" + seen.append(set(idxs)) + + for i in range(replicas): + for j in range(i + 1, replicas): + assert not (seen[i] & seen[j]), f"ranks {i}/{j} share samples" + + +@pytest.mark.parametrize("replicas", [2, 3, 4, 8]) +def test_equal_batch_count_per_rank(replicas): + counts = {len(batches) for batches in _shards(_skewed_lengths(), replicas)} + assert len(counts) == 1, f"ranks disagree on batch count: {counts}" + + +@pytest.mark.parametrize("replicas", [2, 3]) +def test_batches_respect_token_budget(replicas): + lengths = _skewed_lengths() + for batches in _shards(lengths, replicas): + for b in batches: + assert int(lengths[b].sum()) <= MAX_LEN + + +def test_rotation_actually_moves_the_largest_sample_off_rank_zero(): + lengths = _skewed_lengths() + replicas = 3 + shards = _shards(lengths, replicas) + nbatches = len(shards[0]) + + owners = set() + for i in range(min(nbatches, 30)): + best_rank, best_len = None, -1 + for rank in range(replicas): + if len(shards[rank][i]) == 0: + continue + m = int(lengths[shards[rank][i]].max()) + if m > best_len: + best_rank, best_len = rank, m + owners.add(best_rank) + + assert len(owners) > 1, "largest sample always lands on the same rank" diff --git a/tests/unit/train/test_cli_args.py b/tests/unit/train/test_cli_args.py index 58a11bfd1..316d17c22 100644 --- a/tests/unit/train/test_cli_args.py +++ b/tests/unit/train/test_cli_args.py @@ -1,18 +1,20 @@ """Tests for CLI arguments.""" -from scripts.train import parse_args +import argparse + +import pytest + from speculators.models.dflash.core import DFlashDraftModel from speculators.models.dspark.core import DSparkDraftModel from speculators.models.eagle3.core import Eagle3DraftModel from speculators.models.metrics import ce_loss, kl_div_loss, tv_loss_fused_or_eager from speculators.models.peagle.core import PEagleDraftModel +from speculators.train.config import TrainConfig -def _parse(monkeypatch, extra: list[str]): - monkeypatch.setattr( - "sys.argv", ["train.py", "--verifier-name-or-path", "dummy"] + extra - ) - return parse_args() +def _parse(monkeypatch, extra: list[str]) -> argparse.Namespace: + cfg = TrainConfig.resolve(["--verifier-name-or-path", "dummy", *extra]) + return argparse.Namespace(**cfg.flatten()) # --------------------------------------------------------------------------- @@ -169,3 +171,23 @@ def test_no_norm_before_fc_flag(monkeypatch): def test_no_norm_output_flag(monkeypatch): args = _parse(monkeypatch, ["--no-norm-output"]) assert args.norm_output is False + + +# --------------------------------------------------------------------------- +# --max-steps +# --------------------------------------------------------------------------- + + +def test_max_steps_default_is_none(monkeypatch): + args = _parse(monkeypatch, []) + assert args.max_steps is None + + +def test_max_steps_explicit(monkeypatch): + args = _parse(monkeypatch, ["--max-steps", "15"]) + assert args.max_steps == 15 + + +def test_max_steps_rejects_non_positive(monkeypatch): + with pytest.raises(SystemExit): + _parse(monkeypatch, ["--max-steps", "0"]) diff --git a/tests/unit/train/test_data.py b/tests/unit/train/test_data.py index e76413551..6e10e3d26 100644 --- a/tests/unit/train/test_data.py +++ b/tests/unit/train/test_data.py @@ -469,8 +469,7 @@ def test_dataset_fallback_when_sample_lengths_json_malformed(tmp_path: Path): assert len(dataset.approx_lengths) == 2 -def test_arrow_dataset_default_split_ratio_does_not_crash(tmp_path: Path): - """ArrowDataset with default split_ratio=1.0 should support indexing.""" +def test_arrow_dataset_default_train_ratio_does_not_crash(tmp_path: Path): ds = Dataset.from_dict( { "input_ids": [[1, 2, 3]], diff --git a/tests/unit/train/test_draft_config_init.py b/tests/unit/train/test_draft_config_init.py index d2bbb33de..cee249229 100644 --- a/tests/unit/train/test_draft_config_init.py +++ b/tests/unit/train/test_draft_config_init.py @@ -11,6 +11,7 @@ ``--draft-config`` is mutually exclusive with the decoder-shaping flags. """ +import argparse import json from pathlib import Path from types import SimpleNamespace @@ -22,16 +23,16 @@ from transformers.models.qwen3.configuration_qwen3 import Qwen3Config from scripts.train import ( - DECODER_SHAPING_FLAGS, _build_from_config_only, build_draft_model, create_transformer_layer_config, load_draft_transformer_layer_config, - parse_args, ) from speculators import SpeculatorsConfig, VerifierConfig from speculators.models.eagle3 import Eagle3DraftModel, Eagle3SpeculatorConfig from speculators.proposals.greedy import GreedyTokenProposalConfig +from speculators.train.config import TrainConfig +from speculators.train.config.resolution import DECODER_SHAPING_FLAGS from speculators.utils.loading import is_config_only_dir # --------------------------------------------------------------------------- @@ -51,11 +52,9 @@ } -def _parse(monkeypatch, extra: list[str]): - monkeypatch.setattr( - "sys.argv", ["train.py", "--verifier-name-or-path", "dummy"] + extra - ) - return parse_args() +def _parse(monkeypatch, extra: list[str]) -> argparse.Namespace: + cfg = TrainConfig.resolve(["--verifier-name-or-path", "dummy", *extra]) + return argparse.Namespace(**cfg.flatten()) def _make_eagle3_config(verifier_name_or_path: str | None = "some-verifier"): diff --git a/tests/unit/train/test_split_boundaries.py b/tests/unit/train/test_split_boundaries.py new file mode 100644 index 000000000..acb24b862 --- /dev/null +++ b/tests/unit/train/test_split_boundaries.py @@ -0,0 +1,92 @@ +"""The train and val splits must be exactly complementary for any train_ratio.""" + +from pathlib import Path +from typing import Literal + +import pytest +from datasets import Dataset + +from speculators.train.data import ArrowDataset + + +def _dataset(tmp_path: Path, n: int) -> str: + ds = Dataset.from_dict( + { + "input_ids": [[i, i + 1, i + 2] for i in range(n)], + "loss_mask": [[1, 1, 1]] * n, + "seq_len": [3] * n, + } + ) + path = tmp_path / "data" + ds.save_to_disk(str(path)) + (path / "hidden_states").mkdir() + return str(path) + + +def _split(path: str, ratio: float, split: Literal["train", "val"]) -> ArrowDataset: + return ArrowDataset( + max_len=128, + datapath=path, + on_missing="skip", + train_ratio=ratio, + split=split, + ) + + +@pytest.mark.parametrize("ratio", [0.1, 0.2, 0.25, 0.3, 0.5, 0.7, 0.8, 0.9, 0.95]) +@pytest.mark.parametrize("n", [1000, 12501, 100000]) +def test_splits_are_exactly_complementary(tmp_path, ratio, n): + path = _dataset(tmp_path, n) + train = _split(path, ratio, "train") + val = _split(path, ratio, "val") + + assert len(train) + len(val) == n, "splits must partition the dataset" + assert len(train) == int(n * ratio) + assert val.start_file_idx == len(train), "val must begin where train ends" + assert len(train) > 0 + assert len(val) > 0 + + +@pytest.mark.parametrize("ratio", [0.2, 0.9]) +def test_no_row_appears_in_both_splits(tmp_path, ratio): + n = 12501 + path = _dataset(tmp_path, n) + train = _split(path, ratio, "train") + val = _split(path, ratio, "val") + + train_ids = {tuple(train.data[i]["input_ids"]) for i in range(len(train))} + val_ids = {tuple(val.data[i]["input_ids"]) for i in range(len(val))} + assert not (train_ids & val_ids), "a row is in both train and val" + + +def test_default_takes_whole_dataset(tmp_path): + path = _dataset(tmp_path, 100) + assert len(_split(path, 1.0, "train")) == 100 + + +@pytest.mark.parametrize("ratio", [0.0, -0.1, 1.5]) +def test_rejects_out_of_range_ratio(tmp_path, ratio): + path = _dataset(tmp_path, 100) + with pytest.raises(ValueError, match="train_ratio must be in"): + _split(path, ratio, "train") + + +def test_rejects_val_split_with_no_val_data(tmp_path): + path = _dataset(tmp_path, 100) + with pytest.raises(ValueError, match="leaves no validation split"): + _split(path, 1.0, "val") + + +@pytest.mark.parametrize("ratio", [0.1, 0.5, 0.9]) +def test_rejects_empty_train_from_small_dataset(tmp_path, ratio): + path = _dataset(tmp_path, 1) + with pytest.raises(ValueError, match="train split is empty"): + _split(path, ratio, "train") + + +def test_small_dataset_both_splits_nonempty(tmp_path): + path = _dataset(tmp_path, 2) + train = _split(path, 0.5, "train") + val = _split(path, 0.5, "val") + assert len(train) == 1 + assert len(val) == 1 diff --git a/tests/unit/train/test_trainer_scheduler.py b/tests/unit/train/test_trainer_scheduler.py index c17ed642b..2aa47a541 100644 --- a/tests/unit/train/test_trainer_scheduler.py +++ b/tests/unit/train/test_trainer_scheduler.py @@ -1,6 +1,6 @@ import pytest -from scripts.train import parse_args +from speculators.train.config import TrainConfig from speculators.train.trainer import ( TrainerConfig, _resolve_scheduler_steps, @@ -61,13 +61,10 @@ def test_scheduler_warmup_ratio_must_be_between_zero_and_one(): _resolve_scheduler_steps(make_config(scheduler_warmup_ratio=1.1), 20) -def test_scheduler_type_rejects_unsupported_values(monkeypatch): +def test_scheduler_type_rejects_unsupported_values(): # --verifier-name-or-path is supplied so the only parse failure is the rejected # --scheduler-type choice (not the missing required verifier arg). - monkeypatch.setattr( - "sys.argv", - ["train.py", "--verifier-name-or-path", "x", "--scheduler-type", "constant"], - ) - with pytest.raises(SystemExit): - parse_args() + TrainConfig.resolve( + ["--verifier-name-or-path", "x", "--scheduler-type", "constant"] + ) diff --git a/tests/unit/train/test_val_sync.py b/tests/unit/train/test_val_sync.py new file mode 100644 index 000000000..d94214096 --- /dev/null +++ b/tests/unit/train/test_val_sync.py @@ -0,0 +1,127 @@ +"""Periodic re-synchronisation inside the validation loop.""" + +from types import SimpleNamespace +from typing import cast +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from speculators.train.trainer import Trainer, TrainerConfig + +# --------------------------------------------------------------------------- +# Unit tests for _maybe_val_sync helper +# --------------------------------------------------------------------------- + + +def _barrier_indices(n_batches: int, *, is_distributed: bool, interval: int) -> list: + stand_in = cast("Trainer", SimpleNamespace(is_distributed=is_distributed)) + hit = [] + with ( + patch("speculators.train.trainer._VAL_SYNC_INTERVAL", interval), + patch("speculators.train.trainer.dist.barrier") as barrier, + ): + for i in range(n_batches): + before = barrier.call_count + Trainer._maybe_val_sync(stand_in, i) + if barrier.call_count > before: + hit.append(i) + return hit + + +def test_barriers_on_interval_boundaries_only(): + assert _barrier_indices(200, is_distributed=True, interval=50) == [50, 100, 150] + + +def test_never_barriers_on_the_first_batch(): + assert 0 not in _barrier_indices(10, is_distributed=True, interval=1) + + +def test_no_barrier_when_single_process(): + assert _barrier_indices(200, is_distributed=False, interval=50) == [] + + +def test_interval_zero_disables_syncing(): + assert _barrier_indices(200, is_distributed=True, interval=0) == [] + + +@pytest.mark.parametrize("interval", [1, 7, 50, 500]) +def test_barrier_count_is_deterministic_for_a_given_batch_count(interval): + n = 1864 + first = _barrier_indices(n, is_distributed=True, interval=interval) + second = _barrier_indices(n, is_distributed=True, interval=interval) + assert first == second + assert len(first) == (n - 1) // interval + + +# --------------------------------------------------------------------------- +# Integration: _maybe_val_sync is actually called through val_epoch +# --------------------------------------------------------------------------- + + +def _fake_batch() -> dict[str, torch.Tensor]: + return { + "input_ids": torch.zeros(1, 4, dtype=torch.long), + "hidden_states": torch.zeros(1, 4, 16), + } + + +def _make_val_trainer(n_batches: int, *, is_distributed: bool) -> "Trainer": + model = MagicMock() + model.return_value = ( + None, + torch.tensor(0.5), + {"loss": torch.tensor(0.5)}, + ) + + batches = [_fake_batch() for _ in range(n_batches)] + loader = MagicMock() + loader.__iter__ = MagicMock(return_value=iter(batches)) + loader.__len__ = MagicMock(return_value=len(batches)) + loader.batch_sampler = MagicMock(spec=[]) + + trainer = Trainer.__new__(Trainer) + trainer.model = model + trainer.val_loader = loader + trainer.is_distributed = is_distributed + trainer.rank = 1 + trainer.local_rank = 0 + trainer.device_type = "cpu" + trainer.global_step = 0 + trainer.config = TrainerConfig( + lr=1e-3, + num_epochs=1, + save_path="/tmp", + hidden_states_dtype=torch.float32, + val_call_kwargs=None, + ) + return trainer + + +@pytest.mark.parametrize("n_batches", [120, 200]) +def test_val_epoch_calls_barrier_at_sync_interval(n_batches): + trainer = _make_val_trainer(n_batches, is_distributed=True) + with ( + patch("speculators.train.trainer._VAL_SYNC_INTERVAL", 50), + patch("speculators.train.trainer.dist.barrier") as barrier, + patch("speculators.train.trainer.dist.all_reduce"), + patch("speculators.train.trainer.dist.get_world_size", return_value=1), + ): + trainer.val_epoch(0) + assert barrier.call_count == (n_batches - 1) // 50 + + +def test_val_epoch_no_barrier_when_not_distributed(): + trainer = _make_val_trainer(120, is_distributed=False) + with ( + patch("speculators.train.trainer._VAL_SYNC_INTERVAL", 50), + patch("speculators.train.trainer.dist.barrier") as barrier, + ): + trainer.val_epoch(0) + barrier.assert_not_called() + + +def test_val_epoch_empty_loader_returns_empty_metrics(): + trainer = _make_val_trainer(0, is_distributed=False) + result = trainer.val_epoch(0) + assert result == {} diff --git a/tests/unit/utils/test_argparse_utils.py b/tests/unit/utils/test_argparse_utils.py deleted file mode 100644 index 1bd317e12..000000000 --- a/tests/unit/utils/test_argparse_utils.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Unit tests for the argparse helpers in the Speculators library.""" - -import argparse - -import pytest - -from speculators.utils.argparse_utils import explicitly_provided_dests - - -def _parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser() - parser.add_argument("--num-layers", type=int, default=3) - parser.add_argument("--draft-arch", type=str, default="llama") - parser.add_argument("--flag", action="store_true") - return parser - - -@pytest.mark.smoke -def test_returns_only_options_present_on_argv(monkeypatch): - monkeypatch.setattr("sys.argv", ["prog", "--num-layers", "5"]) - - provided = explicitly_provided_dests(_parser(), ["num_layers", "draft_arch"]) - - assert provided == {"num_layers"} - - -@pytest.mark.smoke -def test_value_equal_to_default_still_counts_as_provided(monkeypatch): - # Passing the default value explicitly must still be detected (the whole point: - # provided-based, not value-vs-default comparison). - monkeypatch.setattr("sys.argv", ["prog", "--draft-arch", "llama"]) - - provided = explicitly_provided_dests(_parser(), ["num_layers", "draft_arch"]) - - assert provided == {"draft_arch"} - - -@pytest.mark.smoke -def test_empty_when_nothing_provided(monkeypatch): - monkeypatch.setattr("sys.argv", ["prog"]) - - provided = explicitly_provided_dests(_parser(), ["num_layers", "draft_arch"]) - - assert provided == set() - - -@pytest.mark.smoke -def test_store_true_flag_detected(monkeypatch): - monkeypatch.setattr("sys.argv", ["prog", "--flag"]) - - provided = explicitly_provided_dests(_parser(), ["flag", "num_layers"]) - - assert provided == {"flag"}