diff --git a/.claude/docs/weight_sync.md b/.claude/docs/weight_sync.md index 8f03928727..2e7ed8518f 100644 --- a/.claude/docs/weight_sync.md +++ b/.claude/docs/weight_sync.md @@ -300,16 +300,16 @@ trainer-send engines. The `sharded_rdt` backend is unaffected: it vendors its own trainer-side ABCs in `sharded_rdt_base.py`, which already match 0.28.0's `Generic[TTrainerInitInfo]` shape. -### DeepGEMM is unavailable under the torch override +### DeepGEMM follows the selected torch runtime vLLM 0.28.0's metadata pins `torch==2.13.0`; we override torch to `2.11.0` because the CUDA extension wheels we build against (flash-attn, causal-conv1d, mamba-ssm, transformer-engine) have no 2.13 builds. vLLM's main extensions are stable-libtorch-ABI and load fine, but `vllm/third_party/deep_gemm/_C` is a version-specific `cpython-312` build linked against torch -2.13's `c10` and fails with an undefined-symbol `ImportError`. vLLM catches this and logs -`Module vllm.third_party.deep_gemm was found but failed to import`, then falls back, so the -engine runs — but the DeepGEMM fused-MoE and sparse-attention-indexer paths are gone. Revisit -when those wheels publish torch 2.13 builds. +2.13's `c10` and fails with an undefined-symbol `ImportError`. SkyRL therefore pins the matching +DeepGEMM source revision as an external package and builds it against the selected torch runtime. +vLLM prefers that package over its incompatible vendored extension, keeping the DeepGEMM +fused-MoE and sparse-attention-indexer paths available. ## Gotchas diff --git a/ci/gpu_ci_run_h100.sh b/ci/gpu_ci_run_h100.sh index 0b8adafb30..afe0b473a3 100755 --- a/ci/gpu_ci_run_h100.sh +++ b/ci/gpu_ci_run_h100.sh @@ -12,5 +12,6 @@ uv run --directory . --isolated --extra dev --extra fsdp pytest -s -vvv -m h100 # Run Megatron h100 tests. uv run --directory . --isolated --extra dev --extra megatron pytest -s -vvv -m h100 \ + tests/backends/skyrl_train/gpu/gpu_ci/test_deep_gemm.py \ tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py \ tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_router_replay.py diff --git a/pyproject.toml b/pyproject.toml index d9af059274..f5e70b1373 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -152,6 +152,7 @@ megatron = [ "causal-conv1d==1.6.2.post1+cu.13.0.torch.2.11; sys_platform == 'linux'", "mamba-ssm==2.3.2.post1+cu.13.0.torch.2.11; sys_platform == 'linux'", "vllm==0.28.0; sys_platform == 'linux'", + "deep-gemm; sys_platform == 'linux'", "vllm-router; sys_platform == 'linux'", # The `nixl` shim provides that namespace and dispatches on `torch.version.cuda`, # so with a cu13 torch it loads `nixl_cu13`. `nixl-cu13` ships the `nixl_cu13` @@ -292,9 +293,9 @@ override-dependencies = [ # `_moe_C_stable_libtorch.abi3.so`), so its main binaries load against 2.11 and # we hold torch there rather than churn every CUDA extension pin. The one # exception is `vllm/third_party/deep_gemm/_C`, a version-specific cpython-312 - # build linked against torch 2.13's c10: it fails to import and vLLM falls back, - # so the DeepGEMM fused-MoE / sparse-attn-indexer paths are unavailable. See - # `.claude/docs/weight_sync.md`. Revisit once those wheels publish 2.13 builds. + # build linked against torch 2.13's c10. The external DeepGEMM source below is + # rebuilt against this selected torch runtime, and vLLM prefers it over the + # incompatible vendored extension. "torch==2.11.0; sys_platform == 'linux'", # torchvision is NOT stable-ABI (unlike vLLM), and each release hard-pins one # torch: 0.28.0 wants torch 2.13.0, 0.26.0 wants 2.11.0. vLLM 0.28.0 pulls @@ -304,10 +305,12 @@ override-dependencies = [ ] [tool.uv.extra-build-dependencies] +deep-gemm = ["torch==2.11.0", "ninja"] flash-attn = [{requirement = "torch", match-runtime = true}] fast-hadamard-transform = ["torch==2.11.0", "ninja"] [tool.uv.extra-build-variables] +deep-gemm = { DG_FORCE_BUILD = "1" } flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE"} fast-hadamard-transform = { FAST_HADAMARD_TRANSFORM_FORCE_BUILD = "TRUE"} @@ -350,6 +353,7 @@ explicit = true [tool.uv.sources] skyrl-gym = { path = "./skyrl-gym", editable = true } +deep-gemm = { git = "https://github.com/deepseek-ai/DeepGEMM", rev = "8b1392b978f5a03c828dd1711090d7fb50958b8a" } # Match torch's CUDA variant (cu130). flashinfer-jit-cache = { index = "flashinfer-cu130", marker = "sys_platform == 'linux'" } flashinfer-cubin = { index = "flashinfer", marker = "sys_platform == 'linux'" } diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/test_deep_gemm.py b/tests/backends/skyrl_train/gpu/gpu_ci/test_deep_gemm.py new file mode 100644 index 0000000000..177a6b2ec9 --- /dev/null +++ b/tests/backends/skyrl_train/gpu/gpu_ci/test_deep_gemm.py @@ -0,0 +1,8 @@ +import pytest + + +@pytest.mark.h100 +def test_deep_gemm_loads_against_selected_torch(): + import deep_gemm + + assert deep_gemm.get_num_sms() > 0 diff --git a/uv.lock b/uv.lock index 0a028ef6fe..3e19ad68bd 100644 --- a/uv.lock +++ b/uv.lock @@ -1555,6 +1555,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/05/7f/798705f5296a58ca505d600456748d1be48078eac8a7050d8a98bc9edb89/decorator-5.3.1-py3-none-any.whl", hash = "sha256:f47fe6fdbd2edd623ecfe36875d37aba411624e2670dd395dddae1358689bb3c", size = 10365, upload-time = "2026-05-18T06:03:26.517Z" }, ] +[[package]] +name = "deep-gemm" +version = "2.6.1+local" +source = { git = "https://github.com/deepseek-ai/DeepGEMM?rev=8b1392b978f5a03c828dd1711090d7fb50958b8a#8b1392b978f5a03c828dd1711090d7fb50958b8a" } + [[package]] name = "deprecated" version = "1.3.1" @@ -9020,6 +9025,7 @@ megatron = [ { name = "causal-conv1d", marker = "sys_platform == 'linux'" }, { name = "datasets" }, { name = "debugpy" }, + { name = "deep-gemm", marker = "sys_platform == 'linux'" }, { name = "fast-hadamard-transform", version = "1.1.0", source = { git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=e7706faf8d1c3b9f241e36860640ad1dac644ede#e7706faf8d1c3b9f241e36860640ad1dac644ede" }, marker = "(python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (python_full_version >= '3.12' and python_full_version < '3.14' and platform_machine == 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (python_full_version < '3.12' and sys_platform == 'linux' and extra == 'extra-5-skyrl-megatron') or (python_full_version >= '3.14' and sys_platform == 'linux' and extra == 'extra-5-skyrl-megatron') or (platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (extra != 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, { name = "fast-hadamard-transform", version = "1.1.0+cu13torch2.11cxx11abitrue", source = { url = "https://github.com/erictang000/fast-hadamard-transform/releases/download/v1.1.0.post2-cu13torch2.11/fast_hadamard_transform-1.1.0%2Bcu13torch2.11cxx11abitrue-cp312-cp312-linux_x86_64.whl" }, marker = "(python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-5-skyrl-megatron') or (python_full_version != '3.12.*' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (python_full_version != '3.12.*' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (python_full_version != '3.12.*' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (python_full_version != '3.12.*' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (platform_machine != 'x86_64' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (platform_machine != 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (platform_machine != 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (platform_machine != 'x86_64' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (extra != 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, { name = "fast-hadamard-transform", version = "1.1.0+cu13torch2.11cxx11abitrue", source = { url = "https://github.com/erictang000/fast-hadamard-transform/releases/download/v1.1.0.post2-cu13torch2.11/fast_hadamard_transform-1.1.0%2Bcu13torch2.11cxx11abitrue-cp313-cp313-linux_x86_64.whl" }, marker = "(python_full_version == '3.13.*' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-5-skyrl-megatron') or (python_full_version != '3.13.*' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (python_full_version != '3.13.*' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (python_full_version != '3.13.*' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (python_full_version != '3.13.*' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (platform_machine != 'x86_64' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (platform_machine != 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (platform_machine != 'x86_64' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (platform_machine != 'x86_64' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (sys_platform != 'linux' and extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (sys_platform != 'linux' and extra != 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu') or (extra != 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, @@ -9176,6 +9182,7 @@ requires-dist = [ { name = "datasets", specifier = ">=4.0.0" }, { name = "datasets", marker = "extra == 'skyrl-train'", specifier = ">=4.0.0" }, { name = "debugpy", marker = "extra == 'skyrl-train'", specifier = "==1.8.0" }, + { name = "deep-gemm", marker = "sys_platform == 'linux' and extra == 'megatron'", git = "https://github.com/deepseek-ai/DeepGEMM?rev=8b1392b978f5a03c828dd1711090d7fb50958b8a" }, { name = "fast-hadamard-transform", marker = "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'megatron'", url = "https://github.com/erictang000/fast-hadamard-transform/releases/download/v1.1.0.post2-cu13torch2.11/fast_hadamard_transform-1.1.0%2Bcu13torch2.11cxx11abitrue-cp312-cp312-linux_x86_64.whl" }, { name = "fast-hadamard-transform", marker = "python_full_version == '3.13.*' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'megatron'", url = "https://github.com/erictang000/fast-hadamard-transform/releases/download/v1.1.0.post2-cu13torch2.11/fast_hadamard_transform-1.1.0%2Bcu13torch2.11cxx11abitrue-cp313-cp313-linux_x86_64.whl" }, { name = "fast-hadamard-transform", marker = "(python_full_version < '3.12' and sys_platform == 'linux' and extra == 'megatron') or (python_full_version >= '3.14' and sys_platform == 'linux' and extra == 'megatron') or (platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'megatron')", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=e7706faf8d1c3b9f241e36860640ad1dac644ede" },