diff --git a/README.md b/README.md index dbb79990..972bb4cf 100644 --- a/README.md +++ b/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///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 diff --git a/scripts/data/README.md b/scripts/data/README.md index 6c5a6ed5..1201b199 100644 --- a/scripts/data/README.md +++ b/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 diff --git a/scripts/data/launch_sglang_server.sh b/scripts/data/launch_sglang_server.sh index f0b66a78..6e5d3584 100644 --- a/scripts/data/launch_sglang_server.sh +++ b/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 diff --git a/scripts/data/prepare_data.sh b/scripts/data/prepare_data.sh index 53e5db7b..563fcafc 100644 --- a/scripts/data/prepare_data.sh +++ b/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}"