rocm-ml

This repository builds a docker image to leverage AMD graphics cards for machine learning using the ROCm drivers. These instructions are designed for Linux. I am using Pop OS! and installed the ROCm drivers with these instructions.

Usage

To use a gated model such as stable diffusion 3, you must define the environmental variable huggingface_api_key. You can do this by adding the following line to your ~/.bashrc file:

export huggingface_api_key=Put_API_Key_Here_With_No_Quotes_Or_Spaces

To build the docker image:

docker build . -t rocm-ml

To run:

sudo -E docker run --env huggingface_api_key -it -p 8888:8888  -v ./notebooks:/notebooks -w /notebooks --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
--device=/dev/kfd --device=/dev/dri --group-add video \
--ipc=host --shm-size 8G rocm-ml 

• This command mounts the notebooks folder within the docker container. I recommend only editing within a jupyter notebook running on the docker container.
• The -E argument for sudo is key to transfer the environmental variable huggingface_api_key from the user's environment to the root user

Command to start jupyter within docker container:

jupyter notebook --ip 0.0.0.0 --no-browser --allow-root --NotebookApp.token='' --NotebookApp.password=''

References

• AMD Instructions for using ROCm via Docker - the docker run command above is based on this page
• Docker run
  • Bind volumes arg (-v)
• Running Jupyter with Docker - The p argument in the docker run command above is from here

Todo

• Try out different text to image models
• Make jupyter server start automatically
• Makefile or docker compose for build process?
• Fix issue with jupyther authentication