[megatron] GLM-5.3-Flash (glm5_next) support: KDA + NoPE-MLA/DSA + mHC hybrid MoE, 4-layer GPU CI entry - #2179
Open
erictang000 wants to merge 18 commits into
Open
erictang000 wants to merge 18 commits into
erictang000 wants to merge 18 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces support for the GLM-5.3-Flash (glm5_next) model in the Megatron backend, implementing custom Megatron-Core extensions such as Kimi Delta Attention (KDA), Manifold-Constrained Hyper-Connections (mHC), and a custom transformer layer that supports MoE. It also updates dependencies (including vLLM and FlashInfer) and includes compatibility fixes for vLLM >= 0.28.1. The review identified two critical runtime issues: an AttributeError in the DSA attention module due to an incorrect attribute name (self.index_topk instead of self.dsa_indexer_topk), and a TypeError in the mHC transformer layer caused by passing an unsupported padding_mask argument to the MLP forward call.
Comment on lines
+41
to
+44
| if max_seqlen > self.index_topk: | ||
| raise NotImplementedError( | ||
| f"GLM-5.3-Flash sparse attention with sequences longer than dsa_indexer_topk=" | ||
| f"{self.index_topk} tokens (got {max_seqlen}) needs the k-pool indexer, which the " |
Contributor
There was a problem hiding this comment.
The attribute self.index_topk does not exist on DSAttention or Glm5NextDSAttention. The correct attribute name in megatron-core is self.dsa_indexer_topk (or self.config.dsa_indexer_topk). Accessing self.index_topk will raise an AttributeError at runtime.
Suggested change
| if max_seqlen > self.index_topk: | |
| raise NotImplementedError( | |
| f"GLM-5.3-Flash sparse attention with sequences longer than dsa_indexer_topk=" | |
| f"{self.index_topk} tokens (got {max_seqlen}) needs the k-pool indexer, which the " | |
| if max_seqlen > self.dsa_indexer_topk: | |
| raise NotImplementedError( | |
| f"GLM-5.3-Flash sparse attention with sequences longer than dsa_indexer_topk=" | |
| f"{self.dsa_indexer_topk} tokens (got {max_seqlen}) needs the k-pool indexer, which the " |
| pre_mlp_layernorm_output, moe_padding_mask, moe_unflatten_mbs = self._maybe_unflatten_for_moe( | ||
| pre_mlp_layernorm_output, padding_mask, packed_seq_params | ||
| ) | ||
| mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=moe_padding_mask) |
Contributor
There was a problem hiding this comment.
Standard Megatron-Core MLP and MoE layers (such as GroupedMLP or TEGroupedMLP) do not accept padding_mask in their forward method. Passing padding_mask=moe_padding_mask will raise a TypeError at runtime. You should remove this keyword argument.
Suggested change
| mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=moe_padding_mask) | |
| mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output) |
 The DeepGEMM platform coverage must be fixed or the unsupported environments explicitly rejected before merging because GLM-5.3-Flash cannot initialize there. Findings
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| "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'", |
There was a problem hiding this comment.
When the Megatron extra is installed on Linux aarch64 or x86_64 Python 3.13, the environment marker excludes the Torch 2.11-compatible deep-gemm package, so vLLM falls back to its Torch 2.13-linked extension and GLM-5.3-Flash inference fails to initialize its required DSA indexer.
Brings in the torch 2.13 upgrade (NovaSky-AI#2175), which replaces the `+cu.13.0.torch.2.11` local-version pins with plain versions from NovaSky-AI/skyrl-wheels and raises requires-python to 3.12. Conflicts were pyproject.toml (5 hunks) and uv.lock. Resolution takes main's torch 2.13 and CUDA-extension pins throughout, and keeps this branch's flashinfer 0.6.18 (required by the pinned vLLM) and the unpinned `vllm` driven by the per-commit dev-wheel source. Drops the `deep-gemm` dependency and its hosted wheel. It existed only because the wheel we could build topped out at torch 2.11 while vLLM's vendored `vllm.third_party.deep_gemm._C` is built against torch 2.13, so the vendored copy failed to import and vLLM prefers an installed `deep_gemm` when one exists. On torch 2.13 the vendored extension loads, so keeping a torch-2.11 build would have been actively harmful: vLLM would prefer it and then fail to import it, taking out the GLM-5.3-Flash DSA indexer that hard-requires DeepGEMM. Verified on the merged tree: `vllm.third_party.deep_gemm._C` imports, `vllm.utils.deep_gemm.has_deep_gemm()` is True, and no standalone `deep_gemm` is installed. `dsa_index_share_recompute.patch` is unchanged here: this branch still pins megatron-core 14346b65a, which the existing patch matches. (It needs refreshing only alongside the megatron-core bump to b3393bbb -- see NovaSky-AI#2180.) `patch_fa4_cute_import` is still required under torch 2.13: the cutlass DSL the pinned vLLM pulls in still breaks `flash_attn.cute`, and the guard still reports marking it unavailable. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Brings in the megatron bump (NovaSky-AI#2180): megatron-core 14346b65a -> b3393bbb and megatron-bridge 2b8cb21d -> 8e5f0e1, plus the refreshed DSA index-share patch and the removal of the expert-bias padding-mask shim. pyproject.toml conflict resolved to main's megatron-core rev, keeping this branch's flashinfer 0.6.18 and the per-commit vLLM dev-wheel source; uv.lock regenerated from main's so the only additions are flashinfer 0.6.18, the vLLM wheel and instanttensor. The new megatron-core ships mHC natively, which requires the adaptations this branch was carrying on top of the old rev: - `mcore_ext/hyper_connection.py`: the 688-line backport of Megatron-LM main's `HyperConnectionModule` is deleted in favour of `RMSNormInputHyperConnectionModule` (49 lines), which only overrides the input normalization to GLM's standard RMSNorm (`rsqrt(mean(x^2) + eps)`) instead of megatron-core's `1 / (rms(x) + 1e-6)`. - `mcore_ext/mhc_transformer_layer.py`: `TransformerBlock` now expands/contracts the mHC residual streams itself when `enable_mhc_connections` is set, so the layer no longer does it at the first/last layer -- otherwise the streams are expanded twice and the mHC mapping sees `[s, n*n*C]`. The layer keeps the per-sub-layer residual update and gains a guard rejecting `'mhc'` in `recompute_modules`, whose managers it does not thread. - `MegatronWorker.init_configs`: megatron-core rejects mHC under `recompute_granularity="full"`, which is what SkyRL's default `gradient_checkpointing=True` resolves to, and its suggested alternative needs the managers above. Downgrades to selective recompute of the remaining modules with a log line. Validated on the merged tree: `test_glm5_next_modules.py` 2 passed with per-module errors unchanged from before the bump (mHC exact to 4e-9, KDA 0.5% of the reference mean), and Megatron-vs-HF TP1 logit parity on GLM-5.3-Flash-4layer stays at the bf16 noise floor (KL(HF||Meg) 0.0086, argmax agree 92.7%, mean |dlogprob| 0.060). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ngle B300 node The existing GLM-5.3-Flash coverage runs the 4-layer slice, which exercises every code path but is not a coherent LM. This adds the real 45-layer `zai-org/GLM-5.3-Flash` checkpoint as a `test_megatron_models` row so the same logprob-parity and weight-sync check can be run against real weights, where generation should also read as sensible text. ~313B params in bf16 (~627 GiB, 97% of it routed experts) with ~17B activated, so it needs a whole 8xB300 node and cannot run in CI. It carries a new `b300` marker, gated exactly like `h100`: auto-skipped unless `-m b300` is passed, so `-m megatron_models` and `-m h100` both leave it alone. The gating loop in the GPU conftest now covers both markers instead of hard-coding h100. Mesh: Megatron TP2 EP8 ETP1 -> DP4 (EP x ETP == TP x DP), vLLM TP8 colocated on the same 8 GPUs. EP is the scaling dimension for a MoE this sparse -- 36 experts/GPU, ~76 GiB -- while TP only covers the ~9B of non-expert weights. PP stays at 1 because megatron-core rejects mHC with pipeline_model_parallel_size > 1, and CP at 1 because KDA has no context-parallel path. Everything else is picked up by the existing `glm-5.3-flash` branches: `language_model_only`, packed sequences, `inference_only_init`, `max_num_seqs=512` (the vLLM KDA triton grid limit) and the 4096-token / 0.5-utilization engine overrides. Also pins fla to its Triton kernels for GLM-5.3-Flash on Blackwell (`FLA_TILELANG=0`), the same workaround the Qwen3.5 rows use for GDN -- KDA runs fla kernels too. Scoped to Blackwell so the H100 rows keep fla's default backend. The real config needs nothing the bridge rejects: all 45 `indexer_types` are `"full"`, so the unsupported cross-layer DSA index sharing never comes up. Thresholds mirror the other large-MoE rows and have not been measured on this checkpoint; expect to tune them on the first run. Not run here -- this box has no B300s. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Eric Tang <erictang000@gmail.com>
…LoRA packing, HTTP diagnostics Three independent fixes found while bringing up GRPO/DAPO for GLM-5.3-Flash on 8xB300. Forward CUDA_HOME into the Ray runtime env. TileLang JITs kernels by shelling out to nvcc and picks its toolkit from CUDA_HOME, defaulting to the pip wheel tree. That tree is internally inconsistent here -- nvidia-cuda-nvcc==13.3 next to nvidia-cuda-runtime==13.0 (CUDART_VERSION 13000) -- and nvidia-cuda-cccl asserts the two match, so every JIT fails with "CUDA compiler and CUDA toolkit headers are incompatible". Pointing CUDA_HOME at a self-consistent system toolkit fixes it, but workers are re-exec'd through the runtime env, so a driver-side export never reaches them. Reproduced standalone: compiling tilelang's common.h for sm_103a fails with the wheel nvcc and succeeds with /usr/local/cuda-13.3. Backport vllm-project/vllm#56327 as a runtime patch. Glm5NextForConditionalGeneration otherwise inherits Glm4v's packed_modules_mapping, which names none of the model's fused projections, so a Megatron-trained adapter's separate q/k/v/b/f_a/g_a weights have nothing to assemble onto. Also honors the replicated_shard_ids that KDA's in_proj_qkvbfg_a already declares, keeping LoRA-B whole on every TP rank for f_a_proj/g_a_proj (matching parallel_mode="duplicated" on the trainer side). Adds one thing the PR does not: KDA splits its fused projection into non-contiguous views, so a LoRA-wrapped f_b_proj trips `assert inputs.is_contiguous()` in the triton lora_shrink; the shrink input is now made contiguous when needed (a [num_tokens, 128] bf16 copy, only when non-contiguous). The patch is guarded, idempotent, and skips a build that already has the fix. Report non-JSON HTTP responses. The bare orjson.JSONDecodeError says only "line 1 column 1 (char 0)" and the retry reason was logged at debug, so an engine returning 502 was invisible from the driver log. Logging status, length and a body snippet turned an opaque failure into "status=502 Backend request failed: error sending request for url (...)" immediately. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…lash Three scripts: GSM8K GRPO on a single colocated 8xB300 node, and DAPO (sync colocated on 3 nodes, plus fully-async with 2 trainer nodes and 1 inference node). The GSM8K one is verified -- 51 steps, 0 errors, rewards 0.89-1.00 at ~215-290 token responses. The model-specific settings are the non-obvious part, so each is commented in place: - LoRA, not full fine-tuning: ~313B params means bf16 weights + grads + fp32 master/m/v is ~5 TiB against 2.2 TiB of HBM, and at EP8/ETP1 on 8 ranks the expert optimizer state gets no sharding at all (expert-DP == 1). - merge_lora=true. vLLM selects a LoRA-aware MoE expert kernel whenever LoRA is enabled globally (fused_moe/oracle/unquantized.py), but only sets the lora_context that kernel asserts on when the MoE layer is itself LoRA-wrapped. Merging sidesteps the whole path; the scripts switch back to adapter sync with one flag once that is usable. - recompute_granularity=selective with [core_attn,moe]. megatron-core rejects mHC with full activation recompute, and its suggested 'mhc' is then rejected by SkyRL's own HyperConnectionTransformerLayer (it does not thread CheckpointWithoutOutputManager, so accepting it would silently drop the recompute). Selective also requires recompute_num_layers=None, which DEFAULT_TRANSFORMER_CONFIG_KWARGS otherwise injects. - use_fused_mhc left off: the fused kernels implement only 1/(rms+eps), incompatible with the eps-inside-sqrt normalization GLM actually uses. - Sequences capped under 2048. The DSA layers index with dsa_indexer_topk=2048 and the Megatron backend has no k-pool indexer, so glm5_next/dsa.py raises at training time for anything longer. - n_samples_per_prompt=12 for the sync DAPO run: dp is pinned to 24/TP2=12 (KDA has no context-parallel path, mHC rejects PP>1), and validate_cfg needs (mini_batch * n_samples) % dp == 0. The DAPO scripts additionally raise the router queue past the request fan-out, throttle SKYRL_GENERATE_CONCURRENCY_PER_ENGINE (the 512 default releases 1536 requests at once and every engine 502s), and raise SKYRL_WORKER_NCCL_TIMEOUT_IN_S -- a DAPO fwd_logprobs pass runs ~37 min, so DP ranks waiting in a collective blow the 600s default and the NCCL watchdog aborts the worker. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2 tasks
… k-pool indexer Points megatron-core at the head of NVIDIA/Megatron-LM#7054 ("GLM5.3 Flash (KDA + mHC + KPool DSA) support + FP8", HollowMan6) instead of implementing the k-pool indexer ourselves. That PR supersedes our #7427: it carries the same KDA and mHC work plus the pooled indexer, was opened two weeks earlier, is not a draft, and its config field names already match this provider (dsa_indexer_kpool, dsa_indexer_kpool_always_select_tail, mhc_norm_eps_inside_sqrt). Its pooling matches the HF/vLLM semantics independently reverse-engineered from vLLM's Triton kernel: a per-dimension softmax over each pool's slots of `gate_score + ape`, weighting the keys. It also handles the two parts that are easy to get wrong -- packed-sequence boundaries (_kpool_compress_keys_per_seg) and the context-parallel gather/reorder of the gate score. megatron-bridge stays on 8e5f0e13: bridge main needs hydra-core>=1.3.4 and SkyRL pins ==1.3.2, so only megatron-core moves, isolating the variable. That is safe because the core rev bridge main validates against (f6c33bde) is an ANCESTOR of the PR head (ahead 60, behind 0). Two SkyRL-side changes are required: - Map index_kpool_compress_ape / index_kpool_compress_gate. megatron-core only creates these when dsa_indexer_kpool > 1 and initializes them randomly (nn.init.normal_ on the gate), so a pooled run without these mappings would silently train against random pooling weights. They are bare nn.Parameters on DSAIndexer rather than module weights, so AutoMapping cannot infer a parallelism type ("Cannot determine parallelism type for module 'DSAIndexer'"); the indexer is duplicated across TP ranks, so they are ReplicatedMapping. - Relax the >dsa_indexer_topk guard to fire only for kpool <= 1, which is now the only regime megatron-core cannot handle. Note index_kpool is an attribute of DSAIndexer, not DSAttention, so the guard reads it from the config. Testing: upstream's own k-pool unit tests pass here (44), the new test_glm5_next_kpool.py passes (6), and the GLM-5.3-Flash 4-layer model-level parity row passes -- Megatron logprobs against vLLM plus the weight-sync round trip, with kpool=4 active and the pooling weights loaded from the checkpoint. All seven SkyRL GLM modules still import across 60 commits of megatron-core drift. Not yet covered: a >dsa_indexer_topk parity row. Below the budget every pool is selectable, so the existing row cannot distinguish correct pooling weights from wrong ones -- only sequences past 2048 exercise pool selection itself. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The existing GLM-5.3-Flash rows generate 128 tokens from ~100-250 token prompts, so every sequence stays under dsa_indexer_topk (2048). In that regime every pool is selectable and the k-pool indexer covers the full causal prefix regardless of what the compression weights hold -- the short row passes even with index_kpool_compress_gate/ape left randomly initialized, which is exactly what an unmapped bridge entry produces (megatron-core does nn.init.normal_ on the gate). Only past the budget does scoring decide which pools survive, so this is the first row whose logprob comparison is sensitive to those weights at all. Threading a per-row max_generate_length through the parametrize (existing rows pass None and keep the module default) is enough to get there: GSM8K prompts are too short to cross 2048, so the length comes from generation. Measured over three runs on the 4-layer slice: Megatron vs vLLM 0.054 / 0.053 / 0.0546 (threshold 1e-1, unchanged) vLLM pre- vs post-sync 0.351 / 0.349 / 0.3498 (threshold raised to 5e-1) The first is the check that matters, and the pooled path agrees with vLLM just as closely past the budget as the dense path does below it. The second is raised because over a 2048-token greedy generation it stops measuring weight-sync fidelity: one token flipped by a tiny numerical difference makes every later token differ, and all three runs logged "pre/post-sync generation lengths differ". It is reproducible rather than chaotic, so 5e-1 keeps headroom while still failing on a real regression -- the 128-token rows sit around 0.06. Tightening it further means shortening the generation, which would stop the row exercising pool selection at all. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ementing the DSA k-pool indexer Lifts the dsa_indexer_topk (2048) ceiling on Megatron training. megatron-bridge stays pinned at 8e5f0e13 and SkyRL keeps its own glm5_next bridge package; only megatron-core moves. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The k-pool vendoring is the branch's actual configuration now, not a side experiment, and the comment should say why the pin is where it is and when to move it back. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…DSA ceiling megatron-core implements the k-pool indexer, so sequences past dsa_indexer_topk (2048) are selected rather than refused, and the 896 + 1024 cap those scripts carried is gone. Moves to 2048 prompt + 4096 response. Still short of the stock DAPO recipe's 2k + 8k: an 8k response is roughly 8x the per-step cost of the 1024-response runs that measured ~10 min/step, which does not fit an overnight experiment. max_tokens_per_microbatch goes 4096 -> 8192 because a microbatch now has to hold one full sequence; 16384 is still out, having OOM'd at step 1 on the GSM8K run. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
All three GLM-5.3-Flash recipes run merge_lora=true, which is correct but syncs ~599 GiB of merged weights every step instead of a few hundred MB of adapter -- worst in the non-colocated async recipe, where that crosses the network. Documents the blocker chain (vLLM selects a LoRA-aware MoE kernel globally whenever LoRA is enabled, but only sets the lora_context that kernel asserts when the MoE layer is itself wrapped, and wrapping requires 'experts' to pass the target filter), what is already fixed on this branch (the vllm#56327 backport, replicated_shard_ids, the add_shrink contiguity guard), and the leads that are inferred but never executed -- chiefly that the 'experts' entry already sits unused in the GSM8K script, and the likely share_expert_adapters / enable_moe_shared_loras mismatch. Filed under .claude/docs with a routing row so an agent picking this up reads it first. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…nalizing typical response lengths A 45-step run at rank 32 with shared expert adapters moved implied per-sample accuracy only 0.482 -> 0.502 while reward tracked response length at corr -0.93, so the curve could not distinguish learning to solve from learning to be brief. overlong_buffer_len 2048 -> 1024. On a 4096 cap the penalty started at 2048 while the mean response was 2440, so every step's mean was penalized and the penalty's std was 87% of the whole correctness signal's. At 1024 it starts at 3072 and is zero at current lengths, becoming the near-cap guard the recipe intends instead of a tax on the typical rollout. share_expert_adapters=False + normalize_moe_lora=True at rank 64. One adapter shared across all local grouped experts was the capacity bottleneck with 288 experts holding ~97% of the parameters. Normalization divides the expert rank by moe_router_topk (8 here -> expert rank 8), keeping the per-token expert contribution comparable to a dense rank-64 adapter; it requires rank % topk == 0, which 64 satisfies, and the resulting expert rank stays under the rank vLLM would size max_lora_rank from. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
vmnode-6r3vaf61zkut was drained after its GPU0 lost P2P with every peer (nvidia-smi topo -p2p r shows NS across GPU0's row), which made vLLM's TP=8 ncclCommInitRank fail there on every attempt. It cannot be reset in-guest -- the GPUs are passthrough, so nvidia-smi -r is refused even with no process holding the device -- and its NVLinks have been inactive since its fabric manager died on 2026-07-18, so it had been running on PCIe P2P the whole time. sync: 3 -> 2 nodes and 3 -> 2 engines. world 16, dense dp 8, expert dp 2; mini_per_gpu 48, train_per_gpu 192, seqs/step unchanged at 1536. Per-GPU work rises ~1.5x. async: the policy drops from 2 nodes to 1 so the inference engine still gets a whole node to itself (it is the non-colocated recipe). world 8, dense dp 4, expert dp 1; mini_per_gpu 96, train_per_gpu 96. Per-GPU weight sharding is unchanged because EP8 is unchanged. Both carry a note to restore the 3-node sizing once that node is repaired. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Item 3 of 'Still open' in .claude/docs/glm5_3_flash_lora.md: the DAPO recipes could not flip MERGE_LORA because they carried no VLLM_LORA_TARGET_MODULES or conditional LORA_ENGINE_KWARG block -- they were written after the switch to merge_lora=true. Copied across from the GSM8K recipe, including 'experts', which is the whole fix for 'LoRA context must be set'. Gating verified both ways: lora_target_modules is absent from engine_init_kwargs at MERGE_LORA=true and present at false, where _uses_lora_weight_sync then reports enable_lora=True. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds GLM-5.3-Flash (HF
glm5_next,Glm5NextForConditionalGeneration) to the Megatron backend and atest_megatron_modelsH100 CI entry on the 4-layer sliceeatang/GLM-5.3-Flash-4layer(mirror ofCharyZeng/GLM-5.3-Flash-4layer; real truncated weights, so the logprob distribution stays peaked and the vLLM/Megatron comparison is meaningful even though the slice is not a coherent LM).The architecture is unknown to both pinned upstreams: 45 layers of 3:1 KDA linear-attention / NoPE-MLA DeepSeek-sparse-attention (kpool indexer), 288-expert sigmoid MoE with one shared expert (dense MLP in the first layer), clamped SwiGLU, and Manifold-Constrained Hyper-Connections (mHC) on every block, shipped as a VL checkpoint.
Written from the HF modeling code against upstream Megatron-LM / Megatron-Bridge APIs, reusing what upstream already ships: megatron-core's
DSAttention+AbsorbedMLASelfAttention(the NoPE case,qk_pos_emb_head_dim=0, works as is),TEGroupedMLPwithactivation_func_clamp_valuefor the clamped SwiGLU, and — since #2180 bumped megatron-core tob3393bbb— megatron-core's ownHyperConnectionModuleandTransformerBlockmHC stream handling.Components
Model-specific code lives in two packages, split by where it should eventually land upstream.
workers/megatron/mcore_ext/→ Megatron-LM (megatron.core)hyper_connection.py(49)RMSNormInputHyperConnectionModule: megatron-core'sHyperConnectionModulewith the input normalization overridden to GLM's standard RMSNorm,x * rsqrt(mean(x^2) + rms_norm_eps), instead of upstream'sx / (rms(x) + 1e-6).TransformerConfigcarries the input-norm knobs (mhc_norm_eps/mhc_norm_eps_inside_sqrt).mhc_transformer_layer.py(199)HyperConnectionTransformerLayerwith MoE MLP support — megatron-core's own mHC layer raisesNotImplementedErrorfor MoE sub-layers (mHC + MoE is only reachable there by wrapping MoE as aHybridStacklayer), which GLM-5.3-Flash needs on all but the first layer.TransformerBlockowns the block-boundary stream expand/contract; this layer implements only the per-sub-layer residual update. Rejects'mhc'inrecompute_modules, whose managers it does not thread.kda.py(327)KimiDeltaAttention: KDA linear attention (Kimi Linear / GLM-5.3-Flash) from TE column/row/duplicated linears, three depthwise causal convs, fp32A_log/dt_bias, flachunk_kdawith the in-kernel safe gate (lower_bound * sigmoid(exp(A_log) * (f + dt_bias))) and flaFusedRMSNormGatedoutput gate. TP shards heads (conv/A_log/dt_biasalong heads; low-rankf_a/g_aduplicated; replicatedo_normgrads summed across TP); packed thd viacu_seqlens;sharded_state_dictlike GDN. Reuses the GDN config fields pluskda_gate_lower_bound. No CP, no inference cache.experimental_attention_variant="kda"candidate next togdn/gdn2.workers/megatron/glm5_next/→ Megatron-Bridge (models/glm/)provider.py(55)Glm5NextModelProvider(MLAModelProvider): mHC fields named after megatron-core'sTransformerConfig, pluskda_gate_lower_bound,mhc_norm_eps*anddsa_indexer_kpool*. The attention pattern uses the genericlinear_attention_freqlist (1 = KDA, 0 = DSA).layer_specs.py(113)build_glm5_next_layer_spec: per-layer KDA-or-DSA × dense-or-MoEHyperConnectionTransformerLayerspecs, built from the publicget_dsa_module_spec_for_backend/get_moe_module_spec_for_backend/get_mlp_module_spec_for_backendhelpers.dsa.py(49)Glm5NextDSAttention(DSAttention): the k-pool indexer selects every visible token whenever a sequence has at mostindex_topktokens, so megatron-core's token-level indexer withdsa_indexer_topk=index_topkis exact in that regime and is reused; longer sequences raise instead of silently attending to a different subset.bridge.py(318)Glm5NextBridgeregistered forGlm5NextForConditionalGeneration→GPTModel,model_type="glm5_next": config mapping offtext_config(skipping thehead_dim=0RoPE width) and 1:1 parameter mappings undermodel.language_model.*— KDA (separateq/k/v_conv1d,A_log,dt_bias), mHC (hc_*_fn/hc_*_basereplicated,hc_*_scale[3]↔ threealpha_*scalars via a small custom mapping, the shape Megatron-Bridge's DeepSeek-V4 bridge already uses). All 3050 language-model tensors map; no dropped conversion tasks. The vision tower is not bridged (language_model_only=Truerequired).SkyRL glue
patches/megatron/patch_fa4_cute_import.py+workers/megatron/__init__.py: flash-attn 2.8.x's FA4flash_attn.cuteraisesAttributeErroragainst the cutlass DSL the pinned vLLM pulls in, which megatron-core'sexcept (ImportError, PackageNotFoundError)probe does not catch, soimport megatron.coreaborts. The guard marksflash_attn.cuteunavailable only when it is actually broken. Still required on torch 2.13. Delete once megatron-core's probe catchesExceptionor the pins agree.MegatronWorker.init_configs: megatron-core rejects mHC underrecompute_granularity="full", which is what SkyRL's defaultgradient_checkpointing=Trueresolves to. Its suggested alternative (selective recompute with'mhc'inrecompute_modules) needs the mHC recompute managers the layer above does not thread, so this downgrades to selective recompute of the remaining modules with a log line.model_bridges.py: imports the bridge to register it.test_megatron_models.py:glm-5.3-flash-4layer_h100_tp2_ep4(Megatron TP2 EP4 ETP1 → DP2, vLLM TP4 colocated,language_model_only, packed sequences,inference_only_init,max_num_seqs=512).test_glm5_next_modules.py(new, single GPU):RMSNormInputHyperConnectionModulevs HFGlm5NextTextHyperConnectionat the model's activation scale, andKimiDeltaAttentionon packed sequences vs per-sequence HFGlm5NextTextLinearAttention.supported_models.mdx,.claude/docs/backends/megatron.md.Dependencies
0.28.1rc1.dev359+g98ed0856f([Model] add GLM-5.3-Flash support vllm-project/vllm#53906, the GLM-5.3-Flash model). Built for CUDA 13 / torch 2.13, matching the torch pinned on main after [chore] Upgrade to torch 2.13 #2175, so its vendoredvllm.third_party.deep_gemmextension loads and the DSA indexer (which hard-requires DeepGEMM) works with no separately installeddeep_gemm. Verified on this branch:vllm.third_party.deep_gemm._Cimports,vllm.utils.deep_gemm.has_deep_gemm()is True, no standalonedeep_gemmpresent.vLLM 0.28.1 fallout handled:
vllm.entrypoints.openai.cli_argsmoved tovllm.entrypoints.launchers.cli_args(imported with a fallback);/inference/v1/generate, which SkyRL's generation client calls, is now gated behindVLLM_ENABLE_SCALE_OUT_ENDPOINTS=1(set beforebuild_app); vLLM registers a nativesharded_rdtweight-transfer engine, so the SkyRL shim no-ops and its unit test accepts either engine.Validation
Megatron vs HF
transformers5.16.1 (sdpa, bf16) logits on the 4-layer slice, 82-token GSM8K prompt+answer, packed thd input (standalone torchrun parity script).On this branch (torch 2.13, megatron-core
b3393bbb), HF reference recomputed on the same toolchain:For scale, #2156 measured HF-bf16 vs HF-fp32 on this same slice at 0.0071 KL / 95.7% argmax, so this sits at the bf16 noise floor. The slice predicts near-randomly (HF logprob per target token ≈ −14.15), which makes argmax agreement the noisiest of the three statistics — most positions are near-ties that flip on tiny numeric differences while KL barely moves.
test_glm5_next_modules.pyon the same tree: 2 passed — mHC exact to 4e-9 against the HF module, KDA within 0.5% of the reference mean on packed sequences.test_logprobs_matching_roundtrip[glm-5.3-flash-4layer_h100_tp2_ep4]on Anyscale 4×H100:and again in the H100 CI suite alongside the other model entries (
logprob diff mean: 0.063314,vLLM logprob diff mean: 0.250699). Both of those runs predate the merge with main — they were on torch 2.11 with the megatron bump applied. The suites need re-running on this branch's tree; earlier TP2/EP2 parity (KL 0.0066, argmax 98.8%) is from that same pre-merge toolchain.Known issue
test_logprobs_matching_roundtrip[qwen3.5-0.8b-dense_tp2]OOMs on 22 GiB L4 in the megatron_models suite: vLLM reserves 19.83 GiB for KV cache at the defaultgpu_memory_utilization=0.9, then flashinfer 0.6.18'stop_k_top_p_sampling_from_logitsneeds a ~970 MiB vocab-wide softmax during sampler warm-up. Introduced by the flashinfer bump here — that entry has no gmu override and passes on main's 0.6.16.post3. Needs either a gmu override for small cards orVLLM_USE_FLASHINFER_SAMPLER=0; not reproducible on H100 (80 GiB hides it).Limitations / follow-ups
index_topk(2048) tokens; longer sequences raise (glm5_next/dsa.py). Needs aDSAttentionhook letting an indexer own the top-k selection (pool scoring → expand to tokens → append the query's tail pool).indexer_typescontaining"shared"; the 4-layer slice is all"full") raises: megatron-core'sdsa_indexer_topk_freq/skip_topk_offsetarithmetic runs over all layers and would pick KDA layers as sources in this hybrid.layernorm/mlp/mhcrecompute.A_log/dt_bias/mHC params) but this PR exercises forward + weight sync only.Upstreaming plan
Megatron-LM (megatron-core)
HyperConnectionModule: addmhc_norm_eps/mhc_norm_eps_inside_sqrtfor a standard-RMSNorm input norm, with the fused-kernel path extended or guarded. Retiresmcore_ext/hyper_connection.py.HyperConnectionTransformerLayer: allow MoE MLP sub-layers (currentlyNotImplementedError). Retiresmcore_ext/mhc_transformer_layer.py.KimiDeltaAttentionasexperimental_attention_variant="kda"(module + spec next togated_delta_net,kda_gate_lower_boundinTransformerConfig). Retiresmcore_ext/kda.py.Megatron-Bridge
models/glm/glm5_next_{provider,spec,bridge}.py= this PR'sglm5_next/package, once (1)–(3) exist upstream — the bridge only needs its layer/provider imports repointed._HCAlphaMappingpattern (hereHyperConnectionScaleMapping) intoparam_mapping.pyso both bridges share it.Once megatron-core carries (1)–(3),
mcore_ext/is deleted andglm5_next/imports frommegatron.core; once Megatron-Bridge carries the bridge,glm5_next/goes too andmodel_bridges.pyloses one import.🤖 Generated with Claude Code