diff --git a/examples/specdec_bench/run.py b/examples/specdec_bench/run.py index 9ed7ec44272..d9d8b832c59 100644 --- a/examples/specdec_bench/run.py +++ b/examples/specdec_bench/run.py @@ -15,6 +15,7 @@ import argparse import asyncio +import json import yaml from specdec_bench import datasets, metrics, models, runners @@ -58,6 +59,13 @@ } +def parse_runtime_params(value): + if value.lstrip().startswith("{"): + return json.loads(value) + with open(value) as f: + return yaml.safe_load(f) + + async def tqdm_gather(*fs, return_exceptions=False, **kwargs): if not return_exceptions: return await tqdm.gather(*fs, **kwargs) @@ -242,6 +250,7 @@ def run_simple(args): ) runner.clear_metrics() + model.stop() if __name__ == "__main__": @@ -318,10 +327,10 @@ def run_simple(args): ) parser.add_argument( "--runtime_params", - type=str, + type=parse_runtime_params, required=False, default=None, - help="Path to the runtime params yaml file", + help="Path to a runtime params YAML file or an inline JSON object", ) parser.add_argument( "--temperature", @@ -359,8 +368,9 @@ def run_simple(args): required=False, default=None, help=( - "DFlash block size (num_speculative_tokens). Use instead of --draft_length " - "for DFLASH: block_size = draft_length + 1." + "Draft block size for DFlash and DSpark. To draft N tokens " + "(num_speculative_tokens=N), use --block_size N+1 for DFlash and " + "--block_size N for DSpark." ), ) parser.add_argument( @@ -404,10 +414,7 @@ def run_simple(args): ) args = parser.parse_args() - if args.runtime_params is not None: - with open(args.runtime_params) as f: - args.runtime_params = yaml.safe_load(f) - else: + if args.runtime_params is None: args.runtime_params = {} if args.dataset is None: assert ( diff --git a/examples/specdec_bench/specdec_bench/models/vllm.py b/examples/specdec_bench/specdec_bench/models/vllm.py index d286d8fed6a..9c754fac7ab 100644 --- a/examples/specdec_bench/specdec_bench/models/vllm.py +++ b/examples/specdec_bench/specdec_bench/models/vllm.py @@ -119,21 +119,24 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs # (handled in speculative.py:562-573). specdec["model"] = draft_model_dir elif kwargs.get("speculative_algorithm") == "DFLASH": + # DFlash block size includes one anchor, so block 8 maps to 7 speculative tokens in vLLM. specdec = { "method": "dflash", "model": kwargs.get("draft_model_dir"), - "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 8), + "num_speculative_tokens": (kwargs.get("speculative_num_draft_tokens") or 8) - 1, } elif kwargs.get("speculative_algorithm") == "DSPARK": specdec = { "method": "dspark", "model": kwargs.get("draft_model_dir"), - "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 7), - "draft_sample_method": kwargs.get("draft_sample_method", "greedy"), + "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens") or 7, } elif kwargs.get("speculative_algorithm") == "NONE": specdec = None + if specdec is not None and kwargs.get("draft_sample_method") is not None: + specdec["draft_sample_method"] = kwargs["draft_sample_method"] + if specdec is None: num_speculative_tokens = 1 else: @@ -163,6 +166,7 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs temperature=sampling_kwargs.get("temperature", 1.0), top_p=sampling_kwargs.get("top_p", 1.0), top_k=sampling_kwargs.get("top_k", 0), + presence_penalty=sampling_kwargs.get("presence_penalty", 0.0), ) self.loop = asyncio.new_event_loop() asyncio.set_event_loop(self.loop) diff --git a/examples/speculative_decoding/QWEN3_8B_DSPARK_WALKTHROUGH.md b/examples/speculative_decoding/QWEN3_8B_DSPARK_WALKTHROUGH.md new file mode 100644 index 00000000000..9f1b859933b --- /dev/null +++ b/examples/speculative_decoding/QWEN3_8B_DSPARK_WALKTHROUGH.md @@ -0,0 +1,293 @@ +# End-to-End Walkthrough: Qwen3-8B DSpark (Data Synthesis → Streaming Training) + +A complete worked example of training a speculative-decoding drafter, from raw +prompts to an evaluated checkpoint. It uses **Qwen3-8B** as the target and +**DSpark** as the draft algorithm, driven end-to-end by the +[launcher](../../tools/launcher/). + +Qwen3-8B is small enough to run the whole flow on a handful of GPUs, so this +doubles as the recommended first run before scaling to a large MoE target — the +pipeline shape is identical, only node counts and a few model-specific fields +change. + +## Why these two steps + +**Data synthesis.** A drafter is trained to predict what the *target model* +would say. Off-the-shelf SFT corpora contain some other model's answers, so +training on them teaches the drafter the wrong distribution and acceptance +length suffers. Step 1 therefore keeps the prompts but regenerates every +assistant turn with the target model itself. + +**Streaming training.** DSpark trains against the target's hidden states. The +classic route dumps them to disk first, which is slow and enormous. Streaming +mode instead runs a live `vllm serve` and ships hidden states to the trainer +over NIXL RDMA — no dump, no round-trip. + +## What DSpark is + +DSpark = the DFlash backbone + a lightweight sequential (**Markov**) head + a +**confidence** head. The Markov head adds a prefix-dependent transition bias to +the backbone's base logits, which induces a causal block distribution and lets +the draft generate a block semi-autoregressively. The confidence head predicts +per-position acceptance. It trains with a three-term loss: + +```text +dflash_ce_loss_alpha * CE + dflash_l1_loss_alpha * TVD + dflash_confidence_head_alpha * BCE +``` + +Head architecture and loss weights live in +[`dspark.yaml`](../../modelopt_recipes/general/speculative_decoding/dspark.yaml); +the launcher YAMLs below only override data and schedule fields. + +## Prerequisites + +- A Slurm cluster with the launcher configured (see + [launcher docs](../../tools/launcher/docs/configuration.md)). +- `Qwen/Qwen3-8B` present under the launcher's `/hf-local/` mount. +- Two container images: a vLLM image (serve + training) and a TensorRT-LLM image + (dataset build). Both are pinned in the YAMLs. If your cluster's enroot cannot + pull `nvcr.io` (it resolves against Docker Hub and 401s, surfacing only as + `spank_pyxis.so: task_init() failed`), point `container:` at a local `.sqsh`. +- An `HF_TOKEN` with access to the prompt corpus — the dataset Step 1 ships with + is gated, and without a token it fails as `DatasetNotFoundError`, which reads + like a wrong name rather than a missing credential. + +Set your cluster environment once: + +```bash +export SLURM_HOST=$(hostname) # NOT localhost: the launcher stages files + # over SSH even when it submits locally +export SLURM_ACCOUNT= +export SLURM_PARTITION= +export SLURM_HF_LOCAL= +export SLURM_JOB_DIR= # must already exist +export NEMORUN_HOME=$PWD +export HF_TOKEN= + +mkdir -p $SLURM_JOB_DIR +cd tools/launcher && uv pip install -e . # launch.py imports modelopt_launcher +``` + +Pass `identity=` on every `launch.py` call (the key the staging SSH +uses). All commands below omit it for brevity. + +--- + +## Step 1 — Data synthesis + +Regenerate assistant turns with Qwen3-8B over a prompt corpus. + +```bash +cd tools/launcher +uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_synth.yaml --yes +``` + +[`hf_synth.yaml`](../../tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml) +runs a Slurm **array job**. Each task starts its own `vllm serve` on one node and +calls [`common/query.py`](../../tools/launcher/common/query.py) on its own shard +of the dataset. `query.py` replays each sample's user turns, generates a fresh +assistant reply, **discards the original assistant turns**, and writes +`shard_{id}.jsonl` plus a `shard_{id}.done` sentinel. + +Budget for this: a shard is tens of thousands of samples at 1-2 samples/s, so +**hours per shard**. Size the array and the job time limit accordingly. + +The knobs worth knowing: + +| Field | Meaning | +|---|---| +| `--data` | Input prompt corpus (HF id or local path). Ships pointing at `nvidia/Speculative-Decoding-Multilingual-Prompt-v2`. | +| `--save` | Output dir. Set `output_dir` in `global_vars`. | +| `--num-shards` / `array` | Shard count and the array range. Keep them consistent. | +| `--num-proc` | Per-worker request concurrency. Defaults to 32; not set in the YAML. | +| `--tensor-parallel-size` | Must match `gpus_per_node`. | + +Practical notes, mostly learned the hard way: + +- **Resume is per whole shard.** A shard is skipped only when *both* its + `.jsonl` and its `.done` sentinel exist; output is written once at the end. An + interrupted shard leaves nothing behind and restarts from zero. Write `--save` + to a persistent mount, not to scratch. +- **Concurrency vs. yield.** Pushing `--num-proc` very high can *lower* total + yield: requests start timing out under the stampede and those samples come out + user-only. If you see a yield dip, lower it before raising it. +- **Reasoning models need context headroom.** With `--max-model-len` too tight, + any prompt longer than the generation cap gets a 400 back and is dropped. This + loss is *systematic* — it removes the longest, hardest prompts and quietly + biases the corpus short. Keep `max-model-len >= max_prompt + max_output`. The + shipped 4096 is *not* enough headroom for this corpus at a 1024-token cap; + raise it, or accept the drop. Grep the job log for `Error code: 400` to size + it — those 400s are also mangled by an unrelated `APIStatusError` unpickling + bug that kills a result-handler thread without stopping the run. +- **Check yield before training.** Count records that actually have an assistant + turn. A corpus that is silently 60% prompt-only will train, and the result will + just be mysteriously bad. + +Prefer to skip synthesis on a first pass? `task_0` of the training YAML builds a +conversation set from public SFT data, and the pipeline runs standalone. Expect a +lower acceptance length — that gap is exactly what this step buys. + +--- + +## Step 2 — Streaming DSpark training + +```bash +cd tools/launcher +uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_streaming_dspark.yaml --yes +``` + +[`hf_streaming_dspark.yaml`](../../tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dspark.yaml) +is a three-task pipeline: + +| Task | Nodes | What it does | +|---|---|---| +| `task_0` | 1 | Build input conversations → `/scratchspace/data/train.jsonl` | +| `task_1` | 2 | Node 0 `vllm serve`, node 1 trainer; exports to `/scratchspace/export` | +| `task_2` | 1 | vLLM smoke test — acceptance length | + +**To train on your synthesized corpus** (the point of Step 1), edit the YAML: set +`skip: true` on `task_0`, and point `data.data_path` in `task_1` at the synthesis +`output_dir`: + +```yaml + task_0: + skip: true + ... + + task_1: + args: + ... + - data.data_path=/hf-local/modelopt/qwen3-8b-synth-v1 +``` + +CLI overrides use dotted paths for scalar fields, e.g.: + +```bash +uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_streaming_dspark.yaml \ + pipeline.task_1.slurm_config.nodes=4 --yes +``` + +Always preview the resolved config before submitting: + +```bash +uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_streaming_dspark.yaml --dryrun --yes -v +``` + +### The fields that actually matter + +**Capture ids.** `EAGLE_CAPTURE_IDS` selects which target layers the serve +captures. The draft consumes one evenly-spaced layer per draft layer, plus the +final hidden state — so this list is tied to +`dflash_architecture_config.num_hidden_layers` and must be recomputed if you +change draft depth. For Qwen3-8B (36 layers) with 5 draft layers, +`build_target_layer_ids(36,5) = [1,9,17,25,33]`; vLLM's ids are those **+1**, +plus the final layer: + +```text +EAGLE_CAPTURE_IDS: "[2,10,18,26,34,36]" +``` + +Getting the final id wrong is the classic silent failure: capturing the +second-to-last layer instead of the true final one trains fine, shows a healthy +loss curve, and caps acceptance length. If loss looks good but acceptance +plateaus, check this first. + +No spaces in the value — nemo_run emits `export FOO=value` unquoted, so a space +splits the variable. + +**Draft dims are not inherited.** The draft is an independent Qwen3 model; it +does *not* pick up the base model's GQA/FFN dims. `dspark.yaml` sets them +explicitly for Qwen3-8B. Retargeting to a different base means updating +`num_attention_heads`, `num_key_value_heads`, `head_dim` and `intermediate_size` +to match — otherwise you silently train a wrong-shaped draft. + +**`answer_only_loss=false`.** The streaming corpus is prompt-only (the serve +generates the response and we capture *its* hidden states), so there is no +assistant span to mask. Train over the full sequence. + +**`dflash_block_size`** is the semi-AR generation block and must divide +`training_seq_len`. + +**`report_to=none`.** `dspark.yaml` defaults to tensorboard, which hard-fails if +tensorboard isn't installed in the serve container. + +**Batch size and LR stay at the recipe defaults.** `dspark.yaml` ships +`per_device_train_batch_size=1`, `learning_rate=6e-4`, `warmup_ratio=0.04` — +tuned for a from-scratch draft on a single GPU. The Kimi-K2.6 and MiniMax-M3 +examples override these to a larger batch and a gentle `1e-4` because they run 8 +GPUs per node and warm-start a large backbone. Don't copy those numbers to +Qwen3-8B: at `training_seq_len=4096` on one GPU, a batch of 4 will OOM. + +### Scaling up to a large target + +The same pipeline drives large MoE targets — compare this YAML against +[`MiniMaxAI/MiniMax-M3`](../../tools/launcher/examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml) +and +[`moonshotai/Kimi-K2.6`](../../tools/launcher/examples/moonshotai/Kimi-K2.6/hf_streaming_dspark_multi_node.yaml). +What changes: + +| Concern | Qwen3-8B | Large MoE target | +|---|---|---| +| Topology | 2 nodes × 1 GPU | `SERVE_NODES` serve replicas + DDP trainers, 8 GPU/node | +| Draft dims | recipe defaults already match | must be set explicitly to match the base | +| `rope_theta` | default `1e6` matches base | pin the base's value onto the draft | +| Mask token | `151669` | a free id in that model's vocab | +| `trust_remote_code` | not needed | needed at serve *and* export | +| Chat template | not needed (`answer_only_loss=false`) | a `{% generation %}`-tagged copy if masking | + +`SERVE_NODES` splits the allocation: nodes `0..SERVE_NODES-1` each run an +independent whole-node serve replica, the rest are DDP trainers, and each trainer +rank fetches its shard round-robin across replicas. See the header of +[`train_eagle_streaming.sh`](../../tools/launcher/common/eagle3/train_eagle_streaming.sh) +for the full knob list. + +On clusters using EFA rather than InfiniBand, NIXL needs +`NIXL_BACKENDS=LIBFABRIC`. Note that UCX segfaults at agent init on EFA nodes, so +LIBFABRIC is required there even for single-node runs (see the commented block in +the YAML). + +--- + +## Step 3 — Evaluate + +`task_2` serves the exported drafter and reports acceptance length. + +**Do not use the training-time AR eval.** `dspark.yaml` sets +`estimate_ar: false` deliberately: that eval runs the DFlash backbone *only*, +without applying the Markov head, so its number describes the backbone rather +than the model you trained. Acceptance length comes from the vLLM smoke test or +the offline [specdec_bench](../specdec_bench/) harness. + +`task_2` needs a vLLM build whose DSpark path supports a non-DeepSeek backbone. +On an older image the failure is not a clean rejection of the method name: the +loader routes into the DeepSeek-V4 DSpark model and dies mid-load on a config +field the Qwen3 drafter has no reason to carry (`AttributeError: 'Qwen3Config' +object has no attribute 'hc_mult'`). Pin a newer nightly, or set `skip: true` +and use specdec_bench on `/scratchspace/export`. + +Reading the numbers: training accuracy, acceptance length, and accepted-fraction +are three different things and are easy to confuse. Acceptance length is the +deployment-facing one. + +--- + +## Troubleshooting + +| Symptom | Likely cause | +|---|---| +| Training starts but every batch is dropped | Corpus schema — entries without a usable `conversations`/`messages` list are skipped silently. Inspect one record. | +| Loss curve healthy, acceptance length plateaus | Wrong final capture id (see above). | +| Trainer OOM on the serve node | Lower `SERVE_GPU_MEM_UTIL`, `SERVE_MAX_MODEL_LEN`, `SERVE_MAX_NUM_SEQS`. | +| Trainer stuck at "waiting for serve address" | The serve died. Its log is `/scratchspace/vllm_serve..log`, not the task log — read that for the real error. | +| Serve never becomes ready | Raise `SERVE_READY_TIMEOUT`; large models load slowly on cold cache. | +| `export FOO=value` splitting | A space in an env value. Remove it. | +| Synthesis yield well under 100% | Lower `--num-proc`, or raise `--max-model-len`. | + +## Reference + +- [`dspark.yaml`](../../modelopt_recipes/general/speculative_decoding/dspark.yaml) — head arch and loss weights +- [`hf_synth.yaml`](../../tools/launcher/examples/Qwen/Qwen3-8B/hf_synth.yaml) — synthesis +- [`hf_streaming_dspark.yaml`](../../tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dspark.yaml) — streaming training +- [`hf_streaming_dflash.yaml`](../../tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash.yaml) — DFlash, same transport +- [README](README.md) — speculative decoding overview, other algorithms +- [SLURM_prepare_data.md](SLURM_prepare_data.md) — the alternative `server_generate.py` synthesis path diff --git a/examples/speculative_decoding/README.md b/examples/speculative_decoding/README.md index 169b83fb9e5..cba6c7af669 100644 --- a/examples/speculative_decoding/README.md +++ b/examples/speculative_decoding/README.md @@ -396,3 +396,7 @@ python scripts/export_hf_checkpoint.py \ ### Results See [doc/dflash.md](doc/dflash.md) for design details, benchmark results, and open items. + +### DSpark + +For end-to-end Qwen3.5 DSpark synthesis, streaming training, export, and benchmark commands, see the [Qwen3.5 DSpark guide](doc/dspark.md). diff --git a/examples/speculative_decoding/doc/assets/qwen3.5-35b-a3b-dspark-swe-context-distribution-t0.png b/examples/speculative_decoding/doc/assets/qwen3.5-35b-a3b-dspark-swe-context-distribution-t0.png new file mode 100644 index 00000000000..7dfffada57f Binary files /dev/null and b/examples/speculative_decoding/doc/assets/qwen3.5-35b-a3b-dspark-swe-context-distribution-t0.png differ diff --git a/examples/speculative_decoding/doc/dspark.md b/examples/speculative_decoding/doc/dspark.md new file mode 100644 index 00000000000..433f30cfe41 --- /dev/null +++ b/examples/speculative_decoding/doc/dspark.md @@ -0,0 +1,260 @@ +# Qwen3.5 DSpark + +This guide provides complete, reproducible steps to train DSpark drafters for Qwen3.5-9B and Qwen3.5-35B-A3B. See the [Qwen3-8B walkthrough](../QWEN3_8B_DSPARK_WALKTHROUGH.md) for the training mechanics. + +Hyperparameters used in this guide: + +| Parameter | Qwen3.5-9B | Qwen3.5-35B-A3B | 35B SWE fine-tuning | +|-----------------|-----------------------------------------:|-----------------------------:|---------------------------:| +| Input | Nemotron PT v2 prompts | Nemotron PT v2 prompts | Agent traces | +| Synthesis | 192 shards · mixed thinking/non-thinking | 384 shards · thinking only | none | +| Training Nodes | 8 target + 8 trainer nodes | 16 target + 16 trainer nodes | 8 target + 8 trainer nodes | +| Sequence length | 4K | 4K | 32K | +| Block size | 8 | 8 | 8 | +| Anchors | 512 | 512 | 4096 | +| Draft layers | 6 layers · FFN 12288 | 6 layers · FFN 6144 | same as 35B base | +| Captures | `[2,8,13,19,24,30,32]` | `[2,9,16,24,31,38,40]` | same as 35B base | +| Attention | SWA 4096 | SWA 4096 | same as 35B base | +| Learning rate | 6e-4 → 3e-5 · cosine | 6e-4 → 3e-5 · cosine | 1e-4 → 3e-5 · cosine | +| Global batch | 512 | 512 | 256 | +| Epochs | 5 | 5 | ~5 | + +## 1. Setup + +Requirements: a Slurm cluster and an `HF_TOKEN` with access to [Nemotron Post-Training Dataset v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2). Dataset synthesis and training use the `vllm/vllm-openai:v0.27.1` container. Benchmarks use `vllm/vllm-openai:nightly` with Model Runner V2. + +The workflow uses one shared root for the Model-Optimizer checkout, caches, and job outputs: + +```text +dspark-training/ +├── Model-Optimizer/ +├── .cache/ +└── outputs/nemorun/experiments/cicd/ +``` + +```bash +# Set the DSpark training root +export DSPARK_TRAINING_ROOT="/path/to/dspark-training" +mkdir -p "$DSPARK_TRAINING_ROOT" +cd "$DSPARK_TRAINING_ROOT" + +# Set the Slurm account and partition +export SLURM_ACCOUNT="" +export SLURM_PARTITION="" + +# Set the Hugging Face and W&B credentials +export HF_TOKEN="" +export WANDB_API_KEY="" +: "${HF_TOKEN:?Export HF_TOKEN before continuing}" +: "${WANDB_API_KEY:?Export WANDB_API_KEY before continuing}" + +# Clone Model-Optimizer and check out a specific branch when needed +git clone --recurse-submodules https://github.com/NVIDIA/Model-Optimizer.git + +# Set the cache and output directories +export SLURM_HOST="localhost" +export SLURM_HF_LOCAL="$DSPARK_TRAINING_ROOT/.cache/hf" +export UV_CACHE_DIR="$DSPARK_TRAINING_ROOT/.cache/uv" +export NEMORUN_HOME="$DSPARK_TRAINING_ROOT/outputs/nemorun" +export SLURM_JOB_DIR="$NEMORUN_HOME/experiments" +mkdir -p "$SLURM_HF_LOCAL" "$SLURM_JOB_DIR" "$NEMORUN_HOME" "$UV_CACHE_DIR" + +# Enter the launcher directory, install uv, and sync dependencies +cd "$DSPARK_TRAINING_ROOT/Model-Optimizer/tools/launcher" +command -v uv >/dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh +export PATH="$HOME/.local/bin:$PATH" +uv sync +``` + +Download HF models and dataset: + +```bash +# Download Qwen3.5-9B and DFlash drafters +uv run --with huggingface-hub hf download Qwen/Qwen3.5-9B --revision c202236235762e1c871ad0ccb60c8ee5ba337b9a --local-dir "$SLURM_HF_LOCAL/Qwen/Qwen3.5-9B" +uv run --with huggingface-hub hf download z-lab/Qwen3.5-9B-DFlash --revision 5fc3b3d474760f18c516db87d84c37edbfd3ede6 --local-dir "$SLURM_HF_LOCAL/z-lab/Qwen3.5-9B-DFlash" +# Download Qwen3.5-35B-A3B and DFlash drafters +uv run --with huggingface-hub hf download Qwen/Qwen3.5-35B-A3B --revision 59d61f3ce65a6d9863b86d2e96597125219dc754 --local-dir "$SLURM_HF_LOCAL/Qwen/Qwen3.5-35B-A3B" +uv run --with huggingface-hub hf download z-lab/Qwen3.5-35B-A3B-DFlash --revision 52cb554b4995dede3e2e1bdb129cdb1f3529332b --local-dir "$SLURM_HF_LOCAL/z-lab/Qwen3.5-35B-A3B-DFlash" + +# Prepare the synthesis prompts from Nemotron Post-Training Dataset v2 +HF_HOME="$SLURM_HF_LOCAL" uv run --with datasets python ../../examples/dataset/make_nemotron_ptv2_dataset.py --mode generate --output-dir "$SLURM_HF_LOCAL/nvidia/Nemotron-Post-Training-Dataset-v2" +``` + +## 2. Qwen3.5-9B + +### 2.a Dataset synthesis + +Synthesize the dataset with the target model using prompts from Nemotron Post-Training Dataset v2. + +```bash +# Run the two-shard synthesis pipeclean +uv run launch.py --yaml examples/Qwen/Qwen3.5-9B/hf_synth.yaml pipeline.global_vars.output_dir=/hf-local/modelopt/qwen3.5-9b-dspark-pipeclean pipeline.task_0.slurm_config.array=0-1 --yes + +# Synthesize the full dataset (192 shards, one Slurm job per shard) +uv run launch.py --yaml examples/Qwen/Qwen3.5-9B/hf_synth.yaml --yes +``` + +Each completed shard writes a non-empty `shard_.jsonl` in `$DSPARK_TRAINING_ROOT/.cache/hf/modelopt/qwen3.5-9b-dspark-synthesis/`. + +### 2.b Drafter training + +Run a two-node pipeclean on the synthesized dataset before full training. A repeated full-training submission resumes the latest checkpoint and `singleton` serializes submissions. + +```bash +# Run pipeclean training (2 nodes) +# Run two optimizer steps, then cancel it +WANDB_MODE=disabled SERVE_NODES=1 uv run launch.py --yaml examples/Qwen/Qwen3.5-9B/hf_streaming_dspark_multi_node.yaml pipeline.global_vars.output_dir=/cicd/qwen3.5-9b-dspark-pipeclean pipeline.task_1.slurm_config.nodes=2 pipeline.task_2.skip=true --yes + +# Run or resume production training (16 nodes) +uv run launch.py --yaml examples/Qwen/Qwen3.5-9B/hf_streaming_dspark_multi_node.yaml pipeline.task_2.skip=true --yes + +# Export and validate the trained drafter (T=0 and T=1 serving smoke tests) +uv run launch.py --yaml examples/Qwen/Qwen3.5-9B/hf_streaming_dspark_multi_node.yaml pipeline.task_1.skip=true --yes +``` + +Checkpoints and the serving export are written to `$DSPARK_TRAINING_ROOT/outputs/nemorun/experiments/cicd/qwen3.5-9b-dspark/{training,export}`. + +### 2.c Benchmark + +Each benchmark runs Base, DSpark7, MTP3, MTP7, and DFlash8 at T=0 and T=1 on vLLM Model Runner V2. + +```bash +# Benchmark concurrency 1 +uv run launch.py --yaml examples/Qwen/Qwen3.5-9B/specdec_bench_tp1_c1.yaml --yes + +# Benchmark concurrency 32 +uv run launch.py --yaml examples/Qwen/Qwen3.5-9B/specdec_bench_tp1_c32.yaml --yes +``` + +Results are written under `$DSPARK_TRAINING_ROOT/outputs/nemorun/experiments/cicd/` in the `qwen3.5-9b-dspark-benchmark-tp1-c1/` and `qwen3.5-9b-dspark-benchmark-tp1-c32/` directories. + +Qwen3.5-9B DSpark7 acceptance length on SPEED-Bench: + +| Category | DSpark7 (T1) | DSpark7 (T0) | +| :---- | :---- | :---- | +| Coding | 4.3773 | 4.8827 | +| Humanities | 3.1494 | 3.8773 | +| Math | 4.1374 | 4.6412 | +| Multilingual | 3.8842 | 4.3879 | +| QA | 3.2359 | 4.0704 | +| RAG | 4.1099 | 4.5759 | +| Reasoning | 3.7000 | 4.1205 | +| Roleplay | 2.6752 | 3.8322 | +| STEM | 3.3827 | 3.9522 | +| Summarization | 3.6285 | 4.2045 | +| Writing | 2.8436 | 3.1263 | +| Overall AL | 3.5567 | 4.1519 | + +## 3. Qwen3.5-35B-A3B + +### 3.a Dataset synthesis + +Synthesize the dataset with the target model using prompts from Nemotron Post-Training Dataset v2. + +```bash +# Run the one-shard synthesis pipeclean +uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/hf_synth.yaml pipeline.global_vars.output_dir=/hf-local/modelopt/qwen3.5-35b-a3b-dspark-pipeclean pipeline.task_0.slurm_config.array=0-0 --yes + +# Synthesize the full dataset (384 shards, one Slurm job per shard) +uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/hf_synth.yaml --yes +``` + +Each completed shard writes a non-empty `shard_.jsonl` in `$DSPARK_TRAINING_ROOT/.cache/hf/modelopt/qwen3.5-35b-a3b-dspark-synthesis/`. + +### 3.b Drafter training + +The 35B recipe uses 16 target-server nodes and 16 trainer nodes. + +```bash +# Run pipeclean training (2 nodes) +# Run two optimizer steps, then cancel it +WANDB_MODE=disabled SERVE_NODES=1 uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_multi_node.yaml pipeline.global_vars.output_dir=/cicd/qwen3.5-35b-a3b-dspark-pipeclean pipeline.task_1.slurm_config.nodes=2 pipeline.task_2.skip=true --yes + +# Run or resume production training (32 nodes) +uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_multi_node.yaml pipeline.task_2.skip=true --yes + +# Export and validate the trained drafter (T=0 and T=1 serving smoke tests) +uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_multi_node.yaml pipeline.task_1.skip=true --yes +``` + +Checkpoints and the serving export are written to `$DSPARK_TRAINING_ROOT/outputs/nemorun/experiments/cicd/qwen3.5-35b-a3b-dspark/{training,export}`. + +### 3.c Benchmark + +Each benchmark runs Base, DSpark7, MTP3, MTP7, and DFlash8 at T=0 and T=1 on vLLM Model Runner V2. + +```bash +# Benchmark concurrency 1 +uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/specdec_bench_tp2_c1.yaml --yes + +# Benchmark concurrency 32 +uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/specdec_bench_tp2_c32.yaml --yes +``` + +Results are written under `$DSPARK_TRAINING_ROOT/outputs/nemorun/experiments/cicd/` in the `qwen3.5-35b-a3b-dspark-benchmark-tp2-c1/` and `qwen3.5-35b-a3b-dspark-benchmark-tp2-c32/` directories. + +Qwen3.5-35B-A3B DSpark7 acceptance length on SPEED-Bench: + +| Category | DSpark7 (T1) | DSpark7 (T0) | +| :---- | :---- | :---- | +| Coding | 4.3720 | 4.8880 | +| Humanities | 3.0543 | 3.6332 | +| Math | 4.0852 | 4.5721 | +| Multilingual | 3.7786 | 4.1805 | +| QA | 3.2111 | 3.7132 | +| RAG | 4.0790 | 4.4432 | +| Reasoning | 3.6061 | 3.9967 | +| Roleplay | 2.5711 | 3.3469 | +| STEM | 3.3168 | 3.8496 | +| Summarization | 3.5660 | 4.0488 | +| Writing | 2.7125 | 2.9998 | +| Overall AL | 3.4866 | 3.9702 | + +### 3.d Optional SWE fine-tuning + +Use fine-tuning to improve acceptance length (AL) for specific domains, such as SWE and long-context workloads. + +Prepare fine-tuning data: + +1. Prepare rollout traces from the target domain. SWE traces can come from evaluation or RL rollouts. +2. Convert each conversation into one JSONL record using the following schema: + + ```json + {"conversation_id":"...","messages":[...],"token_ids":[...],"loss_mask":[0,1,...]} + ``` + + If the trace contains recorded target token IDs and an assistant loss mask, preserve them to skip additional tokenization and masking (RL traces may already contain these values, and `loss_mask=1` marks assistant tokens). Otherwise, provide `messages`, and the training loader will tokenize them and derive the loss mask. It is not necessary to truncate the sequence length during conversion because the loader automatically truncates tokens and masks together to the configured sequence length. + + See [prepare-swe-data.py](../../../tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/prepare-swe-data.py) for a Prime-RL trace conversion example. + +3. Store the JSONL records under `$DSPARK_TRAINING_ROOT/.cache/hf/modelopt/qwen3.5-35b-a3b-dspark-finetuning/`. + +Run DSpark fine-tuning: + +Set `DRAFT_CHECKPOINT` to a completed checkpoint from 3.b. The first fine-tuning run loads its model weights without the optimizer state. Repeated submissions resume the latest fine-tuning checkpoint. + +```bash +export DRAFT_CHECKPOINT="/cicd/qwen3.5-35b-a3b-dspark/training/checkpoint-" +``` + +```bash +# Run pipeclean training (2 nodes) +# Run two optimizer steps, then cancel it +WANDB_MODE=disabled SERVE_NODES=1 uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_finetuning_multi_node.yaml pipeline.global_vars.draft_model="$DRAFT_CHECKPOINT" pipeline.global_vars.output_dir=/cicd/qwen3.5-35b-a3b-dspark-finetuning-pipeclean pipeline.task_1.slurm_config.nodes=2 pipeline.task_2.skip=true --yes + +# Run or resume production training (16 nodes) +uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_finetuning_multi_node.yaml pipeline.global_vars.draft_model="$DRAFT_CHECKPOINT" pipeline.task_2.skip=true --yes + +# Export and validate the trained drafter (T=0 and T=1 serving smoke tests) +uv run launch.py --yaml examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_finetuning_multi_node.yaml pipeline.task_1.skip=true --yes +``` + +Checkpoints and the serving export are written to `$DSPARK_TRAINING_ROOT/outputs/nemorun/experiments/cicd/qwen3.5-35b-a3b-dspark-finetuning/{training,export}`. + +### 3.e Benchmark + +Benchmark the fine-tuned drafter on domain data, such as long-context SWE traces, by replaying each turn's model-call prompt. Plot AL against context length for the original and fine-tuned drafters under matched serving and sampling settings. (Long context benchmark requires the DSpark hybrid prefix-cache fix in [`jinzex/vllm:jinzex/dspark`](https://github.com/jinzex/vllm/tree/jinzex/dspark)). + +SWE fine-tuning improved DSpark7 acceptance length across all measured context lengths. + +![Qwen3.5-35B-A3B speculative acceptance distribution over 200 SWE traces](assets/qwen3.5-35b-a3b-dspark-swe-context-distribution-t0.png) diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index 0691ca06f52..7cb7bed67fb 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/examples/speculative_decoding/eagle_utils.py @@ -53,6 +53,17 @@ __all__ = ["EagleOfflineDataCollator", "OfflineSupervisedDataset"] +def _aggregate_accuracy_counts(counts): + """Compute per-position and total accuracy across microbatches and distributed ranks.""" + counts = torch.stack(counts).sum(dim=0) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.all_reduce(counts, op=torch.distributed.ReduceOp.SUM) + correct, valid = counts + per_position_acc = correct / valid.clamp_min(1.0) + total_acc = correct.sum() / valid.sum().clamp_min(1.0) + return per_position_acc.unsqueeze(0).cpu().numpy(), total_acc.item() + + def make_speculative_data_module( tokenizer: transformers.PreTrainedTokenizer, data_args, @@ -216,11 +227,16 @@ def compute_loss(self, *args, **kwargs): """Override compute_loss to save train accs and per-component losses in trainer state.""" if not hasattr(self.state, "training_accs"): self.state.training_accs = [] + if not hasattr(self.state, "training_acc_counts"): + self.state.training_acc_counts = [] if not hasattr(self.state, "component_losses"): self.state.component_losses = {"eagle": [], "preservation": []} kwargs.pop("num_items_in_batch", None) loss, outputs = super().compute_loss(return_outputs=True, *args, **kwargs) - if hasattr(outputs, "train_acc") and any(outputs.train_acc): + if hasattr(outputs, "train_acc_correct"): + counts = torch.stack((outputs.train_acc_correct, outputs.train_acc_valid)).detach() + self.state.training_acc_counts.append(counts) + elif hasattr(outputs, "train_acc") and any(outputs.train_acc): self.state.training_accs.append(outputs.train_acc) # Track per-component losses for key, attr in [ @@ -371,37 +387,57 @@ def __init__(self, ar_validate_steps: int = 1000, estimate_ar: bool = False): self.estimate_ar = estimate_ar def on_log(self, args, state, control, **kwargs): - """Log training acc and estimate AR during log step.""" - if not hasattr(state, "training_accs") or len(state.training_accs) == 0: + """Log training accuracy and optionally estimate acceptance length.""" + # Normalize current raw-count metrics and legacy accuracy metrics to + # [parallel branch, draft position]. + accuracy_counts = getattr(state, "training_acc_counts", []) + legacy_accs = getattr(state, "training_accs", []) + if accuracy_counts: + per_position_acc, total_acc = _aggregate_accuracy_counts(accuracy_counts) + elif legacy_accs: + per_position_acc = np.mean(legacy_accs, axis=0) + total_acc = None + else: return control - average_acc = np.mean(state.training_accs, axis=0) - # Always print accuracy to console + + # Print one accuracy summary for the completed logging window. try: - acc_str = ", ".join(f"{a:.4f}" for a in np.array(average_acc).flatten()) - print_rank_0(f"Step {state.global_step} Training Acc: [{acc_str}]") + acc_str = ", ".join(f"{a:.4f}" for a in np.array(per_position_acc).flatten()) + if total_acc is not None: + print_rank_0( + f"Step {state.global_step} Training Acc: " + f"total={total_acc:.4f}, per-position=[{acc_str}]" + ) + else: + print_rank_0(f"Step {state.global_step} Training Acc: [{acc_str}]") except Exception: - print_rank_0(f"Step {state.global_step} Training Acc: {average_acc}") - # Log accuracy to HF Trainer's logs dict (picked up by TensorBoard) - logs = kwargs.get("logs") or {} - for i, draft_acc in enumerate(average_acc): + print_rank_0(f"Step {state.global_step} Training Acc: {per_position_acc}") + + # Publish total and per-position metrics through the Trainer log dict. + logs = kwargs.get("logs") + if logs is None: + logs = {} + if total_acc is not None: + logs["train_acc/total"] = total_acc + for i, draft_acc in enumerate(per_position_acc): for j, step_acc in enumerate(draft_acc): logs[f"train_acc/parallel_{i}_step_{j}"] = float(step_acc) + if self.estimate_ar: - # Calculate mean training AR since last log - # NOTE: This is only an estimate of the real AR. + # Estimate acceptance length from conditional per-position accuracies. est_ar = 1 acc_cumprod = 1 - for step_acc in average_acc[0]: + for step_acc in per_position_acc[0]: acc_cumprod *= step_acc est_ar += acc_cumprod - # Parallel draft tokens only used after all eagle tokens - for draft_acc in average_acc[1:]: + # Parallel branches contribute only after all sequential EAGLE tokens are accepted. + for draft_acc in per_position_acc[1:]: acc_cumprod *= draft_acc[-1] est_ar += acc_cumprod print_rank_0(f"Step {state.global_step} Estimated Training AR: {est_ar:.4f}") logs["estimated_training_ar"] = est_ar - # log to wandb + # Forward the same logging window to W&B when enabled. if wandb is not None and wandb.run is not None and is_master(): if logs: wandb.log({k: v for k, v in logs.items() if v is not None}, step=state.global_step) @@ -412,8 +448,9 @@ def on_log(self, args, state, control, **kwargs): if vals: wandb.log({f"{key}_loss": np.mean(vals)}, step=state.global_step) - # reset training_accs and component_losses + # Start a fresh accumulation window after logging. state.training_accs = [] + state.training_acc_counts = [] if hasattr(state, "component_losses"): state.component_losses = {"eagle": [], "preservation": []} return control diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 46848203606..ad2ca31b08a 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -204,7 +204,16 @@ def train(): if last_checkpoint: print_rank_0(f"Last checkpoint detected: {last_checkpoint}") - checkpoint = training_args.resume_from_checkpoint or last_checkpoint + if training_args.resume_model_only: + # With resume_model_only, initialize the first run from the configured checkpoint; + # fully resume the latest local checkpoint thereafter. + checkpoint = last_checkpoint or training_args.resume_from_checkpoint + model_only_init = last_checkpoint is None + else: + # Without resume_model_only, use the requested checkpoint when set; otherwise resume + # the latest local checkpoint. + checkpoint = training_args.resume_from_checkpoint or last_checkpoint + model_only_init = False use_offline_training = recipe.data.mode != "online" @@ -212,6 +221,8 @@ def train(): # weights load via from_pretrained; FSDP sharded checkpoints load the base model and # resume through the Trainer. checkpoint_is_hf = _is_hf_format_checkpoint(checkpoint) + if model_only_init and not checkpoint_is_hf: + raise ValueError("resume_model_only requires a consolidated Hugging Face checkpoint.") if checkpoint_is_hf: assert checkpoint is not None # guaranteed by checkpoint_is_hf @@ -355,7 +366,9 @@ def train(): fsdp2_buffer_patch.log_param_dtypes(trainer.model) print_rank_0("Start training...") - trainer.train(resume_from_checkpoint=checkpoint) + # With resume_model_only, initialize the first run from the configured checkpoint with fresh + # Trainer state (skipping optimizer state loading); otherwise resume full training state. + trainer.train(resume_from_checkpoint=None if model_only_init else checkpoint) trainer.save_state() trainer.save_model(training_args.output_dir) diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index fad15abc2db..ce64f7adea1 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -412,12 +412,11 @@ def _export_config(self): else: config["layer_types"] = ["full_attention"] * draft_config.num_hidden_layers - # Sliding-window attention: all draft layers use SWA. vLLM's - # _resolve_layer_attention reads dflash_config.use_swa + swa_window_size; with - # layer_types left all "full_attention" it applies a sliding window to every draft - # layer (window from swa_window_size / top-level sliding_window). + # Sliding-window attention: all draft layers use SWA. Keep both the explicit + # layer types and the nested vLLM fields in the exported config. swa_window = getattr(self.model, "dflash_swa_window_size", None) if swa_window is not None: + config["layer_types"] = ["sliding_attention"] * draft_config.num_hidden_layers config["sliding_window"] = swa_window config["dflash_config"].update( { @@ -523,18 +522,30 @@ def _export_config(self): class DSparkExporter(DFlashExporter): """Draft model exporter for DSpark (DFlash backbone + sequential Markov head). - Same z-lab-compatible format as DFlash, plus the DSpark head weights - (``markov_w1.*`` / ``markov_w2.*`` / ``gate_proj.*`` / ``joint_proj.*`` / - ``confidence_proj.*``, already captured by the inherited ``dflash_module.`` - stripping) and the extra config fields the loader needs to rebuild the head - (``projector_type``, ``markov_rank``, ``markov_head_type``, - ``use_confidence_head``, ``shift_label``). + ModelOpt currently uses a Qwen3 draft backbone regardless of target architecture. + Head tensors and top-level config fields follow the DeepSpec/vLLM layout. """ + def _extract_state_dict(self, full_state_dict: dict): + """Map ModelOpt head names to the DeepSpec/vLLM DSpark layout.""" + export_sd = super()._extract_state_dict(full_state_dict) + prefixes = { + "markov_w1.": "markov_head.markov_w1.", + "markov_w2.": "markov_head.markov_w2.", + "gate_proj.": "markov_head.gate_proj.", + "joint_proj.": "markov_head.joint_proj.", + "confidence_proj.": "confidence_head.proj.", + } + for old_prefix, new_prefix in prefixes.items(): + for key in [key for key in export_sd if key.startswith(old_prefix)]: + export_sd[new_prefix + key[len(old_prefix) :]] = export_sd.pop(key) + return export_sd + def _export_config(self): - """Extend the DFlash config with the DSpark head fields.""" + """Add the fields consumed by the Qwen3 DSpark runtime.""" config = super()._export_config() draft_config = self.model.dflash_config + use_confidence_head = bool(getattr(draft_config, "use_confidence_head", False)) config["dflash_config"].update( { @@ -542,7 +553,19 @@ def _export_config(self): "shift_label": getattr(draft_config, "shift_label", True), "markov_rank": draft_config.markov_rank, "markov_head_type": getattr(draft_config, "markov_head_type", "vanilla"), - "use_confidence_head": bool(getattr(draft_config, "use_confidence_head", False)), + "use_confidence_head": use_confidence_head, + } + ) + config["architectures"] = ["Qwen3DSparkModel"] + config.update( + { + "mask_token_id": config["dflash_config"]["mask_token_id"], + "target_layer_ids": config["dflash_config"]["target_layer_ids"], + "num_anchors": self.model.dflash_num_anchors, + "markov_rank": draft_config.markov_rank, + "markov_head_type": getattr(draft_config, "markov_head_type", "vanilla"), + "enable_confidence_head": use_confidence_head, + "confidence_head_with_markov": use_confidence_head, } ) return config diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index ceb76c51f91..7cf818793d5 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -408,9 +408,8 @@ def modify(self, config): "rms_norm_eps", ] for attr in _setdefault_attrs: - if not hasattr(self.dflash_config, attr) or getattr(self.dflash_config, attr) is None: - if hasattr(base_config, attr): - setattr(self.dflash_config, attr, getattr(base_config, attr)) + if attr not in config.dflash_architecture_config and hasattr(base_config, attr): + setattr(self.dflash_config, attr, getattr(base_config, attr)) # RoPE base settings are ENFORCED to match the base model (not setdefault): the # DFlash draft injects the target's KV into every layer, so its RoPE base must @@ -421,10 +420,18 @@ def modify(self, config): # DFlash uses standard Qwen3 RotaryEmbedding; the long-context YaRN scaling is # added only at export via dflash_export_rope_scaling.) for attr in ("rope_theta", "rope_type", "rope_interleaved"): - if not hasattr(base_config, attr): + base_val = getattr(base_config, attr, None) + if base_val is None and attr == "rope_theta": + rope_parameters = getattr(base_config, "rope_parameters", None) + if isinstance(rope_parameters, dict): + base_val = rope_parameters.get(attr) + if base_val is None: continue - base_val = getattr(base_config, attr) - user_val = getattr(self.dflash_config, attr, None) + draft_rope_parameters = getattr(self.dflash_config, "rope_parameters", None) + if attr == "rope_theta" and isinstance(draft_rope_parameters, dict): + user_val = draft_rope_parameters.get(attr) + else: + user_val = getattr(self.dflash_config, attr, None) if user_val is not None and user_val != base_val: logger.warning( "DFlash: ignoring dflash_architecture_config.%s=%r and enforcing the " @@ -434,7 +441,10 @@ def modify(self, config): user_val, base_val, ) - setattr(self.dflash_config, attr, base_val) + if attr == "rope_theta" and isinstance(draft_rope_parameters, dict): + draft_rope_parameters[attr] = base_val + else: + setattr(self.dflash_config, attr, base_val) self.dflash_config.head_dim = getattr( self.dflash_config, diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index 4183a63f1d1..483d45a78aa 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -103,6 +103,15 @@ def _chunk(a, b): return torch.cat(outs, dim=0) +@torch.no_grad() +def _accuracy_counts(logits, target_ids, eval_mask): + """Return correct/valid token counts for each draft position.""" + keep = eval_mask > 0.5 + correct = ((logits.argmax(dim=-1) == target_ids) & keep).sum(dim=(0, 1), dtype=torch.float32) + valid = keep.sum(dim=(0, 1), dtype=torch.float32) + return correct, valid + + @DSparkDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) class HFDSparkModel(HFDFlashModel): """DFlash model with the DSpark sequential (Markov) + confidence head. @@ -208,6 +217,7 @@ def _compute_dspark_loss( ) weight_mask = weight_mask * orig_loss_mask + accuracy_correct, accuracy_valid = _accuracy_counts(final_logits, target_ids, weight_mask) binary_eval_mask = weight_mask.view(-1) # Exponential position decay (exp(-k/gamma); position 0 gets weight 1). @@ -235,7 +245,7 @@ def _compute_dspark_loss( if valid_count <= 1.0: loss = flat_final.sum() * 0.0 metrics = {"ce_loss": 0.0, "l1_loss": 0.0, "confidence_loss": 0.0, "base_accuracy": 0.0} - return loss, 0.0, metrics + return loss, accuracy_correct, accuracy_valid, metrics # Term 1: cross-entropy on the corrected (final) logits. ce_per_token = F.cross_entropy(flat_final, flat_targets, reduction="none") @@ -264,9 +274,6 @@ def _compute_dspark_loss( with torch.no_grad(): eval_count = binary_eval_mask.sum() + 1e-6 keep = binary_eval_mask > 0.5 - accuracy = ( - ((flat_final.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count - ).item() base_accuracy = ( ((flat_base.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count ).item() @@ -276,7 +283,7 @@ def _compute_dspark_loss( "confidence_loss": float(confidence_loss.detach().item()), "base_accuracy": base_accuracy, } - return loss, accuracy, metrics + return loss, accuracy_correct, accuracy_valid, metrics def forward( self, @@ -373,7 +380,13 @@ def forward( if n_blocks == 0 or not block_keep_mask.any(): # Zero loss that still flows through all draft params for DDP sync. dummy = sum(p.sum() for p in self.dflash_module.parameters()) * 0.0 - return ModelOutput(loss=dummy, logits=None, train_acc=[[0.0]]) + zeros = torch.zeros(block_size, device=device) + return ModelOutput( + loss=dummy, + logits=None, + train_acc_correct=zeros, + train_acc_valid=zeros, + ) # 4. Build draft inputs (inherited helpers). noise_embedding = self._build_noise_embedding( @@ -403,7 +416,7 @@ def forward( final_logits, confidence_logits = self._apply_markov_head( hidden, backbone_logits, input_ids, anchor_positions, n_blocks ) - loss, accuracy, metrics = self._compute_dspark_loss( + loss, accuracy_correct, accuracy_valid, metrics = self._compute_dspark_loss( backbone_logits, final_logits, confidence_logits, @@ -414,7 +427,13 @@ def forward( target_model_logits, ) - return ModelOutput(loss=loss, logits=None, train_acc=[[accuracy]], dspark_metrics=metrics) + return ModelOutput( + loss=loss, + logits=None, + train_acc_correct=accuracy_correct, + train_acc_valid=accuracy_valid, + dspark_metrics=metrics, + ) @torch.no_grad() def pseudo_speculative_generate(self, input_ids, steps=1): diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py index ad6082eb71a..38290faea82 100644 --- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py +++ b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py @@ -213,7 +213,7 @@ def __init__( (seeded by ``training_args.seed``), so no shuffle is needed here. Args: - entries: Untokenized per-sample dicts from the input jsonl. Schema is + entries: Per-sample dicts from the input jsonl. Schema is subclass-defined (see :meth:`_tokenize_entry`); passed to :meth:`_fetch`. tokenizer: HF tokenizer; used for client-side tokenization and the server/client loss-mask alignment in :meth:`_fetch`. @@ -283,11 +283,28 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: def _tokenize_entry(self, entry: dict) -> dict | None: """Tokenize a single entry. - Returns ``None`` for entries missing ``cid`` / ``conversations``-or-``messages``, or when - right-truncation to ``max_seq_len`` drops the entire supervised span - (``answer_only_loss`` mode with the assistant turn at the tail). + Pretokenized entries may provide aligned ``token_ids`` / ``loss_mask`` lists. + Otherwise, tokenize ``conversations`` or ``messages``. Returns ``None`` when + required data is absent or right-truncation drops the supervised span. """ cid = entry.get("conversation_id") or entry.get("uuid") + token_ids = entry.get("token_ids") + loss_mask = entry.get("loss_mask") + if token_ids is not None or loss_mask is not None: + if not self.config.answer_only_loss: + raise ValueError("pretokenized loss_mask requires answer_only_loss=True") + if not isinstance(token_ids, list) or not isinstance(loss_mask, list): + raise ValueError("pretokenized entries require token_ids and loss_mask lists") + if len(token_ids) != len(loss_mask): + raise ValueError("pretokenized token_ids and loss_mask lengths must match") + if self.config.max_seq_len is not None: + token_ids = token_ids[: self.config.max_seq_len] + loss_mask = loss_mask[: self.config.max_seq_len] + loss_mask = torch.tensor(loss_mask, dtype=torch.long) + if cid is None or not token_ids or int(loss_mask.sum()) == 0: + return None + return {"cid": str(cid), "token_ids": token_ids, "loss_mask": loss_mask} + # Prefer ``conversations``, fall back to ``messages`` (the documented default format; # see examples README). The order matters: some corpora (e.g. Spec-Decoding-Dataset-v2) # carry a degenerate user-only ``messages`` stub (no assistant turn) alongside the real diff --git a/modelopt/torch/speculative/plugins/hf_training_args.py b/modelopt/torch/speculative/plugins/hf_training_args.py index 38d3f483e6b..872ddf72d95 100644 --- a/modelopt/torch/speculative/plugins/hf_training_args.py +++ b/modelopt/torch/speculative/plugins/hf_training_args.py @@ -107,6 +107,7 @@ class TrainingArguments(BaseModel): model_config = ConfigDict(extra="allow") + resume_model_only: bool = False training_seq_len: int = 2048 estimate_ar: bool = False ar_validate_steps: int = 1000 diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py index 2b5fe989c03..92b08d635f2 100644 --- a/modelopt/torch/speculative/plugins/modeling_fakebase.py +++ b/modelopt/torch/speculative/plugins/modeling_fakebase.py @@ -191,6 +191,10 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM ), orig_config, ) + rope_theta = getattr(base_cfg, "rope_theta", None) + rope_parameters = getattr(base_cfg, "rope_parameters", {}) + if rope_theta is None and isinstance(rope_parameters, dict): + rope_theta = rope_parameters.get("rope_theta") # Extract necessary info for spec training from base config config = FakeBaseConfig( num_hidden_layers=getattr(base_cfg, "num_hidden_layers", None), @@ -203,7 +207,7 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM num_key_value_heads=getattr(base_cfg, "num_key_value_heads", None), intermediate_size=getattr(base_cfg, "intermediate_size", None), rms_norm_eps=getattr(base_cfg, "rms_norm_eps", 1e-6), - rope_theta=getattr(base_cfg, "rope_theta", None), + rope_theta=rope_theta, final_norm_type=_select_final_norm_type( getattr(base_cfg, "model_type", None), base_cfg ), diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py index 10395c71872..d9aa672c773 100644 --- a/modelopt/torch/speculative/plugins/modeling_final_norm.py +++ b/modelopt/torch/speculative/plugins/modeling_final_norm.py @@ -87,6 +87,8 @@ def extra_repr(self): "qwen2": "rmsnorm", "qwen3": "rmsnorm", "qwen3_moe": "rmsnorm", + "qwen3_5_text": "gemma_rmsnorm", # Qwen3.5 scales by (1 + weight). + "qwen3_5_moe_text": "gemma_rmsnorm", "deepseek_v3": "rmsnorm", "kimi_k2": "rmsnorm", # Kimi-K2 / K2-Thinking (DeepSeek-V3 arch) report model_type "kimi_k2" "kimi_k25": "rmsnorm", # Kimi-K2.5 / K2.6 / K2.7 all report model_type "kimi_k25" diff --git a/modelopt_recipes/general/speculative_decoding/dspark.yaml b/modelopt_recipes/general/speculative_decoding/dspark.yaml index cb7b09f40e9..6b5aef86f53 100644 --- a/modelopt_recipes/general/speculative_decoding/dspark.yaml +++ b/modelopt_recipes/general/speculative_decoding/dspark.yaml @@ -73,6 +73,7 @@ dflash: dflash_self_logit_distillation: false # gamma for exponential loss decay (block_size=16 -> 7). dflash_loss_decay_factor: 7.0 + dflash_loss_objective: decay # Qwen3 has no native mask token; 151669 is an unused id used by the reference. dflash_mask_token_id: 151669 # Three-term loss weights (DeepSpec defaults: L1/TVD-dominant). diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py index cf6dfe1a6bc..2b0db5766b6 100644 --- a/tests/unit/torch/speculative/plugins/test_fakebase.py +++ b/tests/unit/torch/speculative/plugins/test_fakebase.py @@ -67,6 +67,12 @@ def test_fakebase_local_happy_path(fake_checkpoint): assert model.embed_tokens.weight.shape == torch.Size([_VOCAB_SIZE, _HIDDEN_SIZE]) +def test_fakebase_reads_transformers5_rope_theta(fake_checkpoint, fake_config): + fake_config.rope_parameters = {"rope_theta": 10_000_000} + model = FakeBaseModel.from_source(str(fake_checkpoint)) + assert model.config.rope_theta == 10_000_000 + + def test_fakebase_missing_index_raises(tmp_path, fake_config): with pytest.raises(FileNotFoundError, match="safetensors"): FakeBaseModel.from_source(str(tmp_path)) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ab5c5a57d21..088dc53f36d 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -804,8 +804,7 @@ def test_export_swa_fields(self, tmp_path): with open(export_dir / "config.json") as f: cfg = json.load(f) - # vLLM _resolve_layer_attention reads these; all-full layer_types + use_swa=True - # → non-causal sliding window on every draft layer. + assert cfg["layer_types"] == ["sliding_attention"] * NUM_DRAFT_LAYERS assert cfg["sliding_window"] == 256 assert cfg["dflash_config"]["use_swa"] is True assert cfg["dflash_config"]["swa_window_size"] == 256 diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index 3686788ad90..30041192429 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -32,12 +32,13 @@ from _test_utils.torch.transformers_models import get_tiny_llama from safetensors.torch import load_file, save_file from transformers import AutoModelForCausalLM +from transformers.models.qwen3.configuration_qwen3 import Qwen3Config import modelopt.torch.opt as mto import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel -from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel +from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel, _accuracy_counts from modelopt.torch.speculative.plugins.modeling_dflash import ( DFlashModule, build_target_layer_ids, @@ -73,12 +74,26 @@ def _get_dspark_config( "markov_rank": MARKOV_RANK, "markov_head_type": head_type, "use_confidence_head": use_confidence_head, - "pure_draft_prefix_len": 1, "shift_label": True, } return config +def test_accuracy_counts_per_position(): + """Accuracy is represented as raw correct/valid counts for each block position.""" + pred_ids = torch.tensor([[[0, 1, 2, 0], [1, 1, 0, 2]]]) + target_ids = torch.tensor([[[0, 2, 2, 1], [1, 0, 0, 2]]]) + eval_mask = torch.tensor([[[1, 1, 1, 0], [1, 0, 1, 1]]]) + logits = torch.nn.functional.one_hot(pred_ids, num_classes=3).float() + + correct, valid = _accuracy_counts(logits, target_ids, eval_mask) + + torch.testing.assert_close(correct, torch.tensor([2.0, 0.0, 2.0, 1.0])) + torch.testing.assert_close(valid, torch.tensor([2.0, 1.0, 2.0, 1.0])) + assert correct.sum() == 5 + assert valid.sum() == 6 + + class TestDSparkConvert: """Test DSpark model conversion routing and head construction.""" @@ -170,6 +185,9 @@ def test_forward_loss_metrics_and_grads(self, head_type): assert out.loss.requires_grad assert out.loss.dim() == 0 + assert out.train_acc_correct.shape == (BLOCK_SIZE,) + assert out.train_acc_valid.shape == (BLOCK_SIZE,) + assert torch.all(out.train_acc_correct <= out.train_acc_valid) # three-term loss bookkeeping for key in ("ce_loss", "l1_loss", "confidence_loss", "base_accuracy"): assert key in out.dspark_metrics @@ -267,7 +285,7 @@ def spy(*args, **kwargs): class TestDSparkExporter: - """Test the DSpark checkpoint export format (z-lab-compatible layout).""" + """Test the DeepSpec/vLLM Qwen3 DSpark checkpoint layout.""" def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False): model = get_tiny_llama(num_hidden_layers=4) @@ -288,36 +306,70 @@ def _export(self, tmp_path, head_type="vanilla", use_confidence_head=False): @pytest.mark.parametrize("head_type", HEAD_TYPES) def test_export_weight_keys_match_reference(self, tmp_path, head_type): - """Exported weights carry the head tensors under reference names, no prefix.""" + """Exported head tensors use the Qwen3DSparkModel module names.""" sd = load_file(str(self._export(tmp_path, head_type=head_type) / "model.safetensors")) for key in sd: assert "dflash_module." not in key assert "rotary_emb" not in key - assert "markov_w1.weight" in sd - assert "markov_w2.weight" in sd - assert ("gate_proj.weight" in sd) == (head_type == "gated") - assert ("joint_proj.weight" in sd) == (head_type == "rnn") + assert "markov_head.markov_w1.weight" in sd + assert "markov_head.markov_w2.weight" in sd + assert ("markov_head.gate_proj.weight" in sd) == (head_type == "gated") + assert ("markov_head.joint_proj.weight" in sd) == (head_type == "rnn") def test_export_includes_confidence_weights(self, tmp_path): """The confidence head weights are exported when enabled.""" sd = load_file(str(self._export(tmp_path, use_confidence_head=True) / "model.safetensors")) - assert "confidence_proj.weight" in sd + assert "confidence_head.proj.weight" in sd + assert "confidence_head.proj.bias" in sd def test_export_config_has_dspark_fields(self, tmp_path): - """config.json carries the dflash_config DSpark head fields.""" + """config.json matches the Qwen3DSparkModel runtime contract.""" export_dir = self._export(tmp_path, head_type="gated") with open(export_dir / "config.json") as f: cfg = json.load(f) - assert cfg["architectures"] == ["DFlashDraftModel"] + assert cfg["architectures"] == ["Qwen3DSparkModel"] dc = cfg["dflash_config"] assert dc["projector_type"] == "dspark" + assert dc["shift_label"] is True assert dc["markov_rank"] == MARKOV_RANK assert dc["markov_head_type"] == "gated" assert dc["use_confidence_head"] is False - assert dc["shift_label"] is True - assert "mask_token_id" in dc - assert "target_layer_ids" in dc + assert cfg["markov_rank"] == MARKOV_RANK + assert cfg["markov_head_type"] == "gated" + assert cfg["enable_confidence_head"] is False + assert cfg["confidence_head_with_markov"] is False + assert cfg["num_anchors"] == 512 + assert cfg["mask_token_id"] == cfg["dflash_config"]["mask_token_id"] + assert cfg["target_layer_ids"] == cfg["dflash_config"]["target_layer_ids"] + + def test_export_config_matches_training_config(self): + """Export preserves every training draft config field after canonicalization.""" + model = get_tiny_llama(num_hidden_layers=4) + model.config.rope_theta = None + model.config.rope_parameters = { + "rope_type": "default", + "rope_theta": 10_000_000, + } + mtsp.convert(model, [("dflash", _get_dspark_config())]) + + training_config = model.dflash_config.to_dict() + exported_config = model.get_exporter()._export_config() + runtime_config = Qwen3Config( + **{**exported_config, **exported_config["dflash_config"]} + ).to_dict() + runtime_config.setdefault("attention_sink_bias", False) + metadata = {"architectures", "dtype", "transformers_version"} + + mismatches = { + field: ( + training_config[field], + runtime_config.get(field, ""), + ) + for field in training_config.keys() - metadata + if field not in runtime_config or training_config[field] != runtime_config[field] + } + assert not mismatches, json.dumps(mismatches, indent=2) class TestDraftAttentionPattern: @@ -705,7 +757,6 @@ def test_nested_head_shape_mismatch_reported(self, tmp_path): path = export_dir / "model.safetensors" sd = load_file(str(path)) sd["markov_head.markov_w1.weight"] = torch.zeros(3, MARKOV_RANK) - del sd["markov_w1.weight"] save_file(sd, str(path)) with pytest.raises(ValueError, match="shape mismatch"): self._make_model(init_checkpoint=export_dir) diff --git a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py index efd77267b3c..2d185499eda 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py +++ b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py @@ -506,6 +506,40 @@ def _fast_tokenizer_with_template(template: str, seq: int = 8) -> MagicMock: _CONV = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] +def test_pretokenized_entry_preserves_mask_and_truncates_in_lockstep(): + ds = StreamingDataset( + [{"conversation_id": "trace", "token_ids": [1, 2, 3, 4], "loss_mask": [0, 1, 0, 1]}], + tokenizer=MagicMock(), + config=StreamingConfig(answer_only_loss=True, max_seq_len=3), + ) + + sample = ds._tokenize_entry(ds.entries[0]) + + assert sample["cid"] == "trace" + assert sample["token_ids"] == [1, 2, 3] + assert sample["loss_mask"].tolist() == [0, 1, 0] + ds.tokenizer.apply_chat_template.assert_not_called() + + +def test_pretokenized_entry_rejects_misaligned_mask(): + ds = StreamingDataset( + [{"conversation_id": "trace", "token_ids": [1, 2], "loss_mask": [1]}], + tokenizer=MagicMock(), + config=StreamingConfig(answer_only_loss=True), + ) + + with pytest.raises(ValueError, match="lengths must match"): + ds._tokenize_entry(ds.entries[0]) + + +def test_pretokenized_entry_requires_answer_only_loss(): + entry = {"conversation_id": "trace", "token_ids": [1], "loss_mask": [1]} + ds = StreamingDataset([entry], tokenizer=MagicMock()) + + with pytest.raises(ValueError, match="answer_only_loss=True"): + ds._tokenize_entry(entry) + + def test_answer_only_loss_rejects_template_without_generation_tags(): """A fast tokenizer whose template lacks {% generation %} tags fails loudly. diff --git a/tests/unit/torch/speculative/plugins/test_modeling_final_norm.py b/tests/unit/torch/speculative/plugins/test_modeling_final_norm.py index 882e1c85416..8b685f3ad4c 100644 --- a/tests/unit/torch/speculative/plugins/test_modeling_final_norm.py +++ b/tests/unit/torch/speculative/plugins/test_modeling_final_norm.py @@ -21,6 +21,7 @@ pytest.importorskip("transformers") from modelopt.torch.speculative.plugins.modeling_final_norm import ( + _FinalGemmaRMSNorm, _FinalRMSNorm, _maybe_apply_base_final_norm, _select_final_norm_type, @@ -34,6 +35,8 @@ [ ("llama", "rmsnorm"), ("qwen3", "rmsnorm"), + ("qwen3_5_text", "gemma_rmsnorm"), + ("qwen3_5_moe_text", "gemma_rmsnorm"), ("deepseek_v3", "rmsnorm"), ("kimi_k2", "rmsnorm"), ("kimi_k25", "rmsnorm"), @@ -60,6 +63,18 @@ def test_final_rmsnorm_dtype_and_shape(): assert out.shape == x.shape +def test_qwen35_final_norm_matches_transformers(): + qwen35 = pytest.importorskip("transformers.models.qwen3_5.modeling_qwen3_5") + reference = qwen35.Qwen3_5RMSNorm(_HIDDEN_SIZE, eps=1e-6).to(torch.bfloat16) + actual = _FinalGemmaRMSNorm(_HIDDEN_SIZE, eps=1e-6, dtype=torch.bfloat16) + weight = torch.randn(_HIDDEN_SIZE, dtype=torch.bfloat16) + with torch.no_grad(): + reference.weight.copy_(weight) + actual.weight.copy_(weight) + hidden = torch.randn(2, 4, _HIDDEN_SIZE, dtype=torch.bfloat16) + torch.testing.assert_close(actual(hidden), reference(hidden), rtol=0, atol=0) + + def test_maybe_apply_norm_postnorm_is_noop(): """base_hidden_prenorm falsy or absent -> hidden returned unchanged, norm never touched.""" hidden = torch.randn(2, 4, _HIDDEN_SIZE) diff --git a/tools/launcher/common/eagle3/train_eagle_streaming.sh b/tools/launcher/common/eagle3/train_eagle_streaming.sh index 72b7388c3b1..4316f550c25 100755 --- a/tools/launcher/common/eagle3/train_eagle_streaming.sh +++ b/tools/launcher/common/eagle3/train_eagle_streaming.sh @@ -137,7 +137,7 @@ gpus_on_node() { nvidia-smi --query-gpu=count --format=csv,noheader,nounits | he # $1 = optional override (SERVE_ADVERTISE_IP / TRAINER_ADVERTISE_IP) resolve_routable_ip() { local ip="$1" - [ -z "$ip" ] && ip=$(getent hosts "${SLURMD_NODENAME:-$(hostname)}" 2>/dev/null | awk '{print $1}' | head -1) + [ -z "$ip" ] && ip=$(getent ahostsv4 "${SLURMD_NODENAME:-$(hostname)}" 2>/dev/null | awk '{print $1}' | head -1) [ -z "$ip" ] && ip=$(hostname -I | tr ' ' '\n' | grep -vE '^(127\.|169\.254\.|fe80:|::1)' | head -1) [ -z "$ip" ] && ip=$(hostname -I | awk '{print $1}') echo "$ip" diff --git a/tools/launcher/common/query.py b/tools/launcher/common/query.py index 27c41953d8a..30ff101d964 100644 --- a/tools/launcher/common/query.py +++ b/tools/launcher/common/query.py @@ -22,11 +22,12 @@ # ruff: noqa: D101, D102, D103, D107, PLR1722 import argparse +import json import os import re from datasets import load_dataset -from openai import OpenAI +from openai import BadRequestError, OpenAI early_termination = False @@ -46,7 +47,15 @@ def __init__(self, args): self.args = args self._pid = os.getpid() self.client = OpenAI(base_url=args.base_url) - self.generate(messages=[{"role": "user", "content": "Hello! /no_think"}], verbose=True) + # Exercise the selected no-thinking control path during server warmup. + if args.thinking_control == "chat-template-kwargs": + self.generate( + messages=[{"role": "user", "content": "Hello!"}], + verbose=True, + enable_thinking=False, + ) + else: + self.generate(messages=[{"role": "user", "content": "Hello! /no_think"}], verbose=True) def _ensure_client(self): """Reinitialize the HTTP client if we've been forked into a new process. @@ -59,15 +68,39 @@ def _ensure_client(self): self._pid = os.getpid() self.client = OpenAI(base_url=self.args.base_url) - def generate(self, messages, verbose=False, **chat_template_kwargs): + def generate( + self, + messages, + verbose=False, + sample_id="", + sampling_params=None, + **chat_template_kwargs, + ): global early_termination self._ensure_client() try: + sampling_params = sampling_params or {} + arg_temperature = self.args.temperature if self.args.temperature is not None else 0.0 + chat_template_kwargs = { + key: value for key, value in chat_template_kwargs.items() if value is not None + } + # vLLM exposes top_k and chat-template controls as OpenAI API extensions. + request_kwargs = {} + extra_body = {} + if chat_template_kwargs: + extra_body["chat_template_kwargs"] = chat_template_kwargs + if "top_k" in sampling_params: + extra_body["top_k"] = sampling_params["top_k"] + if extra_body: + request_kwargs["extra_body"] = extra_body completion = self.client.chat.completions.create( model=self.args.model, messages=messages, - temperature=self.args.temperature, + temperature=sampling_params.get("temperature", arg_temperature), + top_p=sampling_params.get("top_p", 1.0), + presence_penalty=sampling_params.get("presence_penalty", 0.0), max_tokens=self.args.max_tokens, + **request_kwargs, ) new_message = completion.choices[0].message.content if verbose: @@ -76,11 +109,18 @@ def generate(self, messages, verbose=False, **chat_template_kwargs): print("[NEW] {:10}: {:64}\n\n".format("assistant", new_message)) new_message = {"role": "assistant", "content": new_message} + except BadRequestError as e: + # Skip overlength rows; all other request errors must fail the shard. + if e.param == "input_tokens": + print(f"[SKIP] {sample_id}: {e}") + return None + print(e) + raise RuntimeError(str(e)) from None except Exception as e: print(e) if "Connection error" in str(e): early_termination = True - raise # always propagate so datasets.map() halts the shard + raise RuntimeError(str(e)) from None return new_message @@ -103,7 +143,9 @@ def generate(self, messages, verbose=False, **chat_template_kwargs): "--num-samples", "--num_samples", type=int, default=None, help="maximum samples to process." ) parser.add_argument("--num-proc", type=int, default=32, help="number of processes (concurrency).") -parser.add_argument("--temperature", type=float, default=0.0, help="temperature.") +parser.add_argument("--temperature", type=float, default=None, help="temperature (default: 0).") +parser.add_argument("--sampling-params", type=json.loads, default=None) +parser.add_argument("--non-thinking-sampling-params", type=json.loads, default=None) parser.add_argument( "--max-tokens", type=int, default=None, help="maximum tokens to generate per response." ) @@ -114,8 +156,26 @@ def generate(self, messages, verbose=False, **chat_template_kwargs): help="maximum total length (prompt + output). Stops synthesizing remaining turns " "when context exceeds this limit.", ) +parser.add_argument( + "--thinking-control", + choices=["soft-switch", "chat-template-kwargs"], + default="soft-switch", + help="Disable thinking with /no_think or the server's chat_template_kwargs API.", +) +parser.add_argument( + "--response-mode", + choices=["mixed", "thinking", "non-thinking"], + default="mixed", + help="Response mode for synthesized shards; mixed disables thinking on even shards.", +) args = parser.parse_args() +if args.temperature is not None and any( + params is not None and "temperature" in params + for params in (args.sampling_params, args.non_thinking_sampling_params) +): + parser.error("--temperature cannot be combined with temperature in a sampling profile") + llm = LLM(args) if args.data is None: @@ -133,6 +193,7 @@ def synthesize(data): raise ValueError( "No conversations or messages in the data. Only OAI chat data is supported." ) + sample_id = data.get("uuid") or data.get("conversation_id") or "" # Handle generation specific kwargs. enable_thinking = data.get("enable_thinking", True) @@ -146,7 +207,7 @@ def synthesize(data): if role == "system": current_messages.append(msg) elif role == "user": - if not enable_thinking: + if not enable_thinking and args.thinking_control == "soft-switch": # Copy to avoid mutating the original dataset row. msg = dict(msg) msg["content"] = msg["content"] + " /no_think" @@ -159,13 +220,30 @@ def synthesize(data): est_tokens = ctx_chars // 3 # rough char-to-token estimate if est_tokens + args.max_tokens > max_total: # Drop this user turn — context too long for another generation + print(f"[SKIP] {sample_id}: estimated context exceeds {max_total} tokens") current_messages.pop() break - new_message = llm.generate(current_messages, verbose=False) + # Thinking and non-thinking rows use their own sampling profiles. + new_message = llm.generate( + current_messages, + verbose=False, + sample_id=sample_id, + sampling_params=( + args.non_thinking_sampling_params + if not enable_thinking and args.non_thinking_sampling_params is not None + else args.sampling_params + ), + enable_thinking=( + enable_thinking if args.thinking_control == "chat-template-kwargs" else None + ), + ) if new_message is None: + current_messages.pop() break + # Preserve the mode so the training template can handle unfinished thinking. + new_message["enable_thinking"] = enable_thinking last_full_message = new_message if enable_thinking: @@ -200,7 +278,12 @@ def synthesize(data): current_messages[i] = last_full_message break - return {"messages": current_messages} + # Preserve the dataset schema for failed rows, then filter them after mapping. + synthesis_ok = last_full_message is not None + output_messages = [dict(msg) for msg in (current_messages if synthesis_ok else messages)] + for msg in output_messages: + msg.setdefault("enable_thinking", enable_thinking) + return {"messages": output_messages, "_synthesis_ok": synthesis_ok} # Support both HF Hub repo IDs and local file paths (.jsonl, .json, .parquet, etc.) @@ -250,9 +333,20 @@ def synthesize(data): print(len(shard), file_path) num_proc = min(args.num_proc, len(shard)) - if shard_id % 2 == 0: + if args.response_mode == "non-thinking" or ( + args.response_mode == "mixed" and shard_id % 2 == 0 + ): shard = shard.map(disable_thinking_column, num_proc=num_proc) - updated_shard = shard.map(synthesize, num_proc=num_proc) + # Reuse completed map-worker caches and omit rows without a generated response. + cache_dir = os.path.join(args.save, ".cache") + os.makedirs(cache_dir, exist_ok=True) + updated_shard = shard.map( + synthesize, + num_proc=num_proc, + cache_file_name=os.path.join(cache_dir, f"shard_{shard_id}.arrow"), + ) + updated_shard = updated_shard.filter(lambda row: row["_synthesis_ok"]) + updated_shard = updated_shard.remove_columns("_synthesis_ok") updated_shard.to_json(file_path) with open(done_path, "w") as done_file: done_file.write("done\n") diff --git a/tools/launcher/common/specdec/export_latest_and_vllm_smoke_test.sh b/tools/launcher/common/specdec/export_latest_and_vllm_smoke_test.sh new file mode 100644 index 00000000000..fd5b39d40fd --- /dev/null +++ b/tools/launcher/common/specdec/export_latest_and_vllm_smoke_test.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -e + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +TRAINING_DIR=${DRAFT_TRAINING_DIR:?DRAFT_TRAINING_DIR must be set} +EXPORT_DIR=${DRAFT_MODEL:?DRAFT_MODEL must be set} + +LATEST_CHECKPOINT="" +while IFS= read -r checkpoint; do + if [ -s "$checkpoint/config.json" ] && [ -s "$checkpoint/model.safetensors" ] && \ + [ -s "$checkpoint/modelopt_state.pth" ] && [ -s "$checkpoint/trainer_state.json" ]; then + LATEST_CHECKPOINT=$checkpoint + break + fi +done < <(find "$TRAINING_DIR" -maxdepth 1 -type d -name 'checkpoint-*' -print | sort -Vr) + +if [ -z "$LATEST_CHECKPOINT" ]; then + echo "ERROR: No complete checkpoint found in $TRAINING_DIR" >&2 + exit 1 +fi + +echo "Exporting latest checkpoint: $LATEST_CHECKPOINT -> $EXPORT_DIR" +EXPORT_ARGS=(--model_path "$LATEST_CHECKPOINT" --export_path "$EXPORT_DIR") +[ "${EXPORT_TRUST_REMOTE_CODE:-0}" = "1" ] && EXPORT_ARGS+=(--trust_remote_code) +python3 -m pip install --no-cache-dir 'omegaconf>=2.3.0' 'pulp<4.0' scipy +python3 -m pip install --no-cache-dir --no-deps -e modules/Model-Optimizer +python3 \ + modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ + "${EXPORT_ARGS[@]}" + +SMOKE_PROFILE=greedy SMOKE_SAMPLING_FIELDS='"temperature": 0' \ + bash "$SCRIPT_DIR/vllm_smoke_test.sh" +SMOKE_PROFILE=sampled SMOKE_SAMPLING_FIELDS='"temperature": 1.0, "top_p": 0.95, "top_k": 20, "presence_penalty": 1.5' \ + bash "$SCRIPT_DIR/vllm_smoke_test.sh" diff --git a/tools/launcher/common/specdec/vllm_smoke_test.sh b/tools/launcher/common/specdec/vllm_smoke_test.sh index 8028a570a8e..9d22bf01c2b 100644 --- a/tools/launcher/common/specdec/vllm_smoke_test.sh +++ b/tools/launcher/common/specdec/vllm_smoke_test.sh @@ -26,8 +26,13 @@ # NUM_SPEC_TOKENS — number of speculative tokens (default: 15) # TP_SIZE — tensor parallel size (default: 1) # VLLM_PORT — server port (default: 8000) +# MAX_MODEL_LEN — optional maximum context length +# ENFORCE_EAGER — set to "1" to disable compile and CUDA graphs +# SERVER_STARTUP_TIMEOUT — server startup timeout in seconds (default: 600) # REASONING_PARSER — reasoning parser (e.g., "qwen3" for Qwen3.5) # DISABLE_PREFIX_CACHING — set to "1" to disable prefix caching +# SMOKE_PROFILE — profile label printed in results (default: "greedy") +# SMOKE_SAMPLING_FIELDS — JSON sampling fields appended to each request SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" source ${SCRIPT_DIR}/../service_utils.sh 2>/dev/null || true @@ -59,6 +64,7 @@ METHOD=${SPEC_METHOD:-eagle} NUM_SPEC=${NUM_SPEC_TOKENS:-15} PORT=${VLLM_PORT:-8000} TP=${TP_SIZE:-1} +SERVER_STARTUP_TIMEOUT=${SERVER_STARTUP_TIMEOUT:-600} echo "=== vLLM Speculative Decoding Smoke Test ===" echo "Method: ${METHOD}" @@ -82,6 +88,12 @@ fi if [ "${DISABLE_PREFIX_CACHING:-}" = "1" ]; then OPTIONAL_ARGS="${OPTIONAL_ARGS} --no-enable-prefix-caching" fi +if [ -n "${MAX_MODEL_LEN:-}" ]; then + OPTIONAL_ARGS="${OPTIONAL_ARGS} --max-model-len ${MAX_MODEL_LEN}" +fi +if [ "${ENFORCE_EAGER:-}" = "1" ]; then + OPTIONAL_ARGS="${OPTIONAL_ARGS} --enforce-eager" +fi # Start vLLM server (capture output for regression check parsing) VLLM_LOG=$(mktemp /tmp/vllm_server_XXXXXX.log) @@ -105,7 +117,7 @@ SERVER_PID=$! # Wait for server echo "Waiting for vLLM server..." -for i in $(seq 1 180); do +for i in $(seq 1 "$SERVER_STARTUP_TIMEOUT"); do if curl -s http://localhost:${PORT}/health > /dev/null 2>&1; then echo "Server ready after ${i}s" break @@ -122,8 +134,11 @@ fi # Run quick test prompts using chat completions API MAX_TOKENS=${MAX_OUTPUT_TOKENS:-1024} +SMOKE_PROFILE=${SMOKE_PROFILE:-greedy} +SMOKE_SAMPLING_FIELDS=${SMOKE_SAMPLING_FIELDS:-'"temperature": 0'} echo "" -echo "=== Test Prompts (max_tokens=${MAX_TOKENS}) ===" +echo "=== Test Prompts (${SMOKE_PROFILE}, max_tokens=${MAX_TOKENS}) ===" +echo "Sampling: {${SMOKE_SAMPLING_FIELDS}}" PASS=0 FAIL=0 TOTAL_TOKENS=0 @@ -142,17 +157,17 @@ for PROMPT in \ START=$(date +%s%N) RESULT=$(curl -s http://localhost:${PORT}/v1/chat/completions \ -H "Content-Type: application/json" \ - -d "{\"model\": \"${MODEL}\", \"messages\": [{\"role\": \"user\", \"content\": \"${PROMPT}\"}], \"max_tokens\": ${MAX_TOKENS}, \"temperature\": 0}" \ + -d "{\"model\": \"${MODEL}\", \"messages\": [{\"role\": \"user\", \"content\": \"${PROMPT}\"}], \"max_tokens\": ${MAX_TOKENS}, ${SMOKE_SAMPLING_FIELDS}}" \ 2>/dev/null) END=$(date +%s%N) - ELAPSED=$(echo "scale=2; ($END - $START) / 1000000000" | bc 2>/dev/null || echo "0") + ELAPSED=$(awk "BEGIN {printf \"%.2f\", ($END - $START) / 1000000000}") TOKENS=$(echo "$RESULT" | python3 -c "import json,sys; r=json.load(sys.stdin); print(r.get('usage',{}).get('completion_tokens',0))" 2>/dev/null) if [ -n "$TOKENS" ] && [ "$TOKENS" -gt 0 ] 2>/dev/null; then - TPS=$(echo "scale=1; $TOKENS / $ELAPSED" | bc 2>/dev/null || echo "?") + TPS=$(awk "BEGIN {printf \"%.1f\", $TOKENS / $ELAPSED}") echo " PASS: ${TOKENS} tokens in ${ELAPSED}s (${TPS} tok/s) — \"${PROMPT:0:50}...\"" PASS=$((PASS + 1)) TOTAL_TOKENS=$((TOTAL_TOKENS + TOKENS)) - TOTAL_TIME=$(echo "$TOTAL_TIME + $ELAPSED" | bc 2>/dev/null || echo "0") + TOTAL_TIME=$(awk "BEGIN {printf \"%.2f\", $TOTAL_TIME + $ELAPSED}") else echo " FAIL: \"${PROMPT}\"" echo " Response: $(echo "$RESULT" | head -c 200)" @@ -163,7 +178,7 @@ done echo "" echo "Results: ${PASS} passed, ${FAIL} failed" if [ "$TOTAL_TOKENS" -gt 0 ] 2>/dev/null; then - AVG_TPS=$(echo "scale=1; $TOTAL_TOKENS / $TOTAL_TIME" | bc 2>/dev/null || echo "?") + AVG_TPS=$(awk "BEGIN {printf \"%.1f\", $TOTAL_TOKENS / $TOTAL_TIME}") echo "Total: ${TOTAL_TOKENS} tokens in ${TOTAL_TIME}s (${AVG_TPS} tok/s avg)" fi diff --git a/tools/launcher/common/vllm/query.sh b/tools/launcher/common/vllm/query.sh index c3a2a784bb0..3145aa4dfbc 100755 --- a/tools/launcher/common/vllm/query.sh +++ b/tools/launcher/common/vllm/query.sh @@ -137,9 +137,10 @@ done pip3 install -q datasets openai 2>/dev/null || true echo "Running: python3 common/query.py http://localhost:8000/v1 ${MODEL} ${QUERY_ARGS[*]}" python3 common/query.py http://localhost:8000/v1 "${MODEL}" "${QUERY_ARGS[@]}" +QUERY_EXIT=$? echo "Main process exit" kill $SERVER_PID wait $SERVER_PID 2>/dev/null || true -exit 0 +exit $QUERY_EXIT diff --git a/tools/launcher/core.py b/tools/launcher/core.py index 44f3efab86d..02b37eaa22a 100644 --- a/tools/launcher/core.py +++ b/tools/launcher/core.py @@ -62,6 +62,9 @@ def get_default_env(experiment_title=None): "LAUNCH_SCRIPT": "python", **specdec_s3, } + # Forward the launcher-shell topology override into the Slurm job. + if serve_nodes := os.getenv("SERVE_NODES"): + slurm_env["SERVE_NODES"] = serve_nodes local_env = { "TRITON_CACHE_DIR": os.getenv("TRITON_CACHE_DIR", f"/{title}/triton-cache"), "HF_HOME": os.getenv("HF_HOME", f"/{title}/hf-cache"), @@ -91,6 +94,12 @@ def set_slurm_config_type(cls): SandboxTask2, SandboxTask3, SandboxTask4, + SandboxTask5, + SandboxTask6, + SandboxTask7, + SandboxTask8, + SandboxTask9, + SandboxTask10, ): task_cls.__dataclass_fields__["slurm_config"].type = cls task_cls.__annotations__["slurm_config"] = cls @@ -156,6 +165,36 @@ class SandboxTask4(SandboxTask): """Task slot 4 in a pipeline.""" +@dataclass +class SandboxTask5(SandboxTask): + """Task slot 5 in a pipeline.""" + + +@dataclass +class SandboxTask6(SandboxTask): + """Task slot 6 in a pipeline.""" + + +@dataclass +class SandboxTask7(SandboxTask): + """Task slot 7 in a pipeline.""" + + +@dataclass +class SandboxTask8(SandboxTask): + """Task slot 8 in a pipeline.""" + + +@dataclass +class SandboxTask9(SandboxTask): + """Task slot 9 in a pipeline.""" + + +@dataclass +class SandboxTask10(SandboxTask): + """Task slot 10 in a pipeline.""" + + def create_task_from_yaml(yaml_file, factory_lookup): """Create a SandboxTask from a YAML config file. @@ -241,6 +280,12 @@ class SandboxPipeline: task_2: SandboxTask2 = None task_3: SandboxTask3 = None task_4: SandboxTask4 = None + task_5: SandboxTask5 = None + task_6: SandboxTask6 = None + task_7: SandboxTask7 = None + task_8: SandboxTask8 = None + task_9: SandboxTask9 = None + task_10: SandboxTask10 = None tasks: list[SandboxTask] = None assets: list[str] = None # HF repo paths (relative to hf_local) to verify before submission @@ -259,7 +304,7 @@ def __post_init__(self): """Collect tasks from slots/configs and resolve <> references.""" if self.tasks is None: self.tasks = [] - for i in range(5): + for i in range(11): task = getattr(self, f"task_{i}", None) if task is not None: self.tasks += [task] @@ -557,6 +602,11 @@ def build_slurm_executor( if segment is not None: optional_kwargs["segment"] = segment + additional_parameters = dict(getattr(slurm_config, "additional_parameters", None) or {}) + dependency = getattr(slurm_config, "dependency", None) + if dependency: + additional_parameters["dependency"] = dependency + executor = run.SlurmExecutor( account=slurm_config.account, partition=slurm_config.partition, @@ -573,9 +623,7 @@ def build_slurm_executor( retries=0, packager=packager, srun_args=slurm_config.srun_args, - # Copy into a fresh dict so the requeue mutation below doesn't leak back into - # the shared slurm_config.additional_parameters. - additional_parameters=dict(getattr(slurm_config, "additional_parameters", None) or {}), + additional_parameters=additional_parameters, **optional_kwargs, ) if getattr(slurm_config, "requeue", False): @@ -840,7 +888,7 @@ def run_jobs( if task.reqs or task.reqs_file: pkgs = ["-r", shlex.quote(task.reqs_file)] if task.reqs_file else [] pkgs += [shlex.quote(tok) for tok in shlex.split(task.reqs or "")] - install = "python -m pip install " + " ".join(pkgs) + install = "python3 -m pip install " + " ".join(pkgs) # On Slurm, srun runs this inline on every rank (ntasks_per_node), so install # once per node on local rank 0 behind a filesystem barrier — concurrent pip on # one node corrupts the env. The marker lives in the working dir (/nemo_run/code, diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dspark.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dspark.yaml new file mode 100644 index 00000000000..9fb723008cb --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dspark.yaml @@ -0,0 +1,161 @@ +# DSpark streaming speculative-decoding training for Qwen3-8B. +# +# DSpark = the DFlash backbone + a lightweight sequential (Markov) head + a +# confidence head, generating a causal block semi-autoregressively; see +# dspark.yaml for the head architecture and the three-term loss +# (ce_alpha*CE + l1_alpha*TVD + conf_alpha*confidence_BCE). +# +# Runs the shared streaming pipeline (common/eagle3/train_eagle_streaming.sh): +# a live `vllm serve` runs the target model and ships its hidden states to the +# trainer over NIXL RDMA — no offline hidden-state dump, no disk round-trip. +# Same transport as hf_streaming_dflash.yaml; the recipe changes +# (dflash.yaml -> dspark.yaml), which is what adds the Markov + confidence heads. +# +# Qwen3-8B is the small end of this pipeline — it fits on 2 nodes / 1 GPU each, +# which makes it the recommended first run before scaling to a large MoE target +# (see the Kimi-K2.6 and MiniMax-M3 multi-node examples, same pipeline). +# +# 3-step pipeline: +# task_0: Build input conversations (jsonl) +# task_1: Streaming train — node 0 vllm serve, node 1 trainer +# task_2: vLLM smoke test with DSpark speculative decoding +# +# Tasks share /scratchspace to pass artifacts; the export lands in +# /scratchspace/export. +# +# NOTE ON AR EVAL: dspark.yaml keeps estimate_ar=false and this yaml disables the +# periodic AR validation. The training-time eval runs the DFlash backbone only +# (the Markov head is not applied), so an AR number would describe the backbone, +# not the trained DSpark model. Measure acceptance with task_2 or specdec_bench. +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_streaming_dspark.yaml --yes + +job_name: Qwen3-8B_DSpark_streaming +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/Qwen/Qwen3-8B + + # Step 1: Build /scratchspace/data/train.jsonl. + # + # To train on a target-synthesized corpus instead (recommended — matching the + # target's own output distribution raises acceptance length), run hf_synth.yaml + # first, set `skip: true` here, and point data.data_path at its output_dir. + # eagle_utils also accepts a directory of *.jsonl shards directly. + task_0: + script: common/eagle3/make_dataset.sh + args: + - -f modules/Model-Optimizer/examples/dataset/example_data_config.yaml + - --full-conversations + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 + + # Step 2: Streaming DSpark training — node 0 vllm serve, node 1 trainer. + # + # The draft dims, Markov/confidence head config and loss weights come from + # dspark.yaml, whose defaults are already Qwen3-tuned (num_attention_heads=32, + # num_key_value_heads=8, head_dim=128, intermediate_size=12288). Unlike the + # Kimi/M3 examples there is no dims override block here for exactly that + # reason — retargeting to another base means setting them explicitly, since + # the draft does NOT inherit the base's GQA/FFN dims. + # + # Qwen3-8B also needs no rope_theta pin (the draft default 1e6 matches the + # base) and no trust_remote_code (stock architecture, no custom modeling file). + task_1: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml + - model.model_name_or_path=<> + # Streaming is selected by data.streaming_server_url, which the launcher + # script injects once the serve is up — not by a data.mode= key. + - data.data_path=/scratchspace/data/train.jsonl + - training.output_dir=/scratchspace/dspark + - training.training_seq_len=4096 + - training.disable_tqdm=true + # Backbone-only AR eval is misleading for DSpark (see header) — disable it. + - training.ar_validate_steps=500000 + - training.num_train_epochs=1 + - training.save_steps=1000 + - training.logging_steps=20 + # Batch size and LR are left at the dspark.yaml defaults + # (per_device_train_batch_size=1, learning_rate=6e-4, warmup_ratio=0.04), + # which target a from-scratch draft on a single GPU. The Kimi/M3 examples + # override these to a large batch and a gentle 1e-4 because they run 8 + # GPUs per node and warm-start a big backbone — do not copy those values + # here; at training_seq_len=4096 on 1 GPU a batch of 4 will OOM. + # The streaming corpus is prompt-only (the serve generates the response and + # we capture ITS hidden states), so there is no assistant span to mask -> + # train over the full sequence. This also means no {% generation %}-tagged + # chat template is needed here, unlike the online/offline examples. + - training.answer_only_loss=false + # The vLLM serve container has no tensorboard -> trainer init crash. + - training.report_to=none + # Semi-AR generation block; must divide training_seq_len. dspark.yaml ships + # 16 (the Kimi/M3 runs use 8). No dflash_loss_decay_factor override here: + # dflash_loss_objective defaults to 'dpace', which derives per-position + # weights from draft confidence and ignores the decay gamma outright. + - dflash.dflash_block_size=16 + - dflash.dflash_num_anchors=512 + # Qwen3 has no native mask token; 151669 is an unused id. + - dflash.dflash_mask_token_id=151669 + - dflash.dflash_architecture_config.num_hidden_layers=5 + environment: + - HF_MODEL_CKPT: <> + # 5 aux capture ids = the draft's target_layer_ids+1, plus the true final + # hidden (36). build_target_layer_ids(36,5)=[1,9,17,25,33] -> +1. + # Capturing the second-to-last layer instead of the true final one trains + # fine and silently caps acceptance length, so keep the final id honest. + # No spaces: nemo_run emits `export FOO=value` unquoted, so a space splits. + - EAGLE_CAPTURE_IDS: "[2,10,18,26,34,36]" + - SERVE_TP: "1" + - STREAMING_NUM_WORKERS: "4" + # DSpark exports a custom modeling file; export must trust remote code. + - EXPORT_EXTRA_ARGS: "--trust_remote_code" + - SERVE_MAX_MODEL_LEN: "4160" + - SERVE_MAX_NUM_SEQS: "32" + - SERVE_READY_TIMEOUT: "1800" + # RDMA transport is UCX (InfiniBand) by default. On AWS EFA, uncomment — + # and note UCX SEGFAULTS at agent init on EFA nodes (it detects the EFA + # devices), so LIBFABRIC is required there even for single-node runs: + # - NIXL_BACKENDS: "LIBFABRIC" + # - FI_PROVIDER: "efa" + # - NCCL_IB_DISABLE: "1" + slurm_config: + _factory_: "slurm_factory" + nodes: 2 + ntasks_per_node: 1 + gpus_per_node: 1 + # Needs the aux-capture fix (vllm#46788), in-tree in recent nightlies; + # without it the final capture id is off by one and caps acceptance length. + container: vllm/vllm-openai:latest + + # Step 3: vLLM smoke test (DSpark, uses the exported checkpoint from training). + # This is the acceptance-length number that reflects the trained Markov head — + # the training-time eval cannot produce it (see the AR-eval note in the header). + # + # REQUIRES a vLLM build with DSpark speculative-decoding support. If your image + # predates it, `--speculative-config '{"method": "dspark", ...}'` is rejected at + # startup; pin a newer nightly, or set `skip: true` here and evaluate the + # exported drafter with the offline specdec_bench harness instead. + task_2: + script: common/specdec/vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_MODEL: /scratchspace/export + - SPEC_METHOD: "dspark" + - NUM_SPEC_TOKENS: "7" + - MIN_ACCEPTANCE_LENGTH: "1.2" + slurm_config: + _factory_: "slurm_factory" + container: vllm/vllm-openai:nightly + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 diff --git a/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/chat_template_train.jinja b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/chat_template_train.jinja new file mode 100644 index 00000000000..5467436ea79 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/chat_template_train.jinja @@ -0,0 +1,161 @@ +{# Official Qwen3.5-35B-A3B template with assistant output wrapped in generation tags. #} +{%- set image_count = namespace(value=0) %} +{%- set video_count = namespace(value=0) %} +{%- macro render_content(content, do_vision_count, is_system_content=false) %} + {%- if content is string %} + {{- content }} + {%- elif content is iterable and content is not mapping %} + {%- for item in content %} + {%- if 'image' in item or 'image_url' in item or item.type == 'image' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain images.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set image_count.value = image_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Picture ' ~ image_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|image_pad|><|vision_end|>' }} + {%- elif 'video' in item or item.type == 'video' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain videos.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set video_count.value = video_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Video ' ~ video_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|video_pad|><|vision_end|>' }} + {%- elif 'text' in item %} + {{- item.text }} + {%- else %} + {{- raise_exception('Unexpected item type in content.') }} + {%- endif %} + {%- endfor %} + {%- elif content is none or content is undefined %} + {{- '' }} + {%- else %} + {{- raise_exception('Unexpected content type.') }} + {%- endif %} +{%- endmacro %} +{%- if not messages %} + {{- raise_exception('No messages provided.') }} +{%- endif %} +{%- if tools and tools is iterable and tools is not mapping %} + {{- '<|im_start|>system\n' }} + {{- "# Tools\n\nYou have access to the following functions:\n\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n" }} + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n' }} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {%- if content %} + {{- '\n\n' + content }} + {%- endif %} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- else %} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" %} + {%- set content = render_content(message.content, false)|trim %} + {%- if not(content.startswith('') and content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if ns.multi_step_tool %} + {{- raise_exception('No user query found in messages.') }} +{%- endif %} +{%- for message in messages %} + {%- set content = render_content(message.content, true)|trim %} + {%- if message.role == "system" %} + {%- if not loop.first %} + {{- raise_exception('System message must be at the beginning.') }} + {%- endif %} + {%- elif message.role == "user" %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- set unfinished_thinking = message.enable_thinking is defined and message.enable_thinking and '' not in content %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- set reasoning_content = reasoning_content|trim %} + {{- '<|im_start|>' + message.role + '\n' }} + {%- generation %} + {%- if loop.index0 > ns.last_query_index and unfinished_thinking %} + {{- '\n' + content }} + {%- elif loop.index0 > ns.last_query_index %} + {{- '\n' + reasoning_content + '\n\n\n' + content }} + {%- else %} + {{- content }} + {%- endif %} + {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- if loop.first %} + {%- if content|trim %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n\n' }} + {%- endif %} + {%- else %} + {{- '\n\n\n' }} + {%- endif %} + {%- if tool_call.arguments is defined %} + {%- for args_name, args_value in tool_call.arguments|items %} + {{- '\n' }} + {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %} + {{- args_value }} + {{- '\n\n' }} + {%- endfor %} + {%- endif %} + {{- '\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- endgeneration %} + {%- elif message.role == "tool" %} + {%- if loop.previtem and loop.previtem.role != "tool" %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if not loop.last and loop.nextitem.role != "tool" %} + {{- '<|im_end|>\n' }} + {%- elif loop.last %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- else %} + {{- raise_exception('Unexpected message role.') }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n' }} + {%- endif %} +{%- endif %} diff --git a/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_finetuning_multi_node.yaml b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_finetuning_multi_node.yaml new file mode 100644 index 00000000000..f9b1402cde7 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_finetuning_multi_node.yaml @@ -0,0 +1,95 @@ +# Qwen3.5-35B-A3B DSpark continuation training on pretokenized conversations. + +job_name: qwen3.5-35b-a3b-dspark-finetuning + +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3.5-35B-A3B + hf_data: /hf-local/modelopt/qwen3.5-35b-a3b-dspark-finetuning + draft_model: /cicd/qwen3.5-35b-a3b-dspark/training/checkpoint- + output_dir: /cicd/qwen3.5-35b-a3b-dspark-finetuning + + task_1: + script: common/eagle3/train_eagle_streaming.sh + reqs: wandb==0.22.3 + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - data.data_path=<> + - data.chat_template=examples/Qwen/Qwen3.5-35B-A3B/chat_template_train.jinja + - training.output_dir=<>/training + - training.resume_from_checkpoint=<> + - training.resume_model_only=true + - training.training_seq_len=32768 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=5 + - training.per_device_train_batch_size=1 + - training.gradient_accumulation_steps=4 + - training.learning_rate=1.0e-4 + - training.warmup_ratio=0.04 + - training.lr_scheduler_type=cosine_with_min_lr + - training.lr_scheduler_kwargs.min_lr=3.0e-5 + - training.save_steps=200 + - training.logging_steps=10 + - training.estimate_ar=true + - training.answer_only_loss=true + - training.report_to=wandb + - dflash.dflash_block_size=8 + - dflash.dflash_num_anchors=4096 + - dflash.dflash_loss_decay_factor=4.0 + - dflash.dflash_mask_token_id=248077 + - dflash.dflash_swa_window_size=4096 + - dflash.dflash_architecture_config.num_hidden_layers=6 + - dflash.dflash_architecture_config.num_attention_heads=32 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.head_dim=128 + - dflash.dflash_architecture_config.intermediate_size=6144 + environment: + - HF_MODEL_CKPT: <> + - EAGLE_CAPTURE_IDS: "[2,9,16,24,31,38,40]" + - SERVE_NODES: "8" + - SERVE_TP: "8" + - STREAMING_NUM_WORKERS: "4" + - HS_POOL_SLOTS: "32" + - HS_MAX_TOKENS: "32768" + - WANDB_PROJECT: modelopt-dspark-training + - WANDB_NAME: qwen3.5-35b-a3b-dspark-finetuning + - WANDB_RUN_ID: qwen35-35b-a3b-dspark-finetuning + - WANDB_RESUME: allow + - WANDB_DIR: <>/wandb + - EXPORT_PATH: <>/export + - EXPORT_EXTRA_ARGS: --trust_remote_code + - SERVE_MAX_MODEL_LEN: "32832" + - SERVE_MAX_NUM_SEQS: "32" + - SERVE_READY_TIMEOUT: "1800" + slurm_config: + _factory_: slurm_factory + nodes: 16 + ntasks_per_node: 1 + gpus_per_node: 8 + dependency: singleton + time: 04:00:00 + container: vllm/vllm-openai:v0.27.1 + + task_2: + script: common/specdec/export_latest_and_vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_TRAINING_DIR: <>/training + - DRAFT_MODEL: <>/export + - EXPORT_TRUST_REMOTE_CODE: "1" + - SPEC_METHOD: dspark + - NUM_SPEC_TOKENS: "7" + - TP_SIZE: "8" + - MAX_MODEL_LEN: "32832" + - MAX_OUTPUT_TOKENS: "128" + - MIN_ACCEPTANCE_LENGTH: "1.2" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + time: 02:00:00 + container: vllm/vllm-openai:v0.27.1 diff --git a/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_multi_node.yaml b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_multi_node.yaml new file mode 100644 index 00000000000..49fd28ba7c1 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/hf_streaming_dspark_multi_node.yaml @@ -0,0 +1,91 @@ +# Qwen3.5-35B-A3B DSpark training: 16 target nodes + 16 trainer nodes. + +job_name: qwen3.5-35b-a3b-dspark-training + +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3.5-35B-A3B + hf_data: /hf-local/modelopt/qwen3.5-35b-a3b-dspark-synthesis + output_dir: /cicd/qwen3.5-35b-a3b-dspark + + task_1: + script: common/eagle3/train_eagle_streaming.sh + reqs: wandb==0.22.3 + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - data.data_path=<> + - data.chat_template=examples/Qwen/Qwen3.5-35B-A3B/chat_template_train.jinja + - training.output_dir=<>/training + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=5 + - training.per_device_train_batch_size=1 + - training.gradient_accumulation_steps=4 + - training.learning_rate=6.0e-4 + - training.warmup_ratio=0.04 + - training.lr_scheduler_type=cosine_with_min_lr + - training.lr_scheduler_kwargs.min_lr=3.0e-5 + - training.save_steps=200 + - training.logging_steps=10 + - training.estimate_ar=true + - training.answer_only_loss=true + - training.report_to=wandb + - dflash.dflash_block_size=8 + - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_decay_factor=4.0 + - dflash.dflash_mask_token_id=248077 + - dflash.dflash_swa_window_size=4096 + - dflash.dflash_architecture_config.num_hidden_layers=6 + - dflash.dflash_architecture_config.num_attention_heads=32 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.head_dim=128 + - dflash.dflash_architecture_config.intermediate_size=6144 + environment: + - HF_MODEL_CKPT: <> + - EAGLE_CAPTURE_IDS: "[2,9,16,24,31,38,40]" + - SERVE_NODES: "16" + - SERVE_TP: "8" + - STREAMING_NUM_WORKERS: "4" + - HS_POOL_SLOTS: "32" + - WANDB_PROJECT: modelopt-dspark-training + - WANDB_NAME: qwen3.5-35b-a3b-dspark + - WANDB_RUN_ID: qwen35-35b-a3b-dspark + - WANDB_RESUME: allow + - WANDB_DIR: <>/wandb + - EXPORT_PATH: <>/export + - EXPORT_EXTRA_ARGS: --trust_remote_code + - SERVE_MAX_MODEL_LEN: "4160" + - SERVE_MAX_NUM_SEQS: "32" + - SERVE_READY_TIMEOUT: "1800" + slurm_config: + _factory_: slurm_factory + nodes: 32 + ntasks_per_node: 1 + gpus_per_node: 8 + dependency: singleton + time: 04:00:00 + container: vllm/vllm-openai:v0.27.1 + + task_2: + script: common/specdec/export_latest_and_vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_TRAINING_DIR: <>/training + - DRAFT_MODEL: <>/export + - EXPORT_TRUST_REMOTE_CODE: "1" + - SPEC_METHOD: dspark + - NUM_SPEC_TOKENS: "7" + - TP_SIZE: "8" + - MAX_MODEL_LEN: "4160" + - MAX_OUTPUT_TOKENS: "128" + - MIN_ACCEPTANCE_LENGTH: "1.2" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + time: 02:00:00 + container: vllm/vllm-openai:v0.27.1 diff --git a/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/hf_synth.yaml b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/hf_synth.yaml new file mode 100644 index 00000000000..0ca6cfc77c2 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/hf_synth.yaml @@ -0,0 +1,57 @@ +# Synthesize Qwen3.5-35B-A3B conversations dataset for DSpark training. + +job_name: qwen3.5-35b-a3b-dspark-synthesis + +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3.5-35B-A3B + hf_data: /hf-local/nvidia/Nemotron-Post-Training-Dataset-v2/default.jsonl + output_dir: /hf-local/modelopt/qwen3.5-35b-a3b-dspark-synthesis + + task_0: + script: common/vllm/query.sh + args: + - --model + - <> + - --tensor-parallel-size + - "8" + - --trust-remote-code + - --language-model-only + - --enable-prefix-caching + - --mamba-cache-mode + - align + - --gpu-memory-utilization + - "0.95" + - --max-model-len + - "8192" + - -- + - --data + - <> + - --save + - <> + - --num-shards + - "384" + - --shard-id + - $SLURM_ARRAY_TASK_ID + - --thinking-control + - chat-template-kwargs + - --response-mode + - thinking + - "--sampling-params '{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}'" + - --num-proc + - "64" + - --max-tokens + - "4096" + - --max-total-length + - "8192" + environment: + - VLLM_STARTUP_TIMEOUT: "1800" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + time: 04:00:00 + container: vllm/vllm-openai:v0.27.1 + array: "0-383%192" + requeue: true diff --git a/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/prepare-swe-data.py b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/prepare-swe-data.py new file mode 100644 index 00000000000..e9be7a57e7a --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/prepare-swe-data.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Convert Prime-RL train traces into exact pretokenized DSpark conversations. + +Usage: + python prepare-swe-data.py +""" + +from __future__ import annotations + +import argparse +import json +import os +from collections import Counter +from pathlib import Path + + +def branch_nodes(nodes: list[dict]) -> list[dict]: + """Return the single root-to-leaf branch stored by these SWE rollouts.""" + parents = {node["parent"] for node in nodes if node["parent"] is not None} + leaves = [index for index in range(len(nodes)) if index not in parents] + if len(leaves) != 1: + raise ValueError(f"expected one branch, found {len(leaves)}") + path = [] + node_id = leaves[0] + while node_id is not None: + path.append(nodes[node_id]) + node_id = nodes[node_id]["parent"] + return path[::-1] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("rollout_root", type=Path) + parser.add_argument("output_dir", type=Path) + parser.add_argument("--first-step", type=int, default=1) + parser.add_argument("--last-step", type=int, default=300) + parser.add_argument("--records-per-shard", type=int, default=1024) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.output_dir.exists(): + raise FileExistsError(args.output_dir) + staging = args.output_dir.with_name(f".{args.output_dir.name}.tmp-{os.getpid()}") + staging.mkdir(parents=True) + + counts = Counter() + policy_versions = Counter() + tasks = Counter() + trace_ids: set[str] = set() + shard = None + shard_index = -1 + + try: + for step in range(args.first_step, args.last_step + 1): + path = args.rollout_root / f"step_{step}" / "train" / "all" / "traces.jsonl" + if not path.is_file(): + raise FileNotFoundError(path) + counts["input_bytes"] += path.stat().st_size + with path.open() as source: + for line in source: + counts["train_traces"] += 1 + trace = json.loads(line) + if not trace.get("ok", False): + counts["excluded"] += 1 + continue + trace_id = trace["id"] + if trace_id in trace_ids: + raise ValueError(f"duplicate trace id: {trace_id}") + trace_ids.add(trace_id) + + nodes = branch_nodes(trace["nodes"]) + token_ids = [token for node in nodes for token in node["token_ids"]] + loss_mask = [int(mask) for node in nodes for mask in node["mask"]] + if len(token_ids) != len(loss_mask) or not any(loss_mask): + raise ValueError(f"invalid token/mask data: {trace_id}") + + if counts["converted"] % args.records_per_shard == 0: + if shard is not None: + shard.close() + shard_index += 1 + shard = (staging / f"shard_{shard_index:05d}.jsonl").open("w") + record = { + "conversation_id": trace_id, + "token_ids": token_ids, + "loss_mask": loss_mask, + "step": step, + "task_id": trace["task"]["data"]["name"], + "policy_version": trace["info"]["policy_version"], + } + shard.write(json.dumps(record, separators=(",", ":")) + "\n") + counts["converted"] += 1 + counts["tokens"] += len(token_ids) + counts["supervised_tokens"] += sum(loss_mask) + policy_versions[str(record["policy_version"])] += 1 + tasks[record["task_id"]] += 1 + if step % 25 == 0 or step == args.last_step: + print(f"step {step}: {counts['converted']} converted", flush=True) + if shard is not None: + shard.close() + shard = None + + manifest = { + "steps": [args.first_step, args.last_step], + "counts": dict(counts), + "num_shards": shard_index + 1, + "num_tasks": len(tasks), + "policy_versions": dict(sorted(policy_versions.items(), key=lambda item: int(item[0]))), + "task_trace_counts": dict(sorted(tasks.items())), + } + (staging / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + staging.rename(args.output_dir) + print(json.dumps(manifest["counts"], indent=2)) + print(f"output: {args.output_dir}") + except BaseException: + if shard is not None: + shard.close() + raise + + +if __name__ == "__main__": + main() diff --git a/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/specdec_bench_tp2_c1.yaml b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/specdec_bench_tp2_c1.yaml new file mode 100644 index 00000000000..cb8b94f31c3 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/specdec_bench_tp2_c1.yaml @@ -0,0 +1,322 @@ +# SPEED-Bench: Qwen3.5-35B-A3B Base / DSpark7 / MTP3 / MTP7 / DFlash8 (TP2 / C1 / T0+T1 / MRV2) + +job_name: qwen3.5-35b-a3b-specdec-bench-tp2-c1 + +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3.5-35B-A3B + draft_model: /cicd/qwen3.5-35b-a3b-dspark/export + output_dir: /cicd/qwen3.5-35b-a3b-dspark-benchmark-tp2-c1 + + # Draft semantics: DSpark7 drafts 7 speculative tokens; DFlash8 drafts 7 speculative tokens (includes 1 anchor). + + # Prepare SPEED-Bench data + task_0: + inline: >- + python3 -m pip install -r modules/Model-Optimizer/examples/specdec_bench/requirements.txt && + (test -f /hf-local/nvidia/SPEED-Bench/speed/qualitative/test.parquet || + python3 modules/Model-Optimizer/examples/specdec_bench/prepare_data.py + --dataset speed --config qualitative --output_dir /hf-local/nvidia/SPEED-Bench) + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # Base T1 C1 + task_1: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm NONE + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/base/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # Base T0 C1 + task_2: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm NONE + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/base/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # DSpark7 T1 C1 + task_3: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DSPARK + - --draft_model_dir <> + - --block_size 7 + - --max_seq_len 16384 + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5},\"engine_args\":{\"draft_sample_method\":\"probabilistic\"}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dspark7/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # DSpark7 T0 C1 + task_4: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DSPARK + - --draft_model_dir <> + - --block_size 7 + - --max_seq_len 16384 + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dspark7/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # MTP3 T1 C1 + task_5: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 3 + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp3/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # MTP3 T0 C1 + task_6: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 3 + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp3/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # MTP7 T1 C1 + task_7: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 7 + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp7/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # MTP7 T0 C1 + task_8: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 7 + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp7/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # DFlash8 T1 C1 + task_9: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DFLASH + - --draft_model_dir /hf-local/z-lab/Qwen3.5-35B-A3B-DFlash + - --block_size 8 + - --max_seq_len 16384 + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5},\"engine_args\":{\"draft_sample_method\":\"probabilistic\"}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dflash8/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # DFlash8 T0 C1 + task_10: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DFLASH + - --draft_model_dir /hf-local/z-lab/Qwen3.5-35B-A3B-DFlash + - --block_size 8 + - --max_seq_len 16384 + - --tp_size 2 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dflash8/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 04:00:00 + container: vllm/vllm-openai:nightly diff --git a/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/specdec_bench_tp2_c32.yaml b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/specdec_bench_tp2_c32.yaml new file mode 100644 index 00000000000..2b8c1c36590 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-35B-A3B/specdec_bench_tp2_c32.yaml @@ -0,0 +1,322 @@ +# SPEED-Bench: Qwen3.5-35B-A3B Base / DSpark7 / MTP3 / MTP7 / DFlash8 (TP2 / C32 / T0+T1 / MRV2) + +job_name: qwen3.5-35b-a3b-specdec-bench-tp2-c32 + +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3.5-35B-A3B + draft_model: /cicd/qwen3.5-35b-a3b-dspark/export + output_dir: /cicd/qwen3.5-35b-a3b-dspark-benchmark-tp2-c32 + + # Draft semantics: DSpark7 drafts 7 speculative tokens; DFlash8 drafts 7 speculative tokens (includes 1 anchor). + + # Prepare SPEED-Bench data + task_0: + inline: >- + python3 -m pip install -r modules/Model-Optimizer/examples/specdec_bench/requirements.txt && + (test -f /hf-local/nvidia/SPEED-Bench/speed/qualitative/test.parquet || + python3 modules/Model-Optimizer/examples/specdec_bench/prepare_data.py + --dataset speed --config qualitative --output_dir /hf-local/nvidia/SPEED-Bench) + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # Base T1 C32 + task_1: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm NONE + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/base/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # Base T0 C32 + task_2: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm NONE + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/base/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # DSpark7 T1 C32 + task_3: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DSPARK + - --draft_model_dir <> + - --block_size 7 + - --max_seq_len 16384 + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5},\"engine_args\":{\"draft_sample_method\":\"probabilistic\"}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dspark7/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # DSpark7 T0 C32 + task_4: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DSPARK + - --draft_model_dir <> + - --block_size 7 + - --max_seq_len 16384 + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dspark7/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # MTP3 T1 C32 + task_5: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 3 + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp3/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # MTP3 T0 C32 + task_6: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 3 + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp3/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # MTP7 T1 C32 + task_7: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 7 + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp7/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # MTP7 T0 C32 + task_8: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 7 + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp7/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # DFlash8 T1 C32 + task_9: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DFLASH + - --draft_model_dir /hf-local/z-lab/Qwen3.5-35B-A3B-DFlash + - --block_size 8 + - --max_seq_len 16384 + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5},\"engine_args\":{\"draft_sample_method\":\"probabilistic\"}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dflash8/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # DFlash8 T0 C32 + task_10: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DFLASH + - --draft_model_dir /hf-local/z-lab/Qwen3.5-35B-A3B-DFlash + - --block_size 8 + - --max_seq_len 16384 + - --tp_size 2 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dflash8/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + time: 02:00:00 + container: vllm/vllm-openai:nightly diff --git a/tools/launcher/examples/Qwen/Qwen3.5-9B/chat_template_train.jinja b/tools/launcher/examples/Qwen/Qwen3.5-9B/chat_template_train.jinja new file mode 100644 index 00000000000..2dde46e080e --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-9B/chat_template_train.jinja @@ -0,0 +1,161 @@ +{# Official Qwen3.5-9B template with assistant output wrapped in generation tags. #} +{%- set image_count = namespace(value=0) %} +{%- set video_count = namespace(value=0) %} +{%- macro render_content(content, do_vision_count, is_system_content=false) %} + {%- if content is string %} + {{- content }} + {%- elif content is iterable and content is not mapping %} + {%- for item in content %} + {%- if 'image' in item or 'image_url' in item or item.type == 'image' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain images.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set image_count.value = image_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Picture ' ~ image_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|image_pad|><|vision_end|>' }} + {%- elif 'video' in item or item.type == 'video' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain videos.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set video_count.value = video_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Video ' ~ video_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|video_pad|><|vision_end|>' }} + {%- elif 'text' in item %} + {{- item.text }} + {%- else %} + {{- raise_exception('Unexpected item type in content.') }} + {%- endif %} + {%- endfor %} + {%- elif content is none or content is undefined %} + {{- '' }} + {%- else %} + {{- raise_exception('Unexpected content type.') }} + {%- endif %} +{%- endmacro %} +{%- if not messages %} + {{- raise_exception('No messages provided.') }} +{%- endif %} +{%- if tools and tools is iterable and tools is not mapping %} + {{- '<|im_start|>system\n' }} + {{- "# Tools\n\nYou have access to the following functions:\n\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n" }} + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n' }} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {%- if content %} + {{- '\n\n' + content }} + {%- endif %} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- else %} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" %} + {%- set content = render_content(message.content, false)|trim %} + {%- if not(content.startswith('') and content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if ns.multi_step_tool %} + {{- raise_exception('No user query found in messages.') }} +{%- endif %} +{%- for message in messages %} + {%- set content = render_content(message.content, true)|trim %} + {%- if message.role == "system" %} + {%- if not loop.first %} + {{- raise_exception('System message must be at the beginning.') }} + {%- endif %} + {%- elif message.role == "user" %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- set unfinished_thinking = message.enable_thinking is defined and message.enable_thinking and '' not in content %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- set reasoning_content = reasoning_content|trim %} + {{- '<|im_start|>' + message.role + '\n' }} + {%- generation %} + {%- if loop.index0 > ns.last_query_index and unfinished_thinking %} + {{- '\n' + content }} + {%- elif loop.index0 > ns.last_query_index %} + {{- '\n' + reasoning_content + '\n\n\n' + content }} + {%- else %} + {{- content }} + {%- endif %} + {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- if loop.first %} + {%- if content|trim %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n\n' }} + {%- endif %} + {%- else %} + {{- '\n\n\n' }} + {%- endif %} + {%- if tool_call.arguments is defined %} + {%- for args_name, args_value in tool_call.arguments|items %} + {{- '\n' }} + {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %} + {{- args_value }} + {{- '\n\n' }} + {%- endfor %} + {%- endif %} + {{- '\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- endgeneration %} + {%- elif message.role == "tool" %} + {%- if loop.previtem and loop.previtem.role != "tool" %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if not loop.last and loop.nextitem.role != "tool" %} + {{- '<|im_end|>\n' }} + {%- elif loop.last %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- else %} + {{- raise_exception('Unexpected message role.') }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n' }} + {%- endif %} +{%- endif %} diff --git a/tools/launcher/examples/Qwen/Qwen3.5-9B/hf_streaming_dspark_multi_node.yaml b/tools/launcher/examples/Qwen/Qwen3.5-9B/hf_streaming_dspark_multi_node.yaml new file mode 100644 index 00000000000..ac1500dcc60 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-9B/hf_streaming_dspark_multi_node.yaml @@ -0,0 +1,91 @@ +# Qwen3.5-9B DSpark training: 8 target nodes + 8 trainer nodes. + +job_name: qwen3.5-9b-dspark-training + +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3.5-9B + hf_data: /hf-local/modelopt/qwen3.5-9b-dspark-synthesis + output_dir: /cicd/qwen3.5-9b-dspark + + task_1: + script: common/eagle3/train_eagle_streaming.sh + reqs: wandb==0.22.3 + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - data.data_path=<> + - data.chat_template=examples/Qwen/Qwen3.5-9B/chat_template_train.jinja + - training.output_dir=<>/training + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=5 + - training.per_device_train_batch_size=1 + - training.gradient_accumulation_steps=8 + - training.learning_rate=6.0e-4 + - training.warmup_ratio=0.04 + - training.lr_scheduler_type=cosine_with_min_lr + - training.lr_scheduler_kwargs.min_lr=3.0e-5 + - training.save_steps=200 + - training.logging_steps=10 + - training.estimate_ar=true + - training.answer_only_loss=true + - training.report_to=wandb + - dflash.dflash_block_size=8 + - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_decay_factor=4.0 + - dflash.dflash_mask_token_id=248077 + - dflash.dflash_swa_window_size=4096 + - dflash.dflash_architecture_config.num_hidden_layers=6 + - dflash.dflash_architecture_config.num_attention_heads=32 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.head_dim=128 + - dflash.dflash_architecture_config.intermediate_size=12288 + environment: + - HF_MODEL_CKPT: <> + - EAGLE_CAPTURE_IDS: "[2,8,13,19,24,30,32]" + - SERVE_NODES: "8" + - SERVE_TP: "8" + - STREAMING_NUM_WORKERS: "4" + - HS_POOL_SLOTS: "32" + - WANDB_PROJECT: modelopt-dspark-training + - WANDB_NAME: qwen3.5-9b-dspark + - WANDB_RUN_ID: qwen35-9b-dspark + - WANDB_RESUME: allow + - WANDB_DIR: <>/wandb + - EXPORT_PATH: <>/export + - EXPORT_EXTRA_ARGS: --trust_remote_code + - SERVE_MAX_MODEL_LEN: "4160" + - SERVE_MAX_NUM_SEQS: "32" + - SERVE_READY_TIMEOUT: "1800" + slurm_config: + _factory_: slurm_factory + nodes: 16 + ntasks_per_node: 1 + gpus_per_node: 8 + dependency: singleton + time: 04:00:00 + container: vllm/vllm-openai:v0.27.1 + + task_2: + script: common/specdec/export_latest_and_vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_TRAINING_DIR: <>/training + - DRAFT_MODEL: <>/export + - EXPORT_TRUST_REMOTE_CODE: "1" + - SPEC_METHOD: dspark + - NUM_SPEC_TOKENS: "7" + - TP_SIZE: "1" + - MAX_MODEL_LEN: "4160" + - MAX_OUTPUT_TOKENS: "128" + - MIN_ACCEPTANCE_LENGTH: "1.2" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:v0.27.1 diff --git a/tools/launcher/examples/Qwen/Qwen3.5-9B/hf_synth.yaml b/tools/launcher/examples/Qwen/Qwen3.5-9B/hf_synth.yaml new file mode 100644 index 00000000000..6a3f6965842 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-9B/hf_synth.yaml @@ -0,0 +1,56 @@ +# Synthesize Qwen3.5-9B conversations dataset for DSpark training. + +job_name: qwen3.5-9b-dspark-synthesis + +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3.5-9B + hf_data: /hf-local/nvidia/Nemotron-Post-Training-Dataset-v2/default.jsonl + output_dir: /hf-local/modelopt/qwen3.5-9b-dspark-synthesis + + task_0: + script: common/vllm/query.sh + args: + - --model + - <> + - --tensor-parallel-size + - "8" + - --trust-remote-code + - --language-model-only + - --enable-prefix-caching + - --mamba-cache-mode + - align + - --gpu-memory-utilization + - "0.95" + - --max-model-len + - "8192" + - -- + - --data + - <> + - --save + - <> + - --num-shards + - "192" + - --shard-id + - $SLURM_ARRAY_TASK_ID + - --thinking-control + - chat-template-kwargs + - "--sampling-params '{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}'" + - "--non-thinking-sampling-params '{\"temperature\":0.7,\"top_p\":0.8,\"top_k\":20,\"presence_penalty\":1.5}'" + - --num-proc + - "64" + - --max-tokens + - "4096" + - --max-total-length + - "8192" + environment: + - VLLM_STARTUP_TIMEOUT: "1800" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + time: 04:00:00 + container: vllm/vllm-openai:v0.27.1 + array: "0-191%96" + requeue: true diff --git a/tools/launcher/examples/Qwen/Qwen3.5-9B/specdec_bench_tp1_c1.yaml b/tools/launcher/examples/Qwen/Qwen3.5-9B/specdec_bench_tp1_c1.yaml new file mode 100644 index 00000000000..252f6ab0b07 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-9B/specdec_bench_tp1_c1.yaml @@ -0,0 +1,322 @@ +# SPEED-Bench: Qwen3.5-9B Base / DSpark7 / MTP3 / MTP7 / DFlash8 (TP1 / C1 / T0+T1 / MRV2) + +job_name: qwen3.5-9b-specdec-bench-tp1-c1 + +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3.5-9B + draft_model: /cicd/qwen3.5-9b-dspark/export + output_dir: /cicd/qwen3.5-9b-dspark-benchmark-tp1-c1 + + # Draft semantics: DSpark7 drafts 7 speculative tokens; DFlash8 drafts 7 speculative tokens (includes 1 anchor). + + # Prepare SPEED-Bench data + task_0: + inline: >- + python3 -m pip install -r modules/Model-Optimizer/examples/specdec_bench/requirements.txt && + (test -f /hf-local/nvidia/SPEED-Bench/speed/qualitative/test.parquet || + python3 modules/Model-Optimizer/examples/specdec_bench/prepare_data.py + --dataset speed --config qualitative --output_dir /hf-local/nvidia/SPEED-Bench) + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # Base T1 C1 + task_1: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm NONE + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/base/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # Base T0 C1 + task_2: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm NONE + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/base/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # DSpark7 T1 C1 + task_3: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DSPARK + - --draft_model_dir <> + - --block_size 7 + - --max_seq_len 16384 + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5},\"engine_args\":{\"draft_sample_method\":\"probabilistic\"}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dspark7/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # DSpark7 T0 C1 + task_4: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DSPARK + - --draft_model_dir <> + - --block_size 7 + - --max_seq_len 16384 + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dspark7/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # MTP3 T1 C1 + task_5: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 3 + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp3/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # MTP3 T0 C1 + task_6: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 3 + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp3/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # MTP7 T1 C1 + task_7: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 7 + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp7/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # MTP7 T0 C1 + task_8: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 7 + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp7/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # DFlash8 T1 C1 + task_9: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DFLASH + - --draft_model_dir /hf-local/z-lab/Qwen3.5-9B-DFlash + - --block_size 8 + - --max_seq_len 16384 + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5},\"engine_args\":{\"draft_sample_method\":\"probabilistic\"}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dflash8/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly + + # DFlash8 T0 C1 + task_10: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DFLASH + - --draft_model_dir /hf-local/z-lab/Qwen3.5-9B-DFlash + - --block_size 8 + - --max_seq_len 16384 + - --tp_size 1 + - --ep_size 1 + - --concurrency 1 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dflash8/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 04:00:00 + container: vllm/vllm-openai:nightly diff --git a/tools/launcher/examples/Qwen/Qwen3.5-9B/specdec_bench_tp1_c32.yaml b/tools/launcher/examples/Qwen/Qwen3.5-9B/specdec_bench_tp1_c32.yaml new file mode 100644 index 00000000000..994c7d83211 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3.5-9B/specdec_bench_tp1_c32.yaml @@ -0,0 +1,322 @@ +# SPEED-Bench: Qwen3.5-9B Base / DSpark7 / MTP3 / MTP7 / DFlash8 (TP1 / C32 / T0+T1 / MRV2) + +job_name: qwen3.5-9b-specdec-bench-tp1-c32 + +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3.5-9B + draft_model: /cicd/qwen3.5-9b-dspark/export + output_dir: /cicd/qwen3.5-9b-dspark-benchmark-tp1-c32 + + # Draft semantics: DSpark7 drafts 7 speculative tokens; DFlash8 drafts 7 speculative tokens (includes 1 anchor). + + # Prepare SPEED-Bench data + task_0: + inline: >- + python3 -m pip install -r modules/Model-Optimizer/examples/specdec_bench/requirements.txt && + (test -f /hf-local/nvidia/SPEED-Bench/speed/qualitative/test.parquet || + python3 modules/Model-Optimizer/examples/specdec_bench/prepare_data.py + --dataset speed --config qualitative --output_dir /hf-local/nvidia/SPEED-Bench) + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # Base T1 C32 + task_1: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm NONE + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/base/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # Base T0 C32 + task_2: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm NONE + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/base/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # DSpark7 T1 C32 + task_3: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DSPARK + - --draft_model_dir <> + - --block_size 7 + - --max_seq_len 16384 + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5},\"engine_args\":{\"draft_sample_method\":\"probabilistic\"}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dspark7/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # DSpark7 T0 C32 + task_4: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DSPARK + - --draft_model_dir <> + - --block_size 7 + - --max_seq_len 16384 + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dspark7/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # MTP3 T1 C32 + task_5: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 3 + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp3/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # MTP3 T0 C32 + task_6: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 3 + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp3/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # MTP7 T1 C32 + task_7: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 7 + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp7/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # MTP7 T0 C32 + task_8: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 7 + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/mtp7/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # DFlash8 T1 C32 + task_9: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DFLASH + - --draft_model_dir /hf-local/z-lab/Qwen3.5-9B-DFlash + - --block_size 8 + - --max_seq_len 16384 + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5},\"engine_args\":{\"draft_sample_method\":\"probabilistic\"}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dflash8/t1 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly + + # DFlash8 T0 C32 + task_10: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench/speed/qualitative + - --engine VLLM + - --speculative_algorithm DFLASH + - --draft_model_dir /hf-local/z-lab/Qwen3.5-9B-DFlash + - --block_size 8 + - --max_seq_len 16384 + - --tp_size 1 + - --ep_size 1 + - --concurrency 32 + - --num_requests 440 + - --output_length 4096 + - "--runtime_params '{\"sampling_kwargs\":{\"temperature\":0.0,\"top_p\":0.95,\"top_k\":20,\"presence_penalty\":1.5}}'" + - --aa_timing + - --show_progress + - --save_dir <>/dflash8/t0 + environment: + - HF_MODEL_CKPT: <> + - VLLM_USE_V2_MODEL_RUNNER: "1" + slurm_config: + _factory_: slurm_factory + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: 02:00:00 + container: vllm/vllm-openai:nightly diff --git a/tools/launcher/launch.py b/tools/launcher/launch.py index 0f3313aaf4a..4979b6da446 100644 --- a/tools/launcher/launch.py +++ b/tools/launcher/launch.py @@ -107,6 +107,8 @@ def _add_package_glob(pattern: str) -> None: _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/modelopt")) _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/modelopt_recipes")) _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/examples")) + _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/pyproject.toml")) + _add_package_path(os.path.join(LAUNCHER_DIR, "modules/Model-Optimizer/LICENSE_HEADER")) packager = run.PatternPackager( include_pattern=_include_pattern, diff --git a/tools/launcher/pyproject.toml b/tools/launcher/pyproject.toml index 8ba825776d4..cbc1d1f2c80 100644 --- a/tools/launcher/pyproject.toml +++ b/tools/launcher/pyproject.toml @@ -24,6 +24,9 @@ modelopt_launcher = [ "examples/**/*.jinja", ] +[tool.uv] +package = true + [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/tools/launcher/slurm_config.py b/tools/launcher/slurm_config.py index 2647a7ba5ea..0ced06a7c94 100644 --- a/tools/launcher/slurm_config.py +++ b/tools/launcher/slurm_config.py @@ -44,6 +44,7 @@ class SlurmConfig: modelopt_install_path: str = "/usr/local/lib/python3.12/dist-packages/modelopt" container_mounts: Optional[list[str]] = None srun_args: Optional[list[str]] = None + dependency: Optional[str] = None array: Optional[str] = None requeue: bool = False nodes: int = 1 @@ -81,6 +82,7 @@ def slurm_factory( "{}:/hf-local".format(os.environ.get("SLURM_HF_LOCAL", "/hf-local")), ], srun_args: list[str] = ["--no-container-mount-home"], + dependency: Optional[str] = None, array: Optional[str] = None, requeue: bool = False, time: str = "04:00:00", @@ -101,6 +103,7 @@ def slurm_factory( modelopt_install_path=modelopt_install_path, container_mounts=container_mounts, srun_args=srun_args, + dependency=dependency, array=array, requeue=requeue, time=time, diff --git a/tools/launcher/tests/test_core.py b/tools/launcher/tests/test_core.py index 463e41f412a..9d873ec4992 100644 --- a/tools/launcher/tests/test_core.py +++ b/tools/launcher/tests/test_core.py @@ -38,6 +38,12 @@ SandboxTask2, SandboxTask3, SandboxTask4, + SandboxTask5, + SandboxTask6, + SandboxTask7, + SandboxTask8, + SandboxTask9, + SandboxTask10, create_task_from_yaml, get_default_env, register_factory, @@ -75,8 +81,8 @@ class TestSandboxPipeline: def test_task_slots_collected(self): t0 = SandboxTask0(script="a.sh") - t1 = SandboxTask1(script="b.sh") - pipeline = SandboxPipeline(task_0=t0, task_1=t1) + t10 = SandboxTask10(script="b.sh") + pipeline = SandboxPipeline(task_0=t0, task_10=t10) assert len(pipeline.tasks) == 2 assert pipeline.tasks[0].script == "a.sh" assert pipeline.tasks[1].script == "b.sh" @@ -195,6 +201,12 @@ class MockSlurmConfig: SandboxTask2, SandboxTask3, SandboxTask4, + SandboxTask5, + SandboxTask6, + SandboxTask7, + SandboxTask8, + SandboxTask9, + SandboxTask10, ): assert task_cls.__annotations__["slurm_config"] is MockSlurmConfig assert task_cls.__dataclass_fields__["slurm_config"].type is MockSlurmConfig @@ -224,6 +236,16 @@ def test_custom_title(self, monkeypatch): assert slurm_env["HF_HOME"] == "/modelopt/hf-cache" assert local_env["HF_HOME"] == "/modelopt/hf-cache" + def test_serve_nodes_forwarded_to_slurm(self, monkeypatch): + monkeypatch.setenv("SERVE_NODES", "8") + slurm_env, _ = get_default_env() + assert slurm_env["SERVE_NODES"] == "8" + + def test_serve_nodes_omitted_when_unset(self, monkeypatch): + monkeypatch.delenv("SERVE_NODES", raising=False) + slurm_env, _ = get_default_env() + assert "SERVE_NODES" not in slurm_env + class TestReportVersions: """Tests for report_versions git info utility.""" diff --git a/tools/launcher/tests/test_core_extended.py b/tools/launcher/tests/test_core_extended.py index 757debc01c4..ff329d4af4b 100644 --- a/tools/launcher/tests/test_core_extended.py +++ b/tools/launcher/tests/test_core_extended.py @@ -433,7 +433,7 @@ def test_reqs_inline_barrier(self, mock_docker, mock_exp, tmp_path): inline = mock_script.call_args[1]["inline"] # `<` and `>` are shlex-quoted so the shell treats them literally. - assert "python -m pip install 'transformers<5' fire" in inline + assert "python3 -m pip install 'transformers<5' fire" in inline # Local rank 0 installs; other ranks wait on a per-job/step/node marker. assert '[ "${SLURM_LOCALID:-0}" -eq 0 ]' in inline marker = ( @@ -481,7 +481,7 @@ def test_reqs_script_wrapped_inline(self, mock_docker, mock_exp, tmp_path): call_kwargs = mock_script.call_args[1] inline = call_kwargs["inline"] - assert "python -m pip install fire" in inline + assert "python3 -m pip install fire" in inline # "--flag value" keeps the shell-word-split convention (expands to two args). assert inline.rstrip().endswith("bash run.sh --flag value") # Wrapped inline: no separate script/args kwargs. diff --git a/tools/launcher/tests/test_slurm_config.py b/tools/launcher/tests/test_slurm_config.py index 96cfc689dbd..c4391fd532e 100644 --- a/tools/launcher/tests/test_slurm_config.py +++ b/tools/launcher/tests/test_slurm_config.py @@ -45,6 +45,7 @@ def test_defaults(self): assert cfg.local is False assert cfg.container_mounts is None assert cfg.srun_args is None + assert cfg.dependency is None assert cfg.array is None def test_custom_values(self): @@ -57,6 +58,7 @@ def test_custom_values(self): container="nvcr.io/nvidia/pytorch:24.01-py3", container_mounts=["/data:/data"], srun_args=["--no-container-mount-home"], + dependency="singleton", ) assert cfg.host == "login.example.com" assert cfg.account == "my_account" @@ -64,6 +66,7 @@ def test_custom_values(self): assert cfg.gpus_per_node == 8 assert cfg.mem == "128G" assert cfg.container_mounts == ["/data:/data"] + assert cfg.dependency == "singleton" def test_nullable_gpus_per_node(self): cfg = SlurmConfig(gpus_per_node=None) diff --git a/tools/launcher/tests/test_slurm_executor.py b/tools/launcher/tests/test_slurm_executor.py index 3ce72ba656b..b84cb087341 100644 --- a/tools/launcher/tests/test_slurm_executor.py +++ b/tools/launcher/tests/test_slurm_executor.py @@ -425,7 +425,7 @@ def test_none_container_mounts_handled(self, mock_tunnel, mock_executor): @patch("core.run.SlurmExecutor") @patch("core.run.SSHTunnel") - def test_requeue_sets_param_and_bumps_retries(self, mock_tunnel, mock_executor): + def test_additional_params_and_requeue(self, mock_tunnel, mock_executor): mock_tunnel.return_value = MagicMock() executor = mock_executor.return_value executor.retries = 0 @@ -445,6 +445,8 @@ def test_requeue_sets_param_and_bumps_retries(self, mock_tunnel, mock_executor): ntasks_per_node=1, gpus_per_node=1, array=None, + dependency="singleton", + additional_parameters={"signal": "USR1@60"}, ) build_slurm_executor( @@ -457,6 +459,10 @@ def test_requeue_sets_param_and_bumps_retries(self, mock_tunnel, mock_executor): packager=MagicMock(), ) + assert mock_executor.call_args.kwargs["additional_parameters"] == { + "signal": "USR1@60", + "dependency": "singleton", + } # requeue=True flags the additional parameter and bumps retries above 0 so # nemo-run's sbatch wrapper actually issues `scontrol requeue` on preemption. assert executor.additional_parameters["requeue"] is True