amd-rocm

AMD ROCm Inference Backend

This folder is the evidence of the AMD side of Gradelytics. The web app itself is provider-agnostic — by default it runs on Fireworks serverless inference — but I also wanted to prove I could run our own Gemma model on AMD hardware and have the exact same frontend talk to it. That's what lives here.

Because I kept the server's HTTP surface OpenAI-compatible, pointing the app at this backend is just a base-URL swap: no client code changes.

What's in here

File	What it is
serve_gemma.py	A small FastAPI server that loads Gemma on the AMD GPU (ROCm) and exposes POST /v1/chat/completions in OpenAI's response shape.
gemma_rocm.ipynb	The notebook I ran on the AMD box — GPU/ROCm checks, model load, a sample generation, and launching the server behind a public tunnel.
requirements.txt	Python deps. torch must be the ROCm build, not the CUDA one.
Dockerfile	Reproducible image on rocm/pytorch with the GPU passed through.

Why ROCm (and the one gotcha)

Everything runs on AMD via ROCm/HIP. The thing that tripped me up first: on an AMD card, PyTorch still uses the cuda device names — that's just how the ROCm backend is shipped. So torch.cuda.is_available() returning True here means the AMD GPU is up, and I confirm it properly by printing torch.version.hip and rocm-smi.

The failure mode to avoid is installing the default torch wheel, which is the NVIDIA/CUDA build and errors with "Found no NVIDIA driver on your system" on an AMD box. The fix is to install torch from AMD's ROCm wheel index:

pip install --index-url https://download.pytorch.org/whl/rocm6.2 torch
pip install -r requirements.txt

Run it

# 1. Make sure the AMD GPU is visible
rocm-smi

# 2. Start the OpenAI-compatible server (needs HF access to Gemma)
huggingface-cli login
MODEL_ID=google/gemma-2-2b-it python serve_gemma.py

# 3. In another shell, expose it publicly (no install needed — just SSH)
ssh -R 80:localhost:8000 serveo.net
# → prints a public https://<id>.serveousercontent.com URL

Or containerized, with the GPU passed through:

docker build -t gradelytics-amd .
docker run --rm -it \
  --device=/dev/kfd --device=/dev/dri \
  --group-add video --security-opt seccomp=unconfined \
  -p 8000:8000 -e MODEL_ID=google/gemma-2-2b-it \
  gradelytics-amd

Verify

curl -X POST https://<your-tunnel>.serveousercontent.com/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"model":"gemma","messages":[{"role":"user","content":"hello"}]}'

A choices[0].message.content field in the JSON response means Gemma answered from the AMD GPU. To drive the whole app through it, I point Gradelytics at https://<your-tunnel>.serveousercontent.com/v1 and every AI call runs on AMD.