ci(rocm): AMD ROCm GPU CI on Buildkite#356
Conversation
![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 support for AMD ROCm GPU suites by adding a Buildkite pipeline configuration, updating the Qwen2.5 0.5B fully async short test with ROCm-specific configurations, and adapting the training execution utility to run directly on ROCm to bypass a Ray job submission race condition. The feedback highlights a potential shell parsing issue where environment variables in amd_exports are joined without escaping, and suggests using shlex.quote to safely handle values containing spaces or shell metacharacters.
| **extra_env_vars, | ||
| **_parse_extra_env_vars(config.extra_env_vars), | ||
| } | ||
| amd_exports = " ".join(f"{k}={v}" for k, v in amd_env.items()) |
There was a problem hiding this comment.
When constructing amd_exports by joining key-value pairs with spaces, any environment variable value containing spaces or shell metacharacters will cause shell parsing errors or unexpected behavior when executed via exec_command. Using shlex.quote ensures that the values are safely escaped for the shell.
import shlex
amd_exports = " ".join(f"{k}={shlex.quote(str(v))}" for k, v in amd_env.items())
Add U.is_rocm() (torch.version.hip gate) and a ROCm branch in execute_train that runs the train script directly against the ray head instead of 'ray job submit' (avoids the ROCm 'No available agent' race). Adapt the Qwen2.5-0.5B fully-async GRPO test for ROCm: HF->Megatron convert to a container-local /tmp path (no shared-model race), --ref-load it, drop bridge, add AMD megatron flags, lower vLLM mem. The ROCm checkpoint-writer fix comes from upstream (vime/utils/rocm_checkpoint_writer.py). Signed-off-by: indianspeedster <cp3793@nyu.edu>
8f9b77a to
959dbff
Compare
There was a problem hiding this comment.
Pull request overview
Adds an AMD ROCm CI path (Buildkite) and a ROCm-specific execution branch so GPU tests can run reliably in ROCm containers (working around Ray job submission flakiness), plus adapts the fully-async short Qwen2.5 0.5B test to the ROCm conversion/no-bridge flow.
Changes:
- Add
.buildkite/pipeline-rocm.yamlwith a manual-gated ROCm GPU test group running inside a prebuilt ROCm Docker image and usinggpu_lock_execwithHIP_VISIBLE_DEVICES. - Add a ROCm gate (
U.is_rocm()) and a ROCm-specific training execution path that runs the train script directly against the Ray head (instead ofray job submit). - Update
test_qwen2.5_0.5B_fully_async_short.pyto perform HF→Megatron conversion on ROCm and adjust vLLM/train flags accordingly.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
vime/utils/external_utils/command_utils.py |
Adds is_rocm() and a ROCm-specific execute_train path that bypasses ray job submit. |
tests/test_qwen2.5_0.5B_fully_async_short.py |
Switches ROCm runs to use HF→Megatron conversion and ROCm-friendly vLLM/train args. |
.buildkite/pipeline-rocm.yaml |
New Buildkite pipeline defining manual-gated ROCm GPU suites running inside the ROCm container. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| amd_env = { | ||
| "no_proxy": f"127.0.0.1,{master_addr}", | ||
| "PYTHONUNBUFFERED": "1", | ||
| "PYTHONPATH": f"{repo_base_dir}:/root/Megatron-LM/", | ||
| "CUDA_DEVICE_MAX_CONNECTIONS": "1", | ||
| "NCCL_NVLS_ENABLE": "0", | ||
| "MASTER_ADDR": master_addr, | ||
| "WANDB_MODE": "disabled", | ||
| **extra_env_vars, | ||
| **_parse_extra_env_vars(config.extra_env_vars), | ||
| } | ||
| amd_exports = " ".join(f"{k}={v}" for k, v in amd_env.items()) | ||
| exec_command( | ||
| f"export {amd_exports} && " | ||
| f"{cmd_megatron_model_source}" | ||
| f"python3 {train_script} {model_args} {train_args}" | ||
| ) |
Declarative ROCm Buildkite pipeline (.buildkite/pipeline-rocm.yaml) in the style of vllm-omni's .buildkite/test-amd*.yaml: one step per test on a self-hosted AMD agent queue, behind a manual block gate. Each step runs the prebuilt ROCm image via docker (/dev/kfd+/dev/dri) and arbitrates GPUs with gpu_lock_exec --target-env-name HIP_VISIBLE_DEVICES. short suite: fully_async hard-fails (validated); gsm8k soft-fails until the vLLM/ROCm NaN-logprob issue is fixed (mirrors vllm-omni's grade/soft_fail convention). Signed-off-by: indianspeedster <cp3793@nyu.edu>
474a611 to
5490d6c
Compare
Adds an AMD ROCm CI path for vime on Buildkite.
.buildkite/pipeline-rocm.yaml— declarative ROCm pipeline behind a manual block gate; each test runs the prebuilt ROCm image via docker and arbitrates GPUs withgpu_lock_execonHIP_VISIBLE_DEVICES.U.is_rocm()gate + a ROCm branch inexecute_train(runs the train script directly against the ray head).test_qwen2.5_0.5B_fully_async_short.pyadapted for the ROCm path (HF→Megatron convert, no bridge).short suite: fully_async hard-fails (validated end-to-end on gfx950);
Draft — Buildkite agent provisioning (queue
amd_gfx950) is a separate follow-up.