Skip to content

README.md

Latest commit

 

History

History
72 lines (51 loc) · 2 KB

README.md

File metadata and controls

72 lines (51 loc) · 2 KB
Raw

Experiments for KernelBench

Environment Setup

Initialize Repository

After cloning the repository, update the submodule, and apply the patches.

git submodule update --init --recursive
git -C original apply ../patches/add-ollama-server.patch

Launch Docker Container

Since the latest stable version of PyTorch has not yet supported ROCm 7.0, we have to use a nightly version (v2.10) instead. There is a pre-built Docker image rocm/pytorch-nightly on Docker Hub.

cont_name=$(whoami)-kernelbench
mount_dir=$(dirname $(realpath $(pwd)))  # /path/to/ml_tuning

docker run -it \
    -h $cont_name \
    --name $cont_name \
    --device=/dev/kfd \
    --device=/dev/dri \
    --dns=165.204.219.249 \
    --security-opt seccomp=unconfined \
    -v $mount_dir:/mnt \
    -w /mnt \
    rocm/pytorch-nightly

Setup Ollama Server

Install Ollama in the Docker container.

curl -fsSL https://ollama.com/install.sh | sh

Serve Ollama in the background (e.g. tmux, nohup)

# with tmux
apt update && apt install -y tmux
tmux new -ds ollama 'ollama serve'

# or with nohup
nohup ollama serve &> /tmp/ollama.log &

Download model in another session, take CodeLlama-7B for example.

ollama pull codellama:7b

Install KernelBench

Since we are using the nightly build of torch pre-installed in the container, we need to carefully specify the versions of other packages (for example, numpy<1.23,>=1.18). Here we provide requirements.lock for that instead of using requirements.txt in the original repository.

pip install -r requirements.lock
pip install -e original

Run a single problem (e.g. matmul with transposed A) with local LLM (e.g. codellama-7b) on Ollama.

python3 original/scripts/generate_and_eval_single_sample.py dataset_src="huggingface" level=1 problem_id=16 server_type="ollama" model_name="codellama:7b"