[megatron] Support qwen3-vl/qwen3-omni cp

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @Jintao-Huang, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the Megatron framework by introducing comprehensive support for Context Parallelism (CP) to the Qwen3-VL and Qwen3-Omni models. The core objective is to enable efficient distribution of data and computations across parallel contexts, which is crucial for scaling these large multimodal models. This involved adapting how inputs, especially visual embeddings and decoder inputs, are processed and split, alongside a generalization of the underlying utility functions for parallel data handling.

Highlights

• Context Parallelism Support: Implemented robust support for Context Parallelism (CP) specifically for the Qwen3-VL and Qwen3-Omni models within the Megatron framework.
• Generalized Input Splitting: Refactored and unified the token splitting logic into a single, more flexible split_cp_inputs utility function, replacing previous specialized functions and allowing for splitting along any dimension.
• Qwen3-VL Integration: Modified qwen3_vl.py to correctly handle visual embeddings and masks, incorporating packed_seq_params and the new split_cp_inputs for CP-aware processing.
• Streamlined Decoder Input Handling: Updated mm_gpt_model.py to use the new split_cp_inputs for processing decoder inputs, simplifying the context parallel data flow.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for context parallelism (CP) in qwen3-vl and qwen3-omni models. The core change is a refactoring that replaces specialized data splitting functions with a single, more generic split_cp_inputs utility. This new function is then used to correctly partition model inputs, including visual features, across context parallel ranks. While the overall approach and refactoring are good, there are a couple of minor style issues with local imports and a potential bug related to leftover code that could process decoder_input incorrectly.

[gemini-code-assist]

This simplification introduces a potential issue. The previous code had special handling for decoder_input, which requires splitting along dim=0 due to its (sequence_length, batch_size, hidden_size) shape. The new code splits all tensors along dim=-1, which is incorrect for decoder_input as it would split the hidden dimension.

While MultimodalGPTModel now handles this splitting separately, the 'decoder_input' key is still checked for in this function (on line 103), making this code path potentially active for other cases. This is confusing and could lead to silent errors.

To fix this, please either restore the special handling for decoder_input or, if it's confirmed to be dead code, remove 'decoder_input' from the keys list on line 103.

Suggested change

	batch[key] = split_cp_inputs(val, packed_seq_params.cu_seqlens_q, -1)
	if key == 'decoder_input':
	batch[key] = split_cp_inputs(val, packed_seq_params.cu_seqlens_q, 0)
	else:
	batch[key] = split_cp_inputs(val, packed_seq_params.cu_seqlens_q, -1)

[gemini-code-assist]

For better code organization and to avoid potential circular import issues, it's recommended to move imports to the top of the file. Please move from ...trainers.utils import split_cp_inputs to the top-level imports.

[gemini-code-assist]

To improve code readability and maintainability, please move this import statement to the top of the file with other imports.