vllm.distributed.weight_transfer.sharded_rdt_trainer ¶

@dataclass
class ShardedRDTTrainerInitInfo(TrainerInitInfo):
    """Trainer init info for the sharded-RDT backend.

    Identical on every rank except ``rank`` (rank 0 is the sender). Carries only
    the must-agree wire params; the sender forwards them verbatim onto the
    worker-side init info so the two cannot drift. Server-actor names are
    generated per rank and all-gathered by the engine, not supplied here.
    """

    backend: ClassVar[str] = "sharded_rdt"

    num_consumers: int
    """Total inference-worker (consumer) count across the whole fleet
    (DP*TP*PP*PCP), for the M:N block assignment / free ref-count."""
    workers_per_replica: int = 0
    """Consumers per inference DEPLOYMENT (``num_consumers // num_replicas``).

    Fixes the slot-sharing groups: consumers whose ids differ by a multiple of
    this are the same worker of different deployments, so they bake identical
    plans and pull byte-identical chunks, and ONE registered serve slot can
    serve all of them. 0 disables sharing, and so does the single-deployment
    value ``num_consumers``, which makes every group a singleton and the serve
    path identical to the unshared one.
    """
    trainer_actor_namespace: str | None = None
    """Ray namespace the engine spawns its serve actors in. The inference
    workers (which run in their own EngineCore subprocess with its own
    ``ray.init``) resolve those actors by name, so this must be the namespace
    they can see. Forwarded to the worker-side init info."""
    num_rdt_buffers: int = 2
    """Serve/receive ring depth K (must match the worker)."""
    buffer_presize_gb: float = 0.0
    """Serve-buffer pre-size floor in GiB (avoids NIXL desc-cache churn)."""
    gather_lookahead: int = DEFAULT_GATHER_LOOKAHEAD
    """Gathered-but-unfreed groups the gather loop may run ahead by. Bounds
    trainer-resident memory at ``gather_lookahead + 1`` groups."""
    stall_timeout_s: float = DEFAULT_STALL_TIMEOUT_S
    """Seconds of no publish/serve/free progress before the producer fails the
    sync (see ``DEFAULT_STALL_TIMEOUT_S``)."""

class ShardedRDTTrainerWeightTransferEngine(
    TrainerWeightTransferEngine[ShardedRDTTrainerInitInfo]
):
    """Trainer-side engine for the pull-based sharded-RDT backend.

    Lives on every trainer rank. Owns a per-rank `_RDTProducerServer` actor
    (the NIXL serve surface). `send_weights` gathers this rank's weights
    group-by-group from the `WeightSource`, shares each group into the server
    over CUDA IPC, and — on the sender — drives the inference-side handshake so
    the workers pull. Non-sender ranks only gather (staying in the collective).
    """

    init_info_cls = ShardedRDTTrainerInitInfo

    def __init__(
        self,
        *,
        client: VLLMWeightSyncClient,
        source: WeightSource,
        is_sender: bool = True,
        init_info: ShardedRDTTrainerInitInfo,
    ) -> None:
        super().__init__(client=client, source=source, is_sender=is_sender)
        self._init_info = init_info
        self._server: Any = None  # the _RDTProducerServer actor handle
        self._server_name: str | None = None
        # Group-major metadata / partition, computed at trainer_init.
        self._meta: list[ParamMeta] = []
        self._groups: list[list[str]] = []
        # Ownership resolved at trainer_init from the fleet's held names: the
        # distinct owner sets, the per-name index into them, this rank's groups
        # and its held set. Routing itself is consumer-side; the trainer only
        # ships the table.
        self._owner_sets: list[list[int]] = []
        self._name_owner_class: list[int] = []
        self._owned_idx: list[int] = []
        self._held_names: set[str] | None = None
        # Strong refs to gathered tensors we've shared into the server, keyed by
        # group index. CUDA-IPC exports must outlive the importer, so we hold
        # them until the server reports the group freed. See send_weights.
        self._inflight: dict[int, dict[str, torch.Tensor]] = {}
        # [RDT-EXPORT-RING] Reusable packed slots, one per resident group. See
        # _pack_group_for_export. Persistent across syncs: freeing and
        # reallocating them each sync fragmented the caching allocator badly
        # enough to grow reserved memory ~1.5 GB per sync.
        self._export_ring: list[torch.Tensor | None] = []
        self._export_ring_args: list[tuple | None] = []
        # Which live group holds which slot, and the slots free to hand out. A
        # modular counter is unsafe here: the credit gate bounds the NUMBER of
        # unfreed groups, not their order, so it can hand a later group the slot
        # a still-live group is being served from.
        self._slot_of_group: dict[int, int] = {}
        self._free_slots: list[int] = []
        # Set by trainer_init from the server-name all-gather; retained so a
        # restarted consumer can be re-initialized without another collective.
        self._server_names: list[str] | None = None

    def _rpc(self, method: str, *args: Any) -> Any:
        """Call one of the server actor's methods and block for the result.
        The single seam through which the engine talks to its server, so tests
        can inject a local (non-Ray) fake server."""
        import ray

        return ray.get(getattr(self._server, method).remote(*args))

    @staticmethod
    def _contig_stride(shape) -> tuple:
        stride: list[int] = []
        acc = 1
        for s in reversed(list(shape)):
            stride.append(acc)
            acc *= int(s)
        return tuple(reversed(stride))

    def _pack_group_for_export(self, held: list, slot_idx: int) -> tuple:
        """Copy ``held`` into ring slot ``slot_idx``; return ``(storages, views,
        refs)`` in the same shape the per-storage path returns, but with ONE
        storage. Views are built exactly as ``publish_group`` rebuilds them, so
        the two sides cannot disagree about layout.

        Slot safety is not the credit gate alone: that bounds how MANY groups
        are unfreed, not the order they free in, and barriers do complete out of
        order. The caller takes ``slot_idx`` from a pool keyed by live group
        (``_slot_of_group`` / ``_free_slots``); ``_drop_inflight`` returns a slot
        only once its group is freed everywhere.
        """
        offsets: list[int] = []
        cur = 0
        for _name, t in held:
            cur = (cur + 15) & ~15  # 16B keeps every element size we ship aligned
            offsets.append(cur)
            cur += t.numel() * t.element_size()
        need = max(16, (cur + 15) & ~15)

        slot = self._export_ring[slot_idx]
        device = held[0][1].device
        if slot is None or slot.numel() < need or slot.device != device:
            self._export_ring[slot_idx] = self._export_ring_args[slot_idx] = None
            slot = torch.empty(need, dtype=torch.uint8, device=device)
            self._export_ring[slot_idx] = slot

        ust = slot.untyped_storage()
        sid = ust.data_ptr()
        reduce_args = self._export_ring_args[slot_idx]
        if reduce_args is None:
            base = torch.empty(0, dtype=torch.uint8, device=device)
            base.set_(ust, 0, (ust.nbytes(),))
            _rebuild, reduce_args = reduce_tensor(base)
            self._export_ring_args[slot_idx] = reduce_args

        views: dict[str, tuple] = {}
        refs: dict[str, torch.Tensor] = {}
        for (name, t), off in zip(held, offsets):
            shape, esz = tuple(t.shape), t.element_size()
            stride = self._contig_stride(shape)
            dst = torch.as_strided(slot.view(t.dtype), shape, stride, off // esz)
            dst.copy_(t)
            refs[name] = dst
            views[name] = (
                sid,
                str(t.dtype).split(".")[-1],
                list(shape),
                list(stride),
                off // esz,
            )
        # The producer sidecar reads this slot from another process over CUDA
        # IPC, while these copies are merely enqueued on this rank's stream;
        # nothing orders the two. Without this sync it packs bytes the copy has
        # not written yet -- silent corruption, and only for ring-packed tensors.
        if slot.is_cuda:
            torch.cuda.current_stream(slot.device).synchronize()
        return {sid: reduce_args}, views, refs

    def _publish_async(self, group_idx: int, entries):
        """Fire publish_group WITHOUT waiting on the RPC (the gather loop
        overlaps the publish's server-side rebuild with the next group's
        gather) and return a handle that ``_await_publish`` resolves. Ray actor
        handle in production; a plain (non-Ray) fake server runs inline."""
        method = self._server.publish_group
        remote = getattr(method, "remote", None)
        if remote is not None:
            return remote(group_idx, entries)
        return method(group_idx, entries)

    def _await_publish(self, ref) -> None:
        """Resolve one async publish so a server-side rebuild error surfaces at
        window depth instead of at end_sync. Publishes carry nothing back —
        freed groups flow only through wait_freed/end_sync (one channel; see
        publish_group)."""
        import ray

        if isinstance(ref, ray.ObjectRef):
            ray.get(ref)

    # ---------------- construction ----------------

    @classmethod
    def trainer_init(
        cls,
        init_info: ShardedRDTTrainerInitInfo,
        *,
        client: VLLMWeightSyncClient,
        source: WeightSource | None = None,
    ) -> Self:
        if source is None:
            raise ValueError(
                "Sharded RDT trainer weight transfer requires a WeightSource."
            )
        engine = cls(
            client=client,
            source=source,
            is_sender=init_info.is_sender,
            init_info=init_info,
        )

        # The VMM-based expandable_segments allocator makes IPC storage opens ~9x
        # slower on both export and rebuild, and CUDA-IPC publish is this engine's
        # hot path. Frameworks enabling it should disable it around send_weights.
        # Only the env var is visible; runtime changes are not introspectable.
        import os as _os

        if "expandable_segments:True" in _os.environ.get("PYTORCH_CUDA_ALLOC_CONF", ""):
            logger.warning(
                "Sharded-RDT trainer: PYTORCH_CUDA_ALLOC_CONF enables "
                "expandable_segments; CUDA-IPC weight publishing will be "
                "several times slower. Disable expandable segments around "
                "weight sync (allocate gather buffers in classic segments)."
            )

        engine._meta = list(source.metadata())
        names = [m.name for m in engine._meta]
        engine._groups = layerwise_groups(names)
        flat = [n for g in engine._groups for n in g]
        if flat != names:
            raise ValueError(
                "Sharded RDT requires a WeightSource whose metadata order is "
                "group-contiguous. Reorder the source so all names sharing a "
                "layer index are adjacent."
            )

        world, rank = engine._world_and_rank()
        engine._resolve_ownership(world, rank)
        engine._spawn_server(sorted(engine._held_names or []))

        # Every rank's server must exist before the sender's init RPC (the worker
        # init calls reserve_serve_buffer back on ALL producer servers). The
        # all-gather of server names doubles as that barrier.
        server_names = engine._all_gather_server_names(world, rank)
        # Retained so a RESTARTED consumer can be re-initialized without another
        # all-gather (see get_worker_init_payload). The uuid actor names stay valid
        # for the run, because producers never restart.
        engine._server_names = server_names

        if engine.is_sender:
            worker_init = engine._build_worker_init_info(server_names)
            engine.client.init_weight_transfer_engine(asdict(worker_init))
        return engine

    def get_worker_init_payload(self) -> dict:
        """The consumer-side init payload, rebuilt on demand. Pure — no collective.

        Reads only retained state, so a restarted inference engine can rejoin at a
        sync boundary: the driver can ask any time, including mid-run with every
        rank in its own training step. A collective here would deadlock, since the
        ranks are not at a matching one.

        Raises:
            RuntimeError: called before ``trainer_init`` cached the server names.
        """
        if self._server_names is None:
            raise RuntimeError(
                "get_worker_init_payload requires trainer_init to have run (the "
                "producer server names are gathered there)."
            )
        return asdict(self._build_worker_init_info(self._server_names))

    def _world_and_rank(self) -> tuple[int, int]:
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            return torch.distributed.get_world_size(), torch.distributed.get_rank()
        return 1, self._init_info.rank

    def _resolve_ownership(self, world: int, rank: int) -> None:
        """Resolve which rank holds which name, and this rank's publish plan.

        A source may hold only part of the model — pipeline stages, expert
        parallelism, or any mix — so each rank declares its held names and the
        fleet all-gathers them. The consumers route per name, so the wire carries
        the transposed result: the distinct owner sets, and a per-name index into
        them.

        The masks are positional over metadata order, which is why the metadata
        digest is checked first: a rank whose names disagree would transpose into
        the wrong owners entirely.
        """
        assert self.source is not None  # guaranteed by trainer_init
        num_groups = len(self._groups)
        names = [m.name for m in self._meta]
        held = self.source.held_names()
        if held is None:
            held_set = set(names)
        else:
            held_set = {str(n) for n in held}
            unknown = sorted(held_set - set(names))
            if unknown:
                raise ValueError(
                    f"WeightSource.held_names() lists {len(unknown)} name(s) not "
                    f"in metadata(), e.g. {unknown[:3]}."
                )
            if not held_set:
                raise ValueError(
                    "WeightSource.held_names() is empty; a rank with nothing to "
                    "serve cannot take part in the gather"
                )

        # One collective carries this rank's holdings as a bitmask over metadata
        # order, plus the digest of the metadata it indexes into. The digest is
        # what makes partial ownership safe: only the sender's metadata reaches
        # the consumers, so a rank describing just its own share would silently
        # serve the wrong weights.
        digest = self._meta_digest()
        mask = bytearray((len(names) + 7) // 8)
        for i, n in enumerate(names):
            if n in held_set:
                mask[i >> 3] |= 1 << (i & 7)
        per_rank = self._all_gather_owned(world, (digest, bytes(mask)))
        mismatched = [r for r, (d, *_rest) in enumerate(per_rank) if d != digest]
        if mismatched:
            raise ValueError(
                f"WeightSource.metadata() disagrees across trainer ranks "
                f"(rank {rank} digest {digest}, differing ranks {mismatched[:4]}). "
                "Every rank must describe the WHOLE model, even when it holds "
                "only some of it."
            )
        masks = [m for _d, m in per_rank]

        # Transpose to per-name owners, then dedup into classes. Numbering by
        # FIRST APPEARANCE in metadata order keeps it a pure function of
        # rank-identical inputs, which is what lets a rejoining consumer rebuild
        # the identical table from get_worker_init_payload.
        owner_sets: list[list[int]] = []
        class_of_owners: dict[tuple[int, ...], int] = {}
        name_owner_class: list[int] = []
        for i, n in enumerate(names):
            owners = tuple(r for r, m in enumerate(masks) if m[i >> 3] & (1 << (i & 7)))
            if not owners:
                raise ValueError(
                    f"no trainer rank holds {n!r}; every name in metadata() must "
                    "be held by at least one rank or it can never be served."
                )
            ci = class_of_owners.get(owners)
            if ci is None:
                ci = len(owner_sets)
                class_of_owners[owners] = ci
                owner_sets.append(list(owners))
            name_owner_class.append(ci)

        self._owner_sets = owner_sets
        self._name_owner_class = name_owner_class
        # Groups holding anything here: exactly what iteration must cover.
        self._owned_idx = [
            gi
            for gi in range(num_groups)
            if any(n in held_set for n in self._groups[gi])
        ]
        self._held_names = held_set

    def _validate_held_yields(self, gi: int, names, tensors) -> None:
        """Check ``held_names()`` against what the source actually yields — the
        one place both sit side by side.

        Without it, a source that claims a name but yields ``None`` for it dies
        by stall watchdog 300s later: consumers route pulls here, the pull
        passes the served-names guard, and the cache wait never completes. With
        it, that is an immediate error naming the weight. One set lookup per
        name per sync."""
        if self._held_names is None:
            return
        for name, tensor in zip(names, tensors):
            if (tensor is None) == (name in self._held_names):
                claim = "does not hold" if tensor is not None else "holds"
                have = "a real tensor" if tensor is not None else "None"
                raise RuntimeError(
                    f"held_names() disagrees with the yielded tensors: group {gi} "
                    f"name {name!r} yielded {have} but this rank {claim} it. A "
                    "WeightSource must yield a real tensor for every held name "
                    "and None for the rest (see WeightSource.held_names)."
                )

    def _meta_digest(self) -> str:
        """Stable digest of this rank's metadata (name order + count)."""
        import hashlib

        h = hashlib.sha256()
        h.update(f"{len(self._meta)}\n".encode())
        for m in self._meta:
            h.update(m.name.encode())
            h.update(b"\n")
        return h.hexdigest()[:16]

    def _all_gather_owned(self, world: int, mine: tuple) -> list[tuple]:
        """All-gather each rank's (metadata digest, held-name bitmask)."""
        if world <= 1 or not (
            torch.distributed.is_available() and torch.distributed.is_initialized()
        ):
            return [mine]
        gathered: list[Any] = [None] * world
        torch.distributed.all_gather_object(gathered, mine)
        return gathered

    def _all_gather_server_names(self, world: int, rank: int) -> list[str]:
        assert self._server_name is not None
        if world <= 1 or not (
            torch.distributed.is_available() and torch.distributed.is_initialized()
        ):
            return [self._server_name]
        gathered: list[str | None] = [None] * world
        torch.distributed.all_gather_object(gathered, self._server_name)
        return [n for n in gathered if n is not None]

    def _spawn_server(self, served_names: list[str]) -> None:
        import ray
        from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy

        # The trainer is usually a separate install from the workers, so it gets
        # its own check rather than trusting the consumer side's.
        check_ray_rdt_version()

        ii = self._init_info
        self._server_name = f"vllm_rdt_producer_{uuid.uuid4().hex[:12]}_rk{ii.rank}"
        node_id = ray.get_runtime_context().get_node_id()
        # Pin the server to this rank's physical GPU: num_gpus=0 so Ray does not
        # allocate a second, CUDA_VISIBLE_DEVICES so CUDA IPC to the rank's
        # gathered tensors works. max_concurrency > 1 serves pulls while control
        # calls are in flight; enable_tensor_transport gives the NIXL serve.
        gpu_ids = ray.get_gpu_ids()
        # The server is the trainer rank's process twin: forward the env the
        # rank runs under (library paths etc.) so it imports torch/vllm the
        # same way, then pin it to the rank's physical GPU for CUDA IPC.
        import os

        env_vars = {
            k: os.environ[k]
            for k in (
                "LD_LIBRARY_PATH",
                "LD_PRELOAD",
                "NCCL_CUMEM_ENABLE",
                "VLLM_NCCL_SO_PATH",
                "PATH",
            )
            if k in os.environ
        }
        if gpu_ids:
            # A num_gpus=0 actor sharing the rank's GPU: Ray would set
            # CUDA_VISIBLE_DEVICES="" and hide every GPU, so tell it not to touch
            # the var and pin the device ourselves.
            env_vars["RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"] = "1"
            env_vars["CUDA_VISIBLE_DEVICES"] = str(gpu_ids[0])
        runtime_env = {"env_vars": env_vars} if env_vars else {}
        server_cls = ray.remote(_RDTProducerServer).options(
            name=self._server_name,
            namespace=ii.trainer_actor_namespace,
            num_cpus=0,
            num_gpus=0,
            # Thread budget: under EP-local routing every consumer pulls from every
            # producer, and with issue-ahead each can park K produce calls in the
            # cache-wait — C*K blocked calls. A smaller pool queues the control
            # plane behind them and deadlocks: the parked produces wait for a
            # publish that can never get a thread (the sync-2 wedge at 235B tp8,
            # 16 parked vs 8 threads). Size for the worst case plus control-plane
            # slack.
            max_concurrency=ii.num_consumers * ii.num_rdt_buffers + 4,
            enable_tensor_transport=True,
            scheduling_strategy=NodeAffinitySchedulingStrategy(
                node_id=node_id, soft=False
            ),
            runtime_env=runtime_env,
        )
        self._server = server_cls.remote(
            served_names=served_names,
            num_rdt_buffers=ii.num_rdt_buffers,
            buffer_presize_gb=ii.buffer_presize_gb,
            gather_lookahead=ii.gather_lookahead,
            stall_timeout_s=ii.stall_timeout_s,
            workers_per_replica=ii.workers_per_replica,
        )
        ray.get(self._server.ping.remote())
        # Pre-barrier NIXL warmup: must complete before the server-name
        # all-gather so no rank can be inside a collective yet (see warmup_nixl).
        ray.get(self._server.warmup_nixl.remote())

    def _build_worker_init_info(self, server_names: list[str]):
        from vllm.distributed.weight_transfer.sharded_rdt_engine import (
            ShardedRDTWeightTransferInitInfo,
        )

        group_lens = [len(g) for g in self._groups]
        names = [m.name for m in self._meta]
        dtype_names = [str(m.dtype).split(".")[-1] for m in self._meta]
        shapes = [list(m.shape) for m in self._meta]
        return ShardedRDTWeightTransferInitInfo(
            trainer_actor_names=server_names,
            trainer_actor_namespace=self._init_info.trainer_actor_namespace,
            produce_method_name=PRODUCE_METHOD_NAME,
            names=names,
            dtype_names=dtype_names,
            shapes=shapes,
            group_lens=group_lens,
            owner_sets=self._owner_sets,
            name_owner_class=self._name_owner_class,
            num_consumers=self._init_info.num_consumers,
            num_rdt_buffers=self._init_info.num_rdt_buffers,
            buffer_presize_gb=self._init_info.buffer_presize_gb,
        )

    # ---------------- per-round ----------------

    def send_weights(self, live_consumer_ids: Collection[int] | None = None) -> None:
        """Gather this rank's weights and publish them for the consumers to pull.

        ``live_consumer_ids`` restricts the sync to the consumers still alive;
        ``None`` serves the whole provisioned set. The provisioned geometry is
        FROZEN for the run — a degraded sync only lowers the live count handed to
        ``begin_sync``. Every rank must get the SAME live set, since they share
        the gather collectives, so the caller computes it once for all of them.
        """
        assert self.source is not None
        if live_consumer_ids is None:
            live_ids = list(range(self._init_info.num_consumers))
            live_count = len(live_ids)
        else:
            live_ids = sorted(set(int(c) for c in live_consumer_ids))
            live_count = len(live_ids)
            logger.warning(
                "[rdt-degraded] serving %d/%d live consumers; every group's "
                "free barrier counts to the live total",
                live_count,
                self._init_info.num_consumers,
            )
        self._send_weights_inner(live_count, live_ids)

    def _send_weights_inner(self, live_count: int, live_ids: list[int]) -> None:
        if not self.is_sender:
            self._run_gather_loop(
                update_future=None, live_count=live_count, live_ids=live_ids
            )
            return

        self.client.start_weight_update()

        from vllm.distributed.weight_transfer.sharded_rdt_engine import (
            ShardedRDTWeightTransferUpdateInfo,
        )

        empty_update = asdict(ShardedRDTWeightTransferUpdateInfo())
        with ThreadPoolExecutor(max_workers=1) as exe:
            # The workers block inside update_weights until they've pulled every
            # group, so it runs concurrently with the gather/publish loop.
            future = exe.submit(self.client.update_weights, empty_update)
            self._run_gather_loop(
                update_future=future, live_count=live_count, live_ids=live_ids
            )
            future.result()  # surface inference-side errors

        self.client.finish_weight_update()

    def _run_gather_loop(
        self, update_future, live_count: int, live_ids: list[int] | None = None
    ) -> None:
        """Gather this rank's weights group-by-group and publish each into the
        server over CUDA IPC. A gathered group is published — serveable —
        immediately; the loop gates BEFORE the next gather while more than
        `gather_lookahead` groups are unfreed (the per-group free barrier: a
        credit releases when every live consumer has signaled the group). So
        the loop self-paces to the consumers' pull rate with at most
        `gather_lookahead + 1` groups resident. Runs on every rank; only the
        sender has an `update_future` to fail fast on."""
        assert self.source is not None  # guaranteed by trainer_init
        # The live IDS ride along with the count: the free barrier needs only
        # the count, but the slot-sharing rendezvous needs to know WHICH
        # consumers are live.
        self._rpc("begin_sync", live_count, live_ids)
        # One generator resume per GROUP: `iter_groups` yields (names, tensors)
        # per owned group in metadata order. Every owner must reach a group in the
        # same order or their shared gather collective mismatches.
        groups = self.source.iter_groups()
        # Publishes fire without an inline ray.get and are harvested this window
        # deep, so the RPC and server-side rebuild overlap the NEXT group's
        # gather/export. The window also surfaces a server-side error at depth 2
        # rather than at end_sync.
        _PUBLISH_WINDOW = 2
        pending_publish: list = []
        # The memory bound: gate BEFORE gathering while more
        # than `bound` groups are unfreed, so at most bound + 1 are resident. The
        # count is exact — `_inflight` gains its entry right after each gather and
        # only wait_freed/end_sync shrinks it, so a group freed server-side but not
        # yet collected still holds its refs and still counts. The publish window
        # drains first, so every resident group is pullable while we wait.
        #
        # 0 is legal and means NO pipelining: the loop waits for group N to be
        # freed by every live consumer before gathering N+1, so exactly one group
        # is resident and the sync costs the full serialized
        # gather + publish + pull + free RTT per group. 1 (the default) overlaps
        # N+1's gather with N's pulls.
        bound = max(0, self._init_info.gather_lookahead)

        # [RDT-EXPORT-RING] `bound + 1` slots is exactly the residency the credit
        # gate enforces. Small storages are packed into a slot and exported once
        # per slot instead of once per storage (518 -> 7 cudaIpcGetMemHandle
        # calls per rank per sync at 235B, each 0.5-1.5ms).
        nslots = bound + 1
        # data_ptr -> (storage weakref, reduce_args), this sync only. `storages`
        # is per-group, so a source that reuses a buffer across groups would
        # otherwise re-export it every group. The weakref is load-bearing: a
        # freed allocation can be reissued at the same address, and serving that
        # from a stale handle is silently wrong bytes.
        handle_cache: dict[int, tuple] = {}
        if len(getattr(self, "_export_ring", [])) != nslots:
            self._export_ring = [None] * nslots
            self._export_ring_args = [None] * nslots
        # Slots are handed out per group and returned by _drop_inflight, so the
        # pool starts full every sync.
        self._slot_of_group = {}
        self._free_slots = list(range(nslots))

        try:
            for gi in self._owned_idx:
                group = self._groups[gi]
                while len(self._inflight) > bound:
                    while pending_publish:
                        self._await_publish(pending_publish.pop(0))
                    self._drop_inflight(self._rpc("wait_freed"))
                names, tensors = next(groups)
                if list(names) != list(group):
                    raise RuntimeError(
                        f"WeightSource group yielded {len(names)} names starting "
                        f"{names[:2]!r} but expected {len(group)} starting "
                        f"{group[:2]!r}; iteration order must match metadata."
                    )
                self._validate_held_yields(gi, names, tensors)
                # Share each unique STORAGE once and describe every name as an
                # as_strided view onto it. ``None`` means a name this rank does not
                # hold (a foreign expert); the source keeps it in the list so the
                # order check stays rank-uniform, and it is dropped here, before
                # the IPC export, matching the sidecar's served_names.
                storages: dict[int, tuple] = {}
                views: dict[str, tuple] = {}
                refs: dict[str, torch.Tensor] = {}
                small: list = []
                for name, tensor in zip(names, tensors):
                    if tensor is None:
                        continue
                    tensor = tensor.detach()
                    if not tensor.is_cuda:
                        tensor = tensor.cuda()
                    # Threshold on the STORAGE, not the tensor: a source may hand
                    # out many small views of one large stacked buffer, and a
                    # per-tensor test would pack the whole thing, costing +1
                    # group of peak memory for a storage already coalesced.
                    if tensor.untyped_storage().nbytes() <= _EXPORT_RING_MAX_BYTES:
                        small.append((name, tensor))
                        continue
                    tensor = tensor.contiguous()
                    refs[name] = tensor  # keep the export alive
                    ust = tensor.untyped_storage()
                    sid = ust.data_ptr()
                    if sid not in storages:
                        cached = handle_cache.get(sid)
                        if cached is not None and not cached[0].expired():
                            reduce_args = cached[1]
                        else:
                            base = torch.empty(
                                0, dtype=torch.uint8, device=tensor.device
                            )
                            base.set_(ust, 0, (ust.nbytes(),))
                            _rebuild, reduce_args = reduce_tensor(base)
                            handle_cache[sid] = (StorageWeakRef(ust), reduce_args)
                        storages[sid] = reduce_args
                    views[name] = (
                        sid,
                        str(tensor.dtype).split(".")[-1],
                        list(tensor.shape),
                        list(tensor.stride()),
                        tensor.storage_offset(),
                    )
                if small:
                    # Take a slot no live group holds. The gate above already
                    # guarantees one exists, so this is a backstop -- but it
                    # must drain publishes first, like the credit gate.
                    while not self._free_slots:
                        while pending_publish:
                            self._await_publish(pending_publish.pop(0))
                        self._drop_inflight(self._rpc("wait_freed"))
                    slot_idx = self._free_slots.pop(0)
                    self._slot_of_group[gi] = slot_idx
                    st, vw, rf = self._pack_group_for_export(small, slot_idx)
                    storages.update(st)
                    views.update(vw)
                    refs.update(rf)
                del small
                del tensors
                # The loop leaves its last `tensor` / `ust` / `base` bound
                # through the credit gate and into the next gather, and `base`
                # is a whole-storage view -- so one tensor and its entire
                # allocation stay alive past `_drop_inflight`, unseen by the
                # `_inflight` accounting the residency bound rests on.
                # Assignment, not `del`: unbound if no direct-path tensor.
                tensor = ust = base = None
                if not refs:
                    # A group with nothing held here cannot occur (every group
                    # carries replicated names), but publishing an empty group
                    # would park a credit nobody's pull is waiting on; skip.
                    continue
                # Hold our refs before publishing; drop them only when the
                # server reports the group freed (IPC export must outlive import).
                self._inflight[gi] = refs
                pending_publish.append(self._publish_async(gi, (storages, views)))
                while len(pending_publish) >= _PUBLISH_WINDOW:
                    self._await_publish(pending_publish.pop(0))
                if update_future is not None and update_future.done():
                    # update_weights returned/failed early — surface now instead
                    # of blocking further gathers.
                    update_future.result()
            while pending_publish:
                self._await_publish(pending_publish.pop(0))
            freed = self._rpc("end_sync")
            self._drop_inflight(freed)
        except BaseException as e:
            with contextlib.suppress(Exception):
                self._rpc("set_gather_error", repr(e))
            self._inflight.clear()
            raise
        finally:
            # Keep the slots, drop the handles: the sidecar's base cache is
            # cleared per sync, so a handle must not outlive the sync either.
            self._export_ring_args = [None] * len(getattr(self, "_export_ring", []))

    def _drop_inflight(self, freed_keys: list) -> None:
        """Release a freed group's refs and return its export-ring slot."""
        free_slots = getattr(self, "_free_slots", None)
        slot_of_group = getattr(self, "_slot_of_group", None)
        for k in freed_keys:
            gi = int(k)
            self._inflight.pop(gi, None)
            if slot_of_group is None:
                continue
            slot = slot_of_group.pop(gi, None)
            if slot is not None and free_slots is not None and slot not in free_slots:
                free_slots.append(slot)

    # ---------------- misc ----------------

    def shutdown(self) -> None:
        if self._server is None:
            return
        import ray

        with contextlib.suppress(Exception):
            ray.get(self._server.shutdown.remote())
            ray.kill(self._server)
        self._server = None
        self._inflight.clear()