Poziomka pretraining env setup

gistfile1.txt

# update apt and install basic tools
apt update
apt -y install zip tmux nano
apt -y install build-essential
apt -y install python-is-python3


# install CUDA 12.9
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.9.0/local_installers/cuda-repo-ubuntu2204-12-9-local_12.9.0-575.51.03-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-12-9-local_12.9.0-575.51.03-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-12-9-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-9

# Set CUDA environment variables
export CUDA_HOME=/usr/local/cuda-12.9
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH


# Add these lines to ~/.bashrc for persistence
echo 'export CUDA_HOME=/usr/local/cuda-12.9' >> ~/.bashrc
echo 'export PATH=$CUDA_HOME/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc

# we want tmux to be active so that we can disconnect ssh without stopping the training
tmux new -s megatron

pip install huggingface_hub
pip install hf_transfer
export HF_HUB_ENABLE_HF_TRANSFER=1
export PIP_BREAK_SYSTEM_PACKAGES=1

# I use wandb for experiment tracking

pip install wandb

cd /home/ubuntu
mkdir training
cd training

git clone https://github.com/adamo1139/Ling-V2.git
cd Ling-V2

pip install torch==2.8.0 torchvision --index-url https://download.pytorch.org/whl/cu129
pip install datasets pyarrow bitsandbytes ninja pandas regex safetensors tokenizers tiktoken transformers pybind11 psutil psutils
python -m pip install --upgrade pip wheel setuptools
pip install packaging==24.2 "setuptools>=77.0"

sudo apt install libnccl-dev

pip install megatron-core==0.13.0 torch==2.8.0
sudo apt-get -y install cudnn9-cuda-12
pip install --no-build-isolation transformer-engine[pytorch]

cd Ling-V2
# apply megatron patch
bash training/megatron/apply_patch.sh
# apply te patch
bash training/te/apply_te_patch.sh

bash deploy_parquet_support.sh

#apex

git clone https://github.com/NVIDIA/apex.git
cd apex
NVCC_APPEND_FLAGS="--threads 4" APEX_PARALLEL_BUILD=8 APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .
cd ..

# TE compilation can fail without nvidia-mathdx
pip install nvidia-mathdx

apt install libnccl2 libnccl-dev

NVTE_FRAMEWORK=pytorch pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable


#flash attn 3

git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/
git checkout 27f501d
cd hopper/
python setup.py install
python_path=`python -c "import site; print(site.getsitepackages()[0])"`
mkdir -p $python_path/flash_attn_3
wget -P $python_path/flash_attn_3 https://raw.githubusercontent.com/Dao-AILab/flash-attention/27f501dbe011f4371bff938fe7e09311ab3002fa/hopper/flash_attn_interface.py

# run training

bash examples/pretrain/run_pretrain_poziomka_X.sh