From 4b1a3b0bf163f544da716e71c3d88e7349725de1 Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Thu, 9 Jul 2026 11:45:32 -0700 Subject: [PATCH] example(eagle3): Qwen3-0.6B EAGLE3 online/offline launcher examples Adds examples/Qwen/Qwen3-0.6B/hf_online_eagle3.yaml + hf_offline_eagle3.yaml, adapted from the Qwen3-8B EAGLE3 examples for the small, ungated Qwen3-0.6B model on a single GPU (make_dataset/query -> dump -> train_eagle -> common/specdec_bench VLLM benchmark). Uses the co-versioned examples/specdec_bench and a VLLM backend. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenhan Yu --- .../Qwen/Qwen3-0.6B/hf_offline_eagle3.yaml | 104 ++++++++++++++++++ .../Qwen/Qwen3-0.6B/hf_online_eagle3.yaml | 73 ++++++++++++ 2 files changed, 177 insertions(+) create mode 100644 tools/launcher/examples/Qwen/Qwen3-0.6B/hf_offline_eagle3.yaml create mode 100644 tools/launcher/examples/Qwen/Qwen3-0.6B/hf_online_eagle3.yaml diff --git a/tools/launcher/examples/Qwen/Qwen3-0.6B/hf_offline_eagle3.yaml b/tools/launcher/examples/Qwen/Qwen3-0.6B/hf_offline_eagle3.yaml new file mode 100644 index 00000000000..4214ae9de4d --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-0.6B/hf_offline_eagle3.yaml @@ -0,0 +1,104 @@ +# EAGLE3 offline speculative decoding pipeline for Qwen3-8B. +# +# 4-step pipeline: +# task_0: Data synthesis — query TRT-LLM server to generate prompt samples +# task_1: Dump hidden states — run target model to capture hidden states +# task_2: Offline training — train the EAGLE3 draft head +# task_3: Benchmark — evaluate speculative decoding speedup via VLLM +# +# All tasks share /scratchspace to pass artifacts between steps. +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-0.6B/hf_offline_eagle3.yaml --yes +# uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3-0.6B/hf_offline_eagle3.yaml --yes + +job_name: Qwen3-0.6B_EAGLE3_offline +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/Qwen/Qwen3-0.6B + + # Step 1: Data synthesis via TRT-LLM server + # Args before "--" go to trtllm-serve; args after "--" go to tools/query.py. + task_0: + script: common/tensorrt_llm/query.sh + args: + - --model <> + - --tp_size 8 + - --ep_size 8 + - --max_num_tokens 32000 + - --port 8000 + - --host 0.0.0.0 + - --trust_remote_code + - -- + - --data /hf-local/modelopt/Speculative-Decoding-Prompt-Samples + - --save /scratchspace/data + environment: + - HF_LOCAL: /hf-local + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 8 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0 + + # Step 2: Dump hidden states from target model + task_1: + script: common/eagle3/dump_offline_data.sh + args: + - --input-data /scratchspace/data + - --output-dir /scratchspace/offline_hidden_states + - --max-seq-len 8192 + - --tp 8 + - --moe-ep 8 + environment: + - HF_MODEL_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 8 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0 + + # Step 3: Train EAGLE3 draft head (offline, single task) + task_2: + script: common/eagle3/train_eagle.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/eagle3.yaml + - model.model_name_or_path=<> + - data.offline_data_path=/scratchspace/offline_hidden_states + - training.output_dir=/scratchspace/eagle3 + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0 + + # Step 4: Benchmark speculative decoding (VLLM backend) + task_3: + script: common/specdec_bench/quick_check.sh + args: + - --draft_model_dir /scratchspace/export + - --draft_length 3 + - --output_length 4096 + - --engine VLLM + - --tp_size 8 + - --ep_size 1 + - --speculative_algorithm EAGLE3 + - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl + - --concurrency 32 + environment: + - HF_MODEL_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: vllm/vllm-openai:latest diff --git a/tools/launcher/examples/Qwen/Qwen3-0.6B/hf_online_eagle3.yaml b/tools/launcher/examples/Qwen/Qwen3-0.6B/hf_online_eagle3.yaml new file mode 100644 index 00000000000..5c066d68cd9 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-0.6B/hf_online_eagle3.yaml @@ -0,0 +1,73 @@ +# EAGLE3 offline speculative decoding pipeline for Qwen3-8B. +# +# 4-step pipeline: +# task_0: Data synthesis — query TRT-LLM server to generate prompt samples +# task_1: Dump hidden states — run target model to capture hidden states +# task_2: Offline training — train the EAGLE3 draft head +# task_3: Benchmark — evaluate speculative decoding speedup via VLLM +# +# All tasks share /scratchspace to pass artifacts between steps. +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-0.6B/hf_offline_eagle3.yaml --yes +# uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3-0.6B/hf_offline_eagle3.yaml --yes + +job_name: Qwen3-0.6B_EAGLE3_online +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/Qwen/Qwen3-0.6B + + 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.0rc10 + + task_1: + script: common/eagle3/train_eagle.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/eagle3.yaml + - model.model_name_or_path=<> + - data.data_path=/scratchspace/data/train.jsonl + - training.output_dir=/scratchspace/eagle3 + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=1 + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + task_2: + script: common/specdec_bench/quick_check.sh + args: + - --draft_model_dir /scratchspace/export + - --draft_length 3 + - --output_length 4096 + - --engine VLLM + - --tp_size 1 + - --ep_size 1 + - --speculative_algorithm EAGLE3 + - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl + - --concurrency 32 + environment: + - HF_MODEL_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: vllm/vllm-openai:latest