Perf: avoid intermediate tensor allocs via in-place div & optimized top-k flow

[hjh0119]

No description provided.

[hjh0119]

/gemini review

[gemini-code-assist]

Summary of Changes

Hello, 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 focuses on enhancing the performance and reducing the memory footprint of various trainer modules by optimizing tensor operations. The core change involves replacing out-of-place tensor divisions with in-place operations and restructuring top-k processing to minimize the creation of temporary tensors. These changes aim to make the training process more memory-efficient without altering the numerical results.

Highlights

• Performance Optimization: Converted several tensor division operations from creating new tensors to performing in-place modifications using div_(), thereby reducing intermediate tensor allocations.
• Optimized Top-K Flow: Refactored the sequence of operations in top-k scenarios within generalized_jsd_loss and _jsd_topk functions to perform gather or topk operations before division, followed by in-place scaling, further minimizing memory footprint.
• Memory Footprint Reduction: Applied these optimizations across swift/megatron/trainers/gkd_trainer.py, swift/rlhf_trainers/gkd_trainer.py, and swift/rlhf_trainers/grpo_trainer.py to improve overall efficiency.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• swift/megatron/trainers/gkd_trainer.py
  • Converted temperature scaling for student_logits_masked and teacher_logits_masked to use in-place division.
  • Modified _jsd_topk to apply temperature scaling in-place after gathering student logits, avoiding an intermediate tensor.
• swift/rlhf_trainers/gkd_trainer.py
  • Refactored generalized_jsd_loss to optimize top-k processing for both API and local teacher scenarios, using in-place division for temperature scaling.
  • Replaced out-of-place temperature scaling with in-place div_ for final student_logits and teacher_logits.
• swift/rlhf_trainers/grpo_trainer.py
  • Updated _get_logps_via_sp to apply temperature scaling to logits using an in-place division operation.

Activity

• No specific activity has been recorded for this pull request yet.

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 aims to improve performance by avoiding intermediate tensor allocations, primarily by using in-place division and optimizing operation order. While many of the changes effectively achieve this, I've identified two instances where using in-place operations introduces unintended side effects by modifying input tensors. This alters the functions' behavior and could lead to bugs if the calling code reuses these tensors. My review includes suggestions to revert these specific changes to preserve correctness while retaining the other valuable performance optimizations.

[hjh0119]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces performance optimizations to the GKD trainers by avoiding intermediate tensor allocations. The changes primarily involve using in-place division (div_) and reordering operations to apply temperature scaling after masking or top-k selection, which reduces memory usage. These optimizations are consistently applied in both swift/megatron/trainers/gkd_trainer.py and swift/rlhf_trainers/gkd_trainer.py. The logic appears sound and the changes align with the performance improvement goals of the PR.

[gemini-code-assist]

Code Review

The pull request introduces performance optimizations by avoiding intermediate tensor allocations through the use of in-place division (.div_()) and optimizing the top-k flow in the generalized_jsd_loss and _jsd_topk functions. These changes reduce memory footprint and improve computational efficiency, aligning well with the stated objective of the pull request. The modifications are consistently applied across both swift/megatron/trainers/gkd_trainer.py and swift/rlhf_trainers/gkd_trainer.py.