diff --git a/.github/workflows/gpu-ci.yml b/.github/workflows/gpu-ci.yml index 3e2cb3e4..1863939f 100644 --- a/.github/workflows/gpu-ci.yml +++ b/.github/workflows/gpu-ci.yml @@ -48,9 +48,26 @@ jobs: chmod 600 ~/.ssh/id_ed25519 ssh-keygen -y -f ~/.ssh/id_ed25519 > /dev/null && echo "key OK" || echo "key BROKEN" + # Follow-up (#191): to test multiple architectures, add a matrix and pass + # GPU_ID + TARGET_SM (+ KERNEL_ALIGN_FORCE_SM90 for Hopper) through to the script. + # run_gpu_ci.sh reads all three, normalizes TARGET_SM, asserts the pod matches it, + # and forwards KERNEL_ALIGN_FORCE_SM90 into the remote build: + # strategy: + # matrix: + # include: + # - { gpu_id: "NVIDIA RTX A4000", target_sm: "8.6" } # Ampere + # - { gpu_id: "NVIDIA A100 80GB PCIe", target_sm: "8.0" } + # - { gpu_id: "NVIDIA H100 PCIe", target_sm: "9.0", force_sm90: "1" } # build Hopper TMA/WGMMA kernels + # - { gpu_id: "NVIDIA B200", target_sm: "10.0" } + # Per-arch jobs must NOT fall back to a different-capability GPU: the script + # fails fast when the pod arch != requested TARGET_SM, so keep fallback within + # the same compute capability (or unset it for these jobs). - name: Run GPU tests on RunPod env: RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} PR_REPO_URL: ${{ github.event.pull_request.head.repo.clone_url }} PR_SHA: ${{ github.event.pull_request.head.sha }} + # GPU_ID: ${{ matrix.gpu_id }} + # TARGET_SM: ${{ matrix.target_sm }} + # KERNEL_ALIGN_FORCE_SM90: ${{ matrix.force_sm90 }} run: bash ci/run_gpu_ci.sh diff --git a/ci/run_gpu_ci.sh b/ci/run_gpu_ci.sh index a521780e..54eb56e6 100644 --- a/ci/run_gpu_ci.sh +++ b/ci/run_gpu_ci.sh @@ -1,13 +1,20 @@ #!/usr/bin/env bash set -uo pipefail -# TP=2 -PRIMARY_GPU_ID="NVIDIA RTX A4000" -PRIMARY_GPU_COUNT=2 +# TP=2 (override via env; GPU_ID/GPU_COUNT accepted as matrix-friendly aliases) +PRIMARY_GPU_ID="${PRIMARY_GPU_ID:-${GPU_ID:-NVIDIA RTX A4000}}" +PRIMARY_GPU_COUNT="${PRIMARY_GPU_COUNT:-${GPU_COUNT:-2}}" # TP=1 -FALLBACK_GPU_ID="NVIDIA A40" -FALLBACK_GPU_COUNT=1 +FALLBACK_GPU_ID="${FALLBACK_GPU_ID:-NVIDIA A40}" +FALLBACK_GPU_COUNT="${FALLBACK_GPU_COUNT:-1}" + +# Optional arch override; asserted against the pod's real cap in the remote build so a +# cross-arch resource fallback cannot build mismatched SASS. +TARGET_SM="${TARGET_SM:-}" + +# Forwarded to the remote build; setup.py compiles the Hopper (sm90) kernels only when "1". +KERNEL_ALIGN_FORCE_SM90="${KERNEL_ALIGN_FORCE_SM90:-}" CI_IMAGE="${CI_IMAGE:-runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04}" DISK_GB=40 @@ -127,17 +134,64 @@ if ! "$PY" -c "import torch" >/dev/null 2>&1; then done fi echo "[remote] Using interpreter: $PY" -export TORCH_CUDA_ARCH_LIST=8.6 export FORCE_CUDA=1 export MAX_JOBS=8 +export KERNEL_ALIGN_FORCE_SM90="'"${KERNEL_ALIGN_FORCE_SM90}"'" + +# normalize_sm: compact (90) or dotted (9.0) compute cap -> torch dotted form, keeping +PTX. +normalize_sm() { + sm_in="$1"; sm_ptx="" + case "$sm_in" in *+PTX) sm_ptx="+PTX"; sm_in="${sm_in%+PTX}";; esac + case "$sm_in" in + *.*) : ;; + [0-9][0-9]|[0-9][0-9][0-9]) sm_major="${sm_in%?}"; sm_in="${sm_major}.${sm_in#$sm_major}" ;; + *) return 1 ;; + esac + case "$sm_in" in [0-9]*.[0-9]|[0-9]*.[0-9][0-9]) echo "${sm_in}${sm_ptx}" ;; *) return 1 ;; esac +} + +ACTUAL_SM=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1 | tr -d "[:space:]") +[ -z "$ACTUAL_SM" ] && ACTUAL_SM=$("$PY" -c "import torch;a,b=torch.cuda.get_device_capability();print(f\"{a}.{b}\")" 2>/dev/null || true) +[ -z "$ACTUAL_SM" ] && { echo "[remote] FATAL: cannot determine GPU compute capability"; exit 3; } + +REQUESTED_SM="'"${TARGET_SM}"'" +if [ -n "$REQUESTED_SM" ]; then + NORM_REQ=$(normalize_sm "$REQUESTED_SM") || { echo "[remote] FATAL: unsupported TARGET_SM=$REQUESTED_SM"; exit 3; } + NORM_REQ_BASE="${NORM_REQ%+PTX}" + if [ "$NORM_REQ_BASE" != "$ACTUAL_SM" ]; then + echo "[remote] FATAL: requested TARGET_SM=$REQUESTED_SM (sm_$NORM_REQ_BASE) but provisioned GPU is sm_$ACTUAL_SM." + echo "[remote] Refusing to build mismatched kernels (likely a cross-arch resource fallback)." + exit 3 + fi + BUILD_SM="$NORM_REQ_BASE" +else + BUILD_SM=$(normalize_sm "$ACTUAL_SM") || { echo "[remote] FATAL: unsupported detected arch $ACTUAL_SM"; exit 3; } +fi +# BUILD_SM is always bare here (both paths strip +PTX); +PTX gives forward-compat JIT. +export TORCH_CUDA_ARCH_LIST="${BUILD_SM}+PTX" +echo "[remote] Detected GPU sm_$ACTUAL_SM; building _C for TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST" + cd /workspace git clone '"${PR_REPO_URL:-https://github.com/RL-Align/RL-Kernel.git}"' repo cd repo git fetch origin '"${PR_SHA}"' git checkout --detach '"${PR_SHA}"' -"$PY" -m pip install -e . -"$PY" -m pip install pytest +"$PY" -c "import torch;print(f\"[remote] image torch {torch.__version__} cuda {torch.version.cuda}\")" +# Pin torch (cu124, matching the CI image) so the extension is built against the exact +# runtime torch, not a non-deterministic bare-install upgrade of the 2.4.0 in the image. +TORCH_SPEC="${TORCH_SPEC:-torch==2.4.1}" +TORCH_INDEX_URL="${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu124}" +"$PY" -m pip install --no-cache-dir "$TORCH_SPEC" --index-url "$TORCH_INDEX_URL" +"$PY" -c "import torch;print(f\"[remote] pinned torch {torch.__version__} cuda {torch.version.cuda}\")" +# --no-build-isolation: torch must be visible to setup.py, else the extension is silently skipped. +# --no-deps: keep the pinned torch; do not let the editable install re-resolve it. +"$PY" -m pip install --no-build-isolation --no-deps -e . +"$PY" -m pip install --no-cache-dir numpy tabulate accelerate transformers pytest nvidia-smi +# Fail fast if _C did not build or cannot launch, instead of silently using native fallbacks. +"$PY" scripts/ci_smoke.py +# Enforce _C in the pytest suite too (test_extension_smoke.py skips unless this is set). +export RL_KERNEL_REQUIRE_EXT=1 '"${TEST_CMD}" echo "[ci] Launching remote test suite on GPU pod (Distributed Execution Mode: TP=${GPU_COUNT})..." diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md index d4563517..71fe227d 100644 --- a/docs/getting_started/installation.md +++ b/docs/getting_started/installation.md @@ -5,24 +5,44 @@ CUDA toolchain; ROCm builds require a compatible ROCm environment. ## From Source +For CUDA (or ROCm) source builds, install PyTorch first (matching your CUDA/ROCm +runtime), then build with build isolation disabled so `setup.py` can access +PyTorch's extension build utilities and compile the native kernels (`rl_engine._C`): + ```bash git clone https://github.com/RL-Align/RL-Kernel.git cd RL-Kernel -pip install -e . +# Optional: pin the compile target. If unset, the build targets your GPU's arch. +# export TORCH_CUDA_ARCH_LIST="9.0+PTX" # e.g. Hopper; or "8.6+PTX", "12.0+PTX" +pip install --no-build-isolation -e . ``` +Without `--no-build-isolation`, PyTorch is invisible to the isolated build +environment, the extension is silently skipped, and the library falls back to the +slower pure-PyTorch kernels. Confirm the compiled extension is present with: + +```bash +python -c "from rl_engine import _C; print('compiled extension OK')" +``` + +A CPU-only install (plain `pip install -e .` on a machine with no GPU) remains +supported and runs on the pure-PyTorch backends. + ## Optional Backends +The extras add optional dependencies on top of the compiled package, so they use +the same `--no-build-isolation` flag as the source build above. + ```bash -pip install -e ".[cuda]" +pip install --no-build-isolation -e ".[cuda]" ``` ```bash -pip install -e ".[rocm]" +pip install --no-build-isolation -e ".[rocm]" ``` ```bash -pip install -e ".[vllm]" +pip install --no-build-isolation -e ".[vllm]" ``` Install the vLLM extra only on rollout or benchmark environments that need the diff --git a/scripts/ci_smoke.py b/scripts/ci_smoke.py new file mode 100644 index 00000000..0cbe1cda --- /dev/null +++ b/scripts/ci_smoke.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""CI fail-fast smoke check for the compiled CUDA extension (``rl_engine._C``). + +Exits non-zero with a clear message when either failure mode from issue #191 +occurs, so GPU CI cannot pass while silently running only native fallbacks: + + * Bug A - the extension did not build at all (e.g. PEP 517 build isolation hid + torch from ``setup.py``), so ``from rl_engine import _C`` raises ImportError. + * Bug B - the extension built for the wrong GPU architecture; the import + succeeds (``dlopen`` does not check arch) but the first kernel launch raises + ``cudaErrorNoKernelImageForDevice`` once the stream is synchronized. + +Uses only ``fused_logp`` - the op registered unconditionally in ``csrc/ops.cpp`` - +so it does not require ``KERNEL_ALIGN_FORCE_SM90=1`` / a Hopper build. +""" +import sys + +import torch + + +def main() -> int: + if not torch.cuda.is_available(): + print("[smoke] FATAL: CUDA is not available in this CI environment", file=sys.stderr) + return 2 + + print(f"[smoke] torch: {torch.__version__} (cuda {torch.version.cuda})") + print(f"[smoke] device: {torch.cuda.get_device_name()}") + cc = torch.cuda.get_device_capability() + print(f"[smoke] capability: sm_{cc[0]}{cc[1]}") + + # (1) Import check -> catches Bug A (no .so was compiled at all). + try: + from rl_engine import _C + except ImportError as exc: + print( + "[smoke] FATAL: compiled extension rl_engine._C is missing - the CUDA " + "kernels were not built.\n" + " Likely PEP 517 build isolation hid torch from setup.py; install " + "with `pip install --no-build-isolation -e .`.\n" + f" Underlying error: {exc}", + file=sys.stderr, + ) + return 1 + print(f"[smoke] _C file: {getattr(_C, '__file__', None)}") + + # (2) Real launch + synchronize -> catches Bug B (arch mismatch only surfaces + # on launch, asynchronously, so the sync is required to raise it here). + try: + logits = torch.randn(4, 32, device="cuda", dtype=torch.float32) + token_ids = torch.randint(0, 32, (4,), device="cuda", dtype=torch.long) + out = _C.fused_logp(logits, token_ids) + torch.cuda.synchronize() + # Broad on purpose: an arch mismatch raises a CUDA RuntimeError (not ImportError), + # and any launch failure whatsoever must fail the smoke check loudly. + except Exception as exc: + print( + "[smoke] FATAL: rl_engine._C built but fused_logp failed to launch on " + f"sm_{cc[0]}{cc[1]}.\n" + " The extension was likely compiled for a different architecture; " + "set TORCH_CUDA_ARCH_LIST / TARGET_SM to match this GPU.\n" + f" Underlying error: {type(exc).__name__}: {exc}", + file=sys.stderr, + ) + return 1 + + if tuple(out.shape) != (4,): + print( + f"[smoke] FATAL: unexpected fused_logp output shape {tuple(out.shape)}", file=sys.stderr + ) + return 1 + + print(f"[smoke] OK: rl_engine._C built and fused_logp ran on sm_{cc[0]}{cc[1]}.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_extension_smoke.py b/tests/test_extension_smoke.py new file mode 100644 index 00000000..5c81a733 --- /dev/null +++ b/tests/test_extension_smoke.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Regression guard for issue #191. + +On a GPU host the compiled extension ``rl_engine._C`` MUST be present and +launchable; a missing or arch-mismatched build must fail loudly here instead of +silently degrading to the pure-PyTorch fallbacks. On a CPU host (no CUDA) the +test skips, so the CPU CI job - which legitimately has no ``_C`` - stays green. + +Enforcement is opt-in via ``RL_KERNEL_REQUIRE_EXT=1``, which the GPU CI +orchestrator sets right before pytest (after it has built the extension). That +way the check does not fail in environments not expected to have ``_C`` compiled - +a plain ``pytest`` run, or CI running an older orchestrator that has not built it +yet (the compiled kernels are always enforced separately by scripts/ci_smoke.py). +""" +import os + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA GPU") +@pytest.mark.skipif( + os.environ.get("RL_KERNEL_REQUIRE_EXT") != "1", + reason="extension enforcement is CI-only (orchestrator sets RL_KERNEL_REQUIRE_EXT=1)", +) +def test_compiled_extension_present_and_launches(): + # Import directly from the package, NOT from rl_engine.kernels.ops.base, which + # deliberately swallows the ImportError and falls back to _C = None. + try: + from rl_engine import _C + except ImportError as exc: # Bug A: the extension was never built + pytest.fail( + "rl_engine._C is missing on a GPU host - the CUDA extension was not " + "built. Install with `pip install --no-build-isolation -e .`. " + f"Underlying error: {exc}" + ) + + logits = torch.randn(4, 32, device="cuda", dtype=torch.float32) + token_ids = torch.randint(0, 32, (4,), device="cuda", dtype=torch.long) + out = _C.fused_logp(logits, token_ids) + torch.cuda.synchronize() # Bug B: an arch mismatch surfaces on synchronize + assert tuple(out.shape) == (4,)