[megatron] GLM-5.3-Flash (glm5_next) support + 4-layer GPU CI entry

.claude/docs/backends/megatron.md

@@ -8,6 +8,13 @@ SkyRL uses Megatron-Bridge for HF-to-Megatron model conversion. Installed from g
 - `MegatronConfig` in `skyrl/train/config.py`
 - `MegatronWorker` in `skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py`.
 - Custom bridges in `skyrl/backends/skyrl_train/workers/megatron/model_bridges.py` (e.g., `GLM47FlashBridge`).
+- GLM-5.3-Flash (`glm5_next`: KDA linear attention + DSA sparse MLA with kpool indexer + mHC hyper-connections)
+  lives in `skyrl/backends/skyrl_train/workers/megatron/glm5_next/` -- its own layer classes, block spec, provider and
+  bridge, since neither megatron-core nor Megatron-Bridge support it. Needs `language_model_only=True` (VL checkpoint),
+  packed sequences (KDA), and the vLLM side needs DeepGEMM for the DSA indexer.
+- `skyrl/backends/skyrl_train/workers/megatron/__init__.py` applies `patch_fa4_cute_import` before any megatron import:
+  flash-attn 2.8.x's FA4 `flash_attn.cute` module is incompatible with the cutlass DSL vLLM >= 0.28 pins and would
+  otherwise abort `import megatron.bridge`.
 
 ## Parallelism Strategies
 

docker/pyproject.amd.toml

@@ -16,8 +16,10 @@ dependencies = [
     "pillow>=11.3.0",
     "rich>=14.1.0",
     "safetensors>=0.6.2",
-    "tokenizers>=0.21.2",
-    "transformers>=5.6.1,<=5.8.0",
+    # Floor tracks the transformers pin below: transformers 5.16.1 requires
+    # tokenizers>=0.23.1.
+    "tokenizers>=0.23.1",
+    "transformers>=5.6.1,<=5.16.1",
     "typer>=0.17.4",
     "peft==0.18.1",
     "hf_transfer",

docs/content/docs/getting-started/supported_models.mdx

@@ -18,6 +18,7 @@ The following models are fully supported in SkyRL's Megatron backend via Megatro
 - Qwen-3.5 (Dense: [0.8B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_qwen3.5.sh)/2B/4B/9B/27B, MoE: [35B-A3B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_qwen3.5_35b_a3b.sh)/122B-A10B/397B-A17B)
 - Nemotron-3 (Dense: [Nano-4B-BF16](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/nemotron_3/run_nemotron_3_nano_4b_gsm8k.sh), MoE: Nano-30B-A3B-BF16)
 - [GLM-4.7-Flash](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_grpo_glm4_7_30b.sh)/GLM-4.7
+- GLM-5.3-Flash (KDA + DSA hybrid MoE with mHC hyper-connections; SkyRL ships its own Megatron layers and bridge under `skyrl/backends/skyrl_train/workers/megatron/glm5_next`, validated on the [4-layer slice](https://huggingface.co/CharyZeng/GLM-5.3-Flash-4layer) in GPU CI; requires `trainer.policy.language_model_only=true`)
 - Qwen3 (Dense: [0.6B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron.sh)/1.7B/4B/8B/32B, MoE: [30B-A3B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_dapo_qwen3_30b_a3b.sh)/[235B-A22B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_qwen3-235b-a22b.sh))
 - [Moonlight-16B-A3B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_moonlight.sh)
 

pyproject.toml

@@ -16,8 +16,10 @@ dependencies = [
     "pillow>=11.3.0",
     "rich>=14.1.0",
     "safetensors>=0.6.2",
-    "tokenizers>=0.21.2",
-    "transformers>=5.6.1,<=5.8.0",
+    # Floor tracks the transformers pin below: transformers 5.16.1 requires
+    # tokenizers>=0.23.1.
+    "tokenizers>=0.23.1",
+    "transformers>=5.6.1,<=5.16.1",
     "typer>=0.17.4",
     "peft==0.18.1",
     "hf_transfer",
@@ -113,8 +115,9 @@ skyrl-train = [
 
 fsdp = [
     "skyrl[skyrl-train]",
-    "vllm==0.28.0; sys_platform == 'linux'",
+    "vllm; sys_platform == 'linux'",
     "vllm-router; sys_platform == 'linux'",
+    "deep-gemm; sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version >= '3.12' and python_full_version < '3.13'",
     # 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`
     # module, but vLLM imports `nixl._api`. Its metadata hard-depends on `nixl-cu12`
@@ -125,9 +128,9 @@ fsdp = [
     "causal-conv1d==1.6.2.post1+cu.13.0.torch.2.11; sys_platform == 'linux'",
     "flash-attn==2.8.3+cu.13.0.torch.2.11; sys_platform == 'linux'",
     "torch==2.11.0; sys_platform == 'linux'",
-    "flashinfer-python==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-jit-cache==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-cubin==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-python==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-jit-cache==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-cubin==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
     "torchvision; sys_platform == 'linux'",
 ]
 
@@ -149,23 +152,24 @@ megatron = [
     "flash-linear-attention; sys_platform == 'linux'",
     "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'",
+    "vllm; sys_platform == 'linux'",
     "vllm-router; sys_platform == 'linux'",
+    "deep-gemm; sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version >= '3.12' and python_full_version < '3.13'",
     # 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`
     # module, but vLLM imports `nixl._api`. Its metadata hard-depends on `nixl-cu12`
     # too; that variant is overridden out below (it would drag the CUDA-12 runtime
     # into the image for a module the shim would never load).
     "nixl; sys_platform == 'linux'",
     "torch==2.11.0; sys_platform == 'linux'",
-    "flashinfer-python==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-python==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
     "torchvision; sys_platform == 'linux'",
     # megatron-bridge requires Python 3.12+; pin megatron-core to the same
     # constraint so both packages are consistently available (or absent).
     "megatron-bridge; sys_platform == 'linux' and python_version >= '3.12'",
     "megatron-core; sys_platform == 'linux' and python_version >= '3.12'",
-    "flashinfer-jit-cache==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-cubin==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-jit-cache==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-cubin==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
     "nvidia-modelopt; sys_platform == 'linux'",
     "fast-hadamard-transform; sys_platform == 'linux'",
 ]
@@ -249,16 +253,16 @@ no-build-isolation-package = [
 override-dependencies = [
     "nvidia-resiliency-ext; sys_platform == 'never'",
     "transformer-engine[pytorch]==2.16.0; sys_platform == 'linux'",
-    "transformers>=5.6.1,<=5.8.0; sys_platform == 'linux'",
+    "transformers>=5.6.1,<=5.16.1; sys_platform == 'linux'",
     "megatron-core>=0.16.0; sys_platform == 'linux' and python_version >= '3.12'",
     "ml_dtypes>=0.5.0; sys_platform == 'linux'",
     # `nixl` hard-depends on both nixl-cu12 and nixl-cu13; drop the cu12 variant
     # so it doesn't pull a second, unused CUDA runtime in alongside the cu13 one.
     "nixl-cu12; sys_platform == 'never'",
     # Megatron-Bridge pins flashinfer-python==0.6.8.post1, which conflicts with
-    # our pin, so all three flashinfer packages are held here. vLLM 0.28.0 needs
-    # >=0.6.16: it imports `flashinfer.autotuner.set_autotune_process_group`,
-    # which 0.6.14 does not have, and the engine fails to start without it.
+    # our pin, so all three flashinfer packages are held here. The pinned vLLM
+    # (0.28.1rc1.dev359, see `vllm` in [tool.uv.sources]) requires exactly 0.6.18; vLLM
+    # 0.28.0 already needed >=0.6.16 (`flashinfer.autotuner.set_autotune_process_group`).
     # flashinfer hard-errors when an installed cubin's version differs from its
     # own, so the three must move together.
     # They must also be >=0.6.13: older flashinfer rejects the `layout_code` vLLM's
@@ -268,9 +272,9 @@ override-dependencies = [
     # flashinfer-cubin is not published to PyPI past 0.6.13, so it (like
     # flashinfer-jit-cache) comes from flashinfer's own index -- see the
     # `flashinfer` index and [tool.uv.sources] below.
-    "flashinfer-python==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-jit-cache==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-cubin==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-python==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-jit-cache==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-cubin==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
     # mamba-ssm 2.3.2.post1 pins tilelang==0.1.8 and apache-tvm-ffi<=0.1.9, while
     # vLLM 0.28.0 pins tilelang==0.1.12 and apache-tvm-ffi==0.1.11. Both overrides
     # hold these at the versions the lock already resolved to before mamba-ssm moved
@@ -397,6 +401,25 @@ megatron-bridge = {git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge", rev =
 # `git ls-tree <bridge-rev> 3rdparty/Megatron-LM`. This one is on the Megatron-LM
 # main branch (https://github.com/NVIDIA/Megatron-LM/tree/main), dated 8/12/26.
 megatron-core = {git = "https://github.com/NVIDIA/Megatron-LM", rev = "14346b65a2d0790e451919858f7771078105c5f0", marker = "sys_platform == 'linux'"}
+# vLLM at 98ed0856f31fa3aaf5e27464e2b4ef5a8ee6b2f5 (vllm-project/vllm#53906, "[Model] add GLM-5.3-Flash
+# support", 2026-09-03), i.e. 0.28.1rc1.dev359 -- the first vLLM with the GLM-5.3-Flash (glm5_next: KDA +
+# DSA-kpool + mHC) model. Per-commit wheels are published by vLLM's CI at wheels.vllm.ai; this one is
+# built for CUDA 13 / torch 2.13 but, like 0.28.0, its main binaries are stable-libtorch-ABI so they load
+# against the torch 2.11 we hold (see the torch override above). Move back to a plain PyPI pin once a
+# release containing #53906 ships.
+vllm = [
+    { url = "https://wheels.vllm.ai/98ed0856f31fa3aaf5e27464e2b4ef5a8ee6b2f5/vllm-0.28.1rc1.dev359%2Bg98ed0856f-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
+    { url = "https://wheels.vllm.ai/98ed0856f31fa3aaf5e27464e2b4ef5a8ee6b2f5/vllm-0.28.1rc1.dev359%2Bg98ed0856f-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
+]
+# DeepGEMM at deepseek-ai/DeepGEMM@8b1392b (the commit the pinned vLLM vendors), built against torch 2.11 +
+# CUDA 13 (nvcc 13.0.88, cp312, x86_64, glibc >= 2.35) and hosted like the fast-hadamard-transform wheels
+# below. vLLM ships its own copy under `vllm.third_party.deep_gemm`, but that `_C` is a version-specific
+# cpython extension linked against torch 2.13 and fails to import under the torch 2.11 we hold; vLLM
+# prefers an installed `deep_gemm` package when one exists. Needed because vLLM's GLM-5.3-Flash DSA
+# indexer hard-requires DeepGEMM (DeepSeek-V3.2-style sparse-attention models too). Kernels are JIT
+# compiled at runtime, so the host still needs nvcc/nvrtc. Rebuild if the torch pin moves; drop once
+# vLLM's vendored extension loads against our torch.
+deep-gemm = { url = "https://github.com/erictang000/DeepGEMM/releases/download/v2.6.1-8b1392b-cu13torch2.11/deep_gemm-2.6.1%2B8b1392b-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version >= '3.12' and python_full_version < '3.13'" }
 # NOTE (sumanthrh): This custom wheel of vllm-router includes a fix for /chat/completions endpoint: https://github.com/vllm-project/router/pull/162  on top of 0.1.14
 # TODO (sumanthrh): Remove this after 0.1.15 vllm-router release
 vllm-router = [

skyrl/backends/skyrl_train/inference_servers/render_server.py

@@ -80,7 +80,11 @@ def _run_render_server(model_path: str, port: int, log_file: Optional[str]) -> N
         build_and_serve_renderer,
         setup_server,
     )
-    from vllm.entrypoints.openai.cli_args import make_arg_parser
+
+    try:  # vLLM >= 0.28.1 moved the CLI args to entrypoints.launchers
+        from vllm.entrypoints.launchers.cli_args import make_arg_parser
+    except ImportError:
+        from vllm.entrypoints.openai.cli_args import make_arg_parser
     from vllm.utils.argparse_utils import FlexibleArgumentParser
 
     async def _serve() -> None:

skyrl/backends/skyrl_train/inference_servers/utils.py

@@ -68,7 +68,11 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace:
     """Build CLI args for vLLM server from config."""
     from vllm import AsyncEngineArgs
     from vllm.config import WeightTransferConfig
-    from vllm.entrypoints.openai.cli_args import FrontendArgs
+
+    try:  # vLLM >= 0.28.1 (entrypoints.launchers); the openai.cli_args path is the pre-0.28.1 location
+        from vllm.entrypoints.launchers.cli_args import FrontendArgs
+    except ImportError:
+        from vllm.entrypoints.openai.cli_args import FrontendArgs
     from vllm.platforms import current_platform
     from vllm.utils.argparse_utils import FlexibleArgumentParser
 

skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py

@@ -592,6 +592,9 @@ async def _build_and_serve_vllm_server(
     # One uvicorn per port (no api_server_count fan-out), matching vLLM's own
     # single-server path, so SO_REUSEPORT stays off.
     sock = create_server_socket(sock_addr, reuse_port=False)
+    # vLLM >= 0.28.1 gates the token-in/token-out ``/inference/v1/generate`` route (which SkyRL's
+    # generation client calls) behind this env var, read when the app is built.
+    os.environ.setdefault("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1")
     app = build_app(cli_args)
 
     # Initialize the engine (this loads the model - takes time)
@@ -646,7 +649,11 @@ def _build_standalone_cli_args(argv: Optional[List[str]] = None) -> Namespace:
     ``--worker-extension-cls``, ...).
     """
     from vllm import AsyncEngineArgs as _AsyncEngineArgs
-    from vllm.entrypoints.openai.cli_args import FrontendArgs
+
+    try:  # vLLM >= 0.28.1 moved the CLI args to entrypoints.launchers
+        from vllm.entrypoints.launchers.cli_args import FrontendArgs
+    except ImportError:
+        from vllm.entrypoints.openai.cli_args import FrontendArgs
     from vllm.platforms import current_platform
     from vllm.utils.argparse_utils import FlexibleArgumentParser
 

skyrl/backends/skyrl_train/patches/megatron/patch_fa4_cute_import.py

@@ -0,0 +1,53 @@
+"""Keep a broken FlashAttention-4 (``flash_attn.cute``) install from taking down ``megatron.core``.
+
+flash-attn 2.8.x ships the FA4 ``flash_attn.cute`` package alongside FA2. It is written against
+a specific ``nvidia-cutlass-dsl`` release; vLLM >= 0.28 pins a newer cutlass DSL (4.6.x) whose
+``cutlass.cute.core`` no longer has the symbols FA4 imports (``ThrMma``), so importing
+``flash_attn.cute`` raises ``AttributeError`` rather than ``ImportError``.
+
+``megatron.core.transformer.attention`` probes FA4 with ``from flash_attn.cute import
+flash_attn_varlen_func`` guarded only by ``except ImportError`` (megatron-core 0.20), so the
+``AttributeError`` escapes and every ``import megatron.bridge`` fails. SkyRL never uses FA4 (TE
+dispatches to FA2 or cuDNN fused attention), so when the probe import fails for any reason
+other than a plain ``ImportError`` we register ``flash_attn.cute`` as unavailable
+(``sys.modules[name] = None`` makes any later import of it raise ``ImportError``), which is
+exactly the state megatron-core handles.
+
+No-op when ``flash_attn.cute`` imports cleanly or is genuinely absent. Call it before the first
+``megatron`` import in a process; ``skyrl.backends.skyrl_train.workers.megatron`` does so on
+package import.
+
+DELETE THIS PATCH once megatron-core's FA4 probe catches ``Exception`` or the flash-attn /
+cutlass DSL pins agree again.
+"""
+
+import importlib
+import sys
+
+from loguru import logger
+
+_MODULE = "flash_attn.cute"
+_PATCHED_FLAG = "_skyrl_fa4_cute_import_patched"
+
+
+def patch_fa4_cute_import() -> bool:
+    """Return True if ``flash_attn.cute`` was marked unavailable, False if nothing was needed."""
+    if getattr(sys, _PATCHED_FLAG, False):
+        return sys.modules.get(_MODULE, True) is None
+    setattr(sys, _PATCHED_FLAG, True)
+    if _MODULE in sys.modules:
+        return sys.modules[_MODULE] is None
+    try:
+        importlib.import_module(_MODULE)
+        return False
+    except ImportError:
+        return False  # genuinely absent: megatron-core already handles this
+    except Exception as exc:  # e.g. AttributeError from a cutlass DSL mismatch
+        for name in [m for m in sys.modules if m == _MODULE or m.startswith(_MODULE + ".")]:
+            del sys.modules[name]
+        sys.modules[_MODULE] = None
+        logger.warning(
+            f"`{_MODULE}` (FlashAttention-4) failed to import with {type(exc).__name__}: {exc}. "
+            "Marking it unavailable so megatron-core falls back to FA2 / cuDNN attention."
+        )
+        return True

skyrl/backends/skyrl_train/weight_sync/sharded_rdt/rdt_vllm_register.py

@@ -22,7 +22,10 @@
 function and ImportError is swallowed on platforms without vLLM.
 
 REMOVAL: delete this module + its imports once SkyRL's pinned vLLM registers the
-``sharded_rdt`` engine in ``WeightTransferEngineFactory`` natively.
+``sharded_rdt`` engine in ``WeightTransferEngineFactory`` natively. vLLM >= 0.28.1
+(the pinned dev wheel) does register ``sharded_rdt`` natively, so on that wheel this
+module is a no-op and vLLM's own engine is used; the vendored engine below stays as the
+fallback for older wheels until the two implementations are reconciled.
 """
 
 import logging

skyrl/backends/skyrl_train/workers/megatron/__init__.py

@@ -0,0 +1,13 @@
+"""Megatron backend workers.
+
+The package import runs :func:`patch_fa4_cute_import` first: with the pinned flash-attn and
+the cutlass DSL that vLLM >= 0.28 pulls in, ``import megatron.bridge`` otherwise aborts on a
+broken FlashAttention-4 probe (see the patch module). Every Megatron entry point in SkyRL
+lives under this package, so this is the earliest common place to apply it.
+"""
+
+from skyrl.backends.skyrl_train.patches.megatron.patch_fa4_cute_import import (
+    patch_fa4_cute_import,
+)
+
+patch_fa4_cute_import()

skyrl/backends/skyrl_train/workers/megatron/glm5_next/__init__.py

@@ -0,0 +1,14 @@
+"""GLM-5.3-Flash (``glm5_next``) support for the Megatron backend.
+
+GLM-5.3-Flash (HF ``Glm5NextForConditionalGeneration``) is a 45-layer hybrid: 34 KDA
+linear-attention layers + 11 DSA sparse-MLA layers (NoPE, kpool-compressed lightning
+indexer), a 288-expert sigmoid MoE (first 3 layers dense, one shared expert) and mHC
+hyper-connections on every block. Neither the pinned megatron-core nor Megatron-Bridge
+support it, so this package provides the layer implementations (``kda``, ``dsa``, ``mhc``,
+``layer``), the block spec / provider, and the HF<->Megatron bridge.
+
+Importing :mod:`.bridge` registers the bridge with Megatron-Bridge's ``AutoBridge``
+(``model_bridges.py`` does so for the Megatron worker). The package itself imports nothing
+heavy so the pure-torch pieces (``kpool_indexer``) stay importable without megatron-core.
+Requires ``flash-linear-attention`` (KDA kernels) at runtime.
+"""