[megatron] feat: migrate fused logprob/entropy from GPTModel.forward monkey-patch to Megatron output_processor hook

[chengcuiping]

Motivation

verl/models/mcore/model_forward_fused.py currently replaces the entire GPTModel.forward (_fused_GPTModel_forward via patch_fused_forward) for one reason only: to smuggle the training-loop temperature into the _postprocess boundary, where linear_cross_entropy computes the fused logprob/entropy. The in-code TODO already flags this:

TODO: Currently we still need to patch forward because we need to pass temperature explicitly to self._postprocess.

Megatron-LM PR #4686 (core ≥ 0.18) added an output_processor / output_processor_context extension point at exactly that _postprocess boundary. Using it lets us delete the whole forward copy: temperature travels via output_processor_context, and a small callback runs the fused kernel at the hook. This is the downstream half of Megatron issue #4590.

Scope (single live Megatron fused path)

The migration targets verl's single live Megatron fused path. In current main the fused logprob/entropy path was consolidated into the new MegatronEngine (workers/engine/megatron/transformer_impl.py): model-build time calls patch_fused_forward, and forward time calls fused_forward_model_engine. The legacy fused_forward_model (get_mcore_forward_fused_fn) is now an unused export (no caller in-tree). Both the legacy and engine callers invoke the same monkey-patched GPTModel.forward, so migrating that single patch to the output_processor hook covers the live path, with a one-line temperature→output_processor_context change at the engine call site. Megatron's build_schedule_plan() / 1F1B postprocess path is tracked separately in #4590 and is not part of this PR.

Changes (single file: verl/models/mcore/model_forward_fused.py)

Gated behind env var VERL_FUSED_USE_OP_HOOK (default 0 = legacy patch path, byte-identical to today):

1. Add use_output_processor_hook(), FusedOutputProcessorContext (carries temperature), and the fused_output_processor callback (SP gather + linear_cross_entropy at the _postprocess boundary → CausalLMOutputForPPO).
2. patch_fused_forward: no-op in hook mode (forward is not replaced); runs a one-time capability probe (below).
3. unpatch_fused_forward: symmetric no-op in hook mode (nothing was backed up).
4. fused_forward_model (legacy) and fused_forward_model_engine (live engine path): in hook mode, drop temperature= from the model(...) call and pass output_processor=fused_output_processor, output_processor_context=FusedOutputProcessorContext(temperature=...) instead. The two call sites take the identical change.
5. Both paths coexist; the legacy _fused_GPTModel_forward is retained (default path).

Two review-worthy correctness points (handled)

• Callback keyword must be context: Megatron's _postprocess forwards context=output_processor_context, so the callback parameter is context (not output_processor_context); the call site uses the forward's own param name output_processor_context=. Different names, each correct.
• output_weight semantics are inverted vs the legacy patch: the hook passes output_weight=None when embeddings are not tied (real weight is in output_layer.weight), and the shared weight when tied — opposite of the legacy patch. The callback resolves weight = output_weight if output_weight is not None else output_layer.weight.

Compatibility

• Hook mode requires megatron-core >= 0.18.0 (the release shipping PR [trainer, fsdp, megatron] feat: support one_step_off_policy on Ascend NPU #4686's output_processor; merged to Megatron-Core main 2026-05-12 as 97f3bce88, first released in core_v0.18.0 on 2026-06-23).
• Fail-fast: hook mode validates capability at model-build time via _assert_output_processor_supported() — it inspects inspect.signature(GPTModel.forward) and, if the output_processor parameter is absent, raises a clear RuntimeError naming PR [trainer, fsdp, megatron] feat: support one_step_off_policy on Ascend NPU #4686 / megatron-core>=0.18.0, instead of a late opaque TypeError during the forward. (Signature introspection, not a version-string compare, so dev/local builds like 0.18.0+<sha> are handled correctly; the version number appears only in the error text.)
• Default off is a hard guarantee: with VERL_FUSED_USE_OP_HOOK unset, behavior is byte-identical to today; users who have not upgraded Megatron (< 0.18) are unaffected.
• Prerequisite (already merged, not part of this diff): Megatron-Core ≥ 0.18 support in verl's get_model landed in [megatron] fix: guard ModelType.encoder_and_decoder for Megatron-Core >= 0.18 compatibility #6927 ([megatron] fix: guard ModelType.encoder_and_decoder ..., merged 2026-07-06). This PR builds on top of it and does not re-include those changes.

Validation (numerically bit-identical / max_abs = 0)

Downstream A/B (hook vs legacy patch), same weights + same batch, two independent processes toggled by VERL_FUSED_USE_OP_HOOK, on Megatron 0.18.0+97f3bce88 / torch 2.9.1+cu128 / TE 2.10 / flash-attn 2.8.3, bf16:

Scenario	Coverage	log_probs / entropy (/pg_loss)
A/B synthetic batch	TP=1/2 × tied/non-tied (both output_weight branches) + SP gather	max_abs = 0
Regression via real get_model build chain	manual wrap vs real get_model	bit-identical
Real training-shape batch (classic form)	variable-length padding, THD padding_causal, GRPO advantages; incl. pg_loss	max_abs = 0
After installing vLLM	no contamination	bit-identical
Engine live path (real gsm8k tokens, nested/THD)	TP=2 tie=1 (and TP=1 tie=1), fused_forward_model_engine + build-time patch_fused_forward; log_probs + entropy	max_abs = 0

• The engine-path row exercises the actual live call form (nested/jagged inputs via preprocess_thd_engine, label = input_ids.clone()), not the legacy entry.
• temperature via output_processor_context, SP-gather timing, and the output_weight inversion are all verified.
• Memory: hook vs patch forward memory is identical (Δ = 0).
• Known-untested: fp8 padding (use_fp8_padding / config.fp8 ∈ {e4m3, hybrid}) was not exercised (all validation in bf16). The callback is fp8-agnostic, but the fp8 input-padding path is not covered here.

Notes

• Benefit: removes the "copy the entire forward" maintenance burden and automatically follows Megatron's native _preprocess/_postprocess evolution (the legacy patch's hard-coded preproc_output[:5] is a latent hazard on 0.18+).

[CLAassistant]

All committers have signed the CLA.

[gemini-code-assist]

Code Review

This pull request introduces an alternative hook-based path (output_processor) in verl/models/mcore/model_forward_fused.py for Megatron-Core, which avoids monkey-patching GPTModel.forward by passing a callback and context directly to the native forward method. The hook path is controlled via the VERL_FUSED_USE_OP_HOOK environment variable and includes capability checks. The feedback suggests removing a redundant assertion checking if weight is None in fused_output_processor, as the subsequent call to linear_cross_entropy would naturally fail with a clear error.

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]

Avoid adding an assertion to check if weight is None since its subsequent use in linear_cross_entropy will naturally fail with a clear error if it is indeed None.

Suggested change

	weight = output_weight if output_weight is not None else output_layer.weight
	assert weight is not None, "cannot resolve output weight for fused linear_cross_entropy"
	weight = output_weight if output_weight is not None else output_layer.weight

References

1. Avoid adding an assertion to check if a variable is None if its subsequent use in an operation would naturally fail with a clear error (e.g., TypeError) if it were None.

[chengcuiping]

Quick follow-up: besides the offline bit-identical A/B in the PR description, I ran a small live end-to-end smoke on Megatron-Core 0.18 to make sure the hook path survives the full RL loop, not just an isolated forward.

Setup was a dense Qwen2.5-0.5B on 4×A100, megatron engine (v1) with TP=2/PP=1, use_fused_kernels=True, full param/grad/optimizer offload, vLLM rollout, GRPO on gsm8k, 5 steps with a checkpoint at step 4. I ran it twice — once with VERL_FUSED_USE_OP_HOOK=0 (legacy patch) and once with =1 (the hook) — so the whole loop gets exercised both ways: rollout → weight resharding → fused logprob/entropy → optimizer update → distributed checkpoint.

Both paths completed all 5 steps and saved checkpoints cleanly, with no NaN/inf and comparable, sane metrics (entropy ~0.43–0.66, training log-ppl ~0.40–0.62 on both). The three integration points I was worried about all held up under the hook: rollout↔actor resharding (actor/rollout prob correlation ~0.9996 each step), offload interacting with the hook path, and dist-checkpoint save/load.

I also added a one-off build-time check to confirm the hook is genuinely what's running: under VERL_FUSED_USE_OP_HOOK=1 the GPTModel has no forward_backup on any rank (i.e. patch_fused_forward was a no-op and the native forward is untouched), while under =0 it does. So the live run really is going through output_processor, not the monkey-patch.

One honest caveat: the per-step metrics differ slightly between the two runs, but that's rollout non-reproducibility (vLLM across independent processes samples different trajectories), not the hook — the two runs simply don't see the same data. The exact numerical equivalence is what the offline A/B in the PR description covers (max_abs=0 on identical weights + batch); this live run is about integration, not re-proving the numerics.

Heads-up on one prerequisite unrelated to this PR: bringing up the megatron backend on Megatron-Core 0.18 currently also needs mbridge patched, because mbridge 0.15.1 (latest) still references the removed ModelType.encoder_and_decoder and crashes at model build — upstream of any forward, so it's orthogonal to this change. Filed separately (mbridge: ISEEKYAN/mbridge#155, verl dependency note: #6936).