Optimize performance

Signed-off-by: Hollow Man <hollowman@opensuse.org>

Author: HollowMan6

3rdparty/Megatron-Bridge-workspace/Megatron-Bridge

@@ -54,6 +54,13 @@ policy:
         nixl:
           device: cuda
           topology: auto
+          transfer_mode: staged
+          buffer_count: 2
+          background_progress: false
+          load_batch_bucket_count: 1
+          direct_min_bytes: 2147483648
+          direct_stripe_count: 1
+          metadata_batch_size: 1
           backend_name: UCX
           backend_init_params:
             ucx_error_handling_mode: peer
@@ -226,16 +233,20 @@ The built-in NIXL backend is registered as `nixl`. It uses:
 
 - NIXL agents for memory registration and transfer
 - ZMQ messages for bucket metadata and transfer notifications
-- two reusable transfer buffers per worker for pipelined bucket movement
+- reusable transfer buffers per worker for pipelined bucket movement
+- optional direct source descriptors for contiguous policy tensor chunks
 - `split_weight_chunks()` and `merge_weight_chunk_batches()` for tensors larger
   than one bucket
 
-The NIXL backend chooses one of two topologies:
+The NIXL backend chooses one of three topologies:
 
 - If `train_world_size >= rollout_world_size`, each rollout rank is paired with
   a policy rank. Extra policy ranks are idle and do not materialize weights.
-- If `rollout_world_size > train_world_size`, policy rank 0 sends into a chain
-  of rollout ranks that forward buckets.
+- If `rollout_world_size > train_world_size`, `auto` uses a rollout-side
+  binary tree rooted at policy rank 0. Tree forwarding avoids the fully
+  serialized rollout chain.
+- `leader_chain` keeps the older single-policy-sender chain for debugging and
+  benchmark comparison.
 
 Set `engine_kwargs.nixl.topology=leader_chain` to force the single-policy-sender
 chain even when policy workers can cover rollout workers. This reduces
@@ -248,6 +259,71 @@ read handles alive for the lifetime of the worker. Reusing these objects avoids
 repeated multi-GB memory registration and transfer-handle initialization in
 long-lived Ray/vLLM actors.
 
+`transfer_mode=staged` is the conservative default: policy tensors are copied
+into reusable registered buckets, and small tensors are packed to reduce
+metadata overhead. `transfer_mode=auto` uses direct NIXL descriptors only for
+large contiguous policy tensors at least as large as the staged bucket size.
+Sub-bucket tensors stay staged because direct descriptors lose the coalescing
+benefit of staged buckets and can make source-side registration dominate. Use
+`transfer_mode=direct` to force direct descriptors for all contiguous tensors
+when benchmarking the raw source-direct path. Increase `buffer_count` when
+sender-side wait time shows that two staged buffers are not enough to keep
+reads in flight. The receiver also uses extra buffers to prefetch future reads
+before yielding the current transfer bucket to vLLM, while preserving buffer
+lifetime until vLLM has consumed the yielded views.
+
+`direct_stripe_count` splits each direct-transfer bucket into smaller direct
+read stripes. This is an opt-in direct/auto-mode tuning knob for clusters where
+one large NIXL read does not keep the available RDMA lanes busy. Pair it with
+enough receiver `buffer_count` headroom and set `metadata_batch_size` above `1`
+so the additional stripe descriptors do not turn into one ZMQ publish per
+stripe. Staged mode keeps using the normal bucket size and is unaffected by
+direct striping.
+
+Bucket metadata is still carried on the Python control path, but the NIXL
+backend can publish several metadata records in one control message. Receivers
+unwrap the batch before starting reads, so the data path and ordering contract
+remain one logical bucket at a time.
+
+When `background_progress=true`, each receiver read has a lightweight progress
+thread that polls NIXL state while the main thread is inside the synchronous
+vLLM load path. This reduces exposed wait time for reads that were started
+before the current load batch was yielded. It is disabled by default and should
+be enabled only when cluster benchmarks show stable tail latency.
+
+`load_batch_bucket_count` can coalesce multiple received transfer buckets into
+one vLLM load call. NIXL uses explicit bucket leases for this path: chunks that
+were copied into a reconstructed tensor release their transfer bucket
+immediately, while chunks that are exposed as tensor views keep their bucket
+leased until the outer vLLM load batch returns. This preserves view lifetime
+while reducing Python load-call overhead. Benchmark values above `1` with
+enough `buffer_count` headroom, because coalescing too many buckets can reduce
+receive prefetch overlap.
+
+For CUDA-backed batches, NIXL returns a list subclass that implements
+`record_cuda_load_complete()`. vLLM calls this after `load_weights()` queues its
+copies. NIXL records CUDA events on the batch devices and waits on those events
+before recycling transfer buffers, which avoids a hard
+`torch.cuda.current_stream().synchronize()` after every load batch.
+
+Buffer recycling uses an async release queue. If the receiver already has
+another read in flight, the batch lease is released by a background task after
+the load-complete CUDA event is ready. If releasing the current bucket is needed
+to prefetch the next bucket, the release stays inline so the receive generator
+cannot run out of work and terminate early.
+
+On the sender, staged buckets use CUDA events instead of a device-wide sync
+before every metadata publish. The sender enqueues the staged copy, records
+events for the involved CUDA devices, and keeps filling later buffers. Metadata
+publication still happens in original bucket order; a pending bucket is exposed
+only after its copy event is complete, and buffer reuse waits for the previous
+read notification for that buffer.
+
+When the runtime NIXL Python agent exposes a native `progress()` method, the
+wrapper invokes it under the agent lock before polling transfer state or
+notifications. Runtimes without that method keep the existing
+`check_xfer_state()`/notification polling behavior.
+
 ### Fault-Tolerance Boundary
 
 The NIXL backend is restart-safe, not actor-healing. It is designed so a failed

docs/design-docs/checkpoint-engines.md

@@ -36,6 +36,13 @@ policy:
           device: cuda
           cleanup_after_load: false
           topology: auto
+          transfer_mode: staged
+          buffer_count: 2
+          background_progress: false
+          load_batch_bucket_count: 1
+          direct_min_bytes: 2147483648
+          direct_stripe_count: 1
+          metadata_batch_size: 1
           backend_name: UCX
           backend_init_params:
             ucx_error_handling_mode: none
@@ -57,7 +64,14 @@ participating worker. `2048` MiB is a good starting point for large models.
 |---|---|
 | `device` | Transfer-buffer device. Use `cuda` for the fast GPU-RDMA path when the NIXL/UCX runtime supports CUDA memory registration. Use `cpu` only as a host-pinned fallback when CUDA buffers are unavailable. |
 | `cleanup_after_load` | Whether vLLM should run garbage collection and `torch.cuda.empty_cache()` after loading each refit. Disabling this avoids extra steady-state overhead when memory is stable. |
-| `topology` | Transfer topology. `auto` keeps the paired topology when policy workers cover rollout workers and falls back to `leader_chain` otherwise. Use `leader_chain` to force one policy sender and rollout-side NIXL forwarding. |
+| `topology` | Transfer topology. `auto` keeps the paired topology when policy workers cover rollout workers and falls back to `leader_tree` otherwise. Use `leader_chain` to force the older single chain for comparison. |
+| `transfer_mode` | Source-side transfer mode. `staged` copies policy tensors into reusable NIXL buckets. `direct` exposes contiguous policy tensor chunks directly through NIXL and stages only tensors that cannot be directly exposed. `auto` stages sub-bucket tensors but uses direct descriptors for contiguous tensors at or above `max(direct_min_bytes, update_weights_bucket_megabytes)`. |
+| `buffer_count` | Number of reusable staged transfer buffers. `2` is the original ping-pong path. `3` or higher allows more outstanding staged reads before a sender reuses a buffer. |
+| `background_progress` | Whether the receiver starts a lightweight background progress thread for in-flight NIXL reads. This lets NIXL progress continue while vLLM synchronously loads the previous batch. Keep this `false` unless the target cluster shows stable tail latency with it enabled. |
+| `load_batch_bucket_count` | Number of received transfer buckets to coalesce into one vLLM load batch. Keep this at `1` unless enough buffers are available to coalesce without starving receive prefetch; benchmark `2` with `buffer_count >= 4`. |
+| `direct_min_bytes` | Minimum tensor size for `transfer_mode: auto` direct transfer. Auto mode also floors this at the staged bucket size, because direct descriptors for sub-bucket tensors lose the coalescing benefit of staged buckets. Start at the bucket size, for example `2147483648` with 2048 MiB buckets, and tune against the staged baseline. |
+| `direct_stripe_count` | Number of stripes to split each direct-transfer bucket into. This only affects `transfer_mode: auto` or `direct`; `1` preserves the existing one-read-per-bucket behavior. Try `2` or `4` only with enough `buffer_count` headroom for concurrent receiver reads. |
+| `metadata_batch_size` | Number of bucket metadata records to publish together on the control path. `1` preserves one ZMQ object per bucket. Values above `1` mainly help direct-striped runs where many direct stripe descriptors are emitted together. |
 | `backend_name` | NIXL backend plugin name, usually `UCX`. |
 | `backend_init_params` | Optional NIXL backend initialization parameters. Values are converted to strings before NIXL receives them. |
 
@@ -84,6 +98,9 @@ policy:
           device: cuda
           cleanup_after_load: false
           topology: auto
+          transfer_mode: staged
+          buffer_count: 2
+          direct_min_bytes: 2147483648
           backend_name: UCX
           backend_init_params:
             ucx_error_handling_mode: peer
@@ -141,6 +158,8 @@ uv run --extra mcore --extra vllm examples/run_grpo.py \
   ++policy.generation.checkpoint_engine.engine_kwargs.nixl.device=cpu \
   ++policy.generation.checkpoint_engine.engine_kwargs.nixl.cleanup_after_load=false \
   ++policy.generation.checkpoint_engine.engine_kwargs.nixl.topology=auto \
+  ++policy.generation.checkpoint_engine.engine_kwargs.nixl.transfer_mode=staged \
+  ++policy.generation.checkpoint_engine.engine_kwargs.nixl.buffer_count=2 \
   ++policy.generation.checkpoint_engine.engine_kwargs.nixl.backend_name=UCX \
   ++policy.generation.checkpoint_engine.engine_kwargs.nixl.backend_init_params.ucx_error_handling_mode=none
 ```
@@ -206,12 +225,19 @@ policy:
     checkpoint_engine:
       enabled: true
       backend: nixl
-      update_weights_bucket_megabytes: 2048
+      update_weights_bucket_megabytes: 1024
       engine_kwargs:
         nixl:
           device: cuda
           cleanup_after_load: false
           topology: auto
+          transfer_mode: staged
+          buffer_count: 3
+          background_progress: false
+          load_batch_bucket_count: 1
+          direct_min_bytes: 2147483648
+          direct_stripe_count: 1
+          metadata_batch_size: 1
           backend_name: UCX
           backend_init_params:
             ucx_error_handling_mode: none
@@ -247,18 +273,37 @@ Use measured timings to tune from there:
 - Prefer RDMA transports; logs should show `rc_mlx5`, not TCP-only transport.
 - Sweep bucket size around `1024` to `2048` MiB instead of assuming larger is
   always faster. On the latest tested 30B MoE CUDA-buffer refit, eight rails
-  with `2048` MiB buckets was best; for CPU fallback, `1536` MiB buckets were
-  best.
+  with `1024` MiB buckets and three staged buffers reduced steady vLLM receive
+  time from about `0.90s` to about `0.82s`; for CPU fallback, `1536` MiB
+  buckets were best.
 - Try multiple UCX rails when the fabric supports it. On the same 30B MoE
-  setup, eight CUDA rails reduced steady transfer/update to about `1.46s`,
-  compared with about `3.7s` for four CUDA rails and about `3.1s` for NCCL
-  broadcast. CPU fallback regressed with eight rails; keep it on four rails
-  unless a local sweep proves otherwise.
+  setup, eight CUDA rails plus three staged buffers reduced steady
+  transfer/update to about `1.36s` to `1.38s`, compared with about `3.7s` for
+  four CUDA rails and about `3.1s` for NCCL broadcast. CPU fallback regressed
+  with eight rails; keep it on four rails unless a local sweep proves otherwise.
 - Pass the RDMA NICs through NIXL UCX `backend_init_params.device_list` as well
   as process-level `UCX_NET_DEVICES` when using CUDA transfer buffers. On the
   tested 30B MoE setup, adding the eight-device list plus
-  `engine_config: MAX_RMA_RAILS=8` reduced steady transfer/update from about
-  `3.7s` to about `1.46s`.
+  `engine_config: MAX_RMA_RAILS=8`, then using `1024` MiB buckets with three
+  buffers, reduced steady transfer/update from about `3.7s` to about `1.36s`
+  to `1.38s`.
+- Benchmark `transfer_mode: auto` against the staged baseline. Auto mode keeps
+  small tensors packed in staged buckets, but exposes large contiguous policy
+  tensor chunks directly to NIXL to avoid the policy-side bucket copy. If direct
+  descriptor setup dominates on a model, keep `transfer_mode: staged`.
+- For direct or auto transfer, benchmark `direct_stripe_count: 2` or `4`
+  together with `metadata_batch_size` set to the same value. This can expose
+  multiple reads for a large tensor chunk while batching the Python control
+  messages. It is not expected to help `transfer_mode: staged`.
+- Benchmark `buffer_count: 3` against `2`. More staged buffers let the sender
+  keep additional reads outstanding before reusing a buffer, and the receiver
+  can prefetch future reads before yielding the current bucket to vLLM. Each
+  extra buffer reserves another `update_weights_bucket_megabytes` on every
+  participating worker.
+- Use `topology: leader_tree` when rollout workers outnumber policy workers.
+  `auto` selects this topology for that shape. Keep `leader_chain` only as a
+  comparison point or for debugging because the chain serializes rollout-side
+  forwarding.
 - Do not rely only on `UCX_MAX_RNDV_RAILS` for NIXL refit. NIXL checkpoint-engine
   transfers are UCX RMA READ operations, and the tested runtime still printed a
   single `rc_mlx5` device in the `rma(...)` log line even when the device-list

docs/guides/checkpoint-engine-refit.md

@@ -339,6 +339,13 @@ policy:
           device: cuda
           cleanup_after_load: true
           topology: auto
+          transfer_mode: staged
+          buffer_count: 2
+          background_progress: false
+          load_batch_bucket_count: 1
+          direct_min_bytes: 2147483648
+          direct_stripe_count: 1
+          metadata_batch_size: 1
           backend_name: UCX
     colocated:
       # true: generation shares training GPUs

examples/configs/grpo_math_1B.yaml

@@ -69,6 +69,13 @@ policy:
           device: cuda
           cleanup_after_load: true
           topology: auto
+          transfer_mode: staged
+          buffer_count: 2
+          background_progress: false
+          load_batch_bucket_count: 1
+          direct_min_bytes: 2147483648
+          direct_stripe_count: 1
+          metadata_batch_size: 1
           backend_name: UCX
 
 cluster:

examples/configs/grpo_math_8B.yaml

@@ -422,7 +422,15 @@ async def _update_weights_from_checkpoint_engine_async(self) -> bool:
             load_time += time.time() - load_start
 
             sync_start = time.time()
-            torch.cuda.current_stream().synchronize()
+            record_cuda_load_complete = getattr(
+                weight_batch,
+                "record_cuda_load_complete",
+                None,
+            )
+            if callable(record_cuda_load_complete):
+                record_cuda_load_complete()
+            else:
+                torch.cuda.current_stream().synchronize()
             sync_time += time.time() - sync_start
             del weight_batch
 

nemo_rl/models/generation/vllm/vllm_backend.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 
 import asyncio
+import concurrent.futures
 import copy
 import gc
+import inspect
 import threading
 import time
 import uuid
@@ -44,6 +46,27 @@
 from nemo_rl.models.generation.vllm.vllm_worker import BaseVllmGenerationWorker
 
 
+async def _resolve_collective_rpc_result(result: Any) -> Any:
+    while inspect.isawaitable(result):
+        result = await result
+
+    if isinstance(result, concurrent.futures.Future):
+        return await _resolve_collective_rpc_result(await asyncio.wrap_future(result))
+
+    if isinstance(result, ray.ObjectRef):
+        return await asyncio.to_thread(ray.get, result)
+
+    if isinstance(result, list):
+        if all(isinstance(item, ray.ObjectRef) for item in result):
+            return await asyncio.to_thread(ray.get, result)
+        return [await _resolve_collective_rpc_result(item) for item in result]
+
+    if isinstance(result, tuple):
+        return tuple([await _resolve_collective_rpc_result(item) for item in result])
+
+    return result
+
+
 def _replace_prefix_tokens(
     tokenizer,
     model_prefix_token_ids: list[int],
@@ -832,9 +855,7 @@ async def init_collective_async(
 
     async def _checkpoint_engine_rpc_async(self, method_name: str, *args: Any) -> Any:
         result = await self.llm.collective_rpc(method_name, args=args)
-        if asyncio.iscoroutine(result):
-            return await result
-        return result
+        return await _resolve_collective_rpc_result(result)
 
     async def init_checkpoint_engine_async(
         self,

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -98,7 +98,13 @@ async def send_weights(
     async def receive_weight_batches(
         self,
     ) -> AsyncGenerator[list[tuple[str, torch.Tensor]], None]:
-        """Receive model weights in batches that share a transfer buffer."""
+        """Receive model weights in batches that share a transfer buffer.
+
+        Backends that yield CUDA tensors backed by reusable transfer buffers may
+        return a list subclass with ``record_cuda_load_complete()``. vLLM calls
+        that hook after queuing its load operation so the backend can fence
+        buffer reuse with a CUDA event instead of forcing a device-wide sync.
+        """
         raise NotImplementedError