feat(vLLM): support direct vLLM TP rollout engine and vLLM native IPC weight sync

[cboss6]

Summary

This PR adds a direct, text-only vLLM rollout engine and a full-weight
synchronization path for colocated FSDP training and vLLM tensor-parallel
inference.

The PR has two first-class deliverables:

1. Direct vLLM Rollout Engine — a new UniRL rollout backend that owns a
   spawned vLLM 0.27 runtime, TP placement, text generation, sampled-token
   log-probabilities, lifecycle management, and partial sleep/wake.
2. Native WTE Weight Synchronization — a lazy FSDP-to-vLLM TP publication
   path built on vLLM's native packed CUDA IPC Weight Transfer Engine.

The implementation is intentionally built around vLLM 0.27's native Weight
Transfer Engine (WTE) instead of introducing another UniRL-owned CUDA IPC
protocol. UniRL supplies the framework-specific scheduling, lazy canonical
weight export, topology validation, rank consensus, publication receipts, and
fail-stop policy. vLLM continues to own:

• packed CUDA IPC handle transport;
• physical-GPU handle selection;
• TP slicing;
• model-specific parameter fusion;
• layerwise reload;
• the model's native load_weights() implementation.

This work was extracted from
UniRL PR #430.
PR #430 covers the broader Qwen3-MoE train-inference parity experiment; this PR
isolates the reusable direct-vLLM rollout and weight-sync infrastructure.

The final PR does not modify the existing SGLang, TensorWeightSync,
BucketedWeightSender, or vLLM-Omni legacy IPC behavior.

Why these capabilities are delivered together

origin/main has vLLM-Omni integrations, but it does not have a direct
text-only VLLMRolloutEngine. Adding only a trainer-side WTE adapter would
therefore leave no direct vLLM runtime in UniRL to receive native updates.
Conversely, adding a rollout engine without a post-optimizer weight publication
path would only support inference from the initial checkpoint, not iterative RL
training.

In colocated RL training, the Actor and rollout engine time-share the same
physical GPUs:

1. FSDP owns the trainable model, gradients, and optimizer state while training.
2. vLLM owns TP-local inference weights and KV cache while generating.
3. After an optimizer step, the updated Actor weights must be published into
   vLLM before the next rollout.

Copying an entire model through CPU memory or checkpoint files introduces:

• GPU-to-CPU and CPU-to-GPU copies;
• host-memory and disk staging;
• a large rank-0 bottleneck;
• another framework-owned TP redistribution implementation;
• high peak memory and long synchronization latency.

Sending FSDP local shards directly is also incorrect. FSDP and vLLM do not
necessarily shard a parameter in the same way:

• FSDP normally shards storage by its training layout;
• vLLM uses row-, column-, vocab-, and model-specific parallel loaders;
• Qwen3-MoE expert weights may be packed differently in training and inference;
• vLLM kernels may require fused resident layouts.

The correct boundary is therefore a lazy stream of complete canonical
HuggingFace-format tensors. vLLM's native loader then performs the correct TP
slice and model-specific fusion.

Supported configuration

The currently validated scope is deliberately narrow:

• model: Qwen3-30B-A3B / qwen3_moe;
• training: single-node FSDP4;
• rollout: direct vLLM 0.27 TP4;
• compute and wire dtype: BF16;
• quantization: disabled;
• pipeline parallelism: 1;
• expert parallelism: 1;
• topology: one Actor rank and one vLLM TP rank per physical GPU;
• synchronization: full fine-tuning only;
• transport: vLLM native packed CUDA IPC;
• failure policy: fail-stop, without fallback to CPU staging or legacy ZMQ.

The reference recipe is:

examples/ar/qwen3_moe_grpo_30b_a3b_gsm8k_fsdp_vllm_tp4.yaml

Part I: Direct vLLM Rollout Engine

Purpose and boundary

VLLMRolloutEngine implements UniRL's BaseRolloutEngine contract for direct
text-only vLLM inference. It is distinct from VLLMOmniRolloutEngine:

• direct vLLM serves text-only Qwen3-MoE generation through vllm.LLM;
• vLLM-Omni serves multimodal/diffusion pipelines through its existing
  orchestrator;
• the two engines keep independent process, loading, and weight-sync
  implementations.

The direct engine provides:

• typed configuration through VLLMEngineConfig;
• TP-aware UniRL placement;
• a clean spawned Python interpreter for vLLM;
• synchronous request/response control over a local Pipe;
• generation through the existing UniRL text adapter;
• sampled-token processed log-probabilities;
• weight-only and KV-cache wake-up phases;
• native WTE control methods;
• fail-stop connection state and process-tree cleanup.

Rollout architecture

flowchart LR
    subgraph UniRL["UniRL worker process"]
        Manager["RolloutManager / ARTrainer"]
        Engine["VLLMRolloutEngine"]
        Adapter["TextLMAdapter"]
        State["VLLMConnectionState<br/>IDLE / INFLIGHT / BROKEN"]
        Pipe["multiprocessing.Pipe"]

        Manager -->|"generate(sample)"| Engine
        Engine --> Adapter
        Engine --> State
        Engine --> Pipe
    end

    subgraph Child["Spawned vLLM runtime"]
        Runtime["engine_process_main"]
        LLM["vllm.LLM"]
        Executor["vLLM multiprocess TP executor"]
        TP0["TP worker 0"]
        TPn["TP workers 1..N"]

        Pipe --> Runtime
        Runtime --> LLM --> Executor
        Executor --> TP0
        Executor --> TPn
    end

    TP0 -->|"plain token IDs, text, log-probs"| Runtime
    Runtime -->|"versioned response"| Pipe
    Engine -->|"build_response()"| Adapter
    Adapter -->|"version-stamped Sample"| Manager

Loading

TP placement and process isolation

The new engine opts into UniRL's rollout-TP placement contract through
_accepts_rollout_tp_kwargs = True.

Handle therefore:

1. reads config.tp_size;
2. builds TP rank metadata;
3. verifies that every TP group is on one node;
4. collects each outer Worker's physical CUDA_VISIBLE_DEVICES token;
5. injects the complete ordered TP device list into TP rank 0;
6. starts one vLLM runtime that owns the complete TP group.

Only UniRL TP rank 0 creates the child process and communicates with
vllm.LLM; sibling UniRL ranks still participate in FSDP and weight-sync
collectives.

flowchart TD
    Config["VLLMEngineConfig.tp_size = 4"]
    Handle["UniRL Handle TP layout"]
    Devices["ordered physical devices<br/>GPU0, GPU1, GPU2, GPU3"]
    Rank0["UniRL rollout TP rank 0"]
    Other["UniRL rollout TP ranks 1..3"]
    Child["one spawned vLLM runtime"]
    VTP["vLLM TP4 executor"]

    Config --> Handle --> Devices
    Devices --> Rank0 --> Child --> VTP
    Handle --> Other
    Other -->|"no child process; distributed no-op methods"| Other

Loading

Spawning isolates vLLM's global runtime state, CUDA environment, multiprocessing
configuration, and EngineCore process tree from the long-lived UniRL Worker.

Generation path

sequenceDiagram
    participant Trainer as ARTrainer
    participant Engine as VLLMRolloutEngine
    participant Adapter as TextLMAdapter
    participant Runtime as spawned runtime
    participant LLM as vllm.LLM

    Trainer->>Engine: generate(Sample)
    Engine->>Adapter: build_inputs()
    Adapter-->>Engine: prompt IDs + sampling payloads
    Engine->>Engine: truncate to max_prompt_length
    Engine->>Runtime: Pipe generate request
    Runtime->>LLM: generate(prompts, SamplingParams)
    LLM-->>Runtime: RequestOutput list
    Runtime->>Runtime: extract selected-token processed log-probs
    Runtime-->>Engine: plain serializable outputs
    Engine->>Adapter: build_response()
    Engine->>Engine: stamp committed model version
    Engine-->>Trainer: generated Sample

Loading

The Pipe transports request metadata and generated outputs. It is not the
weight data plane.

Runtime control protocol

Every request carries:

• a protocol version;
• a monotonic request ID;
• a command name;
• command-specific payload.

Every response echoes the same identity and carries:

• ok;
• a typed result;
• a formatted exception and traceback on failure.

VLLMRolloutEngine permits only one in-flight request on the synchronous Pipe.
The connection state is:

stateDiagram-v2
    [*] --> IDLE
    IDLE --> INFLIGHT: send request
    INFLIGHT --> IDLE: validated response
    INFLIGHT --> BROKEN: timeout / EOF / protocol error
    INFLIGHT --> BROKEN: sleep, wake, or weight-update failure
    BROKEN --> [*]: terminate process tree

Loading

An invalid response ID, command, protocol version, result type, timeout, or
broken Pipe marks the runtime unusable and triggers process-tree cleanup.

Sleep and wake lifecycle

The direct rollout supports vLLM sleep levels and tagged wake-up:

• initial bootstrap sleep uses level 1 so checkpoint weights are preserved for
  the first rollout;
• later native-sync sleeps use level 2 so vLLM can discard old weights;
• wake_up(tags=["weights"]) restores only the storage required for weight
  publication;
• a later full wake_up() restores KV-cache resources before generation.

flowchart LR
    Initial["checkpoint-loaded runtime"]
    Sleep1["bootstrap sleep level 1<br/>preserve weights"]
    First["first rollout"]
    Sleep2["steady-state sleep level 2<br/>discard stale weights"]
    WeightWake["wake weights"]
    Reload["native WTE reload"]
    KVWake["wake KV cache"]
    Next["next rollout"]

    Initial --> Sleep1 --> First --> Sleep2 --> WeightWake --> Reload --> KVWake --> Next
    Next --> Sleep2

Loading

Native WTE control surface

The rollout engine exposes a narrow adapter around vLLM's native lifecycle:

• init_native_weight_transfer();
• start_native_weight_update();
• update_native_weights();
• finish_native_weight_update();
• release_native_ipc();
• poison().

These methods do not implement tensor loading. They validate publication state,
forward control to the spawned runtime, validate the child response, and mark
the connection broken after an unsafe failure.

Rollout-engine files

• unirl/rollout/engine/vllm/__init__.py
• unirl/rollout/engine/vllm/config.py
• unirl/rollout/engine/vllm/engine.py
• unirl/rollout/engine/vllm/runtime.py
• unirl/rollout/engine/vllm/worker_extension.py
• unirl/distributed/group/handle.py
• unirl/rollout/engine/base.py

Part II: Native WTE Weight Synchronization

Design principles

Reuse community-owned mechanisms

The data plane is based on vLLM's native:

• WeightSource;
• VLLMWeightSyncClient;
• IPCTrainerWeightTransferEngine;
• worker-side IPCWeightTransferEngine;
• packed IPC producer/consumer format;
• layerwise reload and load_weights() lifecycle.

The design also follows the central idea used by Verl's colocated vLLM sync:

• materialize FSDP full tensors lazily;
• pair Actor rank r with vLLM TP rank r on the same GPU;
• transfer CUDA IPC handles and metadata, not tensor bytes through the control
  plane;
• let vLLM perform TP-aware loading;
• do not add a second Actor-to-vLLM NCCL weight broadcast.

Keep UniRL policy outside vLLM

UniRL remains responsible for:

• training/rollout memory scheduling;
• Qwen3-MoE canonical export planning;
• physical CUDA UUID validation;
• expected and actual structural manifests;
• rank-synchronous failure propagation;
• monotonic model-version publication;
• prefix-cache invalidation;
• poisoning the runtime after a partial failure.

Fail closed

The update mutates resident vLLM storage during layerwise loading; it is not an
atomic weight swap and cannot roll back a partially loaded model. Safety is
therefore provided by fail-stop publication:

• no generation is allowed during an update;
• the version is committed only after all TP workers finish and attest;
• any failed update poisons and destroys the rollout runtime;
• the code never retries a failed native update through legacy ZMQ or CPU
  staging.

High-level architecture

flowchart LR
    subgraph Trainer["UniRL trainer process group"]
        AR["ARTrainer"]
        FSDP["FSDP Actor ranks"]
        Sync["IPCWeightSync"]
        Source["VLLMNativeWeightSource"]
        NativeTrainer["VLLMNativeIPCTrainerEngine"]
        Client["VLLMNativeWeightSyncClient"]

        AR -->|"selective offload + sync()"| Sync
        FSDP -->|"lazy canonical iterator"| Source
        Sync --> Source
        Sync --> NativeTrainer
        Source --> NativeTrainer
        NativeTrainer --> Client
    end

    subgraph Control["Spawned-runtime control plane"]
        Rollout["VLLMRolloutEngine"]
        Pipe["versioned synchronous Pipe protocol"]
        Runtime["engine_process_main"]

        Client -->|"init / start / update / finish / release"| Rollout
        Rollout --> Pipe --> Runtime
    end

    subgraph VLLM["vLLM 0.27 native weight-transfer stack"]
        LLM["vllm.LLM"]
        WorkerWTE["IPCWeightTransferEngine on TP workers"]
        Consumer["packed IPC consumer"]
        Loader["layerwise reload + model.load_weights()"]
        Resident["TP-local resident weights"]

        Runtime -->|"LLM.init/update/finish_weight_*"| LLM
        LLM --> WorkerWTE --> Consumer --> Loader --> Resident
    end

    NativeTrainer -.->|"merged UUID-to-handle metadata"| WorkerWTE
    FSDP ==>|"same-GPU CUDA IPC payload"| Consumer

Loading

Rank and data-plane topology

Every Actor rank materializes the same canonical tensor stream because
DTensor.full_tensor() is a collective. Each rank copies its materialized
tensor into a reusable packed buffer on its own GPU. vLLM's trainer engine
all-gathers the IPC handles and rank 0 sends a merged UUID-to-handle map through
the control plane.

Each vLLM TP worker selects the handle matching its physical CUDA UUID.

flowchart TB
    subgraph FSDP["FSDP4 collective domain"]
        A0["Actor rank 0<br/>GPU UUID 0"]
        A1["Actor rank 1<br/>GPU UUID 1"]
        A2["Actor rank 2<br/>GPU UUID 2"]
        A3["Actor rank 3<br/>GPU UUID 3"]
    end

    Gather["DTensor.full_tensor()<br/>FSDP NCCL all-gather"]
    A0 --> Gather
    A1 --> Gather
    A2 --> Gather
    A3 --> Gather
    Gather --> A0
    Gather --> A1
    Gather --> A2
    Gather --> A3

    subgraph IPC["Same-GPU CUDA IPC; no Actor-to-vLLM NCCL"]
        A0 ==>|"UUID 0 handle"| T0["vLLM TP rank 0"]
        A1 ==>|"UUID 1 handle"| T1["vLLM TP rank 1"]
        A2 ==>|"UUID 2 handle"| T2["vLLM TP rank 2"]
        A3 ==>|"UUID 3 handle"| T3["vLLM TP rank 3"]
    end

    subgraph Load["vLLM-native TP loading"]
        T0 --> L0["load_weights: TP slice/fusion"]
        T1 --> L1["load_weights: TP slice/fusion"]
        T2 --> L2["load_weights: TP slice/fusion"]
        T3 --> L3["load_weights: TP slice/fusion"]
    end

Loading

Large-scale cross-GPU communication occurs only in the FSDP
full_tensor() collective. The Actor-to-vLLM boundary is same-GPU CUDA IPC.

Engine selection and legacy compatibility

IPCWeightSync remains the public full-weight IPC handler. It selects its
implementation from the concrete rollout engine:

flowchart TD
    Init["IPCWeightSync.__init__"]
    Name["_receiver_component_name()"]
    Resolve["_resolve_weight_sync_engine()"]
    Native["VLLM_NATIVE_WTE"]
    Legacy["LEGACY_ZMQ"]

    Init --> Name --> Resolve
    Resolve -->|"component_name() == vllm"| Native
    Resolve -->|"all existing non-direct-vLLM cases"| Legacy

    Native --> NativeSync["_sync_with_vllm_native_engine()"]
    Legacy --> LegacySync["_sync_with_legacy_zmq()<br/>existing main behavior"]

Loading

Only the direct VLLMRolloutEngine selects the native WTE path. Existing
vLLM-Omni and other legacy users continue through the original bucketed ZMQ
implementation. The following files are unchanged from origin/main:

• unirl/distributed/weight_sync/transfer/bucketed_transfer.py;
• unirl/distributed/weight_sync/transfer/sgl_compat.py;
• unirl/distributed/weight_sync/full/tensor.py;
• unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py.

IPCWeightSync class design

Class relationships

classDiagram
    direction LR

    class FullWeightSync {
        +_iter_full_tensors()
        +_iter_buckets()
    }

    class IPCWeightSync {
        -_weight_sync_engine
        -_gpu_memory_headroom
        -_control_group
        -_next_model_version
        +uses_gpu_streaming()
        +setup()
        +sync()
        +cleanup()
        +shutdown()
        -_receiver_component_name()
        -_resolve_weight_sync_engine()
        -_dist_ready()
        -_global_rank
        -_world_size
        -_all_gather()
        -_broadcast_from_rank_zero()
        -_consensus()
        -_target_topology()
        -_run_preflight()
        -_shape_metadata()
        -_expected_tensor_metadata()
        -_poison_rollout()
        -_sync_with_vllm_native_engine()
        -_sync_with_legacy_zmq()
        -_release_resources()
    }

    class VLLMNativeWeightSource {
        +metadata()
        +__iter__()
        +actual_metadata
    }

    class VLLMNativeIPCTrainerEngine {
        +initialize()
        +send_weights()
        -_send_planned_packed()
        -_plan_chunks()
        -_validate_tensor()
    }

    class VLLMNativeWeightSyncClient {
        +init_weight_transfer_engine()
        +start_weight_update()
        +update_weights()
        +finish_weight_update()
        +release_ipc_imports()
    }

    class IPCTrainerWeightTransferEngine {
        +_all_gather_and_merge_handles()
        +_do_send()
    }

    class VLLMRolloutEngine {
        +init_native_weight_transfer()
        +start_native_weight_update()
        +update_native_weights()
        +finish_native_weight_update()
        +release_native_ipc()
        +poison()
    }

    FullWeightSync <|-- IPCWeightSync
    IPCWeightSync --> VLLMNativeWeightSource
    IPCWeightSync --> VLLMNativeIPCTrainerEngine
    IPCWeightSync --> VLLMNativeWeightSyncClient
    IPCTrainerWeightTransferEngine <|-- VLLMNativeIPCTrainerEngine
    VLLMNativeWeightSyncClient --> VLLMRolloutEngine

Loading

Member-function responsibilities

Selection and setup

• __init__()

  • initializes existing full-weight export behavior;
  • resolves native WTE versus legacy ZMQ once;
  • stores packed-buffer size, memory safety margin, control timeout, and the
    next model version.

• _receiver_component_name()

  • obtains the concrete receiver through tensor_weight_sync_target when the
    rollout is a wrapper;
  • uses the public component_name() accessor rather than reading a private
    class field across modules.

• _resolve_weight_sync_engine()

  • maps direct vllm to VLLM_NATIVE_WTE;
  • preserves LEGACY_ZMQ for existing non-direct-vLLM behavior.

• uses_gpu_streaming()

  • exposes the resolved memory policy to ARTrainer;
  • is defined only on IPCWeightSync, so Tensor/NCCL/LoRA sync paths do not
    receive a new RPC or state.

• setup()

  • performs normal Remote setup;
  • creates a Gloo control group for native WTE;
  • configures direct vLLM to use sleep level 2 after the initial checkpoint
    rollout, so discarded inference weights are reconstructed by sync.

Rank control and failure consensus

• _dist_ready()

  • checks whether a multi-rank PyTorch process group is active.

• _global_rank and _world_size

  • expose stable Actor-rank identity for native synchronization.

• _all_gather()

  • gathers rank-local control reports through the dedicated Gloo group.

• _broadcast_from_rank_zero()

  • distributes rank-0 topology, publication ID, and finish result to all Actor
    ranks.

• _consensus(error, phase)

  • converts one rank-local exception into the same failure on every Actor rank;
  • prevents healthy ranks from entering a later FSDP collective after another
    rank has already failed.

Topology and memory preflight

• _target_topology()

  • reads vLLM TP fanout and physical worker CUDA UUIDs from
    VLLMRolloutEngine;
  • validates that one UUID is available for every TP worker.

• _run_preflight()

  • verifies one node, one visible CUDA device per Actor rank, equal
    FSDP/TP world sizes, and exact Actor/worker physical UUID ordering;
  • rejects a tensor larger than the packed buffer;
  • checks a conservative GPU-memory budget before the first lazy
    materialization.

The memory check is:

required_free_memory =
    3 * packed_buffer_size
  + largest_layer_size
  + gpu_memory_headroom

The three buffer-sized terms conservatively cover overlapping full-tensor,
trainer packed-buffer, and receiver/loading storage. largest_layer_size
covers deferred layerwise processing. gpu_memory_headroom_mb is an additional
safety margin for allocator fragmentation, CUDA/NCCL state, and temporary
kernel workspaces. It is a preflight threshold, not a reservation.

Canonical export plan

• _shape_metadata()

  • creates structural metadata without reading tensor payload bytes.

• _expected_tensor_metadata()

  • independently plans the expected Qwen3-MoE HuggingFace key stream;
  • expands packed gate_up_proj and down_proj expert tensors;
  • validates required embedding, norm, attention, router, and every expert
    projection key;
  • rejects unsupported LoRA merge, expert transforms, remapping, and non-Qwen3
    models.

The expected plan is intentionally independent from the actual lazy iterator.
If the exporter consistently omits one expert on every rank, comparing Actor
ranks only would not detect the bug; comparing against the independent plan
does.

Publication

• sync()

  • dispatches to native WTE only for direct vLLM;
  • otherwise executes the unchanged legacy ZMQ implementation.

• _sync_with_vllm_native_engine()

  • builds the expected manifest and publication header;
  • constructs VLLMNativeWeightSource,
    VLLMNativeWeightSyncClient, and
    VLLMNativeIPCTrainerEngine;
  • initializes WTE, streams planned chunks, verifies the actual Actor manifest,
    finalizes the update, releases IPC imports, validates receipts, and advances
    the model version;
  • poisons the rollout on any incomplete publication.

• _sync_with_legacy_zmq()

  • is the existing origin/main path for existing rollout engines;
  • is not used by direct vLLM native WTE.

• _poison_rollout()

  • destroys the spawned runtime after native publication failure.

• cleanup(), shutdown(), _release_resources()

  • destroy the dedicated Gloo group during normal or worker-driven teardown.

Internal method dependency graph

flowchart TD
    Construct["__init__()"]
    Receiver["_receiver_component_name()"]
    Resolve["_resolve_weight_sync_engine()"]
    Policy["uses_gpu_streaming()"]
    Setup["setup()"]
    Sync["sync()"]

    Construct --> Resolve
    Resolve --> Receiver
    Construct --> Policy
    Setup --> Dist["_dist_ready()"]
    Setup --> Sleep["rollout.set_weight_sync_sleep_level()"]

    Sync -->|"native"| Native["_sync_with_vllm_native_engine()"]
    Sync -->|"legacy"| Legacy["_sync_with_legacy_zmq()"]

    Native --> Expected["_expected_tensor_metadata()"]
    Expected --> Shape["_shape_metadata()"]
    Expected --> Manifest["structural_manifest()"]
    Native --> Target["_target_topology()"]
    Target --> Consensus["_consensus()"]
    Target --> Broadcast["_broadcast_from_rank_zero()"]
    Native --> Preflight["_run_preflight()"]
    Preflight --> Gather["_all_gather()"]
    Preflight --> Consensus
    Native --> Source["VLLMNativeWeightSource"]
    Native --> Client["VLLMNativeWeightSyncClient"]
    Native --> Engine["VLLMNativeIPCTrainerEngine"]
    Engine --> Consensus
    Native --> Validate["validate_publication_result()"]
    Native -->|"failure"| Poison["_poison_rollout()"]

    Cleanup["cleanup() / shutdown()"] --> Release["_release_resources()"]

Loading

vLLM-native adapter layer

VLLMNativeWeightSource

VLLMNativeWeightSource implements vLLM's WeightSource contract:

• metadata() returns ordered name/dtype/global-shape information without
  materializing FSDP full tensors;
• __iter__() invokes FullWeightSync._iter_full_tensors() lazily;
• actual_metadata records the exact stream consumed by WTE for post-transfer
  manifest verification.

The source is re-iterable through an iterator factory, as required by vLLM.

VLLMNativeWeightSyncClient

The native vLLM client protocol assumes four control operations:

init_weight_transfer_engine
start_weight_update
update_weights
finish_weight_update

VLLMNativeWeightSyncClient maps those operations onto the spawned
VLLMRolloutEngine.

Raw PyTorch IPC reducer objects cannot cross vLLM EngineCore's msgspec channel.
The client therefore follows vLLM's trusted local IPC-client pattern:

1. remove ipc_handles from the ordinary update payload;
2. pickle and base64-encode only the handle metadata;
3. send the encoded metadata through the local Pipe/runtime boundary;
4. let vLLM decode and consume it through its native WTE.

Tensor payload bytes never pass through Pipe, pickle, base64, or CPU memory.

VLLMNativeIPCTrainerEngine

This class subclasses vLLM's native IPCTrainerWeightTransferEngine. It reuses
the upstream:

• physical GPU UUID;
• handle all-gather and merge;
• _do_send() update format;
• worker-side native IPC receiver;
• packed consumer;
• layerwise load lifecycle.

UniRL overrides packed iteration for one reason: vLLM 0.27's generic packed
producer obtains the next tensor from the iterator before deciding to flush the
current buffer. With FSDP, obtaining the next tensor performs another
full_tensor() collective and can temporarily keep two complete tensors alive.

VLLMNativeIPCTrainerEngine._plan_chunks() uses cheap metadata first. The
runtime then materializes exactly the tensors belonging to one planned chunk,
copies them into the reusable buffer, sends the chunk, waits for rank consensus,
and only then consumes the next chunk.

flowchart LR
    Meta["WeightSource.metadata()<br/>no full_tensor"] --> Plan["_plan_chunks()"]
    Plan --> C1["chunk 1 names/shapes"]
    Plan --> C2["chunk 2 names/shapes"]
    Plan --> CN["chunk N names/shapes"]

    C1 --> M1["materialize only chunk 1"]
    M1 --> B["reusable packed CUDA buffer"]
    B --> S1["_do_send() + consensus"]
    S1 --> M2["materialize only chunk 2"]
    M2 --> B
    B --> S2["_do_send() + consensus"]
    S2 --> More["..."]

Loading

This is a bounded wrapper around vLLM native WTE, not a replacement data plane.

End-to-end synchronization sequence

sequenceDiagram
    autonumber
    participant Trainer as ARTrainer
    participant Backend as FSDPBackend
    participant Sync as IPCWeightSync
    participant Source as VLLMNativeWeightSource
    participant WTE as VLLMNativeIPCTrainerEngine
    participant Client as VLLMNativeWeightSyncClient
    participant Rollout as VLLMRolloutEngine
    participant Runtime as spawned vLLM runtime
    participant LLM as vllm.LLM native WTE
    participant TP as vLLM TP workers

    Trainer->>Backend: offload optimizer; clear consumed gradients
    Trainer->>Rollout: wake_up(tags=["weights"])
    Trainer->>Sync: sync()

    Sync->>Sync: _expected_tensor_metadata()
    Sync->>Sync: structural_manifest(expected)
    Sync->>Rollout: read TP fanout and worker CUDA UUIDs
    Sync->>Sync: _run_preflight()
    Note over Sync: Verify single node, FSDP4 == TP4,<br/>rank-aligned UUIDs, buffer size, free memory

    Sync->>Source: construct metadata + lazy iterator factory
    Sync->>Client: bind rollout, header, cache policy
    Sync->>WTE: construct packed native trainer engine

    WTE->>Client: init_weight_transfer_engine(packed=true)
    Client->>Rollout: init_native_weight_transfer()
    Rollout->>Runtime: Pipe init_native_weight_transfer
    Runtime->>LLM: init_weight_transfer_engine()

    WTE->>Client: start_weight_update()
    Client->>Rollout: start_native_weight_update(header)
    Rollout->>Runtime: Pipe start_native_weight_update
    Runtime->>LLM: start_weight_update()

    loop each metadata-planned chunk on every Actor rank
        WTE->>Source: consume next planned tensors
        Source->>Backend: _iter_full_tensors()
        Backend->>Backend: DTensor.full_tensor()
        Source-->>WTE: canonical CUDA tensors
        WTE->>WTE: validate order/shape/dtype
        WTE->>WTE: copy into reusable packed CUDA buffer
        WTE->>WTE: upstream all_gather_and_merge_handles()
        WTE->>Client: update_weights(merged UUID handles)
        Client->>Rollout: update_native_weights(encoded handle metadata)
        Rollout->>Runtime: Pipe update_native_weights
        Runtime->>LLM: update_weights()
        LLM->>TP: native IPC consumer selects same-GPU handle
        TP->>TP: layerwise load_weights / TP slice / fusion
        WTE->>Sync: _consensus(materialize and transfer)
    end

    WTE->>Sync: verify actual Actor manifests
    Sync->>Sync: all ranks must equal expected manifest

    WTE->>Client: finish_weight_update(model_version)
    Client->>Rollout: finish_native_weight_update()
    Rollout->>Runtime: Pipe finish_native_weight_update
    Runtime->>LLM: finish_weight_update()
    Runtime->>TP: collect publication receipts
    TP-->>Runtime: rank, world size, UUID, manifest, resident count
    Runtime->>Runtime: validate receipts
    Runtime->>LLM: update_weight_version()
    Runtime->>LLM: reset_prefix_cache()
    Runtime-->>Rollout: committed result
    Rollout->>Rollout: validate untrusted child result
    Rollout-->>Client: committed result

    WTE-->>Sync: finish consensus
    Sync->>Client: release_ipc_imports()
    Client->>Rollout: release_native_ipc()
    Rollout->>Runtime: Pipe release_native_ipc
    Runtime->>TP: ipc_collect()
    Sync->>Sync: validate_publication_result()

    Sync-->>Trainer: sync complete
    Trainer->>Backend: offload Actor model shards
    Trainer->>Rollout: wake_up() for KV cache
    Note over Trainer,TP: Next rollout observes the committed model version

Loading

Trainer memory schedule

ARTrainer queries IPCWeightSync.uses_gpu_streaming() once after constructing
the remote sync object. Existing sync implementations do not implement this
method and therefore keep their original path.

For direct vLLM native WTE:

flowchart LR
    Step["optimizer step complete"]
    OptOffload["offload optimizer<br/>clear gradients"]
    ActorResident["keep FSDP local parameter shards on GPU"]
    WeightWake["wake vLLM weights only"]
    Sync["native lazy IPC sync"]
    ActorOffload["offload Actor model"]
    KVWake["wake vLLM KV cache"]
    Generate["next rollout"]

    Step --> OptOffload --> ActorResident --> WeightWake --> Sync
    Sync --> ActorOffload --> KVWake --> Generate

Loading

The Actor parameter shards must remain resident until synchronization finishes,
because each lazy full_tensor() is an FSDP collective over those shards.
Optimizer state and consumed gradients are not needed by export and are removed
first to create GPU headroom.

Native GPU-streaming sync requires:

• enable_fsdp_offload=true;
• a rollout engine whose wake_up() supports weight-only tags;
• the SPMD one-Actor-rank-per-TP-rank layout;
• no rollout_anchor_device mode.

Publication identity and validation

Each publication header binds:

• protocol version;
• transport identifier;
• random publication ID;
• monotonic model version;
• TP world size;
• expected structural manifest;
• ordered Actor/worker physical CUDA UUIDs.

The structural manifest hashes ordered:

• tensor name;
• shape;
• dtype;
• element count;
• byte count.

It deliberately does not hash tensor payload bytes, avoiding a GPU-to-CPU copy
solely for checksums.

Validation occurs at multiple boundaries:

1. Actor ranks validate the expected schema before materialization.
2. Every emitted tensor is checked against its planned ParamMeta.
3. Every Actor rank's actual consumed manifest must match the expected
   manifest.
4. Each worker receipt must match publication ID, model version, TP world,
   expected fingerprint, TP rank, and CUDA UUID.
5. The spawned runtime validates receipts before publishing its local version.
6. The parent rollout validates the child response again across the process
   boundary.
7. IPCWeightSync broadcasts and validates the final result on every Actor
   rank.

Publication state and failure behavior

stateDiagram-v2
    [*] --> Idle
    Idle --> Initialized: init native WTE
    Initialized --> Started: start_weight_update
    Started --> Loading: first packed chunk
    Loading --> Loading: next chunk + rank consensus
    Loading --> Finalizing: source manifest verified
    Finalizing --> Committed: TP receipts + version + cache reset
    Committed --> Released: IPC imports collected
    Released --> Idle

    Initialized --> Poisoned: init/start failure
    Started --> Poisoned: materialization/transfer failure
    Loading --> Poisoned: schema/receiver/consensus failure
    Finalizing --> Poisoned: receipt/finalize/cache failure
    Committed --> Poisoned: release/final validation failure
    Poisoned --> BrokenRuntime
    BrokenRuntime --> [*]

Loading

There is no fallback transition from Poisoned to legacy ZMQ. A failed native
publication may have already mutated some resident storage, so reusing that
runtime would be unsafe.

Rank-consensus failure gate

FSDP materialization requires every Actor rank to enter collectives in exactly
the same order. A receiver failure on one rank must therefore be reported
before another rank advances to the next tensor.

flowchart TD
    Materialize["materialize planned chunk on all Actor ranks"]
    LocalResult["capture rank-local exception or success"]
    Gloo["Gloo all_gather_object"]
    AnyFailure{"any rank failed?"}
    Raise["raise the same phase/rank error everywhere"]
    Transfer["send merged IPC handles"]
    TransferConsensus["Gloo transfer consensus"]
    Next["consume next planned chunk"]

    Materialize --> LocalResult --> Gloo --> AnyFailure
    AnyFailure -->|"yes"| Raise
    AnyFailure -->|"no"| Transfer --> TransferConsensus
    TransferConsensus -->|"all succeeded"| Next
    TransferConsensus -->|"any failed"| Raise

Loading

The control group is Gloo rather than the FSDP NCCL group so error propagation
does not depend on entering the next weight-materialization collective.

Relationship to vLLM

This PR does not duplicate vLLM worker-side weight loading.

Reused directly from vLLM 0.27:

• WeightTransferConfig(backend="ipc");
• IPCTrainerWeightTransferEngine;
• IPC handle all-gather and UUID mapping;
• packed worker consumer;
• start_weight_update() / update_weights() /
  finish_weight_update();
• layerwise reload;
• model-native load_weights();
• TP slicing and model-specific fusion.

UniRL's VLLMNativeIPCTrainerEngine is a bounded trainer-side wrapper. Its
customization is limited to:

• metadata-first chunk planning;
• per-phase rank consensus;
• Actor manifest verification before finish.

References:

• vLLM Weight Transfer documentation
• vLLM FSDP/EP IPC RLHF example

Relationship to Verl

The shared design with Verl's colocated vLLM synchronization is:

• FSDP full tensors are materialized only as the iterator is consumed;
• every Actor rank participates in the same ordered tensor stream;
• Actor rank r and vLLM TP rank r exchange same-GPU IPC handles;
• the control channel carries handles and metadata;
• vLLM owns TP-aware loading;
• no additional Actor-to-vLLM NCCL broadcast is introduced.

UniRL-specific policy includes:

• a spawned synchronous runtime instead of Verl's Ray async-server adapter;
• an independent Qwen3-MoE expected schema;
• end-to-end physical CUDA UUID checks;
• a dedicated Gloo error-consensus group;
• model-version and receipt validation across the process boundary;
• fail-stop process-tree poisoning after partial failure.

Verl supports a broader set of update modes. This PR intentionally fails closed
outside its validated direct-vLLM Qwen3-MoE full-weight configuration.

File-level implementation map

Part I: direct rollout files

• unirl/rollout/engine/vllm/__init__.py

  • lazy exports for the direct vLLM engine.

• unirl/rollout/engine/vllm/config.py

  • typed vLLM configuration and command-specific timeouts.

• unirl/rollout/engine/vllm/engine.py

  • UniRL rollout API, spawned process ownership, sleep/wake state, native
    weight-update client methods, connection poisoning, and process-tree
    cleanup.

• unirl/rollout/engine/vllm/runtime.py

  • owns vllm.LLM, enables native IPC WTE, and serves the synchronous Pipe
    protocol.

• unirl/rollout/engine/vllm/worker_extension.py

  • exposes TP rank/world/UUID capabilities, publication receipts, and IPC
    import cleanup.

Part II: native weight-sync files

• unirl/distributed/weight_sync/transfer/vllm_native_engine.py

  • native vLLM source/client/trainer-engine adapters.

• unirl/distributed/weight_sync/transfer/vllm_native_protocol.py

  • header, manifest, receipt, and committed-result validation.

• unirl/distributed/weight_sync/full/ipc.py

  • preserves legacy ZMQ and adds the direct-vLLM native WTE branch.

Integration and reference files

• examples/ar/qwen3_moe_grpo_30b_a3b_gsm8k_fsdp_vllm_tp4.yaml

  • reference FSDP4/vLLM-TP4 recipe.

• unirl/distributed/group/handle.py

  • generalizes TP rollout placement so direct vLLM receives the same per-rank
    visible-device map already used for TP rollout engines.

• unirl/rollout/engine/base.py

  • exposes component_name() while retaining _component_name as the internal
    class field.

• unirl/train/backend/base_backend.py

  • allows optimizer/model/gradient state to be offloaded independently while
    preserving the old no-argument behavior.

• unirl/trainer/ar.py

  • adds the native GPU-streaming memory schedule without changing existing
    sync behavior when uses_gpu_streaming() is absent or false.

Compatibility and non-goals

Preserved behavior

• Existing SGLang synchronization is unchanged.
• Existing TensorWeightSync and NCCLWeightSync are unchanged.
• Existing vLLM-Omni bucketed ZMQ IPC is unchanged.
• Existing no-argument backend offload retains the original behavior.
• IPCWeightSync remains the public configuration target.

Non-goals

• generic model support beyond Qwen3-MoE;
• LoRA or merged-LoRA updates through native WTE;
• quantized or FP8/QAT weight publication;
• multi-node CUDA IPC;
• PP or EP rollout topologies;
• disaggregated trainer/inference hosts;
• atomic rollback of partially written vLLM weights;
• bitwise train-inference log-probability parity, which remains in PR feat(experimental): add bitwise Qwen3-MoE train–inference parity case based on GRPO + vLLM(TP>1) + FSDP #430.

Validation

Completed checks:

• Python compilation of the modified modules;
• Ruff lint and formatting checks;
• Hydra recipe resolution;
• native versus legacy engine selection;
• metadata-first chunk planning;
• native protocol and structural-manifest validation;
• TP capability and world-size validation;
• AR trainer native/legacy memory sequencing;
• absence of old custom DirectVLLM*, standalone NativeIPCWeightSync,
  _direct_send_large_weight, and send_bucket paths;
• confirmation that SGLang, TensorWeightSync, BucketedWeightSender, and
  vLLM-Omni receiver files match origin/main.

A two-rollout Qwen3-30B-A3B FSDP4/vLLM-TP4 run of the native WTE path completed
successfully with exit code 0 before the final naming and dead-code-only cleanup.
A final post-squash TP4 rerun is pending GPU availability.

Reviewer focus

Please focus review on:

1. whether Actor rank order exactly matches vLLM TP worker CUDA UUID order;
2. whether metadata-first planning preserves lazy, bounded FSDP materialization;
3. whether every materialization/transfer failure reaches rank consensus before
   the next FSDP collective;
4. whether vLLM remains the sole owner of TP slicing, fusion, and
   load_weights();
5. whether publication IDs, model versions, manifests, and receipts are checked
   at the correct process boundaries;
6. whether failure always poisons the runtime rather than continuing with
   partially updated weights;
7. whether IPC imports and the dedicated Gloo group are released on both normal
   and failure paths;
8. whether the native trainer memory estimate is conservative enough for the
   validated hardware.

Related Issue

Test Plan

Compatibility / Risk

Reviewer Notes

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.