diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml index 6d08abba99c..9eeb83f2da8 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml @@ -9,9 +9,10 @@ # See the dflash2.yaml recipe and # modelopt/torch/speculative/plugins/{modeling,hf}_dflash2.py. # -# 2-step pipeline: +# 3-step pipeline: # task_0: Build training conversations (Daring-Anteater multi-turn SFT, 50K) # task_1: Online DFlash2 training + export of the drafter checkpoint +# task_2: vLLM smoke test of the exported drafter, gated on acceptance length # # As configured this is a short convergence check (max_steps=2000), matching the # other Qwen3-8B online examples so it finishes on one node. To reproduce the @@ -35,6 +36,10 @@ job_name: Qwen3-8B_DFlash2_online pipeline: global_vars: hf_model: /hf-local/Qwen/Qwen3-8B + # Training output: checkpoints, exported-checkpoint-*, and the gate verdicts. task_1 + # writes it and task_2 reads it, so override this one key to move both, e.g. to + # /cicd//dflash2_bs16 to rerun task_2 against an earlier experiment. + output_dir: /scratchspace/dflash2_bs16 # Step 1: Build input conversations. example_data_config.yaml enables only the # daring-anteater source (train: 50000) — multi-turn SFT with real assistant @@ -50,7 +55,8 @@ pipeline: nodes: 1 ntasks_per_node: 1 gpus_per_node: 1 - container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + # Same image as task_1, so both tasks see the same Python stack. + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 # Step 2: Online DFlash2 training (the script exports the drafter at the end). # Consumes the conversations built in task_0 (shared via /scratchspace). @@ -61,7 +67,7 @@ pipeline: - model.model_name_or_path=<> - data.data_path=/scratchspace/data/train.jsonl - data.chat_template=examples/Qwen/Qwen3-8B/chat_template_train.jinja - - training.output_dir=/scratchspace/dflash2_bs16 + - training.output_dir=<> - training.per_device_train_batch_size=1 - training.num_train_epochs=1 - training.max_steps=2000 @@ -74,11 +80,12 @@ pipeline: - training.logging_steps=100 - training.disable_tqdm=true - training.answer_only_loss=true - # Draft backbone — identical to the DFlash / DSpark / Domino arms so the - # only difference between them is the correction head. - - dflash.dflash_block_size=16 + # Draft backbone — as in the DFlash / DSpark / Domino arms, except the block size. + # Train on the block vLLM serves, 1 + NUM_SPEC_TOKENS (task_2); vLLM ignores the + # checkpoint's block_size. A 16-token block gave the same AL at the dev budget. + # dflash_loss_decay_factor is not set: the recipe's dpace objective ignores it. + - dflash.dflash_block_size=8 - dflash.dflash_num_anchors=512 - - dflash.dflash_loss_decay_factor=7 - dflash.dflash_mask_token_id=151669 - dflash.dflash_self_logit_distillation=false - dflash.dflash_architecture_config.num_hidden_layers=5 @@ -108,3 +115,25 @@ pipeline: nodes: 1 ntasks_per_node: 1 gpus_per_node: 8 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 + + # Step 3: Serve Qwen3-8B with the newest exported drafter and check it is accepted. + # DRAFT_CKPT_DIR picks exported-checkpoint-final (else the highest step) from task_1. + # The method stays "dflash": vLLM selects the DFlash2 path from the checkpoint's + # architectures: ["DFlash2DraftModel"] (vllm#52816, in v0.31.0). + task_2: + script: common/specdec/vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_CKPT_DIR: <> + - SPEC_METHOD: "dflash" + - NUM_SPEC_TOKENS: "7" + - MIN_ACCEPTANCE_LENGTH: "1.5" + - SMOKE_JSON: <>/smoke_check.json + - TRITON_CACHE_DIR: /tmp/triton-cache + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: vllm/vllm-openai:v0.31.0 diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml index 89c06eeceed..c414ffb0b1e 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml @@ -15,10 +15,10 @@ # task_1: Streaming train — 1 serve node (TP=4) + 1 trainer node (4-rank DDP) # task_2: vLLM smoke test with the exported drafter # -# Site-specific settings left unset on purpose: an aarch64 cluster needs its own image -# (the default is x86), slurm_config.time where the queue's default is short, the NIXL and -# NCCL HCA pinning an InfiniBand fabric needs, and container-visible cache paths (srun -# runs with --no-container-mount-home, and TRITON_CACHE_DIR must be node-local). +# Site-specific settings left unset on purpose: slurm_config.time where the queue's default +# is short, the NIXL and NCCL HCA pinning an InfiniBand fabric needs, and container-visible +# cache paths such as HF_HOME (srun runs with --no-container-mount-home). The pinned images +# are multi-arch (x86_64 and aarch64), and TRITON_CACHE_DIR is set node-local below. # # Usage: # uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml --yes @@ -31,6 +31,10 @@ pipeline: global_vars: hf_model: /hf-local/Qwen/Qwen3-8B + # Training output: checkpoints, the exported drafter (export/), and the smoke verdict. + # task_1 writes it and task_2 reads it, so override this one key to move both, e.g. to + # /cicd//dflash2 to rerun task_2 against an earlier experiment. + output_dir: /scratchspace/dflash2 # Step 1: Build input conversations task_0: @@ -65,8 +69,10 @@ pipeline: # peels the last one off as the (unused) KD target and the draft's fc gets # 4x4096 where it wants 5x4096. - data.final_aux_is_base_hidden=true - - training.output_dir=/scratchspace/dflash2 + - training.output_dir=<> - training.training_seq_len=4096 + # The recipe's 16; train on the block vLLM serves, 1 + NUM_SPEC_TOKENS (task_2). + - dflash.dflash_block_size=8 - training.disable_tqdm=true # The serve prefills the corpus conversation rather than generating, so the # assistant turn is present; Qwen3-8B's stock template just has no @@ -98,12 +104,17 @@ pipeline: - STREAMING_NUM_WORKERS: "4" # DFlash2 uses a custom modeling file; export must trust remote code. - EXPORT_EXTRA_ARGS: "--trust_remote_code" + # Where the drafter is exported (train_eagle_streaming.sh defaults to /scratchspace/export). + - EXPORT_PATH: <>/export + # Node-local: a shared Triton cache breaks when jobs on different nodes write it. + - TRITON_CACHE_DIR: /tmp/triton-cache slurm_config: _factory_: "slurm_factory" nodes: 2 ntasks_per_node: 1 gpus_per_node: 4 - container: vllm/vllm-openai:latest + # A release with DFlash2 serving (vllm#52816). + container: vllm/vllm-openai:v0.31.0 # Step 3: vLLM smoke test (uses the exported checkpoint from training). # The method stays "dflash": vLLM has no separate dflash2 method and selects the @@ -112,13 +123,15 @@ pipeline: script: common/specdec/vllm_smoke_test.sh environment: - HF_MODEL_CKPT: <> - - DRAFT_MODEL: /scratchspace/export + - DRAFT_MODEL: <>/export - SPEC_METHOD: "dflash" - NUM_SPEC_TOKENS: "7" - MIN_ACCEPTANCE_LENGTH: "1.2" + - SMOKE_JSON: <>/smoke_check.json + - TRITON_CACHE_DIR: /tmp/triton-cache slurm_config: _factory_: "slurm_factory" - container: vllm/vllm-openai:nightly + container: vllm/vllm-openai:v0.31.0 nodes: 1 ntasks_per_node: 1 gpus_per_node: 1