Skip to content

feat(moe): add validation-driven dynamic top-k annealing scheduler - #183

Open
Lfan-ke wants to merge 1 commit into
Tencent:mainfrom
Lfan-ke:feat/moe-dynamic-topk-anneal
Open

feat(moe): add validation-driven dynamic top-k annealing scheduler#183
Lfan-ke wants to merge 1 commit into
Tencent:mainfrom
Lfan-ke:feat/moe-dynamic-topk-anneal

Conversation

@Lfan-ke

@Lfan-ke Lfan-ke commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

What

Adds DynamicTopKScheduler: a validation-driven scheduler that anneals the number of active experts (top_k) during training, instead of the auxiliary balance-loss coefficient that existing MoE schedulers tune.

Each validation epoch it drops top_k by one when validation mAP has saturated (sliding-window) and expert usage is balanced (EMA Gini ≤ ceiling), and rolls back + freezes if the step hurts mAP. A converged, balanced model then runs with fewer active experts — lower inference FLOPs, matching the ES-MoE "compute-on-demand" premise, and complementary to post-hoc MoEPruner.

API

  • DynamicTopKScheduler / TopKAnnealConfig / apply_top_k(model, k) in ultralytics/nn/modules/moe/topk_anneal.py.
  • Opt-in training callback: model.add_callback("on_fit_epoch_end", create_topk_anneal_callback(cfg)).
  • Fully backward compatible: disabled by default (enabled=False); the callback is defensive and never interrupts training if MoE modules or optional deps are absent.
  • Reuses mean_usage_gini_from_model and mirrors the apply_balance_loss_coeff pattern.

Tests

tests/test_moe_topk_anneal.py — 9 tests (anneal on saturation+balance, no-anneal when imbalanced / still improving, rollback on mAP drop, floor at min_k, state_dict roundtrip, disabled no-op, callback never raises, val-map extraction). All pass.

Closes #52

Signed-off-by: 林晨 (Leo Cheng) <leo-cheng@vip.qq.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

【2026犀牛鸟开源人才专属】【中高难度】MoE 优化专项:MoE 专家剪枝与动态超参数调度优化

1 participant