[RFC]: Dynamic Expert Load Balance with Zero-like-overhead

[raindaywhu]

Motivation

Currently dynamically experts balancing would stop-the-world. Asynchronously expert load balancing would be better without flowing problems:

1. Host-bound latency:
   There are many cpu operations during EPLB such as eplb-algorithm、creating p2p ops、and log2phy expert converting would spend long cpu time, as ~1s.
2. Communication latency: The transfer time would cost much in the situation without nvlink. As the weight of an expert maybe transfer to multiple new positions, thus N times send/recv for one expert, with result long latency. We had tested that batch_isend_irecv cost more 100ms for 16 experts weight transmission in A2 server of ascend.

We have finished SwiftBalancer -- Dynamic Expert Load Balance with Zero-like-overhead on vllm-ascend 1943(vllm-project/vllm-ascend#2186) and it is well tested.
SwiftBalancer would not stop-the-world anymore, in out test on NPU 1~2ms cost for each layer while benefit 5ms-8ms decode latency with ep_size = 64.
The following updates have been made:
1、expert distribution recording with lower cost.
2、async cpu computing for eplb algo and other python operator.
3、new eplb algo with less expert rebalancing while almost the same effect.

Proposed Change

We would like to expand the current implementation:

The overall workflow involves:

1. Record experts distribution during forward. We using expert_token_num after disptach instead of topk_ids, thus we got much smaller tensor shape to reduce cost of hbm recording and add-operator.
2. Do all-gather for experts distribution. Using all-gather instead of all-reduce as less traffic volume.
3. Wake up eplb worker process with experts distribution when num_iterations comes. Run eplb algorithm in eplb worker.
4. Generate p2p send/recv ops and other operator such as log2phy would cost long cpu time.
5. Lanch ibatch_send_recv in async_stream before forward.
6. After forward, wait for the ibatch_send_recv finish, then do uapte expert map and expert weights.

Feedback Period

No response

CC List.

No response

Any Other Things.

Currently we are working on adding adaptor for working with eplb mainline code.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[njuyuan]

Swiftbalancer dynamically adjusts expert allocation by asynchronously gathering expert heat information and triggering real-time rebalancing when the system's imbalance degree exceeds a predefined threshold. Under the constraint that communication between any pair of devices cannot exceed once, it employs a heuristic algorithm to derive dynamic expert adjustment strategies with high load-balancing efficiency. Compared to existing algorithms, it achieves similar load-balancing performance with significantly fewer expert adjustments.

[lengrongfu]

It is Zero-like-overhead , instead of Zero-overhead, right？

[raindaywhu]

It is Zero-like-overhead , instead of Zero-overhead, right？

Yeah, due to 1ms for HBM expert weights copy

[youkaichao]

cc @abmfy can you help take a look?

[abmfy]

cc @abmfy can you help take a look?

Thanks for the mention!

How does this compare to #22179? That PR also introduces async weight transfer.

Nonetheless, I think the other optimizations here are also worthwhile.

[raindaywhu]

cc @abmfy can you help take a look?

Thanks for the mention!

How does this compare to #22179? That PR also introduces async weight transfer.

Nonetheless, I think the other optimizations here are also worthwhile.

Thanks for your attention! In our test weight transfer would finished during the bubble between near iterations. And I just wondering whether the thread is needed for async stream of this pr.

[cad-x2]

I had a question about the EPLB policies, what policies are currently being used? because I saw two different versions in the code. I specifically want to know what their algorithms are. Also, how long does the execution take in total based on what I'm reading 1-2ms, and up to 7-8ms when you scale up to 64? because you're hiding the 100ms communication time behind the decode, would that be accurate?

cc @abmfy can you help take a look?

Thanks for the mention!
How does this compare to #22179? That PR also introduces async weight transfer.
Nonetheless, I think the other optimizations here are also worthwhile.

Thanks for your attention! In our test weight transfer would finished during the bubble between near iterations. And I just wondering whether the thread is needed for async stream of this pr.

[Skywalker-EP]

We plan to split the entire EPLB functionality into multiple PRs to be merged. The goal is to ensure compatibility with ours and existing EPLB as well as to strive for generalization, which mainly includes:

• New EPLB algorithm, and generalized class abstraction of EPLB Algorithm
  • [EPLB] Reconstruct EPLB algorithm invocation method #24069
  • Abstract eplb algo #26471
• A new processes to manage the asynchronous execution of EPLB algorithms and the orchestration of weight transmission tasks
  • [EPLB] Add Asynchronous Expert Rebalancing #23343 (closed)
  • [Core][EPLB] Refactor step to unify synchronous and asynchronous expert rearrangement #27326
• Optimization of heat collection, where heat is collected from dispatch ops (such as torch_npu.npu_moe_distribute_dispatch, related operators in deepep, etc.)
  • [EPLB]: Optimize export_load_view update #24091
• I. Model adaptation, extracting EPLB-related utility functions from the model to form utility classes, and registering them into the model via monkey patching. II. Generalized Optimization of the EPLB Architecture
  • [Core] Refactor EPLB #23913 (include [EPLB] Model Modify of EPLB #23553)
  • Currently, only the methods compatible with our algorithm have been adapted. The next step is to extract the original EPLB-related common functions
  • Pending
• Asynchronous weight transmission
  • We can check with [EPLB] Optimize EPLB for Async Rearrange Experts  #22179
  • And our scheme [Core][EPLB] Refactor step to unify synchronous and asynchronous expert rearrangement #27326

cc @abmfy @tlrmchlsmth

[mengxingkongzhouhan]

3、new eplb algo with less expert rebalancing while almost the same effect.The implementation of this feature is as follows for PR connection:https://github.com/vllm-project/vllm/pull/23243