hf_streaming_dflash_multi_node.yaml

# DFlash streaming speculative-decoding training for Kimi-K2.5-NVFP4 on
# GB200/Blackwell — MULTI-NODE: both serve and trainer sides scale out.
# nodes=N, SERVE_NODES=K -> K serve replicas (TP=4, whole node) + (N-K) trainer
# nodes. See common/eagle3/train_eagle_streaming.sh for dispatch/sharding/scaling.
#
# Requires GB200: native NVFP4 + 192 GB/GPU fits ~551 GB Kimi-K2.5-NVFP4 at TP=4
# on one 4-GPU node, so each serve replica owns a whole node.
#
# Capture ids: build_target_layer_ids(num_orig=61, num_draft=5)=[1,15,30,44,58]
# -> +1 for embedding = [2,16,31,45,59], append true final layer 61.
# 6 captured = 5 aux layers, matching the 5-layer DFlash draft block.
# (61 = true final hidden; requires a vLLM with the aux-capture fix vllm#46788.
#  Without it use 60 — the 2nd-to-last layer — which caps acceptance length.)
#
# Run ON the cluster login node (paramiko can't reach the cluster through its login proxy):
#   export SLURM_HOST=localhost SLURM_ACCOUNT=<your_account> \
#          SLURM_PARTITION=batch \
#          SLURM_HF_LOCAL=<hf_models_dir> \
#          SLURM_JOB_DIR=<experiments_dir> \
#          NEMORUN_HOME=$PWD
#   uv run launch.py --yaml examples/moonshotai/Kimi-K2.5/hf_streaming_dflash_multi_node.yaml \
#          identity=$HOME/.ssh/id_ecdsa detach=True --yes
#
# The export lands in /scratchspace/export. To benchmark it, point
# specdec_bench.yaml's --draft_model_dir there (or copy it under /hf-local).

job_name: Kimi-K2.5-NVFP4_DFlash_streaming_multi_node
pipeline:
  allow_to_fail: false
  skip: false
  note:

  global_vars:
    hf_model: /hf-local/nvidia/Kimi-K2.5-NVFP4

  # 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
      # The cluster QOS requires whole-node GPU alloc even though make_dataset is CPU-only.
      gpus_per_node: 4
      container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10

  # Streaming DFlash training: 2 serve replicas (TP=4) + 2 trainer nodes.
  task_1:
    script: common/eagle3/train_eagle_streaming.sh
    args:
      - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash.yaml
      - model.model_name_or_path=<<global_vars.hf_model>>
      - model.use_fake_base_for_offline=true
      - model.trust_remote_code=true
      - data.mode=streaming
      - data.data_path=/scratchspace/data/train.jsonl
      - training.output_dir=/scratchspace/dflash
      # Must be divisible by dflash_block_size (8).
      - training.training_seq_len=4096
      - training.disable_tqdm=true
      - training.ar_validate_steps=500000
      - training.num_train_epochs=1
      - training.max_steps=2000
      # Kimi's slow tokenizer can't emit assistant masks the standard way; the mask
      # is recovered from token ids (modelopt.torch.utils.loss_mask).
      - training.answer_only_loss=true
      # vLLM container has no tensorboard (dflash.yaml's default) -> init crash.
      - training.report_to=none
      # Kimi has no dedicated mask token; 163838 is a reserved slot used as the mask.
      - dflash.dflash_mask_token_id=163838
    environment:
      - HF_MODEL_CKPT: <<global_vars.hf_model>>
      # See header for derivation.
      - EAGLE_CAPTURE_IDS: "[2,16,31,45,59,61]"
      - SERVE_NODES: "2"
      - SERVE_TP: "4"
      # Per-rank in-flight fetches; keep low so the cold NVFP4-MoE serve isn't flooded past its execute-model timeout (kills EngineCore).
      - STREAMING_NUM_WORKERS: "1"
      # Kimi's custom-modeling base needs --trust_remote_code at export.
      - EXPORT_EXTRA_ARGS: "--trust_remote_code"
      # > training_seq_len: room for prompt + the 1-token decode fetch (else vLLM 400s at the cap).
      - SERVE_MAX_MODEL_LEN: "4160"
      - SERVE_MAX_NUM_SEQS: "4"
      - SERVE_GPU_MEM_UTIL: "0.8"
      - SERVE_READY_TIMEOUT: "2400"
      - SERVE_EXTRA_ARGS: "--trust-remote-code"
      # Cold NVFP4-MoE kernels stall the first serving step past vLLM's default execute-model timeout; raise it (seconds).
      - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200"
      - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200"
    slurm_config:
      _factory_: "slurm_factory"
      nodes: 4
      # Pin nodes into one NVL72 block (essential for cross-node trainer DDP).
      segment: 4
      ntasks_per_node: 1
      gpus_per_node: 4
      # Pin 0.22.0: 0.22.1 regressed Kimi serve (profile_run runs the ViT and hits
      # a `fmax()` crash in kimi_k25_vit.py); 0.17.0 fails SpeculativeConfig validation.
      container: vllm/vllm-openai:v0.22.0-aarch64