Skip to content

[rl] Stateful Trainer Send: New Abstractions [1/N] - #48042

Merged
Isotr0py merged 7 commits into
vllm-project:mainfrom
hao-aaron:trainer-send-pr1-abstractions
Jul 17, 2026
Merged

[rl] Stateful Trainer Send: New Abstractions [1/N] #48042
Isotr0py merged 7 commits into
vllm-project:mainfrom
hao-aaron:trainer-send-pr1-abstractions

Conversation

@hao-aaron

Copy link
Copy Markdown
Contributor

Context

This is the first of a three-PR split of the trainer-side weight-transfer rework
originally proposed as one large PR. The split is:

  • PR 1 (this one): introduce the new trainer-side abstractions
    (WeightSource / ModuleSource, VLLMWeightSyncClient,
    TrainerWeightTransferEngine, WeightTransferTrainerFactory). Purely
    additive — the existing worker-side WeightTransferEngine and every backend
    engine (NCCL / IPC / sparse NCCL) are untouched and remain valid as-is.
  • PR 2: migrate the IPC backend onto the new abstractions end-to-end
    (trainer engine + config subclass + update-info slimming + worker
    read-from-config + IPC examples/tests/docs). First PR with an end-to-end
    correctness test of the new trainer path.
  • PR 3: migrate NCCL + sparse NCCL, remove trainer_send_weights
    from the worker ABC, and port the remaining examples/tests/docs.

Splitting this way keeps each backend's coupled change (the config carries the
static "must-agree" wire params, which the worker now reads instead of the
per-round update-info) together with its own examples and tests, so PR 2 and
PR 3 are each independently reviewable and independently end-to-end testable.

Motivation

Design doc

The old WeightTransferEngine ABC treated the trainer side as stateless — a
pair of @staticmethods (trainer_init, trainer_send_weights) plus a
per-backend *TrainerSendWeightsArgs dataclass the caller had to thread back in
every round. That put transfer state in the caller, defeated the type system
with trainer_args: dict | Any, and made every example open-code the backend
concurrency by hand.

This PR lays the groundwork to make the trainer side symmetric to the worker
side — a stateful engine that owns its state, pulls weights from a
WeightSource, and drives the handshake through a transport-agnostic client —
without yet migrating any backend. Nothing in this PR changes existing
behavior; it only adds new, unused-until-PR-2 abstractions plus their unit
tests.

What changed

All changes are additive. No existing class, method, or signature is removed or
behaviorally modified.

New trainer-side engine ABC (base.py)

  • TrainerWeightTransferEngine — stateful ABC, generic over
    (config, init_info), built via a trainer_init classmethod factory and
    driven by a parameter-free send_weights(). Optional shutdown().
  • TrainerInitInfo — base trainer init info carrying an explicit rank
    with an is_sender property; rank 0 is always the sender. Rank is passed
    explicitly (not read from a global process group) because that is ambiguous
    once several groups — FSDP / TP / PP / EP — exist.

Weight source (base.py)

  • WeightSource ABC — a re-iterable source of the trainer's weights with
    two channels:
    • metadata() -> list[ParamMeta](name, wire dtype, full shape) for every
      param without transferring (cheap when shapes are known locally; may
      cache for producers that must materialize to learn shapes).
    • iteration — yields fully-materialized (name, tensor) pairs one at a time;
      every trainer rank must iterate the same source in lockstep (materializing
      is often a collective).
  • ModuleSource(module) — the common case over module.named_parameters();
    handles plain and FSDP-sharded modules with no special casing.
  • ParamMeta (frozen name/dtype/shape) and
    materialize_full_tensor() helper — gathers FSDP DTensor shards via
    full_tensor() at send time (once), while metadata() reads global
    shape/dtype without gathering. Custom producers (Megatron export, MoE
    re-fusing) subclass WeightSource.

Control plane (base.py, clients.py)

  • VLLMWeightSyncClient — a @runtime_checkable structural Protocol with
    four synchronous methods (init_weight_transfer_engine, start_weight_update,
    update_weights, finish_weight_update). Any object with those methods works;
    no import/subclassing required. Backend-specific concurrency lives in the
    engine, not the client.
  • Built-ins HTTPVLLMWeightSyncClient (RLHF HTTP routes) and
    RayVLLMWeightSyncClient (fans out to one or more AsyncLLM/LLM Ray
    actors).
  • HTTP transport can't carry raw CUDA IPC handles, so the client pickles +
    base64-encodes them into ipc_handles_pickled (_json_safe_update_info);
    the worker deserializes gated on VLLM_ALLOW_INSECURE_SERIALIZATION=1.
    clients.py is added to the pickle-import allowlist in
    tools/pre_commit/check_forbidden_imports.py.

Factory (factory.py)

  • WeightTransferTrainerFactory — lazy-import registry parallel to
    WeightTransferEngineFactory (separate registry; trainer and worker never
    instantiate each other's engines). Its registry is intentionally empty in
    this PR
    — the nccl/ipc trainer engines register in PR 2 / PR 3 alongside
    the concrete classes, so the registry never points at classes that don't yet
    exist.

Misc

  • AsyncLLM.init_weight_transfer_engine / update_weights now accept
    dict | Request (backward-compatible widening; needed by the Ray client
    path).

Explicitly not in this PR

Deferred to the per-backend PRs because they are coupled trainer+worker changes:

  • The concrete NCCLTrainerWeightTransferEngine / IPCTrainerWeightTransferEngine.
  • NCCLWeightTransferConfig / IPCWeightTransferConfig and the update-info
    slimming (moving packed / buffer sizes off the per-round update-info onto the
    config, which the worker then reads).
  • EngineArgs dispatch of dict/CLI config to the right subclass.
  • Removing trainer_send_weights from the worker ABC (PR 3).
  • All examples/rl/*, integration tests, and doc rewrites.

Breaking changes

None. This PR only adds new symbols.

Testing

pytest tests/distributed/test_weight_transfer.py

Adds GPU-free unit tests for the new abstractions: ModuleSource
metadata/iteration/re-iterability, VLLMWeightSyncClient structural
conformance (recording / HTTP / Ray), HTTP client JSON-safety (IPC handles
pickled + base64-encoded, NCCL metadata passed through unchanged),
WeightTransferTrainerFactory register/dispatch/error paths, and
TrainerWeightTransferEngine base construction. The end-to-end trainer
transfer tests land with their backend in PR 2 / PR 3.

Duplication check

Not a duplicate of an existing open PR — no other open PR reworks the
trainer-side weight-transfer ABC.

This PR was written with AI assistance. Every line has been human-reviewed.

Signed-off-by: haoaaron <ahao@anyscale.com>
@hao-aaron
hao-aaron requested review from hmellor and njhill as code owners July 8, 2026 20:23

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@mergify mergify Bot added the v1 label Jul 8, 2026
@kouroshHakha kouroshHakha added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 8, 2026
@mergify

mergify Bot commented Jul 12, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hao-aaron.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Jul 12, 2026
Comment thread tests/distributed/test_weight_transfer.py Outdated
Comment thread vllm/v1/engine/async_llm.py Outdated

async def init_weight_transfer_engine(
self, request: WeightTransferInitRequest
self, request: WeightTransferInitRequest | dict

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why did you have to change the type here?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

dict was passed by the new ray client, but i changed it to create the typed payload

Comment thread vllm/distributed/weight_transfer/factory.py Outdated

@SumanthRH SumanthRH left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks mostly good! Left some nits!

hao-aaron and others added 4 commits July 13, 2026 11:48
Co-authored-by: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com>
Signed-off-by: Aaron Hao <ahao@anyscale.com>
Co-authored-by: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com>
Signed-off-by: Aaron Hao <ahao@anyscale.com>
Signed-off-by: haoaaron <ahao@anyscale.com>
@mergify mergify Bot removed the needs-rebase label Jul 13, 2026
is_sender: bool = True,
) -> None:
self.config = config
self.is_sender = is_sender

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here may violate the rule of Single Source of Truth, we can just get the is_sender value directly from TrainerInitInfo

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall LGTM

@hao-aaron hao-aaron Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would need to pass init_info in to fix, i think its fine as is for now

@kwen2501

kwen2501 commented Jul 17, 2026

Copy link
Copy Markdown

Does PR #48042 make M2N integration easier?

The new TrainerWeightTransferEngine with a stateful send_weights() is a better fit for nccl-m2n's one-time init (ncclM2nInit) + per-round reshard pattern. So yes, it helps the trainer side. Also, WeightTransferTrainerFactory.register("nccl-m2n", ...) gives a clean registration point.

What is still needed for integration?

ParamMeta in the new abstraction only carries (name, dtype, shape). TrainerInitInfo only carries rank. Neither carries the mesh topology (e.g., EP=16/DP=2/PP=4 on the trainer side, TP=8/EP=1/DP=16/PP=1 on the generator side) that ncclM2nInit needs to set up the cross-mesh reshard.

What we can do:

  • An M2NTrainerInitInfo subclass of TrainerInitInfo that carries trainer and generator mesh topology, and/or
  • An M2NParamMeta subclass that extends ParamMeta with per-param shard descriptors.

@aoshen02

Copy link
Copy Markdown
Collaborator

Does PR #48042 make M2N integration easier?

The new TrainerWeightTransferEngine with a stateful send_weights() is a better fit for nccl-m2n's one-time init (ncclM2nInit) + per-round reshard pattern. So yes, it helps the trainer side. Also, WeightTransferTrainerFactory.register("nccl-m2n", ...) gives a clean registration point.

What is still needed for integration?

ParamMeta in the new abstraction only carries (name, dtype, shape). TrainerInitInfo only carries rank. Neither carries the mesh topology (e.g., EP=16/DP=2/PP=4 on the trainer side, TP=8/EP=1/DP=16/PP=1 on the generator side) that ncclM2nInit needs to set up the cross-mesh reshard.

What we can do:

  • An M2NTrainerInitInfo subclass of TrainerInitInfo that carries trainer and generator mesh topology, and/or
  • An M2NParamMeta subclass that extends ParamMeta with per-param shard descriptors.

Merge it first, aaron please take a look. @hao-aaron

@Isotr0py
Isotr0py merged commit fb1d8cc into vllm-project:main Jul 17, 2026
93 checks passed
plasticchris pushed a commit to plasticchris/vllm that referenced this pull request Jul 20, 2026
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: Aaron Hao <ahao@anyscale.com>
Co-authored-by: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com>
edwinlim0919 pushed a commit to chaeminlim-mb/vllm that referenced this pull request Jul 29, 2026
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: Aaron Hao <ahao@anyscale.com>
Co-authored-by: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com>
zhangxinyuehfad added a commit to zhangxinyuehfad/vllm-ascend that referenced this pull request Jul 31, 2026
… vLLM version compat

vLLM main (post-0.26.0) removed IPCTrainerSendWeightsArgs and the static
trainer_send_weights path via:

  vllm-project/vllm#48042 — Stateful Trainer Send: New Abstractions [1/N]
  vllm-project/vllm#48981 — Stateful Trainer Send: IPC [2/N]

Add a vllm_version_is conditional branch:

- 0.26.0: preserves the existing static NPUIPCTrainerSendWeightsArgs +
  NPUIPCWeightTransferEngine.trainer_send_weights path unchanged.
- main:   introduces NPUIPCTrainerInitInfo (backend='npu_ipc'),
  NPUIPCTrainerWeightTransferEngine (subclass of upstream's
  IPCTrainerWeightTransferEngine), and delegates HTTP transport to
  HTTPVLLMWeightSyncClient.  Register the trainer-side engine via
  WeightTransferTrainerFactory only on main.

Signed-off-by: hfadzxy <starmoon_zhang@163.com>
zhangxinyuehfad added a commit to zhangxinyuehfad/vllm-ascend that referenced this pull request Aug 1, 2026
… vLLM version compat

vLLM main (post-0.26.0) removed IPCTrainerSendWeightsArgs and the static
trainer_send_weights path via:

  vllm-project/vllm#48042 — Stateful Trainer Send: New Abstractions [1/N]
  vllm-project/vllm#48981 — Stateful Trainer Send: IPC [2/N]

Add a vllm_version_is conditional branch:

- 0.26.0: preserves the existing static NPUIPCTrainerSendWeightsArgs +
  NPUIPCWeightTransferEngine.trainer_send_weights path unchanged.
- main:   introduces NPUIPCTrainerInitInfo (backend='npu_ipc'),
  NPUIPCTrainerWeightTransferEngine (subclass of upstream's
  IPCTrainerWeightTransferEngine), and delegates HTTP transport to
  HTTPVLLMWeightSyncClient.  Register the trainer-side engine via
  WeightTransferTrainerFactory only on main.

Signed-off-by: hfadzxy <starmoon_zhang@163.com>
zhangxinyuehfad added a commit to zhangxinyuehfad/vllm-ascend that referenced this pull request Aug 1, 2026
… vLLM version compat

vLLM main (post-0.26.0) removed IPCTrainerSendWeightsArgs and the static
trainer_send_weights path via:

  vllm-project/vllm#48042 — Stateful Trainer Send: New Abstractions [1/N]
  vllm-project/vllm#48981 — Stateful Trainer Send: IPC [2/N]

Add a vllm_version_is conditional branch:

- 0.26.0: preserves the existing static NPUIPCTrainerSendWeightsArgs +
  NPUIPCWeightTransferEngine.trainer_send_weights path unchanged.
- main:   introduces NPUIPCTrainerInitInfo (backend='npu_ipc'),
  NPUIPCTrainerWeightTransferEngine (subclass of upstream's
  IPCTrainerWeightTransferEngine), and delegates HTTP transport to
  HTTPVLLMWeightSyncClient.  Register the trainer-side engine via
  WeightTransferTrainerFactory only on main.

Signed-off-by: hfadzxy <starmoon_zhang@163.com>
zhangxinyuehfad added a commit to zhangxinyuehfad/vllm-ascend that referenced this pull request Aug 1, 2026
… vLLM version compat

vLLM main (post-0.26.0) removed IPCTrainerSendWeightsArgs and the static
trainer_send_weights path via:

  vllm-project/vllm#48042 — Stateful Trainer Send: New Abstractions [1/N]
  vllm-project/vllm#48981 — Stateful Trainer Send: IPC [2/N]

Add a vllm_version_is conditional branch:

- 0.26.0: preserves the existing static NPUIPCTrainerSendWeightsArgs +
  NPUIPCWeightTransferEngine.trainer_send_weights path unchanged.
- main:   introduces NPUIPCTrainerInitInfo (backend='npu_ipc'),
  NPUIPCTrainerWeightTransferEngine (subclass of upstream's
  IPCTrainerWeightTransferEngine), and delegates HTTP transport to
  HTTPVLLMWeightSyncClient.  Register the trainer-side engine via
  WeightTransferTrainerFactory only on main.

Signed-off-by: hfadzxy <starmoon_zhang@163.com>
zhangxinyuehfad added a commit to zhangxinyuehfad/vllm-ascend that referenced this pull request Aug 3, 2026
… vLLM version compat

vLLM main (post-0.26.0) removed IPCTrainerSendWeightsArgs and the static
trainer_send_weights path via:

  vllm-project/vllm#48042 — Stateful Trainer Send: New Abstractions [1/N]
  vllm-project/vllm#48981 — Stateful Trainer Send: IPC [2/N]

Add a vllm_version_is conditional branch:

- 0.26.0: preserves the existing static NPUIPCTrainerSendWeightsArgs +
  NPUIPCWeightTransferEngine.trainer_send_weights path unchanged.
- main:   introduces NPUIPCTrainerInitInfo (backend='npu_ipc'),
  NPUIPCTrainerWeightTransferEngine (subclass of upstream's
  IPCTrainerWeightTransferEngine), and delegates HTTP transport to
  HTTPVLLMWeightSyncClient.  Register the trainer-side engine via
  WeightTransferTrainerFactory only on main.

Signed-off-by: hfadzxy <starmoon_zhang@163.com>
zhangxinyuehfad added a commit to zhangxinyuehfad/vllm-ascend that referenced this pull request Aug 3, 2026
… vLLM version compat

vLLM main (post-0.26.0) removed IPCTrainerSendWeightsArgs and the static
trainer_send_weights path via:

  vllm-project/vllm#48042 — Stateful Trainer Send: New Abstractions [1/N]
  vllm-project/vllm#48981 — Stateful Trainer Send: IPC [2/N]

Add a vllm_version_is conditional branch:

- 0.26.0: preserves the existing static NPUIPCTrainerSendWeightsArgs +
  NPUIPCWeightTransferEngine.trainer_send_weights path unchanged.
- main:   introduces NPUIPCTrainerInitInfo (backend='npu_ipc'),
  NPUIPCTrainerWeightTransferEngine (subclass of upstream's
  IPCTrainerWeightTransferEngine), and delegates HTTP transport to
  HTTPVLLMWeightSyncClient.  Register the trainer-side engine via
  WeightTransferTrainerFactory only on main.

Signed-off-by: hfadzxy <starmoon_zhang@163.com>
linfeng-yuan pushed a commit to vllm-project/vllm-ascend that referenced this pull request Aug 3, 2026
### What this PR does / why we need it?

| Files | Upstream vLLM change | vllm-ascend adaptation |
|-------|---------------------|------------------------|
| `vllm_ascend/__init__.py` |
[vllm#48841](vllm-project/vllm#48841) — added
`from triton.experimental import gluon`, `from triton.experimental.gluon
import language as gl` and `from triton.language.core import _aggregate`
to `triton_utils/__init__.py`, requiring Triton 3.6+ API absent from
`triton-ascend 3.2.1` | Pre-register `triton.experimental.gluon` /
`.language` as `sys.modules` stubs; stub
`triton.language.core._aggregate`. Gated on `os.getenv("VLLM_VERSION")
!= "0.26.0"`, mirroring `vllm_version_is` env var path. |
| `tests/ut/patch/platform/test_patch_structured_output.py` |
[vllm#49665](vllm-project/vllm#49665) — changed
`VLLMValidationError` base class from `ValueError` to `VLLMClientError`,
breaking `pytest.raises(ValueError)` assertions | Updated 3 assertions
from `pytest.raises(ValueError, ...)` to
`pytest.raises(VLLMValidationError, ...)` |
| `tests/e2e/pull_request/one_card/test_guided_decoding.py` |
[vllm#49665](vllm-project/vllm#49665) — same
`VLLMValidationError` base-class change; on 0.26.0 the upstream
validation may still raise `ValueError` | Version-gated assertion:
`pytest.raises(ValueError, ...)` on 0.26.0 vs
`pytest.raises(VLLMValidationError, ...)` on main |
|
`vllm_ascend/ops/vocab_parallel_embedding.py`<br>`vllm_ascend/_310p/ops/vocab_parallel_embedding.py`
| [vllm#49731](vllm-project/vllm#49731) — added
`*, disable_tp: bool = False` keyword arg to `ParallelLMHead.__init__()`
and `VocabParallelEmbedding.__init__()` | Version-gated `__init__` with
`vllm_version_is("0.26.0")`: else branch accepts `disable_tp` kwarg and
forwards it (Ascend manages TP via `lmhead_tp_enable()`) |
|
`vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py`<br>`vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py`
| [vllm#26593](vllm-project/vllm#26593) —
`KVConnectorBase_V1.__init__()` sets `self._kv_transfer_config`,
required by `requires_kv_delivery` property | Version-gated `__init__`:
else branch sets `self._kv_transfer_config` on vllm main |
| `vllm_ascend/ops/fused_moe/fused_moe.py` |
[vllm#50089](vllm-project/vllm#50089) — added
`fused_output_is_reduced` kwarg to
`_maybe_reduce_shared_expert_output()` and `output_is_reduced` kwarg to
`_maybe_reduce_final_output()` | Version-gated both methods: else branch
accepts new kwargs but ignores them (Ascend handles reduction
independently); `trunc_size` may now be `None` |
| `vllm_ascend/patch/worker/patch_distributed.py` |
[vllm#47288](vllm-project/vllm#47288) — added
`use_all2all: bool = False` param to `GroupCoordinator.__init__()` |
Version-gated via `_IS_VLLM_026` env var check (avoids vllm-ascend
import cycle). Else branch stores `self.use_all2all` (unused, no all2all
on Ascend) |
| `vllm_ascend/distributed/device_communicators/npu_communicator.py` |
[vllm#47288](vllm-project/vllm#47288) — added
`use_all2all: bool = False` param to `DeviceCommunicatorBase.__init__()`
| Version-gated `NPUCommunicator.__init__` with
`vllm_version_is("0.26.0")`: else branch accepts and forwards
`use_all2all` (NPU keeps the no-op `_NpuAll2AllManager`) |
| `vllm_ascend/worker/npu_input_batch.py` |
[vllm#48018](vllm-project/vllm#48018) — added
`use_replayssm: bool` kwarg (Mamba replay-SSM);
[vllm#40996](vllm-project/vllm#40996) — added
`slot_mapping_modes` kwarg (DCP hybrid attention) | Accepts both kwargs
unconditionally; stores `self.use_replayssm` and
`self.slot_mapping_modes` only on main (`not vllm_version_is("0.26.0")`)
for interface alignment |
| `vllm_ascend/worker/v2/spec_decode/eagle/aclgraph.py` |
[vllm#49364](vllm-project/vllm#49364) — renamed
`skip_attn` parameter to `full_cudagraph` in cudagraph mode logic |
Version-gated call: 0.26.0 keeps `skip_attn=(cg_mode != PIECEWISE)`;
main passes `full_cudagraph=(cg_mode != PIECEWISE)` |
|
`vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py`<br>`vllm_ascend/distributed/weight_transfer/__init__.py`
| [vllm#48042](vllm-project/vllm#48042) +
[vllm#48981](vllm-project/vllm#48981) — replaced
static `IPCTrainerSendWeightsArgs`/`trainer_send_weights` with a
stateful `IPCTrainerWeightTransferEngine` driven by
`WeightTransferTrainerFactory.trainer_init(...).send_weights()`;
`packed` moved onto trainer init-info; per-round `update_info` slimmed |
Whole module version-gated: 0.26.0 keeps the static
`NPUIPCWeightTransferEngine`; main defines stateful
`NPUIPCTrainerWeightTransferEngine` (`_send_unpacked` instance method,
`packed` on `NPUIPCWeightTransferInitInfo`, no-arg
`npu_generate_uuid()`, `is_sender`/`_all_gather_and_merge_handles`).
`register_engine()` also registers the trainer engine in
`WeightTransferTrainerFactory` on main only |
| `examples/rl/rlhf_http_npu_ipc.py` |
[vllm#48981](vllm-project/vllm#48981) — ported
RL examples to `WeightTransferTrainerFactory.trainer_init(...)` +
`engine.send_weights()` with `HTTPVLLMWeightSyncClient`/`ModuleSource` |
Version-gated example: 0.26.0 uses static
`NPUIPCWeightTransferEngine.trainer_send_weights`; main uses the
stateful trainer engine path |
| `tests/e2e/pull_request/one_card/test_npu_ipc_weight_transfer.py` |
[vllm#48981](vllm-project/vllm#48981) — IPC
weight-transfer E2E now exercised via the stateful trainer engine |
Version-gated test: 0.26.0 keeps the manual `_post(start/finish)` +
static `trainer_send_weights`; main drives
`WeightTransferTrainerFactory.trainer_init(...).send_weights()` (engine
owns the lifecycle) |
| `tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py` |
[vllm#48981](vllm-project/vllm#48981) —
trainer-side IPC APIs changed from static methods to instance methods on
the stateful engine | Version-gated tests (`IS_VLLM_026`):
`_send_unpacked`, `packed` placement, `npu_generate_uuid` arity, and
`start/finish_weight_update` no-op behavior asserted per version |
| `vllm_ascend/worker/v2/spec_decode/dflash/speculator.py` |
[vllm#50000](vllm-project/vllm#50000) —
`_prepare_dflash_inputs_kernel` gained `temperature`/`seeds` sampling
params for probabilistic draft sampling | Version-gated
`_prepare_dflash_inputs_kernel_ascend`: 0.26.0 branch keeps the old
signature; main branch adds the four `temperature`/`seeds` pointers and
the corresponding stores to stay aligned with upstream while keeping
Ascend's own kernel |
| `.github/vllm-main-verified.commit` | — | Updated verified main commit
hash to `0351e9aa1fdf1a51329d1906881528dfe61fc88e` |

### Does this PR introduce _any_ user-facing change?

### How was this patch tested?


- vLLM version: v0.26.0
- vLLM main:
vllm-project/vllm@d02df74

---------

Signed-off-by: hfadzxy <starmoon_zhang@163.com>
hao-aaron added a commit to hao-aaron/vllm that referenced this pull request Aug 3, 2026
Migrate the dense NCCL and sparse NCCL backends onto the stateful trainer
engine, completing the trainer-side weight-transfer rework started in vllm-project#48042
(abstractions) and vllm-project#48981 (IPC).

- `NCCLTrainerWeightTransferEngine`: rank 0 holds the `PyNcclCommunicator` and
  owns the concurrency, running the inference-side `update_weights` on a side
  thread while it broadcasts (both rendezvous inside the same NCCL calls), with
  a `future.done()` early-error check so a request rejected before any NCCL call
  surfaces instead of hanging the broadcast.
- Wire params ride the init info, mirroring IPC: `NCCLTrainerInitInfo` carries
  the rendezvous fields plus the must-agree `packed` / `packed_buffer_size_bytes`
  / `packed_num_buffers`, and the sender ships them to the worker at
  `trainer_init`, so the two sides cannot disagree. `NCCLWeightTransferUpdateInfo`
  is slimmed to per-round `names` / `dtype_names` / `shapes`, and the worker reads
  `self.packed` recorded at the init handshake.
- `SparseNCCLTrainerWeightTransferEngine` is modeled as a delta backend: sparse
  patches differ every round, so they are not a stable `WeightSource`. The engine
  takes no `source`; each round's patches go to `send_weights(patches)`.
  `SparseWeightPatch` gains `full_shape`, required on this path.
- `source` is now optional on `TrainerWeightTransferEngine` / the factory to
  support delta backends; full-resync backends (NCCL, IPC) reject `None`
  themselves.
- Multi-rank trainers: every rank builds the engine and calls `send_weights()`.
  Non-senders hold no communicator and skip the client RPCs and the broadcast,
  but still iterate the `WeightSource` and run `metadata()` to stay in the
  trainer-side collectives (e.g. the FSDP `full_tensor()` all-gather).
- Removes the static NCCL trainer API (`trainer_send_weights`,
  `NCCLTrainerSendWeightsArgs`, and the `NCCLWeightTransferEngine.trainer_init`
  re-export); ports the NCCL / HTTP / FSDP-EP / sparse examples to
  `WeightTransferTrainerFactory.trainer_init(...).send_weights()`.

Signed-off-by: haoaaron <ahao@anyscale.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
aditi-amd pushed a commit to aditi-amd/vllm that referenced this pull request Aug 4, 2026
Signed-off-by: haoaaron <ahao@anyscale.com>
Signed-off-by: Aaron Hao <ahao@anyscale.com>
Co-authored-by: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com>
Signed-off-by: root <root@smci355-ccs-aus-m02-09.cs-aus.dcgpu>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready ONLY add when PR is ready to merge/full CI is needed v1

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants