feat(moe): add validation-driven dynamic top-k annealing scheduler - #183
Open
Lfan-ke wants to merge 1 commit into
Open
feat(moe): add validation-driven dynamic top-k annealing scheduler#183Lfan-ke wants to merge 1 commit into
Lfan-ke wants to merge 1 commit into
Conversation
Signed-off-by: 林晨 (Leo Cheng) <leo-cheng@vip.qq.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_kby 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-hocMoEPruner.API
DynamicTopKScheduler/TopKAnnealConfig/apply_top_k(model, k)inultralytics/nn/modules/moe/topk_anneal.py.model.add_callback("on_fit_epoch_end", create_topk_anneal_callback(cfg)).enabled=False); the callback is defensive and never interrupts training if MoE modules or optional deps are absent.mean_usage_gini_from_modeland mirrors theapply_balance_loss_coeffpattern.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