Skip to content

fix(moe): preserve DynamicRoutingLayer compatibility with released checkpoints - #158

Merged
isLinXu merged 1 commit into
Tencent:mainfrom
SidKC:fix/released-checkpoint-router-compat
Jul 22, 2026
Merged

fix(moe): preserve DynamicRoutingLayer compatibility with released checkpoints#158
isLinXu merged 1 commit into
Tencent:mainfrom
SidKC:fix/released-checkpoint-router-compat

Conversation

@SidKC

@SidKC SidKC commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

说明

  • 为旧版本反序列化得到的 DynamicRoutingLayer 增加兼容处理:当实例中没有 in_channels 属性时,从 routing_network 的第一个卷积层读取输入通道数;
  • 新创建的 router 仍然优先使用已有的 in_channels,保持当前校验逻辑不变;
  • 补充旧 checkpoint 对象的输出一致性、state 一致性和错误通道输入测试。

问题背景

官方 v26.02 YOLO-Master-EsMoE-N.pt 中包含四个较早版本序列化的 DynamicRoutingLayer。这些对象生成时还没有 self.in_channels 属性,而当前 forward() 会直接读取该属性。因此 checkpoint 本身可以加载,但第一次执行 forward 时会报错:

AttributeError: 'DynamicRoutingLayer' object has no attribute 'in_channels'

所需的输入通道数实际上已经保存在 routing_network 的第一个 Conv2d 中;当前模块也已有 _get_router_in_channels(),可以安全读取这一结构信息。

修改方式

in_channels 存在时继续使用原值;只有旧对象缺少该属性时,才回退到 _get_router_in_channels(self.routing_network)。该修改不会增加或改变任何 parameter、buffer,也不会改变 checkpoint key。

验证结果

  • 旧 checkpoint 兼容性定向测试:2 passed
  • tests/test_moe_router_boundaries.py36 passed
  • tests/test_mixture_fixes.py10 passed
  • 官方 v26.02 checkpoint SHA256:29e1b93f09b16c8cf7c402f36dcaafc19d4812155631ed45b769e941e4c88c32
  • 在官方 checkpoint 上复现了修复前的 AttributeError
  • 确认四个 legacy router 均缺少该实例属性,fallback 得到的输入通道数为 64/128/128/256
  • 与显式补齐 metadata 的对照实现比较,各层 routing logits、probabilities 以及完整 detector raw output 的 max_abs_diff 均为 0.0
  • 修复前后的 state-dict keys 和 tensors 完全一致。

影响范围

这是一个针对已发布 pickle 模型对象的向后兼容修复,不会改变 routing weights、Top-K 行为、训练逻辑、导出逻辑或 checkpoint schema。

Related to #52

@SidKC
SidKC marked this pull request as ready for review July 22, 2026 04:22
Copilot AI review requested due to automatic review settings July 22, 2026 04:22

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR restores backward compatibility for legacy pickled DynamicRoutingLayer objects found in released checkpoints by avoiding direct reliance on a missing in_channels instance attribute during forward() input validation.

Changes:

  • Update DynamicRoutingLayer.forward() to derive expected input channels via a fallback when in_channels is absent.
  • Add regression tests to ensure legacy objects (missing in_channels) preserve outputs/state and still enforce channel-shape validation.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File Description
ultralytics/nn/modules/moe/routers.py Adds compatibility fallback for determining expected router input channels when legacy checkpoints lack in_channels.
tests/test_moe_router_boundaries.py Adds targeted tests covering legacy missing-attribute behavior (output/state invariance and channel-check enforcement).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +444 to +445
expected_channels = getattr(self, "in_channels", _get_router_in_channels(self.routing_network))
_validate_router_input(x, expected_channels, "DynamicRoutingLayer")
@isLinXu

isLinXu commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

LGTM

@isLinXu
isLinXu merged commit 504b362 into Tencent:main Jul 22, 2026
17 checks passed
@SidKC
SidKC deleted the fix/released-checkpoint-router-compat branch July 24, 2026 07:14
kub-inst pushed a commit to kub-inst/YOLO-Master that referenced this pull request Sep 4, 2026
…ter-compat

fix(moe): preserve DynamicRoutingLayer compatibility with released checkpoints
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.

3 participants