Add AMD ROCm Docker image and Tinker training examples#1892
Add AMD ROCm Docker image and Tinker training examples#1892eddierichter-amd wants to merge 7 commits into
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 AMD ROCm support for SkyRL's Tinker-compatible API, adding a dedicated Dockerfile, pyproject configuration, and example scripts for running a Tinker server and GRPO/smoke clients on AMD GPUs. However, several critical issues were identified in the changes: removing LayerwiseReloadWorkerMixin from NewInferenceWorkerWrap will break weight updates, and defining finish_weight_update directly risks collisions with newer vLLM versions. Additionally, non-existent package versions are specified for boto3, s3transfer, and transformers, and an unsupported --active flag is passed to uv run in the server startup script.
| _empty_cuda_cache() | ||
|
|
||
| def skyrl_finish_weight_update(self) -> None: | ||
| """ | ||
| Finalize the current weight update. | ||
|
|
||
| For checkpoint-format weights, runs layerwise postprocessing | ||
| (quantization repacking, attention weight processing, etc.). | ||
| Must be called after all update_weights_* calls are done. | ||
|
|
||
| Keep the SkyRL-prefixed name to avoid colliding with vLLM Worker | ||
| methods such as finish_weight_update. | ||
| """ | ||
| super().skyrl_finish_weight_update() | ||
| _empty_cuda_cache() |
There was a problem hiding this comment.
Can you make this modification directly in LayerwiseReloadWorkerMixin itself?
| def _empty_cuda_cache() -> None: | ||
| """Release unused CUDA/ROCm cached blocks after full-weight sync.""" | ||
| if not torch.cuda.is_available(): | ||
| return |
There was a problem hiding this comment.
| def _empty_cuda_cache() -> None: | |
| """Release unused CUDA/ROCm cached blocks after full-weight sync.""" | |
| if not torch.cuda.is_available(): | |
| return | |
| def _empty_cuda_cache_rocm() -> None: | |
| """Release unused ROCm cached blocks after full-weight sync.""" | |
| is_rocm = torch.version.hip is not None | |
| if not torch.cuda.is_available() or not is_rocm: | |
| return |
Since this is only needed for RoCm, it would be good to add this guard.
2 participants
This PR adds an AMD ROCm path for running SkyRL’s Tinker-compatible API with FSDP training and vLLM ROCm inference.
Changes
docker/Dockerfile.amdfor a ROCm-based SkyRL image.docker/pyproject.amd.tomlto avoid CUDA-specific dependencies in the AMD image.examples/train/amd/run_tinker_server_amd.shto launch the Tinker API on AMD nodes.tinker_hello_world.py: small LoRA smoke test.grpo_client.py: GSM8K GRPO/PPO training example.examples/train/amd/README.mdwith build/run instructions.