Repository navigation
[PoC] eplb experts exchange async - #14974
qiaoning13256 wants to merge 1 commit into
Conversation
Summary of ChangesHello @qiaoning13256, 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 introduces a Proof of Concept solution to address performance bottlenecks caused by synchronous expert location updates in the EPLB (Expert Placement Load Balancing) system. By refactoring the update mechanism to be asynchronous, utilizing Python's generator capabilities and threading, the system can avoid blocking critical forward threads. This change also incorporates MUSA stream management for GPU operations and leverages specific MoE process groups for more efficient distributed communication, ultimately aiming to improve the overall responsiveness and tail latency of the system. Highlights
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an asynchronous mechanism for expert weight exchange to improve performance, which is a valuable optimization. The changes are a good proof-of-concept, but there are a few critical issues and areas for improvement. I've identified two critical bugs where yield is used instead of yield from, which will break the asynchronous flow. Additionally, I've pointed out the use of global variables for thread communication, which is not robust and should be refactored. There is also some code duplication that could be cleaned up. Addressing these points will make the implementation more correct and maintainable.
| yield self._model_runner.update_expert_location( | ||
| expert_location_metadata, | ||
| update_layer_ids=update_layer_ids, | ||
| ) |
There was a problem hiding this comment.
The update_expert_location method is now a generator. You should use yield from to delegate to it, instead of yield. Using yield will only yield the generator object itself, not its produced values, which will break the asynchronous execution flow.
| yield self._model_runner.update_expert_location( | |
| expert_location_metadata, | |
| update_layer_ids=update_layer_ids, | |
| ) | |
| yield from self._model_runner.update_expert_location( | |
| expert_location_metadata, | |
| update_layer_ids=update_layer_ids, | |
| ) |
| ) | ||
| else: | ||
| self.expert_location_updater.update( | ||
| yield self.expert_location_updater.update( |
There was a problem hiding this comment.
The update method of expert_location_updater is a generator. You should use yield from to delegate to it. Using yield will only yield the generator object itself, which is not the intended asynchronous behavior.
| yield self.expert_location_updater.update( | |
| yield from self.expert_location_updater.update( |
| async_ongoing = False | ||
| buffer2weight_copy_infos: List[Tuple[int, int]] = [] |
There was a problem hiding this comment.
Using global variables async_ongoing and buffer2weight_copy_infos for communication between threads is not a robust design. It makes the code harder to reason about, test, and maintain, especially in a concurrent environment. Consider encapsulating the asynchronous update logic and its state within a class. This would avoid global state and provide a cleaner, more maintainable API.
For example:
class AsyncExpertUpdater:
def __init__(self, ...):
self.ongoing = False
self.copy_infos = []
# other necessary state
def run(self):
# logic from async_update
...
def _update_worker(self):
# logic from update_expert_weights_single_layer
...| num_tensors = len(routed_experts_weights) | ||
| local_expert_location_range = ( | ||
| rank * num_local_physical_experts, | ||
| (rank + 1) * num_local_physical_experts, | ||
| ) | ||
|
|
||
| def _get_tensor(tensors, tensor_index: int, expert_location: int) -> torch.Tensor: | ||
| return tensors[tensor_index][_get_local_expert_location(expert_location)] | ||
|
|
||
| def _get_local_expert_location(expert_location: int) -> int: | ||
| assert ( | ||
| local_expert_location_range[0] | ||
| <= expert_location | ||
| < local_expert_location_range[1] | ||
| ) | ||
| return expert_location % num_local_physical_experts |
There was a problem hiding this comment.
The helper functions _get_tensor and _get_local_expert_location and their surrounding logic are duplicated from update_expert_weights_single_layer_. This code duplication increases maintenance overhead. Please refactor this logic into a shared utility function or a helper class to avoid redundancy.
|
Thanks @qiaoning13256. Closing this because it is still a draft and has not been updated in 245 days. Reopen it if the work is still relevant. Some directories moved recently, so an older branch may need retargeting: |
Motivation
EPLB experts exchange synchronously will block main forward threads, and negative influence TTFT/TOPT p99.
Reference vLLM vllm-project/vllm#22179 we introduce a PoC solution.
Modifications
Accuracy Tests
Benchmarking and Profiling
Checklist