diff --git a/docker/Dockerfile.gb10 b/docker/Dockerfile.gb10 new file mode 100644 index 0000000000..471e2ce6f1 --- /dev/null +++ b/docker/Dockerfile.gb10 @@ -0,0 +1,208 @@ +# Dockerfile for NVIDIA DGX Spark (GB10 / sm_121a, aarch64). +# +# slime's stock Dockerfile targets x86_64 + H100/B200/GB200 and does not work +# on GB10 (ptxas in CUDA 12.9 lacks sm_121a; published sgl_kernel wheels have +# no sm_121 variant and ABI-mismatch against NGC libtorch). +# +# This Dockerfile rebases on the NGC vLLM container which ships CUDA 13.2, +# PyTorch 2.11.0a0 (compiled with compute_120 for Blackwell PTX forward-compat), +# Triton 3.6.0, and flash-attn 2.7.4.post1 — all verified working on GB10. +# +# Digest-pinned for reproducibility. Update the digest together with the tag. +# Companion notes: NOTES_GB10.md (15 resolved blockers documented). + +ARG BASE_IMAGE=nvcr.io/nvidia/vllm:26.03-py3 +ARG BASE_DIGEST=sha256:13e327dad79e6e417f6687fec2ba76b0386d597082ec0ee003c1e964ec6ad0e7 +FROM ${BASE_IMAGE}@${BASE_DIGEST} + +# ============================================================================= +# Common env: GB10 reports cc 12.1 but NGC torch only lists compute_120. PTX +# forward-compat carries sm_120 PTX → sm_121 at runtime. Set the same for all +# downstream compilation layers. +# ============================================================================= +ENV TORCH_CUDA_ARCH_LIST="12.0+PTX" + +# ============================================================================= +# System packages missing from the NGC base but needed downstream: +# - libz3-dev : tilelang's Z3 SMT autoscheduler dlopens libz3.so +# ============================================================================= +RUN apt-get update && apt-get install -y --no-install-recommends \ + libz3-dev \ + && rm -rf /var/lib/apt/lists/* + +# ============================================================================= +# Header shims for CUDA 13 / NVTX removal +# ============================================================================= +# (1) NVTX3: CUDA 13 dropped bundled NVTX headers; the pypi placeholder is empty. +# Clone the upstream NVIDIA/NVTX repo and install just the C include tree. +RUN git clone --depth 1 https://github.com/NVIDIA/NVTX.git /tmp/nvtx_src \ + && cp -r /tmp/nvtx_src/c/include/nvtx3 /usr/local/cuda/include/nvtx3 \ + && rm -rf /tmp/nvtx_src + +# (2) cuda_profiler_api.h: CUDA 13 removed the public header; libcudart.so.13 +# still exports the two symbols. TE 2.10 only #includes the header without +# actually calling the APIs in the affected TUs, so a tiny shim is sufficient. +COPY docker/patch/gb10/cuda_profiler_api.h /usr/local/cuda/include/cuda_profiler_api.h + +# ============================================================================= +# cuDNN engines_precompiled library +# ============================================================================= +# NGC deliberately ships only runtime-compiled cuDNN libs to save space, but +# TE's bundled cudnn-frontend requires CUDNN::cudnn_engines_precompiled. The +# pypi wheel nvidia-cudnn-cu13 9.20.0.48 contains the missing .so. +RUN PIP_CONSTRAINT= pip install --no-deps nvidia-cudnn-cu13==9.20.0.48 \ + && ln -sf /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/lib/libcudnn_engines_precompiled.so.9 \ + /usr/lib/aarch64-linux-gnu/libcudnn_engines_precompiled.so.9 \ + && ln -sf /usr/lib/aarch64-linux-gnu/libcudnn_engines_precompiled.so.9 \ + /usr/lib/aarch64-linux-gnu/libcudnn_engines_precompiled.so + +# ============================================================================= +# CMake 4.3.1: needed for CUDA 13 family-specific `f` arch suffix (120f/121f). +# NGC's /etc/pip/constraint.txt pins cmake==3.31.6; we must clear PIP_CONSTRAINT +# for this install. CMAKE_POLICY_VERSION_MINIMUM=3.5 lets CMake 4.x accept +# bundled third-party CMakeLists that declare cmake_minimum_required < 3.5. +# ============================================================================= +RUN PIP_CONSTRAINT= pip install --no-deps --ignore-installed --force-reinstall cmake==4.3.1 +ENV CMAKE_POLICY_VERSION_MINIMUM=3.5 + +# ============================================================================= +# Build prerequisites for the source installs below +# ============================================================================= +RUN pip install --no-deps pybind11 pathspec pyproject_metadata scikit-build-core + +# ============================================================================= +# sgl-kernel for sm_121a (SGLang's CUDA kernel library) +# ============================================================================= +# The published sgl_kernel wheels are cu12x (libnvrtc.so.12 ABI), plus the cu130 +# arm64 wheel only has the sm_100 variant, not sm_121. Also ABI-mismatch against +# NGC torch. Build from source with the GB10 arch whitelist patch. +# +# Patch: docker/patch/gb10/patch_sgl_kernel.py adds SGL_KERNEL_GB10_ONLY CMake +# option that restricts gencode to sm_120a + sm_121a (skipping sm_90a/sm_100a/ +# sm_103a/sm_110a). This is essential — compiling all seven archs simultaneously +# OOM-kills cicc on 128 GB Spark hosts. +ARG SGLANG_COMMIT=v0.5.9 +RUN git clone --depth 1 --branch ${SGLANG_COMMIT} https://github.com/sgl-project/sglang.git /root/src/sglang +COPY docker/patch/gb10/patch_sgl_kernel.py /tmp/patch_sgl_kernel.py +RUN python3 /tmp/patch_sgl_kernel.py /root/src/sglang/sgl-kernel/CMakeLists.txt \ + && rm /tmp/patch_sgl_kernel.py + +# Conservative parallelism (8 parallel × 2 nvcc-internal threads = 16 active) +# keeps peak memory ~25 GB on the 128 GB Spark. ~25 minutes on 20-core GB10. +RUN cd /root/src/sglang/sgl-kernel \ + && CMAKE_BUILD_PARALLEL_LEVEL=8 \ + CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=2 -DSGL_KERNEL_GB10_ONLY=ON" \ + pip wheel . --no-build-isolation --no-deps -w /tmp/wheels \ + && pip install --no-deps /tmp/wheels/sgl_kernel-*.whl \ + && rm -rf /tmp/wheels /root/.cache/pip + +# ============================================================================= +# SGLang runtime + the small pure-python deps it needs. +# --no-deps preserves NGC torch / transformers / flashinfer / quack-kernels / +# cutlass-dsl / xgrammar / outlines that were carefully pinned in the NGC base. +# Pulling sglang's full dep set here would downgrade torch to 2.9.1+cu129 and +# break everything. +# ============================================================================= +RUN pip install --no-deps \ + "sglang==0.5.9" \ + "torch_memory_saver==0.0.9" \ + IPython traitlets stack_data executing asttokens pure_eval prompt_toolkit wcwidth \ + orjson hf_transfer modelscope litellm + +# SGLang router: zhuzilin's fork ships x86_64-only. Upstream 0.3.2 has an arm64 +# wheel that satisfies slime's version-parse comparisons (`<=0.2.1`, `<0.3.0`). +# Only loss: the `'slime' in version` check in wandb_utils.py won't match, +# which is a non-critical logging branch. +RUN pip install --no-deps sglang-router==0.3.2 + +# ============================================================================= +# TransformerEngine 2.10 for sm_121f +# ============================================================================= +# GB10 needs the CUDA 13 "f" (family-specific) arch suffix — TE's ptx.cuh has +# static_asserts that reject plain sm_120 / sm_121 in favor of sm_120f / sm_121f. +# CMake 4 is required to pass this through CMAKE_CUDA_ARCHITECTURES. +RUN TORCH_CUDA_ARCH_LIST="12.0+PTX" \ + NVTE_CUDA_ARCHS="120f;121f" \ + NVTE_FRAMEWORK=pytorch \ + MAX_JOBS=8 CMAKE_BUILD_PARALLEL_LEVEL=8 \ + pip install --no-build-isolation --no-deps \ + "git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.10" + +# TE pytorch subpackage pulls onnxscript → onnx_ir → onnx → ml_dtypes at import +# time (not setup time). Install the runtime chain so `import transformer_engine.pytorch` works. +RUN pip install --no-deps onnxscript onnx_ir onnx ml_dtypes + +# ============================================================================= +# apex (NVIDIA fused optimizers / norms) +# ============================================================================= +# Pinned commit matches slime upstream Dockerfile. +RUN NVCC_APPEND_FLAGS="--threads 4" MAX_JOBS=8 \ + pip install --no-cache-dir --no-build-isolation --no-deps \ + --config-settings "--build-option=--cpp_ext --cuda_ext --parallel 8" \ + "git+https://github.com/NVIDIA/apex.git@10417aceddd7d5d05d7cbf7b0fc2daad1105f8b4" + +# ============================================================================= +# Megatron-LM (+ slime's megatron.patch) +# ============================================================================= +# slime pins Megatron commit 3714d81d... +ARG MEGATRON_COMMIT=3714d81d418c9f1bca4594fc35f9e8289f652862 +RUN git clone https://github.com/NVIDIA/Megatron-LM.git /root/src/Megatron-LM \ + && cd /root/src/Megatron-LM \ + && git checkout ${MEGATRON_COMMIT} +COPY docker/patch/latest/megatron.patch /tmp/megatron.patch +RUN cd /root/src/Megatron-LM \ + && git apply /tmp/megatron.patch --3way \ + && ! grep -R -n '^<<<<<<< ' . \ + && rm /tmp/megatron.patch \ + && pip install --no-deps --no-build-isolation -e . + +# Megatron.training / .rl / .legacy are sibling dirs meant to be on PYTHONPATH +# (slime's own docs confirm this). +ENV PYTHONPATH="/root/src/Megatron-LM" + +# ============================================================================= +# Remaining slime deps: bridges, modelopt, FLA, tilelang, antlr, omegaconf chain +# ============================================================================= +RUN pip install --no-deps \ + "git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c" \ + flash-linear-attention==0.4.1 \ + tilelang==0.1.8 + +RUN pip install --no-deps --no-build-isolation \ + "git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl" \ + "nvidia-modelopt>=0.37.0" + +# antlr4-python3-runtime must be pinned to 4.9.3 — omegaconf's bundled grammar +# was generated with antlr 4.9 serialized format and runtime 4.13 breaks on it +# with "Could not deserialize ATN with version 3 (expected 4)". +RUN pip install --no-deps "antlr4-python3-runtime==4.9.3" omegaconf wandb + +# Megatron requires numpy 1.x; NGC ships 2.x. +RUN PIP_CONSTRAINT= pip install --no-deps --force-reinstall "numpy<2" + +# Reward function runtime deps (grade_answer_verl / deepscaler math verifier) +RUN pip install --no-deps pylatexenc math_verify word2number + +# Eval / data-prep runtime deps +RUN pip install --no-deps pyarrow accelerate + +# ============================================================================= +# slime itself +# ============================================================================= +COPY . /root/slime +RUN cd /root/slime \ + && pip install --no-deps --no-build-isolation -e . + +# slime's int4 QAT kernel (fused per-channel fake-quant + dequant) +RUN cd /root/slime/slime/backends/megatron_utils/kernels/int4_qat \ + && pip install --no-build-isolation --no-deps . + +WORKDIR /root/slime + +# Sanity: load the full arg parser on container build +RUN python train.py --help > /tmp/help.txt \ + && test $(wc -l < /tmp/help.txt) -gt 3000 \ + && echo "slime train.py --help OK ($(wc -l < /tmp/help.txt) lines)" \ + && rm /tmp/help.txt + +CMD ["/bin/bash"] diff --git a/docker/NOTES_GB10.md b/docker/NOTES_GB10.md new file mode 100644 index 0000000000..d5417321f5 --- /dev/null +++ b/docker/NOTES_GB10.md @@ -0,0 +1,99 @@ +# slime on NVIDIA DGX Spark (GB10) — Porting Notes + +Target platform: +- Chip: NVIDIA GB10 (Grace + consumer Blackwell, SM 12.1 / sm_121a) +- Arch: aarch64 +- OS: Ubuntu 24.04, NVIDIA driver 580.142 (CUDA 13.x forward-compat) +- Unified memory: 128 GB (CPU+GPU shared) + +## Why a new Dockerfile is needed + +slime's published Docker images (`slimerl/slime:*`) are x86_64-only. The arm64 base +it derives from (`slimerl/sglang:v0.5.9`) ships with CUDA 12.9 and stock PyTorch +2.9.1+cu129 — neither of which knows about the `sm_121a` target used by GB10: +the Triton PTX pipeline crashes with `ptxas fatal: Value 'sm_121a' is not defined`. + +The `ENABLE_CUDA_13=1` branch in the upstream Dockerfile is aimed at GB200/GB300 +(sm_100a) and still uses an x86-only router wheel and amd64 base, so it does not +directly apply to GB10. + +This port rebases slime on `nvcr.io/nvidia/vllm:26.03-py3` (arm64), which ships: +- CUDA 13.2 (ptxas understands sm_121a ✅) +- PyTorch 2.11.0a0 compiled with `compute_120` (PTX forward-compat to sm_121 ✅) +- Triton 3.6.0 (verified: Triton kernel JITs on GB10 ✅) +- flash-attn 2.7.4.post1 preinstalled + +### Base image pinning (for reproducibility) + +Pull from NGC and verify digest: + +```bash +docker pull nvcr.io/nvidia/vllm:26.03-py3 +docker inspect nvcr.io/nvidia/vllm:26.03-py3 --format '{{range .RepoDigests}}{{.}}{{end}}' +# Expected digest: +# nvcr.io/nvidia/vllm@sha256:13e327dad79e6e417f6687fec2ba76b0386d597082ec0ee003c1e964ec6ad0e7 +``` + +All downstream steps pin this digest. Product page: +https://catalog.ngc.nvidia.com/orgs/nvidia/containers/vllm?version=26.03-py3 + +## Known blockers and their resolutions + +| # | Blocker | Root cause | Resolution | +|---|---------|-----------|------------| +| 1 | `ptxas fatal: sm_121a` | CUDA 12.9 ptxas predates sm_121a (added in CUDA 13.0) | Rebase on NGC CUDA 13.2 image | +| 2 | `libnvrtc.so.12` missing from sgl_kernel wheel | Published sgl_kernel wheel is cu12x | Use cu130 wheel OR build from source | +| 3 | `sgl_kernel/sm100/...abi3.so: undefined symbol _ZN3c104cuda29c10_cuda_check_implementationEiPKcS2_ib` | Wheel built against stock libtorch, NGC libtorch has different C++ ABI | Build sgl_kernel from source against NGC torch | +| 4 | sgl_kernel arm64 wheels only build sm_100 variant | sgl-project CI doesn't target GB10 | Build sm_121 variant from source | +| 5 | sgl-kernel source build: cicc OOM on sm_90+100+103+110+120+121 concurrent compile | Cutlass FP8 templates × 7 arches × 12 parallel → >128GB | `SGL_KERNEL_GB10_ONLY=ON` CMake option (see `docker/patch/gb10/sgl-kernel-arch.patch`), drops to sm_120a+121a | +| 6 | TE 2.10 build: `CUDNN::cudnn_engines_precompiled` target not found | NGC vLLM image ships only runtime-compiled cuDNN libs | Install `nvidia-cudnn-cu13==9.20.0.48` pypi wheel (has the precompiled engines), symlink into `/usr/lib/aarch64-linux-gnu/` | +| 7 | TE 2.10 build: `nvtx3/nvToolsExt.h: No such file` | CUDA 13 dropped bundled NVTX headers; `nvidia-nvtx-cu13` pypi wheel is an empty 0.0.1 placeholder | Clone `NVIDIA/NVTX` github, copy `c/include/nvtx3/*` to `/usr/local/cuda/include/nvtx3/` | +| 8 | TE 2.10 build: `ptx.cuh` static_assert "Compile for smXXXf instead of smXXX" | Blackwell family-specific TMA features require CUDA 13's new `f` (family-specific) arch suffix, not plain integer | Set `NVTE_CUDA_ARCHS="120f;121f"` — but CMake ≤3.31 rejects `f` suffix | +| 9 | CMake 3.31 rejects `CMAKE_CUDA_ARCHITECTURES=120f;121f` | `f` suffix for CUDA 13 Blackwell family is only supported in CMake ≥4.0 | Upgrade to `cmake==4.3.1` via pip (must override NGC `/etc/pip/constraint.txt` with `PIP_CONSTRAINT=`) and set `CMAKE_POLICY_VERSION_MINIMUM=3.5` for old bundled deps | +| 10 | TE 2.10 build: `cuda_profiler_api.h: No such file` | CUDA 13 removed the public header for `cudaProfilerStart/Stop`; the symbols still exist in `libcudart.so.13`. TE's 3 `fused_softmax` TUs `#include` the header but don't call the APIs | Install a 20-line shim header at `/usr/local/cuda/include/cuda_profiler_api.h` declaring the two functions extern. Stored as `docker/patch/gb10/cuda_profiler_api.h` | +| 11 | slime `train.py --help`: `'tuple' object has no attribute 'strip'` | Typo in `slime/utils/arguments.py:1073`: `help=("string",)` (trailing comma → tuple) instead of `help=("string")` | Remove trailing comma — simple one-line slime fix, upstream-able | +| 12 | `sglang_router` x86_64-only wheel from `zhuzilin/sgl-router` fork | slime Dockerfile pins `zhuzilin/sgl-router` release (no arm64 builds); slime's `'slime' in version` assertion is only in `wandb_utils.py` (non-critical path) | Install upstream `sglang-router==0.3.2` from PyPI (has arm64 wheel). Accept wandb path fallback | +| 13 | `antlr4-python3-runtime==4.13.2` → `Could not deserialize ATN with version 3` | Omegaconf's bundled grammar was generated with antlr 4.9 serialized format; runtime 4.13 only reads format v4 | Pin `antlr4-python3-runtime==4.9.3` | +| 14 | `megatron.training` not importable after `pip install -e Megatron-LM` | Megatron-LM's setup.py only packages `megatron-core`; `megatron.training`, `megatron.rl`, `megatron.legacy` are sibling dirs meant to be on `PYTHONPATH` | `export PYTHONPATH=/root/src/Megatron-LM:$PYTHONPATH` (slime docs confirm this) | +| 15 | `libz3.so` missing for tilelang | tilelang uses Z3 SMT solver for autoscheduling; NGC vllm base doesn't include libz3 | `apt-get update && apt-get install -y libz3-dev` (libz3-4 package alias needs update first) | + +## Build journal + +### sgl-kernel OOM during build → arch whitelist patch + +First build attempt on GB10 with stock sgl-kernel v0.5.9 CMake: cicc +(NVIDIA's CUDA frontend compiler) was OOM-killed while compiling +`csrc/gemm/fp8_gemm_kernel.cu`, `fp8_blockwise_gemm_kernel.cu`, and +`nvfp4_scaled_mm_kernels.cu`. + +Root cause: for `CUDA_VERSION >= 13.0 && aarch64`, sgl-kernel unconditionally +emits seven gencodes per TU — `sm_90, sm_90a, sm_100a, sm_103a, sm_110a, +sm_120a, sm_121a`. Cutlass FP8/NVFP4 gemm template instantiation uses +~10–15 GB of RAM per TU per arch. Combined with 12 parallel nvcc jobs, peak +memory exceeded the 128 GB unified memory limit. + +Only `sm_121a` (and `sm_120a` as PTX fallback) actually runs on GB10. Hopper +(sm_90a) and datacenter Blackwell (sm_100a/sm_103a) binaries are dead weight. + +Fix: `docker/patch/gb10/sgl-kernel-arch.patch` adds a CMake option +`SGL_KERNEL_GB10_ONLY`. When set, the other gencodes and FA3 (sm_90a-only) +are skipped. Default OFF preserves upstream behavior. See the patch python +script `docker/patch/gb10/patch_sgl_kernel.py` for a programmatic apply. + +### sgl-kernel build success (M1) + +With `SGL_KERNEL_GB10_ONLY=ON`, `CMAKE_BUILD_PARALLEL_LEVEL=8`, and +`SGL_KERNEL_COMPILE_THREADS=2`, the build completed in 24 minutes with RAM +never exceeding ~25 GB. Produced `sgl_kernel-0.3.21-cp310-abi3-linux_aarch64.whl` +(74 MB), ABI-compatible with NGC libtorch 2.11.0a0 (verified via clean import +and functional rmsnorm kernel on GB10). + +## Work items + +- [x] M1: Build sgl_kernel from source for sm_121 + NGC torch ABI +- [x] M2: TransformerEngine 2.10 (sm_121f), apex, Megatron-LM all built + verified +- [x] M3: slime installed with megatron.patch applied; `train.py --help` prints 3714-line arg list +- [x] M4: End-to-end RL loop (Qwen2.5-0.5B + dapo-math-17k + GRPO colocated on 1 GB10 GPU) runs clean in 2m10s + +Reproducible image: `slime-gb10:m4-success` (36.3 GB, committed post-smoke) +- Step 0 train metrics logged; full rollout → reward → policy-update → weight-sync cycle confirmed. diff --git a/docker/patch/gb10/cuda_profiler_api.h b/docker/patch/gb10/cuda_profiler_api.h new file mode 100644 index 0000000000..276b80b656 --- /dev/null +++ b/docker/patch/gb10/cuda_profiler_api.h @@ -0,0 +1,26 @@ +#ifndef __CUDA_PROFILER_API_H__ +#define __CUDA_PROFILER_API_H__ + +/* CUDA 13 dropped the public cuda_profiler_api.h header, but the runtime still + * exports cudaProfilerStart / cudaProfilerStop from libcudart.so. This shim + * declares them against the public types so legacy callers (e.g., TE 2.10, + * which only includes the header without actually calling the APIs in the + * affected TUs) continue to compile. + * + * Installed by slime GB10 port. See docker/patch/gb10/ and NOTES_GB10.md. + */ + +#include + +#if defined(__cplusplus) +extern "C" { +#endif + +extern __host__ cudaError_t CUDARTAPI cudaProfilerStart(void); +extern __host__ cudaError_t CUDARTAPI cudaProfilerStop(void); + +#if defined(__cplusplus) +} +#endif + +#endif /* __CUDA_PROFILER_API_H__ */ diff --git a/docker/patch/gb10/patch_sgl_kernel.py b/docker/patch/gb10/patch_sgl_kernel.py new file mode 100644 index 0000000000..75454c906e --- /dev/null +++ b/docker/patch/gb10/patch_sgl_kernel.py @@ -0,0 +1,51 @@ +"""Patch sgl-kernel CMakeLists.txt to add a SGL_KERNEL_GB10_ONLY build option. + +On GB10 (DGX Spark, sm_121a aarch64 with unified 128 GB memory), the stock +sgl-kernel CMake emits gencodes for sm_90a + sm_100a + sm_103a + sm_110a + +sm_120a + sm_121a. Cutlass FP8 gemm template instantiation per extra arch uses +10-15 GB RAM, and 12 parallel nvcc jobs OOM-kill cicc. + +This patch wraps the non-GB10 gencode blocks in `if (NOT SGL_KERNEL_GB10_ONLY)` +and adds a GB10-only branch that emits just sm_120a + sm_121a. Default OFF +preserves upstream behavior. +""" +import sys +from pathlib import Path + +p = Path(sys.argv[1] if len(sys.argv) > 1 else "/root/src/sglang/sgl-kernel/CMakeLists.txt") +s = p.read_text() + +marker = 'option(SGL_KERNEL_ENABLE_SM100A "Enable SM100A" OFF)' +new_opt = ( + marker + + "\n" + + 'option(SGL_KERNEL_GB10_ONLY "Build only for GB10 (sm_121a + sm_120a). Skips sm_90a/sm_100a/sm_103a/sm_110a and FA3." OFF)' +) +assert marker in s, "marker 1 (SM100A option line) not found" +s = s.replace(marker, new_opt, 1) + +old = 'if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A)' +assert old in s, "marker 2 (CUDA 12.8 block) not found" +s = s.replace(old, "if (NOT SGL_KERNEL_GB10_ONLY)\n\n" + old, 1) + +old2 = 'if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_FP4)' +assert old2 in s, "marker 3 (FP4 block) not found" +insert_before = """endif() # NOT SGL_KERNEL_GB10_ONLY + +if (SGL_KERNEL_GB10_ONLY) + # GB10 (DGX Spark) is sm_121a (consumer Blackwell, aarch64). It cannot run + # Hopper (sm_90a) nor datacenter Blackwell (sm_100a) binaries. Compiling + # them wastes ~6x build time and ~10-15 GB RAM per TU (cutlass templates), + # which OOM-kills cicc on 128 GB Spark hosts. + list(APPEND SGL_KERNEL_CUDA_FLAGS + "-gencode=arch=compute_120a,code=sm_120a" + "-gencode=arch=compute_121a,code=sm_121a" + "--compress-mode=size" + ) +endif() + +""" +s = s.replace(old2, insert_before + old2, 1) + +p.write_text(s) +print(f"edits applied to {p}") diff --git a/docker/patch/gb10/sgl-kernel-arch.patch b/docker/patch/gb10/sgl-kernel-arch.patch new file mode 100644 index 0000000000..a507e04c09 --- /dev/null +++ b/docker/patch/gb10/sgl-kernel-arch.patch @@ -0,0 +1,40 @@ +--- CMakeLists.txt.orig 2026-04-15 07:39:19.526964507 +0000 ++++ CMakeLists.txt 2026-04-15 07:39:19.548828896 +0000 +@@ -188,6 +188,7 @@ + option(SGL_KERNEL_ENABLE_FA3 "Enable FA3" OFF) + option(SGL_KERNEL_ENABLE_SM90A "Enable SM90A" OFF) + option(SGL_KERNEL_ENABLE_SM100A "Enable SM100A" OFF) ++option(SGL_KERNEL_GB10_ONLY "Build only for GB10 (sm_121a + sm_120a). Skips sm_90a/sm_100a/sm_103a/sm_110a and FA3." OFF) + + if (SGL_KERNEL_ENABLE_BF16) + list(APPEND SGL_KERNEL_CUDA_FLAGS +@@ -216,6 +217,8 @@ + + endif() + ++if (NOT SGL_KERNEL_GB10_ONLY) ++ + if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A) + list(APPEND SGL_KERNEL_CUDA_FLAGS + "-gencode=arch=compute_100a,code=sm_100a" +@@ -249,6 +252,20 @@ + ) + endif() + ++endif() # NOT SGL_KERNEL_GB10_ONLY ++ ++if (SGL_KERNEL_GB10_ONLY) ++ # GB10 (DGX Spark) is sm_121a (consumer Blackwell, aarch64). It cannot run ++ # Hopper (sm_90a) nor datacenter Blackwell (sm_100a) binaries. Compiling ++ # them wastes ~6x build time and ~10-15 GB RAM per TU (cutlass templates), ++ # which OOM-kills cicc on 128 GB Spark hosts. ++ list(APPEND SGL_KERNEL_CUDA_FLAGS ++ "-gencode=arch=compute_120a,code=sm_120a" ++ "-gencode=arch=compute_121a,code=sm_121a" ++ "--compress-mode=size" ++ ) ++endif() ++ + if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_FP4) + list(APPEND SGL_KERNEL_CUDA_FLAGS + "-DENABLE_NVFP4=1" diff --git a/scripts/run-qwen2.5-0.5B-gb10-smoke.sh b/scripts/run-qwen2.5-0.5B-gb10-smoke.sh new file mode 100755 index 0000000000..294c6836fc --- /dev/null +++ b/scripts/run-qwen2.5-0.5B-gb10-smoke.sh @@ -0,0 +1,113 @@ +#!/bin/bash +# Minimal GRPO smoke test for slime on NVIDIA DGX Spark (GB10, single GPU). +# Goal: exercise the full rollout → reward → policy-update loop for one tiny +# step and exit cleanly. Used only to validate the GB10 port; not a training +# recipe. +# +# Prerequisites: +# - /root/Qwen2.5-0.5B-Instruct (HF checkpoint) +# - /root/Qwen2.5-0.5B-Instruct_torch_dist (from tools/convert_hf_to_torch_dist.py) +# - /root/dapo-math-17k/dapo-math-17k.jsonl + +set -ex + +# clean any leftover ray/sglang +pkill -9 sglang 2>/dev/null || true +ray stop --force 2>/dev/null || true +pkill -9 ray python 2>/dev/null || true +sleep 2 + +export PYTHONBUFFERED=1 + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen2.5-0.5B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen2.5-0.5B-Instruct/ + --ref-load /root/Qwen2.5-0.5B-Instruct_torch_dist/ + --save /tmp/slime_smoke_save/ + --save-interval 9999 +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + + --num-rollout 1 + --rollout-batch-size 2 + --n-samples-per-prompt 2 + --num-steps-per-rollout 1 + --global-batch-size 4 + + --rollout-max-response-len 256 + --rollout-temperature 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 2048 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 1 + --sglang-mem-fraction-static 0.4 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +ray start --head --node-ip-address 127.0.0.1 --num-gpus 1 --disable-usage-stats + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json='{ + "env_vars": { + "PYTHONPATH": "/root/src/Megatron-LM", + "CUDA_DEVICE_MAX_CONNECTIONS": "1" + } + }' \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 1 \ + --colocate \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index a634d1f003..8b5d301952 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1070,7 +1070,7 @@ def add_wandb_arguments(parser): default=None, help=( "Log statistics of the category of reward, such as why the reward function considers it as failed. " - "Specify the key in the reward dict using this argument.", + "Specify the key in the reward dict using this argument." ), ) parser.add_argument(