[asolergi-nv]

What does this PR do ?

NVIDIA/Megatron-LM#4103 Introduced changes to the get_batch_on_this_cp_rank function. This PR is updating M-Bridge to the new signature.

Changelog

• Add specific line by line info of high level changes in this PR.

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the 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)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

[claude]

Light Code Review

Missing mock updates (potential test breakage)

This PR updates all source call sites and some test mocks for the new is_hybrid_cp parameter on get_batch_on_this_cp_rank, but misses updating 5 test mocks that use lambda x: x - these will raise TypeError if actually invoked with the new keyword argument:

• tests/unit_tests/training/test_vlm_step.py lines 134, 188, 289, 354
• tests/unit_tests/training/test_audio_lm_step.py lines 132, 183, 282

These should be updated to lambda x, **kwargs: x to match the fix applied in test_dgpt_step.py:473.

PR hygiene

The PR title and description are not filled in - consider adding a one-line summary of the upstream MCore change being adapted to (MLM!4103) and listing the affected call sites.

---

Suggested test cases

No perf tests impacted.

Unit tests that should be verified:

• tests/unit_tests/utils/test_slice_batch_for_cp.py (updated in this PR)
• tests/unit_tests/diffusion/model/nemotron_labs_diffusion/test_dgpt_step.py (updated in this PR)
• tests/unit_tests/training/test_vlm_step.py (not updated - verify these still pass)
• tests/unit_tests/training/test_audio_lm_step.py (not updated - verify these still pass)

[asolergi-nv]

Fixed in #3944