v223rc_exp2 / parallel_train_eval.sh

Add files using upload-large-folder tool

e881749 verified 3 months ago

8.4 kB

#!/usr/bin/env bash
set -uo pipefail

# --------------------------------------------------
# Paths and configuration
# --------------------------------------------------

EXP_ROOT="${EXP_ROOT:-/workspace/v223rc_exp2}"
QA_SCRIPT="${QA_SCRIPT:-/workspace/run_all_qa_hotswap.sh}"
LMF_BIN="${LMF_BIN:-lmf}"

LOG_ROOT="${LOG_ROOT:-/workspace/train_logs}"
STATUS_ROOT="${STATUS_ROOT:-/workspace/train_status}"

NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}"
BASE_QA_PORT="${BASE_QA_PORT:-8000}"
# --------------------------------------------------
# Configurations to skip
# --------------------------------------------------

# Skip training for these configs
SKIP_TRAIN=(
    # lr1b{1..5..1}
    # lr2b{1..4..1}
    # lr3b1
    # lr3b3
    # lr3b4
    # lr3b6
    # lr4b{1..3..1}
    # lr5b{1..3..1}
    # lr5b5
    # lr6b{1..3..1}
)

# Skip QA evaluation for these configs
SKIP_EVAL=(
    # lr1b{2..5..1}
    # lr2b{2..4..1}
    # lr3b3
    # lr3b4
    # lr3b6
    # lr4b{2..3..1}
    # lr5b{1..3..1}
    # lr5b5
    # lr6b{1..3..1}
)

contains() {
    local item="$1"
    shift
    for x in "$@"; do
        [[ "$x" == "$item" ]] && return 0
    done
    return 1
}
# This must match model_name_or_path in the training YAML files.
BASE_MODEL_PATH="${BASE_MODEL_PATH:-/workspace/allenai/Olmo-3-1125-32B}"

mkdir -p "$LOG_ROOT" "$STATUS_ROOT"

# --------------------------------------------------
# Validation
# --------------------------------------------------

if [[ ! -d "$EXP_ROOT" ]]; then
    echo "ERROR: Experiment directory does not exist: $EXP_ROOT" >&2
    exit 1
fi

if [[ ! -f "$QA_SCRIPT" ]]; then
    echo "ERROR: QA script not found: $QA_SCRIPT" >&2
    exit 1
fi

if [[ ! -x "$QA_SCRIPT" ]]; then
    echo "ERROR: QA script is not executable: $QA_SCRIPT" >&2
    echo "Run: chmod +x '$QA_SCRIPT'" >&2
    exit 1
fi

if ! command -v "$LMF_BIN" >/dev/null 2>&1; then
    echo "ERROR: Command not found: $LMF_BIN" >&2
    echo "PATH=$PATH" >&2
    exit 1
fi

if ! command -v nvidia-smi >/dev/null 2>&1; then
    echo "ERROR: nvidia-smi not found." >&2
    exit 1
fi

if [[ "$NUM_GPUS" -lt 1 ]]; then
    echo "ERROR: No GPUs detected." >&2
    exit 1
fi

echo "Detected GPUs:  $NUM_GPUS"
echo "Experiment root: $EXP_ROOT"
echo "QA script:       $QA_SCRIPT"
echo "Base model:      $BASE_MODEL_PATH"
echo "Logs:            $LOG_ROOT"
echo "Statuses:        $STATUS_ROOT"

# --------------------------------------------------
# Build ordered job list
# --------------------------------------------------

declare -a JOBS=()

for b in {0..7}; do
    for lr in {0..7}; do
        # if [[ "$lr" -eq 3 && "$b" -eq 2 ]]; then
        #     echo "Skipping baseline: lr3b2"
        #     continue
        # fi

        config="lr${lr}b${b}"
        yaml_path="${EXP_ROOT}/${config}.yaml"

        if [[ ! -f "$yaml_path" ]]; then
            echo "WARNING: Missing configuration: $yaml_path" >&2
            # exit 1
            continue
        fi

        JOBS+=("${config}|${yaml_path}")
    done
done

TOTAL_JOBS="${#JOBS[@]}"

echo "Total jobs: $TOTAL_JOBS"

# --------------------------------------------------
# Train and evaluate one configuration
# --------------------------------------------------

run_job() {
    local gpu_id="$1"
    local config="$2"
    local yaml_path="$3"

    local qa_port=$((BASE_QA_PORT + gpu_id))
    local log_file="${LOG_ROOT}/${config}.log"
    local server_log="${LOG_ROOT}/${config}_server.log"

    local running_file="${STATUS_ROOT}/${config}.running"
    local success_file="${STATUS_ROOT}/${config}.success"
    local failed_file="${STATUS_ROOT}/${config}.failed"

    rm -f "$running_file" "$success_file" "$failed_file"

    {
        echo "Config: $config"
        echo "GPU: $gpu_id"
        echo "YAML: $yaml_path"
        echo "Started: $(date --iso-8601=seconds)"
    } >"$running_file"

    echo "[GPU $gpu_id] Starting $config"

    if (
        set -euo pipefail

        # This process and all its children see exactly one GPU.
        export CUDA_VISIBLE_DEVICES="$gpu_id"

        echo "=================================================="
        echo "Config:               $config"
        echo "Physical GPU:         $gpu_id"
        echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
        echo "YAML:                 $yaml_path"
        echo "QA port:              $qa_port"
        echo "Started:              $(date --iso-8601=seconds)"
        echo "=================================================="

        python - <<'PY'
import os
import torch

print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("CUDA available:", torch.cuda.is_available())
print("Visible GPU count:", torch.cuda.device_count())

if torch.cuda.is_available():
    print("Visible GPU name:", torch.cuda.get_device_name(0))
PY

            if contains "$config" "${SKIP_TRAIN[@]}"; then
        echo
        echo "===== Skipping training: $config ====="
    else
        echo
        echo "===== Training: $config ====="
    
        "$LMF_BIN" train "$yaml_path"
    fi
    
    if contains "$config" "${SKIP_EVAL[@]}"; then
        echo
        echo "===== Skipping QA evaluation: $config ====="
    else
        echo
        echo "===== QA evaluation: $config ====="
    
        CONFIGS="$config" \
        BASE_MODEL_PATH="/workspace/allenai/Olmo-3.1-32B-Instruct" \
        PORT="$qa_port" \
        SERVER_URL="http://127.0.0.1:${qa_port}/v1/chat/completions" \
        ADAPTER_ADMIN_URL="http://127.0.0.1:${qa_port}" \
        SERVER_LOG="$server_log" \
        START_SERVER=1 \
            "$QA_SCRIPT"
    fi

        echo
        echo "Completed: $(date --iso-8601=seconds)"
    ) >"$log_file" 2>&1; then
        mv "$running_file" "$success_file"
        echo "[GPU $gpu_id] Completed $config"
        return 0
    else
        local exit_code=$?

        {
            echo "Exit code: $exit_code"
            echo "Failed: $(date --iso-8601=seconds)"
            echo "Main log: $log_file"
            echo "Server log: $server_log"
        } >>"$running_file"

        mv "$running_file" "$failed_file"

        echo "[GPU $gpu_id] FAILED $config" >&2
        echo "[GPU $gpu_id] See: $log_file" >&2
        return "$exit_code"
    fi
}

# --------------------------------------------------
# Persistent worker: one worker per GPU
# --------------------------------------------------

gpu_worker() {
    local gpu_id="$1"
    local job_index
    local job
    local config
    local yaml_path
    local failures=0

    echo "[GPU $gpu_id] Worker started"

    # GPU 0 gets jobs 0, 4, 8, ...
    # GPU 1 gets jobs 1, 5, 9, ...
    # etc.
    for ((job_index = gpu_id;
          job_index < TOTAL_JOBS;
          job_index += NUM_GPUS)); do

        job="${JOBS[$job_index]}"
        IFS='|' read -r config yaml_path <<<"$job"

        if ! run_job "$gpu_id" "$config" "$yaml_path"; then
            failures=$((failures + 1))
        fi
    done

    echo "[GPU $gpu_id] Worker finished; failures=$failures"

    if [[ "$failures" -gt 0 ]]; then
        return 1
    fi
}

# --------------------------------------------------
# Start workers
# --------------------------------------------------

declare -a worker_pids=()

for ((gpu_id = 0; gpu_id < NUM_GPUS; gpu_id++)); do
    gpu_worker "$gpu_id" &
    worker_pids+=("$!")
done

echo
echo "Started ${#worker_pids[@]} GPU workers."
echo "Monitor with: watch -n 1 nvidia-smi"
echo "Follow logs:  tail -F ${LOG_ROOT}/*.log"
echo

overall_status=0

for pid in "${worker_pids[@]}"; do
    if ! wait "$pid"; then
        overall_status=1
    fi
done

# --------------------------------------------------
# Summary
# --------------------------------------------------

success_count="$(
    find "$STATUS_ROOT" -maxdepth 1 -type f -name '*.success' | wc -l
)"

failed_count="$(
    find "$STATUS_ROOT" -maxdepth 1 -type f -name '*.failed' | wc -l
)"

echo
echo "=================================================="
echo "Parallel training and QA finished"
echo "Successful: $success_count"
echo "Failed:     $failed_count"
echo "Total:      $TOTAL_JOBS"
echo "Logs:       $LOG_ROOT"
echo "=================================================="

if [[ "$failed_count" -gt 0 ]]; then
    echo
    echo "Failed configurations:"

    for file in "$STATUS_ROOT"/*.failed; do
        [[ -e "$file" ]] || continue
        basename "$file" .failed
    done
fi

exit "$overall_status"

	#!/usr/bin/env bash
	set -uo pipefail

	# --------------------------------------------------
	# Paths and configuration
	# --------------------------------------------------

	EXP_ROOT="${EXP_ROOT:-/workspace/v223rc_exp2}"
	QA_SCRIPT="${QA_SCRIPT:-/workspace/run_all_qa_hotswap.sh}"
	LMF_BIN="${LMF_BIN:-lmf}"

	LOG_ROOT="${LOG_ROOT:-/workspace/train_logs}"
	STATUS_ROOT="${STATUS_ROOT:-/workspace/train_status}"

	NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L \| wc -l)}"
	BASE_QA_PORT="${BASE_QA_PORT:-8000}"
	# --------------------------------------------------
	# Configurations to skip
	# --------------------------------------------------

	# Skip training for these configs
	SKIP_TRAIN=(
	# lr1b{1..5..1}
	# lr2b{1..4..1}
	# lr3b1
	# lr3b3
	# lr3b4
	# lr3b6
	# lr4b{1..3..1}
	# lr5b{1..3..1}
	# lr5b5
	# lr6b{1..3..1}
	)

	# Skip QA evaluation for these configs
	SKIP_EVAL=(
	# lr1b{2..5..1}
	# lr2b{2..4..1}
	# lr3b3
	# lr3b4
	# lr3b6
	# lr4b{2..3..1}
	# lr5b{1..3..1}
	# lr5b5
	# lr6b{1..3..1}
	)

	contains() {
	local item="$1"
	shift
	for x in "$@"; do
	[[ "$x" == "$item" ]] && return 0
	done
	return 1
	}
	# This must match model_name_or_path in the training YAML files.
	BASE_MODEL_PATH="${BASE_MODEL_PATH:-/workspace/allenai/Olmo-3-1125-32B}"

	mkdir -p "$LOG_ROOT" "$STATUS_ROOT"

	# --------------------------------------------------
	# Validation
	# --------------------------------------------------

	if [[ ! -d "$EXP_ROOT" ]]; then
	echo "ERROR: Experiment directory does not exist: $EXP_ROOT" >&2
	exit 1
	fi

	if [[ ! -f "$QA_SCRIPT" ]]; then
	echo "ERROR: QA script not found: $QA_SCRIPT" >&2
	exit 1
	fi

	if [[ ! -x "$QA_SCRIPT" ]]; then
	echo "ERROR: QA script is not executable: $QA_SCRIPT" >&2
	echo "Run: chmod +x '$QA_SCRIPT'" >&2
	exit 1
	fi

	if ! command -v "$LMF_BIN" >/dev/null 2>&1; then
	echo "ERROR: Command not found: $LMF_BIN" >&2
	echo "PATH=$PATH" >&2
	exit 1
	fi

	if ! command -v nvidia-smi >/dev/null 2>&1; then
	echo "ERROR: nvidia-smi not found." >&2
	exit 1
	fi

	if [[ "$NUM_GPUS" -lt 1 ]]; then
	echo "ERROR: No GPUs detected." >&2
	exit 1
	fi

	echo "Detected GPUs: $NUM_GPUS"
	echo "Experiment root: $EXP_ROOT"
	echo "QA script: $QA_SCRIPT"
	echo "Base model: $BASE_MODEL_PATH"
	echo "Logs: $LOG_ROOT"
	echo "Statuses: $STATUS_ROOT"

	# --------------------------------------------------
	# Build ordered job list
	# --------------------------------------------------

	declare -a JOBS=()

	for b in {0..7}; do
	for lr in {0..7}; do
	# if [[ "$lr" -eq 3 && "$b" -eq 2 ]]; then
	# echo "Skipping baseline: lr3b2"
	# continue
	# fi

	config="lr${lr}b${b}"
	yaml_path="${EXP_ROOT}/${config}.yaml"

	if [[ ! -f "$yaml_path" ]]; then
	echo "WARNING: Missing configuration: $yaml_path" >&2
	# exit 1
	continue
	fi

	JOBS+=("${config}\|${yaml_path}")
	done
	done

	TOTAL_JOBS="${#JOBS[@]}"

	echo "Total jobs: $TOTAL_JOBS"

	# --------------------------------------------------
	# Train and evaluate one configuration
	# --------------------------------------------------

	run_job() {
	local gpu_id="$1"
	local config="$2"
	local yaml_path="$3"

	local qa_port=$((BASE_QA_PORT + gpu_id))
	local log_file="${LOG_ROOT}/${config}.log"
	local server_log="${LOG_ROOT}/${config}_server.log"

	local running_file="${STATUS_ROOT}/${config}.running"
	local success_file="${STATUS_ROOT}/${config}.success"
	local failed_file="${STATUS_ROOT}/${config}.failed"

	rm -f "$running_file" "$success_file" "$failed_file"

	{
	echo "Config: $config"
	echo "GPU: $gpu_id"
	echo "YAML: $yaml_path"
	echo "Started: $(date --iso-8601=seconds)"
	} >"$running_file"

	echo "[GPU $gpu_id] Starting $config"

	if (
	set -euo pipefail

	# This process and all its children see exactly one GPU.
	export CUDA_VISIBLE_DEVICES="$gpu_id"

	echo "=================================================="
	echo "Config: $config"
	echo "Physical GPU: $gpu_id"
	echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
	echo "YAML: $yaml_path"
	echo "QA port: $qa_port"
	echo "Started: $(date --iso-8601=seconds)"
	echo "=================================================="

	python - <<'PY'
	import os
	import torch

	print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES"))
	print("CUDA available:", torch.cuda.is_available())
	print("Visible GPU count:", torch.cuda.device_count())

	if torch.cuda.is_available():
	print("Visible GPU name:", torch.cuda.get_device_name(0))
	PY

	if contains "$config" "${SKIP_TRAIN[@]}"; then
	echo
	echo "===== Skipping training: $config ====="
	else
	echo
	echo "===== Training: $config ====="

	"$LMF_BIN" train "$yaml_path"
	fi

	if contains "$config" "${SKIP_EVAL[@]}"; then
	echo
	echo "===== Skipping QA evaluation: $config ====="
	else
	echo
	echo "===== QA evaluation: $config ====="

	CONFIGS="$config" \
	BASE_MODEL_PATH="/workspace/allenai/Olmo-3.1-32B-Instruct" \
	PORT="$qa_port" \
	SERVER_URL="http://127.0.0.1:${qa_port}/v1/chat/completions" \
	ADAPTER_ADMIN_URL="http://127.0.0.1:${qa_port}" \
	SERVER_LOG="$server_log" \
	START_SERVER=1 \
	"$QA_SCRIPT"
	fi

	echo
	echo "Completed: $(date --iso-8601=seconds)"
	) >"$log_file" 2>&1; then
	mv "$running_file" "$success_file"
	echo "[GPU $gpu_id] Completed $config"
	return 0
	else
	local exit_code=$?

	{
	echo "Exit code: $exit_code"
	echo "Failed: $(date --iso-8601=seconds)"
	echo "Main log: $log_file"
	echo "Server log: $server_log"
	} >>"$running_file"

	mv "$running_file" "$failed_file"

	echo "[GPU $gpu_id] FAILED $config" >&2
	echo "[GPU $gpu_id] See: $log_file" >&2
	return "$exit_code"
	fi
	}

	# --------------------------------------------------
	# Persistent worker: one worker per GPU
	# --------------------------------------------------

	gpu_worker() {
	local gpu_id="$1"
	local job_index
	local job
	local config
	local yaml_path
	local failures=0

	echo "[GPU $gpu_id] Worker started"

	# GPU 0 gets jobs 0, 4, 8, ...
	# GPU 1 gets jobs 1, 5, 9, ...
	# etc.
	for ((job_index = gpu_id;
	job_index < TOTAL_JOBS;
	job_index += NUM_GPUS)); do

	job="${JOBS[$job_index]}"
	IFS='\|' read -r config yaml_path <<<"$job"

	if ! run_job "$gpu_id" "$config" "$yaml_path"; then
	failures=$((failures + 1))
	fi
	done

	echo "[GPU $gpu_id] Worker finished; failures=$failures"

	if [[ "$failures" -gt 0 ]]; then
	return 1
	fi
	}

	# --------------------------------------------------
	# Start workers
	# --------------------------------------------------

	declare -a worker_pids=()

	for ((gpu_id = 0; gpu_id < NUM_GPUS; gpu_id++)); do
	gpu_worker "$gpu_id" &
	worker_pids+=("$!")
	done

	echo
	echo "Started ${#worker_pids[@]} GPU workers."
	echo "Monitor with: watch -n 1 nvidia-smi"
	echo "Follow logs: tail -F ${LOG_ROOT}/*.log"
	echo

	overall_status=0

	for pid in "${worker_pids[@]}"; do
	if ! wait "$pid"; then
	overall_status=1
	fi
	done

	# --------------------------------------------------
	# Summary
	# --------------------------------------------------

	success_count="$(
	find "$STATUS_ROOT" -maxdepth 1 -type f -name '*.success' \| wc -l
	)"

	failed_count="$(
	find "$STATUS_ROOT" -maxdepth 1 -type f -name '*.failed' \| wc -l
	)"

	echo
	echo "=================================================="
	echo "Parallel training and QA finished"
	echo "Successful: $success_count"
	echo "Failed: $failed_count"
	echo "Total: $TOTAL_JOBS"
	echo "Logs: $LOG_ROOT"
	echo "=================================================="

	if [[ "$failed_count" -gt 0 ]]; then
	echo
	echo "Failed configurations:"

	for file in "$STATUS_ROOT"/*.failed; do
	[[ -e "$file" ]] \|\| continue
	basename "$file" .failed
	done
	fi

	exit "$overall_status"