Note that this expert id recording involves copying tensors from gpu to cpu during the model execution so will affect throughput and latency significantly. I will add an argument to turn off the recording by default.

Hi @yuhsaun-t, thank you for your contribution. Could you please provide more details about the output of dump_record? The output format is not clear to me. Also, can we use non-blocking copy (reference) to address the blocking issue?

Hi @yuhsaun-t, thank you for your contribution. Could you please provide more details about the output of dump_record? The output format is not clear to me. Also, can we use non-blocking copy (reference) to address the blocking issue?

Hello @ch-wan , sure, here is an attached output json from the dump. Right now there will be one json file for each gpu rank, but the content is the same. I can update the code to dump only on rank 0. Yes, I will update to use the non-blocking copy.
expert_distribution_1738710857.859208.json

Hello @ch-wan , I have updated the PR based on the reviews, and added unittests and docs.
Now the dumped file looks like this:

expert_id,count
33,12
4,16
32,8
24,12
40,13
37,11
11,17

Can you review the PR again? Thanks!

@yuhsaun-t Thank you very much for your great effort! This PR would be very useful for analyzing experts' dynamic workloads during MoE serving. I have added some comments. Could you please take a look?

Another question comes to my mind when I double-check this PR. When DP is enabled, the recorders from different workers will collect different experts distribution. We may need all-reduce to synchronize their results, and only the master worker can dump results.

The current implementation makes it so that the server will dump one csv file for one rank. I think we can keep it this way so that the server does not have to synchronize on the fly to save performance. We can process the dumped csv files later and aggregate them into one file. Does that sound good to you?

THank you for your great effort! I have approved the change.