From 30b4e7f479674a9c4d8889d4857294d3bd5e6849 Mon Sep 17 00:00:00 2001 From: Aaron Hao Date: Thu, 30 Jul 2026 06:59:08 -0700 Subject: [PATCH] [rl] Stateful Trainer Send: IPC [2/N] (#48981) Signed-off-by: haoaaron --- examples/rl/rlhf_http_ipc.py | 51 +-- examples/rl/rlhf_ipc.py | 40 +- examples/rl/rlhf_ipc_fsdp_ep.py | 149 +++---- tests/distributed/test_weight_transfer.py | 85 +++- vllm/distributed/weight_transfer/base.py | 49 ++- vllm/distributed/weight_transfer/factory.py | 30 +- .../distributed/weight_transfer/ipc_engine.py | 405 ++++++++---------- .../weight_transfer/packed_tensor.py | 3 +- 8 files changed, 391 insertions(+), 421 deletions(-) diff --git a/examples/rl/rlhf_http_ipc.py b/examples/rl/rlhf_http_ipc.py index 0a0efcbee3..8c9526e16e 100644 --- a/examples/rl/rlhf_http_ipc.py +++ b/examples/rl/rlhf_http_ipc.py @@ -46,10 +46,12 @@ import torch from openai import OpenAI from transformers import AutoModelForCausalLM -from vllm.distributed.weight_transfer.ipc_engine import ( - IPCTrainerSendWeightsArgs, - IPCWeightTransferEngine, +from vllm.distributed.weight_transfer import ( + HTTPVLLMWeightSyncClient, + ModuleSource, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo BASE_URL = "http://localhost:8000" MODEL_NAME = "facebook/opt-125m" @@ -72,28 +74,6 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list return results -def init_weight_transfer_engine(base_url: str) -> None: - """Initialize weight transfer via HTTP endpoint (no-op for IPC).""" - url = f"{base_url}/init_weight_transfer_engine" - payload = {"init_info": dict()} - response = requests.post(url, json=payload, timeout=60) - response.raise_for_status() - - -def start_weight_update(base_url: str) -> None: - """Start a weight update via HTTP endpoint.""" - url = f"{base_url}/start_weight_update" - response = requests.post(url, json={}, timeout=60) - response.raise_for_status() - - -def finish_weight_update(base_url: str) -> None: - """Finish a weight update via HTTP endpoint.""" - url = f"{base_url}/finish_weight_update" - response = requests.post(url, json={}, timeout=60) - response.raise_for_status() - - def pause_generation(base_url: str) -> None: """Pause generation via HTTP endpoint.""" url = f"{base_url}/pause" @@ -159,23 +139,20 @@ def main(): print("Initializing weight transfer (IPC backend)...") - # Initialize weight transfer on vLLM server (no-op for IPC, but still required) - init_weight_transfer_engine(BASE_URL) + # The trainer engine drives the inference side over HTTP. init for IPC is a + # no-op rendezvous; the same client carries start/update/finish. + engine = WeightTransferTrainerFactory.trainer_init( + init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender + client=HTTPVLLMWeightSyncClient(BASE_URL), + source=ModuleSource(train_model), + ) # Pause generation before weight sync pause_generation(BASE_URL) - # Start weight update, broadcast via IPC, then finish - start_weight_update(BASE_URL) - print("Broadcasting weights via CUDA IPC (HTTP)...") - trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL) - IPCWeightTransferEngine.trainer_send_weights( - iterator=train_model.named_parameters(), - trainer_args=trainer_args, - ) - - finish_weight_update(BASE_URL) + # One call drives start_weight_update / update_weights / finish_weight_update. + engine.send_weights() # Resume generation after weight sync resume_generation(BASE_URL) diff --git a/examples/rl/rlhf_ipc.py b/examples/rl/rlhf_ipc.py index cb85428987..6ed0343ade 100644 --- a/examples/rl/rlhf_ipc.py +++ b/examples/rl/rlhf_ipc.py @@ -29,10 +29,12 @@ from transformers import AutoModelForCausalLM from vllm import LLM, SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.weight_transfer.ipc_engine import ( - IPCTrainerSendWeightsArgs, - IPCWeightTransferEngine, +from vllm.distributed.weight_transfer import ( + ModuleSource, + RayVLLMWeightSyncClient, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo class MyLLM(LLM): @@ -65,23 +67,19 @@ class TrainModel: self.llm_handle = llm_handle def init_weight_transfer(self): - # IPC backend doesn't need initialization info - ray.get( - self.llm_handle.init_weight_transfer_engine.remote(dict(init_info=dict())) + """Build the trainer-side IPC engine (no rendezvous needed for IPC).""" + self.engine = WeightTransferTrainerFactory.trainer_init( + init_info=IPCTrainerInitInfo(rank=0, packed=False), # rank 0 = sender + client=RayVLLMWeightSyncClient(self.llm_handle), + source=ModuleSource(self.train_model), ) - def broadcast_weights( - self, llm_handle: ray.actor.ActorHandle, packed: bool = False - ): - """Broadcast weights to the inference engine using IPC.""" - self.llm_handle = llm_handle - trainer_args = IPCTrainerSendWeightsArgs( - send_mode="ray", llm_handle=llm_handle, packed=packed - ) - IPCWeightTransferEngine.trainer_send_weights( - iterator=self.train_model.named_parameters(), - trainer_args=trainer_args, - ) + def broadcast_weights(self): + """Broadcast weights to the inference engine using IPC. + + Drives start/update/finish on the inference side internally. + """ + self.engine.send_weights() ray.init() @@ -138,10 +136,8 @@ for output in outputs: ray.get(llm.sleep.remote(level=0)) ray.get(train_model.init_weight_transfer.remote()) -# Start weight update, sync weights, then finish -ray.get(llm.start_weight_update.remote()) -ray.get(train_model.broadcast_weights.remote(llm)) -ray.get(llm.finish_weight_update.remote()) +# One call drives start_weight_update / update_weights / finish_weight_update. +ray.get(train_model.broadcast_weights.remote()) ray.get(llm.wake_up.remote(tags=["scheduling"])) diff --git a/examples/rl/rlhf_ipc_fsdp_ep.py b/examples/rl/rlhf_ipc_fsdp_ep.py index 77ac6b4cfc..c4c3559489 100644 --- a/examples/rl/rlhf_ipc_fsdp_ep.py +++ b/examples/rl/rlhf_ipc_fsdp_ep.py @@ -14,8 +14,10 @@ Layout (4 GPUs, TP=1, DP=4, EP): * A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors, waits for initialization, and orchestrates generation / weight-sync. -Uses the built-in ``ray`` send_mode: each FSDP worker calls -``trainer_send_weights`` targeting its colocated LLM actor. +Every FSDP rank builds an ``IPCTrainerWeightTransferEngine`` (via ``trainer_init``) +and calls ``send_weights()``; all ranks join the IPC handle all-gather, and only +rank 0 (the sender) ships the merged handles and drives the DP LLM actors through +its ``RayVLLMWeightSyncClient``. This example was run on 4xH100. """ @@ -23,7 +25,6 @@ This example was run on 4xH100. from __future__ import annotations import os -from dataclasses import asdict import ray import torch @@ -31,19 +32,26 @@ import torch.distributed as dist from huggingface_hub import snapshot_download from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from torch.distributed._tensor import DTensor from torch.distributed.fsdp import fully_shard from transformers import AutoModelForCausalLM from vllm import LLM, SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.weight_transfer.ipc_engine import ( - IPCTrainerSendWeightsArgs, - IPCWeightTransferEngine, - IPCWeightTransferInitInfo, +from vllm.distributed.weight_transfer import ( + ModuleSource, + RayVLLMWeightSyncClient, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo from vllm.utils.network_utils import get_ip, get_open_port +# The inference side only needs the backend; the packed wire params are +# trainer-side and get propagated to the workers at the init handshake. +WEIGHT_TRANSFER_CONFIG = WeightTransferConfig(backend="ipc") +# Packed IPC transfer with a 1 GB buffer (matches the per-chunk buffer size). +PACKED = True +PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 + TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42")) VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42")) @@ -125,72 +133,34 @@ class FSDPTrainWorker: def get_weight_metadata(self): return self.weight_names, self.weight_dtype_names, self.weight_shapes - def gather_and_broadcast_weights_ipc(self, llm_handle, packed: bool = True): - """All-gather full params; all ranks create IPC handles, rank 0 sends. + def setup_engine(self, llm_handles): + """Build the trainer IPC engine on every FSDP rank. - All ranks must call trainer_send_weights so they participate in the - all_gather_object collective inside _all_gather_and_merge_handles. - Only rank 0 actually sends the payload to vLLM (gated by _is_rank_zero). + Called on all ranks: rank 0 becomes the sender (drives the inference + side and performs the no-op IPC init handshake against all DP LLM + actors); the other ranks hold a null-client engine and only join the + IPC handle all-gather during send_weights. """ - - def _full_param_iter(): - # HF's Qwen3MoeExperts (and other recent HF MoE impls) packs - # all experts into two fused 3-D tensors per layer: - # experts.gate_up_proj shape (E, 2*I, H) - # experts.down_proj shape (E, H, I) - # vLLM's Qwen3MoE load_weights still expects the older - # per-expert HF layout (experts..gate_proj.weight, - # experts..up_proj.weight, experts..down_proj.weight), - # so we un-fuse on the fly. Split order matches HF's forward: - # gate, up = linear(x, gate_up_proj[i]).chunk(2, dim=-1) - # → rows [:I] of gate_up_proj[i] are gate, rows [I:] are up. - params = self.model.state_dict() - for name in list(params.keys()): - param = params.pop(name) - if isinstance(param, DTensor): - tensor = param.full_tensor().detach().contiguous() - else: - tensor = param.detach().contiguous() - del param - - if name.endswith(".experts.gate_up_proj") and tensor.dim() == 3: - prefix = name[: -len(".gate_up_proj")] - num_experts, two_inter, _ = tensor.shape - inter = two_inter // 2 - for i in range(num_experts): - expert = tensor[i] - yield ( - f"{prefix}.{i}.gate_proj.weight", - expert[:inter].contiguous(), - ) - yield ( - f"{prefix}.{i}.up_proj.weight", - expert[inter:].contiguous(), - ) - del tensor - elif name.endswith(".experts.down_proj") and tensor.dim() == 3: - prefix = name[: -len(".down_proj")] - num_experts = tensor.shape[0] - for i in range(num_experts): - yield ( - f"{prefix}.{i}.down_proj.weight", - tensor[i].contiguous(), - ) - del tensor - else: - yield name, tensor - - trainer_args = IPCTrainerSendWeightsArgs( - send_mode="ray", - llm_handle=llm_handle, - packed=packed, - packed_buffer_size_bytes=1024 * 1024 * 1024, # 1 GB - ) - IPCWeightTransferEngine.trainer_send_weights( - iterator=_full_param_iter(), - trainer_args=trainer_args, + self.engine = WeightTransferTrainerFactory.trainer_init( + init_info=IPCTrainerInitInfo( + rank=self.rank, # FSDP rank; sender is 0 + packed=PACKED, + packed_buffer_size_bytes=PACKED_BUFFER_SIZE_BYTES, + ), + client=RayVLLMWeightSyncClient(llm_handles), + source=ModuleSource(self.model), ) + def gather_and_broadcast_weights_ipc(self): + """All-gather full params across FSDP ranks; rank 0 sends to vLLM. + + Called on all ranks concurrently. `send_weights` gathers each param and + contributes to the IPC handle all-gather on every rank; only rank 0 (the + sender) drives start/update/finish on the inference side (the other + ranks' client RPCs no-op). + """ + self.engine.send_weights() + @ray.remote(num_cpus=1) class DataParallelInferenceEngine: @@ -224,7 +194,7 @@ class DataParallelInferenceEngine: distributed_executor_backend="ray", enable_expert_parallel=True, gpu_memory_utilization=0.35, - weight_transfer_config=WeightTransferConfig(backend="ipc"), + weight_transfer_config=WEIGHT_TRANSFER_CONFIG, enable_sleep_mode=True, load_format="dummy", dp_rank=r, @@ -267,22 +237,6 @@ class DataParallelInferenceEngine: rank_idx += 1 return ordered - def init_weight_transfer(self): - ray.get( - [ - actor.init_weight_transfer_engine.remote( - dict(init_info=asdict(IPCWeightTransferInitInfo())) - ) - for actor in self.llm_actors - ] - ) - - def start_weight_update(self): - ray.get([actor.start_weight_update.remote() for actor in self.llm_actors]) - - def finish_weight_update(self): - ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors]) - def sleep(self, level: int = 0): ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors]) @@ -370,8 +324,10 @@ def main(): print("-" * 60) # --- Weight transfer --- - print("[transfer] Initializing IPC weight transfer...") - ray.get(inference_engine.init_weight_transfer.remote()) + # The rank-0 FSDP worker owns the trainer engine and drives the inference + # side (init/start/update/finish) for all DP actors via the Ray client. + print("[transfer] Initializing IPC weight transfer (all FSDP ranks)...") + ray.get([w.setup_engine.remote(llm_actors) for w in fsdp_workers]) # Two-phase sleep/wake pattern: # 1. sleep(level=1) — offload weights to CPU, discard KV cache @@ -384,18 +340,11 @@ def main(): print("[sync] Waking weights (KV cache stays free)...") ray.get(inference_engine.wake_up.remote(tags=["weights"])) - print("[sync] Starting weight update...") - ray.get(inference_engine.start_weight_update.remote()) - + # All FSDP ranks participate in the IPC handle all-gather; rank 0's engine + # additionally drives start_weight_update / update_weights / + # finish_weight_update on the inference side. print("[sync] Packed IPC transfer FSDP → vLLM...") - ray.get( - [ - w.gather_and_broadcast_weights_ipc.remote(llm_actors, packed=True) - for w in fsdp_workers - ] - ) - - ray.get(inference_engine.finish_weight_update.remote()) + ray.get([w.gather_and_broadcast_weights_ipc.remote() for w in fsdp_workers]) print("[sync] Weight transfer complete.") print("[sync] Waking KV cache + scheduling...") diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index eeeceb9599..ddc6dec04e 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -27,10 +27,13 @@ from vllm.distributed.weight_transfer import ( WeightTransferTrainerFactory, ) from vllm.distributed.weight_transfer.base import ( + TrainerInitInfo, WeightTransferInitRequest, WeightTransferUpdateRequest, ) from vllm.distributed.weight_transfer.ipc_engine import ( + IPCTrainerInitInfo, + IPCTrainerWeightTransferEngine, IPCWeightTransferEngine, IPCWeightTransferInitInfo, IPCWeightTransferUpdateInfo, @@ -1071,7 +1074,7 @@ def inference_receive_ipc_tensor( import torch - _set_ray_assigned_device() + device = _set_ray_assigned_device() from vllm.config.parallel import ParallelConfig from vllm.config.weight_transfer import WeightTransferConfig @@ -1088,6 +1091,8 @@ def inference_receive_ipc_tensor( for name, tensor in weights: self.received.append((name, tensor.clone())) + # Trainer sends unpacked IPC handles; the worker learns packed=False from + # the init handshake below (IPCWeightTransferInitInfo defaults to False). config = WeightTransferConfig(backend="ipc") vllm_config = MagicMock() parallel_config = MagicMock(spec=ParallelConfig) @@ -1099,9 +1104,7 @@ def inference_receive_ipc_tensor( vllm_config.model_config = MagicMock() recorder = Recorder() - engine = IPCWeightTransferEngine( - config, vllm_config, _get_ray_assigned_device(), recorder - ) + engine = IPCWeightTransferEngine(config, vllm_config, device, recorder) # Transport-only test: bypass the set_current_vllm_config context that # receive_weights enters, since vllm_config here is a mock. import vllm.config as _vllm_config_mod @@ -1211,6 +1214,7 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises(): torch.device("cuda:0"), MagicMock(spec=torch.nn.Module), ) + # No init handshake here, so the engine keeps its default packed=False. dummy_tensor = torch.ones(10, 10, device="cuda:0") _, ipc_handle = reduce_tensor(dummy_tensor) @@ -1264,8 +1268,8 @@ class _DummyTrainerEngine(TrainerWeightTransferEngine): """Minimal concrete trainer engine to exercise base-class + factory.""" @classmethod - def trainer_init(cls, config, init_info, *, client, source): - return cls(config, client=client, source=source) + def trainer_init(cls, init_info, *, client, source): + return cls(client=client, source=source) def send_weights(self): pass @@ -1369,18 +1373,17 @@ class TestModuleSource: class TestTrainerFactory: """WeightTransferTrainerFactory registry mechanics.""" - def test_builtin_registry_has_no_trainer_backends_yet(self): - # Concrete backends register in the per-backend migration PRs. - assert WeightTransferTrainerFactory._registry == {} + def test_registry_has_ipc(self): + # IPC is the first backend migrated to the trainer engine; NCCL / + # sparse NCCL register in later PRs. + assert "ipc" in WeightTransferTrainerFactory._registry def test_register_and_dispatch(self): saved = dict(WeightTransferTrainerFactory._registry) try: WeightTransferTrainerFactory.register_engine("dummy", _DummyTrainerEngine) engine = WeightTransferTrainerFactory.trainer_init( - "dummy", - WeightTransferConfig(backend="dummy"), - MagicMock(), + MagicMock(backend="dummy"), # backend read from the init info client=RecordingClient(), source=ModuleSource(_module_with(("w", torch.zeros(2)))), ) @@ -1395,20 +1398,26 @@ class TestTrainerFactory: def test_unknown_backend_raises(self): with pytest.raises(ValueError, match="Invalid weight transfer backend"): WeightTransferTrainerFactory.trainer_init( - "nope", - WeightTransferConfig(backend="nope"), - MagicMock(), + MagicMock(backend="nope"), client=RecordingClient(), source=ModuleSource(_module_with(("w", torch.zeros(2)))), ) + def test_ipc_init_info_declares_backend(self): + assert IPCTrainerInitInfo.backend == "ipc" + + def test_trainer_init_info_subclass_must_set_backend(self): + with pytest.raises(TypeError, match="class-level `backend`"): + + class _NoBackend(TrainerInitInfo): + pass + class TestTrainerEngineBase: """Base-class construction (no GPU).""" def test_source_stored_and_sender_by_default(self): engine = _DummyTrainerEngine( - WeightTransferConfig(backend="nccl"), client=RecordingClient(), source=ModuleSource(_module_with(("w", torch.zeros(2)))), ) @@ -1417,10 +1426,52 @@ class TestTrainerEngineBase: def test_shutdown_default_is_noop(self): engine = _DummyTrainerEngine( - WeightTransferConfig(backend="nccl"), client=RecordingClient(), source=ModuleSource(_module_with(("w", torch.zeros(2)))), is_sender=False, ) assert engine.is_sender is False engine.shutdown() # must not raise + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 1, + reason="Need at least 1 GPU (CUDA IPC handles).", +) +def test_ipc_trainer_send_weights_drives_client_in_order(): + """send_weights issues start -> update -> finish and ships per-round metadata; + the packed wire param rides the init info, not the per-round update_info.""" + client = RecordingClient() + engine = IPCTrainerWeightTransferEngine( + client=client, + source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))), + packed=False, + ) + + engine.send_weights() + + assert client.order == ["start", "update", "finish"] + assert client.last_update_info is not None + assert client.last_update_info["names"] == ["w"] + assert client.last_update_info["shapes"] == [[4]] + assert "packed" not in client.last_update_info + + +def test_ipc_trainer_init_ships_packed_to_worker(): + """trainer_init drives the inference-side init handshake and propagates the + must-agree `packed` flag to the worker.""" + if torch.accelerator.device_count() < 1: + pytest.skip("Need at least 1 GPU (CUDA IPC handles).") + + client = RecordingClient() + engine = WeightTransferTrainerFactory.trainer_init( + init_info=IPCTrainerInitInfo(rank=0, packed=True), # backend from init info + client=client, + source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))), + ) + + assert isinstance(engine, IPCTrainerWeightTransferEngine) + assert engine.is_sender is True + assert engine.packed is True + assert client.order == ["init"] + assert client.last_init_info == {"packed": True} diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index adddf41ff4..3fd21101e7 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -5,7 +5,15 @@ from abc import ABC, abstractmethod from collections.abc import Iterator from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Generic, + Protocol, + TypeVar, + runtime_checkable, +) import torch from typing_extensions import Self @@ -18,7 +26,7 @@ from vllm.config.weight_transfer import WeightTransferConfig TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo") TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo") -TConfig = TypeVar("TConfig", bound="WeightTransferConfig") +TTrainerInitInfo = TypeVar("TTrainerInitInfo", bound="TrainerInitInfo") # A trainer supplies its parameters as a `WeightSource` (defined below): a # re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()` @@ -108,7 +116,7 @@ class WeightTransferInitInfo(ABC): # noqa: B024 @dataclass -class TrainerInitInfo(WeightTransferInitInfo): +class TrainerInitInfo: """Base trainer-side init info: which trainer rank drives the transfer. `rank` is this trainer process's rank, provided **explicitly** by the @@ -118,10 +126,25 @@ class TrainerInitInfo(WeightTransferInitInfo): while every rank still runs the trainer-side collectives. Backend subclasses add their own (positional) fields; `rank` is keyword-only so that ordering never conflicts. + + Every concrete subclass sets a class-level `backend` string (the same key it + registers under in `WeightTransferTrainerFactory`). The factory reads it to + dispatch, so callers pass only the init info/ It is a `ClassVar` + (a fixed per-backend constant), so it is not an ``__init__`` field. """ + backend: ClassVar[str] + rank: int = field(kw_only=True) + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if not getattr(cls, "backend", None): + raise TypeError( + f"{cls.__name__} must set a class-level `backend` string " + "(the WeightTransferTrainerFactory registry key)." + ) + @property def is_sender(self) -> bool: return self.rank == 0 @@ -373,7 +396,7 @@ class VLLMWeightSyncClient(Protocol): def finish_weight_update(self, weight_version: str | None = None) -> None: ... -class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): +class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]): """Trainer-side weight transfer engine. Symmetric to `WeightTransferEngine` but lives in the training process. @@ -382,6 +405,13 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): plans) on `self`. The `WeightSource` is required at `trainer_init`, then replayed each round by the no-argument `send_weights()`. + Unlike the worker engine, the trainer side does not take a + `WeightTransferConfig`: the backend is selected from the init info's + `backend` `ClassVar` (so callers pass only the init info), and the static + wire params (packed, buffer sizes) ride the backend-specific + `TrainerInitInfo`, which the sender also propagates to the worker at the init + handshake. + Multi-rank trainers: `trainer_init` and `send_weights` are called on *every* trainer rank. Rank 0 is the sender, resolved once at `trainer_init` into `is_sender`. Non-sender ranks still run every @@ -392,22 +422,18 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): Subclasses should define: init_info_cls: Type of backend-specific trainer init info - config_cls: Type of backend-specific config """ - # Subclasses should override these class attributes - init_info_cls: type[TInitInfo] - config_cls: type[TConfig] + # Subclasses should override this class attribute + init_info_cls: type[TTrainerInitInfo] def __init__( self, - config: TConfig, *, client: "VLLMWeightSyncClient", source: "WeightSource", is_sender: bool = True, ) -> None: - self.config = config self.is_sender = is_sender # The real client is held on every rank; each engine only *calls* it when # `is_sender`, so non-sender ranks never touch the wire. @@ -418,8 +444,7 @@ class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): @abstractmethod def trainer_init( cls, - config: TConfig, - init_info: TInitInfo, + init_info: TTrainerInitInfo, *, client: "VLLMWeightSyncClient", source: "WeightSource", diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py index 4ea27c5ef5..1172b46480 100644 --- a/vllm/distributed/weight_transfer/factory.py +++ b/vllm/distributed/weight_transfer/factory.py @@ -18,9 +18,9 @@ if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( + TrainerInitInfo, VLLMWeightSyncClient, WeightSource, - WeightTransferInitInfo, ) logger = init_logger(__name__) @@ -164,9 +164,7 @@ class WeightTransferTrainerFactory: @classmethod def trainer_init( cls, - backend: str, - config: "WeightTransferConfig", - init_info: "WeightTransferInitInfo", + init_info: "TrainerInitInfo", *, client: "VLLMWeightSyncClient", source: "WeightSource", @@ -176,16 +174,22 @@ class WeightTransferTrainerFactory: Called on every trainer rank (multi-rank trainers construct on all ranks; the sender is resolved inside the engine's ``trainer_init``). + The trainer side takes no `WeightTransferConfig` and no separate + `backend` argument: the backend is read from ``init_info.backend`` (a + `ClassVar` on each `TrainerInitInfo` subclass), and the static wire + params ride `init_info`. + Args: - backend: Backend name (must be registered). - config: Backend-specific weight transfer config. - init_info: Backend-specific trainer init info. + init_info: Backend-specific trainer init info. Its `backend` + selects the engine; it also carries the wire params (e.g. + `packed`). client: Inference-side control-plane client. source: `WeightSource` of `(name, tensor)` pairs to send each round. Raises: - ValueError: If the backend is not registered. + ValueError: If `init_info.backend` is not registered. """ + backend = init_info.backend if backend not in cls._registry: available = list(cls._registry.keys()) raise ValueError( @@ -200,7 +204,6 @@ class WeightTransferTrainerFactory: ) return engine_cls.trainer_init( - config=config, init_info=init_info, client=client, source=source, @@ -228,3 +231,12 @@ WeightTransferEngineFactory.register_engine( "vllm.distributed.weight_transfer.sparse_nccl_engine", "SparseNCCLWeightTransferEngine", ) + + +# Trainer-side engines. Backends register here as they migrate to the stateful +# trainer engine; NCCL / sparse NCCL keep their static trainer path until then. +WeightTransferTrainerFactory.register_engine( + "ipc", + "vllm.distributed.weight_transfer.ipc_engine", + "IPCTrainerWeightTransferEngine", +) diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index b191e8c50d..d512c9898c 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -3,19 +3,21 @@ """IPC-based weight transfer engine using CUDA IPC for communication.""" import pickle -from collections.abc import Callable, Iterator from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar import pybase64 as base64 -import ray -import requests import torch from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor +from typing_extensions import Self from vllm import envs from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( + TrainerInitInfo, + TrainerWeightTransferEngine, + VLLMWeightSyncClient, + WeightSource, WeightTransferEngine, WeightTransferInitInfo, WeightTransferUpdateInfo, @@ -31,47 +33,35 @@ from vllm.distributed.weight_transfer.packed_tensor import ( @dataclass -class IPCTrainerSendWeightsArgs: - """Arguments for IPC trainer_send_weights method.""" +class IPCWeightTransferInitInfo(WeightTransferInitInfo): + """Worker-side init info for IPC weight transfer. No rendezvous needed. + + `packed` is a must-agree wire param: the trainer ships it here at the init + handshake so the worker decodes with the same setting the trainer encoded + with. The consumer rebuilds from the IPC handle + `tensor_sizes`, so it does + not need the buffer size (producer-only).""" - send_mode: str | Callable[["IPCWeightTransferUpdateInfo"], None] - """How to send updates to vLLM. Either a string ('ray' or 'http') for - built-in transports, or a callable that receives an - IPCWeightTransferUpdateInfo and performs the send.""" - llm_handle: Any = None - """Ray actor handle or list of handles (required for 'ray' send_mode).""" - url: str | None = None - """Base URL for HTTP endpoint (required for 'http' send_mode).""" packed: bool = False - """Whether to use packed tensor transfer for bounded-memory chunking.""" - packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES - """Size in bytes for each packed tensor buffer when packed=True.""" - - def __post_init__(self): - """Validate that required arguments are provided for the selected mode.""" - if callable(self.send_mode): - return - if self.send_mode == "ray" and self.llm_handle is None: - raise ValueError("llm_handle is required for 'ray' send_mode") - if self.send_mode == "http" and self.url is None: - raise ValueError("url is required for 'http' send_mode") - if self.send_mode not in ("ray", "http"): - raise ValueError( - f"send_mode must be 'ray', 'http', or a callable, " - f"got {self.send_mode!r}" - ) @dataclass -class IPCWeightTransferInitInfo(WeightTransferInitInfo): - """Initialization info for IPC weight transfer backend. No init needed for IPC.""" +class IPCTrainerInitInfo(TrainerInitInfo): + """Trainer-side init info for IPC weight transfer. No rendezvous needed; + `rank` (from `TrainerInitInfo`) identifies this trainer process — rank 0 + ships the merged IPC handles. All ranks still join the handle all-gather. - pass + `packed` / `packed_buffer_size_bytes` are the transfer's wire params. The + trainer propagates `packed` to the worker at `trainer_init` so the two sides + cannot disagree. `backend` is the factory dispatch key.""" + + backend: ClassVar[str] = "ipc" + packed: bool = False + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES @dataclass class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo): - """Update info for IPC weight transfer backend.""" + """Per-round update info for the IPC weight transfer backend.""" names: list[str] dtype_names: list[str] @@ -85,8 +75,6 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo): tensor_sizes: list[int] | None = None """Per-parameter sizes in bytes within the packed buffer. Required when packed=True, unused otherwise.""" - packed: bool = False - """Whether this update uses packed tensor format.""" def __post_init__(self): if self.ipc_handles_pickled is not None: @@ -119,17 +107,14 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo): f"`shapes` should be of the same size as `names`: " f"got {len(self.shapes)} and {len(self.names)}" ) - if ( - not self.packed - and isinstance(self.ipc_handles, list) - and len(self.ipc_handles) != num_params - ): + # Unpacked transfers carry a per-parameter handle dict (a list); packed + # transfers carry a single dict for the whole buffer, so the list check + # only applies to the list case. + if isinstance(self.ipc_handles, list) and len(self.ipc_handles) != num_params: raise ValueError( f"`ipc_handles` should be of the same size as `names`: " f"got {len(self.ipc_handles)} and {len(self.names)}" ) - if self.packed and self.tensor_sizes is None: - raise ValueError("`tensor_sizes` is required when packed=True") class IPCWeightTransferEngine( @@ -154,57 +139,21 @@ class IPCWeightTransferEngine( device: torch.device, model: torch.nn.Module, ) -> None: - """ - Initialize the IPC weight transfer engine. - - Args: - config: The configuration for the weight transfer engine - vllm_config: The full vLLM config - device: The device this worker's model lives on - model: The local model instance which will receive the weights - """ super().__init__(config, vllm_config, device, model) - - def parse_update_info( - self, update_dict: dict[str, Any] - ) -> IPCWeightTransferUpdateInfo: - """Parse update dict, deserializing pickled IPC handles if present. - - HTTP transport sends IPC handles as a base64-encoded pickle under the - key ``ipc_handles_pickled``. This method deserializes them back into - ``ipc_handles`` before constructing the typed dataclass, keeping - serialization concerns out of the dataclass itself. - - Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the - payload is deserialized via ``pickle.loads``. - """ - pickled = update_dict.pop("ipc_handles_pickled", None) - if pickled is not None: - if update_dict.get("ipc_handles") is not None: - raise ValueError( - "Cannot specify both `ipc_handles` and `ipc_handles_pickled`" - ) - - if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION: - raise ValueError( - "Refusing to deserialize `ipc_handles_pickled` without " - "VLLM_ALLOW_INSECURE_SERIALIZATION=1" - ) - - update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled)) - - return super().parse_update_info(update_dict) + # Set from the trainer-supplied init info at the handshake; defaults are + # only for the (unreachable) receive-before-init case. + self.packed = False def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None: """ - Initialize the weight transfer mechanism. - This is called once at the beginning of training. - No initialization needed for IPC backend. + Initialize the weight transfer mechanism. No data-plane rendezvous is + needed for IPC; this just records the trainer-supplied wire params so the + worker decodes exactly as the trainer encoded. Args: - init_info: IPC initialization info (empty) + init_info: IPC initialization info (carries `packed`). """ - pass + self.packed = init_info.packed def start_weight_update(self) -> None: """Initialize layerwise reloading for the incoming checkpoint weights.""" @@ -226,6 +175,10 @@ class IPCWeightTransferEngine( """ Receive weights from the trainer via CUDA IPC handles and load them. + Whether the transfer is packed is read from `self.packed`, set at the + init handshake from the trainer's init info, so it is guaranteed to + match how the trainer encoded. + Args: update_info: IPC update info containing parameter names, dtypes, shapes, and IPC handles. Each IPC handle is a mapping between physical @@ -238,8 +191,9 @@ class IPCWeightTransferEngine( # rebuilt on the device the model lives on. device_index = self.device.index - if update_info.packed: - assert update_info.tensor_sizes is not None + if self.packed: + if update_info.tensor_sizes is None: + raise ValueError("`tensor_sizes` is required when packed=True") assert isinstance(update_info.ipc_handles, dict) weights = packed_ipc_consumer( ipc_handle=update_info.ipc_handles, @@ -285,65 +239,108 @@ class IPCWeightTransferEngine( pass @staticmethod - def trainer_send_weights( - iterator: Iterator[tuple[str, torch.Tensor]], - trainer_args: dict[str, Any] | IPCTrainerSendWeightsArgs, - ) -> None: - """Send weights from trainer to inference workers via CUDA IPC. + def trainer_send_weights(*args: Any, **kwargs: Any) -> None: + """Removed. Use the stateful `IPCTrainerWeightTransferEngine` instead. - Supports two transport modes ('ray' and 'http') and two transfer - strategies: - - Non-packed (default): all weights in a single API call. - - Packed (packed=True): chunked transfer with bounded GPU memory. - - For multi-GPU training, all ranks must call this method in - parallel. IPC handles are all-gathered across ranks and merged - so that each vLLM worker can find its own GPU UUID. Only rank 0 - sends the payload to vLLM. - - .. note:: - This method calls ``update_weights`` internally. The caller must - call ``start_weight_update`` before and ``finish_weight_update`` - after this method. - - Args: - iterator: Iterator of (name, tensor) pairs. For multi-GPU, - each rank should yield the full tensor on its own GPU - (e.g. via FSDP full_tensor()). - trainer_args: IPCTrainerSendWeightsArgs or equivalent dict. + Transitional stub kept only to satisfy the (still abstract) + `WeightTransferEngine.trainer_send_weights`; that member is dropped from + the worker ABC once every backend has migrated to the trainer engine. """ - args = ( - IPCTrainerSendWeightsArgs(**trainer_args) - if isinstance(trainer_args, dict) - else trainer_args + raise NotImplementedError( + "The static IPC trainer path has been replaced by " + "IPCTrainerWeightTransferEngine. Build it via " + "WeightTransferTrainerFactory.trainer_init(IPCTrainerInitInfo(...), " + "client=..., source=...) and drive it with send_weights()." ) - device_index = torch.accelerator.current_device_index() - gpu_uuid = str(torch.cuda.get_device_properties(device_index).uuid) - if args.packed: - IPCWeightTransferEngine._send_packed(iterator, args, gpu_uuid) - else: - IPCWeightTransferEngine._send_unpacked(iterator, args, gpu_uuid) - @staticmethod - def _is_rank_zero() -> bool: - """Return True if this is rank 0 or no distributed group exists.""" - if not torch.distributed.is_initialized(): - return True - return torch.distributed.get_rank() == 0 - @staticmethod +class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitInfo]): + """Trainer-side CUDA IPC weight transfer engine. + + Called on every trainer rank. For multi-rank (e.g. FSDP) trainers all ranks + iterate the source (materializing each tensor) and contribute to the + IPC-handle all-gather; only the sender (rank 0) ships the merged handles to + the inference side. IPC transfer is straight-line (no concurrent broadcast + like NCCL): `update_weights` *is* the transfer, and it rides the client, so + it no-ops on non-senders. + + `packed` / `packed_buffer_size_bytes` come from `IPCTrainerInitInfo`; the + sender propagates `packed` to the worker at `trainer_init`. + """ + + init_info_cls = IPCTrainerInitInfo + + def __init__( + self, + *, + client: VLLMWeightSyncClient, + source: WeightSource, + is_sender: bool = True, + packed: bool = False, + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, + ) -> None: + super().__init__(client=client, source=source, is_sender=is_sender) + self.packed = packed + self.packed_buffer_size_bytes = packed_buffer_size_bytes + self.device_index = torch.accelerator.current_device_index() + self.gpu_uuid = str(torch.cuda.get_device_properties(self.device_index).uuid) + + @classmethod + def trainer_init( + cls, + init_info: IPCTrainerInitInfo, + *, + client: VLLMWeightSyncClient, + source: WeightSource, + ) -> Self: + engine = cls( + client=client, + source=source, + is_sender=init_info.is_sender, + packed=init_info.packed, + packed_buffer_size_bytes=init_info.packed_buffer_size_bytes, + ) + # IPC needs no data-plane rendezvous. The sender ships the must-agree + # `packed` flag so the worker decodes exactly as this trainer encodes. + if engine.is_sender: + engine.client.init_weight_transfer_engine({"packed": init_info.packed}) + return engine + + def send_weights(self) -> None: + source = self.source + if self.is_sender: + self.client.start_weight_update() + # Unpacked returns strong refs to its IPC-shared copies; they must stay + # alive across `finish_weight_update` too. + weight_refs = self._send(source) + if self.is_sender: + self.client.finish_weight_update() + self._post_send_sync() + del weight_refs + + # ---- data plane (runs on all ranks; only the sender ships) ---- + + def _send(self, source: WeightSource) -> list[torch.Tensor] | None: + if self.packed: + self._send_packed(source) + return None + return self._send_unpacked(source) + def _all_gather_and_merge_handles( + self, handles: list[dict[str, tuple]], ) -> list[dict[str, tuple]]: """All-gather and merge IPC handle dicts across ranks in one call. Each rank contributes a list of {gpu_uuid: ipc_args} dicts (one per parameter or one per chunk). A single all_gather_object - collects every rank's full list, then rank 0 merges per-index so + collects every rank's full list, then the sender merges per-index so each dict maps every GPU UUID to its args. - Non-rank-0 returns a list of empty dicts. - No-op (returns handles unchanged) when no distributed group exists. + The all-gather runs over the *default* process group; this assumes the + default group is exactly the set of colocated trainer ranks and that the + sender is a member. No-op (returns handles unchanged) when no distributed + group exists. """ if ( not torch.distributed.is_initialized() @@ -357,7 +354,7 @@ class IPCWeightTransferEngine( torch.distributed.barrier() torch.cuda.synchronize() - if torch.distributed.get_rank() == 0: + if self.is_sender: merged: list[dict[str, tuple]] = [] for param_idx in range(len(handles)): m: dict[str, tuple] = {} @@ -378,25 +375,23 @@ class IPCWeightTransferEngine( torch.distributed.barrier() torch.cuda.ipc_collect() - @staticmethod - def _send_unpacked( - iterator: Iterator[tuple[str, torch.Tensor]], - args: IPCTrainerSendWeightsArgs, - gpu_uuid: str, - ) -> None: - """Send all weights in a single API call (non-packed mode).""" + def _send_unpacked(self, source: WeightSource) -> list[torch.Tensor]: + """Iterate the source, build one IPC handle per param, all-gather the + handles across ranks, and (sender) ship them in one update call. + + Returns the strong refs to every contiguous copy. reduce_tensor's args + do NOT keep storage alive, and non-contiguous inputs allocate fresh + storage in .contiguous(); the caller must keep these alive until the + post-send barrier (past `finish`) so the consumer's IPC views stay valid. + """ names: list[str] = [] dtype_names: list[str] = [] shapes: list[list[int]] = [] ipc_handles: list[dict[str, tuple]] = [] - # Hold strong refs to every contiguous copy until the send + post-send - # sync completes. reduce_tensor's returned args do NOT keep storage - # alive, and non-contiguous inputs allocate fresh storage in - # .contiguous() that would otherwise be GC'd before the consumer opens - # the IPC handle. weight_refs: list[torch.Tensor] = [] - for name, tensor in iterator: + for name, tensor in source: + # FSDP shards were gathered by the source; ensure contiguity here. names.append(name) dtype_names.append(str(tensor.dtype).split(".")[-1]) shapes.append(list(tensor.shape)) @@ -404,98 +399,64 @@ class IPCWeightTransferEngine( weight = tensor.detach().contiguous() weight_refs.append(weight) _, ipc_args = reduce_tensor(weight) - ipc_handles.append({gpu_uuid: ipc_args}) + ipc_handles.append({self.gpu_uuid: ipc_args}) - ipc_handles = IPCWeightTransferEngine._all_gather_and_merge_handles(ipc_handles) + ipc_handles = self._all_gather_and_merge_handles(ipc_handles) + self._do_send( + names=names, + dtype_names=dtype_names, + shapes=shapes, + ipc_handles=ipc_handles, + ) + return weight_refs - if IPCWeightTransferEngine._is_rank_zero(): - IPCWeightTransferEngine._do_send( - args=args, - names=names, - dtype_names=dtype_names, - shapes=shapes, - ipc_handles=ipc_handles, - ) - - IPCWeightTransferEngine._post_send_sync() - - @staticmethod - def _send_packed( - iterator: Iterator[tuple[str, torch.Tensor]], - args: IPCTrainerSendWeightsArgs, - gpu_uuid: str, - ) -> None: + def _send_packed(self, source: WeightSource) -> None: """Send weights in bounded-memory chunks (packed mode).""" - post_iter_func: Callable = lambda item: item[1] - for chunk in packed_ipc_producer( - iterator=iterator, - gpu_uuid=gpu_uuid, - post_iter_func=post_iter_func, - buffer_size_bytes=args.packed_buffer_size_bytes, + iterator=iter(source), + gpu_uuid=self.gpu_uuid, + post_iter_func=lambda item: item[1], + buffer_size_bytes=self.packed_buffer_size_bytes, ): - ipc_handle = IPCWeightTransferEngine._all_gather_and_merge_handles( - [chunk.ipc_handle] - )[0] + ipc_handle = self._all_gather_and_merge_handles([chunk.ipc_handle])[0] + self._do_send( + names=chunk.names, + dtype_names=chunk.dtype_names, + shapes=chunk.shapes, + ipc_handles=ipc_handle, + tensor_sizes=chunk.tensor_sizes, + ) + # Per-chunk barrier: the producer reuses a single IPC buffer across + # chunks, but only the sender waits for the consumers (via _do_send). + # Without syncing every rank here, non-sender ranks race ahead and + # overwrite their buffer while their colocated worker is still + # reading the current chunk, silently corrupting the transfer. + self._post_send_sync() - if IPCWeightTransferEngine._is_rank_zero(): - IPCWeightTransferEngine._do_send( - args=args, - names=chunk.names, - dtype_names=chunk.dtype_names, - shapes=chunk.shapes, - ipc_handles=ipc_handle, - tensor_sizes=chunk.tensor_sizes, - packed=True, - ) - - IPCWeightTransferEngine._post_send_sync() - - @staticmethod def _do_send( - args: IPCTrainerSendWeightsArgs, + self, names: list[str], dtype_names: list[str], shapes: list[list[int]], ipc_handles: list[dict[str, tuple]] | dict[str, tuple], tensor_sizes: list[int] | None = None, - packed: bool = False, ) -> None: - """Send a single update payload via the configured transport.""" + """Build one update payload and ship it via the client. Only the sender + ships (non-sender ranks already contributed to the handle all-gather). + + Emits raw `ipc_handles`; transports that cannot carry them natively + (HTTP/JSON) pickle them in their client (see `HTTPVLLMWeightSyncClient`). + """ + if not self.is_sender: + return update_fields: dict[str, Any] = { "names": names, "dtype_names": dtype_names, "shapes": shapes, - "packed": packed, + "ipc_handles": ipc_handles, } if tensor_sizes is not None: update_fields["tensor_sizes"] = tensor_sizes - update_fields["ipc_handles"] = ipc_handles update_info = IPCWeightTransferUpdateInfo(**update_fields) - - if callable(args.send_mode): - args.send_mode(update_info) - elif args.send_mode == "ray": - handles = ( - args.llm_handle - if isinstance(args.llm_handle, list) - else [args.llm_handle] - ) - ray.get( - [ - h.update_weights.remote(dict(update_info=asdict(update_info))) - for h in handles - ] - ) - elif args.send_mode == "http": - pickled_handles = base64.b64encode(pickle.dumps(ipc_handles)).decode( - "utf-8" - ) - http_fields = {k: v for k, v in update_fields.items() if k != "ipc_handles"} - http_fields["ipc_handles_pickled"] = pickled_handles - - url = f"{args.url}/update_weights" - payload = {"update_info": http_fields} - response = requests.post(url, json=payload, timeout=300) - response.raise_for_status() + self.client.update_weights(asdict(update_info)) diff --git a/vllm/distributed/weight_transfer/packed_tensor.py b/vllm/distributed/weight_transfer/packed_tensor.py index ce42c691f4..b34c58d53d 100644 --- a/vllm/distributed/weight_transfer/packed_tensor.py +++ b/vllm/distributed/weight_transfer/packed_tensor.py @@ -10,8 +10,7 @@ from typing import Any import torch from torch.multiprocessing.reductions import reduce_tensor -# Default values for packed tensor configuration. -# These are imported by NCCLWeightTransferUpdateInfo and trainer_send_weights. +# Default values for packed tensor transfer. DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 # 1GB DEFAULT_PACKED_NUM_BUFFERS = 2