Respect CUDA_VISIBLE_DEVICES in data scripts

README.md

@@ -36,7 +36,7 @@ bash scripts/train/train.sh
 
 `train.sh` launches `train.py`, which spawns one worker per visible GPU. Select the algorithm and target model by pointing `config_path` at one of the configs under [config/](./config/) (e.g. `config/dspark/dspark_qwen3_4b.py`); see the script header for the full list of configs, how to override `config_path` / `target_cache_dir`, and how to use `--opts` to override individual config fields. Checkpoints are written to `~/checkpoints/<project_name>/<exp_name>/step_*`.
 
-Hardware: the default configs and scripts assume a single node with 8 GPUs. For fewer GPUs, reduce `CUDA_VISIBLE_DEVICES`.
+Hardware: the default configs and scripts assume a single node with 8 GPUs. For fewer GPUs, reduce `CUDA_VISIBLE_DEVICES`; the data server launcher and training script derive their worker counts from the visible GPU list by default.
 
 
 ## Evaluation

scripts/data/README.md

@@ -22,7 +22,7 @@ train_datasets/qwen3_4b/perfectblend_train_regen.jsonl
 ~/.cache/deepspec/qwen3_4b_target_cache
 ```
 
-The example scripts assume a single machine with eight visible GPUs by default. For fewer GPUs, edit `num_workers` and `CUDA_VISIBLE_DEVICES` in the shell scripts.
+The example scripts assume a single machine with eight visible GPUs by default. For fewer GPUs, set `CUDA_VISIBLE_DEVICES` before running the shell scripts; `num_workers` is derived from the visible GPU list unless you override it explicitly.
 
 ## Step 1: Download And Split Data
 
@@ -60,6 +60,14 @@ By default this starts eight `Qwen/Qwen3-4B` workers on ports `30000` to `30007`
 logs/sglang_qwen3_4b/
 ```
 
+For fewer GPUs, pass the device list to the launcher:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3 bash scripts/data/launch_sglang_server.sh
+```
+
+This starts four workers on ports `30000` to `30003`. Use the same `CUDA_VISIBLE_DEVICES`, `num_workers`, and `start_port` values when running `prepare_data.sh` so the data generation step connects to the right servers.
+
 In another terminal, regenerate the assistant answers:
 
 ```bash

scripts/data/launch_sglang_server.sh

@@ -5,15 +5,27 @@ set -euo pipefail
 #   pip install "sglang[all]"
 # See https://docs.sglang.ai/get_started/install.html for details.
 
-model_path=Qwen/Qwen3-4B
-num_workers=8
-start_port=30000
-start_nccl_port=31000
-host=0.0.0.0
-dtype=bfloat16
-mem_frac=0.9
-log_dir=logs/sglang_qwen3_4b
-heartbeat_interval=300
+model_path=${model_path:-Qwen/Qwen3-4B}
+gpu_list=${gpu_list:-${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}}
+IFS=',' read -r -a gpu_ids <<< "${gpu_list}"
+num_workers=${num_workers:-${#gpu_ids[@]}}
+start_port=${start_port:-30000}
+start_nccl_port=${start_nccl_port:-31000}
+host=${host:-0.0.0.0}
+dtype=${dtype:-bfloat16}
+mem_frac=${mem_frac:-0.9}
+log_dir=${log_dir:-logs/sglang_qwen3_4b}
+heartbeat_interval=${heartbeat_interval:-300}
+
+if (( num_workers < 1 )); then
+    echo "num_workers must be at least 1" >&2
+    exit 1
+fi
+
+if (( num_workers > ${#gpu_ids[@]} )); then
+    echo "num_workers (${num_workers}) cannot exceed available GPUs in gpu_list (${gpu_list})" >&2
+    exit 1
+fi
 
 get_host_ip() {
     local host_ip=""
@@ -93,13 +105,15 @@ cleanup() {
 
 trap cleanup INT TERM EXIT
 
-for ((gpu_id = 0; gpu_id < num_workers; gpu_id++)); do
-    port=$((start_port + gpu_id))
-    nccl_port=$((start_nccl_port + gpu_id))
-    log_file=${log_dir}/worker_${host_ip}_gpu_${gpu_id}_port_${port}.log
+for ((worker_id = 0; worker_id < num_workers; worker_id++)); do
+    gpu_id=${gpu_ids[$worker_id]}
+    gpu_log_id=${gpu_id//\//_}
+    port=$((start_port + worker_id))
+    nccl_port=$((start_nccl_port + worker_id))
+    log_file=${log_dir}/worker_${host_ip}_gpu_${gpu_log_id}_port_${port}.log
 
-    echo "Starting sglang worker ip=${host_ip} gpu=${gpu_id} port=${port} nccl_port=${nccl_port} log=${log_file}"
-    CUDA_VISIBLE_DEVICES=${gpu_id} sglang serve \
+    echo "Starting sglang worker ip=${host_ip} worker=${worker_id} gpu=${gpu_id} port=${port} nccl_port=${nccl_port} log=${log_file}"
+    CUDA_VISIBLE_DEVICES="${gpu_id}" sglang serve \
         --model-path "${model_path}" \
         --host "${host}" \
         --port "${port}" \
@@ -113,10 +127,11 @@ for ((gpu_id = 0; gpu_id < num_workers; gpu_id++)); do
 done
 
 echo "Workers launched:"
-for ((gpu_id = 0; gpu_id < num_workers; gpu_id++)); do
-    port=$((start_port + gpu_id))
-    nccl_port=$((start_nccl_port + gpu_id))
-    echo "  http://${host_ip}:${port} nccl_port=${nccl_port}"
+for ((worker_id = 0; worker_id < num_workers; worker_id++)); do
+    gpu_id=${gpu_ids[$worker_id]}
+    port=$((start_port + worker_id))
+    nccl_port=$((start_nccl_port + worker_id))
+    echo "  worker=${worker_id} gpu=${gpu_id} http://${host_ip}:${port} nccl_port=${nccl_port}"
 done
 echo "Heartbeat interval: ${heartbeat_interval}s"
 print_heartbeat

scripts/data/prepare_data.sh

@@ -1,28 +1,36 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-model_path=Qwen/Qwen3-4B
-config_path=config/dspark/dspark_qwen3_4b.py
+model_path=${model_path:-Qwen/Qwen3-4B}
+config_path=${config_path:-config/dspark/dspark_qwen3_4b.py}
 
-dataset_name=mlabonne/open-perfectblend
-test_size=0.05
-train_split_path=train_datasets/perfectblend_train.jsonl
-eval_data_dir=eval_datasets
+dataset_name=${dataset_name:-mlabonne/open-perfectblend}
+test_size=${test_size:-0.05}
+train_split_path=${train_split_path:-train_datasets/perfectblend_train.jsonl}
+eval_data_dir=${eval_data_dir:-eval_datasets}
 
-train_data_path=train_datasets/qwen3_4b/perfectblend_train_regen.jsonl
-cache_dir=${HOME}/.cache/deepspec/qwen3_4b_target_cache
+train_data_path=${train_data_path:-train_datasets/qwen3_4b/perfectblend_train_regen.jsonl}
+cache_dir=${cache_dir:-${HOME}/.cache/deepspec/qwen3_4b_target_cache}
 
-server_host=127.0.0.1
-num_workers=8
-start_port=30000
-concurrency=32
-temperature=0.7
-top_p=0.8
-top_k=20
-min_p=0
-max_tokens=4096
+server_host=${server_host:-127.0.0.1}
+visible_gpus=${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}
+IFS=',' read -r -a visible_gpu_ids <<< "${visible_gpus}"
+num_workers=${num_workers:-${#visible_gpu_ids[@]}}
+start_port=${start_port:-30000}
+concurrency=${concurrency:-32}
+temperature=${temperature:-0.7}
+top_p=${top_p:-0.8}
+top_k=${top_k:-20}
+min_p=${min_p:-0}
+max_tokens=${max_tokens:-4096}
+local_batch_size=${local_batch_size:-16}
 
-export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}
+if (( num_workers < 1 )); then
+    echo "num_workers must be at least 1" >&2
+    exit 1
+fi
+
+export CUDA_VISIBLE_DEVICES=${visible_gpus}
 export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1}
 export MASTER_PORT=${MASTER_PORT:-29500}
 export RANK=${RANK:-0}
@@ -43,7 +51,7 @@ python scripts/data/download_and_split.py \
 
 mkdir -p "$(dirname "${train_data_path}")"
 
-echo "Step 2/3: generating Qwen3-4B train data: ${train_data_path}"
+echo "Step 2/3: generating ${model_path} train data: ${train_data_path}"
 echo "Start sglang first with: bash scripts/data/launch_sglang_server.sh"
 python scripts/data/generate_train_data.py \
     --model "${model_path}" \
@@ -60,9 +68,9 @@ python scripts/data/generate_train_data.py \
     --output-file-path "${train_data_path}"
 
 echo "Stop sglang before Step 3 if it is using the same GPUs."
-echo "Step 3/3: preparing Qwen3-4B target cache: ${cache_dir}"
+echo "Step 3/3: preparing ${model_path} target cache: ${cache_dir}"
 python scripts/data/prepare_target_cache.py \
     --config "${config_path}" \
     --train-data-path "${train_data_path}" \
     --output-dir "${cache_dir}" \
-    --local-batch-size 16
+    --local-batch-size "${local_batch_size}"