127 lines
4.5 KiB
Bash
127 lines
4.5 KiB
Bash
#!/usr/bin/env bash
|
||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||
# SPDX-License-Identifier: Apache-2.0
|
||
|
||
#
|
||
# Launch data curation across multiple GPUs / nodes.
|
||
#
|
||
# Each GPU runs one independent vLLM worker that processes a disjoint shard
|
||
# of the dataset. No torch.distributed communication is needed — each worker
|
||
# is a standalone process with its own rank derived from environment variables.
|
||
#
|
||
# ── Single node, 8 GPUs (tp=1, 8 workers) ────────────────────────────────
|
||
# bash data_curation/run_curation.sh \
|
||
# --model Qwen/Qwen3-4B \
|
||
# --input data.jsonl \
|
||
# --output-dir output/ \
|
||
# --num-gpus 8
|
||
#
|
||
# ── Single node, 2 GPUs (tp=2, 1 worker) ─────────────────────────────────
|
||
# bash data_curation/run_curation.sh \
|
||
# --model Qwen/Qwen3-8B \
|
||
# --input data.jsonl \
|
||
# --output-dir output/ \
|
||
# --num-gpus 2 \
|
||
# --tensor-parallel-size 2
|
||
#
|
||
# ── Multi-node (2 nodes × 8 GPUs, tp=1, 16 workers) ─────────────────────
|
||
# # On node 0:
|
||
# NODE_RANK=0 NUM_NODES=2 bash data_curation/run_curation.sh \
|
||
# --model Qwen/Qwen3-4B \
|
||
# --input data.jsonl \
|
||
# --output-dir output/ \
|
||
# --num-gpus 8
|
||
#
|
||
# # On node 1:
|
||
# NODE_RANK=1 NUM_NODES=2 bash data_curation/run_curation.sh \
|
||
# --model Qwen/Qwen3-4B \
|
||
# --input data.jsonl \
|
||
# --output-dir output/ \
|
||
# --num-gpus 8
|
||
#
|
||
# Environment variables (optional):
|
||
# NUM_NODES – total number of nodes (default: 1)
|
||
# NODE_RANK – rank of this node (default: 0)
|
||
# ──────────────────────────────────────────────────────────────────────────
|
||
|
||
set -euo pipefail
|
||
|
||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||
|
||
# ── Parse --num-gpus and --tensor-parallel-size from args ─────────────────
|
||
NUM_GPUS=1
|
||
TP=1
|
||
PIPELINE_ARGS=()
|
||
|
||
while [[ $# -gt 0 ]]; do
|
||
case "$1" in
|
||
--num-gpus)
|
||
NUM_GPUS="$2"; shift 2 ;;
|
||
--tensor-parallel-size)
|
||
TP="$2"; PIPELINE_ARGS+=("--tensor-parallel-size" "$2"); shift 2 ;;
|
||
*)
|
||
PIPELINE_ARGS+=("$1"); shift ;;
|
||
esac
|
||
done
|
||
|
||
# ── Compute worker layout ────────────────────────────────────────────────
|
||
NUM_NODES="${NUM_NODES:-1}"
|
||
NODE_RANK="${NODE_RANK:-0}"
|
||
WORKERS_PER_NODE=$(( NUM_GPUS / TP ))
|
||
WORLD_SIZE=$(( WORKERS_PER_NODE * NUM_NODES ))
|
||
|
||
echo "=== Data Curation Launch ==="
|
||
echo " Nodes: ${NUM_NODES} (this node: ${NODE_RANK})"
|
||
echo " GPUs per node: ${NUM_GPUS}"
|
||
echo " TP size: ${TP}"
|
||
echo " Workers per node: ${WORKERS_PER_NODE}"
|
||
echo " World size: ${WORLD_SIZE}"
|
||
echo " Pipeline args: ${PIPELINE_ARGS[*]}"
|
||
echo "============================"
|
||
|
||
# ── Launch workers ───────────────────────────────────────────────────────
|
||
PIDS=()
|
||
for (( LOCAL=0; LOCAL<WORKERS_PER_NODE; LOCAL++ )); do
|
||
GLOBAL_RANK=$(( NODE_RANK * WORKERS_PER_NODE + LOCAL ))
|
||
GPU_START=$(( LOCAL * TP ))
|
||
GPU_END=$(( GPU_START + TP - 1 ))
|
||
|
||
# Build CUDA_VISIBLE_DEVICES string, e.g. "0" or "2,3"
|
||
GPUS=""
|
||
for (( g=GPU_START; g<=GPU_END; g++ )); do
|
||
[[ -n "$GPUS" ]] && GPUS="${GPUS},"
|
||
GPUS="${GPUS}${g}"
|
||
done
|
||
|
||
echo "[Node ${NODE_RANK}] Launching worker rank=${GLOBAL_RANK} on GPU(s) ${GPUS}"
|
||
|
||
CUDA_VISIBLE_DEVICES="${GPUS}" \
|
||
RANK="${GLOBAL_RANK}" \
|
||
WORLD_SIZE="${WORLD_SIZE}" \
|
||
python "${SCRIPT_DIR}/pipeline.py" \
|
||
--rank "${GLOBAL_RANK}" \
|
||
--world-size "${WORLD_SIZE}" \
|
||
"${PIPELINE_ARGS[@]}" \
|
||
> >(sed "s/^/[rank${GLOBAL_RANK}] /") \
|
||
2>&1 &
|
||
|
||
PIDS+=($!)
|
||
done
|
||
|
||
# ── Wait for all workers ─────────────────────────────────────────────────
|
||
echo "Waiting for ${#PIDS[@]} workers to finish..."
|
||
FAILED=0
|
||
for PID in "${PIDS[@]}"; do
|
||
if ! wait "$PID"; then
|
||
echo "Worker PID ${PID} failed!"
|
||
FAILED=1
|
||
fi
|
||
done
|
||
|
||
if [[ $FAILED -eq 1 ]]; then
|
||
echo "Some workers failed. Check logs above."
|
||
exit 1
|
||
fi
|
||
|
||
echo "All workers finished successfully."
|