Skip to content

fix(moe): keep pruned expert count in model.yaml so pruned models survive retraining - #192

Merged
isLinXu merged 1 commit into
Tencent:mainfrom
Lfan-ke:fix/moe-prune-yaml-sync
Jul 30, 2026
Merged

fix(moe): keep pruned expert count in model.yaml so pruned models survive retraining#192
isLinXu merged 1 commit into
Tencent:mainfrom
Lfan-ke:fix/moe-prune-yaml-sync

Conversation

@Lfan-ke

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

Copy link
Copy Markdown
Contributor

Problem

MoEPruner physically reduces each ES_MOE's expert list and router, but the saved model.yaml still describes the original expert count. When a pruned checkpoint is then used for further training (the prune -> LoRA / full fine-tune recovery workflow), Trainer.setup_model rebuilds the model from model.yaml with ES_MOE's default expert count, and BaseModel.load()'s intersect_dicts silently drops the reduced expert/router weights on shape mismatch. The model re-inflates to the default expert count with randomly initialized experts, so "recovery" ends up training a reconstructed non-pruned model rather than the pruned one.

Fix

After surgery, MoEPruner writes each pruned ES_MOE's expert count into its model.yaml args ([out_channels, num_experts]; ES_MOE clamps top_k = min(top_k, num_experts)). The pruned architecture now round-trips through a YAML rebuild, so retraining/fine-tuning operates on the actual pruned model.

Verification

Added tests/test_moe_prune_yaml_sync.py: reproduces the re-inflation when the yaml is not synced, and confirms the sync preserves the pruned expert count with all weights round-tripping.

…s pruned architecture

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

isLinXu commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator

LGTM

@isLinXu
isLinXu merged commit f9798d9 into Tencent:main Jul 30, 2026
17 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