[ckpt] feat: implement large tensor slicing in vllm rollout and CheckpointEngine for weight updating

[jianjunzhong]

What does this PR do?

This PR does the following:

1. Adds support for large tensors slicing in the naive(vllm)/nccl/hccl/nixl backends of CheckpointEngine.
2. Abstracts common logic from nccl/hccl/nixl into the CheckpointEngine and CollectiveCheckpointEngine classes.
3. Adds a unit test for the correctness of naive weight updates.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)
• If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.

[gemini-code-assist]

Code Review

This pull request introduces chunked weight handling for large tensors in vLLM rollouts, which is a valuable feature for working with large models. The changes involve slicing oversized weights on the sender side and reassembling them on the receiver side. A new test is also added to verify the correctness of this new functionality. My review identified a critical bug in the chunking logic that could lead to a ZeroDivisionError, and a high-severity issue in the new test related to a hardcoded, user-specific file path. I have provided code suggestions to address both of these points. Apart from these issues, the implementation appears solid.

[gemini-code-assist]

The model path is hardcoded to a user-specific local directory. This makes the test non-portable and difficult for other developers to run. It's recommended to use a model from the Hugging Face Hub that can be downloaded automatically, or at least make this path configurable via an environment variable. Using a smaller model for this test would also make it faster and more efficient.

Suggested change

	config.actor_rollout_ref.model.path = os.path.expanduser("~/models/Qwen/Qwen3-VL-2B-Instruct")
	config.actor_rollout_ref.model.path = os.path.expanduser(os.environ.get("VERL_TEST_MODEL_PATH", "hf-internal-testing/tiny-random-LlamaForCausalLM"))

[wuxibin89]

CheckpointEngine should also slice large tensor as well.

[jianjunzhong]

CheckpointEngine should also slice large tensor as well.

OK, I will also implement large tensor slicing in CheckpointEngine.

[pengwu22]

should merge after #5309
so that the newly-added logic can be thoroughly unit-tested.

[jianjunzhong]

should merge after #5309 so that the newly-added logic can be thoroughly unit-tested.

Thank you for the suggestion. Since similar logic has already been implemented and tested in CheckpointEngine, I think it would be reasonable to merge this PR first.

[gemini-code-assist]

Code Review

This pull request introduces support for large tensor slicing during weight updates in the CheckpointEngine and vllm rollout, refactors common logic into a new CollectiveCheckpointEngine base class, and adds a correctness test for the naive backend. While these changes enhance functionality and code structure, a high-severity security issue was identified: the use of pickle (via zmq.recv_pyobj) for deserializing metadata over the network poses a significant risk of Remote Code Execution (RCE). It is strongly recommended to switch to a safer serialization format like JSON. Additionally, there is a significant piece of duplicated code for the tensor slicing logic that could benefit from refactoring to improve long-term maintainability.

[gemini-code-assist]

The use of zmq.Socket.recv_pyobj() is insecure as it relies on the pickle module for deserialization. pickle is known to be vulnerable to arbitrary code execution if the input data is controlled by an attacker. Since this data is received over the network from the master process, an attacker who can connect to the ZMQ port or spoof the master can achieve Remote Code Execution (RCE) on the worker processes.

Recommendation: Use a safer serialization format such as JSON or msgpack. For non-serializable types like torch.Size or torch.dtype, convert them to strings or lists before sending and reconstruct them on the receiving end.

[jianjunzhong]

I think this PR is ready for review. cc @wuxibin89 @pengwu22