Skip to content

fix: enforce weights-only loading for MemDecode router checkpoints - #4988

Merged
lvhan028 merged 1 commit into
InternLM:mainfrom
lvhan028:fix/memdecode-safe-checkpoint-loading
Sep 21, 2026
Merged

lvhan028 merged 1 commit into
InternLM:mainfrom
lvhan028:fix/memdecode-safe-checkpoint-loading

Conversation

@lvhan028

Copy link
Copy Markdown
Collaborator

No description provided.

Copilot AI lite review requested due to automatic review settings September 18, 2026 09:03

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟢 Approval recommended

The change is small, consistent with existing weights_only=True usage elsewhere in the repo, and directly addresses safer checkpoint loading without altering broader control flow.

Pull request overview

This PR hardens MemDecode router checkpoint loading by forcing PyTorch .pt router checkpoints to be deserialized using PyTorch’s weights-only pathway, reducing the risk of unsafe/unintended pickle object loading.

Changes:

  • Load .pt router checkpoints with torch.load(..., weights_only=True) when extracting router config.
  • Load .pt router checkpoints with torch.load(..., weights_only=True) when extracting the router state dict.
File summaries
File Description
lmdeploy/pytorch/memdecode/fusion.py Enforces weights-only deserialization for .pt router checkpoint reads (config + state dict).
Review details
  • Files reviewed: 1/1 changed files
  • Comments generated: 0
  • Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

@RunningLeon RunningLeon left a comment

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.

LGTM

@lvhan028
lvhan028 merged commit 82eb64e into InternLM:main Sep 21, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants