Repository navigation
GLM 5.3 Flash RL Support - #2786
Merged
Merged
Conversation
… spec) - kda.py: Glm5NextKDA wrapping fla.ops.kda.chunk_kda behind HuggingfaceAttention (hybrid CP), safe gate g = gate_lower_bound * sigmoid(exp(A_log) * (f + dt_bias)) computed fp32 outside the kernel by the pure function kda_gate; packed qkv short conv; fp32 A_log/dt_bias via mark_keep_in_fp32. - dsa.py: no-rope absorbed sparse MLA (qk_pos_emb_head_dim 0, plain 1/sqrt(256) scale), kpool compression params on the module, zero 64-dim tail padding for the tilelang SparseMLA kernels, DSA-ordinal indexer replay streams. - ops/kpool_indexer.py: torch reference of sglang's kpool selection (learned softmax 4->1 pooling, pool top-k via the tilelang indexer fwd, short-seq shortcut, deterministic always-select-tail reconstruction), replay-manager hooked at the token-expanded pool budget. - glm5_next.py: layer_types dispatch on get_gpt_decoder_block_spec with Megatron-native mHC enabled at spec time (no CLI flags exist), mean output contract patch + hc_head_* demotion, kpool/indexer config patching.
Glm5NextBridge extends GlmMoeDsaBridge: text_config unwrap, KDA weight names (1:1 plus the packed conv1d concat), kpool compress gate/ape, hc_* tensors with the [3]-scale split into alpha_pre/post/res, fp32-source dtype preservation, and a bypass of the DeepseekV32 indexer rope-interleave swap (qk_rope is 0). MTP stays unmapped/untrained. Register model_type glm5_next and mark it DSA in miles.utils.hf_config.is_dsa.
Inverse of the bridge for the weight-update direction: KDA 1:1 names, packed conv1d split into q/k/v_conv1d, indexer + kpool names handled locally so the deepseekv3 rope-interleave half-swap can never run, and hc_*_scale reassembled from the three alpha parameters via a per-(layer, site) buffer. Everything DeepseekV3-shaped delegates to convert_deepseekv3_to_hf.
glm5.3-flash.py exposes model_args(nlayers=45) with the sigmoid-router MoE recipe (288 experts, top-8, routed scaling 2.5, 3 dense lead layers), MLA dims with qk-pos-emb-head-dim 0, and the glm5_next spec pointer; the 8-layer cut derives via load_sibling_model_args. run_glm5_3_flash.py: typer launcher with GLM-5.3-Flash (8x4, TP8 PP4 EP16 placeholder) and GLM-5.3-Flash-8layer (2x4, TP2 PP2 EP2) topologies, dapo-math-17k GRPO recipe, optional routing-replay via --enable-r3.
test_kda_wrapper.py (GPU): safe-gate math vs fla's fused_kda_gate, chunk_kda fwd/bwd smoke through Glm5NextKDA, and a bit-identity double-run determinism check. test_bridge_names.py (CPU): generates the expected 8-layer HF tensor list and audits Glm5NextBridge's name mapping for completeness both ways.
…der UAF) Same fix as the qwen4exp branch: python-subclass collective overrides route through pybind trampolines whose returned Work is use-after-freed at deferred grad-sync wait; base-class C++ dispatch avoids the wrapper.
…next transformers has no native glm5_next; a thin PretrainedConfig subclass gives the spec and bridge attribute access to the raw json (nested sub-configs included) without importing sglang.
The GLM-5.3 sglang tree makes resolved ServerArgs read-only; assigning an unchanged host tripped the guard on every engine init.
The container flashinfer lacks the NoPE sparse kwarg the trtllm decode path needs; tilelang on CUDA requires a bfloat16 KV cache.
The sglang capture host buffer is zero-initialized, so an engine whose MoE runner bypasses the topk capture hook (e.g. flashinfer_trtllm on SM100, the auto-resolved default for GLM-5.3 / DeepSeek-family archs) returns an all-zero routed_experts payload. Replaying it routes every token to expert 0 x topk, so the Megatron dispatcher's routing_map has one active expert per token while num_out_tokens stays tokens*topk, and the dispatch alltoall crashes with 'Split sizes doesn't match total dim 0 size' - far from the actual cause. Assert at payload decode instead, in both extraction paths, with a message naming the likely backend misconfiguration. (The sglang-side fix downgrades such backends when --enable-return-routed-experts is set.)
…norm_eps) Megatron's native_proj_rms computes r = 1/(rms(x) + 1e-6), which at GLM-5.3's hidden scale (mean-square down to ~5e-6 on the residual streams) is numerically an eps-free norm. The HF/sglang reference computes r = rsqrt(mean(x^2) + 1e-5) with the checkpoint's rms_norm_eps inside the sqrt, where it dominates the denominator for small tokens. The resulting h_pre/h_post/h_res differed by tens of percent per token, and offline layerwise parity (single prompt, 3328 tokens, TP1) showed this to be the dominant train-vs-serve divergence: it is created at layer 0 and carried through every layer. _patch_reference_proj_rms swaps the per-instance op and eps to the reference convention at spec-build time, glm5_next-only, same pattern as the mean output-contract patch. Offline parity: train/serve next-token logprob abs diff 0.257 -> 0.029 (mean), final-hidden rel err 0.122 -> 0.024, layers 0-2 to bf16-noise level; actual expert-selection overlap 98-99%.
The training side (kpool_indexer replay hook, DSA-ordinal stream_idx in Glm5NextDSAAttention, --use-rollout-indexer-replay plumbing) was already in place; this wires the last pieces now that sglang-miles captures the kpool top-k: - rollout_indexer_topk_num_streams: derive the expected indexer stream count from the HF config (num_hidden_layers minus KDA layers on hybrid models) and assert the engine's indexer_topk_num_layers matches, so a rollout engine that counts KDA layers into the payload fails fast instead of replaying streams onto the wrong layers. - run_glm5_3_flash.py: --enable-indexer-replay adds --use-rollout-indexer-replay (kept separate from --enable-r3's routing replay, mirroring glm5.2).
…ayloads The GLM-5.3 e2e died three minutes into the run at fill_replay_data with 'rollout_indexer_topk is required in rollout_data' because the rollout router (sgl-router-for-miles) silently stripped the return_indexer_topk body field from /generate requests (its typed GenerateRequest spec knows return_routed_experts but not return_indexer_topk, and drops unknown keys). Raise at convert_samples_to_train_data instead, on the rollout manager right after generation, with a message naming the proxy-strip failure mode for both replay streams. Eval flows never reach this conversion, so requests against capture-disabled eval engines stay fine.
Trim narrative module docstrings to short constraint statements, drop standalone comments, shorten fail-fast assert messages to one sentence, and import RouterGatingLinearFunction by name instead of an alias. No behavior change.
Pooling (per-dim softmax gate + ape) becomes one triton kernel over an oversized pool grid with in-kernel sequence binary search, bitwise-equal to the torch reference (sequential fp32 reductions, libdevice exp/div_rn/ mul_rn) and free of the per-sequence python loop and its host sync. Token expansion + tail append + padding fuse into a single kernel launch when indexer replay is off, with the staged topk_fn + append_tail_and_pad pair kept byte-identical for the replay paths. Scoring stays on the fp32 tilelang lighting indexer and the pool top-k on torch.topk (37% of the remaining path; a triton radix select is not clearly better and torch.topk keeps tie order shared with the reference). GB300 (job 2708): pooled keys 100% bitwise vs torch, end-to-end selection byte-equal on every case incl. non-divisible and sub-kpool sequences, bit-deterministic over repeated runs, boundary-tie rate 0.016%; 1.89x at T=8192 single-seq and 2.56x packed-mixed vs the torch path. The torch reference now lives in tests/glm5_next/test_kpool_indexer.py as the oracle.
Zhichenzzz
requested review from
fzyzcjy,
guapisolo,
maocheng23 and
yueming-yuan
as code owners
August 27, 2026 19:10
"Gloo" in name.title() implies "GLOO" in name.upper() for ASCII class names, so the first operand of the or never changes the result.
Each literal was one flag group split into adjacent pieces on a single line, which reads like separate arguments. The AST of both files is unchanged, so the argv is byte-identical.
Keep the 16x4 / 8x4 full-model layouts (EP now follows the GPU count) and the 8-GPU 4-layer slice; drop the 6x4 and 8-layer paths, the indexer-replay switch, the dead SGLANG_SKIP_CHECKPOINT_LOAD_CHECK env, and the hardcoded /tmp cache and offload paths. run_id follows U.create_run_id() like the other launchers.
HJSang
reviewed
Sep 24, 2026
| q_compressed = q_compressed.detach() | ||
| hidden_states = hidden_states.detach() | ||
|
|
||
| index_q, _ = self.wq_b(q_compressed) |
Collaborator
There was a problem hiding this comment.
q_compressed here is un-normalized: q_layernorm is IdentityOp and the q RMSNorm (HF q_a_layernorm) is fused inside linear_q_up_proj. The reference feeds the indexer the normalized tensor. In HF modeling_glm5_next.py that's q_resid = self.q_a_layernorm(self.q_a_proj(x)), then self.wq_b(q_resid), and SGLang's forward_mla / forward_mha do the same. The 1/rms factor doesn't change the top-k, but γ does. γ is mostly near-uniform, with a 4–9% spread, but some channels are near zero (min 0.09 in layer 23, −0.008 in layer 43), so the trainer can select differently from the engine past index_topk. Suggest mirroring the kv path:
q_normed = F.rms_norm(q_compressed.float(), (q_compressed.shape[-1],),
weight=self.linear_q_up_proj.layer_norm_weight.float(),
eps=self.config.layernorm_epsilon).to(q_compressed.dtype)
index_q, _ = self.wq_b(q_normed.detach())The GLM-5 base (glm5.py) has the same pattern.
Collaborator
|
LGTH, please check codex comment (the first one is because of my just-merged refactor; only need to solve this naming conflict. The second one might need checking) [P1] Import the existing HF config module — glm5_next.py:15 |
| config.glm5_next_full_attn_layers = full_attn_layers(text_config) | ||
|
|
||
|
|
||
| def _patch_mean_output_contract() -> None: |
Collaborator
There was a problem hiding this comment.
can we use the hc_head_contraction in Megatron#89 to avoid the patch here?
Collaborator
There was a problem hiding this comment.
Because I'm not sure about the context of this change, qq here: did we check that the changes here are correct and needed? I see Qwen 3.8-next's PR also has this change; can we make it a separate PR?
main moved miles/utils/hf_config.py to miles/utils/hf_utils/config.py (#3620).
The q_a_layernorm is fused into linear_q_up_proj, so wq_b was reading the un-normed q_a_proj output while sglang's kpool indexer reads q_a_layernorm(q).
sglang clamps gate to <= 10 and up to [-10, 10] in the dense, shared and routed MLPs; the trainer ran plain SwiGLU.
…n slot Megatron #89's slot replaces the monkeypatch of learned_output_contract and TransformerBlock.__init__; hc_head_* are no longer allocated at all.
main's _forward_remaining_collectives (#2714) already covers the torch 2.13 backend-map gap; GLM 4-layer runs at DP2 with non-zero grads complete on main's version.
yueming-yuan
approved these changes
Sep 24, 2026
Routing replay needs materialized topk ids; sglang now rejects the fused TRT-LLM runner that auto resolves to on SM100 instead of downgrading it.
This was referenced Sep 30, 2026
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
4 participants
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.
RL training support for zai-org/GLM-5.3-Flash (model_type
glm5_next): 45 layers = 34 KDA linear-attention + 11 DSA sparse-attention layers, 288-expert MoE (sigmoid top-8), mHC hyper-connections, NoPE MLA, kpool-compressed lightning indexer. docker image is indocker pull radixark/miles:glm53nextmiles_plugins/models/glm5_next/): per-layer KDA/DSA block spec — KDA on flachunk_kda+fused_kda_gate, DSA as no-rope absorbed SparseMLA with a triton kpool indexer, mHC matching the GLM-5.3 reference conventions.miles_plugins/mbridge/glm5_next.py): HF<->Megatron weight mapping; conversion and per-rollout weight sync verified with--check-weight-update-equal.