fix(megatron): Handle MTP position_ids correctly

[tdene]

What does this PR do ?

Fixes currently failures with MTP models.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

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 run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[tdene]

/ok to test 4be6ca0

[cspades]

The fix makes sense for the HybridModel MTP assert criterion, we need to pass position_ids now.

For the if condition, I think model_slices_context_parallel_inputs essentially refers to self._apply_context_parallel_sharding which shards the expanded (i.e. after the forward (training) or the wrapper/engine (_forward, but MInf does not support CP>1 rn) injects the encoded multimodal tokens) sequence later cause we doesn't know how many image tokens we will insert into the LM decoder input. A lot of models don't require position_ids but Omni or other models like it which handle the CP sharding (at least for BSHD, not clear for THD) would want this passed.

[tdene]

Team review (5 specialist agents + an adversarial pass) of 4be6ca0. The fix is correct and minimal: the model-derived mtp_enabled gate is what makes logprob/top-k forwards work (a mask-keyed or cfg-keyed gate would not have — 37/39 exemplar YAMLs lack mtp_num_layers and Bridge enables MTP from the HF checkpoint); the full-row-vs-CP-local choice mirrors input_ids/mtp_loss_mask and is required by Omni's same-width check (modeling_nemotron_omni.py L699-701); the CP=2 zigzag layout matches what MLM's _roll_tensor_packed_seq assumes (checked with the real packer on CPU); and GPTModel numerics cannot change — positions only feed the embedding (ignored unless learned_absolute, language_model_embedding.py L47) and the MTP roll, and mrope models stay on the untouched self-packing branch. 6 inline comments: stale docstring, single-source mtp_enabled, a CP=2 real-packer test, the merge gate for the never-run 2-GPU test, and two handoff assertions.

PR description — the template is unfilled. Suggested text (edit to taste):

What broke. #4022 bumped Megatron-Bridge to 5ed9799 → Megatron-LM 1e7598cbf, which contains NVIDIA/Megatron-LM#6912. HybridModel.forward now asserts position_ids is not None whenever its MTP block runs (hybrid_model.py:594), and it runs on every non-inference forward because NeMo-RL never passes compute_mtp_loss=False. The packed-sequence path always passed position_ids=None, so every hybrid MTP model with sequence packing died in get_logprobs → model_forward — the nightly grpo-nemotron3-super-120BA12B-16n8g-megatron has been red since. GPTModel tolerates None (roll_tensor(None) is a no-op), so only hybrids were affected.
Fix 1 — data.py. When the model has MTP layers (mtp_enabled, from model.config.mtp_num_layers > 0), pack per-sample arange positions with the same _pack_sequences_for_megatron call and pad arguments as input_ids, so positions reset per packed segment and follow the same CP layout (full THD row for models that CP-slice their own inputs, CP-local shard otherwise). Gated on the model rather than on mtp_loss_mask because logprob / top-k passes carry no mask yet still run the MTP block.
Fix 2 — train.py. model_forward dropped position_ids for every multimodal batch (added in #1115 without a recorded rationale; VLM wrappers derive their own). Nemotron Omni (model_slices_context_parallel_inputs=True) CP-slices caller position_ids with the tokens and its MTP block asserts them, so they are kept for that model.
Numerics. Unchanged for non-MTP models (position_ids stays None under packing). MTP models — including any GPTModel with mtp_num_layers > 0 — now receive per-segment positions where they got None; these are only rolled by the MTP block and ignored by RoPE/NoPE embeddings, so no loss change is expected.
Validation. First commit: 5 training steps on the 16-node Super recipe (internal CI, pre-rebase). Unit: test_megatron_data.py (per-segment positions for logprob/train/disabled), test_train.py (multimodal gate), test_megatron_split_state.py (flag wiring). New 2-GPU test_nemotron_omni_cp2_mtp_multimodal_logprob_forward runs an MTP Omni multimodal forward through the NeMo-RL iterator. Not yet validated on this head: that 2-GPU test (L0 pending) and the Super recipe on the final SHA.

Follow-ups (pre-existing, not for this PR)

• HybridEP prepad × MTP: configure_hybridep_packed_input_padding keys the "MTP disabled" check on the YAML mtp_num_layers, but Bridge can enable MTP from the HF config (nemotron_h_bridge.py L276, L325) and _apply_mtp_config copies the YAML value only when the key is present. With the key omitted, prepad re-pads the tokens but mtp_loss_mask keeps the old width and crashes in process_mtp_loss; the new position_ids would be short too but are clamped and ignored, so this PR does not widen it. Note the guard runs inside _apply_moe_config (setup.py L996), before _apply_mtp_config (L999), so reading model_cfg.mtp_num_layers at the current site would wrongly reject the two shipped Super perf recipes that disable MTP via YAML — the fix is to hoist the guard call after _apply_mtp_config and check the effective model_cfg.mtp_num_layers. No in-tree recipe hits it.
• Megatron value worker: a critic built from a hybrid checkpoint that declares MTP is unsupported independent of this PR — megatron_value_worker.py never attaches mtp_loss_mask, so the MTP loss would flow into the 1-wide value head. Threading mtp_enabled there would only replace today's loud assert with silent gradient pollution; the right fix is a setup-time disable/raise for critics (note mtp_num_layers=0 alone is not enough for hybrids: Bridge's finalize still joins one depth while mtp_hybrid_override_pattern is set and mtp_use_repeated_layer is True). All shipped Megatron PPO critics are dense GPTs.
• Perf: HybridModel.forward(compute_mtp_loss=True) is the default, so every logprob/top-k forward on Super runs the MTP depths, vocab projection and CE for nothing (logprobs are unaffected: MTPLossAutoScaler.forward returns output unchanged). Passing False on forward-only HybridModel passes is a cheap win; Omni swallows **kwargs (needs a Bridge change) and GPTModel.forward has no such parameter.
• Optional: the three pack-then-select blocks in data.py (mtp_loss_mask, media mask, position_ids) are now rule-of-three; a module-level _pack_token_aligned_like_input_ids(tensor, ...) -> Tensor helper would give the full-row/CP-local rule one home (~-49 lines). Fine to leave as is.

Upstream — the assert this PR works around is stricter than the code needs: with no CP-layout conversion, prepare_cp_layout passes the caller's position_ids through untouched, roll_tensor(None) returns (None, None), and the embedding ignores positions for position_embedding_type="none" — exactly how GPTModel runs the same MTP block with None today (gpt_model.py L716-L732). Worth an issue on NVIDIA/Megatron-LM referencing #6912; this fix should stay regardless (real, layout-correct positions are valid on every MLM).

Linter: pre-commit run (ruff, ruff-isort, ruff-format, pyrefly, configs-minimize) passes on all 7 files, matching the CI Lint check. Threads: none to respond to.

Generated by Claude Code

[tdene]

/ok to test 9f18bd9

[tdene]

/ok to test d896d58