debug

Signed-off-by: Ye Yu <yeyu@nvidia.com>

Author: yeyu-nvidia

examples/speculative_decoding/launch_train.sh

@@ -106,6 +106,10 @@ while [ $# -gt 0 ]; do
       if [[ "$1" != *=* ]]; then shift; fi
       NUM_NODES="${1#*=}"
       ;;
+    --head_node_ip*
+      if [[ "$1" != *=* ]]; then shift; fi
+      HEAD_NODE_IP="${1#*=}"
+      ;;
     *)
       >&2 printf "Error: Invalid argument ${1#*=}\n"
       exit 1
@@ -190,12 +194,12 @@ else
   DRAFT_VOCAB_CACHE_ARGS=""
 fi
 
-if [[ "$HEAD_NODE_IP" != "" ]]; then
+if [[ "$NUM_NODES" != 1 ]]; then
   MULTI_NODE_ARGS="--num_processes $TOTAL_GPU \
                    --num_machines $NUM_NODES \
                    --machine_rank $SLURM_PROCID \
                    --rdzv_backend c10d \
-                   --main_process_ip $(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1) \ 
+                   --main_process_ip $HEAD_NODE_IP \ 
                    --main_process_port 29500"
 else
   MULTI_NODE_ARGS=""