#!/bin/bash # # MAMBA-2 Learning Rate Fix: Restore E10 + Reduce LR to 3e-5 # # This script implements the recommendation from MAMBA2_LR_ANALYSIS_E10_E14.md # # Usage: # 1. Run locally to add --resume-from flag: ./scripts/fix_mamba2_lr.sh local # 2. Deploy to Runpod: ./scripts/fix_mamba2_lr.sh deploy # set -e MODE="${1:-help}" POD_ID="${2:-}" # Colors for output RED='\033[0;31m' GREEN='\033[0;32m' YELLOW='\033[1;33m' BLUE='\033[0;34m' NC='\033[0m' # No Color echo -e "${BLUE}╔═══════════════════════════════════════════════════════════╗${NC}" echo -e "${BLUE}║ MAMBA-2 Learning Rate Fix: E10 Restore + LR=3e-5 ║${NC}" echo -e "${BLUE}╚═══════════════════════════════════════════════════════════╝${NC}" case "$MODE" in local) echo -e "${YELLOW}Step 1: Adding --resume-from flag to train_mamba2_parquet.rs${NC}" # Check if flag already exists if grep -q "resume-from\|resume_from" ml/examples/train_mamba2_parquet.rs; then echo -e "${GREEN}✓ --resume-from flag already exists${NC}" else echo -e "${YELLOW}⚠ --resume-from flag not found. Manual addition required.${NC}" echo -e "${YELLOW}Add to TrainingConfig struct (line ~100):${NC}" echo -e " pub resume_from: Option," echo -e "" echo -e "${YELLOW}Add to Default impl (line ~120):${NC}" echo -e " resume_from: None," echo -e "" echo -e "${YELLOW}Add to arg parsing (line ~430):${NC}" echo ' "--resume-from" if i + 1 < args.len() => {' echo ' config.resume_from = Some(PathBuf::from(&args[i + 1]));' echo ' info!("Resume from checkpoint: {:?}", config.resume_from);' echo ' }' echo -e "" echo -e "${YELLOW}Add before training loop (line ~620):${NC}" echo ' if let Some(resume_path) = &config.resume_from {' echo ' model.load_checkpoint(resume_path.to_str().unwrap())' echo ' .await' echo ' .context("Failed to load resume checkpoint")?;' echo ' info!("✓ Resumed from checkpoint: {:?}", resume_path);' echo ' }' echo -e "" echo -e "${RED}Press Enter after making these changes...${NC}" read -r fi echo -e "${YELLOW}Step 2: Building updated binary${NC}" cargo build -p ml --example train_mamba2_parquet --release --features cuda if [ $? -eq 0 ]; then echo -e "${GREEN}✓ Build successful${NC}" echo -e "${YELLOW}Step 3: Testing checkpoint loading (dry-run)${NC}" # Note: This will fail if no checkpoint exists, which is expected locally ./target/release/examples/train_mamba2_parquet \ --epochs 1 \ --resume-from ml/checkpoints/mamba2_parquet/best_epoch_10.ckpt \ 2>&1 | grep -i "resume\|checkpoint" || true echo -e "" echo -e "${GREEN}✓ Local setup complete${NC}" echo -e "${YELLOW}Next step: Upload binary to Runpod${NC}" echo -e " scp target/release/examples/train_mamba2_parquet runpod:/runpod-volume/binaries/" else echo -e "${RED}✗ Build failed. Fix errors above.${NC}" exit 1 fi ;; deploy) if [ -z "$POD_ID" ]; then echo -e "${RED}✗ Error: Pod ID required${NC}" echo -e "Usage: $0 deploy " exit 1 fi echo -e "${YELLOW}Deploying to Runpod pod: $POD_ID${NC}" echo -e "${YELLOW}Step 1: Stopping current training${NC}" runpodctl exec "$POD_ID" "pkill -f train_mamba2_parquet || true" sleep 2 echo -e "${YELLOW}Step 2: Verifying E10 checkpoint exists${NC}" CHECKPOINT_PATH="/runpod-volume/models/mamba2_parquet/best_epoch_10.ckpt" if runpodctl exec "$POD_ID" "test -f $CHECKPOINT_PATH && ls -lh $CHECKPOINT_PATH"; then echo -e "${GREEN}✓ E10 checkpoint found${NC}" else echo -e "${RED}✗ E10 checkpoint not found at $CHECKPOINT_PATH${NC}" echo -e "${YELLOW}Available checkpoints:${NC}" runpodctl exec "$POD_ID" "ls -lh /runpod-volume/models/mamba2_parquet/" exit 1 fi echo -e "${YELLOW}Step 3: Uploading updated binary${NC}" # Note: Adjust path if different scp target/release/examples/train_mamba2_parquet runpod:/runpod-volume/binaries/ echo -e "${YELLOW}Step 4: Starting training with LR=3e-5${NC}" TRAINING_CMD="/runpod-volume/binaries/train_mamba2_parquet \ --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet \ --epochs 50 \ --learning-rate 3e-5 \ --checkpoint-dir /runpod-volume/models/mamba2_parquet \ --resume-from $CHECKPOINT_PATH" echo -e "${BLUE}Training command:${NC}" echo -e "$TRAINING_CMD" echo -e "" # Start training in background runpodctl exec "$POD_ID" "nohup $TRAINING_CMD > /runpod-volume/logs/train_mamba2_lr3e5.log 2>&1 &" echo -e "" echo -e "${GREEN}✓ Training started successfully${NC}" echo -e "${YELLOW}Monitor progress:${NC}" echo -e " runpodctl logs $POD_ID -f" echo -e " tail -f /runpod-volume/logs/train_mamba2_lr3e5.log" echo -e "" echo -e "${YELLOW}Success Criteria (check at E15-20):${NC}" echo -e " ✓ Val loss < 43.9M (better than E10 baseline)" echo -e " ✓ No spikes > 2% in validation loss" echo -e " ✓ Training loss decreasing" ;; check) if [ -z "$POD_ID" ]; then echo -e "${RED}✗ Error: Pod ID required${NC}" echo -e "Usage: $0 check " exit 1 fi echo -e "${YELLOW}Checking training progress on pod: $POD_ID${NC}" # Get last 20 epochs from log echo -e "${YELLOW}Last 10 epochs:${NC}" runpodctl exec "$POD_ID" "grep -E 'Epoch [0-9]+/50:' /runpod-volume/logs/train_mamba2_lr3e5.log | tail -10" echo -e "" echo -e "${YELLOW}Validation loss trend:${NC}" runpodctl exec "$POD_ID" "grep -oP 'Val Loss = \K[0-9.]+' /runpod-volume/logs/train_mamba2_lr3e5.log | tail -10 | nl" echo -e "" echo -e "${YELLOW}Latest checkpoint:${NC}" runpodctl exec "$POD_ID" "ls -lht /runpod-volume/models/mamba2_parquet/*.ckpt | head -5" ;; help|*) echo -e "${YELLOW}Usage:${NC}" echo -e " $0 local - Add --resume-from flag and build locally" echo -e " $0 deploy - Deploy to Runpod and restart training" echo -e " $0 check - Check training progress" echo -e "" echo -e "${YELLOW}Example workflow:${NC}" echo -e " 1. ./scripts/fix_mamba2_lr.sh local" echo -e " 2. ./scripts/fix_mamba2_lr.sh deploy abc123xyz" echo -e " 3. ./scripts/fix_mamba2_lr.sh check abc123xyz" echo -e "" echo -e "${YELLOW}See MAMBA2_LR_ANALYSIS_E10_E14.md for full analysis${NC}" ;; esac