[megatron] fix megatron-swift seq_cls

[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 addresses a bug in the megatron-swift implementation, specifically concerning sequence classification when sequence parallelism is active. The fix ensures that intermediate hidden states are correctly gathered across parallel regions before being processed by the final output layer, resolving potential issues with incorrect classification results and improving the accuracy of sequence classification tasks.

Highlights

• Sequence Parallelism Fix: Corrected the application of gather_from_sequence_parallel_region in gpt_model.py to apply to hidden_states before the output layer, instead of logits after, when sequence parallelism is enabled. This ensures proper handling of sequence classification outputs.

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 fixes an issue in sequence classification tasks when using sequence parallelism. The change correctly reorders the operations to gather the hidden_states before applying the output layer. My review points out a performance consideration with this approach and suggests exploring a more optimized solution to maintain efficiency while ensuring correctness.

[gemini-code-assist]

This change appears to fix a correctness issue for sequence classification with sequence parallelism by gathering hidden_states before the output_layer. While this ensures correctness, it may introduce a performance regression.

The new implementation involves:

1. An all_gather operation on the potentially large hidden_states tensor.
2. Redundant computation of the output_layer on all tensor parallel ranks, since each now computes the full logits.

The previous implementation gathered the smaller logits tensor, which was more efficient in terms of communication and avoided redundant computation.

If sequence classification relies only on the hidden state of a single token (e.g., the last one), could we consider a more optimized approach? For example:

1. Compute the logit only on the rank that holds the required token's hidden state.
2. Broadcast the resulting logit to all other ranks.

This would be more efficient in both communication and computation. If feasible, it would be a valuable optimization.