#!/bin/bash
# ============================================================================
# Environment Setup Script
# ============================================================================
# Run this ONCE on your cluster to set up all dependencies.
# Creates two conda environments:
#   1. megatron — for pre-training (Megatron-LM + MoE dependencies)
#   2. sft — for SFT/RL with TRL and evaluation
# ============================================================================

set -euo pipefail

echo "============================================================"
echo "MoE Training Pipeline — Environment Setup"
echo "============================================================"

# ============================================================================
# 1. Megatron-LM Environment (Pre-training)
# ============================================================================
echo ""
echo "--- Setting up Megatron-LM environment ---"

conda create -n megatron python=3.11 -y
conda activate megatron

# PyTorch (match your CUDA version)
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124

# Megatron-LM (clone and install)
cd /path/to/software  # adjust this path
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
pip install -e .

# Megatron-Core (if not bundled)
pip install megatron-core

# MoE dependencies
pip install megablocks                  # Block-sparse MoE kernels
pip install grouped-gemm                # Grouped GEMM for expert parallelism

# Flash Attention 2 (critical for performance)
pip install flash-attn --no-build-isolation

# Apex (NVIDIA mixed precision) — build from source for best compatibility
cd /path/to/software
git clone https://github.com/NVIDIA/apex.git
cd apex
pip install -v --disable-pip-version-check --no-cache-dir \
    --no-build-isolation \
    --config-settings "--build-option=--cpp_ext" \
    --config-settings "--build-option=--cuda_ext" .

# TransformerEngine (FP8 training on H100)
pip install transformer-engine[pytorch]

# Data processing
pip install transformers datasets tokenizers
pip install sentencepiece tiktoken  # tokenizer backends

# Monitoring
pip install tensorboard wandb

# Networking (multi-node)
pip install packaging ninja

echo "Megatron environment ready!"
conda deactivate

# ============================================================================
# 2. SFT/RL Environment (Fine-tuning + Evaluation)
# ============================================================================
echo ""
echo "--- Setting up SFT/RL environment ---"

conda create -n sft python=3.11 -y
conda activate sft

pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124
pip install flash-attn --no-build-isolation

# HuggingFace stack
pip install transformers>=4.55.0
pip install trl>=0.17.0
pip install datasets
pip install accelerate
pip install peft
pip install bitsandbytes  # for QLoRA

# Evaluation
pip install evalplus        # HumanEval/MBPP
pip install lighteval       # General benchmarks
pip install vllm            # Fast inference for evaluation

# Data curation
pip install openai anthropic  # API access for synthetic data generation

# Monitoring
pip install tensorboard wandb trackio

echo "SFT/RL environment ready!"
conda deactivate

# ============================================================================
# 3. Data Curation Environment (CPU-only)
# ============================================================================
echo ""
echo "--- Setting up data curation environment ---"

conda create -n datacuration python=3.11 -y
conda activate datacuration

pip install datasets transformers tokenizers
pip install openai anthropic  # for synthetic data generation
pip install tiktoken sentencepiece

echo "Data curation environment ready!"
conda deactivate

echo ""
echo "============================================================"
echo "Setup Complete!"
echo "============================================================"
echo ""
echo "Environments:"
echo "  conda activate megatron    — for pre-training"
echo "  conda activate sft         — for SFT, RL, evaluation"
echo "  conda activate datacuration — for data processing"
echo ""
echo "Next steps:"
echo "  1. Edit paths in scripts/pretrain_megatron.sh"
echo "  2. Run data curation pipeline"
echo "  3. Submit pre-training job: sbatch slurm/pretrain.sbatch"