feat(docker): add an AMD ROCm image, for RDNA2 through RDNA4 and CDNA - #10820
feat(docker): add an AMD ROCm image, for RDNA2 through RDNA4 and CDNA#10820LeoBorcherding wants to merge 3 commits into
Conversation
Reapplies PR unslothai#6231 on top of main. Dockerfile.rocm, Dockerfile.studio-rocm, entrypoint-rocm.sh, smoke_test_rocm.py, test_locally-rocm.sh and the publish workflow are unchanged from that branch. build.sh and run.sh were rewritten by unslothai#5748 after unslothai#6231 was opened, so the --rocm hooks are reimplemented against the current files instead of merged: run.sh now reuses main's numeric --group-add gid lookup, and skips the NVIDIA toolkit prompt on the ROCm path.
Two defects found by building the image: The torch layer used --index-strategy unsafe-best-match with pypi as a fallback, which compares versions across both indexes. Every ROCm index below 7.2 trails pypi's torch (rocm6.2 tops out at 2.5.1, rocm6.4 at 2.9.1), so the best match is pypi's CUDA build: the PR's own default of rocm6.2 resolves to a plain torch==2.14.0 pulling nvidia-cublas and nccl, and then fails the torch.version.hip assertion. Dropped the fallback and moved the --rocm default to 7.2.4, which is also the floor for RDNA4. unsloth-zoo and cut-cross-entropy both declare a bare triton on Linux, the NVIDIA-targeted pypi wheel. It installs over torch's triton-rocm in the same site-packages/triton/, leaving triton 3.8.0 by metadata and triton.__version__ 3.7.1 in one tree. Uninstall it and reinstall the pin torch names.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
for more information, see https://pre-commit.ci
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e7816f7dea
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ROCM_VERSION=${{ github.event.inputs.rocm_version || '6.2.4' }} | ||
| PYTHON_VERSION=3.12 | ||
| TORCH_INDEX_URL=${{ github.event.inputs.torch_index_url || 'https://download.pytorch.org/whl/rocm6.2' }} |
There was a problem hiding this comment.
The push and scheduled paths in docker-publish-rocm.yml always pass ROCm 6.2 defaults, even though Dockerfile.rocm states that gfx1150/1151/1200/1201 require ROCm 7.2 and build.sh --rocm consequently defaults to 7.2.4. Thus the image published as unsloth/unsloth-rocm:latest lacks kernels for the advertised RDNA4 and Strix devices; align these workflow defaults with the supported 7.2 build.
Useful? React with 👍 / 👎.
| ENV UNSLOTH_STUDIO_HOME=/opt/unsloth-studio \ | ||
| DEBIAN_FRONTEND=noninteractive \ | ||
| # torch is already installed in the base image (ROCm wheels). | ||
| # Prevent install.sh from re-running get_torch_index_url() which would | ||
| # probe the absent build-host GPU and silently install a CPU wheel on top. | ||
| SKIP_TORCH=true |
There was a problem hiding this comment.
On the documented GPU-free Studio build, this variable does not skip torch installation: install.sh unconditionally resets SKIP_TORCH=false and only honors --no-torch or UNSLOTH_NO_TORCH. It therefore probes an absent GPU and installs CPU torch into Studio's separate unsloth_studio venv; the base-image entrypoint still validates the ROCm torch in /opt/unsloth-venv, so startup can pass while Studio training uses the CPU build. The Studio install needs an explicit matching ROCm index/family rather than this ignored variable.
Useful? React with 👍 / 👎.
| ARG BASE_TAG=test | ||
| FROM unsloth-rocm:${BASE_TAG} |
There was a problem hiding this comment.
Following the new entry point bash docker/build.sh --rocm produces unsloth-rocm:latest, but a subsequent Studio build without an override looks for unsloth-rocm:test. On a normal checkout with no separately tagged :test image, Docker attempts to resolve that nonexistent base and the Studio build fails before running any instructions; use the same default tag as build.sh or make the Studio build path pass the produced tag.
Useful? React with 👍 / 👎.
| UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || github.sha || 'main' }} | ||
| UNSLOTH_ZOO_REF=${{ github.event.inputs.unsloth_zoo_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || 'main' }} |
There was a problem hiding this comment.
On every v* tag event this expression forces UNSLOTH_ZOO_REF to the Unsloth release tag, but Unsloth's release tags are not guaranteed to exist in unsloth-zoo. The existing CUDA publish workflow explicitly probes for the matching zoo tag and falls back to main; without that check, any release lacking a synchronized zoo tag makes the Git dependency installation fail and prevents the ROCm release image from being published.
Useful? React with 👍 / 👎.
| ROCM_VERSION="${ROCM_VERSION:-7.2.4}" | ||
| TORCH_INDEX_URL="${TORCH_INDEX_URL:-https://download.pytorch.org/whl/rocm7.2}" |
There was a problem hiding this comment.
For gfx1150/gfx1151, the advertised build.sh --rocm path installs the generic pytorch.org ROCm 7.2 wheel. The repository's own install.sh deliberately reroutes these architectures away from generic indexes below ROCm 7.13 to repo.amd.com/rocm/whl/<gfx>/, documenting a _grouped_mm segfault and subsequent Strix kernel bugs in the 7.2 build. Because this Dockerfile bypasses that installer routing, the default Strix image uses the known-broken wheel and can crash during supported training workloads.
Useful? React with 👍 / 👎.
| # SUPPORTED maps gfx arch string -> (family, example cards). | ||
| SUPPORTED = { | ||
| "gfx906": ("Vega20 / CDNA0", "Radeon VII"), | ||
| "gfx908": ("CDNA1", "Instinct MI100"), | ||
| "gfx90a": ("CDNA2", "Instinct MI200 / MI210 / MI250"), |
There was a problem hiding this comment.
When the image is built through the new default build.sh --rocm path, it uses ROCm 7.2, yet this table reports gfx906 as supported. The gfx906 routing in this repository's install.sh states that ROCm 6.3 is the last wheel family whose BLAS kernels run on gfx906 and that newer generic wheels fail at the first BLAS call. An MI50 or Radeon VII user therefore passes the entrypoint with a reassuring “Supported” message and then fails in training; gfx906 needs a separate 6.3 build or a version-aware rejection here.
Useful? React with 👍 / 👎.
| UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || github.sha || 'main' }} | ||
| UNSLOTH_ZOO_REF=${{ github.event.inputs.unsloth_zoo_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || 'main' }} |
There was a problem hiding this comment.
Scheduled builds pass unsloth-zoo@main, and default workflow-dispatch builds pass both repositories as @main, directly into a cached Docker RUN command. When the repository SHA is unchanged, BuildKit can reuse the previous install layer even after either upstream branch advances, so a green nightly or manually triggered image can silently contain stale Unsloth Zoo or Unsloth code. The local build.sh and existing CUDA workflow resolve these names to commit SHAs specifically to invalidate this cache correctly; the ROCm workflow needs the same preparation.
Useful? React with 👍 / 👎.
| images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} | ||
| tags: | | ||
| type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) && github.event.inputs.unsloth_ref == '' }} | ||
| type=ref,event=tag | ||
| type=schedule,pattern=nightly | ||
| type=sha,prefix=sha-,format=short |
There was a problem hiding this comment.
On a default-branch workflow_dispatch, clearing unsloth_ref enables the stable latest tag even if unsloth_zoo_ref, rocm_version, or torch_index_url is overridden. A test dispatch for a custom zoo commit, legacy ROCm version, or experimental wheel index can therefore replace the public unsloth/unsloth-rocm:latest image. Gate latest on every build-affecting input remaining at its stable default, as the existing CUDA publish workflow does.
Useful? React with 👍 / 👎.
| RUN apt-get update && apt-get install -y --no-install-recommends \ | ||
| python3 python3-pip python3-venv \ | ||
| curl wget git \ | ||
| cmake build-essential libcurl4-openssl-dev \ |
There was a problem hiding this comment.
The Dockerfile declares and the publish workflow passes PYTHON_VERSION=3.12, but this installation uses the unversioned Ubuntu 22.04 python3 packages and later creates the venv with that executable, so the resulting image actually uses Ubuntu's Python 3.10. Changing the build argument has no effect, and dependency markers or user code expecting the advertised 3.12 environment resolve differently; install python${PYTHON_VERSION} and create the venv from that interpreter.
Useful? React with 👍 / 👎.
| print("ERROR: torch.cuda.is_available() is False despite rocm-smi working.") | ||
| print() | ||
| print("This image was built against ROCm 6.2. The host ROCm stack must be") | ||
| print("6.2 or newer. Check the host (NOT the container) with:") | ||
| print(" rocm-smi --version") |
There was a problem hiding this comment.
For images produced by the new default build.sh --rocm path, torch and the base runtime are ROCm 7.2, but a failed availability probe tells users the image was built against 6.2 and that any 6.2-or-newer host is sufficient. This sends users with an incompatible 6.x host away from the required driver upgrade and makes the most important runtime failure diagnostic incorrect; derive the displayed version from torch.version.hip or the build argument.
Useful? React with 👍 / 👎.
|
Checked the --rocm hooks in docker/build.sh and docker/run.sh and the existing NVIDIA path stays gated behind the non-ROCm branch, which is the right shape. Since Dockerfile.studio-rocm has not been built and you flagged the SKIP_TORCH and BASE_TAG issues yourself, could you drop it to a follow-up so this lands as the base ROCm image only? |
|
Part of this is not yours. Two failures are from this PR's own new files, and both gates are telling you the fix:
Rebase once #10832 is in. |
2 participants
adds an AMD ROCm image alongside the CUDA one from #5748. same
docker/layout, samebuild.sh/run.shentry points, same smoke-test shape.this is #6231 resynced. that PR targeted
docker-blackwell-build, which was deleted when #5748 merged, so github auto-closed it and won't reopen it. the branch was also 2764 commits behind, andbuild.sh/run.shwere rewritten underneath it, so the--rocmhooks are reimplemented against the current files rather than merged.what's here
docker/Dockerfile.rocm: ROCm torch wheels plus the bitsandbytes pre-release that carries the 4-bit decode fix.[huggingface]extra, SDPA fallback, since xformers has no ROCm wheel.docker/Dockerfile.studio-rocm: Studio variant on the ROCm base.docker/entrypoint-rocm.sh: preflight for/dev/kfd, rocm-smi, HIP torch, gfx arch.docker/smoke_test_rocm.py: ROCm smoke test with a 5-step LoRA..github/workflows/docker-publish-rocm.yml: GPU-free amd64 build and publish.build.sh --rocmandrun.sh --rocm.run.shreuses main's numeric--group-addgid lookup instead of passing the group by name, and skips the NVIDIA toolkit prompt on the ROCm path.two things building it turned up
the old default installed a CUDA torch. the torch layer used
--index-strategy unsafe-best-matchwith pypi as a fallback, which compares versions across both indexes. every ROCm wheel index below 7.2 trails pypi, so pypi won:rocm6.4 does the same thing (
2.14.0vs2.9.1+rocm6.4). the build then failed its ownassert torch.version.hip is not None. dropped the pypi fallback, and moved the default to ROCm 7.2.4: 6.2 tops out at torch 2.5.1, which current unsloth-zoo and transformers won't sit on, and RDNA4 has no kernels before 7.x anyway.the image shipped two tritons in one directory. unsloth-zoo and cut-cross-entropy each declare a bare
tritonon linux, which is the NVIDIA-targeted pypi wheel, and it installs over torch'striton-rocmin the samesite-packages/triton/:neither version describes what's actually on disk. uninstalling the pypi one and reinstalling the pin torch names fixes it:
index selection can't fix this,
download.pytorch.org/whl/rocm7.2/triton/is the pypi wheel mirrored rather than a ROCm build. the durable fix is probably in unsloth-zoo's metadata, where that dependency should be conditional.hardware and OS coverage
no hosted github runner has an AMD GPU, and the self-hosted Strix Halo runner on ooba's fork is the wrong arch and OS for this, so a green check here says nothing about AMD behaviour. this table is the test report.
measured at 8ee07d6.
docker/is unchanged on main since that commit.HSA_OVERRIDE_GFX_VERSIONwhere it was run:
run.sh --rocmend to end, and 58 existing docker tests. it can prove nothing about a GPUwindows can't test the GPU half of this, and the reason is structural. Docker Desktop's WSL2 backend has no
/dev/kfdand no/dev/dri, only/dev/dxg. the host's own WSL ubuntu does see the card through DXG because its ROCm shipslibrocdxg.so; therocm/dev-ubuntu-22.04base images don't ship that shim, and mounting the host's copy in still segfaults rocminfo. so no flag combination gets this image onto a GPU there. the entrypoint currently tells such a user tomodprobe amdgpu, which is impossible on that host, and that message should say so.bottom line
solid: the image builds, the wheel resolution is correct on every ROCm index, the triton tree is consistent, and every failure path refuses clearly rather than half-working.
biggest gap: not one line of this has run on a GPU.
Dockerfile.studio-rocmhas never been built even once, and has two problems that don't need hardware to see.ENV SKIP_TORCH=truedoes nothing, install.sh derives that from--no-torch/UNSLOTH_NO_TORCHand sets it false unconditionally, so install.sh re-runs torch selection against an absent build-host GPU. andARG BASE_TAG=testdoesn't match the:latestthatbuild.sh --rocmproduces. happy to fix both, or drop the Studio image from this PR and do it separately.closes #6230 when it lands. supersedes #6231.