[megatron] fix: make MTP work with recompute_granularity=full on megatron-core 0.18.x

[gaohongkui]

What does this PR do?

Makes Multi-Token Prediction usable together with recompute_granularity='full' on every
released megatron-core, and makes the existing signature-based branch in
patch_mtp_layer_checkpointed_forward visible instead of silent.

The bug

megatron-core 0.18.x is internally inconsistent: MultiTokenPredictionLayer.forward calls

self._checkpointed_forward(..., padding_mask=padding_mask, ...)

while _checkpointed_forward itself does not declare padding_mask. Any MTP run with
recompute_granularity == 'full' dies at the first forward:

TypeError: MultiTokenPredictionLayer._checkpointed_forward() got an unexpected
           keyword argument 'padding_mask'

Upstream cause: two landings that do not compose — NVIDIA/Megatron-LM#2645 added the call-site
kwarg, #4593 refactored the method without it. Tracked as NVIDIA/Megatron-LM#4933, still open.

It is fixed on megatron-core main, but core_v0.18.2 was tagged 2026-07-20 — two months after
that issue was filed — and is still the newest release, so every released megatron-core hits
this. I have proposed a backport as NVIDIA/Megatron-LM#6367; until a release carries it, verl
users cannot combine MTP with full activation recomputation without patching megatron themselves.

Why the existing patch does not cover it

patch_mtp_layer_checkpointed_forward skips any layer whose _checkpointed_forward does not
start with forward_func. That gate matches megatron-core 0.14–0.17 only; 0.18+ renamed the
first parameter to hidden_states, so on every 0.18+ install the patch silently does nothing —
when target_layers is non-empty but patched_count == 0, neither log line is printed, so there
is no way to tell whether it applied.

Skipping is in fact correct on 0.18+: megatron now keeps non-tensor args out of the
checkpoint natively, via its own custom_forward closure. But it should say so rather than be
inferred from missing output.

Changes

• _patch_padding_mask_kwarg() — when the layer's _checkpointed_forward lacks padding_mask,
  rebind it to accept and drop the kwarg. padding_mask is dropped rather than forwarded because
  there is no parameter to forward it to on these versions, and verl never constructs one, so it
  is always None in practice. A non-None value raises NotImplementedError rather than being
  silently ignored — ignoring it would treat padded positions as real tokens. Idempotent per
  layer (after installation the shim's own signature declares padding_mask, so a second call is
  a no-op).
• patch_mtp_layer_checkpointed_forward() — log the skip-by-signature and shim counts.

Style follows the existing precedent a few lines above in the same file, which already probes
signature(self.mtp.forward) for padding_mask before passing it.

Once a megatron-core release carries #6367, the signature check makes this shim skip itself
automatically, and it can be deleted.

Not a duplicate

Searched verl issues and PRs for patch_mtp_layer_checkpointed_forward,
mtp_patch checkpointed_forward, forward_func signature mtp — no hits. The forward_func gate
is still present on main (verl/models/mcore/mtp_patch.py), so nothing has superseded this.

Test

tests/models/test_mtp_checkpointed_forward_shim_on_cpu.py — 5 CPU cases, no GPU, no distributed
init, covering all three megatron-core signature generations:

pytest tests/models/test_mtp_checkpointed_forward_shim_on_cpu.py -q

• megatron-core main signature (declares padding_mask) → shim not installed, kwarg still
  reaches the method
• 0.18.x signature → reproduces the bare TypeError first, then asserts the shim accepts
  padding_mask=None and forwards the remaining args unchanged
• non-None mask → NotImplementedError instead of silent corruption
• 0.14–0.17 forward_func signature → left to the existing recompute patch
• shim is idempotent

All 5 pass locally.

Also exercised end to end on 4x8 H20 with Qwen3.6-35B-A3B (MoE, 40 layers, 256 experts),
TP2/PP2/CP4/EP8, 64K context, mtp_num_layers=1, recompute_granularity='full', megatron-core
0.18.2 + torch 2.11: without the shim training aborts at the first forward with the TypeError
above; with it, MTP trains and mtp_losses/* is reported normally (30 steps, loss 0.447 → 0.318).

AI assistance (Claude) was used for the investigation and to draft the patch and tests; the diff
was reviewed line by line before submitting.

[HollowMan6]

#7328 should have already fixed this (bug 2), can you try the main branch again?

[gaohongkui]

Thanks for the pointer — I checked, and I think #7328 fixes a different failure. I've also rebased this branch onto current main (now including #7328 at 98be7e97); it rebased cleanly and the crash path is untouched.

The two are different failure modes:

	#7328 bug 2	this PR
symptom	async CUDA illegal memory access, surfacing at a later sync point	synchronous TypeError at the Python call
mechanism	use-after-free — resize_(0) on a checkpoint input whose storage a pending backward still reads	signature mismatch — a kwarg the method does not declare
site	apply_patch_megatron_recomputation_backward in patch.py	MultiTokenPredictionLayer._checkpointed_forward
trigger	mHC + MTP (shared storage via torch.chunk view)	megatron-core 0.18.x + recompute_granularity=full

Why it still reproduces on main. In megatron-core 0.18.0 and 0.18.2, multi_token_prediction.py declares:

1093:  def _checkpointed_forward(
1095:      hidden_states: Tensor,
1096:      decoder_input: Tensor,
1097:      attention_mask: Optional[Tensor] = None,
1098:      context: Optional[Tensor] = None,
1099:      context_mask: Optional[Tensor] = None,
       # no padding_mask parameter

and then calls it with one anyway:

1301:  hidden_states = self._checkpointed_forward(
1305:      padding_mask=padding_mask,

verl's existing patch cannot intercept that, because of the guard in patch_mtp_layer_checkpointed_forward (still present on main):

if not params or params[0] != "forward_func":
    continue

On 0.18.x params[0] == "hidden_states", so the layer is skipped and the patch is never installed. _patch_padding_mask_kwarg appears 0 times on main, and #7328's mtp_patch.py diff only adds the mhc_multistream parameter and forwards it into mtp_kwargs — it does not touch _checkpointed_forward.

Note that main does already guard padding_mask at two other call sites with exactly this pattern — self.mtp.forward (mtp_patch.py L107-110) and layer._get_embeddings (L259) — so this PR is applying the established house pattern to the one remaining site.

Cheapest way to check without a cluster: this PR ships tests/models/test_mtp_checkpointed_forward_shim_on_cpu.py, which reproduces the TypeError on CPU with a stand-in object matching the 0.18.x signature, and covers the 0.14–0.17, 0.18.x and megatron-main signature generations plus the non-None mask case. It runs in seconds, no GPU or distributed init.

Happy to also run the full multi-node MTP + recompute_granularity=full repro on top of current main if you'd like that on the record — just say so.

[HollowMan6]

I was using the latest commit of megatron-lm dev branch and can't reproduce this: https://github.com/NVIDIA/Megatron-LM/tree/dev Is this caused by some specific commit in megatron-lm main?

[gaohongkui]

Thanks for the approval. CI ran on the rebase and I've triaged the failures — one was mine and is now fixed in 07c6b79f.

cpu_unit_tests — mine, fixed. My own test was wrong:

test_shim_not_installed_when_padding_mask_already_supported
assert layer._checkpointed_forward is before
AssertionError: assert _checkpointed_forward is _checkpointed_forward

Attribute access on a method that has not been rebound builds a fresh bound-method object every time, so is fails even when the shim correctly did nothing. Since the shim rebinds by assigning an instance attribute, "was not rebound" is now expressed directly as "_checkpointed_forward" not in vars(layer). The idempotency test is unaffected — after patching the attribute holds a plain function, so is is valid there.

I verified the fix by extracting the real _patch_padding_mask_kwarg source from mtp_patch.py with ast.get_source_segment and running all five test cases against it: 5/5 pass. (I can't run the file directly on my machine — pytest.importorskip("megatron.core") makes it skip silently, which is exactly why I didn't catch this before pushing. My mistake.)

pre-commit (3.12) — not from this PR. It fails on tests/trainer/ppo/test_reinforce_pp_multiturn_on_cpu.py (a ruff-format diff on a torch.tensor([...]) literal). This PR touches only verl/models/mcore/mtp_patch.py and tests/models/test_mtp_checkpointed_forward_shim_on_cpu.py, so that file is outside the diff — it looks like it came in with 8a1bf6d5 (#7300) at main's HEAD.

The four e2e_ppo_trainer_* failures — I could not determine the cause. I don't want to hand-wave them as unrelated, so to be precise about what I do and don't know: on this same run e2e_ppo_trainer_megatron-deepseek reports both a fail and a pass across matrix instances, and megatron-qwen3 shows fail/fail/pending, which suggests non-determinism rather than a deterministic break. I read the failing NPU job's log and could not find the actual error — it is dominated by transformers deprecation-alias noise. If you can tell whether these are known-flaky right now, I'm happy to act on it; and if you suspect they are mine, I'll dig further.

For what it's worth on blast radius: the shim only rebinds _checkpointed_forward on layers where params[0] != "forward_func" and the signature lacks padding_mask — i.e. megatron-core 0.18.x MTP layers only. On megatron-core main (which declares the parameter) and on 0.14–0.17 (forward_func first) it is a no-op, which is what tests 1 and 4 pin down.

[gaohongkui]

I was using the latest commit of megatron-lm dev branch and can't reproduce this. Is this caused by some specific commit in megatron-lm main?

That matches what I'd expect: dev and main already declare the parameter, the released 0.18.x tags do not. So it is not caused by a commit in main — main is the side that is already fine, and that fix has not made it into a release yet.

I checked the signature of MultiTokenPredictionLayer._checkpointed_forward across refs by parsing the file with ast (rather than grepping, which is easy to get wrong here — the call site three lines later also contains the string):

ref	first param	padding_mask in signature	padding_mask= passed at the call site	outcome
core_v0.18.0 (tag)	hidden_states	no	yes	TypeError
core_v0.18.2 (tag, newest release)	hidden_states	no	yes	TypeError
core_r0.18.0 (release branch)	hidden_states	no	yes	TypeError
upstream/main @6518b75ec (2026-08-09)	hidden_states	yes (L1296)	yes	fine
upstream/dev @43124b60c (2026-08-08)	hidden_states	yes	yes	fine

So the affected population is anyone on a released megatron-core 0.18.x — which is what pip install megatron-core gives you today, core_v0.18.2 being the newest (2026-07-20). Tracked upstream as NVIDIA/Megatron-LM#4933, and I opened NVIDIA/Megatron-LM#6367 to backport the declaration onto core_r0.18.0; it is currently blocked on copy-pr-bot needing an internal vetter.

I deliberately won't claim which commit fixed it on main — my clone of those branches is shallow, so I can't attribute it, and I'd rather not guess.

That is also why the shim self-disables rather than always patching: it only rebinds when the signature lacks padding_mask, so on dev/main it is a no-op (test 1) and on 0.14–0.17 it is left to the existing forward_func patch (test 4). If you'd prefer this to be gated on a version check instead of signature introspection, I'm happy to change it — though introspection is what mtp_patch.py already uses for self.mtp.forward (L107-110) and layer._get_embeddings (L259).

---

Separately, on CI: pre-commit (3.12) is failing on tests/trainer/ppo/test_reinforce_pp_multiturn_on_cpu.py, which is not in this PR's diff (this PR touches only verl/models/mcore/mtp_patch.py and tests/models/test_mtp_checkpointed_forward_shim_on_cpu.py). It is a ruff-format diff on the torch.tensor([...]) literals introduced with 8a1bf6d5 (#7300), and it reproduces on the current main file, so it is not something my rebase caused. I don't see an open PR fixing it. I'd rather not touch an unrelated test file inside this PR without your say-so — but if you'd like, say the word and I'll include the one-file ruff format fix here so this can go green.

[gaohongkui]

padding_mask entered the _checkpointed_forward signature in ffd66a3e6 ("Roll input IDs for MTP labels", NVIDIA/Megatron-LM#3457, 2026-06-03): 4c6360260 (2026-05-22) still lacks it, ffd66a3e6 is the first commit that has it. That commit threads the parameter through the signature, custom_forward, and both checkpoint_handler() entry points — and adds a test called test_packed_sequences_with_full_recompute. So upstream megatron treats packed sequences + recompute_granularity=full as a combination that must work, and fixed it deliberately rather than incidentally.

It just hasn't shipped: core_v0.18.0 (2026-06-22) and core_v0.18.2 (2026-07-20) are the only 0.18.x releases and neither has it.

So the one question that decides this PR is yours to answer: does verl intend to work against released megatron-core 0.18.x? If yes, this shim covers the gap until a release picks up ffd66a3e6, and it disables itself the moment one does. If verl effectively targets megatron-core dev/main, then it isn't needed and I'll close it — no hard feelings, and the CPU test would still be worth nothing to you in that case, so I'd close the whole thing rather than leave dead weight.

[HollowMan6]

Okay, if NVIDIA/Megatron-LM#3457 can fix all these, then this PR is not needed, just need to note that the Megatron fixing PR is needed

[gaohongkui]

Understood, and agreed — I've opened #7346 with the doc note instead.

One thing worth flagging while doing it: docs/advance/mtp.md pins megatron dev at 23e092f41 (2025-12-09), which predates #3457 by about six months. So following that doc as written still hits this crash with recompute_granularity=full. #7346 adds the constraint to that bullet; I left the pin itself alone since I haven't validated MTP + CP on a newer dev commit, but bumping it would be simpler if you have one you trust.

Thanks for the review — closing this one out was the right call.