Since the vime launch, the AMD team has been working closely with the vime team to bring compatibility support to ROCm, validating the end-to-end pipeline on AMD Instinct hardware, upstreaming ROCm-specific fixes, and shipping a prebuilt container so AMD users can get started without building from source.

We are excited to announce ROCm support for vime, the vLLM ecosystem’s reinforcement learning framework, on AMD Instinct MI355X GPUs. This support enables large-scale RL post-training workflows to run natively on AMD hardware, bringing the full vime pipeline to the ROCm ecosystem.

This blog covers: an overview of vime and its architecture, why AMD Instinct GPUs are well-suited for RL workloads, the current state of ROCm support and validated features, and a walkthrough of running an end-to-end RL training job on ROCm.

vime Explained

vime was announced in June 2026 by the vLLM team and has quickly become a focal point for RL post-training in the vLLM ecosystem. With ROCm support, these workflows now run natively on AMD Instinct GPUs out of the box.

vime architecture
vime architecture.

vime adopts slime’s three-stage, decoupled train-inference design, with the difference being that the rollout backend is vLLM instead of SGLang:

  • Training (Megatron): the main training loop, responsible for parameter updates and synchronizing weights to the rollout side.
  • Rollout (vLLM + Router): inference sampling, producing training samples with reward or verifier signals.
  • Data Buffer: connects the training and rollout sides, managing prompt injection and custom rollout logic.

This entire pipeline has now been validated end-to-end on ROCm.

Why AMD Instinct™ GPUs

RL post-training is among the most memory-intensive workloads in modern ML. Each training step requires holding both training-side weights (in Megatron format) and inference-side KV cache (for vLLM rollouts). In colocated mode, these compete for the same device memory pool. AMD Instinct MI300X and MI355X GPUs are exceptionally well-suited to this workload profile for several reasons.

  • Large unified HBM capacity: The MI300X provides 192 GB of HBM3 per GPU; the MI355X raises this to 288 GB. This headroom makes it easier to fit large models for training without requiring aggressive tensor parallelism to distribute memory pressure, reducing topology complexity and improving cluster utilization.

  • High memory bandwidth: HBM3 delivers over 5 TB/s of aggregate bandwidth on MI300X, and with HBM3E, 8 TB/s on MI355X. RL rollout, autoregressive token generation at scale, is fundamentally memory-bandwidth-bound: each decode step loads the full KV cache and model weights from HBM. Higher bandwidth shortens step latency and improves throughput on the rollout phase that dominates overall step time in most RL pipelines.

  • Open software ecosystem: ROCm is AMD’s open-source GPU compute platform, built on open standards (HIP, LLVM, MIOpen). vLLM and PyTorch both support ROCm natively, meaning vime inherits the full vLLM rollout stack without a separate code path. Teams running ROCm already have a familiar toolchain they can extend.

Training in Details

Enabling vime on ROCm required validating and integrating several components of the stack. Here is what happens under the hood when you run vime on AMD GPUs.

  • Megatron-LM training backend: vime uses Megatron-LM as the training engine. On ROCm, the Megatron stack builds cleanly using a ROCm-compatible fork and a small patch that guards CUDA fused-kernel initialization on non-CUDA builds. The training loop runs with ROCm-compatible Megatron patches and ROCm-specific launch flags. Gradient accumulation uses the native PyTorch path, which is fully supported on ROCm. Checkpoint conversion from HuggingFace format to Megatron’s torch_dist format runs on a single GPU and completes cleanly, producing a layout that Megatron loads correctly at the start of each training run.

  • Colocated weight synchronization: In colocated mode, Megatron and vLLM share the same GPU pool rather than running on separate node partitions. After each optimizer step, Megatron synchronizes updated weights to the vLLM engine via IPC, so the rollout workers always generate from the latest policy without a network round-trip. On ROCm, the torch.cuda.get_device_properties(i).uuid interface returns stable, process-consistent device UUIDs, so vime’s UUID-keyed IPC routing works correctly without modification.

  • GPU visibility and Ray integration: ROCm uses HIP_VISIBLE_DEVICES to control GPU assignment. The vime launch script sets this alongside CUDA_VISIBLE_DEVICES so both the Megatron training actor and the vLLM subprocess see consistent device ordinals throughout the job. Ray’s AMD GPU manager is configured to not override these visibility masks, so the job driver and all Ray actors, including the Megatron training actor and its worker processes, operate on the correct set of GPUs without contention. The container is also started with a raised file descriptor limit (--ulimit nofile=1048576:1048576), which Ray requires when spawning the full set of actor workers at scale.

Getting Started: Run vime on AMD GPUs

vime provides a ROCm-ready workflow with a prebuilt container, so you can run the full RL pipeline with minimal setup.

Launch the vime ROCm container

# Pull the ROCm image
docker pull vllm/vime-rocm
# Start the container
docker run -d --name vime --ulimit nofile=1048576:1048576 \
  --ipc=host --network=host --device=/dev/kfd --device=/dev/dri \
  --security-opt seccomp=unconfined --group-add video --privileged \
  -e WANDB_API_KEY=$wandb_key vllm/vime-rocm
# The launch script enables W&B online mode, so a valid WANDB_API_KEY is required.

# Enter the container
docker exec -it vime bash

The container includes vLLM and Megatron-LM preinstalled, along with the vime codebase at /root/vime.

Download model and dataset

# Download model weights (Qwen3-8B)
hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B
# Download training dataset (dapo-math-17k)
hf download zhuzilin/dapo-math-17k --repo-type dataset --local-dir /root/dapo-math-17k

Convert weights to Megatron format

Load the model configuration for Qwen3-8B, then run the conversion.

cd /root/vime && source scripts/models/qwen3-8B.sh
HIP_VISIBLE_DEVICES=0 PYTHONPATH=/root/vime:/root/Megatron-LM \
  torchrun --nproc-per-node=1 tools/convert_hf_to_torch_dist.py "${MODEL_ARGS[@]}" \
  --no-gradient-accumulation-fusion --attention-backend flash \
  --hf-checkpoint /root/Qwen3-8B --save /root/Qwen3-8B_torch_dist

Note: On ROCm, use HIP_VISIBLE_DEVICES in place of CUDA_VISIBLE_DEVICES to select GPUs.

Launch RL training

NUM_ROLLOUT=100 VISIBLE_GPUS=0,1 bash scripts/run-qwen3-8B-amd.sh

This launches a full RL pipeline with colocated training and inference:

  • vLLM rollout workers
  • GRPO training loop
  • On-policy rollout → train → weight-sync cycle

Configuration notes:

  • VISIBLE_GPUS - two free GPU indices; the script masks execution to these GPUs and avoids clashes. Uses TP=2, single vLLM engine, colocate mode, DP=1.
  • NUM_ROLLOUT - number of training steps (default is 3 for a smoke test).
  • Each run requires approximately 230 GB across the two selected GPUs. Launch only on GPUs with sufficient free memory.

After finishing a run, if rerunning with a different NUM_ROLLOUT, clear the save directory to avoid checkpoint mismatch: rm -rf /root/Qwen3-8B_vime/

Performance Results

With the above runbook, we’re able to test several models such as the Qwen3-4B, Qwen3-8B (dense), and Qwen3-30B-A3B (MoE) models. Below is some performance data we obtained from running the Qwen3-8B example above:

tps
Throughput tending slightly upward with Qwen3-8B model on MI355X.

As shown, throughput on MI355X sustains approximately 4,100 tokens_per_gpu_per_second across 100 training steps, with a slight upward trend over time. This improvement reflects the policy learning to produce more predictable outputs as training progresses: shorter or more uniform generations reduce decode variance and allow the vLLM rollout engine to batch more efficiently.

logprob diff
Train-rollout logprob absolute difference holding steady around 0.012 with Qwen3-8B model on MI355X.

The train_rollout_logprob_abs_diff metric, which measures divergence between the training-side log probabilities and the rollout-side log probabilities, holds steady around 0.012 and trends slightly downward across the run. Weight synchronization between the Megatron training backend and the vLLM rollout workers keeps the two sides consistent, preventing the logprob drift that would otherwise corrupt the policy gradient signal. A stable, low logprob diff is a prerequisite for reliable GRPO updates; values this low are on par with reported results on NVIDIA hardware.

raw reward
Raw reward climbing from near 0 to around 0.5~0.6 by step 100 with Qwen3-8B model on MI355X.

raw_reward is measured on sampled training prompts and starts near 0 at step 0, climbing gradually to around 0.5~0.6 by step 100. At initialization, the model has not yet been shaped by RL, so it approaches math problems with its base pretraining distribution; on competition-level problems from the dapo-math-17k dataset, a freshly initialized policy solves very few, yielding rewards close to zero. As training proceeds, the policy receives gradient signal from problems it gets partially or fully correct, and begins to favor reasoning patterns that the verifier rewards. The rising training reward indicates optimization progress on the sampled training prompts; note that the ROCm launcher disables evaluation (EVAL_ARGS=()), so held-out evaluation is needed to assess generalization.

Feature Support Roadmap on AMD

Today, core vime functionality is supported on AMD, including:

  • GRPO training
  • Colocated training and rollout
  • Asynchronous (non-colocated) training with disjoint actor and rollout GPU pools
  • Megatron-LM training backend
  • vLLM rollout backend
  • Qwen3 Dense and MoE model support

Looking ahead, the vime and AMD teams are committed to expanding support for additional capabilities, including:

  • Full vLLM Router and PD disaggregation support
  • FP8 pipeline optimization
  • R3 (Rollout Routing Replay) for AMD MoE workloads
  • Performance optimization for the asynchronous training pipeline (improving train-rollout logprob divergence and addressing memory leak issues)
  • Agentic RL for multi-turn tool calling and multi-agent settings

Our goal is continuous performance and capability improvements aligned with the evolving vime and vLLM roadmap.

Acknowledgments

We would like to thank all the contributors who made this work possible:

AMD contributors & vime community

We are grateful for their contributions, collaboration, and support throughout this work.

References