Skip to content

[Fix] Reject pickle on HTTP /update_weights - #4964

Open
gokay-ai wants to merge 3 commits into
InternLM:mainfrom
gokay-ai:cursor/fix-update-weights-pickle-1fc2
Open

gokay-ai wants to merge 3 commits into
InternLM:mainfrom
gokay-ai:cursor/fix-update-weights-pickle-1fc2

Conversation

@gokay-ai

@gokay-ai gokay-ai commented Sep 14, 2026

Copy link
Copy Markdown

Motivation

POST /update_weights deserialized request-controlled pickle (ForkingPickler.loads) before any tensor validation. On the default api_server bind (0.0.0.0, no api_keys), that is unauthenticated RCE.

Fixes #4698

Modification

  • HTTP /update_weights rejects pickle by default. LMDEPLOY_ALLOW_PICKLE_UPDATE_PARAMS=1 restores trusted same-node XTuner CUDA IPC (FlattenedTensorBucket, metadata-only reuse, empty finished=true). The HTTP handler still does not call pickle.loads.
  • HTTP also accepts load_format="safetensors" or a structured dict of tensors ({dtype, shape, data}).
  • serialize_named_tensors_safetensors clones overlapping CPU storage so Qwen3 tied embed / LM-head weights can be saved.
  • The pickle opt-in is registered in lmdeploy.pytorch.envs so Ray get_all_envs() forwards it to workers. Default remains deny.
  • Docs updated for the safetensors path and the XTuner opt-in.

BC-breaking

Clients that POST pickle from serialize_state_dict to /update_weights get HTTP 400 unless the server sets LMDEPLOY_ALLOW_PICKLE_UPDATE_PARAMS=1. Untrusted clients should switch to serialize_named_tensors_safetensors with load_format='safetensors', or use /update_weights_from_distributed / /update_weights_from_ipc.

Checklist

  1. Pre-commit / lint to be run in CI.
  2. Unit tests cover HTTP rejection, opt-in forwarding, tied-weight cloning, and Ray env registration.
  3. N/A
  4. docs/en/advance/update_weights.md and docs/zh_cn/advance/update_weights.md updated.

@gokay-ai
gokay-ai marked this pull request as ready for review September 14, 2026 11:48
Copilot AI lite review requested due to automatic review settings September 14, 2026 11:48

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.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Comment thread lmdeploy/utils.py Outdated
Comment thread docs/en/advance/update_weights.md
POST /update_weights deserialized request-controlled pickle data before
validation, which is unauthenticated RCE on the default api_server bind.
Reject pickle over HTTP, accept safetensors or structured tensors, and gate
engine pickle loads behind LMDEPLOY_ALLOW_PICKLE_UPDATE_PARAMS=1.

HTTP pickle rejection is covered even when the engine opt-in env var is set,
including that reconstruction callbacks and pickle.loads are never invoked.

Fixes InternLM#4698

Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com>
@cursor
cursor Bot force-pushed the cursor/fix-update-weights-pickle-1fc2 branch from 289c8e5 to a994860 Compare September 16, 2026 10:25
Clone CPU tensors that share storage so Qwen3 tied embed/LM-head
weights can be saved with serialize_named_tensors_safetensors.

Honor LMDEPLOY_ALLOW_PICKLE_UPDATE_PARAMS=1 on HTTP /update_weights so
XTuner FlattenedTensorBucket CUDA IPC (metadata-only reuse and empty
finished=true) works. Default still rejects pickle for untrusted clients.
The HTTP handler still does not pickle.loads.

Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com>

Co-authored-by: GokayAI <gokay-ai@users.noreply.github.com>
@gokay-ai

Copy link
Copy Markdown
Author

Addressed both review threads:

  1. serialize_named_tensors_safetensors now clones CPU tensors that share storage, so Qwen3 tied embed / LM-head weights no longer make safetensors.save raise.

  2. LMDEPLOY_ALLOW_PICKLE_UPDATE_PARAMS=1 on the server now actually restores XTuner same-node CUDA IPC over HTTP /update_weights (FlattenedTensorBucket, metadata-only reuse, empty finished=true). Default still rejects pickle for untrusted clients, and the HTTP handler still does not pickle.loads.

Comment thread lmdeploy/pytorch/engine/model_agent/agent.py
Register LMDEPLOY_ALLOW_PICKLE_UPDATE_PARAMS in pytorch.envs so
get_all_envs() copies it onto Ray workers. Default remains deny.

Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com>
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.

[Bug] /update_weights deserializes request-controlled pickle data before validation

4 participants