Model-Optimizer/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml

# DFlash2 streaming speculative decoding pipeline for Qwen3-8B.
#
# Same streaming transport as hf_streaming_dflash.yaml: a live `vllm serve` captures the
# target model's hidden states and moves them to the trainer over NIXL RDMA (no disk
# round-trip). DFlash2 trains the same block-diffusion backbone plus the grouped sublayer
# convolutions and the candidate selector. See common/eagle3/train_eagle_streaming.sh for
# the dispatch and rendezvous.
#
# Two nodes, all four GPUs of each: node 0 is one vLLM replica at TP=4, node 1 is a
# 4-rank DDP trainer. Scale down by lowering gpus_per_node and SERVE_TP together; scale up
# by raising nodes and SERVE_NODES together, as hf_streaming_dflash_multi_node.yaml does.
#
# 3-step pipeline:
#   task_0: Build input conversations (jsonl)
#   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).
#
# Usage:
#   uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_streaming_dflash2.yaml --yes

job_name: Qwen3-8B_DFlash2_streaming
pipeline:
  allow_to_fail: false
  skip: false
  note:

  global_vars:
    hf_model: /hf-local/Qwen/Qwen3-8B

  # Step 1: Build input conversations
  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 DFlash2 training — node 0 vllm serve (TP=4), node 1 trainer.
  # DFlash2 extracts 5 target layers (build_target_layer_ids(36,5)=[1,9,17,25,33], the
  # draft's fc input); vLLM's capture ids are those +1 -> [2,10,18,26,34]. No final layer
  # (36): this recipe trains against the hard target, so the base hidden is never read.
  # Dropping it must be paired with data.final_aux_is_base_hidden=true below. To train
  # with distillation instead, set dflash_lk_loss_type=ce and add 36 back without it.
  task_1:
    script: common/eagle3/train_eagle_streaming.sh
    args:
      - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash2.yaml
      - model.model_name_or_path=<<global_vars.hf_model>>
      # Streaming never runs the base's transformer layers on the trainer; load only the
      # embeddings, final norm and lm_head rather than all 36 layers on every rank.
      - model.use_fake_base_for_offline=true
      - data.mode=streaming
      - data.data_path=/scratchspace/data/train.jsonl
      # All five captured planes are aux features; without this the streaming dataset
      # 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.training_seq_len=4096
      - 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
      # {% generation %} tags to locate it. Pass chat_template_train.jinja to mask to it.
      - training.answer_only_loss=false
      # 4 per device x 4 ranks = global batch 16 sequences = 65,536 tokens/step.
      - training.per_device_train_batch_size=4
      - training.num_train_epochs=1
      - training.max_steps=5000
      # dflash2.yaml sets report_to=tensorboard, which hard-fails if tensorboard
      # isn't in the serve container; the streaming trainer doesn't need it.
      - training.report_to=none
    environment:
      - HF_MODEL_CKPT: <<global_vars.hf_model>>
      # No spaces: nemo_run emits unquoted `export FOO=value`, so spaces would split.
      - EAGLE_CAPTURE_IDS: "[2,10,18,26,34]"
      # Serve replica nodes (Slurm nodes 0..SERVE_NODES-1); the rest are trainers.
      - SERVE_NODES: "1"
      - SERVE_TP: "4"
      # training_seq_len plus headroom for the single decode step.
      - SERVE_MAX_MODEL_LEN: "4608"
      - SERVE_MAX_NUM_SEQS: "32"
      - SERVE_GPU_MEM_UTIL: "0.9"
      # RDMA pool slot capacity in tokens. Must be >= training_seq_len or long prompts
      # overflow the slot and the producer silently skips capture. 32 slots is 5.4 GiB of
      # pinned host memory against a peak in-flight of 4 ranks x 4 workers.
      - HS_MAX_TOKENS: "4608"
      - HS_POOL_SLOTS: "32"
      - STREAMING_NUM_WORKERS: "4"
      # DFlash2 uses a custom modeling file; export must trust remote code.
      - EXPORT_EXTRA_ARGS: "--trust_remote_code"
    slurm_config:
      _factory_: "slurm_factory"
      nodes: 2
      ntasks_per_node: 1
      gpus_per_node: 4
      container: vllm/vllm-openai:latest

  # Step 3: vLLM smoke test (uses the exported checkpoint from training).
  # The method stays "dflash": vLLM has no separate dflash2 method and selects the
  # DFlash2 path from the checkpoint's architectures: ["DFlash2DraftModel"].
  task_2:
    script: common/specdec/vllm_smoke_test.sh
    environment:
      - HF_MODEL_CKPT: <<global_vars.hf_model>>
      - DRAFT_MODEL: /scratchspace/export
      - SPEC_METHOD: "dflash"
      - 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