[OMNIML-4333] training_support — pensieve-intern agent draft

Author: pensieve-intern

tools/launcher/common/service_utils.sh

@@ -18,8 +18,8 @@
 native_mpi_rank=$OMPI_COMM_WORLD_RANK
 native_mpi_local_rank=$OMPI_COMM_WORLD_LOCAL_RANK
 # Works with Slurm launching with `--mpi=pmix`
-mpi_rank=${PMIX_RANK:-$native_mpi_rank}
-mpi_local_rank=${PMIX_LOCAL_RANK:-$native_mpi_local_rank}
+mpi_rank=${PMIX_RANK:-${native_mpi_rank:-${SLURM_PROCID:-0}}}
+mpi_local_rank=${PMIX_LOCAL_RANK:-${native_mpi_local_rank:-${SLURM_LOCALID:-0}}}
 
 FAIL=0
 FAIL_EXIT=0
@@ -48,8 +48,23 @@ function report_result {
 }
 
 function util_install_extra_dep {
+    local _marker=/tmp/.nmm_extra_dep_installed
+    if [[ -f "$_marker" ]]; then
+        return 0
+    fi
     if [[ "$mpi_local_rank" -eq 0 ]]; then
         pip install diskcache
+        local _nvrx_dir
+        _nvrx_dir="$(mktemp -d)/nvidia-resiliency-ext"
+        git clone --depth 1 https://github.com/NVIDIA/nvidia-resiliency-ext "${_nvrx_dir}" \
+            && pip install "${_nvrx_dir}"
+        touch "$_marker"
+    else
+        local _waited=0
+        while [[ ! -f "$_marker" && $_waited -lt 600 ]]; do
+            sleep 1
+            _waited=$((_waited + 1))
+        done
     fi
 }
 

tools/launcher/examples/Qwen/qwen3-v0339a-demo/step1_synth.yaml

@@ -0,0 +1,34 @@
+# EAGLE3 offline speculative decoding pipeline — Step 1: Data synthesis
+# for qwen3-v0339a-demo
+
+job_name: qwen3-v0339a-demo_EAGLE3_offline_step1
+
+pipeline:
+  allow_to_fail: false
+  skip: false
+  note:
+
+  global_vars:
+    hf_model: /hf-local/Qwen/qwen3-v0339a-demo
+
+  task_0:
+    script: common/tensorrt_llm/query.sh
+    args:
+      - --model <<global_vars.hf_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: 8
+      container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0

tools/launcher/examples/Qwen/qwen3-v0339a-demo/step2_hidden.yaml

@@ -0,0 +1,29 @@
+# EAGLE3 offline speculative decoding pipeline — Step 2: Dump hidden states
+# for qwen3-v0339a-demo
+
+job_name: qwen3-v0339a-demo_EAGLE3_offline_step2
+
+pipeline:
+  allow_to_fail: false
+  skip: false
+  note:
+
+  global_vars:
+    hf_model: /hf-local/Qwen/qwen3-v0339a-demo
+
+  task_0:
+    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: <<global_vars.hf_model>>
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 1
+      ntasks_per_node: 8
+      gpus_per_node: 8
+      container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0

tools/launcher/examples/Qwen/qwen3-v0339a-demo/step3_train.yaml

@@ -0,0 +1,29 @@
+# EAGLE3 offline speculative decoding pipeline — Step 3: Train EAGLE3 draft head
+# for qwen3-v0339a-demo
+
+job_name: qwen3-v0339a-demo_EAGLE3_offline_step3
+
+pipeline:
+  allow_to_fail: false
+  skip: false
+  note:
+
+  global_vars:
+    hf_model: /hf-local/Qwen/qwen3-v0339a-demo
+
+  task_0:
+    script: common/eagle3/train_eagle.sh
+    args:
+      - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/eagle3.yaml
+      - model.model_name_or_path=<<global_vars.hf_model>>
+      - 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: 8
+      container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0

tools/launcher/examples/Qwen/qwen3-v0339a-demo/step4_speed_eval.yaml

@@ -0,0 +1,33 @@
+# EAGLE3 offline speculative decoding pipeline — Step 4: Benchmark
+# for qwen3-v0339a-demo
+
+job_name: qwen3-v0339a-demo_EAGLE3_offline_step4
+
+pipeline:
+  allow_to_fail: false
+  skip: false
+  note:
+
+  global_vars:
+    hf_model: /hf-local/Qwen/qwen3-v0339a-demo
+
+  task_0:
+    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 1
+    environment:
+      - HF_MODEL_CKPT: <<global_vars.hf_model>>
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 1
+      ntasks_per_node: 1
+      gpus_per_node: 8
+      container: vllm/vllm-openai:latest