rocm: reduce image size by using a multi-stage build#2246
Conversation
mikebonnet
requested review from
bmahabirbu,
cgruver,
engelmi,
jhjaggars,
maxamillion,
rhatdan and
swarajpande5
as code owners
Reviewer's guide (collapsed on small PRs)Reviewer's GuideConverts the ROCm container image to a multi-stage build that installs llama/whisper binaries into a slim runtime image under /tmp/install, adjusts the build script to use a temporary install prefix for ROCm and to use GPU_TARGETS instead of AMDGPU_TARGETS, and removes some redundant install logic for whisper and llama binaries. File-Level Changes
Possibly linked issues
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Summary of ChangesHello @mikebonnet, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a multi-stage build strategy for the ROCm container image. The core purpose is to dramatically decrease the final image's footprint by isolating the build artifacts and only transferring the essential components to the production-ready image. This optimization ensures that the container remains lean while still providing full functionality for ROCm-accelerated applications. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey there - I've reviewed your changes and found some issues that need to be addressed.
- In
clone_and_build_llama_cpp,install_prefixis no longer defined after removing the local assignment, but it’s still used in the finalinstallcommand, which will cause the script to fail or use an unintended value—either reintroduce a local assignment fromset_install_prefixor adjust the install path accordingly. - The HIP build flags changed from
-DAMDGPU_TARGETS=...to-DGPU_TARGETS=...; please confirm this matches the current llama.cpp/whisper.cpp CMake option names, as this rename may be ignored or break target selection if the upstream variable is stillAMDGPU_TARGETS.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `clone_and_build_llama_cpp`, `install_prefix` is no longer defined after removing the local assignment, but it’s still used in the final `install` command, which will cause the script to fail or use an unintended value—either reintroduce a local assignment from `set_install_prefix` or adjust the install path accordingly.
- The HIP build flags changed from `-DAMDGPU_TARGETS=...` to `-DGPU_TARGETS=...`; please confirm this matches the current llama.cpp/whisper.cpp CMake option names, as this rename may be ignored or break target selection if the upstream variable is still `AMDGPU_TARGETS`.
## Individual Comments
### Comment 1
<location> `container-images/scripts/build_llama_and_whisper.sh:246-245` </location>
<code_context>
clone_and_build_llama_cpp() {
local llama_cpp_commit="${LLAMA_CPP_PULL_REF:-$DEFAULT_LLAMA_CPP_COMMIT}"
- local install_prefix
- install_prefix=$(set_install_prefix)
git_clone_specific_commit "${LLAMA_CPP_REPO:-https://github.com/ggml-org/llama.cpp}" "$llama_cpp_commit"
cmake_steps "${common_flags[@]}"
install -m 755 build/bin/rpc-server "$install_prefix"/bin/rpc-server
</code_context>
<issue_to_address>
**issue (bug_risk):** Removing the local install_prefix assignment risks installing rpc-server into an unintended path.
This function still calls `install ... "$install_prefix"/bin/rpc-server`, but `install_prefix` is no longer initialized here. Unless it’s guaranteed to be set in the environment, this will likely resolve to `/bin/rpc-server` or fail. Either restore `local install_prefix=$(set_install_prefix)` or remove this manual `install` in favor of `cmake --install` with the prefix configured in `cmake_steps`.
</issue_to_address>
### Comment 2
<location> `container-images/rocm/Containerfile:18` </location>
<code_context>
+ /tmp/install/bin/llama-quantize \
+ /tmp/install/bin/whisper-server \
+ /usr/bin/
+COPY --from=builder /tmp/install/lib64/*.so /usr/lib64/
+
+RUN dnf -y --setopt=install_weak_deps=false install hipblas rocblas rocm-hip rocm-runtime rocsolver && \
</code_context>
<issue_to_address>
**suggestion:** The COPY glob for shared libraries can fail if no .so files exist or miss versioned symlinks.
To avoid this fragility, either copy the entire directory (e.g., `COPY --from=builder /tmp/install/lib64/ /usr/lib64/`) or use a broader pattern such as `*.so*` so that all ROCm shared libraries and their symlinks are reliably included.
```suggestion
COPY --from=builder /tmp/install/lib64/ /usr/lib64/
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
The code changes refactor the ROCm container build to a multi-stage Dockerfile, separating the llama.cpp and whisper.cpp compilation into a builder stage and then copying specific binaries and shared libraries, along with installing ROCm runtime dependencies, into the final image. The build script was updated to correctly configure the install prefix and GPU targets for ROCm. A review comment points out that the rpc-server binary is not being copied to the final image and suggests using a wildcard in the COPY command for binaries to ensure all executables are included robustly.
mikebonnet
force-pushed
the
rocm-multi-stage
branch
2 times, most recently
from
3fbfc50 to
532e5d0
Compare
|
LGTM, assuming you tested this on an AMD system? |
|
Unfortunately I don't have a AMD system to test with. If someone could test |
|
Well i have a amd laptop, and it fails the same way with your image as the quay.io/ramalama/ramalama image, luckily for me, the laptop works with vulkan driver. |
|
.Memory critical error by agent node-0 (Agent handle: 0x3390320) on address 0x7f599f65e000. Reason: Memory in use. |
|
Does the quay.io/ramalama/rocm:0.15 image behave the same way? |
|
Yes |
|
There were reports that it was a firmware problem. Have you tried with different/older firmware? |
I'm almost done setting up an AMD system so can test this soon, tomorrow hopefully |
mikebonnet
force-pushed
the
rocm-multi-stage
branch
2 times, most recently
from
7677364 to
267815b
Compare
|
I have an amd system but getting is there a new manifest i can use? btw regular rocm on 16 works for me! with a 9070xt |
Let me rebase and rebuild, which will create a new manifest. |
mikebonnet
force-pushed
the
rocm-multi-stage
branch
from
267815b to
2ae90c5
Compare
|
@bmahabirbu You can test with |
mikebonnet
force-pushed
the
rocm-multi-stage
branch
2 times, most recently
from
3731a60 to
71efd86
Compare
container-images/rocm/Containerfile
Outdated
| /tmp/install/bin/rpc-server \ | ||
| /tmp/install/bin/whisper-server \ | ||
| /usr/bin/ | ||
| COPY --from=builder /tmp/install/lib64/*.so /usr/lib64/ |
There was a problem hiding this comment.
symlinks could be an issue I think:
# ls -l /usr/lib64/libggml*
-rw-r--r--. 1 root root 992116 Jan 8 16:07 /usr/lib64/libggml-base.a
lrwxrwxrwx. 1 root root 17 Jan 8 16:11 /usr/lib64/libggml-base.so -> libggml-base.so.0
lrwxrwxrwx. 1 root root 21 Jan 8 16:11 /usr/lib64/libggml-base.so.0 -> libggml-base.so.0.9.5
-rwxr-xr-x. 1 root root 728752 Jan 8 16:08 /usr/lib64/libggml-base.so.0.9.5
-rw-r--r--. 1 root root 986336 Jan 8 16:07 /usr/lib64/libggml-cpu.a
-rw-r--r--. 1 root root 57602 Jan 8 16:08 /usr/lib64/libggml.a
lrwxrwxrwx. 1 root root 12 Jan 8 16:11 /usr/lib64/libggml.so -> libggml.so.0
lrwxrwxrwx. 1 root root 16 Jan 8 16:11 /usr/lib64/libggml.so.0 -> libggml.so.0.9.5
-rwxr-xr-x. 1 root root 47104 Jan 8 16:10 /usr/lib64/libggml.so.0.9.5
Would something like this work?
RUN --mount=type=bind,from=builder,source=/tmp/install,target=/tmp/install cp -a /tmp/install/lib64/* /usr/lib64/
or just keep it simple and cp -a /tmp/install/* /usr/
There was a problem hiding this comment.
A lot of the space savings comes from not copying the unnecessary .a files. Are the symlinks required? In my testing it looked like like llama-server worked fine with just the .so files. If it's not working for you, I can figure out a different way to copy the symlinks over.
There was a problem hiding this comment.
Giving the --mount idea a try.
There was a problem hiding this comment.
RUN --mount worked like a charm:
$ podman run -it --rm quay.io/redhat-user-workloads/ramalama-tenant/rocm:on-pr-bcf7372735e92e94ec51339bc1468e43c0a694f9 bash -c 'ls -la /usr/lib64/libggml* /usr/lib64/libllama*'
lrwxrwxrwx. 1 root root 17 Jan 10 08:21 /usr/lib64/libggml-base.so -> libggml-base.so.0
lrwxrwxrwx. 1 root root 21 Jan 10 08:21 /usr/lib64/libggml-base.so.0 -> libggml-base.so.0.9.5
-rwxr-xr-x. 1 root root 728752 Jan 10 08:04 /usr/lib64/libggml-base.so.0.9.5
lrwxrwxrwx. 1 root root 12 Jan 10 08:21 /usr/lib64/libggml.so -> libggml.so.0
lrwxrwxrwx. 1 root root 16 Jan 10 08:21 /usr/lib64/libggml.so.0 -> libggml.so.0.9.5
-rwxr-xr-x. 1 root root 47104 Jan 10 08:19 /usr/lib64/libggml.so.0.9.5
lrwxrwxrwx. 1 root root 13 Jan 10 08:21 /usr/lib64/libllama.so -> libllama.so.0
lrwxrwxrwx. 1 root root 17 Jan 10 08:21 /usr/lib64/libllama.so.0 -> libllama.so.0.0.1
-rwxr-xr-x. 1 root root 3010024 Jan 10 08:20 /usr/lib64/libllama.so.0.0.1
There was a problem hiding this comment.
You can test the latest build with --image quay.io/redhat-user-workloads/ramalama-tenant/rocm:on-pr-bcf7372735e92e94ec51339bc1468e43c0a694f9
There was a problem hiding this comment.
New development, GGML_CPU_ALL_VARIANTS causes the .so files for the backends to be written to (and loaded from) bin/, so the Containerfile now copies those as well.
You can test the latest build with --image quay.io/redhat-user-workloads/ramalama-tenant/rocm:on-pr-61f7101c2a24f5fb2e152d42f879c02447c32962
There was a problem hiding this comment.
LGTM
$ ramalama bench --image quay.io/redhat-user-workloads/ramalama-tenant/rocm:on-pr-61f7101c2a24f5fb2e152d42f879c02447c32962 granite
ggml_cuda_init: found 1 ROCm devices:
Device 0: AMD Radeon RX 9070 XT, gfx1201 (0x1201), VMM: no, Wave Size: 32
load_backend: loaded ROCm backend from /usr/bin/libggml-hip.so
load_backend: loaded RPC backend from /usr/bin/libggml-rpc.so
load_backend: loaded CPU backend from /usr/bin/libggml-cpu-haswell.so
| model | size | params | backend | ngl | test | t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
| granite 3B Q4_K - Medium | 4.65 GiB | 8.17 B | ROCm | 999 | pp512 | 2636.43 ± 436.43 |
| granite 3B Q4_K - Medium | 4.65 GiB | 8.17 B | ROCm | 999 | tg128 | 69.10 ± 0.23 |
build: 9689295 (1)
Prompt processing also seems to be much better (vs older llama.cpp on latest):
Device 0: AMD Radeon RX 9070 XT, gfx1201 (0x1201), VMM: no, Wave Size: 32
| model | size | params | backend | ngl | test | t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
| granite 3B Q4_K - Medium | 4.65 GiB | 8.17 B | ROCm | 999 | pp512 | 595.84 ± 0.36 |
| granite 3B Q4_K - Medium | 4.65 GiB | 8.17 B | ROCm | 999 | tg128 | 69.15 ± 0.22 |
build: b52edd2 (1)
71efd86 to
bcf7372
Compare
Only copy required binaries and libraries from the installation directory into the final image, and install only necessary runtime dependencies. The final image size is reduced by over 2Gb. Signed-off-by: Mike Bonnet <[email protected]>
mikebonnet
force-pushed
the
rocm-multi-stage
branch
from
bcf7372 to
61f7101
Compare
4 participants
Only copy required binaries and libraries from the installation directory into the final image, and install only necessary runtime dependencies. The final image size is reduced by over 2Gb.
Summary by Sourcery
Use a multi-stage build for the ROCm container image and adjust build scripts to support installation into a temporary prefix for ROCm builds.
Enhancements: