diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm index 58613e081455..c29f995f3c68 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm @@ -83,8 +83,6 @@ echo "ntasks_per_node: ${ntasks_per_node}" echo "===========================================" -nsys_on="" -# nsys_on=${full_logdir} # Uncomment this line to enable Nsys profiling numa_bind=true # Only allocate memory from nodes, this only works on GB200 ctx_max_seq_len=$((isl + 10)) gen_max_seq_len=$((isl + osl + 10)) @@ -96,6 +94,9 @@ logdir=${workdir}/slurm-${SLURM_JOB_ID}/benchmark-${isl}-${osl} mkdir -p ${logdir} full_logdir=${logdir}/ctx${num_ctx_servers}_gen${num_gen_servers}_dep${gen_tp_size}_batch${gen_batch_size}_eplb${eplb_num_slots}_mtp${mtp_size} +nsys_on="" +# nsys_on=${full_logdir} # Uncomment this line to enable Nsys profiling + echo "concurrency: ${concurrency}" ctx_gpus=$((num_ctx_servers * ctx_tp_size * ctx_pp_size)) diff --git a/examples/disaggregated/slurm/benchmark/start_worker.sh b/examples/disaggregated/slurm/benchmark/start_worker.sh index 1a5041b1e33c..9c2f48860e36 100644 --- a/examples/disaggregated/slurm/benchmark/start_worker.sh +++ b/examples/disaggregated/slurm/benchmark/start_worker.sh @@ -65,14 +65,14 @@ else nsys_file=${nsys_folder}/nsys_worker_proc_${instance_id}_${SLURM_PROCID} export TLLM_PROFILE_RECORD_GC=1 export TLLM_NVTX_DEBUG=1 - if [ "${role}" = "GEN" ]; then + if [ "${role}" = "GEN" ] && [ "$SLURM_PROCID" = "0" ]; then export TLLM_PROFILE_START_STOP=200-250 nsys_prefix="nsys profile -e \"NSYS_MPI_STORE_TEAMS_PER_RANK=1\" -o ${nsys_file} -f true -t cuda,nvtx,python-gil -c cudaProfilerApi --cuda-graph-trace node --capture-range-end=stop --gpu-metrics-devices=none" echo "nsys_prefix: ${nsys_prefix}" elif [ "${role}" = "CTX" ]; then echo "nsys is not enabled on ctx_gpus" fi - trtllm-llmapi-launch ${numa_bind_cmd} ${nsys_prefix} \ + ${nsys_prefix} trtllm-llmapi-launch ${numa_bind_cmd} \ trtllm-serve ${model_path} \ --host $(hostname) --port ${port} \ --extra_llm_api_options ${config_file}