How to update weights on multi nodes?

[BraveY]

The readme doc has provided following command to update weights on single node.

torchrun --nproc-per-node 8 examples/update.py --update-method all --checkpoint-path /opt/models/Qwen/Qwen3-235B-A22B-Instruct-2507/

But how to update weights on multi nodes? The benchmark in readme docs has provided the multi nodes's results. Could we add multi-machine test commands?
Thanks!

[BraveY]

I tried the following commands without staring vllm
torchrun --nproc-per-node 8 --node-rank 0 --nnodes 2 --master-addr 192.168.0.1 --master-port 29500 examples/update.py --update-method broadcast --checkpoint-path /home/admin/model/DeepSeek-R1/ It works fine.
But logs stucked at the following lines on both nodes.

node1:

2025-09-17 12:25:21.288 | INFO     | checkpoint_engine.ps:register_named_tensors:484 - [rank15] p2p store register tensor memory_pool_my-checkpoint-iter-0_8 with addr 0x7f98e0000000 and capacity 4290886656
2025-09-17 12:25:21.288 | INFO     | checkpoint_engine.ps:register_named_tensors:484 - [rank15] p2p store register tensor memory_pool_my-checkpoint-iter-0_9 with addr 0x7f97e0000000 and capacity 4287625216
2025-09-17 12:25:21.288 | INFO     | checkpoint_engine.ps:register_named_tensors:484 - [rank15] p2p store register tensor memory_pool_my-checkpoint-iter-0_10 with addr 0x7f96e0000000 and capacity 4287625216
2025-09-17 12:25:21.288 | INFO     | checkpoint_engine.ps:register_named_tensors:484 - [rank15] p2p store register tensor memory_pool_my-checkpoint-iter-0_11 with addr 0x7f95e0000000 and capacity 2352584192

Is it because I haven't started vllm yet, or because I have installed mooncake-transfer-engine? When testing on a single machine without the mooncake-transfer-engine library, the log shows waiting for vllm in checkVllmReady while vllm not ready.
Like this:

2025-09-16 21:30:55.154 | WARNING  | __main__:check_vllm_ready:38 - fail to check vllm ready, retry 1 times, error: HTTPConnectionPool(host='localhost', port=19730): Max retries exceeded with url: /health (Caused by NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7f08d0c0af90>: Failed to establish a new connection: [Errno 111] Connection refused'))
2025-09-16 21:31:00.155 | WARNING  | __main__:check_vllm_ready:38 - fail to check vllm ready, retry 2 times, error: HTTPConnectionPool(host='localhost', port=19730): Max retries exceeded with url: /health (Caused by NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7f08d0c0ba70>: Failed to establish a new connection: [Errno 111] Connection refused'))

[BraveY]

It seems to be a problem with the version of mooncake-transfer-engine.
After 40 seconds, the following error was finally obtained:

heckpoint-iter-0_8 with addr 0x7f7840000000 and capacity 4287575040
2025-09-17 12:25:00.838 | INFO     | checkpoint_engine.ps:register_named_tensors:484 - [rank4] p2p store register tensor memory_pool_my-checkpoint-iter-0_9 with addr 0x7f7740000000 and capacity 3920473088
2025-09-17 12:24:40.679 | ERROR    | checkpoint_engine.ps:register_checkpoint:563 - [rank6] fail to register checkpoint my-checkpoint-iter-0 with files []
Traceback (most recent call last):

  File "/home/admin/model/checkpoint-engine/examples/update.py", line 159, in <module>
    update_weights(
    └ <function update_weights at 0x7f491f89b760>

  File "/home/admin/model/checkpoint-engine/examples/update.py", line 92, in update_weights
    ps.register_checkpoint(checkpoint_name, files=checkpoint_files, named_tensors=named_tensors)
    │  │                   │                      │                               └ {'model.layers.25.mlp.experts.23.down_proj.weight_scale_inv': tensor([[0.0001, 0.0001, 0.0001, 0.0001, 0.0001, 0.0001, 0.0001...
    │  │                   │                      └ []
    │  │                   └ 'my-checkpoint-iter-0'
    │  └ <function ParameterServer.register_checkpoint at 0x7f491f89a950>
    └ <checkpoint_engine.ps.ParameterServer object at 0x7f491f885420>

> File "/home/admin/model/checkpoint-engine/checkpoint_engine/ps.py", line 561, in register_checkpoint
    self._register_parameters_to_p2p_store(checkpoint_name)
    │    │                                 └ 'my-checkpoint-iter-0'
    │    └ <function ParameterServer._register_parameters_to_p2p_store at 0x7f491f89af80>
    └ <checkpoint_engine.ps.ParameterServer object at 0x7f491f885420>

  File "/home/admin/model/checkpoint-engine/checkpoint_engine/ps.py", line 873, in _register_parameters_to_p2p_store
    self._p2p_store.register_named_tensors(named_tensors)
    │    │          │                      └ {'memory_pool_my-checkpoint-iter-0_0': tensor([105,  61,  94,  ...,   2, 247,  56], dtype=torch.uint8), 'memory_pool_my-check...
    │    │          └ <function P2PStore.register_named_tensors at 0x7f491f89a560>
    │    └ <checkpoint_engine.ps.P2PStore object at 0x7f491f885c00>
    └ <checkpoint_engine.ps.ParameterServer object at 0x7f491f885420>

  File "/home/admin/model/checkpoint-engine/checkpoint_engine/ps.py", line 487, in register_named_tensors
    assert self.engine.batch_register_memory(buffer_addresses, capacities) == 0
           │    │                            │                 └ [4287639552, 4287628800, 4287653888, 4287625216, 4287625216, 4287592960, 4287625216, 4287578624, 4287621632, 3847054848]
           │    │                            └ [139879031767040, 139874736799744, 139870441832448, 139866146865152, 139861851897856, 139857556930560, 139853261963264, 13984...
           │    └ <mooncake.engine.TransferEngine object at 0x7f492041eef0>
           └ <checkpoint_engine.ps.P2PStore object at 0x7f491f885c00>

AttributeError: 'mooncake.engine.TransferEngine' object has no attribute 'batch_register_memory'. Did you mean: 'register_memory'?
Traceback (most recent call last):
  File "/home/admin/model/checkpoint-engine/checkpoint_engine/ps.py", line 561, in register_checkpoint
    self._register_parameters_to_p2p_store(checkpoint_name)
  File "/home/admin/model/checkpoint-engine/checkpoint_engine/ps.py", line 873, in _register_parameters_to_p2p_store
    self._p2p_store.register_named_tensors(named_tensors)
  File "/home/admin/model/checkpoint-engine/checkpoint_engine/ps.py", line 487, in register_named_tensors
    assert self.engine.batch_register_memory(buffer_addresses, capacities) == 0
AttributeError: 'mooncake.engine.TransferEngine' object has no attribute 'batch_register_memory'. Did you mean: 'register_memory'?

[weixiao-huang]

Please provide the infos below to help us debug you problem

• Which parallel size do you set to your inference engine, TP8 or TP16? You may set --inference-parallel-size to specify correct parallel size. For example, if set TP16 or EP16, specify --inference-parallel-size=16
• use pip show mooncake-transfer-engine to get it's version and try to upgrade it by pip install mooncake-transfer-engine --upgrade. The version we tested is mooncake-transfer-engine==0.3.6. Also, you can uninstall it by pip uninstall mooncake-transfer-engine and remove --update-method all to only test broadcast method instead of P2P.
• examples/update.py will wait vLLM ready before starting to update weights. The timeout is default to 10m due to dist.barrier(). So if your vLLM has not started in 10m after examples/update.py's checking vLLM ready, the update scripts will fail

[BraveY]

Please provide the infos below to help us debug you problem

• Which parallel size do you set to your inference engine, TP8 or TP16? You may set --inference-parallel-size to specify correct parallel size. For example, if set TP16 or EP16, specify --inference-parallel-size=16
• use pip show mooncake-transfer-engine to get it's version and try to upgrade it by pip install mooncake-transfer-engine --upgrade. The version we tested is mooncake-transfer-engine==0.3.6. Also, you can uninstall it by pip uninstall mooncake-transfer-engine and remove --update-method all to only test broadcast method instead of P2P.
• examples/update.py will wait vLLM ready before starting to update weights. The timeout is default to 10m due to dist.barrier(). So if your vLLM has not started in 10m after examples/update.py's checking vLLM ready, the update scripts will fail

Thanks for the reply.
The version used in my test is our inner version 0.3.4.post2. I will uninstall mooncake-transfer-engine to test broadcast method.
Let's say we use tp=16 on two nodes which has 8 GPU on single node, Is my following command correct? Is there any other parameter I missed?
torchrun --nproc-per-node 8 --node-rank 0 --nnodes 2 --master-addr 192.168.0.1 --master-port 29500 examples/update.py --update-method broadcast --checkpoint-path /home/admin/model/DeepSeek-R1/ --inference-parallel-size=16 .

[weixiao-huang]

Your command is correct. Also, you can install the latest mooncake-transfer-engine==0.3.6 to test update by P2P

[weixiao-huang]

It seems batch_register_memory is introduced in mooncake-transfer-engine>=0.3.5. I'll add this version check in pyproject.toml soon, see #13