add slurm support in scripts

Signed-off-by: h-guo18 <[email protected]>

Author: h-guo18

examples/speculative_decoding/collect_hidden_states/run_trtllm_compute_hiddens.sh

@@ -17,7 +17,9 @@
 # This script computes hidden states using TensorRT-LLM and saves them to
 # the specified output directory.
 
+export TLLM_LOG_LEVEL="error";
 python3 collect_hidden_states/compute_hidden_states_trtllm.py \
   --model meta-llama/Llama-3.2-1B-Instruct \
   --input-file synthetic_conversations/daring-anteater.jsonl \
-  --output-dir /mnt/md0/eagle-hidden-states/llama1b/daring_anteater/
+  --output-dir /mnt/md0/eagle-hidden-states/llama1b/daring_anteater/
+  

examples/speculative_decoding/collect_hidden_states/run_trtllm_compute_hiddens_dp.sh

@@ -25,13 +25,20 @@
 INPUT_FILE=synthetic_conversations/daring-anteater.jsonl
 OUTPUT_DIR=/mnt/md0/eagle-hidden-states/llama1b/daring_anteater/
 DP_SIZE=8
+MODEL=meta-llama/Llama-3.2-1B-Instruct
+export TLLM_LOG_LEVEL="error";
 
 split -n l/$DP_SIZE --numeric-suffixes=0 -d --additional-suffix=.jsonl $INPUT_FILE /tmp/part-
 
 for i in $(seq 0 $((DP_SIZE-1)))
 do
-CUDA_VISIBLE_DEVICES=$i python3 collect_hidden_states/compute_hidden_states_trtllm.py --model meta-llama/Llama-3.2-1B-Instruct --input-file /tmp/part-0${i}.jsonl --output-dir $OUTPUT_DIR --dp-rank $i &
+
+export CUDA_VISIBLE_DEVICES=$i;  python3 collect_hidden_states/compute_hidden_states_trtllm.py --model $MODEL --input-file /tmp/part-0${i}.jsonl --output-dir $OUTPUT_DIR --dp-rank $i &
+
+# #On SLURM:
+# PORT=$((10012 + i)); export TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR="tcp://127.0.0.1:$PORT"; trtllm-llmapi-launch python3 collect_hidden_states/compute_hidden_states_trtllm.py --model $MODEL --input-file /tmp/part-0${i}.jsonl --output-dir $OUTPUT_DIR --dp-rank $i 
+
 done
 wait
 
-rm /tmp/part-*.jsonl
+rm /tmp/part-*.jsonl