Skip to content

[Fix] Load memdecode router checkpoints with weights_only - #4820

Open
zcxGGmu wants to merge 1 commit into
InternLM:mainfrom
zcxGGmu:safe-router-load
Open

[Fix] Load memdecode router checkpoints with weights_only#4820
zcxGGmu wants to merge 1 commit into
InternLM:mainfrom
zcxGGmu:safe-router-load

Conversation

@zcxGGmu

@zcxGGmu zcxGGmu commented Aug 4, 2026

Copy link
Copy Markdown

Motivation

Issue #3255 reports unsafe torch.load(...) usage for .pt checkpoint loading. The original load_weight_ckpt helper no longer exists on current main, but the memdecode adaptive-router path still loads user-provided .pt router checkpoints without explicitly setting weights_only=True.

Refs #3255

Modification

  • Load memdecode router checkpoint metadata with torch.load(..., weights_only=True).
  • Load memdecode router state dicts with torch.load(..., weights_only=True).
  • Add regression coverage to ensure both router checkpoint load paths pass weights_only=True explicitly.

BC-breaking (Optional)

No. Supported router checkpoints using tensors and primitive dict metadata continue to load; unsafe pickle globals are rejected earlier by PyTorch's weights-only loader.

Use cases (Optional)

N/A

Checklist

  1. Pre-commit or other linting tools are used to fix the potential lint issues.
    • python -m pre_commit run --files lmdeploy/pytorch/memdecode/fusion.py tests/pytorch/memdecode/test_fusion.py
  2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
    • python -m pytest tests/pytorch/memdecode/test_fusion.py -q
  3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
    • N/A
  4. The documentation has been modified accordingly, like docstring or example tutorials.
    • N/A; this is an internal checkpoint-loading hardening change.

Signed-off-by: zq <zhouquan1511@163.com>
)


def test_router_checkpoint_torch_load_uses_weights_only(monkeypatch, tmp_path):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I don't think we need TC for this. The change is a single parameter addition to torch.load; there's no business logic to regress.

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