[4/8] launcher: clean up the Qwen3-8B DFlash2 launcher examples

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/<experiment_id>/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=<<global_vars.hf_model>>
       - 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=<<global_vars.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: <<global_vars.hf_model>>
+      - DRAFT_CKPT_DIR: <<global_vars.output_dir>>
+      - SPEC_METHOD: "dflash"
+      - NUM_SPEC_TOKENS: "7"
+      - MIN_ACCEPTANCE_LENGTH: "1.5"
+      - SMOKE_JSON: <<global_vars.output_dir>>/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

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/<experiment_id>/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=<<global_vars.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: <<global_vars.output_dir>>/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: <<global_vars.hf_model>>
-      - DRAFT_MODEL: /scratchspace/export
+      - DRAFT_MODEL: <<global_vars.output_dir>>/export
       - SPEC_METHOD: "dflash"
       - NUM_SPEC_TOKENS: "7"
       - MIN_ACCEPTANCE_LENGTH: "1.2"
+      - SMOKE_JSON: <<global_vars.output_dir>>/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