From e15de99cfe84dc0b0b074e53ba574493b53193c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E6=99=A8=20=28Leo=20Cheng=29?= Date: Thu, 30 Jul 2026 13:47:57 +0800 Subject: [PATCH] fix(moe): preserve pruned experts' kernel sizes in model.yaml so retraining reloads their weights MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 林晨 (Leo Cheng) --- tests/test_moe_prune_yaml_sync.py | 54 ++++++++++++++++++++++----- ultralytics/nn/modules/moe/modules.py | 32 +++++++++++++--- ultralytics/nn/modules/moe/pruning.py | 26 ++++++++++++- 3 files changed, 94 insertions(+), 18 deletions(-) diff --git a/tests/test_moe_prune_yaml_sync.py b/tests/test_moe_prune_yaml_sync.py index 1de1be4c..f27c5fc2 100644 --- a/tests/test_moe_prune_yaml_sync.py +++ b/tests/test_moe_prune_yaml_sync.py @@ -1,14 +1,18 @@ # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license -"""Regression test: MoEPruner must keep ``model.yaml`` consistent with the pruned -expert count so a pruned model survives a YAML-based rebuild during retraining +"""Regression tests: MoEPruner must keep ``model.yaml`` consistent with the pruned +ES_MOE blocks so a pruned model survives a YAML-based rebuild during retraining (the prune -> LoRA / full fine-tune recovery workflow). -Without the fix, ``DetectionModel(pruned.yaml)`` rebuilds with ES_MOE's default -expert count and ``intersect_dicts`` silently drops the reduced expert/router -weights, re-inflating the model with randomly initialized experts. -""" -import copy +Two failure modes are covered: +1. Expert *count* re-inflation - rebuild uses ES_MOE's default expert count and + ``intersect_dicts`` drops the reduced expert/router weights. +2. Expert *kernel* mismatch - pruning keeps experts with heterogeneous kernels + (e.g. [5, 9]) but a bare rebuild assigns the defaults [3, 5], so the kept + experts' depthwise weights are dropped on a shape mismatch and re-initialized. +The fix writes the full ES_MOE arg list (count + per-expert kernel sizes) into +``model.yaml`` so both survive the rebuild. +""" import yaml from ultralytics.nn.modules.moe.modules import ES_MOE @@ -22,6 +26,14 @@ def _experts(model): return [m.num_experts for _, m in model.named_modules() if isinstance(m, ES_MOE)] +def _kernels(model): + return [ + [e.conv.depthwise.kernel_size[0] for e in m.experts] + for _, m in model.named_modules() + if isinstance(m, ES_MOE) + ] + + def _build_reduced(num_experts): """Build a model whose ES_MOE blocks each have ``num_experts`` (stand-in for a pruned model).""" d = yaml.safe_load(open(CFG)) @@ -31,6 +43,16 @@ def _build_reduced(num_experts): return DetectionModel(d, ch=3, nc=10, verbose=False) +def _build_with_kernels(num_experts, kernels): + """Stand-in for a pruned model that kept experts with non-default kernel sizes.""" + d = yaml.safe_load(open(CFG)) + for layer in d["backbone"]: + if layer[2] == "ES_MOE": + out_ch = layer[3][0] + layer[3] = [out_ch, num_experts, 8, num_experts, True, 0.4, 15, list(kernels)] + return DetectionModel(d, ch=3, nc=10, verbose=False) + + def test_reinflation_without_sync(): """Reproduce the bug: reduced model + un-synced (full) yaml re-inflates on rebuild.""" reduced = _build_reduced(2) @@ -41,12 +63,24 @@ def test_reinflation_without_sync(): def test_sync_preserves_pruned_experts(): - """The fix: MoEPruner._sync_yaml_num_experts writes the reduced count into yaml, - so a YAML rebuild preserves the pruned architecture.""" + """The fix writes the reduced count into yaml so a rebuild preserves it.""" reduced = _build_reduced(2) reduced.yaml = yaml.safe_load(open(CFG)) # start from the un-synced (buggy) yaml MoEPruner._sync_yaml_num_experts(type("D", (), {})(), reduced) # method only uses self for logging es_args = [layer[3] for layer in reduced.yaml["backbone"] if layer[2] == "ES_MOE"] - assert all(a[-1] == 2 for a in es_args), f"yaml not synced: {es_args}" + assert all(a[1] == 2 for a in es_args), f"num_experts not synced: {es_args}" rebuilt = DetectionModel(reduced.yaml, ch=3, nc=10, verbose=False) assert _experts(rebuilt) == [2, 2, 2, 2], "pruned expert count must survive the rebuild" + + +def test_sync_preserves_expert_kernels(): + """The fix also writes per-expert kernel sizes so heterogeneous kept experts survive.""" + reduced = _build_with_kernels(2, [5, 9]) + before = _kernels(reduced) + assert all(k == [5, 9] for k in before), f"stand-in kernels not built: {before}" + reduced.yaml = yaml.safe_load(open(CFG)) # un-synced yaml would rebuild default [3, 5] + MoEPruner._sync_yaml_num_experts(type("D", (), {})(), reduced) + es_args = [layer[3] for layer in reduced.yaml["backbone"] if layer[2] == "ES_MOE"] + assert all(a[-1] == [5, 9] for a in es_args), f"kernels not synced: {es_args}" + rebuilt = DetectionModel(reduced.yaml, ch=3, nc=10, verbose=False) + assert _kernels(rebuilt) == before, "expert kernel sizes must survive the rebuild" diff --git a/ultralytics/nn/modules/moe/modules.py b/ultralytics/nn/modules/moe/modules.py index 1e5a3d08..8ae21297 100644 --- a/ultralytics/nn/modules/moe/modules.py +++ b/ultralytics/nn/modules/moe/modules.py @@ -404,7 +404,7 @@ class ES_MOE(nn.Module): def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8, top_k=2, use_sparse_inference=True, dynamic_threshold=0.4, - max_kernel_size=15): + max_kernel_size=15, expert_kernel_sizes=None): """ Args: in_channels: Input channels @@ -415,6 +415,10 @@ def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8, use_sparse_inference: Enable sparse Top-K expert computation during inference dynamic_threshold: Optional threshold for pruning low-confidence experts during inference max_kernel_size: Largest odd depthwise kernel assigned to an expert + expert_kernel_sizes: Optional explicit per-expert depthwise kernel sizes + (length must equal ``num_experts``). When ``None`` the kernels are + derived from defaults; pruned checkpoints set this so retraining + rebuilds the exact kept-expert kernels and reloads their weights. """ super(ES_MOE, self).__init__() @@ -440,6 +444,7 @@ def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8, self.in_channels = in_channels self.out_channels = out_channels self.num_experts = num_experts + self.reduction = reduction self.top_k = min(top_k, num_experts) if top_k is not None else num_experts self.use_top_k = (top_k is not None) self.use_sparse_inference = use_sparse_inference @@ -449,12 +454,27 @@ def __init__(self, in_channels, out_channels=None, num_experts=4, reduction=8, # Dynamic routing (Top-K supported) self.routing = DynamicRoutingLayer(in_channels, num_experts, reduction, top_k) - # Expert group (original design) - default_kernel_sizes = [3, 5, 7] - if num_experts <= len(default_kernel_sizes): - ks = [min(k, max_kernel_size) for k in default_kernel_sizes[:num_experts]] + # Expert group (original design). ``expert_kernel_sizes`` lets a pruned + # checkpoint reconstruct its kept experts' heterogeneous kernels so that + # ``YOLO(pruned.pt).train()`` reloads expert weights instead of dropping + # them on a kernel-shape mismatch (prune -> LoRA/full fine-tune recovery). + if expert_kernel_sizes is not None: + if len(expert_kernel_sizes) != num_experts: + raise ValueError( + f"expert_kernel_sizes must have {num_experts} entries, got {len(expert_kernel_sizes)}" + ) + ks = [] + for k in expert_kernel_sizes: + k = int(k) + if k % 2 == 0: + k -= 1 + ks.append(min(k, max_kernel_size)) else: - ks = [min(3 + 2 * i, max_kernel_size) for i in range(num_experts)] + default_kernel_sizes = [3, 5, 7] + if num_experts <= len(default_kernel_sizes): + ks = [min(k, max_kernel_size) for k in default_kernel_sizes[:num_experts]] + else: + ks = [min(3 + 2 * i, max_kernel_size) for i in range(num_experts)] self.experts = nn.ModuleList( [EfficientExpertGroup(in_channels, out_channels, kernel_size=k) for k in ks] ) diff --git a/ultralytics/nn/modules/moe/pruning.py b/ultralytics/nn/modules/moe/pruning.py index 913c4e79..963f9ab2 100644 --- a/ultralytics/nn/modules/moe/pruning.py +++ b/ultralytics/nn/modules/moe/pruning.py @@ -442,8 +442,30 @@ def _sync_yaml_num_experts(self, pruned_model: nn.Module) -> None: continue args = list(seq[j][3]) if isinstance(seq[j][3], list) else [seq[j][3]] out_ch = args[0] if args else getattr(mod, "out_channels", None) - seq[j][3] = [out_ch, int(mod.num_experts)] - LOGGER.info(" Synced post-prune expert counts into model.yaml") + # Recover the kept experts' actual depthwise kernel sizes; pruning keeps + # experts with heterogeneous kernels whose order/values are not the + # default [3, 5, 7...] a bare rebuild would assign. Writing the full + # positional arg list (incl. expert_kernel_sizes) lets YOLO(pruned.pt) + # .train() rebuild the exact experts so their weights survive + # intersect_dicts instead of being dropped on a kernel-shape mismatch. + kernels = [] + for expert in mod.experts: + conv = getattr(expert, "conv", None) + depthwise = getattr(conv, "depthwise", None) if conv is not None else None + ks = getattr(depthwise, "kernel_size", None) + kernels.append(int(ks[0]) if isinstance(ks, (tuple, list)) else int(ks) if ks else 3) + top_k = int(mod.top_k) if getattr(mod, "use_top_k", True) else None + seq[j][3] = [ + out_ch, + int(mod.num_experts), + int(getattr(mod, "reduction", 8)), + top_k, + bool(getattr(mod, "use_sparse_inference", True)), + float(getattr(mod, "dynamic_threshold", 0.4)), + int(getattr(mod, "max_kernel_size", 15)), + kernels, + ] + LOGGER.info(" Synced post-prune expert counts + kernel sizes into model.yaml") def _save_model(self, pruned_model: nn.Module, output_path: str) -> None: """