[Refactor] Follow weight update start/finish contract in WPI engine and example

Signed-off-by: Bill Du <yangspirit@google.com>

Author: yangspirit

docs/training/weight_transfer/wpi.md

@@ -89,7 +89,7 @@ from vllm.distributed.weight_transfer.wpi_engine import (
 )
 
 # Start weight update on inference side
-llm.start_weight_update(is_checkpoint_format=True)
+llm.start_weight_update()
 
 # Prepare send arguments with the trainer context
 trainer_args = WPITrainerSendWeightsArgs(

examples/rl/rlhf_http_wpi.py

@@ -75,6 +75,20 @@ def init_weight_transfer_engine(
     response.raise_for_status()
 
 
+def start_weight_update(base_url: str) -> None:
+    """Start a weight update via HTTP endpoint."""
+    url = f"{base_url}/start_weight_update"
+    response = requests.post(url, json={}, timeout=60)
+    response.raise_for_status()
+
+
+def finish_weight_update(base_url: str) -> None:
+    """Finish a weight update via HTTP endpoint."""
+    url = f"{base_url}/finish_weight_update"
+    response = requests.post(url, json={}, timeout=60)
+    response.raise_for_status()
+
+
 def pause_generation(base_url: str) -> None:
     """Pause generation via HTTP endpoint."""
     url = f"{base_url}/pause"
@@ -137,11 +151,13 @@ def main():
     socket_dir = "/run/wpi/sockets"
     driver_port = 50051
 
-    # Calculate model size in bytes
-    total_model_bytes = sum(
-        p.numel() * p.element_size() for p in train_model.parameters()
-    )
-    print(f"Calculated model size: {total_model_bytes} bytes")
+    # Calculate model size in bytes with alignment padding
+    total_model_bytes = 0
+    for p in train_model.parameters():
+        itemsize = p.element_size()
+        total_model_bytes = (total_model_bytes + itemsize - 1) // itemsize * itemsize
+        total_model_bytes += p.numel() * itemsize
+    print(f"Calculated model size with alignment: {total_model_bytes} bytes")
 
     print("Initializing weight transfer on server...")
 
@@ -165,6 +181,9 @@ def main():
     # Pause generation before weight sync
     pause_generation(BASE_URL)
 
+    # Start weight update, broadcast via WPI, then finish
+    start_weight_update(BASE_URL)
+
     # Broadcast all weights from trainer to vLLM workers
     print("Propagating weights via WPI...")
     param_iter = ((n, p) for n, p in train_model.named_parameters())
@@ -177,6 +196,8 @@ def main():
     # This will pack weights, trigger WPI propagate, and call /update_weights on server
     WPIWeightTransferEngine.trainer_send_weights(param_iter, args)
 
+    finish_weight_update(BASE_URL)
+
     # Resume generation after weight sync
     resume_generation(BASE_URL)
 

vllm/distributed/weight_transfer/wpi_engine.py

@@ -16,13 +16,19 @@
     │  → receive FD    │                │  → receive FD          │
     │  → CUDA import   │                │  → CUDA import         │
     │                  │                │  → connect notify sock │
+    │                  │                │                        │
+    │ /start_weight_upd├───(HTTP/Ray)──→│ start_weight_update()  │
+    │                  │                │                        │
     │ trainer_send_wts │                │                        │
     │  → pack into buf │  NodePropagate │ receive_weights()      │
     │  → propagate ────┼───(NCCL)──────→│  → wait_for_ready      │
     │  → HTTP metadata │                │  → unpack from buffer  │
-    │    /update_weights├───(HTTP)─────→│  → load_weights()      │
+    │    /update_weights├───(HTTP/Ray)─→│  → load_weights()      │
+    │                  │                │                        │
+    │ /finish_weight_up├───(HTTP/Ray)──→│ finish_weight_update() │
     └──────────────────┘                └───────────────────────┘
 
+
 Key differences from NCCL/IPC engines:
 - NCCL communicator is managed by the WPI driver, not by vLLM
 - VRAM buffer is persistent and reused across training steps (no per-update alloc)
@@ -498,6 +504,11 @@ def trainer_send_weights(
         The WPI driver handles NCCL communicator setup, broadcast execution,
         and READY notification to consumers — all transparently.
 
+        .. note::
+            This method calls ``update_weights`` internally. The caller must
+            call ``start_weight_update`` before and ``finish_weight_update``
+            after this method.
+
         Args:
             iterator: Iterator of (name, tensor) tuples from the model.
             trainer_args: Dict or WPITrainerSendWeightsArgs with WPI config.
@@ -552,6 +563,9 @@ def trainer_send_weights(
 
         for name, tensor in iterator:
             weight = tensor.detach().contiguous()
+            itemsize = weight.element_size()
+            # Align offset to itemsize boundary to prevent unaligned view errors
+            offset = (offset + itemsize - 1) // itemsize * itemsize
             nbytes = weight.nbytes
 
             if offset + nbytes > ctx.buffer_size_bytes: