Skip to content

feature(tj): add MoE expert selection stats and gradient conflict monitoring#470

Open
tAnGjIa520 wants to merge 2 commits intoopendilab:mainfrom
tAnGjIa520:moe4
Open

feature(tj): add MoE expert selection stats and gradient conflict monitoring#470
tAnGjIa520 wants to merge 2 commits intoopendilab:mainfrom
tAnGjIa520:moe4

Conversation

@tAnGjIa520
Copy link
Copy Markdown
Contributor

@tAnGjIa520 tAnGjIa520 commented Feb 3, 2026

  • Add MoE expert selection statistics with multi-window sliding buffers and TensorBoard heatmaps
  • Add gradient conflict monitoring for encoder, MoE, shared expert, and individual experts

  • 增加 MoE 专家选择统计,支持多窗口滑动缓冲和 TensorBoard 热力图
  • 增加 encoder、MoE、shared expert 及各个 expert 的梯度冲突监控

实验结果

moe_expert_selection_wasserstein_distance
image

moe_expert_selection_js_divergence
image

gradient_conflict_comparison_moe_vs_nomoe
image

expert_selection_heatmaps
image

@puyuan1996 puyuan1996 added the research Research work in progress label Feb 6, 2026
Comment thread lzero/policy/utils.py Outdated
Comment thread zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py
Comment thread zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py
… and configs

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

research Research work in progress

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants