[RFC]: EPLB Execution Optimization From pr 18343

[david6666666]

Motivation.

#18343

• EPLB Execution
• Parallelize the rearrangement algorithm (calculating new expert mapping, not the communication)
• Shuffle one layer at once and use multiple steps, to lower the impact on inter-token latency
• Investigate should we pre-allocate expert weight buffer used for transferring
• Take locality into consideration in expert weight transmission, e.g. prioritize transferring to GPUs on the same node
• Use cuda.Stream() asynchronously moves the weight to the buffer

Proposed Change.

Advantages over the synchronous solution: TTFT/TPOT fluctuations are small. The synchronous solution will block reasoning when reordering, causing TTFT/TPOT time to increase by hundreds of milliseconds.

Feedback Period.

July 11th ~ July 30th

CC List.

@abmfy

Any Other Things.

No response

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.

[abmfy]

Overall, the non-blocking design looks great to me. Please ping me once the implementation is ready. Thanks for the contribution!

[david6666666]

Overall, the non-blocking design looks great to me. Please ping me once the implementation is ready. Thanks for the contribution!

Thank you very much. Our PR has been verified. I will release the test plan and test results in the next two days, and then continue to streamline the code.

[david6666666]

@abmfy @hmellor the implementation is ready, PTAL, thanks. wait for #20562 and add the new config to the new EPLBConfig class and rebase

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[hmellor]

@david6666666 should this be considered complete now?

[david6666666]

@david6666666 should this be considered complete now?

Yes, I am working on rebase code, add some test and support MTP