Skip to content

fix(moe): preserve pruned experts' kernel sizes in model.yaml so retraining reloads their weights - #194

Merged
isLinXu merged 1 commit into
Tencent:mainfrom
Lfan-ke:fix/moe-prune-expert-kernels
Aug 2, 2026
Merged

fix(moe): preserve pruned experts' kernel sizes in model.yaml so retraining reloads their weights#194
isLinXu merged 1 commit into
Tencent:mainfrom
Lfan-ke:fix/moe-prune-expert-kernels

Conversation

@Lfan-ke

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

Copy link
Copy Markdown
Contributor

Problem

MoEPruner keeps the highest-utilization experts, which in a trained model have heterogeneous depthwise kernel sizes (e.g. a layer keeps experts with kernels [5, 9]). PR #192 encoded the reduced num_experts into model.yaml, but not the kept experts' kernel sizes.

Consequently YOLO(pruned.pt).train() (the prune -> LoRA / full fine-tune recovery workflow) rebuilds each ES_MOE from YAML with the default kernels [3, 5, ...][:num_experts], which mismatch the kept experts' actual kernels. intersect_dicts then drops those depthwise weights on the shape mismatch and re-initializes the experts randomly, so retraining starts from a degraded model.

On a VisDrone YOLO-Master-EsMoE-N pruned to 2 experts/layer this is reproducible: direct-inference val is mAP50 0.4269, but the first retrain epoch starts at mAP50 ~0.002 (random experts). After the fix the same retrain starts from a correctly loaded model.

Fix

ES_MOE gains an optional expert_kernel_sizes argument (defaults to None, so behavior is unchanged when unset). MoEPruner._sync_yaml_num_experts now writes the full ES_MOE positional arg list, including each kept expert's kernel size, into model.yaml. The YAML rebuild then reconstructs the exact kept experts and their weights survive intersect_dicts.

Tests

tests/test_moe_prune_yaml_sync.py gains test_sync_preserves_expert_kernels, which builds a stand-in pruned model with non-default kernels, syncs, rebuilds, and asserts the kernels and weights are preserved. All three tests pass. The change is backward compatible.

…aining reloads their weights

Signed-off-by: 林晨 (Leo Cheng) <leo-cheng@vip.qq.com>
@Lfan-ke

Lfan-ke commented Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

The three failing checks (Tests 3.13 ubuntu/windows, Tests 3.8 torch1.8) are pre-existing on main at the same base commit d3341f3, not introduced here. The CI run for main fails the identical three jobs: https://github.com/Tencent/YOLO-Master/actions/runs/30513506266 .

This change is isolated to MoE pruning YAML sync (ES_MOE.expert_kernel_sizes + MoEPruner._sync_yaml_num_experts), and its regression tests in tests/test_moe_prune_yaml_sync.py pass locally. The 3.8 torch1.8 failure is the legacy-autocast issue already addressed by #185. The 3.13 failures are in the unrelated MoT/DDP test suite.

@isLinXu

isLinXu commented Aug 2, 2026

Copy link
Copy Markdown
Collaborator

LGTM

@isLinXu
isLinXu merged commit a13938c into Tencent:main Aug 2, 2026
15 of 18 checks passed
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.

2 participants