[erictang000]

Summary

Bumps megatron-bridge and megatron-core to latest and retires the patches whose upstream fix has landed.

Pins

• megatron-bridge c4067f50 -> 7171584b (v0.8.0); megatron-core 07147d69 -> ed9c2c8c (v0.21.0), Bridge's 3rdparty/Megatron-LM pin.
• peft 0.18.1 -> 0.21.0: Bridge now requires peft>=0.20.0,<=0.21.0. 0.21.0 resolves a string target_modules ("all-linear") against the model, which fixes the MoE regression behind the 0.18.1 pin ([ci] fix CI bugs introduced after latest uv sync #1763).
• uv.lock (made with uv 0.12.0) changes only megatron-bridge, megatron-core and peft.
• Transformer Engine stays at 2.19.0. Bridge builds against TE 2.20.2, but megatron-core adds no TE >= 2.20 version gates.

Patches retired (per skyrl/backends/skyrl_train/patches/megatron/README.md)

• Add KPool DSA indexer, NoPE MLA support, and FP8 precision controls for the indexer and MLA attention NVIDIA/Megatron-LM#7522 (k-pool DSA indexer) landed. Deleted mcore_ext/dsa_kpool.py, Glm5NextDSAIndexer and _forward_with_kpool_topk.
  • The GLM-5.3-Flash bridge sets dsa_indexer_weights_proj_output_dtype="fp32", which the vendored indexer hard-coded; upstream defaults to bf16.
  • Glm5NextDSAttention keeps only the kpool <= 1 long-sequence guard and the opt-in SKYRL_DSA_INDEXER_TP_SHARD, now a wrapper around megatron-core's _kpool_topk_from_score_chunks.
  • The k-pool tests are repointed at megatron-core and moved under glm5_next/.
• patch_checkpoint_without_output_frozen_inputs.py ([megatron] Recompute output-discarding checkpoints whose inputs need no grad #2541) removed. megatron-core rewrote CheckpointWithoutOutput (refactor(checkpoint): re-architect CheckpointWithoutOutput internals NVIDIA/Megatron-LM#8334): when no tensor input requires grad, checkpoint() now runs a plain call. On the new pin the patch also crashed megatron_worker import, because _load_args_from_ctx no longer exists. The test's frozen-input scenario gives the correct gradient unpatched.

Still carried (landed check fails at the new pins): KDA (#7054), mHC (#7521), DSA index share (#6793), sparse-MLA NoPE (#7617), MoE dispatcher probs, shared-expert LoRA TP (Bridge#6089), packed-per-expert sharded_state_dict (Bridge#6184), vision attention backend, and the GLM-5.3-Flash Bridge model (Bridge#6044).

• mHC-MoE is partly upstream: megatron-core's layer accepts MoE, but it chunks only in training and needs the not-yet-landed mHC input norm.

Other changes

• test_dsa_index_share_recompute stubs precision_aware_checkpoint, which replaced tensor_parallel.checkpoint / te_checkpoint in megatron.core.recompute.
• patches/megatron/README.md:
  • A "Reference values" table of the GLM-5.3-Flash roundtrip logprob diffs, per megatron-core / Bridge / vLLM pin and GPU. dependency_bumps.md and the GLM test rows point at it, and vLLM bumps should add a column too.
  • Fixed verification commands: the old -k missed the KDA/mHC tests and the LoRA row.
  • Documented the previously unlisted packed-per-expert patch.
  • Noted that concurrent vLLM roundtrip sessions on one box collide on port 8000.

Testing

Rebased on main 86c0bd69, which includes vLLM 0.32.0. Run on 4x B200.

GLM-5.3-Flash-4layer, mean logprob diff (threshold)	main (vLLM 0.32, old Megatron)	this PR
tp2_ep4 Megatron vs vLLM (< 0.1)	0.0634	0.0702, 0.0731 (two runs)
kpool_beyond_topk Megatron vs vLLM (< 0.1)	0.0553	0.0536
adapter zero / updated adapter (< 0.1)	0.0656 / 0.0653	0.0632 / 0.0618

On tp2_ep4 this PR is about 0.008 above main with vLLM 0.32; a repeat run reproduced it. It's well inside the threshold, and the same pins with vLLM 0.31 gave 0.0630. Not investigated further.

• GLM-5.3-Flash roundtrip rows vs vLLM 0.32 (test_megatron_models.py, test_megatron_lora_models.py, -k glm-5.3): 3 passed
• All Megatron patch GPU tests (gpu_ci/patches/megatron/), including k-pool kernels, the two-rank TP shard, and KDA/mHC vs HF
• CPU -m megatron suite: 113 passed
• CPU skyrl-train and fsdp suites: passed before the rebase (peft change)
• Every megatron.* import in skyrl/ and tests/ resolves against the new pins
• run_megatron_gpu_ci, run_megatron_gpu_ci_models, run_h100_gpu_ci: Bridge's refactors (GLM moved to HybridModel, Qwen3 MoE builder-backed) touch models not run locally

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request integrates the upstream landing of the k-pool DSA indexer in megatron-core by removing the locally vendored implementation, updating dependencies (including peft, megatron-bridge, and megatron-core), and refactoring Glm5NextDSAttention to wrap the upstream selection with opt-in tensor-parallel sharding. Feedback on the changes highlights two robust improvements in glm5_next/dsa.py: using ellipsis slicing for the attention mask to safely handle higher-dimensional tensors, and adding a None check on self.pg_collection.tp to prevent potential AttributeError in non-distributed environments.

[gemini-code-assist]

Using mask[..., q_lo:q_hi, :] is more robust than checking mask.ndim == 2 and conditionally slicing. If a 4D attention mask (e.g., [batch, heads, seq_len_q, seq_len_kv]) is ever passed, the current mask[:, q_lo:q_hi] would incorrectly slice the heads dimension instead of the query dimension. Using ellipsis (...) automatically handles any number of leading dimensions (2D, 3D, or 4D) correctly.

Suggested change

	if mask is not None:
	mask = mask[q_lo:q_hi] if mask.ndim == 2 else mask[:, q_lo:q_hi]
	if mask is not None:
	mask = mask[..., q_lo:q_hi, :]

[gemini-code-assist]

If self.pg_collection.tp is None (e.g., in non-distributed or single-GPU environments), calling tp_group.size() will raise an AttributeError. Adding a None check ensures robustness.

Suggested change

	tp_group = self.pg_collection.tp
	if tp_group.size() <= 1:
	return None
	tp_group = self.pg_collection.tp
	if tp_group is None or tp_group.size() <= 1:
	return None

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The changes reach Megatron-backed GLM training and shared PEFT adapter workflows.

The PR appears safe to merge, with non-blocking fixes for sharding test coverage and the documented test command.

Findings

1.  Sharding decision goes untested ▶
2.  Test command names missing file ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[GLM bridge] --> B[Upstream k-pool indexer with FP32 projection]
  B --> C[GLM attention wrapper]
  C --> D{Sharding enabled and queries replicated?}
  D -->|No| E[Upstream scoring]
  D -->|Yes| F[Score a slice on each tensor-parallel rank]
  F --> G[Gather selected pools]
  E --> H[Sparse attention]
  G --> H

Loading

_{Reviews (1) · Last reviewed commit: "[megatron] Bump megatron-bridge to 71715..." · Reviewed by Greptile}

[greptile-apps]

Sharding decision goes untested

The updated test calls make_sharded_kpool_topk directly. It never checks the new _indexer_shard_group conditions or the hook installed by Glm5NextDSAttention.forward. This test can therefore pass even if training silently skips the requested sharding.

Add tests for replicated queries and sequence-local queries, and check whether the helper is called in each case.

[greptile-apps]

Test command names missing file

The updated run command points to test_mcore_dsa.py, which does not exist. Replace it with test_dsa_kpool.py so copying the command runs these checks instead of failing with a file-not-found error.

Suggested change

	tests/backends/skyrl_train/gpu/gpu_ci/patches/megatron/glm5_next/test_mcore_dsa.py
	tests/backends/skyrl_train/gpu/gpu_ci/patches/megatron/glm5_next/test_dsa_kpool.py