diff --git a/rl_engine/kernels/ops/pytorch/ffn/__init__.py b/rl_engine/kernels/ops/pytorch/ffn/__init__.py new file mode 100644 index 00000000..5e1549bd --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/ffn/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Composable Qwen-style feed-forward network reference operators.""" + +from .tensor_parallel import ( + DeterministicTensorParallelCommunication, + FFNContext, + TensorParallelFFN, + shard_qwen3_ffn_weights, +) + +__all__ = [ + "DeterministicTensorParallelCommunication", + "FFNContext", + "TensorParallelFFN", + "shard_qwen3_ffn_weights", +] diff --git a/rl_engine/kernels/ops/pytorch/ffn/tensor_parallel.py b/rl_engine/kernels/ops/pytorch/ffn/tensor_parallel.py new file mode 100644 index 00000000..b06b327d --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/ffn/tensor_parallel.py @@ -0,0 +1,484 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tensor-parallel Qwen3-style SwiGLU FFN orchestration. + +The module implements the non-sequence-parallel TP ownership contract used by +the Qwen3 dense MLP: + +* ``gate`` and ``up`` are ColumnParallel (their output/intermediate dimension + is sharded); their SwiGLU result stays local to the TP rank. +* ``down`` is RowParallel (its input/intermediate dimension is sharded); its + same-coordinate hidden outputs are summed once over the explicit TP group. + +The autograd collective mappings are deliberately asymmetric. A RowParallel +forward reduction has an identity backward, while the replicated FFN input is +copied in the forward and reduced in the backward. Consequently the sole +backward TP SUM occurs after the local Gate and Up input-gradient +contributions have been accumulated, exactly as required by the TP FFN +derivation. There is no TP reduction for Down's feature-sharded ``dHidden``. + +CP/SP configuration and CP weight-gradient reductions intentionally do not +belong in this PR3 module. + +For production CUDA execution, callers can attach +DeterministicTensorParallelCommunication to FFNContext. The adapter lazily +binds the deterministic all-reduce from distributed PR #310 and uses it at +precisely the two TP SUM sites above. The default native torch.distributed path +remains available for the CPU/Gloo reference tests. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Optional, Protocol + +import torch +import torch.distributed as dist +from torch import Tensor, nn + +from rl_engine.kernels.ops.pytorch.activation.swiglu import NativeSwiGLUOp + +LocalGemm = Callable[[Tensor, Tensor], Tensor] + + +class TensorParallelCommunication(Protocol): + """TP SUM interface used at the FFN's two explicit collective boundaries.""" + + def all_reduce(self, tensor: Tensor, *, ctx: "FFNContext") -> Tensor: + """Return the TP sum of tensor using ctx.tp_group.""" + + +class DeterministicTensorParallelCommunication: + """Lazy PR #310 deterministic CUDA all-reduce adapter for TP FFN. + + One instance is intended to be owned by a long-lived FFNContext (or shared + by contexts that use the same TP process group). It caches the underlying + DeterministicCollective so CUDA IPC setup occurs once, not once per FFN + reduction. Call close() after the last use and before destroying the + process group. + + The adapter deliberately fails closed: it accepts only CUDA tensors and + requires PR #310's rl_engine.distributed.DeterministicCollective. This + prevents an accidental fallback to an unspecified NCCL reduction order on + a path advertised as deterministic. + """ + + def __init__( + self, + *, + process_group: Any = None, + collective: Any = None, + max_size_bytes: int = 64 * 1024 * 1024, + ) -> None: + if max_size_bytes <= 0: + raise ValueError("max_size_bytes must be positive.") + self._process_group = process_group + self._collective = collective + self._max_size_bytes = int(max_size_bytes) + + def all_reduce(self, tensor: Tensor, *, ctx: "FFNContext") -> Tensor: + """Run PR #310's in-place deterministic TP SUM for tensor.""" + + if not tensor.is_cuda: + raise ValueError( + "DeterministicTensorParallelCommunication requires a CUDA tensor; " + "use the native TP path for CPU/Gloo reference execution." + ) + collective = self._get_collective(tensor, ctx) + self._validate_collective_coordinates(collective, ctx) + reduced = collective.all_reduce(tensor, out=tensor) + if reduced is not tensor: + raise RuntimeError( + "DeterministicCollective.all_reduce must return its aliased out tensor." + ) + return reduced + + def close(self) -> None: + """Release the cached PR #310 CUDA IPC collective, if one was created.""" + + close = getattr(self._collective, "close", None) + if close is not None: + close() + + def _get_collective(self, tensor: Tensor, ctx: "FFNContext") -> Any: + if self._process_group is not None and self._process_group is not ctx.tp_group: + raise ValueError( + "Deterministic TP communication process_group must match FFNContext.tp_group." + ) + if self._collective is None: + try: + from rl_engine.distributed import DeterministicCollective + except ImportError as exc: + raise RuntimeError( + "deterministic TP FFN communication requires distributed PR #310 " + "(rl_engine.distributed.DeterministicCollective)." + ) from exc + self._collective = DeterministicCollective( + group=ctx.tp_group, + device=tensor.device, + max_size_bytes=self._max_size_bytes, + ) + + collective_group = getattr(self._collective, "group", ctx.tp_group) + if collective_group is not ctx.tp_group: + raise ValueError("DeterministicCollective group must match FFNContext.tp_group.") + return self._collective + + @staticmethod + def _validate_collective_coordinates(collective: Any, ctx: "FFNContext") -> None: + collective_size = getattr(collective, "world_size", None) + collective_rank = getattr(collective, "rank", None) + if collective_size != ctx.tp_size or collective_rank != ctx.tp_rank: + raise ValueError( + "DeterministicCollective coordinates must match FFNContext: " + f"collective=({collective_size}, {collective_rank}), " + f"context=({ctx.tp_size}, {ctx.tp_rank})." + ) + + +def _missing_deterministic_gemm(input_: Tensor, weight: Tensor) -> Tensor: + """Fail closed rather than silently losing the batch-invariance contract.""" + + del input_, weight + raise RuntimeError( + "TensorParallelFFN requires an explicit batch-invariant local GEMM. " + "Pass the existing deterministic_gemm CUDA/Triton primitive with " + "signature gemm(a[M, K], b[K, N]) -> [M, N]." + ) + + +@dataclass(frozen=True) +class FFNContext: + """Explicit tensor-parallel configuration owned by an FFN caller. + + ``tp_group`` is never created or inferred by this module. A multi-rank + configuration must supply an initialized explicit process group; this is + important because later PRs add distinct CP and SP groups. + + ``tp_communication`` optionally replaces the native TP SUM at the two + FFN collective boundaries. Use + ``DeterministicTensorParallelCommunication(process_group=tp_group)`` for + the strict CUDA path backed by distributed PR #310. + """ + + tp_group: Any = None + tp_size: Optional[int] = None + tp_rank: Optional[int] = None + tp_communication: Optional[TensorParallelCommunication] = None + + def __post_init__(self) -> None: + if self.tp_group is None: + size = 1 if self.tp_size is None else int(self.tp_size) + rank = 0 if self.tp_rank is None else int(self.tp_rank) + if size != 1 or rank != 0: + raise ValueError( + "FFNContext(tp_group=None) only supports tp_size=1 and tp_rank=0. " + "Supply an explicit initialized tp_group for TP > 1." + ) + else: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError( + "FFNContext(tp_group=...) requires torch.distributed to be initialized." + ) + group_size = dist.get_world_size(group=self.tp_group) + group_rank = dist.get_rank(group=self.tp_group) + size = group_size if self.tp_size is None else int(self.tp_size) + rank = group_rank if self.tp_rank is None else int(self.tp_rank) + if size != group_size: + raise ValueError( + f"ctx.tp_size={size} does not match tp_group world size={group_size}." + ) + if rank != group_rank: + raise ValueError(f"ctx.tp_rank={rank} does not match tp_group rank={group_rank}.") + + if size < 1 or not 0 <= rank < size: + raise ValueError(f"invalid TP coordinates: tp_size={size}, tp_rank={rank}.") + object.__setattr__(self, "tp_size", size) + object.__setattr__(self, "tp_rank", rank) + + @property + def is_tensor_parallel(self) -> bool: + """Whether this context owns more than one tensor-parallel shard.""" + + assert self.tp_size is not None + return self.tp_size > 1 + + +def _all_reduce_sum(tensor: Tensor, ctx: FFNContext) -> Tensor: + """Synchronously sum a tensor over the explicitly configured TP group.""" + + if ctx.is_tensor_parallel: + if ctx.tp_communication is not None: + return ctx.tp_communication.all_reduce(tensor, ctx=ctx) + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=ctx.tp_group) + return tensor + + +class _CopyToTensorParallelRegion(torch.autograd.Function): + """Replicated input: identity forward, TP SUM backward. + + The Gate and Up ColumnParallel projections each produce a local + same-coordinate ``dX`` contribution. Autograd accumulates those local + contributions before this function's backward executes, so this is one + logical TP all-reduce of their combined ``[... , H]`` gradient. + """ + + @staticmethod + def forward(ctx: Any, input_: Tensor, tp_ctx: FFNContext) -> Tensor: + ctx.tp_ctx = tp_ctx + return input_ + + @staticmethod + def backward(ctx: Any, grad_output: Tensor) -> tuple[Tensor, None]: + # Do not mutate a gradient owned by a downstream autograd node. + grad_input = grad_output.contiguous().clone() + return _all_reduce_sum(grad_input, ctx.tp_ctx), None + + +class _ReduceFromTensorParallelRegion(torch.autograd.Function): + """RowParallel output: TP SUM forward, identity backward. + + Its backward must not reduce ``dOutput``. Each rank receives the replicated + top gradient and computes its local feature-sharded ``dHidden`` through + Down's local weight shard. + """ + + @staticmethod + def forward(ctx: Any, input_: Tensor, tp_ctx: FFNContext) -> Tensor: + # The local Down partial is useful to callers through ``forward_local``; + # preserve it by reducing a clone rather than modifying it in place. + return _all_reduce_sum(input_.contiguous().clone(), tp_ctx) + + @staticmethod + def backward(ctx: Any, grad_output: Tensor) -> tuple[Tensor, None]: + return grad_output, None + + +def _copy_to_tensor_parallel_region(input_: Tensor, ctx: FFNContext) -> Tensor: + return _CopyToTensorParallelRegion.apply(input_, ctx) + + +def _reduce_from_tensor_parallel_region(input_: Tensor, ctx: FFNContext) -> Tensor: + return _ReduceFromTensorParallelRegion.apply(input_, ctx) + + +def shard_qwen3_ffn_weights( + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + ctx: FFNContext, +) -> tuple[Tensor, Tensor, Tensor]: + """Return this rank's Qwen ``[out, in]`` Gate/Up/Down weight shards. + + Gate and Up are split along output rows. Down is split along input columns. + The returned tensors are contiguous views copied only as needed by callers + that register them as local parameters. + """ + + if gate_weight.ndim != 2 or up_weight.ndim != 2 or down_weight.ndim != 2: + raise ValueError("gate_weight, up_weight, and down_weight must all be rank-2 tensors.") + + assert ctx.tp_size is not None + assert ctx.tp_rank is not None + intermediate_size, hidden_size = gate_weight.shape + if up_weight.shape != (intermediate_size, hidden_size): + raise ValueError( + "up_weight must have the same [intermediate, hidden] shape as gate_weight; " + f"got {tuple(up_weight.shape)} versus {tuple(gate_weight.shape)}." + ) + if down_weight.shape != (hidden_size, intermediate_size): + raise ValueError( + "down_weight must have shape [hidden, intermediate]; " + f"expected {(hidden_size, intermediate_size)}, got {tuple(down_weight.shape)}." + ) + if intermediate_size % ctx.tp_size != 0: + raise ValueError( + f"intermediate_size={intermediate_size} must divide evenly by tp_size={ctx.tp_size}." + ) + if not (gate_weight.device == up_weight.device == down_weight.device): + raise ValueError("all FFN weights must be on the same device.") + if not (gate_weight.dtype == up_weight.dtype == down_weight.dtype): + raise ValueError("all FFN weights must have the same dtype.") + + local_intermediate = intermediate_size // ctx.tp_size + start = ctx.tp_rank * local_intermediate + stop = start + local_intermediate + return ( + gate_weight[start:stop].contiguous(), + up_weight[start:stop].contiguous(), + down_weight[:, start:stop].contiguous(), + ) + + +class TensorParallelFFN(nn.Module): + """Qwen3 SwiGLU FFN with ColumnParallel Gate/Up and RowParallel Down. + + Parameters use the normal ``torch.nn.Linear`` ``[out, in]`` layout. This + module owns one TP shard only; use :meth:`from_full_weights` in tests or a + model loader to materialize rank-local parameter shards from full weights. + + ``gemm`` has the deterministic-GEMM-compatible signature + ``gemm(a[M, K], b[K, N]) -> [M, N]``. Production BF16 callers must inject + the existing deterministic GEMM primitive (CUDA or Triton). Omitting it + fails closed when ``forward`` is called, because generic ``torch.matmul`` + would silently violate the batch-invariance contract. This class + intentionally does not introduce a separate GEMM arithmetic implementation. + """ + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + *, + ctx: Optional[FFNContext] = None, + gate_weight: Optional[Tensor] = None, + up_weight: Optional[Tensor] = None, + down_weight: Optional[Tensor] = None, + activation: Optional[Callable[[Tensor, Tensor], Tensor]] = None, + gemm: Optional[LocalGemm] = None, + device: Optional[torch.device | str] = None, + dtype: Optional[torch.dtype] = None, + ) -> None: + super().__init__() + if hidden_size < 1 or intermediate_size < 1: + raise ValueError("hidden_size and intermediate_size must both be positive.") + + self.ctx = FFNContext() if ctx is None else ctx + assert self.ctx.tp_size is not None + if intermediate_size % self.ctx.tp_size != 0: + raise ValueError( + f"intermediate_size={intermediate_size} must divide evenly by " + f"ctx.tp_size={self.ctx.tp_size}." + ) + + self.hidden_size = int(hidden_size) + self.intermediate_size = int(intermediate_size) + self.local_intermediate_size = intermediate_size // self.ctx.tp_size + self.activation = NativeSwiGLUOp() if activation is None else activation + self._gemm: LocalGemm = _missing_deterministic_gemm if gemm is None else gemm + + self.gate_weight = self._make_parameter( + gate_weight, + (self.local_intermediate_size, hidden_size), + "gate_weight", + device=device, + dtype=dtype, + ) + self.up_weight = self._make_parameter( + up_weight, + (self.local_intermediate_size, hidden_size), + "up_weight", + device=device, + dtype=dtype, + ) + self.down_weight = self._make_parameter( + down_weight, + (hidden_size, self.local_intermediate_size), + "down_weight", + device=device, + dtype=dtype, + ) + + @staticmethod + def _make_parameter( + value: Optional[Tensor], + shape: tuple[int, int], + name: str, + *, + device: Optional[torch.device | str], + dtype: Optional[torch.dtype], + ) -> nn.Parameter: + if value is None: + parameter = torch.empty(shape, device=device, dtype=dtype) + nn.init.kaiming_uniform_(parameter, a=5**0.5) + return nn.Parameter(parameter) + + if value.shape != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(value.shape)}.") + if device is not None and value.device != torch.device(device): + raise ValueError( + f"{name} is on {value.device}, expected device {torch.device(device)}." + ) + if dtype is not None and value.dtype != dtype: + raise ValueError(f"{name} has dtype {value.dtype}, expected {dtype}.") + return nn.Parameter(value.detach().clone()) + + @classmethod + def from_full_weights( + cls, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + ctx: Optional[FFNContext] = None, + activation: Optional[Callable[[Tensor, Tensor], Tensor]] = None, + gemm: Optional[LocalGemm] = None, + ) -> "TensorParallelFFN": + """Construct a rank-local FFN module from full Qwen-format weights.""" + + context = FFNContext() if ctx is None else ctx + gate_shard, up_shard, down_shard = shard_qwen3_ffn_weights( + gate_weight, up_weight, down_weight, ctx=context + ) + intermediate_size, hidden_size = gate_weight.shape + return cls( + hidden_size, + intermediate_size, + ctx=context, + gate_weight=gate_shard, + up_weight=up_shard, + down_weight=down_shard, + activation=activation, + gemm=gemm, + ) + + def _local_linear(self, input_: Tensor, weight: Tensor) -> Tensor: + if input_.shape[-1] != weight.shape[1]: + raise ValueError( + f"input last dimension {input_.shape[-1]} does not match weight input " + f"dimension {weight.shape[1]}." + ) + input_2d = input_.reshape(-1, input_.shape[-1]) + output_2d = self._gemm(input_2d, weight.t().contiguous()) + expected_shape = (input_2d.shape[0], weight.shape[0]) + if output_2d.shape != expected_shape: + raise RuntimeError( + "gemm returned an invalid shape: expected " + f"{expected_shape}, got {tuple(output_2d.shape)}." + ) + return output_2d.reshape(*input_.shape[:-1], weight.shape[0]) + + def forward_local(self, input_: Tensor) -> Tensor: + """Return the local pre-TP-reduction Down output partial. + + This method is primarily a validation boundary: its result has shape + ``[..., hidden_size]`` but represents only this rank's RowParallel + contribution. ``forward`` performs the required TP SUM over it. + """ + + if input_.ndim < 2: + raise ValueError( + f"input must have at least [tokens, hidden] dimensions, got {tuple(input_.shape)}." + ) + if input_.device != self.gate_weight.device: + raise ValueError( + f"input is on {input_.device}, but FFN weights are on {self.gate_weight.device}." + ) + if input_.dtype != self.gate_weight.dtype: + raise ValueError( + f"input has dtype {input_.dtype}, but FFN weights have dtype " + f"{self.gate_weight.dtype}." + ) + + replicated_input = _copy_to_tensor_parallel_region(input_, self.ctx) + gate = self._local_linear(replicated_input, self.gate_weight) + up = self._local_linear(replicated_input, self.up_weight) + hidden = self.activation(gate, up) + return self._local_linear(hidden, self.down_weight) + + def forward(self, input_: Tensor) -> Tensor: + """Compute the replicated hidden output after the one Down TP SUM.""" + + local_output_partial = self.forward_local(input_) + return _reduce_from_tensor_parallel_region(local_output_partial, self.ctx) diff --git a/tests/test_tensor_parallel_ffn.py b/tests/test_tensor_parallel_ffn.py new file mode 100644 index 00000000..fc027889 --- /dev/null +++ b/tests/test_tensor_parallel_ffn.py @@ -0,0 +1,342 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""PR3 tensor-parallel Qwen-style FFN topology and invariance tests.""" + +from __future__ import annotations + +import os +import queue +import socket +import tempfile +import traceback +from pathlib import Path +from typing import Any + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.activation.swiglu import NativeSwiGLUOp +from rl_engine.kernels.ops.pytorch.ffn.tensor_parallel import ( + DeterministicTensorParallelCommunication, + FFNContext, + TensorParallelFFN, + shard_qwen3_ffn_weights, +) + + +def _gloo_available() -> bool: + return dist.is_available() and dist.is_gloo_available() + + +requires_gloo = pytest.mark.skipif( + not _gloo_available(), reason="tensor-parallel FFN CPU test requires torch.distributed Gloo." +) + + +class _RecordingTPCommunication: + """Gloo stand-in that makes the configured TP path observable.""" + + def __init__(self) -> None: + self.calls: list[tuple[tuple[int, ...], int, int]] = [] + + def all_reduce(self, tensor: torch.Tensor, *, ctx: FFNContext) -> torch.Tensor: + assert ctx.tp_size is not None + assert ctx.tp_rank is not None + self.calls.append((tuple(tensor.shape), ctx.tp_size, ctx.tp_rank)) + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=ctx.tp_group) + return tensor + + +def _configure_gloo_loopback() -> None: + """Avoid hostname-resolution dependence in local CPU topology tests.""" + + if "GLOO_SOCKET_IFNAME" in os.environ: + return + interfaces = {name for _, name in socket.if_nameindex()} + loopback = "lo" if "lo" in interfaces else "lo0" if "lo0" in interfaces else None + if loopback is not None: + os.environ["GLOO_SOCKET_IFNAME"] = loopback + + +def _full_ffn( + input_: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor: + swiglu = NativeSwiGLUOp() + return F.linear(swiglu(F.linear(input_, gate_weight), F.linear(input_, up_weight)), down_weight) + + +def _fixed_row_gemm(input_2d: torch.Tensor, weight_2d: torch.Tensor) -> torch.Tensor: + """CPU test adapter for the fixed-row contract of the det_gemm backends. + + Production uses the existing CUDA/Triton deterministic GEMM primitive via + ``TensorParallelFFN(gemm=...)``. CPU Gloo has no such backend, so this + adapter invokes PyTorch once per output row; changing M therefore cannot + select a different accumulation path for an already-valid row. + """ + + return torch.stack([torch.matmul(row.unsqueeze(0), weight_2d).squeeze(0) for row in input_2d]) + + +def _tp_ffn_worker( + rank: int, + world_size: int, + init_method: str, + result_queue: Any, + use_recording_tp_communication: bool = False, +) -> None: + try: + torch.set_num_threads(1) + _configure_gloo_loopback() + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + tp_communication = _RecordingTPCommunication() if use_recording_tp_communication else None + ctx = FFNContext(tp_group=dist.group.WORLD, tp_communication=tp_communication) + assert ctx.tp_size == world_size + assert ctx.tp_rank == rank + + hidden_size, intermediate_size = 8, 24 + generator = torch.Generator().manual_seed(20260812) + input_base = torch.randn(3, 5, hidden_size, generator=generator) + gate_full = torch.randn(intermediate_size, hidden_size, generator=generator) + up_full = torch.randn(intermediate_size, hidden_size, generator=generator) + down_full = torch.randn(hidden_size, intermediate_size, generator=generator) + grad_output = torch.randn(3, 5, hidden_size, generator=generator) + + # Materialize the full reference before installing the collective spy. + reference_input = input_base.detach().clone().requires_grad_(True) + reference_gate = gate_full.detach().clone().requires_grad_(True) + reference_up = up_full.detach().clone().requires_grad_(True) + reference_down = down_full.detach().clone().requires_grad_(True) + reference_output = _full_ffn(reference_input, reference_gate, reference_up, reference_down) + reference_output.backward(grad_output) + + ffn = TensorParallelFFN.from_full_weights( + gate_full, up_full, down_full, ctx=ctx, gemm=_fixed_row_gemm + ) + tp_input = input_base.detach().clone().requires_grad_(True) + + # Forward uses one TP SUM for Down. During backward the only TP SUM + # is the combined Gate+Up dX at the replicated-input boundary. + observed_collectives: list[tuple[int, ...]] = [] + original_all_reduce = dist.all_reduce + + def traced_all_reduce(tensor: torch.Tensor, *args: Any, **kwargs: Any) -> Any: + observed_collectives.append(tuple(tensor.shape)) + return original_all_reduce(tensor, *args, **kwargs) + + dist.all_reduce = traced_all_reduce # type: ignore[assignment] + try: + output = ffn(tp_input) + output.backward(grad_output) + finally: + dist.all_reduce = original_all_reduce # type: ignore[assignment] + + local_intermediate = intermediate_size // world_size + start = rank * local_intermediate + stop = start + local_intermediate + expected_collective_shape = tuple(input_base.shape) + result_queue.put( + { + "ok": True, + "rank": rank, + "output_error": float((output - reference_output).abs().max().item()), + "input_grad_error": float( + (tp_input.grad - reference_input.grad).abs().max().item() + ), + "gate_grad_error": float( + (ffn.gate_weight.grad - reference_gate.grad[start:stop]).abs().max().item() + ), + "up_grad_error": float( + (ffn.up_weight.grad - reference_up.grad[start:stop]).abs().max().item() + ), + "down_grad_error": float( + (ffn.down_weight.grad - reference_down.grad[:, start:stop]).abs().max().item() + ), + "collectives": observed_collectives, + "expected_collective_shape": expected_collective_shape, + "configured_collectives": ( + [] if tp_communication is None else tp_communication.calls + ), + } + ) + except Exception: + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _tp_batch_invariance_worker( + rank: int, world_size: int, init_method: str, result_queue: Any +) -> None: + try: + torch.set_num_threads(1) + _configure_gloo_loopback() + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + ctx = FFNContext(tp_group=dist.group.WORLD) + hidden_size, intermediate_size = 8, 24 + generator = torch.Generator().manual_seed(99) + input_valid = torch.randn(4, 3, hidden_size, generator=generator) + padding = torch.randn(3, 3, hidden_size, generator=generator) + gate_full = torch.randn(intermediate_size, hidden_size, generator=generator) + up_full = torch.randn(intermediate_size, hidden_size, generator=generator) + down_full = torch.randn(hidden_size, intermediate_size, generator=generator) + ffn = TensorParallelFFN.from_full_weights( + gate_full, up_full, down_full, ctx=ctx, gemm=_fixed_row_gemm + ) + + full_output = ffn(input_valid) + single_output = ffn(input_valid[:1]) + padded_output = ffn(torch.cat((input_valid, padding), dim=0)) + result_queue.put( + { + "ok": True, + "rank": rank, + "slice_equal": bool(torch.equal(full_output[:1], single_output)), + "padding_equal": bool( + torch.equal(full_output, padded_output[: input_valid.shape[0]]) + ), + } + ) + except Exception: + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_tp_workers(worker: Any, *worker_args: Any) -> list[dict[str, Any]]: + world_size = 2 + mp_context = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmp_dir: + init_file = Path(tmp_dir) / "tp_ffn_init" + init_method = init_file.as_uri() + result_queue = mp_context.Queue() + workers = [ + mp_context.Process( + target=worker, + args=(rank, world_size, init_method, result_queue, *worker_args), + ) + for rank in range(world_size) + ] + for process in workers: + process.start() + + results: list[dict[str, Any]] = [] + try: + for _ in workers: + results.append(result_queue.get(timeout=30)) + except queue.Empty: + pytest.fail("timed out waiting for tensor-parallel Gloo FFN workers") + finally: + for process in workers: + process.join(timeout=30) + if process.is_alive(): + process.terminate() + process.join() + + failures = [result for result in results if not result["ok"]] + if failures: + pytest.fail("\n".join(result["traceback"] for result in failures)) + failed_workers = [process for process in workers if process.exitcode != 0] + if failed_workers: + pytest.fail( + f"tensor-parallel Gloo workers failed: {[p.exitcode for p in failed_workers]}" + ) + return results + + +@requires_gloo +def test_tensor_parallel_ffn_matches_full_reference_and_tp_backward_contract() -> None: + """TP=2 FFN agrees with full FFN and reduces only at the documented sites.""" + + results = _run_tp_workers(_tp_ffn_worker) + for result in results: + # The global reference contracts the full intermediate in one GEMM; + # TP first rounds two local contractions and then sums them. The + # topology is therefore checked against the shared fp32 tolerance, + # rather than requiring a different reduction tree to be bitwise equal. + assert result["output_error"] <= 1e-4 + assert result["input_grad_error"] <= 1e-4 + assert result["gate_grad_error"] <= 1e-4 + assert result["up_grad_error"] <= 1e-4 + assert result["down_grad_error"] <= 1e-4 + # One forward Down reduction and one backward combined Gate+Up dX + # reduction. In particular, there is no [B, S, I / TP] all-reduce + # for Down's dHidden. + assert result["collectives"] == [ + result["expected_collective_shape"], + result["expected_collective_shape"], + ] + + +@requires_gloo +def test_tensor_parallel_ffn_routes_collectives_through_configured_communication() -> None: + """Configured TP communication owns exactly the forward and backward SUMs.""" + + results = _run_tp_workers(_tp_ffn_worker, True) + for result in results: + assert result["output_error"] <= 1e-4 + assert result["input_grad_error"] <= 1e-4 + assert result["collectives"] == [ + result["expected_collective_shape"], + result["expected_collective_shape"], + ] + assert result["configured_collectives"] == [ + (result["expected_collective_shape"], 2, result["rank"]), + (result["expected_collective_shape"], 2, result["rank"]), + ] + + +@requires_gloo +def test_tensor_parallel_ffn_is_batch_invariant() -> None: + """Valid rows are bitwise unchanged by TP FFN slicing or batch padding.""" + + results = _run_tp_workers(_tp_batch_invariance_worker) + assert all(result["slice_equal"] and result["padding_equal"] for result in results) + + +def test_ffn_context_rejects_unbound_multi_rank_tp() -> None: + with pytest.raises(ValueError, match="explicit initialized tp_group"): + FFNContext(tp_size=2) + + +def test_qwen_weight_shards_follow_column_and_row_parallel_dimensions() -> None: + ctx = FFNContext() + gate = torch.empty(24, 8) + up = torch.empty(24, 8) + down = torch.empty(8, 24) + gate_shard, up_shard, down_shard = shard_qwen3_ffn_weights(gate, up, down, ctx=ctx) + assert gate_shard.shape == (24, 8) + assert up_shard.shape == (24, 8) + assert down_shard.shape == (8, 24) + + +def test_deterministic_tp_communication_rejects_cpu_tensors() -> None: + communication = DeterministicTensorParallelCommunication() + with pytest.raises(ValueError, match="requires a CUDA tensor"): + communication.all_reduce(torch.zeros(2), ctx=FFNContext()) + + +def test_ffn_refuses_non_deterministic_default_gemm() -> None: + ffn = TensorParallelFFN(8, 24) + with pytest.raises(RuntimeError, match="explicit batch-invariant local GEMM"): + ffn(torch.randn(2, 8))