10f9cfadb7
Add train_launcher.sh that detects local GPU VRAM and routes training to local or Scaleway cloud. Auto-selects batch size per model based on available VRAM tier. Maps model names to actual ml/examples/train_*.rs cargo targets. Document Scaleway GPU instance types, setup procedure, batch size tables, and cost estimates. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
227 lines
7.6 KiB
Bash
Executable File
227 lines
7.6 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
# GPU Training Launcher -- detects GPU and selects local or cloud path
|
|
#
|
|
# Usage:
|
|
# ./scripts/train_launcher.sh --model dqn [--data path/to/data] [--cloud] [--epochs N]
|
|
# ./scripts/train_launcher.sh --model ppo --data data/databento/ES.FUT/ohlcv-1m/ --epochs 50
|
|
# ./scripts/train_launcher.sh --model tft --cloud
|
|
#
|
|
# Models: dqn, ppo, tft, mamba2, cfc
|
|
# Flags:
|
|
# --cloud Force cloud training even if local GPU is available
|
|
# --epochs Number of training epochs (default: 20)
|
|
# --data Path to training data (parquet file or directory)
|
|
# --lr Learning rate (default: 0.0001)
|
|
# --output Output directory for trained models (default: ml/trained_models)
|
|
|
|
MODEL=""
|
|
DATA_PATH=""
|
|
CLOUD=false
|
|
EPOCHS=20
|
|
LEARNING_RATE=0.0001
|
|
OUTPUT_DIR="ml/trained_models"
|
|
EXTRA_ARGS=()
|
|
|
|
while [[ $# -gt 0 ]]; do
|
|
case "$1" in
|
|
--model) MODEL="$2"; shift 2 ;;
|
|
--data) DATA_PATH="$2"; shift 2 ;;
|
|
--cloud) CLOUD=true; shift ;;
|
|
--epochs) EPOCHS="$2"; shift 2 ;;
|
|
--lr) LEARNING_RATE="$2"; shift 2 ;;
|
|
--output) OUTPUT_DIR="$2"; shift 2 ;;
|
|
--help|-h) usage; exit 0 ;;
|
|
*) EXTRA_ARGS+=("$1"); shift ;;
|
|
esac
|
|
done
|
|
|
|
usage() {
|
|
cat <<'USAGE'
|
|
GPU Training Launcher -- detects GPU and routes to local or cloud
|
|
|
|
Usage:
|
|
./scripts/train_launcher.sh --model <model> [OPTIONS]
|
|
|
|
Models: dqn, ppo, tft, mamba2, cfc
|
|
|
|
Options:
|
|
--model <name> Model to train (required)
|
|
--data <path> Training data parquet file or directory
|
|
--epochs <N> Number of training epochs (default: 20)
|
|
--lr <rate> Learning rate (default: 0.0001)
|
|
--output <dir> Output directory (default: ml/trained_models)
|
|
--cloud Force cloud training
|
|
--help Show this help
|
|
|
|
Examples:
|
|
./scripts/train_launcher.sh --model dqn --epochs 100
|
|
./scripts/train_launcher.sh --model ppo --data test_data/ES_FUT_180d.parquet --epochs 50
|
|
./scripts/train_launcher.sh --model tft --cloud
|
|
USAGE
|
|
}
|
|
|
|
if [[ -z "$MODEL" ]]; then
|
|
echo "Error: --model is required (dqn, ppo, tft, mamba2, cfc)"
|
|
echo "Run with --help for usage information"
|
|
exit 1
|
|
fi
|
|
|
|
VALID_MODELS="dqn ppo tft mamba2 cfc"
|
|
if ! echo "$VALID_MODELS" | grep -qw "$MODEL"; then
|
|
echo "Error: Invalid model '$MODEL'. Valid: $VALID_MODELS"
|
|
exit 1
|
|
fi
|
|
|
|
# Map model names to cargo example targets
|
|
# These are the actual ml/examples/train_*.rs binaries
|
|
declare -A EXAMPLE_MAP
|
|
EXAMPLE_MAP[dqn]="train_dqn"
|
|
EXAMPLE_MAP[ppo]="train_ppo_parquet"
|
|
EXAMPLE_MAP[tft]="train_tft_parquet"
|
|
EXAMPLE_MAP[mamba2]="train_mamba2_parquet"
|
|
EXAMPLE_MAP[cfc]="train_liquid_dbn"
|
|
|
|
EXAMPLE_NAME="${EXAMPLE_MAP[$MODEL]}"
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPU Detection
|
|
# ---------------------------------------------------------------------------
|
|
GPU_DETECTED=false
|
|
VRAM_MB=0
|
|
GPU_NAME="none"
|
|
|
|
if command -v nvidia-smi &>/dev/null; then
|
|
VRAM_MB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 || echo "0")
|
|
if [[ "$VRAM_MB" -gt 0 ]]; then
|
|
GPU_DETECTED=true
|
|
GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1 || echo "unknown")
|
|
echo "Local GPU detected: $GPU_NAME (${VRAM_MB}MB VRAM)"
|
|
fi
|
|
else
|
|
echo "No NVIDIA GPU detected (nvidia-smi not found)"
|
|
fi
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Batch size selection based on VRAM and model
|
|
# ---------------------------------------------------------------------------
|
|
# These values are empirically tested; PPO is constrained on 4GB (max 230).
|
|
# See scripts/measure_vram.sh for VRAM profiling methodology.
|
|
declare -A BS_TIER_LOW # < 4 GB
|
|
declare -A BS_TIER_MED # 4-8 GB (e.g. RTX 3050 Ti 4GB)
|
|
declare -A BS_TIER_HIGH # 8-16 GB (e.g. RTX 3070 8GB)
|
|
declare -A BS_TIER_MAX # 16+ GB (e.g. L4 24GB)
|
|
|
|
BS_TIER_LOW[dqn]=32; BS_TIER_MED[dqn]=128; BS_TIER_HIGH[dqn]=256; BS_TIER_MAX[dqn]=512
|
|
BS_TIER_LOW[ppo]=32; BS_TIER_MED[ppo]=230; BS_TIER_HIGH[ppo]=512; BS_TIER_MAX[ppo]=1024
|
|
BS_TIER_LOW[tft]=8; BS_TIER_MED[tft]=32; BS_TIER_HIGH[tft]=64; BS_TIER_MAX[tft]=256
|
|
BS_TIER_LOW[mamba2]=8; BS_TIER_MED[mamba2]=64; BS_TIER_HIGH[mamba2]=128; BS_TIER_MAX[mamba2]=512
|
|
BS_TIER_LOW[cfc]=32; BS_TIER_MED[cfc]=128; BS_TIER_HIGH[cfc]=256; BS_TIER_MAX[cfc]=512
|
|
|
|
BATCH_SIZE=64
|
|
if [[ "$GPU_DETECTED" == true ]]; then
|
|
if [[ "$VRAM_MB" -lt 4096 ]]; then
|
|
BATCH_SIZE="${BS_TIER_LOW[$MODEL]}"
|
|
elif [[ "$VRAM_MB" -lt 8192 ]]; then
|
|
BATCH_SIZE="${BS_TIER_MED[$MODEL]}"
|
|
elif [[ "$VRAM_MB" -lt 16384 ]]; then
|
|
BATCH_SIZE="${BS_TIER_HIGH[$MODEL]}"
|
|
else
|
|
BATCH_SIZE="${BS_TIER_MAX[$MODEL]}"
|
|
fi
|
|
fi
|
|
|
|
echo ""
|
|
echo "Configuration:"
|
|
echo " Model: $MODEL (example: $EXAMPLE_NAME)"
|
|
echo " Epochs: $EPOCHS"
|
|
echo " Batch size: $BATCH_SIZE"
|
|
echo " Learning rate: $LEARNING_RATE"
|
|
echo " Output: $OUTPUT_DIR"
|
|
if [[ -n "$DATA_PATH" ]]; then
|
|
echo " Data: $DATA_PATH"
|
|
fi
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Build the cargo command arguments
|
|
# ---------------------------------------------------------------------------
|
|
build_cargo_args() {
|
|
local args=()
|
|
args+=(--epochs "$EPOCHS")
|
|
args+=(--batch-size "$BATCH_SIZE")
|
|
args+=(--learning-rate "$LEARNING_RATE")
|
|
args+=(--output-dir "$OUTPUT_DIR")
|
|
|
|
if [[ -n "$DATA_PATH" ]]; then
|
|
args+=(--parquet-file "$DATA_PATH")
|
|
fi
|
|
|
|
# Pass through any extra arguments
|
|
if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then
|
|
args+=("${EXTRA_ARGS[@]}")
|
|
fi
|
|
|
|
echo "${args[@]}"
|
|
}
|
|
|
|
CARGO_EXTRA=$(build_cargo_args)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Route to local or cloud
|
|
# ---------------------------------------------------------------------------
|
|
CLOUD_HOST="${FOXHUNT_GPU_HOST:-gpu.fxhnt.ai}"
|
|
CLOUD_DIR="${FOXHUNT_CLOUD_DIR:-/opt/foxhunt}"
|
|
|
|
need_cloud() {
|
|
# Cloud is needed when: forced, no GPU, or insufficient VRAM (< 4GB)
|
|
[[ "$CLOUD" == true ]] || [[ "$GPU_DETECTED" == false ]] || [[ "$VRAM_MB" -lt 4096 ]]
|
|
}
|
|
|
|
if need_cloud; then
|
|
echo ""
|
|
echo "==> Routing to cloud training at $CLOUD_HOST"
|
|
echo ""
|
|
|
|
# Build the remote command
|
|
REMOTE_CMD="cd $CLOUD_DIR && git pull --ff-only && SQLX_OFFLINE=true cargo run --release -p ml --example $EXAMPLE_NAME --features cuda -- $CARGO_EXTRA"
|
|
|
|
echo "To run on cloud GPU:"
|
|
echo ""
|
|
echo " ssh training@$CLOUD_HOST '$REMOTE_CMD'"
|
|
echo ""
|
|
echo "To sync model artifacts back:"
|
|
echo ""
|
|
echo " rsync -avz training@$CLOUD_HOST:$CLOUD_DIR/checkpoints/ ./checkpoints/"
|
|
echo " rsync -avz training@$CLOUD_HOST:$CLOUD_DIR/ml/trained_models/ ./ml/trained_models/"
|
|
echo ""
|
|
|
|
# If we have SSH access, offer to run it directly
|
|
if [[ "$CLOUD" == true ]]; then
|
|
read -rp "Execute on cloud now? [y/N] " CONFIRM
|
|
if [[ "$CONFIRM" =~ ^[Yy]$ ]]; then
|
|
echo "Connecting to $CLOUD_HOST..."
|
|
# shellcheck disable=SC2029
|
|
ssh "training@$CLOUD_HOST" "$REMOTE_CMD"
|
|
echo ""
|
|
echo "Syncing model artifacts..."
|
|
rsync -avz "training@$CLOUD_HOST:$CLOUD_DIR/$OUTPUT_DIR/" "./$OUTPUT_DIR/"
|
|
rsync -avz "training@$CLOUD_HOST:$CLOUD_DIR/checkpoints/" "./checkpoints/"
|
|
echo "Done."
|
|
else
|
|
echo "Skipped. Run the SSH command above manually."
|
|
fi
|
|
fi
|
|
else
|
|
echo ""
|
|
echo "==> Running local training on $GPU_NAME"
|
|
echo ""
|
|
|
|
mkdir -p "$OUTPUT_DIR"
|
|
|
|
export SQLX_OFFLINE=true
|
|
|
|
# shellcheck disable=SC2086
|
|
exec cargo run --release -p ml --example "$EXAMPLE_NAME" --features cuda -- $CARGO_EXTRA
|
|
fi
|