Add AMD ROCm Docker image and Tinker training examples

docker/Dockerfile.amd

@@ -0,0 +1,58 @@
+FROM vllm/vllm-openai-rocm:v0.20.2
+
+ARG DEBIAN_FRONTEND=noninteractive
+ARG RAY_VERSION=2.51.1
+ARG UV_VERSION=0.11.11
+
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+
+ENV HF_HUB_ENABLE_HF_TRANSFER=1 \
+    PATH=/root/.local/bin:${PATH} \
+    RAY_RUNTIME_ENV_HOOK=ray._private.runtime_env.uv_runtime_env_hook.hook \
+    UV_FIND_LINKS=/opt/wheels \
+    UV_PROJECT_ENVIRONMENT=/usr
+
+RUN apt-get update -y && apt-get install -y --no-install-recommends \
+        build-essential \
+        ca-certificates \
+        curl \
+        git \
+        iproute2 \
+        iputils-ping \
+        kmod \
+        libnuma-dev \
+        libxml2 \
+        net-tools \
+        netcat-openbsd \
+        openssh-server \
+        traceroute \
+        tzdata \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN if ! command -v uv >/dev/null 2>&1; then \
+        curl -LsSf "https://astral.sh/uv/${UV_VERSION}/install.sh" | sh; \
+    fi \
+    && mkdir -p /opt/wheels \
+    && uv --version
+
+# The base vLLM ROCm image supplies ROCm builds of torch, vLLM, and flash-attn.
+# Keep those packages in the system environment and install only the SkyRL
+# runtime pieces that are not provided by the base image.
+RUN python -m pip install --no-cache-dir "ray[default]==${RAY_VERSION}" \
+    && python -m pip install --no-cache-dir "flash-linear-attention[rocm]" orjson torchdata \
+    && python -m pip install --no-cache-dir --no-deps boto3==1.43.0 s3transfer==0.17.1 \
+    && python -m pip install --no-cache-dir "pyOpenSSL<26" "cryptography<49,>=2.5" \
+    && python -m pip install --no-cache-dir --upgrade "importlib-metadata>=6.0,<8.8.0"
+
+WORKDIR /workspace/SkyRL
+COPY . /workspace/SkyRL
+
+# Use an AMD dependency view that omits GPU-specific CUDA packages. The ROCm
+# variants of torch/vLLM/flash-attn come from the base image instead.
+RUN cp docker/pyproject.amd.toml pyproject.toml \
+    && uv pip install --system --no-cache-dir -e ".[fsdp,tinker]" \
+    && python -c "import ray, torch, vllm; print('ray', ray.__version__); print('vllm', vllm.__version__); print('torch', torch.__version__, 'hip', torch.version.hip)"
+
+ENTRYPOINT []
+CMD ["/bin/bash"]

docker/pyproject.amd.toml

@@ -0,0 +1,86 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+include = ["skyrl*"]
+
+[project]
+name = "skyrl"
+dynamic = ["version"]
+description = "Unified API for training and inference"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "datasets>=4.0.0",
+    "pillow>=11.3.0",
+    "rich>=14.1.0",
+    "safetensors>=0.6.2",
+    "tokenizers>=0.21.2",
+    "transformers>=5.6.1,<=5.8.0",
+    "typer>=0.17.4",
+    "peft==0.18.1",
+    "hf_transfer",
+    "cloudpathlib>=0.23.0",
+]
+
+[project.optional-dependencies]
+tinker = [
+    "tinker==0.16.1",
+    "fastapi[standard]",
+    "sqlmodel",
+    "sqlalchemy[asyncio]",
+    "aiosqlite",
+    "asyncpg",
+    "psycopg2-binary",
+]
+
+skyrl-train = [
+    "loguru",
+    "tqdm",
+    "ninja",
+    "tensorboard",
+    "func_timeout",
+    "hydra-core==1.3.2",
+    "accelerate",
+    "omegaconf",
+    "peft==0.18.1",
+    "debugpy==1.8.0",
+    "hf_transfer",
+    "wandb",
+    "datasets>=4.0.0",
+    "tensordict",
+    "jaxtyping",
+    "skyrl-gym",
+    "polars",
+    "s3fs",
+    "fastapi",
+    "uvicorn",
+    "vllm-router; sys_platform == 'linux'",
+    "pybind11",
+    "setuptools",
+]
+
+fsdp = [
+    "skyrl[skyrl-train]",
+]
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.dynamic]
+version = {attr = "skyrl.__version__"}
+
+[tool.uv]
+required-environments = [
+    "sys_platform == 'linux'",
+    "sys_platform == 'darwin' and platform_machine == 'arm64'",
+]
+
+[tool.uv.sources]
+skyrl-gym = { path = "./skyrl-gym", editable = true }
+
+vllm-router = [
+    { url = "https://github.com/SumanthRH/router/releases/download/0.1.14.post1/vllm_router-0.1.14-cp38-abi3-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
+    { url = "https://github.com/SumanthRH/router/releases/download/0.1.14.post1/vllm_router-0.1.14-cp38-abi3-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
+]

examples/train/amd/README.md

@@ -0,0 +1,147 @@
+# AMD ROCm Tinker Example
+
+This example is a starting point for running SkyRL's Tinker-compatible API on AMD GPUs with the FSDP backend and vLLM ROCm inference.
+
+The runtime path is intentionally split from the image path:
+
+- `docker/Dockerfile.amd` builds a SkyRL AMD image from `vllm/vllm-openai-rocm:v0.20.2`.
+- This directory contains commands to run inside that image.
+
+The Docker image bakes in Ray and the non-GPU SkyRL dependencies. It relies on the base vLLM ROCm image for ROCm builds of PyTorch, vLLM, and flash-attn.
+
+## Build
+
+From the repository root:
+
+```bash
+docker build -f docker/Dockerfile.amd -t skyrl-amd-rocm .
+```
+
+## Run The Container
+
+Use the ROCm devices and host IPC. `--network host` is convenient for Ray and for reaching the Tinker API from another shell.
+
+```bash
+docker run --rm -it \
+  --network host \
+  --ipc=host \
+  --device=/dev/kfd \
+  --device=/dev/dri \
+  --group-add video \
+  --cap-add=SYS_PTRACE \
+  --security-opt seccomp=unconfined \
+  skyrl-amd-rocm
+```
+
+## Start The Tinker Server
+
+Inside the container:
+
+```bash
+cd /workspace/SkyRL/examples/train/amd
+bash run_tinker_server_amd.sh
+```
+
+The default server binds to `0.0.0.0:9000` and uses:
+
+- `BASE_MODEL=Qwen/Qwen3-4B-Instruct-2507`
+- `BACKEND=fsdp`
+- `POLICY_NUM_GPUS_PER_NODE=1`
+- `INFERENCE_NUM_ENGINES=6`
+
+This default targets an 8-GPU AMD node with one GPU for the FSDP policy worker,
+six vLLM inference engines, and one GPU left for headroom. For smaller nodes or
+faster local debugging, reduce the inference engine count:
+
+```bash
+INFERENCE_NUM_ENGINES=1 \
+bash run_tinker_server_amd.sh
+```
+
+All server arguments can still be overridden through environment variables or by passing flags directly:
+
+```bash
+bash run_tinker_server_amd.sh --help
+```
+
+## Run The Client Smoke Test
+
+In a second shell inside the container:
+
+```bash
+cd /workspace/SkyRL/examples/train/amd
+TINKER_BASE_URL=http://localhost:9000 TINKER_API_KEY=tml-dummy \
+python tinker_hello_world.py
+```
+
+The client is a fixed Tinker smoke test. It creates a rank-32 LoRA training
+client, builds 16 tiny cross-entropy datums, samples once before training, runs
+four `forward_backward` + `optim_step` iterations, syncs trained weights to the
+sampler, samples once more, and prints `PASS` on success.
+
+## Run The GRPO Client
+
+For a fuller Tinker client, run the GRPO-style GSM8K example against the same
+server:
+
+```bash
+cd /workspace/SkyRL/examples/train/amd
+TINKER_BASE_URL=http://localhost:9000 TINKER_API_KEY=tml-dummy \
+python grpo_client.py
+```
+
+The GRPO client samples groups of responses, computes group-relative advantages
+from rule-based rewards, and trains a rank-32 LoRA policy with the public Tinker
+`ppo` loss. SkyRL maps this to its standard PPO-style policy loss internally. If
+`--data-dir` does not already contain
+`train.parquet` and `validation.parquet`, the client prepares a small GSM8K
+subset automatically under `/tmp/skyrl-tinker-grpo/gsm8k`. Prepared subsets are
+shuffled before truncation, and each training step samples a fresh random prompt
+batch from the loaded train pool.
+
+The default client run uses:
+
+- `--max-train-steps 5`
+- `--num-prompts 64`
+- `--group-size 8`
+- `--max-tokens 512`
+- `--max-train-examples 1024`
+- `--max-val-examples 128`
+
+To force regeneration of the shuffled GSM8K subset:
+
+```bash
+TINKER_BASE_URL=http://localhost:9000 TINKER_API_KEY=tml-dummy \
+python grpo_client.py \
+  --reprepare-data
+```
+
+For a faster single-step smoke:
+
+```bash
+TINKER_BASE_URL=http://localhost:9000 TINKER_API_KEY=tml-dummy \
+python grpo_client.py \
+  --max-train-steps 1 \
+  --num-prompts 2 \
+  --group-size 2 \
+  --max-tokens 64 \
+  --max-train-examples 32 \
+  --max-val-examples 8 \
+  --reprepare-data
+```
+
+## Multi-Node Note
+
+Ray is installed in the image so cluster setup can be handled by normal Ray commands or a platform launcher.
+
+Head node:
+
+```bash
+ray start --head --port=6379
+```
+
+Worker nodes:
+
+```bash
+ray start --address="$HEAD_NODE_IP:6379"
+```