From 16fd70e2d93ec1eb29df4a841e95b45c970b5ae8 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 9 Jul 2026 22:37:50 +0000 Subject: [PATCH] examples: add DSpark recipe + Qwen3-8B online DSpark launcher example Adds the DSpark speculative-decoding recipe (modelopt_recipes/general/speculative_decoding/dspark.yaml) and an online-mode launcher example for Qwen3-8B (tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dspark.yaml). DSpark reuses the DFlash online mode/pipeline and adds a lightweight sequential (Markov) head plus an optional confidence head, selected via dflash_architecture_config.projector_type=dspark. The Qwen3-8B example mirrors hf_online_dflash.yaml and only swaps the recipe to dspark.yaml. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../general/speculative_decoding/dspark.yaml | 96 +++++++++++++++++++ .../Qwen/Qwen3-8B/hf_online_dspark.yaml | 84 ++++++++++++++++ 2 files changed, 180 insertions(+) create mode 100644 modelopt_recipes/general/speculative_decoding/dspark.yaml create mode 100644 tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dspark.yaml diff --git a/modelopt_recipes/general/speculative_decoding/dspark.yaml b/modelopt_recipes/general/speculative_decoding/dspark.yaml new file mode 100644 index 00000000000..cb7b09f40e9 --- /dev/null +++ b/modelopt_recipes/general/speculative_decoding/dspark.yaml @@ -0,0 +1,96 @@ +# DSpark speculative-decoding training recipe. +# +# DSpark reuses the DFlash mode/pipeline and adds a lightweight sequential +# (Markov) head plus an optional confidence head, selected via +# dflash_architecture_config.projector_type=dspark. The Markov head adds a +# prefix-dependent transition bias to the backbone base logits, inducing a causal +# block distribution (semi-autoregressive generation). Trained with a three-term +# loss: ce_alpha*CE + l1_alpha*TVD + conf_alpha*confidence_BCE. Online training is +# the default path (data.mode=online). Override fields via an OmegaConf dotlist. + +metadata: + recipe_type: speculative_dflash + description: DSpark training recipe (DFlash backbone + Markov head + confidence head). + +# maps to ModelArguments (main.py) +model: + model_name_or_path: + trust_remote_code: false + use_fake_base_for_offline: false + +# maps to DataArguments (main.py) +data: + mode: online + data_path: + offline_data_path: + # Jinja chat template with {% generation %} tags for answer_only_loss. + chat_template: + +# maps to TrainingArguments (main.py) +training: + # --- commonly modified --- + output_dir: + num_train_epochs: 6 + per_device_train_batch_size: 1 + learning_rate: 6.0e-4 + warmup_ratio: 0.04 + training_seq_len: 3072 + logging_steps: 50 + save_steps: 2000 + cp_size: 1 + dp_shard_size: 1 + disable_tqdm: true + # Keep off: eval runs the DFlash backbone only (Markov head not applied yet), + # so AR would reflect the backbone alone, not the trained model. Compare via + # export + the offline acceptance-length harness instead. + estimate_ar: false + ar_validate_steps: 0 + answer_only_loss: true + + # --- rarely modified --- + do_eval: false + lr_scheduler_type: linear + save_strategy: steps + weight_decay: 0.0 + max_grad_norm: 1.0 + dataloader_drop_last: true + bf16: true + tf32: true + remove_unused_columns: false + # Safe default: the confidence head params are unused when + # dflash_confidence_head_alpha == 0, which would otherwise trip DDP. + ddp_find_unused_parameters: true + ddp_timeout: 1800 + report_to: tensorboard + +# maps to DFlashConfig (modelopt/torch/speculative/config.py). +dflash: + dflash_block_size: 16 + dflash_num_anchors: 256 + dflash_use_torch_compile: false + # DSpark computes the target distribution internally for its TVD/confidence + # terms; the DFlash KD path is unused. + dflash_self_logit_distillation: false + # gamma for exponential loss decay (block_size=16 -> 7). + dflash_loss_decay_factor: 7.0 + # 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). + dflash_ce_loss_alpha: 0.1 + dflash_l1_loss_alpha: 0.9 + dflash_confidence_head_alpha: 1.0 + dflash_architecture_config: + num_hidden_layers: 5 + # Draft attention/MLP dims — set explicitly (the draft is an independent + # Qwen3 model and does NOT inherit these from the base). GQA: 8 KV heads. + num_attention_heads: 32 + num_key_value_heads: 8 + head_dim: 128 + intermediate_size: 12288 + projector_type: dspark + # Markov head: low-rank first-order transition bias. 'vanilla' (memoryless), + # 'gated' (hidden-gated), or 'rnn' (recurrent, closest to Domino's GRU). + markov_rank: 256 + markov_head_type: vanilla + # Build + train the confidence head (per-position acceptance predictor). + use_confidence_head: true diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dspark.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dspark.yaml new file mode 100644 index 00000000000..f09e2be5be1 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dspark.yaml @@ -0,0 +1,84 @@ +# DSpark online speculative decoding training for Qwen3-8B. +# +# DSpark = the DFlash backbone + a lightweight sequential (Markov) head + a +# confidence head. The Markov head adds a prefix-dependent transition bias to the +# backbone base logits, inducing a causal block distribution (semi-autoregressive +# generation); the confidence head predicts per-position acceptance. Trained with +# a 3-term loss: ce_alpha*CE + l1_alpha*TVD + conf_alpha*confidence_BCE (weights in +# dspark.yaml). It reuses the DFlash online mode/pipeline (same training script); +# only the recipe changes (dflash.yaml -> dspark.yaml, which sets +# projector_type=dspark, the confidence head and the loss weights). +# +# 2-step pipeline: +# task_0: Online DSpark training (8 GPU) -> exports a drafter under output_dir +# task_1: vLLM smoke test with DSpark speculative decoding (measures the real +# semi-AR acceptance length) +# +# NOTE ON AR EVAL: unlike the DFlash example there is no HF AR-eval step here. +# dspark.yaml keeps estimate_ar=false because the training-time eval runs the +# DFlash backbone only (the Markov head is not applied), so an AR number would +# reflect the backbone alone, not the trained DSpark model. Measure acceptance +# via the vLLM smoke test (below) or the offline specdec_bench harness instead. +# +# The draft dims, block_size, mask token and loss weights all come from +# dspark.yaml (Qwen3-tuned defaults); only the commonly-tuned data/schedule fields +# are overridden below. dflash_block_size=16 is the semi-AR generation block +# (training_seq_len must be divisible by it). +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_online_dspark.yaml --yes +# uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dspark.yaml --yes + +job_name: Qwen3-8B_DSpark_online +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3-8B + + # Step 1: Online DSpark training (reuses the DFlash online training script; + # only the --config recipe differs). + task_0: + script: common/specdec/dflash_online_training.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml + - model.model_name_or_path=<> + - data.data_path=/hf-local/modelopt/Speculative-Decoding-Dataset-v1-Qwen3-8B/sample-100K-openai.jsonl + - data.chat_template=examples/Qwen/Qwen3-8B/chat_template_train.jinja + - training.output_dir=/scratchspace/dspark_bs16 + - training.per_device_train_batch_size=1 + - training.num_train_epochs=1 + - training.training_seq_len=4096 + - training.save_steps=5000 + - training.logging_steps=100 + - training.disable_tqdm=true + - training.answer_only_loss=true + # Semi-AR generation block (recipe default 16; must divide training_seq_len). + - dflash.dflash_block_size=16 + - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_decay_factor=7 + # Qwen3 has no native mask token; 151669 is an unused id (recipe default). + - dflash.dflash_mask_token_id=151669 + - dflash.dflash_architecture_config.num_hidden_layers=5 + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + + # Step 2: vLLM smoke test (uses the exported drafter from training). + # Requires a vLLM image with DSpark support (method="dspark"). num_speculative_tokens + # is the number of draft tokens verified per step and can be tuned up to + # dflash_block_size (16); 7 mirrors the DFlash example as a conservative start. + task_1: + script: common/specdec/vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_CKPT_DIR: /scratchspace/dspark_bs16 + - SPEC_METHOD: "dspark" + - NUM_SPEC_TOKENS: "7" + - MIN_ACCEPTANCE_LENGTH: "1.4" + slurm_config: + _factory_: "slurm_factory" + container: "vllm/vllm-openai:nightly" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1