train-student.slurm

#!/bin/bash
#SBATCH --job-name=tiger-train
# Override job name with: sbatch --job-name=<name> train.slurm
#SBATCH --partition=studentkillable
#SBATCH --account=gpu-students
#SBATCH --time=24:00:00
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --mem=32G
#SBATCH --gres=gpu:8
#SBATCH --output=logs/train.%j.log
#SBATCH --error=logs/train.%j.log
#SBATCH --export=ALL

set -euo pipefail

cd /home/yandex/APDL2526a/$USER/TIGER-plus

echo "=== Job Parameters ==="
echo "SLURM_JOB_ID      = ${SLURM_JOB_ID:-N/A}"
echo "SLURM_NODELIST    = ${SLURM_NODELIST:-N/A}"
echo "User              = $USER"
echo "Working directory = $(pwd)"
echo "GPUs requested    = ${SLURM_GPUS:-4}"
echo "Started at        = $(date)"
echo "======================"

mkdir -p logs

# ── Environment setup ──
# Use local /tmp for temp files to avoid NFS cleanup errors
export TMPDIR="$PWD/.tmp-${SLURM_JOB_ID:-$$}"
mkdir -p "$TMPDIR"
export TMP="$TMPDIR"
export TEMP="$TMPDIR"

export HF_HOME="$PWD/.hf_cache"
export HF_DATASETS_CACHE="$HF_HOME/datasets"
export HUGGINGFACE_HUB_CACHE="$HF_HOME/hub"
export TORCH_HOME="$PWD/.torch_cache"
mkdir -p "$TORCH_HOME"
export PYTHONUNBUFFERED=1
export TOKENIZERS_PARALLELISM=false
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# ── Auto-resubmit: preserve original job ID across continuation jobs ──
export TIGER_ORIGINAL_JOB_ID="${TIGER_ORIGINAL_JOB_ID:-${SLURM_JOB_ID}}"
export TIGER_JOB_NAME="${TIGER_JOB_NAME:-${SLURM_JOB_NAME:-tiger-train}}"
export TIGER_PARTITION="${TIGER_PARTITION:-${SLURM_JOB_PARTITION:-killable}}"
export TIGER_ACCOUNT="${TIGER_ACCOUNT:-${SLURM_JOB_ACCOUNT:-}}"
export TIGER_NODELIST="${TIGER_NODELIST:-}"
# Stop 30min before 24h wall time so checkpoint can be saved (23h = 82800s)
export TIGER_TIME_LIMIT_SECS="${TIGER_TIME_LIMIT_SECS:-82800}"
echo "TIGER_ORIGINAL_JOB_ID = $TIGER_ORIGINAL_JOB_ID"
echo "TIGER_PARTITION       = $TIGER_PARTITION"
echo "TIGER_TIME_LIMIT_SECS = $TIGER_TIME_LIMIT_SECS"

# ── Load .env (contains WANDB_API_KEY, HUGGINGFACE_HUB_TOKEN) ──
if [[ -f .env ]]; then
    set -a && source .env && set +a
fi

# ── W&B setup ──
export WANDB_PROJECT="${WANDB_PROJECT:-tiger-plus}"
export WANDB_ENTITY="${WANDB_ENTITY:-selective-entropy-knowledge-distillation}"
export WANDB_DIR="$PWD/.wandb"
mkdir -p "$WANDB_DIR"
export WANDB_CACHE_DIR="$PWD/.wandb_cache"
mkdir -p "$WANDB_CACHE_DIR"
export WANDB_DATA_DIR="$PWD/.wandb_data"
mkdir -p "$WANDB_DATA_DIR"
export WANDB_MODE="${WANDB_MODE:-online}"

# ── Python venv ──
# Two venvs live in the project dir:
#   ./venv3_12  - built on Python 3.12 (e.g. c-005)
#   ./venv      - built on Python 3.10 (SLURM nodes)
# Try venv3_12 first, then venv. If neither works, build ./venv for this node.
PY=""
for CANDIDATE in "$PWD/venv3_12/bin/python" "$PWD/venv/bin/python"; do
    if [[ -x "$CANDIDATE" ]] && "$CANDIDATE" -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
        PY="$CANDIDATE"
        break
    fi
done

if [[ -z "$PY" ]]; then
    echo "No working venv found. Building ./venv for $(python3 -V 2>&1)..."
    rm -rf "$PWD/venv"
    python3 -m venv --without-pip "$PWD/venv"
    curl -sS https://bootstrap.pypa.io/get-pip.py | "$PWD/venv/bin/python"
    "$PWD/venv/bin/python" -m pip install --upgrade pip wheel "setuptools<81"
    # Install requirements excluding pip-system-certs (it corrupts pip's CA
    # bundle mid-install, causing TLS errors on fresh venvs)
    "$PWD/venv/bin/python" -m pip install --no-cache-dir --prefer-binary \
        --extra-index-url https://download.pytorch.org/whl/cu124 \
        $(grep -v '^pip-system-certs' "$PWD/requirements.txt" | tr '\n' ' ')
    PY="$PWD/venv/bin/python"
fi
# ── 8-GPU DDP setup ──
# Disable P2P in case GPUs span different PCIe bridges (safe on all GPU types)
export NCCL_P2P_DISABLE=1
export NCCL_SHM_DISABLE=0
export NCCL_NVML_DISABLE=1

# ── Config file ──
CONFIG_FILE="${CONFIG_FILE:-configs/tiger-echoset-8gpu.yml}"

echo "=== TIGER Training (DDP) ==="
echo "Python: $PY ($($PY -V 2>&1))"
echo "Config: $CONFIG_FILE"
echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set}"
echo "NCCL_P2P_DISABLE=$NCCL_P2P_DISABLE"
echo "TMPDIR: $TMPDIR (local, avoids NFS errors)"
echo "=============================="

# ── Download dataset if needed ──
if [[ ! -f "DataPreProcess/EchoSet/train/mix.json" ]]; then
    echo "Dataset not found, downloading EchoSet from HuggingFace..."
    $PY download_datasets.py --datasets echoset
fi

# ── Unset SLURM env vars so Lightning handles DDP spawning internally ──
# Lightning auto-detects SLURM and limits to SLURM_NTASKS processes.
# Since we want Lightning to spawn DDP workers (not srun), clear these.
SAVED_CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}"
export SLURM_JOB_ID_SAVED="${SLURM_JOB_ID}"
for v in SLURM_NTASKS SLURM_JOB_NAME SLURM_LOCALID SLURM_PROCID SLURM_NODEID SLURM_JOB_ID; do
    unset "$v"
done
export CUDA_VISIBLE_DEVICES="$SAVED_CUDA_VISIBLE_DEVICES"

# ── Launch training (Lightning handles DDP spawning - no torchrun) ──
$PY audio_train.py --conf_dir "$CONFIG_FILE"

echo "=== Training finished at $(date) ==="
echo "SUCCESS: Training complete!"