fix(megatron): pass delegation flags on the presharded train path

[rohitrango]

What this fixes

MegatronPolicyWorkerImpl._train_microbatch_body (the presharded / TQ train path) calls get_microbatch_iterator without the three delegation flags:

• delegate_pack_to_model
• delegate_mtp_loss_mask_to_model
• model_slices_context_parallel_inputs

The other three get_microbatch_iterator call sites in the same file (logprob and the non-presharded train path) all pass them. All three default to False, so on the presharded path a model that declares model_slices_context_parallel_inputs = True had that declaration silently ignored.

Failure mode

For a model that slices its own CP inputs (e.g. Nemotron Omni) with CP>1:

_prepare_prepacked_batch_for_megatron handed the model the CP-local slice (65536/4 = 16384 tokens) while PackedSeqParams.cu_seqlens_q_padded stayed global at 65536. The model, which slices its own CP inputs after media insertion, then sliced again. tex.thd_get_partitioned_indices built indices into a 65536-token stream against a 16384-token tensor, so index_select ran off the end:

Assertion `idx_dim >= 0 && idx_dim < index_size &&
           "scatter gather kernel index out of bounds"` failed.

CUDA reports that asynchronously, so it surfaced in an unrelated later kernel (mamba in_proj RMSNorm), which made it look like a model bug rather than a batch-prep bug.

Net effect: CP>1 was unusable on the presharded train path for any model that slices its own CP inputs. Note TQ itself is not involved in the sharding — tq_policy.py marks context_parallel as a replicated axis.

The fix

Pass the same three flags that the other call sites already pass. Three added lines, no behaviour change for models where the flags are False (which is every model that does not opt in).

Verification

Run on a Nemotron Omni blend, 2 nodes x 8 GPU, TP4 x CP4, seq 65536:

• Before: 0 optimizer steps, 1001 scatter/gather assertion failures.
• After: cu_seqlens[-1] == total_tokens == 65536, per-rank indices numel == 16384 and in bounds, training runs cleanly (grad_norm 2.2–3.6 over 4+ steps).

🤖 Generated with Claude Code

[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.

[rohitrango]

/ok to test e7c5fad