[rl] Stateful Trainer Send: IPC [2/N]

[hao-aaron]

Context

Second of the planned three-PR split of the trainer-side weight-transfer rework.

• PR 1 (merged): introduced the new trainer-side abstractions
  (WeightSource / ModuleSource, VLLMWeightSyncClient,
  TrainerWeightTransferEngine, WeightTransferTrainerFactory). Purely
  additive; no backend migrated. [rl] Stateful Trainer Send: New Abstractions [1/N]  #48042
• PR 2 (this one): migrate the IPC backend onto those abstractions
  end-to-end — trainer engine + wire params on the init info + update-info
  slimming + IPC examples/tests. First PR with an end-to-end correctness test of
  the new trainer path.
• PR 3: migrate NCCL + sparse NCCL
• PR 4: remove trainer_send_weights
  from the worker ABC, and port the remaining examples/tests/docs.

NCCL and sparse NCCL are untouched here and keep their existing static trainer
path; the worker ABC still declares trainer_send_weights (removed in PR 3). Docs remain unchanged and are deferred to PR 4.

What changed

IPC engine (ipc_engine.py)

• IPCTrainerWeightTransferEngine — stateful trainer engine driven by a
  parameter-free send_weights(). Straight-line (no concurrent broadcast like
  NCCL): update_weights is the transfer and rides the client. For multi-rank
  trainers all ranks iterate the source and join the IPC-handle all-gather; only
  the sender (rank 0) ships the merged handles and drives the inference-side RPCs,
  with a per-chunk barrier in packed mode so producers don't overwrite a buffer a
  consumer is still reading. It reads packed / packed_buffer_size_bytes from
  its IPCTrainerInitInfo.
• IPCWeightTransferUpdateInfo slimmed: the packed field is gone; per-round
  payloads carry only per-round metadata + IPC handles.
• The worker IPCWeightTransferEngine reads packed from self.packed (set at
  the init handshake), not from a per-round field or the config.
• The old static IPCWeightTransferEngine.trainer_send_weights and
  IPCTrainerSendWeightsArgs are removed; a transitional stub remains on the
  worker engine (raises NotImplementedError) only to satisfy the still-abstract
  worker ABC member, which PR 3 deletes.

Factory (factory.py)

• WeightTransferTrainerFactory registers the ipc trainer engine (lazy import).
  NCCL / sparse register in PR 3.
• changed trainer side factory to not accept backend str parameter, instead read from the init_info, enforce all subclasses of TrainerInitInfo to define backend str.

Examples & tests

• Ported examples/rl/rlhf_ipc.py, rlhf_http_ipc.py, rlhf_ipc_fsdp_ep.py to
  WeightTransferTrainerFactory.trainer_init(...).send_weights() with a
  VLLMWeightSyncClient (Ray / HTTP) and ModuleSource. The manual
  init/start/finish juggling collapses to a single send_weights() call; the
  inference side takes a plain WeightTransferConfig(backend="ipc") and packing
  is set on IPCTrainerInitInfo.
• tests/distributed/test_weight_transfer.py: IPC worker tests use a
  string-only config and learn packed from the init handshake; the
  trainer-factory registry test asserts ipc is registered; added
  test_ipc_trainer_send_weights_drives_client_in_order (start→update→finish
  ordering, packed off the per-round update-info) and
  test_ipc_trainer_init_ships_packed_to_worker (asserts the init handshake
  carries {"packed": ...}).
• Docs deferred: the docs/training/weight_transfer/* rewrite lands with the
  NCCL migration (PR 3), when the trainer-side story is described once for all
  backends.

Breaking changes (IPC only)

• Removed static IPC trainer API: IPCWeightTransferEngine.trainer_send_weights
  / IPCTrainerSendWeightsArgs. Use
  WeightTransferTrainerFactory.trainer_init(backend="ipc", ...).send_weights().
• Weights are supplied as a WeightSource (e.g. ModuleSource(model)).
• IPC packed / packed_buffer_size_bytes move onto IPCTrainerInitInfo and
  are no longer per-round update_info fields; the trainer propagates packed to
  the worker at init.
• NCCL / sparse NCCL trainer APIs are unchanged.

Migration (IPC)

Before:

trainer_args = IPCTrainerSendWeightsArgs(send_mode="ray", llm_handle=llm, packed=False)
ray.get(llm.start_weight_update.remote())
IPCWeightTransferEngine.trainer_send_weights(
    iterator=model.named_parameters(), trainer_args=trainer_args)
ray.get(llm.finish_weight_update.remote())

After:

engine = WeightTransferTrainerFactory.trainer_init(
    backend="ipc",
    config=WeightTransferConfig(backend="ipc"),
    init_info=IPCTrainerInitInfo(rank=0, packed=False),  # packed lives here
    client=RayVLLMWeightSyncClient(llm),   # or HTTPVLLMWeightSyncClient(url)
    source=ModuleSource(model),
)
engine.send_weights()

The inference side is constructed with the same plain config
(LLM(..., weight_transfer_config=WeightTransferConfig(backend="ipc")) or
vllm serve ... --weight-transfer-config '{"backend": "ipc"}'); it learns
packed from the trainer at the init handshake.

Testing

pytest tests/distributed/test_weight_transfer.py
# 58 passed  (8xH100)

Covers IPC update-info validation, IPC handle (de)serialization, the worker
learning packed from the init handshake, and the new IPC trainer engine
(send_weights ordering + trainer_init propagating packed), plus
multi-process IPC integration (Ray + HTTP).

All three examples were run end-to-end on 8xH100 (facebook/opt-125m for
rlhf_ipc.py / rlhf_http_ipc.py; Qwen3-30B-A3B FSDP2 + expert-parallel,
packed, for rlhf_ipc_fsdp_ep.py) — each produces garbage from the
dummy-initialized server before the sync and coherent text after.

[mergify]

Documentation preview: https://vllm--48981.org.readthedocs.build/en/48981/

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[aoshen02]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7181dd4a10

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Merge handles on the designated sender

When IPCTrainerInitInfo.rank differs from the default process-group rank—for example, rank 0 is local to a trainer subgroup—the designated sender reaches _do_send, but _all_gather_and_merge_handles returns merged handles only on torch.distributed.get_rank() == 0. The explicit sender therefore receives empty handle dictionaries while the global rank-0 process does not send, causing workers to fail to find their GPU UUID. The merge destination must follow self.is_sender/the supplied trainer rank, or the API must use and document the corresponding process group.

Useful? React with 👍 / 👎.

[aoshen02]

Find a codex comment make sense to me, could you check:
Unpacked multi-rank mode: non-sender ranks can free their gathered tensors while colocated workers are still reading them via IPC

The old static path ran _post_send_sync() inside _send_unpacked, while weight_refs was still in scope:

before

weight_refs.append(weight)
...
ipc_handles = IPCWeightTransferEngine._all_gather_and_merge_handles(ipc_handles)
if IPCWeightTransferEngine._is_rank_zero():
IPCWeightTransferEngine._do_send(...)
IPCWeightTransferEngine._post_send_sync() # barrier while weight_refs is alive
so every rank held its tensors until the barrier released, and the barrier only released after rank 0's synchronous update_weights returned — i.e. after all workers had finished reading.

The new send_weights() moves the barrier outside the scope that owns the tensors:

after

self._send(source) # _send_unpacked returns -> weight_refs dropped
if self.is_sender:
self.client.finish_weight_update()
self._post_send_sync() # barrier is now after the free
Failure sequence on a multi-rank FSDP trainer (unpacked mode):

Under FSDP, materialize_full_tensor calls full_tensor(), so each yielded tensor is a temporary whose only reference is weight_refs.
On a non-sender rank, _do_send early-returns, _send_unpacked returns, and weight_refs is dropped — the gathered full tensors go back to the caching allocator.
Meanwhile the sender is still inside update_weights, and the vLLM worker colocated with the non-sender GPU is reading exactly that memory through the IPC handle. Any allocation on the non-sender rank in that window can reuse the block → silently corrupted weights, no error.
This is the same hazard the per-chunk barrier comment in _send_packed describes ("non-sender ranks race ahead and overwrite their buffer while their colocated worker is still reading") — the packed path keeps its barrier inside the loop while the producer still holds the buffer, but the unpacked path lost that property in the refactor.

None of the current examples hit it (the two single-process examples send live module params, and the FSDP example uses packed), but the engine docstring explicitly supports multi-rank + unpacked, so it's reachable through the public API.

Suggested fix: move _post_send_sync() back inside _send_unpacked (within weight_refs' scope), and drop it from the tail of send_weights():

def _send_unpacked(self, source: WeightSource) -> None:
weight_refs: list[torch.Tensor] = []
...
ipc_handles = self._all_gather_and_merge_handles(ipc_handles)
self._do_send(...)
self._post_send_sync() # weight_refs still alive across the barrier
The packed path is unaffected (its per-chunk barrier is already correctly placed).

[aoshen02]

Maybe we should add backward compatibility?

[aoshen02]

A bit verbose