fix(megatron): make linear CE fusion work through VLM wrappers

[dafu-wu]

What does this PR do ?

Makes linear CE fusion usable for checkpoints whose top-level module wraps the GPTModel (VLM wrappers), instead of failing in the first forward pass with an opaque TypeError.

Changelog

• nemo_rl/models/megatron/train.py
  • Add find_linear_ce_fusion_target(): returns the inner GPTModel when a wrapper hides it, and None when the invoked module already is the patched GPTModel (after Float16Module/DDP-style .module unwrapping) or when no GPTModel is reachable. The walk checks .language_model at each level before descending, and is depth-bounded.
  • Add _armed_linear_ce_fusion(): sets a per-instance flag on that inner model for the duration of the call and always restores it in a finally, so it cannot leak into later non-fused forwards.
  • model_forward() now passes the kwarg as before for unwrapped models, and arms the inner model instead when a wrapper is in the way.
• nemo_rl/distributed/model_utils.py: the patched forward honors the per-instance flag.
• nemo_rl/models/megatron/setup.py: keep a startup guard, but only reject models where no GPTModel is reachable at all -- those would silently never run the fused path.
• Tests for both the traversal and the arming lifecycle, including restoration when the forward raises.

Background

patch_gpt_model_forward_for_linear_ce_fusion() rebinds GPTModel.forward, and model_forward() passes return_logprobs_for_linear_ce_fusion to the module it invokes. When a checkpoint loads as a VLM wrapper, that module is not the patched GPTModel:

TypeError: Gemma4VLModel.forward() got an unexpected keyword argument
           'return_logprobs_for_linear_ce_fusion'

Gemma-4 is one such case: the HF checkpoint is Gemma4ForConditionalGeneration, so the invoked module is Gemma4VLModel and GPTModel is only its .language_model. Fused logprobs were therefore unusable there -- including for text-only runs, where the wrapper is otherwise inert. That matters at scale, since the fused path exists specifically to avoid materializing the [B, S, vocab] fp32 logits tensor.

Why arming rather than forwarding the kwarg

The wrapper already forwards everything the fused path needs (labels, decoder_input, packed_seq_params) to the inner model -- only the request could not get through. return_logprobs_for_linear_ce_fusion is introduced by this repo's monkey-patch, not part of the upstream GPTModel API, so teaching every wrapper to declare it would push that coupling into models NeMo-RL does not own. Arming the instance keeps it here.

Unwrapped models are unaffected: they keep receiving the kwarg directly.

GitHub Actions CI

A NVIDIA developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc) - No

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.