foxhunt/scripts/fix_mamba2_lr.sh

#!/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 <pod-id>
#

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<PathBuf>,"
            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 <pod-id>"
            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 <pod-id>"
            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 <pod-id>    - Deploy to Runpod and restart training"
        echo -e "  $0 check <pod-id>     - 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