[rl] Stateful Trainer Send: New Abstractions [1/N] - #48042
Conversation
Signed-off-by: haoaaron <ahao@anyscale.com>
|
This pull request has merge conflicts that must be resolved before it can be |
|
|
||
| async def init_weight_transfer_engine( | ||
| self, request: WeightTransferInitRequest | ||
| self, request: WeightTransferInitRequest | dict |
There was a problem hiding this comment.
Why did you have to change the type here?
There was a problem hiding this comment.
dict was passed by the new ray client, but i changed it to create the typed payload
SumanthRH
left a comment
There was a problem hiding this comment.
Looks mostly good! Left some nits!
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>
| is_sender: bool = True, | ||
| ) -> None: | ||
| self.config = config | ||
| self.is_sender = is_sender |
There was a problem hiding this comment.
Here may violate the rule of Single Source of Truth, we can just get the is_sender value directly from TrainerInitInfo
There was a problem hiding this comment.
would need to pass init_info in to fix, i think its fine as is for now
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 ( What is still needed for integration?
What we can do:
|
Merge it first, aaron please take a look. @hao-aaron |
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: haoaaron <ahao@anyscale.com> Signed-off-by: Aaron Hao <ahao@anyscale.com> Co-authored-by: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com>
… 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>
… 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>
… 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>
… 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>
… 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>
… 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>
### 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>
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>
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>
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:
(
WeightSource/ModuleSource,VLLMWeightSyncClient,TrainerWeightTransferEngine,WeightTransferTrainerFactory). Purelyadditive — the existing worker-side
WeightTransferEngineand every backendengine (NCCL / IPC / sparse NCCL) are untouched and remain valid as-is.
(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.
trainer_send_weightsfrom 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
WeightTransferEngineABC treated the trainer side as stateless — apair of
@staticmethods (trainer_init,trainer_send_weights) plus aper-backend
*TrainerSendWeightsArgsdataclass the caller had to thread back inevery round. That put transfer state in the caller, defeated the type system
with
trainer_args: dict | Any, and made every example open-code the backendconcurrency 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 atrainer_initclassmethod factory anddriven by a parameter-free
send_weights(). Optionalshutdown().TrainerInitInfo— base trainer init info carrying an explicitrankwith an
is_senderproperty; rank 0 is always the sender. Rank is passedexplicitly (not read from a global process group) because that is ambiguous
once several groups — FSDP / TP / PP / EP — exist.
Weight source (
base.py)WeightSourceABC — a re-iterable source of the trainer's weights withtwo channels:
metadata() -> list[ParamMeta]—(name, wire dtype, full shape)for everyparam without transferring (cheap when shapes are known locally; may
cache for producers that must materialize to learn shapes).
(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 overmodule.named_parameters();handles plain and FSDP-sharded modules with no special casing.
ParamMeta(frozenname/dtype/shape) andmaterialize_full_tensor()helper — gathers FSDPDTensorshards viafull_tensor()at send time (once), whilemetadata()reads globalshape/dtype without gathering. Custom producers (Megatron export, MoE
re-fusing) subclass
WeightSource.Control plane (
base.py,clients.py)VLLMWeightSyncClient— a@runtime_checkablestructuralProtocolwithfour 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.
HTTPVLLMWeightSyncClient(RLHF HTTP routes) andRayVLLMWeightSyncClient(fans out to one or moreAsyncLLM/LLMRayactors).
base64-encodes them into
ipc_handles_pickled(_json_safe_update_info);the worker deserializes gated on
VLLM_ALLOW_INSECURE_SERIALIZATION=1.clients.pyis added to the pickle-import allowlist intools/pre_commit/check_forbidden_imports.py.Factory (
factory.py)WeightTransferTrainerFactory— lazy-import registry parallel toWeightTransferEngineFactory(separate registry; trainer and worker neverinstantiate each other's engines). Its registry is intentionally empty in
this PR — the
nccl/ipctrainer engines register in PR 2 / PR 3 alongsidethe concrete classes, so the registry never points at classes that don't yet
exist.
Misc
AsyncLLM.init_weight_transfer_engine/update_weightsnow acceptdict | Request(backward-compatible widening; needed by the Ray clientpath).
Explicitly not in this PR
Deferred to the per-backend PRs because they are coupled trainer+worker changes:
NCCLTrainerWeightTransferEngine/IPCTrainerWeightTransferEngine.NCCLWeightTransferConfig/IPCWeightTransferConfigand the update-infoslimming (moving
packed/ buffer sizes off the per-round update-info onto theconfig, which the worker then reads).
EngineArgsdispatch of dict/CLI config to the right subclass.trainer_send_weightsfrom the worker ABC (PR 3).examples/rl/*, integration tests, and doc rewrites.Breaking changes
None. This PR only adds new symbols.
Testing
Adds GPU-free unit tests for the new abstractions:
ModuleSourcemetadata/iteration/re-iterability,
VLLMWeightSyncClientstructuralconformance (recording / HTTP / Ray), HTTP client JSON-safety (IPC handles
pickled + base64-encoded, NCCL metadata passed through unchanged),
WeightTransferTrainerFactoryregister/dispatch/error paths, andTrainerWeightTransferEnginebase construction. The end-to-end trainertransfer 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.