[Bugfix][Hardware][AMD] Fix colocate weight sync on ROCm and add a ROCm 7.14 Dockerfile

[Treemann]

Description

1. Colocate weight sync fails on every GPU but 0 (ROCm)

On ROCm a HIP IPC handle carries the exporter's device ordinal, and hipIpcOpenMemHandle
resolves it in the importer's enumeration; CUDA honours the importer's current device instead.
In colocate mode the trainer exports and the rollout engine imports, so a trainer rank that sees the
whole node exports at ordinal N while an engine launched with HIP_VISIBLE_DEVICES=<its own gpu>
enumerates one device where ordinal N does not exist. The two agree only when N is 0 — hence weight
sync works on GPU 0 and fails elsewhere with hipErrorInvalidValue.

Not a vLLM defect: the engine's IPC consumer is correct, the mask it was launched with is too
narrow. Only the mask's length matters, so listing the engine's own GPUs first keeps it on its
intended devices while putting the trainer's ordinal back in range.

+def _engine_visible_devices(base: int, local_num_gpus: int, num_gpus_per_node: int) -> str:
+    own = [base + i for i in range(local_num_gpus)]
+    try:
+        import torch
+
+        is_hip = torch.version.hip is not None
+    except Exception:
+        is_hip = False
+    if not is_hip:
+        return ",".join(str(g) for g in own)
+    rest = [g for g in range(max(num_gpus_per_node, base + local_num_gpus)) if g not in own]
+    return ",".join(str(g) for g in own + rest)

-    kwargs["_visible_devices"] = ",".join(str(base + i) for i in range(local_num_gpus))
+    kwargs["_visible_devices"] = _engine_visible_devices(base, local_num_gpus, args.num_gpus_per_node)

Gated on torch.version.hip, so on CUDA the helper returns the byte-identical string the current
code produces and nothing changes. The mask only differs when the engine is narrower than the node:

identical  1 engine over both GPUs of a 2-GPU node: before='0,1'  after='0,1'
CHANGED    2 engines of 1 GPU, engine on GPU 1:     before='1'    after='1,0'
CHANGED    8-GPU node, 1-GPU engine on GPU 3:       before='3'    after='3,0,1,2,4,5,6,7'

2. A ROCm 7.14 Dockerfile

docker/Dockerfile.rocm targets ROCm 7.0.2 / gfx950 and builds the stack onto ubuntu:22.04 across
nine stages. docker/Dockerfile.rocm7.14 instead starts from a prebuilt ROCm 7.14 image
(rocm/primus:v26.4, which already ships torch, triton, flash-attention, TransformerEngine and
aiter) and adds only vLLM, Megatron-LM and vime. Single stage, no build context — every source is a
pinned git coordinate with its commit asserted, nothing is COPYed in. Additive;
docker/Dockerfile.rocm is untouched.

Two guards are worth calling out:

• The base image's ROCm torch is never replaced. A pip constraints file is generated from the
  installed versions after the vLLM build and every later pip install runs under it. vime's
  requirements.txt leaves vllm-router, transformers and ray[default] unpinned, and a
  resolver satisfying one of them with a PyPI torch/vllm wheel — those are CUDA-only — breaks
  the image in a way that only surfaces on the first GPU op. A --dry-run check aborts the build if
  the resolver still plans to do so.
• A verification layer fails the build, not the first training run: ROCm torch untouched,
  VLLM_TAG actually installed, all imports resolving, and the vLLM API window vime needs present
  at both ends.

Verification

8× MI308X (gfx942), ROCm 7.14, torch 2.12.0+rocm7.14, vLLM v0.29.0 built from source,
Megatron-LM 1dcf0dafa + docker/amd_patch/latest/, colocate vLLM + Megatron, Qwen2.5-0.5B.

End-to-end A/B, single variable. Same vime revision, vLLM, script and topology; only the fix
differs. The run uses 4 actor GPUs with --rollout-num-gpus-per-engine 1, so each engine is
launched with a mask narrower than the trainer's. It runs one rollout step, far enough to reach the
first weight sync.

Arm	Result
without the fix	POST /update_weights → 500, CUDA error: invalid argument in rebuild_cuda_tensor → UntypedStorage._new_shared_cuda, on 3 of 4 engines (GPU 0's is fine — its ordinal is already 0)
with the fix	exit 0, weight sync completed on all 4 engines

Reproduced twice per arm, from separately built images.

Mechanism, without vime. A two-process probe (one 4 KiB tensor, torch only, ~10 s) isolates the
ordinal resolution. The exporter sees all 8 GPUs and allocates on GPU 3. An importer with mask=3
then fails with invalid device ordinal, while an importer with mask=3,0,1,2,4,5,6,7 succeeds and
returns an exact checksum, so the mapping landed on the intended device. With --target 0 both arms
pass, which is the expected asymmetry.

Build. docker/Dockerfile.rocm7.14 builds from this branch's head (sha assertion passes,
17 layers, ~3 min warm / ~10 min cold) and its verification layer reports:

[ok] torch is a ROCm build / not a CUDA wheel — 2.12.0+rocm7.14.0a20260608, hip 7.14.60850
[ok] vllm 0.29.0   [ok] numpy 1.26.4
[ok] import vime / megatron.core / megatron.training / ray / torch_memory_saver
[info] this vime revision imports vllm.entrypoints.launchers.cli_args
[ok] launchers.cli_args.FrontendArgs   [ok] nccl_engine.NCCLTrainerInitInfo
[ok] ipc_engine.IPCTrainerInitInfo     [ok] factory.WeightTransferTrainerFactory

That same image ran the A/B above, so both commits are verified together. The cli_args assertion
reads which module this vime revision imports rather than hardcoding a path, so it stays
meaningful as main advances.

pre-commit passes every hook on this branch's changes.

pytest tests/ (in the image, excluding tests/ci): 19 failed, 611 passed, 13 skipped, 57 errors. None involve this PR. The two files covering the changed path pass —
tests/utils/test_vllm_engine.py (65 passed, exercises _compute_server_args and
_visible_devices) and tests/utils/test_vllm_config.py (17 passed) — and every file that reported
a failure also passes on its own (test_vllm_arguments.py 28, test_train_data_utils.py 10,
test_qwen3_5_vl_native.py 8; tests/utils/ in random order 152). The whole-suite result is
pre-existing import-order pollution: ValueError: megatron.core.__spec__ is None from find_spec
after another module imported it, plus pytest failing on filenames containing dots
(test_qwen3_0.6B_parallel_check.py → No module named 'test_qwen3_0').

Scope

• ROCm-only by construction; a no-op on CUDA. No CUDA hardware was available, so CUDA safety rests on the torch.version.hip gate and the expression being unchanged, not on a CUDA run.
• The A/B is deliberately minimal. Larger models, multi-node and fully-async mode were exercised during development but not in it.
• VLLM_TAG=v0.29.0 is a compatibility claim about a moving target, since vime's mainline follows vLLM nightly: vime needs NCCLTrainerInitInfo (absent before v0.28.0) and vllm.entrypoints.launchers.cli_args (the Sep-2026 migration). v0.29.0 has both, and the verification layer asserts that at build time instead of leaving it to a comment.
• Separately: docker/Dockerfile.rocm's VLLM_TAG=6e448d0ea (v0.27.1) now satisfies neither end, so the existing ROCm Dockerfile cannot build a working vime against current main. Not addressed here; happy to open an issue.

Prior art

Same principle as ms-swift#9627, which reached it for cross-process RCCL peer resolution: on ROCm, never shrink the visibility mask, only reorder it.

AI assisted contribution

Per CONTRIBUTING.md, this work was developed with AI assistance. The analysis, the helper, the Dockerfile, the probes and this description were produced with an AI agent driving the debug loop on the hardware above. The submitter has reviewed every changed line, built the images and run both A/B arms, and is responsible for the content — including the Scope section. Both commits carry:

Co-authored-by: Claude
Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Treemann <rongzhang.zheng@amd.com>

Checklist

• PR title prefixed; both commits signed off
• AI assistance disclosed and attributed
• pre-commit passes
• pytest — the tests covering this change pass; all failures are pre-existing (see above)
• Docs / linked Issue — n/a (no user-facing flags; one bugfix + an additive Dockerfile)

[gemini-code-assist]

Code Review

This pull request adds a Dockerfile for AMD ROCm 7.14 and updates the GPU visibility mask generation in vllm_engine.py to prevent HIP IPC handle resolution failures on ROCm. Feedback on the implementation highlights a potential issue where device indices could exceed the node's GPU count if the base index is offset, and suggests using a modulo operation to wrap the indices safely.

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #34489063 | 📁 Comparing f5c3c88 against latest (ce92eff)

  🔍 Preview build  

45 files changed · ± 45 modified

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 35 more...