Support GLM-5.2 744B-A40B

[yueming-yuan]

This PR is migrated from (THUDM/slime#2093).

Thanks @zhuzilin!

[gemini-code-assist]

Code Review

This pull request introduces support for GLM-5.2 (744B-A40B) with DSA cross-layer index sharing, adding model architecture modifications, training scripts, and end-to-end tests. The reviewer feedback focuses on optimizing and cleaning up the implementation: calculating the computing layer index in O(1) time instead of using a loop, refactoring duplicated topk_indices logic in the forward pass, removing an unnecessary loop when validating pipeline splits, and adding a descriptive error message to a NotImplementedError in the training script.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The computing layer index can be calculated directly using an O(1) mathematical formula instead of a while loop. This is more efficient and completely avoids any risk of infinite loops.

Suggested change

	layer = layer_number
	while is_skip_topk_layer(layer, skip_topk_offset, topk_freq):
	layer -= 1
	return layer
	if layer_number <= skip_topk_offset:
	return layer_number
	return skip_topk_offset + ((layer_number - skip_topk_offset) // topk_freq) * topk_freq

[gemini-code-assist]

The logic for computing topk_indices is duplicated three times in this method. We can refactor this to compute topk_indices once in the else block, and then conditionally store it in the holder if self.index_share is enabled. This significantly improves maintainability and readability.

        if self.skip_topk:
            # Cross-layer index sharing. The top-k holder lives on the per-microbatch
            # ``packed_seq_params`` object: it is closure-captured by Megatron's
            # activation-checkpoint ``custom_forward``, so the same instance is reused at
            # recompute time. That gives per-microbatch isolation (no cross-microbatch
            # clobber under PP 1F1B) AND recompute safety. A stage always starts on a
            # computing layer (asserted in ``get_glm5_spec``), so a skip layer's source is
            # always in-stage.
            holder = getattr(packed_seq_params, self._HOLDER_ATTR, None)
            if holder is None or self._source_layer not in holder:
                raise AssertionError(
                    "DSA index-share: skip layer "
                    f"(layer_number={self.layer_number}) needs the top-k of its source "
                    f"computing layer (layer_number={self._source_layer}), but that layer "
                    "did not run in this pipeline stage's forward. Cross-PP top-k sharing "
                    "is not supported; ensure every pipeline stage starts on a computing "
                    f"layer (index_topk_freq={self.index_topk_freq}, "
                    f"index_skip_topk_offset={self.skip_topk_offset}). "
                    f"Holder has layers {sorted(holder) if holder is not None else []}."
                )
            topk_indices = holder[self._source_layer]
        else:
            starts, ends = generate_varlen_mask_params(packed_seq_params.cu_seqlens_q)
            index_key = index_key.squeeze(1)
            head_weights = head_weights.unsqueeze(-1)
            starts = scatter_to_sequence_parallel_region(starts, group=parallel_state.get_context_parallel_group())
            ends = scatter_to_sequence_parallel_region(ends, group=parallel_state.get_context_parallel_group())
            _, topk_indices = fused_select_topk(index_query, index_key, head_weights, starts, ends)
            if self.index_share:
                holder = getattr(packed_seq_params, self._HOLDER_ATTR, None)
                if holder is None:
                    holder = {}
                    setattr(packed_seq_params, self._HOLDER_ATTR, holder)
                holder[self.layer_number] = topk_indices

[gemini-code-assist]

Since the check is only performed for local_id == 0, there is no need to loop over all layers. We can directly check the first layer of the stage, which is cleaner and more efficient.

Suggested change

	for local_id in range(num_layers_to_build):
	layer_number = local_id + layer_offset + 1 # Megatron layer_number is 1-indexed
	if local_id == 0 and is_skip_topk_layer(
	layer_number, config.index_skip_topk_offset, config.index_topk_freq
	):
	src = source_compute_layer(layer_number, config.index_skip_topk_offset, config.index_topk_freq)
	raise AssertionError(
	"DSA index-share pipeline split is invalid: this stage starts at global "
	f"layer_number={layer_number} which is a skip layer whose source computing "
	f"layer={src} is on a previous pipeline stage. Cross-layer top-k sharing does "
	"not cross PP boundaries. Choose a pipeline layout where every stage begins on "
	"a computing layer (index_topk_freq="
	f"{config.index_topk_freq}, index_skip_topk_offset={config.index_skip_topk_offset})."
	)
	if num_layers_to_build > 0:
	layer_number = layer_offset + 1 # Megatron layer_number is 1-indexed
	if is_skip_topk_layer(
	layer_number, config.index_skip_topk_offset, config.index_topk_freq
	):
	src = source_compute_layer(layer_number, config.index_skip_topk_offset, config.index_topk_freq)
	raise AssertionError(
	"DSA index-share pipeline split is invalid: this stage starts at global "
	f"layer_number={layer_number} which is a skip layer whose source computing "
	f"layer={src} is on a previous pipeline stage. Cross-layer top-k sharing does "
	"not cross PP boundaries. Choose a pipeline layout where every stage begins on "
	"a computing layer (index_topk_freq="
	f"{config.index_topk_freq}, index_skip_topk_offset={config.index_skip_topk_offset})."
	)

[gemini-code-assist]

Provide a descriptive error message when raising NotImplementedError to make debugging and troubleshooting easier.

Suggested change

	raise NotImplementedError
	raise NotImplementedError(f"Training on {args.num_nodes} nodes is not supported.")

[yushengsu-thu]

cc catches a potential issue:

Full-model checkpoint conversion will crash on the new assertion

run_glm5_2_744b_a40b.py → _prepare_megatron_ckpt (the non-pruned branch) keeps the GLM-5 conversion layout verbatim:

extra_args += (
"--pipeline-model-parallel-size 4 "
"--expert-model-parallel-size 32 "
"--decoder-last-pipeline-num-layers 18 "
)

Conversion runs tools/convert_hf_to_torch_dist.py, which builds the model via get_model(get_model_provider_func(args), ...) (line 110) → get_glm5_spec per pipeline stage. get_glm5_spec loads the GLM-5.2 HF config
(index_topk_freq=4 > 1), so the new assertion is active during conversion, not just training.

With PP=4 and last=18 over 78 layers, Megatron splits [20, 20, 20, 18] → stages start at 1-indexed layers 1, 21, 41, 61. For freq=4, offset=3:

• is_skip_topk_layer(21,3,4) → (18 % 4) = 2 ≠ 0 → skip
• 41 → (38 % 4)=2 → skip · 61 → (58 % 4)=2 → skip

So pp_ranks 1/2/3 each start on a skip layer and raise AssertionError("DSA index-share pipeline split is invalid…"). The documented full-model prepare step aborts.

This is almost certainly an oversight rather than intent: the author did account for the constraint everywhere else — the train layout (PP=8), and even the 5-layer convert path, whose comment explicitly says "nproc=1
also avoids the convert tool's PP auto-bump (which would split onto a skip layer)." The full-model convert split was just inherited from run_glm5_744b_a40b.py (where freq=1, so the assertion never fired).

Why it slips through CI: the 5-layer smoke test converts via the <= 5 branch (PP=1, single GPU), so the multi-stage split is never exercised.

Fix: use an index-share-valid conversion split. Either reuse the training layout (PP=8, first=14/last=16), or a tuned PP=4 split such as first=14/last=24 → [14,20,20,24] → starts 1,15,35,55 (all ≡ computing). Verify the
chosen split against is_skip_topk_layer for each stage start.

[yueming-yuan]

@yushengsu-thu good catch! fixed