[feature] Add dense CP support for Qwen3-Omni

[hbhflw2000]

What does this PR do ?

Add dense context parallelism support for Qwen3-Omni thinker training.

Changelog

• Enable dense CP batch padding and slicing for Qwen3-Omni.
• Preserve full input_ids for mRoPE while using local CP tensors for loss/attention.
• Add unit tests for dense CP behavior and mRoPE handling.

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

Validation:

• ruff check
• ruff format --check
• git diff --check
• Focused Qwen3-Omni unit tests: 12 passed
• Focused Qwen3-Omni mRoPE unit test: 1 passed
• Local 4-GPU A800 smoke: CP=2, EP=2
• 4-node / 32-GPU validation checked via TensorBoard: CP=2, TP=2, PP=2, EP=4

Upcoming follow-up:

• Add packed sequence support for Qwen3-Omni dense CP in a separate PR.

Related to #3317

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

[yaoyu-33]

possible to re-use current utils?

[hbhflw2000]

Thanks, this is a good point. I checked the existing CP helper. slice_batch_for_context_parallel operates after embedding merge and slices inputs_embeds, while this Qwen3-Omni path needs a pre-slice raw-batch normalization step so that Megatron's dense CP zigzag split receives tensors divisible by TP * CP * 2, while the full input_ids are still preserved for thinker-side mRoPE.

So I don't think directly reusing that helper is semantically safe here. The current code already reuses the shared padding primitives, but I agree the framework should probably have a clearer raw-batch sequence normalization utility. I can either add a clarifying comment in this PR, or follow up with a small utility extraction that can be shared by Qwen3-Omni and Qwen3-VL packed paths.

[yaoyu-33]

/ok to test 68bfc03