Bash script to setup axolotl+FA2+BnB+liger-kernel on Runpod MI300X

axolotl_ROCm_setup_v2.sh

#!/bin/bash

# Setup Axolotl with FA2 and BnB ROCm - doctorshotgun Sept 16, 2024
# Runpod image: RunPod Pytorch 2.1.2 ROCm 6.1 runpod/pytorch:2.1.2-py3.10-rocm6.1-ubuntu22.04

# Install torch and flash-attn
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/rocm6.1
#pip install https://download.pytorch.org/whl/nightly/pytorch_triton_rocm-3.0.0%2Bdafe145982-cp310-cp310-linux_x86_64.whl
pip install https://github.com/DocShotgun/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+rocm6.1+torch2.4.0-cp310-cp310-linux_x86_64.whl

# For some reason we need to manually install amdsmi for torch 2.4.0 with ROCm 6.1
cd /opt/rocm/share/amd_smi && pip install .

# Install Axolotl
cd /workspace/
git clone https://github.com/DocShotgun/axolotl && cd axolotl
pip install -e '.[deepspeed]'

# Install Bitsandbytes (multi-backend-refactor branch)
cd /workspace/
git clone https://github.com/bitsandbytes-foundation/bitsandbytes && cd bitsandbytes/
git checkout 2bfa3472ecde8f3e4a0306b017826314c288b7c8
pip install -r requirements-dev.txt
cmake -DCOMPUTE_BACKEND=hip -S .
make
pip install -e .

# Install Liger-Kernel and compatible Triton
pip uninstall liger-kernel -y
pip install liger-kernel@git+https://github.com/DocShotgun/Liger-Kernel

# To begin training, run:
# accelerate launch -m axolotl.cli.train <your_config.yml>
# If you encounter an error related to xformers, you can try editing /src/axolotl/monkeypatch/llama_attn_hijack_flash.py (for llama-type models) to comment out the xformers import and force is_xformers_swiglu_available to return False