ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

docker: fix bnb GEMV floor and install the HF training stack in the ROCm image

Browse source 
Two review items, both reproduced before fixing.

The pinned bitsandbytes commit (0.48.0.dev0, 2025-09-15) predates bnb PR
1887, which enabled the 4-bit GEMV kernels on AMD CDNA; every build at or
below 0.49.2 NaNs at the GEMV decode shape on AMD GPUs, the exact floor
install.sh already enforces for its ROCm wheel. BNB_COMMIT now defaults
to the PR 1887 merge commit (0.50.0.dev0), which carries the fix.

The unsloth install used bare pip install -e ., which pulls only the CLI
dependencies (typer, rich, pydantic, pyyaml, nest-asyncio); the pip show
probe passed on metadata alone while transformers, trl, peft, datasets,
and accelerate were all absent, so the image built green but could not
train. It now installs the huggingfacenotorch extra, which delivers the
full HF training stack without torch or torchvision so the ROCm torch in
the vLLM base image is never clobbered by a PyPI CPU/CUDA build.
This commit is contained in:
Daniel Han 2026-07-18 13:09:44 +00:00
commit 66f912625d
1 changed files with 9 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

11
docker/Dockerfile.rocm
View file

@ -11,7 +11,10 @@ RUN apt-get update \
ARG COMPUTE_BACKEND="hip"
ARG BNB_ROCM_ARCHS="gfx942"
ARG BNB_COMMIT="4b0257482bef447106fcaada67d1c6d081fdc82f"
# Must be at or after bnb PR #1887 (merged as this commit, 0.50.0.dev0): earlier
# builds (<= 0.49.2) NaN at the 4-bit GEMV decode shape on every AMD GPU, the
# same floor install.sh enforces for its ROCm wheel.
ARG BNB_COMMIT="713a3b83e405d32449e1cb392e5e25985ab927c6"
ARG BNB_REPO="https://github.com/bitsandbytes-foundation/bitsandbytes.git"
# Fetch only the pinned commit via init + shallow fetch: minimal, reproducible.
@ -33,10 +36,14 @@ ARG UNSLOTH_REPO="https://github.com/unslothai/unsloth.git"
ARG UNSLOTH_BRANCH="main"
# huggingfacenotorch pulls the full HF training stack (transformers, trl, peft,
# datasets, accelerate) without torch/torchvision, so the base image's ROCm
# torch is never clobbered by a PyPI CPU/CUDA build. Bare `-e .` installs only
# the CLI deps and the image cannot actually train.
RUN git clone --branch ${UNSLOTH_BRANCH} --depth 1 ${UNSLOTH_REPO} unsloth \
&& cd unsloth \
&& rm -rf .git \
&& pip install -e .
&& pip install -e '.[huggingfacenotorch]'
ARG UNSLOTH_ZOO_REPO="https://github.com/unslothai/unsloth-zoo.git"