# DPO Training Scripts

This directory contains scripts for running Direct Preference Optimization (DPO) / Iterative Preference Optimization (IPO) training on SLURM clusters.

## 📋 Overview

The scripts are designed to be **user-agnostic** and work for any user account. They automatically detect:
- Training path (relative to script location)
- User-specific temporary directories using `$USER`
- Python environment (from your activated conda/venv)
- Neon workspace root (automatically finds `neon_sweep`, `neon_2`, etc.)

## 🗂️ Directory Structure

```
dpo/
├── run_ipo_dodo_crow_r1_t2.sh   # Round 1 training script
├── run_ipo_dodo_crow_r2_t2.sh   # Round 2 training script
├── run_ipo_dodo_crow_r3_t2.sh   # Round 3 training script
├── run_ipo_dodo_crow_r4_t2.sh   # Round 4 training script
├── run_ipo_dodo_crow_r5.sh      # Round 5 training script
├── run_ipo_dodo_crow_r6.sh      # Round 6 training script
├── run_ipo_dodo_crow_r7.sh      # Round 7 training script
├── chain_config.yaml            # Configuration for chained job submission
├── run_chained.py               # Python script for automated chain submission
└── README.md                    # This file
```

## 🚀 Quick Start

### Prerequisites

1. **Activate your Python environment** before submitting jobs:
   ```bash
   # Using conda
   conda activate gpt_n
   
   # Using venv
   source /path/to/venv/bin/activate
   ```

2. **Update `chain_config.yaml`** with your paths (optional - smart defaults will work):
   ```yaml
   settings:
     base_dir: "/home/YOUR_USERNAME/Work/neon_sweep/sunoGPT/slurm_scripts/dpo"
     log_dir: "/app/suno/slurm/logs"  # or your preferred log directory
   ```
   
   **Note**: If you don't specify paths, the script will automatically:
   - Use the script's directory for `base_dir`
   - Detect your `neon*` workspace root (e.g., `neon_sweep`, `neon_2`)
   - Use standard log directory

### 🆕 Starting a New Training Run

**IMPORTANT**: Before starting a new DPO chain, you must:

1. **Update the starting checkpoint in `run_ipo_dodo_crow_r1_t2.sh`**:
   ```bash
   # Edit the first round script
   vim run_ipo_dodo_crow_r1_t2.sh
   
   # Update this line with your new starting checkpoint:
   --preload_checkpoint="/app2/suno/checkpoints/YOUR_NEW_CHECKPOINT/last_ckpt_infer.pt"
   ```

2. **Rename all runs in `chain_config.yaml`** to avoid conflicts:
   ```yaml
   jobs:
     - name: "round_1"
       script: "run_ipo_dodo_crow_r1_t2.sh"
       nodes: 16
       model_cache_loss_name: "my_new_run_v51_n16_r1"  # ← Update this
       wandb_run_name: "my_new_run_v51_n16_r1"         # ← Update this
       sft_loss_scale: 0.01
     
     - name: "round_2"
       script: "run_ipo_dodo_crow_r2_t2.sh"
       nodes: 16
       model_cache_loss_name: "my_new_run_v51_n16_r1"  # References previous round
       wandb_run_name: "my_new_run_v52_n16_r2"         # ← Update this
       sft_loss_scale: 0.008
     # ... update all rounds
   ```

3. **Verify data directories** in each script match your intended dataset:
   ```bash
   --data_dir="/app2/suno/data/dpo/YOUR_DATASET"
   ```

### Option 1: Submit Individual Job

```bash
# Navigate to the dpo directory
cd /home/YOUR_USERNAME/Work/neon_sweep/sunoGPT/slurm_scripts/dpo

# Create an sbatch wrapper and submit
sbatch <<EOF
#!/bin/bash
#SBATCH --job-name=dpo
#SBATCH --nodes=16
#SBATCH --ntasks-per-node=8
#SBATCH --cpus-per-task=4
#SBATCH --gres=gpu:8
#SBATCH --output=/app/suno/slurm/logs/run_%x_%j.txt
#SBATCH --error=/app/suno/slurm/logs/run_%x_%j_err.txt
#SBATCH --time=48:00:00

srun -K1 ./run_ipo_dodo_crow_r1_t2.sh
EOF
```

### Option 2: Use the Chain Submission Script

The `run_chained.py` script automates multi-round training by:
- Submitting jobs sequentially
- Extracting checkpoint paths from completed jobs
- Automatically updating the next script with the new checkpoint

#### Basic Usage

```bash
# Using configuration file (recommended)
python run_chained.py --config chain_config.yaml

# Dry run to see what would be executed
python run_chained.py --config chain_config.yaml --dry-run

# Resume from last successful job
python run_chained.py --config chain_config.yaml --resume

# Submit as a single mega-job (all rounds in one SLURM job)
python run_chained.py --config chain_config.yaml --mega-job
```

#### Advanced Options

```bash
# Override sft_loss_scale for all rounds
python run_chained.py --config chain_config.yaml --sft_loss_scale 0.01

# Submit to specific partition
python run_chained.py --config chain_config.yaml --mega-job --partition gpu_partition

# Restore scripts from backups
python run_chained.py --config chain_config.yaml --restore

# Enable debug logging
python run_chained.py --config chain_config.yaml --log-level DEBUG
```

## ⚙️ Configuration

### chain_config.yaml

The configuration file supports smart defaults - you can omit the paths and it will auto-detect your workspace:

```yaml
chain_name: "dodo DPO Chain"
description: "Multi-round DPO training pipeline"

settings:
  # Optional - will auto-detect if not specified
  base_dir: "/home/YOUR_USERNAME/Work/neon_sweep/sunoGPT/slurm_scripts/dpo"
  log_dir: "/app/suno/slurm/logs"
  
  slurm_defaults:
    time: "48:00:00"
    ntasks_per_node: 8
    cpus_per_task: 4
    gpus_per_node: 8

jobs:
  - name: "round_1"
    script: "run_ipo_dodo_crow_r1_t2.sh"
    nodes: 16
    model_cache_loss_name: "dodo_t23_2025-08-28_01-30-09"
    wandb_run_name: "dodo_0828_crow_redo_v51_sft01_n16_r1"
    sft_loss_scale: 0.01
  
  - name: "round_2"
    script: "run_ipo_dodo_crow_r2_t2.sh"
    nodes: 16
    model_cache_loss_name: "dodo_0828_crow_redo_v51_sft01_n16_r1"
    wandb_run_name: "dodo_0828_crow_redo_v52_sft008_n16_r2"
    sft_loss_scale: 0.008
  
  # ... more rounds
```

### Key Configuration Parameters

- **base_dir**: Directory containing the DPO scripts
- **log_dir**: Directory for SLURM logs
- **nodes**: Number of nodes to allocate per job
- **model_cache_loss_name**: Identifier for model checkpointing
- **wandb_run_name**: Weights & Biases run name
- **sft_loss_scale**: SFT loss weight (can be different per round)

## 📝 Training Script Structure

Each training script (`run_ipo_dodo_crow_r*.sh`) follows this structure:

1. **Environment Setup**
   ```bash
   # NCCL and PyTorch distributed settings
   export NCCL_DEBUG=WARN
   export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
   
   # Auto-detect paths
   SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
   TRAIN_PATH="$(cd "$SCRIPT_DIR/../.." && pwd)"
   ```

2. **Temporary Directory Management**
   ```bash
   # User-specific temp directory
   rm -rf /mnt/localdisk/tmp_$USER
   mkdir -p /mnt/localdisk/tmp_$USER
   chmod -R 777 /mnt/localdisk/tmp_$USER
   ```

3. **Training Command**
   ```bash
   python -u train_dpo.py \
       --master_addr=$MASTER_ADDR \
       --master_port=$MASTER_PORT \
       --out_dir="/app2/suno/checkpoints" \
       --data_dir="/app2/suno/data/dpo/crow_t1_v51" \
       # ... more arguments
   ```

4. **Cleanup**
   ```bash
   rm -rf /mnt/localdisk/tmp_$USER
   ```

## 🔧 Customization

### Modifying Training Parameters

Edit the training scripts directly to adjust:
- Learning rates (`--learning_rate`, `--min_lr`)
- Batch sizes (`--batch_size`, `--eval_loss_batch_size`)
- Training iterations (`--max_iters`, `--warmup_iters`)
- DPO parameters (`--dpo_beta`, `--sft_loss_scale`)
- Model architecture (`--n_layer`, `--n_head`, etc.)

### Adding New Rounds

1. Copy an existing script:
   ```bash
   cp run_ipo_dodo_crow_r7.sh run_ipo_dodo_crow_r8.sh
   ```

2. Update parameters in the new script:
   - Data directory (`--data_dir`)
   - Checkpoint path (`--preload_checkpoint`)
   - Run names (`--wandb_run_name`, `--model_cache_loss_name`)

3. Add to `chain_config.yaml`:
   ```yaml
   - name: "round_8"
     script: "run_ipo_dodo_crow_r8.sh"
     nodes: 16
     model_cache_loss_name: "dodo_0828_crow_redo_v58_n16_r8"
     wandb_run_name: "dodo_0828_crow_redo_v58_n16_r8"
   ```

## ✅ Pre-Flight Checklist

Before submitting a new DPO chain, verify:

- [ ] Updated starting checkpoint in `run_ipo_dodo_crow_r1_t2.sh`
- [ ] Renamed all `model_cache_loss_name` in `chain_config.yaml` (unique identifiers)
- [ ] Renamed all `wandb_run_name` in `chain_config.yaml` (for tracking)
- [ ] Verified data directories in scripts (`--data_dir`)
- [ ] Activated correct Python environment
- [ ] Tested with `--dry-run` first

```bash
# Quick verification
python run_chained.py --config chain_config.yaml --dry-run

# If everything looks good, submit
python run_chained.py --config chain_config.yaml
```

## 🐛 Troubleshooting

### Issue: "python: command not found"

**Solution**: Activate your conda/venv environment before submitting:
```bash
conda activate gpt_n  # or your environment name
```

### Issue: "Permission denied" on temp directory

**Solution**: The scripts automatically create and set permissions. If issues persist:
```bash
# Manual cleanup
rm -rf /mnt/localdisk/tmp_$USER
mkdir -p /mnt/localdisk/tmp_$USER
chmod -R 777 /mnt/localdisk/tmp_$USER
```

### Issue: Checkpoint path not found

**Solution**: Verify checkpoint exists or update path in the script:
```bash
# Check if checkpoint exists
ls -lh /app2/suno/checkpoints/2025-11-09_16-16-35/last_ckpt_infer.pt
```

### Issue: SLURM job fails immediately

**Solution**: Check logs and verify paths:
```bash
# View SLURM logs
tail -f /app/suno/slurm/logs/run_dpo_JOBID.txt
tail -f /app/suno/slurm/logs/run_dpo_JOBID_err.txt

# Verify training path
cd /home/YOUR_USERNAME/Work/neon_sweep/sunoGPT
ls -la train_dpo.py
```

## 📊 Monitoring

### Check Job Status

```bash
# View all your jobs
squeue -u $USER

# View specific job
squeue -j JOBID

# View detailed job info
scontrol show job JOBID
```

### Monitor Training Progress

```bash
# Follow SLURM logs
tail -f /app/suno/slurm/logs/run_dpo_JOBID.txt

# Check for checkpoint creation
watch -n 60 'ls -lth /app2/suno/checkpoints/ | head -20'

# Monitor with Weights & Biases
# Visit: https://wandb.ai/your-entity/chirp-dodo-dpo
```

### Cancel Jobs

```bash
# Cancel specific job
scancel JOBID

# Cancel all your jobs
scancel -u $USER
```

## 🔒 Best Practices

1. **Always activate your Python environment** before submitting jobs
2. **Update the starting checkpoint** in round 1 script before each new run
3. **Rename all runs** in config.yaml to avoid checkpoint/W&B naming conflicts
4. **Update configuration files** rather than hardcoding paths
5. **Use chain submission** for multi-round training to ensure consistency
6. **Monitor disk usage** - temp files and checkpoints can be large
7. **Back up important checkpoints** before starting new training
8. **Use meaningful run names** for easier tracking in W&B
9. **Test with `--dry-run`** before submitting actual jobs
10. **Verify data directories** match your intended dataset version

## 📚 Additional Resources

- **Training Code**: `../../train_dpo.py`
- **SLURM Documentation**: [SLURM Official Docs](https://slurm.schedmd.com/)
- **Weights & Biases**: [W&B Documentation](https://docs.wandb.ai/)

## 🤝 Contributing

When modifying these scripts:
1. Keep them user-agnostic (use `$USER`, relative paths)
2. Update this README with any significant changes
3. Test with `--dry-run` before committing
4. Document new parameters or features

## 📧 Support

For issues or questions:
- Check logs in `/app/suno/slurm/logs/`
- Review SLURM job output
- Consult team members or documentation

---

**Last Updated**: 2025-11-10
**Maintainer**: Team (user-agnostic version)