Refactor mem_pool_host deallocation to main thread#9302
Refactor mem_pool_host deallocation to main thread#9302pansicheng wants to merge 10 commits intosgl-project:mainfrom
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @pansicheng, I'm Gemini Code Assist1! 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 refactors the memory deallocation mechanism for mem_pool_host indices, specifically those pre-allocated for prefetch storage. The primary goal is to centralize the freeing of these host indices to the main thread, thereby resolving inconsistencies that arose from multi-threaded deallocation and minimizing the need for cross-tensor parallel (TP) group synchronization. The changes involve introducing a new queue for deallocation requests, enhancing prefetch operation tracking, and implementing dedicated logic on the main thread to manage the memory freeing process.
Highlights
- Centralized Memory Deallocation: The core change centralizes the deallocation of
mem_pool_hostindices, moving this responsibility from multiple threads to the main thread. This aims to resolve inconsistencies and reduce synchronization overhead across different tensor parallel (TP) groups. - New Deallocation Queue: A new
prefetch_free_queuehas been introduced within theCacheControllerto queuePrefetchOperationobjects for deferred memory deallocation. This decouples the freeing process from the prefetch and I/O auxiliary threads. - Enhanced Prefetch Operation Tracking: The
PrefetchOperationclass now includesstorage_hit_count,min_completed_tokens, andmatched_lengthattributes. These are used to precisely track and manage the specific segments of pre-allocated host memory that need to be freed or retained. - Dedicated Deallocation Logic: New methods,
free_operationandcheck_free_prefetch, have been added toHiRadixCache. These methods are responsible for processing theprefetch_free_queue, performing the actual memory deallocation, and ensuring proper synchronization across TP workers before freeing memory.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request refactors the deallocation of host memory for prefetched data to the main thread, which is a good approach to centralize resource management and avoid synchronization issues. The implementation uses a queue to pass deallocation tasks, which is sound. My review includes two main suggestions for improvement: one to address code duplication for better maintainability, and another to optimize the memory reclamation logic for better efficiency.
Co-authored-by: JinYan Su <jinyansu792@gmail.com>
Co-authored-by: JinYan Su <jinyansu792@gmail.com>
| if len(written_indices): | ||
| self.cache_controller.mem_pool_host.update_prefetch(written_indices) | ||
|
|
||
| self.cache_controller.mem_pool_host.free(host_indices[:matched_length]) |
There was a problem hiding this comment.
I think we can still release this here, since matched_length should be synchronized across workers.
There was a problem hiding this comment.
I’ve made the changes, please take a look!
|
to replace this with #9797 |
Motivation
Currently, host indices pre-allocated for prefetch storage are freed by multiple threads, causing inconsistencies in mem_pool_host across different ranks. This requires synchronization operations between tp groups.
This PR proposes to centralize the freeing of host indices in the main thread, minimizing the need for cross-tp synchronization.
Modifications
Accuracy Tests
Benchmarking and Profiling
Checklist