46 lines
1.4 KiB
Bash
46 lines
1.4 KiB
Bash
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
REPO="$(cd "$SCRIPT_DIR/../.." && pwd)"
|
|
IMAGE="${IMAGE:-aha-l2a-qwen3-repro:torch2.9.1-cu128}"
|
|
GPU_COUNT="$(nvidia-smi -L | wc -l)"
|
|
if [[ "$GPU_COUNT" -lt 1 ]]; then
|
|
echo "At least one NVIDIA GPU is required." >&2
|
|
exit 2
|
|
fi
|
|
GPU_AHA="${GPU_AHA:-0}"
|
|
if [[ -z "${GPU_L2A+x}" ]]; then
|
|
GPU_L2A=0
|
|
if [[ "$GPU_COUNT" -ge 2 ]]; then
|
|
GPU_L2A=1
|
|
fi
|
|
fi
|
|
|
|
for gpu in "$GPU_AHA" "$GPU_L2A"; do
|
|
if [[ ! "$gpu" =~ ^[0-9]+$ ]] || (( gpu >= GPU_COUNT )); then
|
|
echo "Invalid GPU index $gpu; nvidia-smi reports $GPU_COUNT visible GPU(s)." >&2
|
|
exit 2
|
|
fi
|
|
done
|
|
|
|
echo "Training mode: one GPU per experimental variant (no DDP/FSDP/DeepSpeed)."
|
|
echo "AHA variant GPU: $GPU_AHA"
|
|
echo "L2A-style variant GPU: $GPU_L2A"
|
|
if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
|
|
echo "The two variants will run sequentially on the same GPU."
|
|
else
|
|
echo "The two variants will run concurrently on two independent GPUs."
|
|
fi
|
|
if [[ "$GPU_COUNT" -gt 2 ]]; then
|
|
echo "Visible GPUs: $GPU_COUNT; this recipe intentionally does not use all GPUs."
|
|
fi
|
|
|
|
docker build -t "$IMAGE" -f "$REPO/recipe/Dockerfile" "$REPO/recipe"
|
|
docker run --rm --gpus all --ipc=host \
|
|
--ulimit memlock=-1 --ulimit stack=67108864 \
|
|
-e GPU_AHA="$GPU_AHA" -e GPU_L2A="$GPU_L2A" \
|
|
-v "$REPO:/workspace" \
|
|
-w /workspace/recipe \
|
|
"$IMAGE" bash scripts/train_both.sh
|