Code on Expert Specialization Experiments

[Tangkexian]

Great work! And I am particularly interested in the section about "Expert Specialization".

I am reaching out to inquire if there is any publicly available code or additional material related to the experiments conducted in this section of the paper? Alternatively, if the code is not available, any additional insights or tips on replicating the experiments would also be extremely helpful.

Thank you !

[XiaoYee]

@Spico197 Please help

[Spico197]

Hi there, thanks for your attention to our work~

You may check the code here to reproduce the visualization on gate loads:

llama-moe/smoe/entrypoint/analysis/gate_load_vis.py

Lines 108 to 116 in b17aff4

	for layer_idx in range(gate_load_sum.shape[0]):
	visualize_expert_load_heatmap(
	gate_load_sum[layer_idx],
	layer_idx,
	name,
	shape=(4, 4),
	save_dir=str(result_dir),
	save_fig=True,
	)