A minimal tool to compile python modules for AMD ROCm GPUs. Inspired by https://github.com/scooter-lacroix/Stan-s-ML-Stack/ and I have borrowed a couple things from there, so thanks!
You can compile for ROCm 6.4.3 or 7.0
Compiles the following wheels
- pytorch
- torchvision
- torchaudio
- tensorflow
- onnxruntime
- migraphx
- triton
- megatron
- flash attention
- bitsandbytes
- vllm
- rocm_smi (python)
- Weights and Biases
Warnings: pytorch, onnxruntime and tensorflow have very long compile times. You may be better off installing the tensorflow_rocm package unless you want a fresh build.
Compiles happen inside the official ROCm docker container based on Ubuntu 24.04 with wheels being saved to a volume for the user to install manually.
Q: Why not just install eveything into the base system?
A: I can think of 6 different ways to manage python packages (system, virtualenv, pipenv, poetry, uv,
conda) then add in distribution package managers on top of that. Plus I prefer a little but of control.
So creating wheels and then letting the user decide on how best to use them seens the best option.
Helpful little features:
- Where possible all packages are built off the latest release tag. This should improve stability as installing from the latest git commit may cause problems due to regressions etc. Unlikely but you never know.
- Compiles use all available CPU cores minus one. Should keep your system responsive so you can recover from any potential issues where things get bogged down
Verification scripts