#!/bin/bash

# Setup Axolotl with FA2 and BnB ROCm - doctorshotgun Aug 6, 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://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/axolotl-ai-cloud/axolotl && cd axolotl
git checkout 70978467a088da3abf3fe45d92d90f6529f19ea9
pip install -e '.[deepspeed]'

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

# 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