File size: 4,390 Bytes
6269460 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | #!/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"
|