diff --git a/afd_plugin/connectors/gpu/p2p.py b/afd_plugin/connectors/gpu/p2p.py index 5534859d..892b483d 100644 --- a/afd_plugin/connectors/gpu/p2p.py +++ b/afd_plugin/connectors/gpu/p2p.py @@ -11,16 +11,21 @@ Topology: The connector creates one AFD NCCL world ordered as ``[F0, F1, ..., A0, A1, ...]``: FFN ranks first, followed by Attention - ranks. Each FFN rank owns one subgroup containing itself and one or more - consecutive Attention ranks, which requires:: - - num_attention_ranks >= num_ffn_ranks - num_attention_ranks % num_ffn_ranks == 0 - - Attention sends hidden states to its mapped FFN rank; the FFN rank - concatenates inputs from its Attention peers, runs FFN work, splits the - output by the recorded sequence lengths, and sends each slice back to the - originating Attention rank. + ranks. The world is partitioned into ``G = min(A, F)`` subgroups by a + balanced block distribution: Attention ``a`` joins subgroup ``a*G//A`` + and FFN ``f`` joins subgroup ``f*G//F``. Every positive ``(A, F)`` pair + is valid, and the role with fewer ranks contributes exactly one member + per subgroup: + + - ``A >= F``: one FFN rank and one or more Attention ranks. Attention + sends its hidden states whole; the FFN rank concatenates the inputs + from its Attention peers, runs FFN work, splits the output by the + recorded sequence lengths, and sends each slice back to its origin. + - ``A < F``: one Attention rank and several FFN ranks. Attention splits + its tokens across the FFN members with ``split_send_sizes``; each FFN + returns its share, and Attention reassembles them in member order. An + FFN that receives no tokens runs its forward on a one-token dummy + batch and returns a zero-size tensor. Control and data planes: DP metadata handling is a pluggable control plane, not part of the @@ -61,6 +66,7 @@ from __future__ import annotations +import socket from collections.abc import Callable, Mapping from datetime import timedelta from typing import TYPE_CHECKING, Any, NamedTuple @@ -92,6 +98,7 @@ DefaultProcessGroupSwitcher, build_rank_mapping, init_afd_process_group, + split_send_sizes, ) if TYPE_CHECKING: @@ -115,9 +122,11 @@ class P2pNcclAFDConnector(AFDConnectorBase): """NCCL-backed Attention <-> FFN connector for CUDA deployments. The P2P topology places FFN ranks before Attention ranks in the AFD world - (``[F0, F1, ..., A0, A1, ...]``), and each FFN rank owns a subgroup with - one or more consecutive Attention ranks. Within a subgroup, the FFN rank - is subgroup rank ``0`` and its Attention peers occupy ranks ``1..ratio``. + (``[F0, F1, ..., A0, A1, ...]``) and partitions it into ``min(A, F)`` + subgroups (see the module docstring). Within a subgroup the FFN members + occupy ranks ``0..subgroup_ffn_count-1`` and the Attention members + follow; historically that is FFN rank ``0`` plus Attention ranks + ``1..ratio``. Hidden states move through per-subgroup ``PyNcclCommunicator`` instances on the current CUDA stream; a separate NCCL process group distributes the @@ -181,6 +190,9 @@ def __init__( self.min_size = self.mapping.min_size self.ratio = self.mapping.ratio self.group_size = len(self.mapping.subgroup_ranks) + self.subgroup_ffn_count = sum( + 1 for rank in self.mapping.subgroup_ranks if rank < self.ffn_size + ) self.dst_list = list(self.mapping.dp_metadata_destinations) text_config = vllm_config.model_config.hf_text_config self.num_hidden_layers = (text_config.num_hidden_layers,) @@ -193,6 +205,8 @@ def __init__( tuple[int, int], _TensorMetadata, ] = {} + # Stages of this step whose incoming shares are all zero. + self._dummy_stages: set[int] = set() self._recv_attn_buffers: dict[ tuple[int, int, tuple[int, ...]], torch.Tensor, @@ -242,9 +256,10 @@ def init_afd_connector(self) -> None: 1. Joins the AFD world process group (FFN ranks first, then Attention ranks) rendezvoused at ``tcp://host:port``. 2. Creates this rank's subgroup ``StatelessProcessGroup`` on - ``port + subgroup_index + 1`` and two ``PyNcclCommunicator`` - instances over it (Attention-to-FFN and FFN-to-Attention), each - registered for use by the P2P custom ops. + ``port + subgroup_index + 1``, at the address of the subgroup's + first member, and two ``PyNcclCommunicator`` instances over it + (Attention-to-FFN and FFN-to-Attention), each registered for use + by the P2P custom ops. 3. On ranks that participate in the DP metadata control plane, joins the ``p2p`` process group that ``control_plane`` uses to distribute per-stage token counts. @@ -265,10 +280,20 @@ def init_afd_connector(self) -> None: timeout=timedelta(minutes=2), ) + # A subgroup's group is served by its first member, which can only + # bind its own node's address. + addresses = _gather_rank_addresses( + afd_pg, + host=self.afd_config.host, + port=self.afd_config.port, + world_size=self.ffn_size + self.attn_size, + ) + subgroup_host = addresses[self.mapping.subgroup_ranks[0]] + with DefaultProcessGroupSwitcher(_get_default_group(), afd_pg): base_port = self.afd_config.port self.a2e_group = StatelessProcessGroup.create( - host=self.afd_config.host, + host=subgroup_host, port=base_port + self.mapping.subgroup_index + 1, rank=self.mapping.rank_in_subgroup, world_size=len(self.mapping.subgroup_ranks), @@ -308,7 +333,11 @@ def send_attn_output( context: AFDTransferContext, **kwargs: Any, ) -> None: - """Send Attention hidden states to this rank's mapped FFN rank. + """Send Attention hidden states to this rank's subgroup FFN member(s). + + Each subgroup FFN member receives one contiguous share of this + rank's tokens, sized by ``split_send_sizes``. A single FFN member + receives the whole tensor. Args: hidden_states: CUDA tensor of shape ``(num_tokens, hidden_size)`` @@ -350,26 +379,27 @@ def send_attn_output( ) if input_ids.dtype != _INPUT_IDS_DTYPE: raise ValueError("P2P AFD input_ids must use torch.int32") - self._send_hidden_states( - hidden_states, - 0, - self.a2e_group, - self.a2e_comm_id, - ) - if router_logits is not None: - self._send_hidden_states( - router_logits, - 0, - self.a2e_group, - self.a2e_comm_id, - ) - if input_ids is not None: - self._send_hidden_states( - input_ids, - 0, - self.a2e_group, - self.a2e_comm_id, - ) + if self.subgroup_ffn_count == 1: + self._send_hidden_states(hidden_states, 0, self.a2e_group, self.a2e_comm_id) + if router_logits is not None: + self._send_hidden_states( + router_logits, 0, self.a2e_group, self.a2e_comm_id + ) + if input_ids is not None: + self._send_hidden_states(input_ids, 0, self.a2e_group, self.a2e_comm_id) + return + # Traced by torch.compile: pass the bounds to the op rather than + # slicing here, and send every share (see _register_p2p_custom_ops). + sizes = split_send_sizes(hidden_states.shape[0], self.subgroup_ffn_count) + start = 0 + for dst, size in enumerate(sizes): + # Wire order per destination, matching recv_attn_output. + self._send_slice(hidden_states, start, size, dst) + if router_logits is not None: + self._send_slice(router_logits, start, size, dst) + if input_ids is not None: + self._send_slice(input_ids, start, size, dst) + start = start + size def recv_ffn_output( self, @@ -377,13 +407,17 @@ def recv_ffn_output( ubatch_idx: int = 0, **kwargs: Any, ) -> torch.Tensor: - """Receive this rank's FFN output slice on the Attention side. + """Receive this rank's FFN output on the Attention side. + + Receives one share per subgroup FFN member, in the sizes and member + order ``send_attn_output`` used, and assembles them into an output + shaped like ``ref_tensor``, restoring the original token order. Args: - ref_tensor: Preallocated CUDA tensor to receive into. Used as a - stable buffer for CUDA graph capture and returned directly - when the subgroup has a single rank and no wire transfer - occurs. + ref_tensor: Preallocated CUDA tensor of shape + ``(num_tokens, hidden_size)``. A single FFN member receives + into it directly, keeping the allocation stable for CUDA + graph capture; several members use its shape and device. ubatch_idx: Stage/microbatch index. Defaults to ``0``. **kwargs: Unused; accepted for interface compatibility. @@ -394,17 +428,26 @@ def recv_ffn_output( RuntimeError: If the connector is not initialized, or no receive is performed for a single-rank subgroup. """ - output = self._recv_hidden_states( - 0, - self.e2a_group, - self.e2a_comm_id, - self.tensor_metadata_list[ubatch_idx], - ref_tensor=ref_tensor, - ) - if output is None: - raise RuntimeError( - "P2P recv_ffn_output requires ref_tensor when no receive is performed", + if self.subgroup_ffn_count == 1: + output = self._recv_hidden_states( + 0, + self.e2a_group, + self.e2a_comm_id, + self.tensor_metadata_list[ubatch_idx], + ref_tensor=ref_tensor, ) + if output is None: + raise RuntimeError( + "P2P recv_ffn_output requires ref_tensor when no receive " + "is performed", + ) + return output + sizes = split_send_sizes(ref_tensor.shape[0], self.subgroup_ffn_count) + output = torch.empty_like(ref_tensor) + start = 0 + for src, size in enumerate(sizes): + self._recv_slice(ref_tensor, output, start, size, src) + start = start + size return output def recv_attn_output( @@ -414,12 +457,14 @@ def recv_attn_output( ) -> AFDA2FTransferPayload: """Receive and concatenate Attention hidden states on the FFN rank. - Receives one tensor from every Attention peer in this rank's - subgroup (subgroup ranks ``1..ratio``), concatenates them along the - token dimension, and records each peer's sequence length in the - returned metadata so ``send_ffn_output`` can split the FFN output - back per peer. When CUDA graphs are enabled, receives reuse the - buffers preallocated by ``control_plane.update_state_from_dp_metadata``. + Receives one tensor from every Attention member of this rank's + subgroup, concatenates them along the token dimension, and records + each peer's sequence length in the returned metadata so + ``send_ffn_output`` can split the FFN output back per peer. + Zero-size shares add no rows; when every share is zero a one-token + dummy batch is returned instead. When CUDA graphs are enabled, + receives reuse the buffers preallocated by + ``control_plane.update_state_from_dp_metadata``. Args: ubatch_idx: Stage/microbatch index to receive. Defaults to ``0``. @@ -432,8 +477,7 @@ def recv_attn_output( lengths. Raises: - RuntimeError: If the connector is not initialized or the subgroup - has no Attention peers. + RuntimeError: If the connector is not initialized. """ routing_spec: AFDExpertRoutingSpec | None = kwargs.get("routing_spec") recv_input_ids: bool = kwargs.get("recv_input_ids", False) @@ -441,25 +485,27 @@ def recv_attn_output( router_logits_list: list[torch.Tensor] = [] input_ids_list: list[torch.Tensor] = [] - for src in range(1, self.group_size): + for src in range(self.subgroup_ffn_count, self.group_size): tensor_metadata = self._recv_attn_tensor_metadata_list.get( (ubatch_idx, src), self.tensor_metadata_list[ubatch_idx], ) + # The sender posts every share; a zero-size one adds no rows. + share_is_empty = tensor_metadata.size[0] == 0 ref_tensor = None if not self.vllm_config.model_config.enforce_eager: ref_tensor = self._recv_attn_buffers.get( (ubatch_idx, src, tuple(tensor_metadata.size)), ) - hidden_states_list.append( - self._recv_hidden_states( - src, - self.a2e_group, - self.a2e_comm_id, - tensor_metadata, - ref_tensor=ref_tensor, - ), + received = self._recv_hidden_states( + src, + self.a2e_group, + self.a2e_comm_id, + tensor_metadata, + ref_tensor=ref_tensor, ) + if not share_is_empty: + hidden_states_list.append(received) if routing_spec is not None: router_metadata = _TensorMetadata( device=tensor_metadata.device, @@ -471,14 +517,14 @@ def recv_attn_output( ], ), ) - router_logits_list.append( - self._recv_hidden_states( - src, - self.a2e_group, - self.a2e_comm_id, - router_metadata, - ), + received = self._recv_hidden_states( + src, + self.a2e_group, + self.a2e_comm_id, + router_metadata, ) + if not share_is_empty: + router_logits_list.append(received) if recv_input_ids: input_ids_metadata = _TensorMetadata( device=tensor_metadata.device, @@ -490,18 +536,52 @@ def recv_attn_output( input_ids_ref = self._recv_attn_input_ids_buffers.get( (ubatch_idx, src, tuple(input_ids_metadata.size)), ) - input_ids_list.append( - self._recv_hidden_states( - src, - self.a2e_group, - self.a2e_comm_id, - input_ids_metadata, - ref_tensor=input_ids_ref, - ), + received = self._recv_hidden_states( + src, + self.a2e_group, + self.a2e_comm_id, + input_ids_metadata, + ref_tensor=input_ids_ref, ) + if not share_is_empty: + input_ids_list.append(received) if not hidden_states_list: - raise RuntimeError("P2P FFN rank has no Attention peers") + # Nothing arrived: run the forward on a dummy batch, as vLLM does + # for an idle DP rank. AFDTransferMetadata rejects zero lengths. + stage_metadata = self.tensor_metadata_list[ubatch_idx] + return AFDA2FTransferPayload( + hidden_states=torch.zeros( + (1, self.hidden_size), + dtype=stage_metadata.dtype, + device=stage_metadata.device, + ), + context=AFDTransferContext( + metadata=AFDTransferMetadata.create_ffn_metadata( + layer_idx=0, + stage_idx=ubatch_idx, + seq_lens=[1], + ), + ), + router_logits=( + None + if routing_spec is None + else torch.zeros( + (1, routing_spec.router_logits_width), + dtype=routing_spec.router_logits_dtype, + device=stage_metadata.device, + ) + ), + input_ids=( + None + if not recv_input_ids + else torch.zeros( + (1,), + dtype=_INPUT_IDS_DTYPE, + device=stage_metadata.device, + ) + ), + ) hidden_states = ( torch.cat(hidden_states_list, dim=0) if len(hidden_states_list) > 1 @@ -544,11 +624,11 @@ def send_ffn_output( ) -> None: """Split the FFN output and send each slice back to its Attention peer. - With a one-to-one mapping (``ratio == 1``) the whole tensor is sent to - the single Attention peer. Otherwise the output is split along the - token dimension by ``metadata.seq_lens`` — falling back to an even - split when the recorded lengths do not cover every peer — and each - slice is sent to the Attention rank that originally produced it. + With a single Attention member the whole tensor is sent to it, or a + zero-size tensor when nothing arrived this step. Otherwise the output + is split along the token dimension by ``metadata.seq_lens``, falling + back to an even split when the recorded lengths do not cover every + peer, and each slice is sent to the Attention rank that produced it. Args: ffn_output: CUDA tensor of shape ``(num_tokens, hidden_size)`` @@ -566,6 +646,14 @@ def send_ffn_output( RuntimeError: If the connector is not initialized. """ metadata = context.metadata + if metadata.stage_idx in self._dummy_stages: + self._send_hidden_states( + ffn_output.narrow(0, 0, 0), + self.subgroup_ffn_count, + self.e2a_group, + self.e2a_comm_id, + ) + return if not torch.compiler.is_compiling() and not metadata.validate_tensor_shape( tuple(ffn_output.shape), ): @@ -573,7 +661,12 @@ def send_ffn_output( f"ffn_output shape {ffn_output.shape!r} does not match metadata", ) if self.ratio == 1: - self._send_hidden_states(ffn_output, 1, self.e2a_group, self.e2a_comm_id) + self._send_hidden_states( + ffn_output, + self.subgroup_ffn_count, + self.e2a_group, + self.e2a_comm_id, + ) return split_sizes = metadata.seq_lens @@ -589,7 +682,7 @@ def send_ffn_output( start = 0 for dst, token_count in zip( - range(1, self.group_size), + range(self.subgroup_ffn_count, self.group_size), split_sizes, strict=False, ): @@ -631,6 +724,40 @@ def _send_hidden_states( ) return + def _send_slice( + self, + base: torch.Tensor, + start: int, + size: int, + dst: int, + ) -> None: + """Send rows ``[start, start + size)`` of ``base`` to subgroup rank ``dst``.""" + if self.a2e_group is None or self.a2e_comm_id is None: + raise RuntimeError("P2P connector is not initialized") + if dst >= self.a2e_group.world_size: + raise ValueError(f"invalid P2P destination rank {dst}") + torch.ops.vllm.afd_p2p_send_slice(base, start, size, dst, self.a2e_comm_id) + + def _recv_slice( + self, + base: torch.Tensor, + out: torch.Tensor, + start: int, + size: int, + src: int, + ) -> None: + """Receive ``size`` rows from subgroup rank ``src`` into ``out``. + + ``base`` is only read, to order this receive after the sends that + declare it mutated; without that dependency the compiler hoisted the + receives above the sends and deadlocked. + """ + if self.e2a_group is None or self.e2a_comm_id is None: + raise RuntimeError("P2P connector is not initialized") + if src >= self.e2a_group.world_size: + raise ValueError(f"invalid P2P source rank {src}") + torch.ops.vllm.afd_p2p_recv_slice(base, out, start, size, src, self.e2a_comm_id) + def _recv_hidden_states( self, src: int, @@ -729,34 +856,30 @@ def update_state_from_dp_metadata( connector.is_warmup = payload.is_warmup connector.tensor_metadata_list = {} connector._recv_attn_tensor_metadata_list = {} + connector._dummy_stages = set() device = torch.device(f"cuda:{connector.local_rank}") dtype = connector.vllm_config.model_config.dtype for stage_idx, dp_metadata in payload.dp_metadata_list.items(): stage_idx = stage_idx if connector.afd_config.role == "ffn": peer_metadata: list[_TensorMetadata] = [] - for src_rank in range(1, connector.group_size): - if src_rank <= 0 or src_rank >= connector.group_size: - raise ValueError(f"invalid Attention subgroup rank {src_rank}") + ffn_member_count = connector.subgroup_ffn_count + for src_rank in range(ffn_member_count, connector.group_size): attention_rank = ( - connector.mapping.subgroup_index * connector.ratio - + src_rank - - 1 + connector.mapping.subgroup_ranks[src_rank] - connector.ffn_size ) - + token_count = _num_tokens_for_attention_rank( + dp_metadata, + attention_rank=attention_rank, + attention_size=connector.attn_size, + ) + share = split_send_sizes(token_count, ffn_member_count)[ + connector.mapping.rank_in_subgroup + ] tensor_metadata = _TensorMetadata( device, dtype, - torch.Size( - [ - _num_tokens_for_attention_rank( - dp_metadata, - attention_rank=attention_rank, - attention_size=connector.attn_size, - ), - connector.hidden_size, - ], - ), + torch.Size([share, connector.hidden_size]), ) connector._recv_attn_tensor_metadata_list[(stage_idx, src_rank)] = ( tensor_metadata @@ -765,6 +888,8 @@ def update_state_from_dp_metadata( num_tokens = sum( tensor_metadata.size[0] for tensor_metadata in peer_metadata ) + if num_tokens == 0: + connector._dummy_stages.add(stage_idx) else: num_tokens = _num_tokens_for_attention_rank( dp_metadata, @@ -902,6 +1027,58 @@ def _register_comm(communicator: PyNcclCommunicator) -> int: return comm_id +def _local_address_toward(host: str, port: int) -> str: + """Address of the local interface that routes to ``host:port``. + + Connecting a UDP socket sends nothing; it only makes the kernel pick the + outgoing interface, whose address is then read back. On a single host + this is ``host`` itself (for example ``127.0.0.1``). + """ + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect((host, port)) + return probe.getsockname()[0] + + +def _gather_rank_addresses( + group: ProcessGroup, + *, + host: str, + port: int, + world_size: int, +) -> list[str]: + """Every rank's routable address, indexed by AFD world rank.""" + gathered: list[str | None] = [None] * world_size + torch.distributed.all_gather_object( + gathered, + _local_address_toward(host, port), + group=group, + ) + return [address or host for address in gathered] + + +def _register_op_once( + op_name: str, + op_func: Callable[..., None], + mutates_args: list[str], + fake_impl: Callable[..., None], +) -> None: + """Register ``op_name`` under ``torch.ops.vllm`` unless it already exists. + + A prior import of this module (for example after an importlib reload) can + leave the op defined in torch's process-global registry while + ``_AFD_CUSTOM_OPS_REGISTERED`` is back to ``False``. Reuse the existing op + instead of re-defining it, which torch rejects. + """ + if hasattr(torch.ops.vllm, op_name): + return + direct_register_custom_op( + op_name=op_name, + op_func=op_func, + mutates_args=mutates_args, + fake_impl=fake_impl, + ) + + def _register_p2p_custom_ops() -> None: """Register the AFD P2P send and receive custom ops. @@ -909,6 +1086,11 @@ def _register_p2p_custom_ops() -> None: fake implementations keeps the transfers traceable by ``torch.compile`` and capturable by CUDA graphs. Registration happens once per process; ops already registered by another connector instance are reused. + + The ``_slice`` variants take a whole tensor plus ``(start, size)`` and + slice inside the op, because a view at a non-zero offset passed to a + mutating custom op reached the op with offset 0 under ``torch.compile``. + Only subgroups with several FFN members call them. """ global _AFD_CUSTOM_OPS_REGISTERED @@ -958,37 +1140,83 @@ def afd_p2p_recv_fake( ) -> None: pass - def register_one( - op_name: str, - op_func: Callable[..., None], - mutates_args: list[str], - fake_impl: Callable[..., None], + def afd_p2p_send_slice_impl( + base: torch.Tensor, + start: int, + size: int, + dst: int, + comm_id: int, ) -> None: - # A prior import of this module (for example after an importlib reload) - # can leave the op defined in torch's process-global registry while - # this module's _AFD_CUSTOM_OPS_REGISTERED flag is back to False. Reuse - # the existing op instead of re-defining it, which torch rejects. - if hasattr(torch.ops.vllm, op_name): - return - direct_register_custom_op( - op_name=op_name, - op_func=op_func, - mutates_args=mutates_args, - fake_impl=fake_impl, + communicator = _AFD_COMMUNICATORS.get(comm_id) + if communicator is None: + raise RuntimeError(f"AFD communicator id {comm_id} is not registered") + communicator.send( + base.narrow(0, start, size).contiguous(), + dst, + stream=torch.cuda.current_stream(base.device), + ) + return None + + def afd_p2p_send_slice_fake( + base: torch.Tensor, + start: int, + size: int, + dst: int, + comm_id: int, + ) -> None: + pass + + def afd_p2p_recv_slice_impl( + base: torch.Tensor, + out: torch.Tensor, + start: int, + size: int, + src: int, + comm_id: int, + ) -> None: + communicator = _AFD_COMMUNICATORS.get(comm_id) + if communicator is None: + raise RuntimeError(f"AFD communicator id {comm_id} is not registered") + communicator.recv( + out.narrow(0, start, size), + src, + stream=torch.cuda.current_stream(out.device), ) - register_one( + def afd_p2p_recv_slice_fake( + base: torch.Tensor, + out: torch.Tensor, + start: int, + size: int, + src: int, + comm_id: int, + ) -> None: + pass + + _register_op_once( op_name="afd_p2p_send", op_func=afd_p2p_send_impl, mutates_args=["tensor"], fake_impl=afd_p2p_send_fake, ) - register_one( + _register_op_once( op_name="afd_p2p_recv", op_func=afd_p2p_recv_impl, mutates_args=["out"], fake_impl=afd_p2p_recv_fake, ) + _register_op_once( + op_name="afd_p2p_send_slice", + op_func=afd_p2p_send_slice_impl, + mutates_args=["base"], + fake_impl=afd_p2p_send_slice_fake, + ) + _register_op_once( + op_name="afd_p2p_recv_slice", + op_func=afd_p2p_recv_slice_impl, + mutates_args=["out"], + fake_impl=afd_p2p_recv_slice_fake, + ) _AFD_CUSTOM_OPS_REGISTERED = True diff --git a/afd_plugin/distributed/__init__.py b/afd_plugin/distributed/__init__.py index 8769568e..a06634c6 100644 --- a/afd_plugin/distributed/__init__.py +++ b/afd_plugin/distributed/__init__.py @@ -6,6 +6,7 @@ AFDRankMapping, build_rank_mapping, resolve_role_rank, + split_send_sizes, topology_from_config, validate_p2p_topology, ) @@ -32,6 +33,7 @@ def __getattr__(name: str): "create_hccl_process_group_options", "init_afd_process_group", "resolve_role_rank", + "split_send_sizes", "topology_from_config", "validate_p2p_topology", ] diff --git a/afd_plugin/distributed/topology.py b/afd_plugin/distributed/topology.py index 55732b26..3bef6efa 100644 --- a/afd_plugin/distributed/topology.py +++ b/afd_plugin/distributed/topology.py @@ -18,8 +18,11 @@ class AFDRankMapping: """Rank mapping for the P2P connector. The P2P world always places FFN ranks first, followed by Attention ranks: - ``[F0, F1, ..., A0, A1, ...]``. Each FFN rank owns one subgroup containing - itself at subgroup rank 0 and one or more consecutive Attention ranks. + ``[F0, F1, ..., A0, A1, ...]``. The world is partitioned into + ``min(A, F)`` subgroups, each holding at least one rank of both roles; + ``subgroup_ranks`` lists a subgroup's FFN members first, then its + Attention members, in world order. ``ratio`` is the number of Attention + members in this rank's subgroup. """ role: str @@ -52,16 +55,10 @@ def topology_from_config(config: AFDConfig) -> tuple[int, int]: def validate_p2p_topology(config: AFDConfig) -> None: attention_size, ffn_size = topology_from_config(config) - if attention_size < ffn_size: + if attention_size <= 0 or ffn_size <= 0: raise ValueError( - "P2pNcclAFDConnector currently requires num_attention_ranks >= " - f"num_ffn_ranks, got {attention_size} < {ffn_size}", - ) - if attention_size % ffn_size != 0: - raise ValueError( - "P2pNcclAFDConnector currently requires num_attention_ranks to be a " - "multiple of num_ffn_ranks, got " - f"{attention_size} and {ffn_size}", + "P2P topology requires positive rank counts, got " + f"A={attention_size}, F={ffn_size}", ) @@ -112,6 +109,30 @@ def resolve_role_rank(vllm_config: VllmConfig, config: AFDConfig) -> int: return role_rank +def split_send_sizes(token_count: int, parts: int) -> tuple[int, ...]: + """Split one rank's token count into the sizes it sends to each peer. + + The sender slices its tensor with the whole tuple and each receiver + reads the element at its own position, so both sides reach identical + sizes without extra communication. The rule matches + ``torch.tensor_split``: sizes differ by at most one and the + ``token_count % parts`` leading shares carry the extra token. A share is + zero when ``token_count < parts``. + + Share ``i`` is the closed form ``(token_count - i + parts - 1) // parts`` + rather than a branch on the remainder, so that ``token_count`` may also + be a ``SymInt``: under ``torch.compile`` a Python branch would freeze at + its trace-time outcome and split every other token count wrongly. + """ + if parts <= 0: + raise ValueError(f"parts must be positive, got {parts}") + if token_count < 0: + raise ValueError(f"token_count must be >= 0, got {token_count}") + return tuple( + (token_count - position + parts - 1) // parts for position in range(parts) + ) + + def build_rank_mapping( config: AFDConfig, role_rank: int, @@ -130,7 +151,6 @@ def build_rank_mapping( f"(rank={role_rank}, size={attention_size})", ) world_rank = ffn_size + role_rank - subgroup_index = role_rank // (attention_size // ffn_size) elif config.role == "ffn": if role_rank >= ffn_size: raise ValueError( @@ -138,20 +158,33 @@ def build_rank_mapping( f"(rank={role_rank}, size={ffn_size})", ) world_rank = role_rank - subgroup_index = role_rank else: raise ValueError(f"unknown AFD role {config.role!r}") - ratio = attention_size // ffn_size + # Balanced block distribution: i * G // N maps N ranks onto G contiguous + # blocks differing in size by at most one, and is onto for both roles. min_size = min(ffn_size, attention_size) - ffn_ranks = list(range(ffn_size)) - attention_ranks = list(range(ffn_size, ffn_size + attention_size)) - subgroup_ranks = tuple( - [ffn_ranks[subgroup_index]] - + [attention_ranks[subgroup_index * ratio + offset] for offset in range(ratio)], - ) + if config.role == "attention": + subgroup_index = role_rank * min_size // attention_size + else: + subgroup_index = role_rank * min_size // ffn_size + subgroup_ffn_ranks = [ + ffn_rank + for ffn_rank in range(ffn_size) + if ffn_rank * min_size // ffn_size == subgroup_index + ] + subgroup_attention_ranks = [ + ffn_size + attention_rank + for attention_rank in range(attention_size) + if attention_rank * min_size // attention_size == subgroup_index + ] + subgroup_ranks = tuple(subgroup_ffn_ranks + subgroup_attention_ranks) rank_in_subgroup = subgroup_ranks.index(world_rank) - p2p_rank = role_rank + min_size if config.role == "attention" else role_rank + ratio = len(subgroup_attention_ranks) + + # FFN ranks take p2p_rank 0..F-1, so Attention starts at ffn_size, the + # offset the receiver side assumes. + p2p_rank = role_rank + ffn_size if config.role == "attention" else role_rank destinations: list[int] = [] if ffn_size <= world_rank < ffn_size + min_size: @@ -181,6 +214,7 @@ def build_rank_mapping( "AFDRankMapping", "build_rank_mapping", "resolve_role_rank", + "split_send_sizes", "topology_from_config", "validate_p2p_topology", ] diff --git a/afd_plugin/v1/worker/ffn_metadata.py b/afd_plugin/v1/worker/ffn_metadata.py index e2e15ea9..995c4314 100644 --- a/afd_plugin/v1/worker/ffn_metadata.py +++ b/afd_plugin/v1/worker/ffn_metadata.py @@ -4,6 +4,8 @@ from __future__ import annotations +from afd_plugin.distributed import split_send_sizes + def aggregate_ffn_token_counts( attention_counts: tuple[int, ...], @@ -12,17 +14,20 @@ def aggregate_ffn_token_counts( ffn_size: int, fallback: int = 1, ) -> tuple[int, ...]: - """Aggregate consecutive Attention-rank counts for each FFN rank. + """Token count each FFN rank receives for a stage, from Attention counts. - For example, ``4A2F`` counts ``(0, 4, 5, 6)`` become ``(5, 11)`` because - every zero-token Attention peer contributes one placeholder row. Missing - peers use the same per-peer fallback, so empty counts become ``(2, 2)``. + An FFN rank sums the shares its subgroup's Attention members send it, + splitting with ``split_send_sizes`` as the connector does. For example, + ``4A2F`` counts ``(0, 4, 5, 6)`` become ``(5, 11)`` because every + zero-token Attention peer contributes one placeholder row, and ``1A2F`` + count ``(5,)`` becomes ``(3, 2)``. No count is below one, since an FFN + rank that receives nothing runs a one-token dummy batch. Missing peers + use the same per-peer fallback. """ fallback_count = max(1, int(fallback)) - fallback_counts = tuple(fallback_count for _ in range(max(0, ffn_size))) - if ffn_size <= 0 or attention_size < ffn_size or attention_size % ffn_size != 0: - return fallback_counts + if ffn_size <= 0 or attention_size <= 0: + return tuple(fallback_count for _ in range(max(0, ffn_size))) expanded_counts = attention_counts if ( @@ -36,19 +41,28 @@ def aggregate_ffn_token_counts( for rank in range(attention_size) ) - group_size = attention_size // ffn_size - return tuple( - sum( - max(1, int(expanded_counts[attention_rank])) - if attention_rank < len(expanded_counts) - else fallback_count - for attention_rank in range( - ffn_rank * group_size, - (ffn_rank + 1) * group_size, - ) + def peer_count(attention_rank: int) -> int: + if attention_rank < len(expanded_counts): + return max(1, int(expanded_counts[attention_rank])) + return fallback_count + + group_count = min(attention_size, ffn_size) + ffn_counts = [] + for ffn_rank in range(ffn_size): + group = ffn_rank * group_count // ffn_size + ffn_members = [ + member + for member in range(ffn_size) + if member * group_count // ffn_size == group + ] + position = ffn_members.index(ffn_rank) + total = sum( + split_send_sizes(peer_count(attention_rank), len(ffn_members))[position] + for attention_rank in range(attention_size) + if attention_rank * group_count // attention_size == group ) - for ffn_rank in range(ffn_size) - ) + ffn_counts.append(max(1, total)) + return tuple(ffn_counts) def project_ffn_token_counts_to_dp( diff --git a/tests/unit/connectors/test_p2p_compiled_send.py b/tests/unit/connectors/test_p2p_compiled_send.py new file mode 100644 index 00000000..7526b5ae --- /dev/null +++ b/tests/unit/connectors/test_p2p_compiled_send.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""The Attention-side transfer traced once must be right for every token count. + +vLLM compiles the model forward once with a symbolic token dimension and +drops all guards (``TorchCompileWithNoGuardsWrapper``), then reuses that +trace for every later batch. ``send_attn_output`` and ``recv_ffn_output`` +run inside that trace, so their share arithmetic must not depend on +anything frozen at trace time: no Python-int sizes from the metadata and no +branch on the remainder (a ``divmod``-based split traced at 64 tokens would +keep sending (n//2, n//2) — two rows for five tokens). + +These tests compile the connector methods with the same options vLLM uses, +trace at 64 tokens, and then call the compiled code with other token counts +through custom ops that record what actually went on the wire. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + import torch +else: + torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.config import AFDConfig # noqa: E402 +from afd_plugin.connectors import ( # noqa: E402 + AFDConnectorFactory, + AFDControlPayload, + AFDDPMetadata, + AFDTransferContext, + AFDTransferMetadata, +) +from afd_plugin.distributed import split_send_sizes # noqa: E402 + +_HIDDEN = 8 +_SENT: list[tuple[int, int]] = [] +_RECEIVED: list[tuple[int, int]] = [] + + +# Custom ops with fake implementations stand in for the NCCL ops so the +# transfer survives tracing and the recorded sizes are runtime values. Like +# the real slice ops they take whole tensors plus (start, size). +@torch.library.custom_op("afd_test::record_send_slice", mutates_args=("base",)) +def _record_send_slice(base: torch.Tensor, start: int, size: int, dst: int) -> None: + _SENT.append((dst, int(size))) + + +@_record_send_slice.register_fake +def _(base, start, size, dst): + return None + + +@torch.library.custom_op("afd_test::record_recv_slice", mutates_args=("out",)) +def _record_recv_slice( + base: torch.Tensor, out: torch.Tensor, start: int, size: int, src: int +) -> None: + _RECEIVED.append((src, int(size))) + out.narrow(0, start, size).fill_(float(src)) + + +@_record_recv_slice.register_fake +def _(base, out, start, size, src): + return None + + +@torch.library.custom_op("afd_test::record_send", mutates_args=("tensor",)) +def _record_send(tensor: torch.Tensor, dst: int) -> None: + _SENT.append((dst, int(tensor.shape[0]))) + + +@_record_send.register_fake +def _(tensor, dst): + return None + + +def _attention_connector(attention, ffn): + text_config = SimpleNamespace(hidden_size=_HIDDEN, num_hidden_layers=2) + vllm_config = SimpleNamespace( + additional_config={}, + model_config=SimpleNamespace( + dtype=torch.float32, + enforce_eager=False, + hf_config=text_config, + hf_text_config=text_config, + ), + parallel_config=SimpleNamespace( + data_parallel_size=attention, + data_parallel_rank=0, + prefill_context_parallel_size=1, + tensor_parallel_size=1, + ), + ) + connector = AFDConnectorFactory.create_connector( + 0, + 0, + vllm_config, + AFDConfig( + role="attention", + connector="P2pNcclAFDConnector", + num_attention_ranks=attention, + num_ffn_ranks=ffn, + ), + ) + connector.control_plane.update_state_from_dp_metadata( + AFDControlPayload( + dp_metadata_list={0: AFDDPMetadata([64] * attention)}, + is_graph_capturing=False, + is_warmup=False, + ), + ) + connector._send_hidden_states = lambda hidden, dst, group, comm_id: ( + torch.ops.afd_test.record_send(hidden, dst) + ) + connector._send_slice = lambda base, start, size, dst: ( + torch.ops.afd_test.record_send_slice(base, start, size, dst) + ) + connector._recv_slice = lambda base, out, start, size, src: ( + torch.ops.afd_test.record_recv_slice(base, out, start, size, src) + ) + return connector + + +def _compile_like_vllm(fn): + return torch.compile( + fn, + fullgraph=True, + dynamic=False, + backend="aot_eager", + options={"guard_filter_fn": torch.compiler.skip_all_guards_unsafe}, + ) + + +def _context(seq_len): + return AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=seq_len, + ), + ) + + +@pytest.mark.parametrize("ffn", [2, 4]) +def test_compiled_send_splits_every_token_count_like_the_receivers_expect(ffn): + torch._dynamo.reset() + connector = _attention_connector(1, ffn) + context = _context(64) + compiled = _compile_like_vllm( + lambda hidden: connector.send_attn_output(hidden, context), + ) + traced = torch.zeros((64, _HIDDEN)) + torch._dynamo.mark_dynamic(traced, 0) + compiled(traced) + for tokens in (64, 5, 7, 1, 128, 0): + _SENT.clear() + compiled(torch.zeros((tokens, _HIDDEN))) + assert list(enumerate(split_send_sizes(tokens, ffn))) == _SENT, tokens + + +def test_compiled_receive_fills_each_share_in_member_order(): + torch._dynamo.reset() + connector = _attention_connector(1, 2) + compiled = _compile_like_vllm( + lambda buffer: connector.recv_ffn_output(ref_tensor=buffer, ubatch_idx=0), + ) + traced = torch.zeros((64, _HIDDEN)) + torch._dynamo.mark_dynamic(traced, 0) + compiled(traced) + for tokens in (5, 7, 1): + _RECEIVED.clear() + output = compiled(torch.full((tokens, _HIDDEN), -1.0)) + first, second = split_send_sizes(tokens, 2) + assert [(0, first), (1, second)] == _RECEIVED, tokens + assert output.shape == (tokens, _HIDDEN) + assert torch.equal(output[:first], torch.zeros((first, _HIDDEN))) + assert torch.equal(output[first:], torch.ones((second, _HIDDEN))) + + +def test_single_member_compiled_send_is_one_whole_send(): + torch._dynamo.reset() + connector = _attention_connector(2, 1) + context = _context(64) + compiled = _compile_like_vllm( + lambda hidden: connector.send_attn_output(hidden, context), + ) + traced = torch.zeros((64, _HIDDEN)) + torch._dynamo.mark_dynamic(traced, 0) + compiled(traced) + for tokens in (5, 1, 128): + _SENT.clear() + compiled(torch.zeros((tokens, _HIDDEN))) + assert [(0, tokens)] == _SENT diff --git a/tests/unit/connectors/test_p2p_connector.py b/tests/unit/connectors/test_p2p_connector.py index 13250390..dab0b73b 100644 --- a/tests/unit/connectors/test_p2p_connector.py +++ b/tests/unit/connectors/test_p2p_connector.py @@ -419,30 +419,33 @@ def test_p2p_tensor_metadata_clamps_idle_attention_rank_to_dummy_token(): assert attention_connector.tensor_metadata_list[0].size == torch.Size([1, 16]) +# The A >= F and divisibility constraints were removed with the unified +# subgroup partition; the shapes this test used to reject now parse cleanly, +# and only non-positive rank counts remain invalid. @pytest.mark.parametrize( - ("raw", "message"), - [ - ( - { - "connector": "P2pNcclAFDConnector", - "num_attention_ranks": 1, - "num_ffn_ranks": 2, - }, - "num_attention_ranks >= num_ffn_ranks", - ), - ( + ("attention", "ffn"), + [(1, 2), (3, 2)], +) +def test_p2p_topology_accepts_previously_constrained_shapes(attention, ffn): + config = afd_config_from_mapping( + { + "connector": "P2pNcclAFDConnector", + "num_attention_ranks": attention, + "num_ffn_ranks": ffn, + }, + ) + assert (config.num_attention_ranks, config.num_ffn_ranks) == (attention, ffn) + + +def test_p2p_topology_rejects_non_positive_rank_counts(): + with pytest.raises(ValueError, match="positive"): + afd_config_from_mapping( { "connector": "P2pNcclAFDConnector", - "num_attention_ranks": 3, + "num_attention_ranks": 0, "num_ffn_ranks": 2, }, - "multiple of num_ffn_ranks", - ), - ], -) -def test_p2p_topology_validation_errors_are_clear(raw, message): - with pytest.raises(ValueError, match=message): - afd_config_from_mapping(raw) + ) def test_p2p_module_exports_connector_class(): @@ -542,11 +545,16 @@ def direct_register_custom_op(**kwargs): assert [call["op_name"] for call in calls] == [ "afd_p2p_send", "afd_p2p_recv", + "afd_p2p_send_slice", + "afd_p2p_recv_slice", + ] + assert [call["mutates_args"] for call in calls] == [ + ["tensor"], + ["out"], + ["base"], + ["out"], ] - assert calls[0]["mutates_args"] == ["tensor"] - assert calls[1]["mutates_args"] == ["out"] - assert callable(calls[0]["fake_impl"]) - assert callable(calls[1]["fake_impl"]) + assert all(callable(call["fake_impl"]) for call in calls) def test_p2p_hidden_state_send_uses_registered_custom_op(monkeypatch): diff --git a/tests/unit/connectors/test_p2p_generalized_topology.py b/tests/unit/connectors/test_p2p_generalized_topology.py new file mode 100644 index 00000000..486a1729 --- /dev/null +++ b/tests/unit/connectors/test_p2p_generalized_topology.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Transport behavior of ``P2pNcclAFDConnector`` on generalized topologies. + +The subgroup partition is now defined for any positive (A, F). The +existing connector tests cover the historical A >= F layouts; this file +covers the new shapes — A < F (one attention splitting across several FFN +members) and non-divisible A > F (uneven subgroups) — plus the historical +4A2F layout run through the same checks as a regression anchor. + +All communication is mocked at the ``_send_hidden_states`` / +``_recv_hidden_states`` seam, the same one the existing tests use, so the +connector's own logic (who sends how many rows to whom, in which order) +runs unchanged on the CPU. The round-trip tests connect all ranks of a +topology through a fake wire (a dict of queues keyed by subgroup and +member ranks) so that sizes, addressing, and ordering are checked end to +end exactly as the real NCCL exchange would pair them. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + import torch +else: + torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.config import AFDConfig # noqa: E402 +from afd_plugin.connectors import ( # noqa: E402 + AFDConnectorFactory, + AFDControlPayload, + AFDDPMetadata, + AFDTransferContext, + AFDTransferMetadata, +) +from afd_plugin.connectors.gpu.p2p import _TensorMetadata # noqa: E402 + +_HIDDEN = 16 + + +def _fake_vllm_config(*, dp_size=1, dp_rank=0, enforce_eager=True): + text_config = SimpleNamespace(hidden_size=_HIDDEN, num_hidden_layers=2) + return SimpleNamespace( + additional_config={}, + model_config=SimpleNamespace( + dtype=torch.float32, + enforce_eager=enforce_eager, + hf_config=text_config, + hf_text_config=text_config, + ), + parallel_config=SimpleNamespace( + data_parallel_size=dp_size, + data_parallel_rank=dp_rank, + prefill_context_parallel_size=1, + tensor_parallel_size=1, + ), + ) + + +def _connector(role, role_rank, attention, ffn, **config_kwargs): + # The factory derives the role rank from the vLLM DP coordinates, so the + # fake config carries this connector's own DP size/rank. + role_size = attention if role == "attention" else ffn + return AFDConnectorFactory.create_connector( + role_rank, + 0, + _fake_vllm_config(dp_size=role_size, dp_rank=role_rank, **config_kwargs), + AFDConfig( + role=role, + connector="P2pNcclAFDConnector", + num_attention_ranks=attention, + num_ffn_ranks=ffn, + ), + ) + + +def _all_connectors(attention, ffn): + attentions = [_connector("attention", a, attention, ffn) for a in range(attention)] + ffns = [_connector("ffn", f, attention, ffn) for f in range(ffn)] + return attentions, ffns + + +def _payload(token_counts): + return AFDControlPayload( + dp_metadata_list={0: AFDDPMetadata(token_counts)}, + is_graph_capturing=False, + is_warmup=False, + ) + + +def _apply(connectors, token_counts): + for connector in connectors: + connector.control_plane.update_state_from_dp_metadata(_payload(token_counts)) + # The dummy batch allocates on the stage metadata's device; point + # it at the CPU so the tests can run without CUDA. + old = connector.tensor_metadata_list[0] + connector.tensor_metadata_list[0] = _TensorMetadata( + torch.device("cpu"), + old.dtype, + old.size, + ) + + +def _attention_context(seq_len): + return AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=seq_len, + ), + ) + + +def _ffn_context(seq_lens): + return AFDTransferContext( + metadata=AFDTransferMetadata.create_ffn_metadata( + layer_idx=0, + stage_idx=0, + seq_lens=seq_lens, + ), + ) + + +def _record_sends(monkeypatch, connector): + """Replace both send seams with a recorder of (rows, dst, tensor). + + Single-member subgroups send whole tensors through ``_send_hidden_states``; + multi-member subgroups send (base, start, size) slices through + ``_send_slice`` — recorded as the sliced rows. + """ + sent = [] + monkeypatch.setattr( + connector, + "_send_hidden_states", + lambda hidden, dst, group, comm_id: sent.append((hidden.shape[0], dst, hidden)), + ) + monkeypatch.setattr( + connector, + "_send_slice", + lambda base, start, size, dst: sent.append( + (size, dst, base.narrow(0, start, size)) + ), + ) + return sent + + +def _fake_recvs(monkeypatch, connector, fill=None): + """Replace the receive seam; returns the (src, rows) call log. + + Received tensors are filled with ``fill(src)`` (default: the source + rank) so callers can tell which source each row came from. + """ + received = [] + + def fake_recv(src, group, comm_id, tensor_metadata, *, ref_tensor=None): + received.append((src, tensor_metadata.size[0])) + value = float(src) if fill is None else fill(src) + tensor = torch.full( + tuple(tensor_metadata.size), + value, + dtype=tensor_metadata.dtype, + ) + if ref_tensor is not None: + # The real op receives in place into the caller's buffer. + ref_tensor.copy_(tensor) + return ref_tensor + return tensor + + def fake_recv_slice(base, out, start, size, src): + received.append((src, size)) + value = float(src) if fill is None else fill(src) + out.narrow(0, start, size).fill_(value) + + monkeypatch.setattr(connector, "_recv_hidden_states", fake_recv) + monkeypatch.setattr(connector, "_recv_slice", fake_recv_slice) + return received + + +# -------------------------------------------------------------------------- +# 1. Construction: every topology is accepted in eager and graph mode alike +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("role", ["attention", "ffn"]) +@pytest.mark.parametrize("enforce_eager", [True, False]) +def test_multi_ffn_subgroup_is_accepted_in_both_modes(role, enforce_eager): + connector = _connector(role, 0, 1, 2, enforce_eager=enforce_eager) + assert connector.subgroup_ffn_count == 2 + + +def test_registration_defines_the_four_transfer_ops(): + """Every topology registers the same ops, so the schemas are pinned here. + + ``base`` is declared mutated on the send and only read by the receive, + which is what keeps receives ordered after sends in a compiled graph. + """ + from afd_plugin.connectors.gpu import p2p + + p2p._register_p2p_custom_ops() + schemas = { + name: str(getattr(torch.ops.vllm, name).default._schema) + for name in ( + "afd_p2p_send", + "afd_p2p_recv", + "afd_p2p_send_slice", + "afd_p2p_recv_slice", + ) + } + assert schemas["afd_p2p_send"] == ( + "vllm::afd_p2p_send(Tensor(a0!) tensor, SymInt dst, SymInt comm_id) -> ()" + ) + assert schemas["afd_p2p_recv"] == ( + "vllm::afd_p2p_recv(Tensor(a0!) out, SymInt src, SymInt comm_id) -> ()" + ) + assert schemas["afd_p2p_send_slice"] == ( + "vllm::afd_p2p_send_slice(Tensor(a0!) base, SymInt start, SymInt size, " + "SymInt dst, SymInt comm_id) -> ()" + ) + assert schemas["afd_p2p_recv_slice"] == ( + "vllm::afd_p2p_recv_slice(Tensor base, Tensor(a1!) out, SymInt start, " + "SymInt size, SymInt src, SymInt comm_id) -> ()" + ) + + +def test_historical_layout_still_allows_graph_mode(): + connector = _connector("attention", 1, 4, 2, enforce_eager=False) + assert connector.subgroup_ffn_count == 1 + + +# -------------------------------------------------------------------------- +# 2. Size pre-computation (update_state_from_dp_metadata) +# -------------------------------------------------------------------------- + + +def test_sizes_1a2f_split_across_two_ffn_members(): + # Group (F0, F1, A0): A0 sits at group rank 2 and splits 5 tokens 3/2. + attentions, ffns = _all_connectors(1, 2) + _apply(attentions + ffns, [5]) + assert attentions[0].tensor_metadata_list[0].size == torch.Size([5, _HIDDEN]) + expected = ffns[0]._recv_attn_tensor_metadata_list[(0, 2)].size + assert expected == torch.Size([3, _HIDDEN]) + expected = ffns[1]._recv_attn_tensor_metadata_list[(0, 2)].size + assert expected == torch.Size([2, _HIDDEN]) + assert ffns[0].tensor_metadata_list[0].size == torch.Size([3, _HIDDEN]) + assert ffns[1].tensor_metadata_list[0].size == torch.Size([2, _HIDDEN]) + assert all(not ffn._dummy_stages for ffn in ffns) + + +def test_sizes_1a4f_with_fewer_tokens_than_members_marks_dummy_stages(): + # 2 tokens over 4 FFN members: shares (1, 1, 0, 0). F2 and F3 receive + # nothing this step and are marked for the dummy batch. + attentions, ffns = _all_connectors(1, 4) + _apply(attentions + ffns, [2]) + shares = [ffn._recv_attn_tensor_metadata_list[(0, 4)].size[0] for ffn in ffns] + assert shares == [1, 1, 0, 0] + assert [sorted(ffn._dummy_stages) for ffn in ffns] == [[], [], [0], [0]] + + +def test_sizes_6a4f_uneven_subgroups(): + # Subgroups: {F0,A0,A1} {F1,A2} {F2,A3,A4} {F3,A5}; token count = rank+1. + attentions, ffns = _all_connectors(6, 4) + _apply(attentions + ffns, [1, 2, 3, 4, 5, 6]) + + def peers(ffn): + return sorted( + (src, md.size[0]) + for (_, src), md in ffn._recv_attn_tensor_metadata_list.items() + ) + + assert peers(ffns[0]) == [(1, 1), (2, 2)] + assert peers(ffns[1]) == [(1, 3)] + assert peers(ffns[2]) == [(1, 4), (2, 5)] + assert peers(ffns[3]) == [(1, 6)] + assert [ffn.tensor_metadata_list[0].size[0] for ffn in ffns] == [3, 3, 9, 6] + + +def test_sizes_4a2f_match_historical_layout(): + attentions, ffns = _all_connectors(4, 2) + _apply(attentions + ffns, [3, 5, 7, 11]) + assert ffns[0]._recv_attn_tensor_metadata_list[(0, 1)].size[0] == 3 + assert ffns[0]._recv_attn_tensor_metadata_list[(0, 2)].size[0] == 5 + assert ffns[1]._recv_attn_tensor_metadata_list[(0, 1)].size[0] == 7 + assert ffns[1]._recv_attn_tensor_metadata_list[(0, 2)].size[0] == 11 + + +# -------------------------------------------------------------------------- +# 3. Attention send +# -------------------------------------------------------------------------- + + +def test_attention_sends_slices_to_ffn_members_in_order(monkeypatch): + attention = _connector("attention", 0, 1, 2) + _apply([attention], [5]) + sent = _record_sends(monkeypatch, attention) + hidden = torch.arange(5 * _HIDDEN, dtype=torch.float32).reshape(5, _HIDDEN) + attention.send_attn_output(hidden, _attention_context(5)) + assert [(rows, dst) for rows, dst, _ in sent] == [(3, 0), (2, 1)] + assert torch.equal(sent[0][2], hidden[:3]) + assert torch.equal(sent[1][2], hidden[3:]) + + +def test_attention_sends_every_share_including_zero_size_ones(monkeypatch): + # 2 tokens over 4 members: shares (1, 1, 0, 0). The zero-size shares are + # sent too — the send path has no branch on the token count so that it + # traces correctly under torch.compile (see test_p2p_compiled_send). + attention = _connector("attention", 0, 1, 4) + _apply([attention], [2]) + sent = _record_sends(monkeypatch, attention) + attention.send_attn_output(torch.zeros((2, _HIDDEN)), _attention_context(2)) + assert [(rows, dst) for rows, dst, _ in sent] == [(1, 0), (1, 1), (0, 2), (0, 3)] + + +def test_attention_with_single_ffn_member_sends_the_whole_tensor(monkeypatch): + # One FFN member: the upstream path — a single send of the tensor + # object itself (no slicing, no view). + attention = _connector("attention", 1, 4, 2) + _apply([attention], [3, 5, 7, 11]) + sent = _record_sends(monkeypatch, attention) + hidden = torch.zeros((5, _HIDDEN)) + attention.send_attn_output(hidden, _attention_context(5)) + assert len(sent) == 1 + assert sent[0][1] == 0 + assert sent[0][2] is hidden + + +# -------------------------------------------------------------------------- +# 4. FFN receive (incl. the dummy batch) +# -------------------------------------------------------------------------- + + +def test_ffn_receives_its_share_from_the_attention_member(monkeypatch): + ffn1 = _connector("ffn", 1, 1, 2) + _apply([ffn1], [5]) + received = _fake_recvs(monkeypatch, ffn1) + payload = ffn1.recv_attn_output(ubatch_idx=0) + assert received == [(2, 2)] # attention member at group rank 2, 2 rows + assert payload.hidden_states.shape == (2, _HIDDEN) + assert payload.context.metadata.seq_lens == [2] + + +def test_ffn_with_nothing_to_receive_runs_on_a_dummy_batch(monkeypatch): + ffn2 = _connector("ffn", 2, 1, 4) + _apply([ffn2], [2]) + received = _fake_recvs(monkeypatch, ffn2) + payload = ffn2.recv_attn_output(ubatch_idx=0) + # The zero-size share is still received (the attention member sent it). + assert received == [(4, 0)] + assert payload.hidden_states.shape == (1, _HIDDEN) + assert torch.equal(payload.hidden_states, torch.zeros((1, _HIDDEN))) + assert payload.context.metadata.seq_lens == [1] + assert 0 in ffn2._dummy_stages + + +# -------------------------------------------------------------------------- +# 5. FFN return send +# -------------------------------------------------------------------------- + + +def test_ffn_returns_whole_output_to_its_attention_member(monkeypatch): + ffn1 = _connector("ffn", 1, 1, 2) + _apply([ffn1], [5]) + sent = _record_sends(monkeypatch, ffn1) + ffn1.send_ffn_output(torch.zeros((2, _HIDDEN)), _ffn_context([2])) + assert [(rows, dst) for rows, dst, _ in sent] == [(2, 2)] + + +def test_ffn_on_a_dummy_stage_returns_a_zero_size_tensor(monkeypatch): + # The dummy result is discarded; a zero-size tensor goes back to the + # attention member (group rank 4) to match its posted receive. + ffn2 = _connector("ffn", 2, 1, 4) + _apply([ffn2], [2]) + _fake_recvs(monkeypatch, ffn2) + payload = ffn2.recv_attn_output(ubatch_idx=0) + sent = _record_sends(monkeypatch, ffn2) + ffn2.send_ffn_output(payload.hidden_states, payload.context) + assert [(rows, dst) for rows, dst, _ in sent] == [(0, 4)] + + +def test_ffn_with_several_attention_members_splits_by_seq_lens(monkeypatch): + ffn0 = _connector("ffn", 0, 6, 4) # subgroup {F0, A0, A1} + _apply([ffn0], [1, 2, 3, 4, 5, 6]) + sent = _record_sends(monkeypatch, ffn0) + ffn0.send_ffn_output(torch.zeros((3, _HIDDEN)), _ffn_context([1, 2])) + assert [(rows, dst) for rows, dst, _ in sent] == [(1, 1), (2, 2)] + + +# -------------------------------------------------------------------------- +# 6. Attention return receive +# -------------------------------------------------------------------------- + + +def test_attention_reassembles_returns_in_member_order(monkeypatch): + attention = _connector("attention", 0, 1, 2) + _apply([attention], [5]) + received = _fake_recvs(monkeypatch, attention) + buffer = torch.full((5, _HIDDEN), -1.0) + output = attention.recv_ffn_output(ref_tensor=buffer, ubatch_idx=0) + assert received == [(0, 3), (1, 2)] + assert output.shape == (5, _HIDDEN) + assert torch.equal(output[:3], torch.zeros((3, _HIDDEN))) + assert torch.equal(output[3:], torch.ones((2, _HIDDEN))) + + +# -------------------------------------------------------------------------- +# 7 + 8. Round trip over a fake wire +# -------------------------------------------------------------------------- + + +def _wire(monkeypatch, connectors): + """Connect every connector through queues keyed by subgroup and ranks. + + A send from member ``s`` to member ``d`` of subgroup ``g`` lands on + ``wire[(g, s, d)]``; a receive on member ``d`` from ``s`` pops the same + queue and checks the tensor has the size the receiver expected. + """ + wire: dict[tuple[int, int, int], list[torch.Tensor]] = {} + for connector in connectors: + group = connector.mapping.subgroup_index + me = connector.mapping.rank_in_subgroup + + def fake_send(hidden, dst, process_group, comm_id, *, group=group, me=me): + wire.setdefault((group, me, dst), []).append(hidden.clone()) + + def fake_recv( + src, + process_group, + comm_id, + tensor_metadata, + *, + ref_tensor=None, + group=group, + me=me, + ): + queue = wire.get((group, src, me)) + assert queue, f"nothing on the wire from {src} to {me} in subgroup {group}" + tensor = queue.pop(0) + assert tuple(tensor.shape) == tuple(tensor_metadata.size) + if ref_tensor is not None: + ref_tensor.copy_(tensor) + return ref_tensor + return tensor + + def fake_send_slice(base, start, size, dst, *, group=group, me=me): + wire.setdefault((group, me, dst), []).append( + base.narrow(0, start, size).clone() + ) + + def fake_recv_slice(base, out, start, size, src, *, group=group, me=me): + queue = wire.get((group, src, me)) + assert queue, f"nothing on the wire from {src} to {me} in subgroup {group}" + tensor = queue.pop(0) + assert tuple(tensor.shape) == (size, out.shape[1]) + out.narrow(0, start, size).copy_(tensor) + + monkeypatch.setattr(connector, "_send_hidden_states", fake_send) + monkeypatch.setattr(connector, "_recv_hidden_states", fake_recv) + monkeypatch.setattr(connector, "_send_slice", fake_send_slice) + monkeypatch.setattr(connector, "_recv_slice", fake_recv_slice) + return wire + + +@pytest.mark.parametrize( + ("attention", "ffn", "tokens"), + [ + (1, 2, [5]), # A < F: one attention splits 3/2 + (2, 3, [5, 3]), # A < F: one subgroup splits, one does not + (6, 4, [1, 2, 3, 4, 5, 6]), # non-divisible A > F: uneven subgroups + (4, 2, [3, 5, 7, 11]), # historical divisible layout (regression) + (1, 4, [2]), # fewer tokens than members: dummy stages + ], +) +def test_round_trip_preserves_every_token(monkeypatch, attention, ffn, tokens): + attentions, ffns = _all_connectors(attention, ffn) + _apply(attentions + ffns, tokens) + wire = _wire(monkeypatch, attentions + ffns) + + # Distinct values per attention rank and per token so any misrouting or + # reordering shows up in the final comparison. + inputs = [ + torch.arange(t * _HIDDEN, dtype=torch.float32).reshape(t, _HIDDEN) + 1000 * a + for a, t in enumerate(tokens) + ] + for connector, hidden in zip(attentions, inputs, strict=True): + connector.send_attn_output(hidden, _attention_context(hidden.shape[0])) + + # Every FFN rank receives what arrived (or a dummy batch), "computes" + # by adding one, and returns the result. + for connector in ffns: + payload = connector.recv_attn_output(ubatch_idx=0) + connector.send_ffn_output(payload.hidden_states + 1, payload.context) + + for connector, hidden in zip(attentions, inputs, strict=True): + output = connector.recv_ffn_output( + ref_tensor=torch.empty_like(hidden), + ubatch_idx=0, + ) + assert torch.equal(output, hidden + 1) + + # Nothing sent but unreceived, nothing expected but unsent. + assert all(not queue for queue in wire.values()) + + +def test_round_trip_zero_shares_travel_as_zero_size_tensors(monkeypatch): + # 1A4F with 2 tokens: F2 and F3 get zero-size shares. Both directions + # still post the transfer (as zero-row tensors), so the pairs appear on + # the wire, the dummy members return zero rows, and nothing is left over. + attentions, ffns = _all_connectors(1, 4) + _apply(attentions + ffns, [2]) + wire = _wire(monkeypatch, attentions + ffns) + hidden = torch.arange(2 * _HIDDEN, dtype=torch.float32).reshape(2, _HIDDEN) + attentions[0].send_attn_output(hidden, _attention_context(2)) + assert [t.shape[0] for t in wire[(0, 4, 2)]] == [0] + assert [t.shape[0] for t in wire[(0, 4, 3)]] == [0] + for connector in ffns: + payload = connector.recv_attn_output(ubatch_idx=0) + connector.send_ffn_output(payload.hidden_states + 1, payload.context) + assert [t.shape[0] for t in wire[(0, 2, 4)]] == [0] + assert [t.shape[0] for t in wire[(0, 3, 4)]] == [0] + output = attentions[0].recv_ffn_output( + ref_tensor=torch.empty_like(hidden), + ubatch_idx=0, + ) + assert torch.equal(output, hidden + 1) + assert all(not queue for queue in wire.values()) diff --git a/tests/unit/connectors/test_p2p_multinode_address.py b/tests/unit/connectors/test_p2p_multinode_address.py new file mode 100644 index 00000000..3aca84c1 --- /dev/null +++ b/tests/unit/connectors/test_p2p_multinode_address.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Subgroup rendezvous address selection for multi-node P2P deployments. + +Each subgroup's StatelessProcessGroup is served by its first member binding +host:port, so members on other nodes must dial that member's own address. +The connector derives every rank's address from the interface that routes +to the configured rendezvous host. On a single host that address equals the +configured one, which keeps the historical single-host layouts unchanged. +""" + +from __future__ import annotations + +import socket + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.connectors.gpu.p2p import _local_address_toward # noqa: E402 + + +def test_loopback_host_yields_loopback_address(): + # The recipes use host=127.0.0.1: every rank must keep binding/dialing + # loopback exactly as before the multi-node change. + assert _local_address_toward("127.0.0.1", 6239) == "127.0.0.1" + + +def test_own_lan_address_yields_itself(): + # host = this machine's own routable address (single host, LAN IP in the + # config): the probe must return that same address. + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect(("192.0.2.1", 9)) # TEST-NET-3, nothing is sent + own = probe.getsockname()[0] + if own.startswith("127."): + pytest.skip("no routable interface on this machine") + assert _local_address_toward(own, 6239) == own + + +def test_probe_sends_nothing_and_needs_no_listener(): + # A closed port on loopback still resolves (UDP connect only sets the + # route), so init does not depend on the master being up yet. + assert _local_address_toward("127.0.0.1", 1) == "127.0.0.1" diff --git a/tests/unit/distributed/test_topology_partition.py b/tests/unit/distributed/test_topology_partition.py new file mode 100644 index 00000000..8b80eb38 --- /dev/null +++ b/tests/unit/distributed/test_topology_partition.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Unified subgroup partition tests for arbitrary positive ``(A, F)``. + +``build_rank_mapping`` partitions the world into ``min(A, F)`` subgroups +(attention ``a`` -> subgroup ``a*G//A``, FFN ``f`` -> subgroup ``f*G//F``). +The historical ``A >= F`` divisible layout is the ``G == F`` special case, +pinned byte-for-byte by ``test_topology_snapshot.py``; this file covers what +the snapshots cannot: the structural properties that must hold on every +positive pair. + +Also guards the DP metadata channel numbering: attention ``p2p_rank`` uses +the ``ffn_size`` offset that the receiver side assumes (``p2p.py``: +``src = p2p_rank % min_size + ffn_size``); under ``A >= F`` it equals the +historical ``min_size`` offset, and under ``A < F`` it is what keeps the +numbering collision-free. +""" + +from __future__ import annotations + +import pytest + +from afd_plugin.config import AFDConfig +from afd_plugin.distributed.topology import ( + build_rank_mapping, + split_send_sizes, + validate_p2p_topology, +) + +# Every positive pair with A, F in 1..6 — divisible, non-divisible, and +# A < F shapes (36 combinations). +_FULL_GRID = [(a, f) for a in range(1, 7) for f in range(1, 7)] + + +def _config(role: str, attention: int, ffn: int) -> AFDConfig: + return AFDConfig( + role=role, + num_attention_ranks=attention, + num_ffn_ranks=ffn, + ) + + +def _mapping(role, attention, ffn, role_rank): + return build_rank_mapping(_config(role, attention, ffn), role_rank) + + +def _all_mappings(attention, ffn): + for role, size in (("ffn", ffn), ("attention", attention)): + for role_rank in range(size): + yield _mapping(role, attention, ffn, role_rank) + + +def _receiver_expected_source_rank(ffn_rank, attention, ffn): + """The p2p_rank the FFN receiver polls for its metadata sender. + + Mirrors the receiver convention in ``connectors/gpu/p2p.py`` + (``recv_dp_metadata_list``): ``src = p2p_rank % min_size + ffn_size`` + with ``p2p_rank == ffn_rank`` on FFN ranks. Kept in sync by review; the + connector cannot be imported here without torch/vLLM. + """ + return ffn_rank % min(attention, ffn) + ffn + + +# -------------------------------------------------------------------------- +# Validation — positive rank counts only +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("attention", "ffn"), [(1, 2), (3, 2), (2, 4), (6, 4)]) +def test_validation_accepts_any_positive_pair(attention, ffn): + # The historical A >= F and divisibility constraints are removed; the + # previously rejected shapes validate cleanly. + validate_p2p_topology(_config("attention", attention, ffn)) + + +@pytest.mark.parametrize(("attention", "ffn"), [(0, 2), (2, 0)]) +def test_validation_rejects_non_positive_pair(attention, ffn): + with pytest.raises(ValueError, match="positive rank counts"): + validate_p2p_topology(_config("attention", attention, ffn)) + + +# -------------------------------------------------------------------------- +# Partition structure on the full grid +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("attention", "ffn"), _FULL_GRID) +def test_partition_properties(attention, ffn): + group_count = min(attention, ffn) + groups: dict[int, tuple[int, ...]] = {} + + for mapping in _all_mappings(attention, ffn): + # Index formulas and membership consistency: every member computes + # the identical member tuple, and sits at its claimed position. + if mapping.role == "ffn": + assert mapping.subgroup_index == mapping.role_rank * group_count // ffn + else: + assert ( + mapping.subgroup_index == mapping.role_rank * group_count // attention + ) + prior = groups.setdefault(mapping.subgroup_index, mapping.subgroup_ranks) + assert prior == mapping.subgroup_ranks + assert mapping.subgroup_ranks[mapping.rank_in_subgroup] == mapping.world_rank + # ratio is this subgroup's attention member count. + attention_members = [rank for rank in mapping.subgroup_ranks if rank >= ffn] + assert mapping.ratio == len(attention_members) + + # Exactly min(A, F) subgroups, each holding at least one rank of both + # roles, members FFN-first in ascending world order, and together they + # partition the world with no overlap. + assert sorted(groups) == list(range(group_count)) + seen: list[int] = [] + for members in groups.values(): + ffn_members = [rank for rank in members if rank < ffn] + attention_members = [rank for rank in members if rank >= ffn] + assert ffn_members and attention_members + assert list(members) == sorted(ffn_members) + sorted(attention_members) + seen.extend(members) + assert sorted(seen) == list(range(ffn + attention)) + + +def test_example_2a3f_partition(): + # The shape that motivated the unified partition: 2 attention, 3 FFN + # (world order [F0, F1, F2, A0, A1]) splits into {F0, F1, A0} and + # {F2, A1}. + members = { + mapping.subgroup_index: mapping.subgroup_ranks + for mapping in _all_mappings(2, 3) + } + assert members == {0: (0, 1, 3), 1: (2, 4)} + + +@pytest.mark.parametrize( + ("attention", "ffn", "expected"), + [ + (4, 2, {0: (0, 2, 3), 1: (1, 4, 5)}), # historical divisible layout + (3, 2, {0: (0, 2, 3), 1: (1, 4)}), # non-divisible A > F + (1, 2, {0: (0, 1, 2)}), # capacity-bound target shape + ], +) +def test_partition_literal_examples(attention, ffn, expected): + members = { + mapping.subgroup_index: mapping.subgroup_ranks + for mapping in _all_mappings(attention, ffn) + } + assert members == expected + + +# -------------------------------------------------------------------------- +# DP metadata channel numbering +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("attention", "ffn"), _FULL_GRID) +def test_metadata_p2p_ranks_are_unique_and_contiguous(attention, ffn): + ranks = sorted( + mapping.p2p_rank + for mapping in _all_mappings(attention, ffn) + if mapping.participates_in_dp_metadata_group + ) + assert ranks == list(range(ffn)) + [ + ffn + attention_rank for attention_rank in range(min(attention, ffn)) + ] + + +@pytest.mark.parametrize(("attention", "ffn"), _FULL_GRID) +def test_sender_p2p_rank_matches_receiver_formula(attention, ffn): + min_size = min(attention, ffn) + for ffn_rank in range(ffn): + representative = ffn_rank % min_size + sender = _mapping("attention", attention, ffn, representative) + assert sender.p2p_rank == _receiver_expected_source_rank( + ffn_rank, + attention, + ffn, + ) + + +@pytest.mark.parametrize(("attention", "ffn"), _FULL_GRID) +def test_dp_metadata_destinations_cover_every_ffn_once(attention, ffn): + destinations: list[int] = [] + for role_rank in range(attention): + mapping = _mapping("attention", attention, ffn, role_rank) + destinations.extend(mapping.dp_metadata_destinations) + assert sorted(destinations) == list(range(ffn)) + + +# -------------------------------------------------------------------------- +# split_send_sizes — the per-peer size arithmetic the transport shares +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("token_count", "parts", "expected"), + [ + (5, 2, (3, 2)), # remainder rides on the leading shares + (6, 3, (2, 2, 2)), # even split + (7, 1, (7,)), # parts=1 = whole count (every A >= F subgroup) + (2, 4, (1, 1, 0, 0)), # fewer tokens than peers -> trailing zeros + (0, 3, (0, 0, 0)), # nothing to send + ], +) +def test_split_send_sizes_literal_cases(token_count, parts, expected): + assert split_send_sizes(token_count, parts) == expected + + +@pytest.mark.parametrize("token_count", range(9)) +@pytest.mark.parametrize("parts", range(1, 5)) +def test_split_send_sizes_properties(token_count, parts): + sizes = split_send_sizes(token_count, parts) + # Conservation, near-evenness, and deterministic front-loading — the + # contract that lets sender and receivers agree without communication. + assert sum(sizes) == token_count + assert len(sizes) == parts + assert max(sizes) - min(sizes) <= 1 + assert list(sizes) == sorted(sizes, reverse=True) + + +@pytest.mark.parametrize(("token_count", "parts"), [(1, 0), (1, -1), (-1, 2)]) +def test_split_send_sizes_rejects_invalid_inputs(token_count, parts): + with pytest.raises(ValueError): + split_send_sizes(token_count, parts) diff --git a/tests/unit/distributed/test_topology_snapshot.py b/tests/unit/distributed/test_topology_snapshot.py new file mode 100644 index 00000000..aff1fdd9 --- /dev/null +++ b/tests/unit/distributed/test_topology_snapshot.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Characterization snapshots of the P2P rank mapping. + +These tests freeze the *current* ``build_rank_mapping`` behavior for legal +``A >= F`` topologies so that the M2N generalization work can prove, at every +step, that existing gather-mode mappings remain byte-for-byte unchanged. + +Two layers of protection: + +1. Literal golden snapshots for representative topologies (1A1F, 2A2F, 4A2F). +2. Structural invariants over a wider grid of legal ``(A, F)`` pairs. +""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from afd_plugin.config import AFDConfig +from afd_plugin.distributed.topology import ( + build_rank_mapping, +) + + +def _config(role: str, attention: int, ffn: int) -> AFDConfig: + return AFDConfig( + role=role, + num_attention_ranks=attention, + num_ffn_ranks=ffn, + ) + + +def _mapping_dict(role: str, attention: int, ffn: int, role_rank: int) -> dict: + mapping = build_rank_mapping(_config(role, attention, ffn), role_rank) + return dataclasses.asdict(mapping) + + +# Captured from main @ 603c111 (pre-M2N). Do not regenerate to make a failing +# test pass: a diff here means gather-mode behavior changed. +_GOLDEN: dict[tuple[int, int], dict[tuple[str, int], dict]] = { + (1, 1): { + ("ffn", 0): { + "role": "ffn", + "role_rank": 0, + "world_rank": 0, + "p2p_rank": 0, + "attention_size": 1, + "ffn_size": 1, + "min_size": 1, + "ratio": 1, + "subgroup_index": 0, + "rank_in_subgroup": 0, + "subgroup_ranks": (0, 1), + "dp_metadata_destinations": (), + }, + ("attention", 0): { + "role": "attention", + "role_rank": 0, + "world_rank": 1, + "p2p_rank": 1, + "attention_size": 1, + "ffn_size": 1, + "min_size": 1, + "ratio": 1, + "subgroup_index": 0, + "rank_in_subgroup": 1, + "subgroup_ranks": (0, 1), + "dp_metadata_destinations": (0,), + }, + }, + (2, 2): { + ("ffn", 0): { + "role": "ffn", + "role_rank": 0, + "world_rank": 0, + "p2p_rank": 0, + "attention_size": 2, + "ffn_size": 2, + "min_size": 2, + "ratio": 1, + "subgroup_index": 0, + "rank_in_subgroup": 0, + "subgroup_ranks": (0, 2), + "dp_metadata_destinations": (), + }, + ("ffn", 1): { + "role": "ffn", + "role_rank": 1, + "world_rank": 1, + "p2p_rank": 1, + "attention_size": 2, + "ffn_size": 2, + "min_size": 2, + "ratio": 1, + "subgroup_index": 1, + "rank_in_subgroup": 0, + "subgroup_ranks": (1, 3), + "dp_metadata_destinations": (), + }, + ("attention", 0): { + "role": "attention", + "role_rank": 0, + "world_rank": 2, + "p2p_rank": 2, + "attention_size": 2, + "ffn_size": 2, + "min_size": 2, + "ratio": 1, + "subgroup_index": 0, + "rank_in_subgroup": 1, + "subgroup_ranks": (0, 2), + "dp_metadata_destinations": (0,), + }, + ("attention", 1): { + "role": "attention", + "role_rank": 1, + "world_rank": 3, + "p2p_rank": 3, + "attention_size": 2, + "ffn_size": 2, + "min_size": 2, + "ratio": 1, + "subgroup_index": 1, + "rank_in_subgroup": 1, + "subgroup_ranks": (1, 3), + "dp_metadata_destinations": (1,), + }, + }, + (4, 2): { + ("ffn", 0): { + "role": "ffn", + "role_rank": 0, + "world_rank": 0, + "p2p_rank": 0, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 0, + "rank_in_subgroup": 0, + "subgroup_ranks": (0, 2, 3), + "dp_metadata_destinations": (), + }, + ("ffn", 1): { + "role": "ffn", + "role_rank": 1, + "world_rank": 1, + "p2p_rank": 1, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 1, + "rank_in_subgroup": 0, + "subgroup_ranks": (1, 4, 5), + "dp_metadata_destinations": (), + }, + ("attention", 0): { + "role": "attention", + "role_rank": 0, + "world_rank": 2, + "p2p_rank": 2, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 0, + "rank_in_subgroup": 1, + "subgroup_ranks": (0, 2, 3), + "dp_metadata_destinations": (0,), + }, + ("attention", 1): { + "role": "attention", + "role_rank": 1, + "world_rank": 3, + "p2p_rank": 3, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 0, + "rank_in_subgroup": 2, + "subgroup_ranks": (0, 2, 3), + "dp_metadata_destinations": (1,), + }, + ("attention", 2): { + "role": "attention", + "role_rank": 2, + "world_rank": 4, + "p2p_rank": 4, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 1, + "rank_in_subgroup": 1, + "subgroup_ranks": (1, 4, 5), + "dp_metadata_destinations": (), + }, + ("attention", 3): { + "role": "attention", + "role_rank": 3, + "world_rank": 5, + "p2p_rank": 5, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 1, + "rank_in_subgroup": 2, + "subgroup_ranks": (1, 4, 5), + "dp_metadata_destinations": (), + }, + }, +} + +_GOLDEN_CASES = [ + (attention, ffn, role, role_rank) + for (attention, ffn), mappings in _GOLDEN.items() + for (role, role_rank) in mappings +] + +_LEGAL_GRID = [ + (1, 1), + (2, 1), + (2, 2), + (3, 1), + (4, 1), + (4, 2), + (4, 4), + (6, 2), + (6, 3), + (8, 2), + (8, 4), +] + + +@pytest.mark.parametrize( + ("attention", "ffn", "role", "role_rank"), + _GOLDEN_CASES, +) +def test_rank_mapping_matches_golden_snapshot(attention, ffn, role, role_rank): + assert ( + _mapping_dict(role, attention, ffn, role_rank) + == _GOLDEN[(attention, ffn)][(role, role_rank)] + ) + + +@pytest.mark.parametrize(("attention", "ffn"), _LEGAL_GRID) +def test_rank_mapping_invariants(attention, ffn): + ratio = attention // ffn + min_size = min(attention, ffn) + + seen_subgroups: dict[int, tuple[int, ...]] = {} + dp_destination_union: list[int] = [] + + for role, size in (("ffn", ffn), ("attention", attention)): + for role_rank in range(size): + mapping = build_rank_mapping(_config(role, attention, ffn), role_rank) + + if role == "ffn": + assert mapping.world_rank == role_rank + assert mapping.p2p_rank == role_rank + assert mapping.subgroup_index == role_rank + else: + assert mapping.world_rank == ffn + role_rank + assert mapping.p2p_rank == role_rank + min_size + assert mapping.subgroup_index == role_rank // ratio + + assert mapping.ratio == ratio + assert mapping.min_size == min_size + assert len(mapping.subgroup_ranks) == 1 + ratio + assert mapping.subgroup_ranks[0] == mapping.subgroup_index + assert ( + mapping.subgroup_ranks[mapping.rank_in_subgroup] == mapping.world_rank + ) + + prior = seen_subgroups.setdefault( + mapping.subgroup_index, + mapping.subgroup_ranks, + ) + assert prior == mapping.subgroup_ranks + + dp_destination_union.extend(mapping.dp_metadata_destinations) + + # The DP metadata senders must cover every FFN rank exactly once. + assert sorted(dp_destination_union) == list(range(ffn)) + # Subgroups must partition the attention ranks. + attention_members = sorted( + rank for ranks in seen_subgroups.values() for rank in ranks[1:] + ) + assert attention_members == list(range(ffn, ffn + attention)) + + +# The historical constraint-characterization test (A < F and non-divisible +# pairs raising ValueError) was removed together with the constraint itself; +# test_topology_partition.py now pins the relaxed validation behavior. diff --git a/tests/unit/v1/worker/test_ffn_metadata.py b/tests/unit/v1/worker/test_ffn_metadata.py new file mode 100644 index 00000000..707e07ca --- /dev/null +++ b/tests/unit/v1/worker/test_ffn_metadata.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""FFN-side token counts feed the EP dispatch sizes; they must follow the +subgroup partition exactly as the connector's receive sizes do.""" + +from __future__ import annotations + +import pytest + +from afd_plugin.v1.worker.ffn_metadata import ( + aggregate_ffn_token_counts, + project_ffn_token_counts_to_dp, +) + + +@pytest.mark.parametrize( + ("attention_counts", "attention", "ffn", "expected"), + [ + # Historical A >= F layouts: unchanged (docstring example and 1:1). + ((0, 4, 5, 6), 4, 2, (5, 11)), + ((3, 5), 2, 2, (3, 5)), + ((), 4, 2, (2, 2)), + # TP-expanded counts (one count per DP rank, two TP ranks each). + ((4, 6), 4, 2, (8, 12)), + # A < F: one attention splits across its FFN members. + ((5,), 1, 2, (3, 2)), + ((5, 3), 2, 4, (3, 2, 2, 1)), + # Fewer tokens than members: zero shares become one-token dummies. + ((2,), 1, 4, (1, 1, 1, 1)), + # Non-divisible A > F: uneven subgroups {A0,A1}{A2}. + ((1, 2, 3), 3, 2, (3, 3)), + ((1, 2, 3, 4, 5, 6), 6, 4, (3, 3, 9, 6)), + ], +) +def test_aggregate_follows_the_subgroup_partition( + attention_counts, attention, ffn, expected +): + assert ( + aggregate_ffn_token_counts( + attention_counts, attention_size=attention, ffn_size=ffn + ) + == expected + ) + + +def test_project_collapses_tp_ranks_to_dp_ranks(): + assert project_ffn_token_counts_to_dp((8, 8, 13, 13), dp_size=2) == (8, 13) + assert project_ffn_token_counts_to_dp((3, 2), dp_size=2) == (3, 2)