[CI/Build][Hardware][AMD] Rebuild the ROCm image on vllm/vllm-openai-rocm:nightly - #432
indianspeedster wants to merge 1 commit into
Conversation
00368f9 to
e70ae47
Compare
![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
This pull request introduces a new ROCm Docker build process by layering on the vllm/vllm-openai-rocm:nightly image, adding necessary training dependencies, and configuring patching logic for vLLM. Documentation and a justfile target were added to support this build. Review feedback identified an invalid SciPy version, an ineffective Docker layer cleanup step, and a potential error in the patch-checking logic.
| # same as docker/Dockerfile. The base's cupy-rocm-7-0 is built against the | ||
| # numpy 2 ABI and stops importing here; nothing in vime, vLLM, aiter or Megatron | ||
| # imports it. | ||
| RUN pip install "numpy==1.26.4" "scipy==1.17.1" |
There was a problem hiding this comment.
The version scipy==1.17.1 does not exist. The latest stable release of SciPy is in the 1.15.x series (as of early 2025), and version 1.17.1 has not been released. This will cause the pip install command to fail and break the Docker build.
Since the original code used "scipy<1.16", you should pin to a valid version like scipy==1.14.1 or scipy==1.13.1 which is compatible with Python 3.12 and NumPy 1.26.4.
RUN pip install "numpy==1.26.4" "scipy==1.14.1"
| # imports it. | ||
| RUN pip install "numpy==1.26.4" "scipy==1.17.1" | ||
|
|
||
| RUN rm -rf /root/.cache/pip |
There was a problem hiding this comment.
Running RUN rm -rf /root/.cache/pip in a separate RUN instruction does not reduce the size of the Docker image. In Docker, each RUN instruction creates a new read-only layer. Deleting files in a subsequent layer only hides them from the final container view but keeps them in the previous layers, resulting in image bloat.
To keep the image size minimal, you should either:
- Re-add the BuildKit cache mount
--mount=type=cache,target=/root/.cache/pipto theRUNinstructions (as was done in the original Dockerfile). - Or use the
--no-cache-dirflag for allpip installcommands and remove this redundantRUN rm -rf /root/.cache/pipstep.
# Redundant cache cleanup removed
| RUN VLLM_SITE="$(python3 -c 'import os, vllm; print(os.path.dirname(os.path.dirname(vllm.__file__)))')" && \ | ||
| cd "$VLLM_SITE" && \ | ||
| grep -q 'disable_tp=self.disable_tp' vllm/model_executor/models/qwen3_omni_moe_thinker.py && \ | ||
| if grep -q 'def merge' vllm/v1/outputs.py; then \ |
There was a problem hiding this comment.
If vllm/v1/outputs.py does not exist in the base image, grep will print an error to stderr: grep: vllm/v1/outputs.py: No such file or directory. While this won't break the build chain because it's wrapped in an if statement (which evaluates to false and exits with 0), it is cleaner and safer to check if the file exists first using [ -f ... ] before running grep.
if [ -f vllm/v1/outputs.py ] && grep -q 'def merge' vllm/v1/outputs.py; then \\
The ROCm image built everything from ubuntu:22.04 — ROCm 7.0.2, PyTorch,
triton, flash-attention, aiter and vLLM all from source. AMD now publishes
vllm/vllm-openai-rocm:nightly with that whole inference half already built
(ROCm 7.2.3, torch 2.12, triton 3.7.1, amd-aiter 0.1.21, flash-attn 2.8.3,
vLLM 0.28.1), so layer on it and add only the training half, mirroring how
docker/Dockerfile layers on vllm/vllm-openai:nightly.
Result: 51.9 GB on disk / 12.7 GB compressed, down from 81.9 / 21.6, and the
build drops from hours to ~15 min (TransformerEngine is the only long compile).
Added on top of the base: TransformerEngine (ROCm/TransformerEngine v2_15_rocm,
gfx950), apex, Megatron-LM plus the docker/amd_patch/ patches,
torch_memory_saver, rsync (the base omits it and command_utils.rsync_simple
shells out to it), the vime requirements and vime itself, including the
int4_qat HIP kernel the old image never built.
Dropped as now redundant: the ROCm/torch/triton/flash-attention/aiter/vLLM
source builds, the cupy and xgrammar/tilelang/fastapi pins, and the
TransformerEngine flash-attn max_version sed (2.8.3 is exactly TE 2.15's cap).
Three things this base needs that the CUDA one does not:
* VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1. This vLLM gates the token-in-token-out
endpoints behind an env var defaulting to off; older nightlies mounted them
unconditionally. Without it every rollout request to /inference/v1/generate
404s. docker/Dockerfile will need this too once its base rolls forward.
* GPU_ARCHS. ROCm apex imports aiter, which resolves the live gfx arch at
module import, so megatron.training cannot be imported without a visible
GPU. Pinning the arch fixes the build-time codegen path; the runtime
dispatch still needs a device, which is why the image's import check stops
at megatron.core.
* A tolerant vLLM patch apply. The ROCm nightly is cut from a different point
on vLLM main than the CUDA base docker/patch/ is written against — ahead of
it on some files, behind on others — so everything goes through
patch --fuzz=3 rather than git apply, and a single .rej sweep is what proves
each hunk landed. Four hunks or files do not fit, each dropped behind a
check that fails the build once the base converges:
- the Qwen3-Omni audio-encoder hunk, already upstream here in a better
form (TP sharding when the head count divides, disable_tp otherwise);
- the output_processor delta-sampling-mask hunk from vllm-project#410, which rewrites
a block built on SamplingMaskLists.merge that this base does not have,
so ROCm keeps emitting sampling masks only on finish;
- vllm-pd-request-metrics.patch in full, which patches
vllm/entrypoints/openai/engine/protocol.py, a module this base splits
across generate/base/ and serve/engine/. The cost is per-request PD
telemetry, which vime reads defensively in observability/trace_utils.py;
- the tests/ and rust/ hunks, which have no target: site-packages ships
neither tree.
vllm-inflight-queue-diagnostics.patch from vllm-project#410 applies in full.
numpy stays pinned to 1.x as in docker/Dockerfile — vime asserts it in
megatron_utils.initialize — which costs the base's cupy-rocm-7-0, built against
the numpy 2 ABI. Nothing in vime, vLLM, aiter or Megatron imports it.
Verified on MI350X by rebuilding the image and replaying the gating ROCm CI
step (.buildkite/pipeline-rocm.yaml, rocm-fully-async-short) against it:
tests/test_qwen2.5_0.5B_fully_async_short.py passes (exit 0), with
train_rollout_logprob_abs_diff 0.0117 and 0.0155 over the two steps against a
0.1 CI threshold, and rollout-vs-ref logprobs agreeing at -0.29962 / -0.29944
and -0.45356 / -0.45456.
Signed-off-by: indianspeedster <cspandey016@gmail.com>
e70ae47 to
f49b639
Compare
1 participant
What
docker/Dockerfile.rocmbuilt everything fromubuntu:22.04— ROCm 7.0.2, PyTorch, triton, flash-attention, aiter and vLLM all from source. AMD now publishesvllm/vllm-openai-rocm:nightlywith that whole inference half already built, so this layers on it and adds only the training half, mirroring howdocker/Dockerfilelayers onvllm/vllm-openai:nightly.Result: 51.9 GB on disk / 12.7 GB compressed, down from 81.9 / 21.6, and the build drops from hours to ~15 min (TransformerEngine is the only long compile).
Added on top of the base: TransformerEngine (
ROCm/TransformerEnginev2_15_rocm, gfx950), apex, Megatron-LM plus thedocker/amd_patch/patches, torch_memory_saver, rsync (the base omits it andcommand_utils.rsync_simpleshells out to it), the vime requirements and vime itself, including theint4_qatHIP kernel the old image never built.Dropped as now redundant: the ROCm/torch/triton/flash-attention/aiter/vLLM source builds, the cupy and xgrammar/tilelang/fastapi pins, and the TransformerEngine flash-attn
max_versionsed (2.8.3 is exactly TE 2.15's cap).Three things this base needs that the CUDA one does not
VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1. This vLLM gates the token-in-token-out endpoints behind an env var defaulting to off; older nightlies mounted them unconditionally. Without it every rollout request to/inference/v1/generate404s.docker/Dockerfilewill need this too once its base rolls forward.GPU_ARCHS. ROCm apex imports aiter, which resolves the live gfx arch at module import, somegatron.trainingcannot be imported without a visible GPU. Pinning the arch fixes the build-time codegen path; the runtime dispatch still needs a device, which is why the image's import check stops atmegatron.core.Patch handling vs.
docker/patch/latest/The ROCm nightly is cut from a different point on vLLM main than the CUDA base the patches are written against — ahead of it on some files, behind on others. So everything goes through
patch --fuzz=3rather thangit apply, and a single.rejsweep is what proves each hunk landed.Four hunks or files do not fit. Each is dropped behind a check that fails the build once the base converges, so none of them can quietly outlive the drift that motivated it:
vllm.patchdisable_tpotherwise)grep -q 'disable_tp=self.disable_tp'output_processordelta-sampling-mask hunk (from #410)SamplingMaskLists.merge, which this base does not have. ROCm keeps emitting sampling masks only on finishdef mergeappears invllm/v1/outputs.pyvllm-pd-request-metrics.patch, in full (from #410)vllm/entrypoints/openai/engine/protocol.py, a module this base splits acrossgenerate/base/andserve/engine/. Cost is per-request PD telemetry, which vime reads defensively inobservability/trace_utils.pytests/andrust/hunksvllm-inflight-queue-diagnostics.patch(from #410) applies in full.numpystays pinned to 1.x as indocker/Dockerfile— vime asserts it inmegatron_utils.initialize— which costs the base'scupy-rocm-7-0, built against the numpy 2 ABI. Nothing in vime, vLLM, aiter or Megatron imports it.Test plan
Built and run on 8x MI350X. The image reports torch 2.12.0+git6bbd260 / hip 7.2.53211, vLLM 0.28.1rc1.dev516+g9ea8f3ffc, TransformerEngine 2.15.0+e7835ed.
The gating ROCm CI step reproduced locally against the built image, exactly as
.buildkite/pipeline-rocm.yamlinvokes it:Result: passes, exit 0.
train_rollout_logprob_abs_diff= 0.0117 (step 0) and 0.0155 (step 1) against the 0.1 CI threshold; rollout-vs-ref logprobs agree at -0.29962 / -0.29944 and -0.45356 / -0.45456./inference/v1/generateserved 200 OK throughout at ~8.4k tok/s generation.