diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..993ccf52 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,6 +76,17 @@ jobs: run: | python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu" + - name: Run WS2 Logprob Contract Tests (CPU-safe) + run: python -m pytest tests/test_logprob_contract.py -v + + - name: Run WS2 Vocab-Parallel Logprob Tests (CPU-safe) + run: python -m pytest tests/test_vocab_parallel_logp.py -v + + - name: Run WS2 Logprob Comparison Tests (CPU-safe) + run: | + python -m pytest tests/test_logprob_comparison.py -v + python -m pytest tests/test_distributed_logprob_comparison.py -v + docs: runs-on: ubuntu-latest steps: diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index bedf3475..23c1586a 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -11,6 +11,14 @@ logical type, and the registry selects the first available backend for the curre 4. Cache successfully constructed operator instances. 5. Skip backends that already failed in the current process. +WS2 TP-aware logprob uses the stricter `KernelRegistry.get_logprob_op(contract)` path. In +addition to platform priority, this path requires a backend capability descriptor and checks +the requested role, dtype, TP/CP layout, padded-vs-real vocab masking, inactive-token +support, vocab-domain LSE export, and deterministic TP merge semantics. Incompatible +candidates produce explicit rejection reasons and are never used as an undeclared fallback. +The contract objects and their normative reduction semantics are documented in +`rl_engine.kernels.logprob_contract`. + ## LogP Priority | Platform | Priority | diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index fbc0e9f1..b060c10f 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -54,6 +54,35 @@ CUDA priority list when the extension exposes `_C.batch_invariant_logp_sm90` (built with `KERNEL_ALIGN_FORCE_SM90=1`) on an SM90 device. On any other build or device, dispatch is unchanged (Triton -> PyTorch). +## Tensor Parallel + +`VocabParallelLogprobOp` +(`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) +defines a cross-TP bitwise contract for TP=1, TP=2, and TP=4 when +`num_vocab_tiles` is fixed and every vocabulary-shard boundary is tile-aligned. +The complete BF16 CUDA/NCCL validation matrix for this contract is tracked by +issue #241 PR4. + +1. Split the padded vocabulary into `num_vocab_tiles` fixed tiles. +2. Each rank computes fp32 `(max, sumexp)` for the tiles it owns. Every tile + is reduced as the same contiguous `[n, tile]` shape, on any rank. +3. All tile partials are shared with `all_gather`. The collective only moves + bytes; it never does math, so it cannot round anything. +4. Every rank merges all tiles in the same fixed order, over the same + `[n, num_vocab_tiles]` shape. `LSE = M + log(sum(s_t * exp(m_t - M)))`. +5. The target logit is copied from the rank that owns it (never summed). +6. `logp = target_logit - LSE`. Inactive rows become `0.0`. + +Usage goes through the contract-aware entry point: + +```python +from rl_engine.kernels.registry import kernel_registry + +result = kernel_registry.get_logprob_op(contract) # LogprobContract from +op = result.op # rl_engine.kernels.logprob_contract +logp, lse = op(local_logits, target_ids, contract=contract, tp_group=tp_group) +``` + ## Benchmarks `benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the @@ -179,6 +208,160 @@ fp16/bf16 backward: checked against fp32 reference with relaxed tolerance CPU-vs-CUDA comparisons use tolerance-based checks; batch-invariance checks within the same backend use exact equality where appropriate. +## TP=1 Comparison Harness + +The single-GPU comparison harness is the TP=1 registration and regression guard +for issue #241. It uses the batch-invariant PyTorch implementation as the +reference and compares exact `pytorch`, `triton`, or `cuda-sm90` backends before +distributed communication is introduced. + +Each backend exposes a diagnostic-only entry point while the production contract +remains unchanged: + +```text +op(logits, target_ids) -> logp +op.forward_with_lse(logits, target_ids) -> (logp, lse) +``` + +The harness reports LSE drift over every logical token row and selected-logprob +drift over active response/action tokens only. Drift summaries contain max, +mean, p95, p99, and the number of compared values. Reports also record requested +and actual backends, implementation, direct-LSE provenance, input shape and +dtype, `tp_world=1`, and `communication=none`. + +Backend selection is exact and does not use registry fallback. In particular, +an explicit `cuda-sm90` comparison fails unless the compiled SM90 extension, +Hopper hardware, input dtype, and vocab row stride satisfy the kernel contract. + +Run the PyTorch TP=1 guard directly from the kernel-specific testing module: + +```bash +python rl_engine/testing/logprob_comparison.py \ + --candidate pytorch \ + --device cpu \ + --dtype fp32 \ + --batch 2 \ + --seq 16 \ + --vocab 257 +``` + +On a GPU, repeat `--candidate` to compare multiple exact backends: + +```bash +python rl_engine/testing/logprob_comparison.py \ + --candidate triton \ + --candidate cuda-sm90 \ + --device cuda \ + --dtype bf16 \ + --batch 2 \ + --seq 16 \ + --vocab 151936 +``` + +The command writes structured JSON to stdout and routes backend diagnostics to +stderr. The harness does not implement vocab sharding, collective communication, +cross-rank LSE merging, or CP reconstruction. + +### SM90 validation + +SM90 validation requires a Hopper GPU, CUDA-enabled PyTorch, and an `nvcc` +toolkit matching `torch.version.cuda`. Build the extension with: + +```bash +export FORCE_CUDA=1 +export KERNEL_ALIGN_FORCE_SM90=1 +export TORCH_CUDA_ARCH_LIST="9.0+PTX" + +python -m pip install --no-build-isolation --no-deps -e . +``` + +Run the focused harness tests, the complete operator suite, and an explicit +SM90 comparison: + +```bash +python -m pytest \ + tests/test_logprob_comparison.py \ + tests/test_operator_inputs.py \ + tests/test_op_checks.py -q + +python -m pytest tests/test_batch_invariant_logp.py -q + +python rl_engine/testing/logprob_comparison.py \ + --candidate cuda-sm90 \ + --device cuda \ + --dtype bf16 \ + --batch 2 \ + --seq 16 \ + --vocab 151936 \ + --prompt-tokens 8 \ + --seed 241 +``` + +The PR2 path was validated on an NVIDIA H800 PCIe with PyTorch 2.11.0+cu128, +CUDA 12.8, and Triton 3.6.0. The focused tests passed 41 cases and the complete +batch-invariant suite passed 67 cases. For BF16 shape `[2, 16, 151936]`, both +LSE and active-token dlogp had maximum absolute drift +`9.5367431640625e-07` against the PyTorch reference, with no backend fallback. + +## Distributed WS2 Drift Report + +The issue #241 PR4 runner materializes one TP/CP topology per `torchrun` +invocation. TP partitions the vocabulary and is the only numerical merge axis; +CP partitions token rows and is recorded in provenance without participating in +the vocab-domain LSE merge. For global rank `r`: + +```text +tp_rank = r % tp_world_size +cp_rank = r // tp_world_size +``` + +Every case generates the same seeded FP32 logical logits, targets, and active +mask. The candidate receives a BF16 token/vocab shard through the explicit +`pytorch-vocab-parallel-logp-ws2` backend, while the independent oracle computes +`torch.logsumexp` over the complete real-vocab FP32 token slice. Distributed +dispatch rejects `auto`, capability fallback, topology mismatches, non-tileable +vocabularies, and incomplete materialization. + +Reports follow the issue #116 fields and contain per-rank and aggregate LSE and +active-token dlogp summaries: max/mean/p95/p99 absolute drift, max relative +drift, worst global token position, target id, target owner rank, #108 tolerance, +and pass/fail. Provenance includes TP/CP topology, dtype, shard bounds, backend +capability, contract fingerprint, reduction spec, merge order, transport, and +the exact launch command. Replicated TP outputs are checked bitwise before one +representative per CP shard is included in aggregate statistics. + +Print the scoped TP=1/2/4 x CP=1/2 launch matrix without starting workers: + +```bash +python rl_engine/testing/distributed_logprob_comparison.py \ + --plan \ + --device cuda \ + --dtype bf16 \ + --output artifacts/ws2-logprob/report.json +``` + +Run one TP=2, CP=2 Qwen3-vocab case on four local GPUs: + +```bash +torchrun --standalone --nproc-per-node=4 \ + rl_engine/testing/distributed_logprob_comparison.py \ + --tp 2 \ + --cp 2 \ + --dtype bf16 \ + --backend pytorch-vocab-parallel-logp-ws2 \ + --real-vocab 151936 \ + --padded-vocab 151936 \ + --num-vocab-tiles 64 \ + --batch 2 \ + --seq 16 \ + --prompt-tokens 8 \ + --output artifacts/ws2-logprob/tp2-cp2.json +``` + +The full matrix requires up to eight ranks for TP=4, CP=2. CPU/Gloo cases are +available for topology and artifact validation; the scoped numerical gate is +BF16 on CUDA/NCCL. + ## Minimal Example ```python @@ -206,11 +389,17 @@ out.sum().backward() python -m pytest tests/test_batch_invariant_logp.py -q -rs ``` -All backends (Native, Triton) are tested in a single file. Coverage includes: +All production backends are tested in a single file. Coverage includes correctness, leading-shape preservation, batch-invariance (bitwise), validation, ignore-index behavior, backward correctness, CUDA smoke cases, registry dispatch, and Triton-specific fp32/fp16/bf16 correctness, large vocab, backward -gradient batch-invariance, and ignored-row zero gradients. +gradient batch-invariance, and ignored-row zero gradients. The focused +`tests/test_logprob_comparison.py` suite covers TP=1 bitwise regression, direct +LSE identity, active-token drift statistics, structured serialization, exact +backend diagnostics, and fail-closed provenance. +`tests/test_distributed_logprob_comparison.py` covers topology planning, TP/CP +rank mapping, token/vocab sharding, explicit backend materialization, #116 JSON +artifacts, and a real four-process TP=2, CP=2 Gloo smoke case. Triton tests skip when Triton or CUDA is unavailable. On Windows, run via WSL/Linux with CUDA. @@ -223,4 +412,11 @@ WSL/Linux with CUDA. - `csrc/cuda/batch_invariant_logp_kernel_sm90.cu` - `rl_engine/kernels/registry.py` - `tests/test_batch_invariant_logp.py` +- `tests/test_logprob_comparison.py` +- `rl_engine/testing/logprob_drift.py` +- `rl_engine/testing/distributed_logprob_comparison.py` +- `tests/test_distributed_logprob_comparison.py` - `benchmarks/benchmark_batch_invariant_logp.py` +- `rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py` +- `rl_engine/kernels/logprob_contract.py` +- `tests/test_vocab_parallel_logp.py` diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py new file mode 100644 index 00000000..1267740a --- /dev/null +++ b/rl_engine/kernels/logprob_contract.py @@ -0,0 +1,655 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for TP-aware selected-token log-probability. + +The objects in this module describe a vocab-parallel logprob invocation: + +``selected_logp[t] = logits[t, target[t]] - logsumexp_vocab(logits[t, :])`` + +Under vocab-parallel tensor parallelism the vocabulary-wide ``logsumexp`` +requires cross-rank reduction. This module only *describes* that invocation +(shard ownership, merge semantics, mask/ignore-index metadata); it does not +shard tensors, launch collectives, or implement the ``(max, sumexp)`` merge. +Keeping description and materialization separate lets dispatch reject an +incompatible backend before any numerically different path is launched. + +Context parallelism is a declared non-merge axis: CP partitions tokens, never +the vocabulary, so the logprob reduction spans TP vocab shards only. CP rank +metadata is carried for provenance and must never widen the merge. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + +# Policy keywords accepted by KernelRegistry.get_logprob_op; a backend id must +# never shadow one of these, or it becomes unselectable by id. +RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) +# Backend tiers; determinism is a separate axis (DeterminismScope). +IMPLEMENTATION_KINDS = frozenset({"production", "reference"}) + + +class LogprobContractError(ValueError): + """Raised when logprob metadata does not describe a valid invocation.""" + + +class LogprobRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class LogprobDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class LogprobMerge(str, Enum): + """Merge primitive for per-shard ``(local_max, local_sumexp)`` partials.""" + + MAX_SUMEXP = "max_sumexp" + + +class MergeAxis(str, Enum): + """The only reduction axis of this contract; CP is a non-merge axis.""" + + TP_VOCAB = "tp_vocab" + + +class ReductionOrder(str, Enum): + GLOBAL_VOCAB_SHARD_INDEX = "global_vocab_shard_index" + + +class ReductionTransport(str, Enum): + """Collectives move partial states only; they never reduce numerically.""" + + ALL_GATHER = "all_gather" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class DeterminismScope(str, Enum): + """Strength of the reduction's determinism guarantee. + + ``fixed_topology``: bitwise-reproducible for one fixed TP degree; results + at different TP degrees are compared against the #108 tolerance table. + + ``cross_tp_bitwise``: additionally bitwise-equal across TP degrees. This + requires the entire reduction to follow a global tile-level structure that + is independent of TP partitioning: a fixed tile decomposition of the + vocabulary plus a fixed merge order and rescaling tree over those tiles, + identical at every TP degree, so the TP degree only selects which rank + computes which tiles and never changes the floating-point grouping. + Fixed shard-order merging alone is not sufficient, because shard + boundaries would still group the combines differently across degrees. + """ + + FIXED_TOPOLOGY = "fixed_topology" + CROSS_TP_BITWISE = "cross_tp_bitwise" + + +class MaskMode(str, Enum): + """How a backend consumes inactive-token information. + + The contract permits inactive targets that do not hold ``ignore_index``, + so an ``ignore_index``-only backend cannot serve a contract with inactive + tokens. + """ + + EXPLICIT_ACTIVE_MASK = "explicit_active_mask" + IGNORE_INDEX = "ignore_index" + + +class TPPlacement(str, Enum): + REPLICATED = "replicated" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise LogprobContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise LogprobContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _plain_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise LogprobContractError(f"{field} must be an integer; got {value!r}") + return value + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical vocab-parallel TP ownership for one logprob invocation. + + ``vocab_shard_bounds`` lists every TP rank's half-open ``[start, end)`` + vocab range, indexed by rank; the full table is required on every rank and + must form a contiguous ``[0, padded_vocab_size)`` partition. + ``padded_vocab_size`` is the shard-covered (weight) vocabulary, + ``real_vocab_size`` the tokenizer vocabulary; padding columns occupy + ``[real_vocab_size, padded_vocab_size)``. + """ + + tp_rank: int + tp_world_size: int + vocab_shard_bounds: tuple[tuple[int, int], ...] + real_vocab_size: int + padded_vocab_size: int + cp_rank: int = 0 + cp_world_size: int = 1 + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + if tp_rank >= tp_world_size: + raise LogprobContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if cp_rank >= cp_world_size: + raise LogprobContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + real_vocab_size = _positive_int(self.real_vocab_size, "real_vocab_size") + padded_vocab_size = _positive_int(self.padded_vocab_size, "padded_vocab_size") + if padded_vocab_size < real_vocab_size: + raise LogprobContractError( + f"padded_vocab_size={padded_vocab_size} must not be smaller than " + f"real_vocab_size={real_vocab_size}" + ) + + try: + bounds = tuple((pair[0], pair[1]) for pair in self.vocab_shard_bounds) + except (TypeError, IndexError) as exc: + raise LogprobContractError( + "vocab_shard_bounds must be an iterable of (start, end) integer pairs" + ) from exc + if len(bounds) != tp_world_size: + raise LogprobContractError( + "vocab_shard_bounds must declare exactly one (start, end) pair per TP rank; " + f"got {len(bounds)} pairs for tp_world_size={tp_world_size}" + ) + expected_start = 0 + for rank, (start, end) in enumerate(bounds): + start = _plain_int(start, f"vocab_shard_bounds[{rank}][0]") + end = _plain_int(end, f"vocab_shard_bounds[{rank}][1]") + if end <= start: + raise LogprobContractError( + f"vocab_shard_bounds[{rank}] must satisfy end > start; got [{start}, {end})" + ) + if start != expected_start: + raise LogprobContractError( + "vocab_shard_bounds must form a contiguous [0, padded_vocab_size) " + f"partition in TP-rank order; rank {rank} starts at {start}, " + f"expected {expected_start}" + ) + expected_start = end + if expected_start != padded_vocab_size: + raise LogprobContractError( + "vocab_shard_bounds must cover padded_vocab_size exactly; " + f"covered {expected_start}, declared {padded_vocab_size}" + ) + object.__setattr__(self, "vocab_shard_bounds", bounds) + + @property + def local_vocab_start(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][0] + + @property + def local_vocab_end(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][1] + + @property + def local_vocab_size(self) -> int: + start, end = self.vocab_shard_bounds[self.tp_rank] + return end - start + + def owner_rank(self, token_id: int) -> int: + """Return the unique TP rank owning ``token_id``; error outside real vocab.""" + + token_id = _plain_int(token_id, "token_id") + if token_id < 0 or token_id >= self.real_vocab_size: + raise LogprobContractError( + f"token_id={token_id} is outside the real vocabulary " + f"[0, {self.real_vocab_size}); mask it as inactive instead" + ) + for rank, (start, end) in enumerate(self.vocab_shard_bounds): + if start <= token_id < end: + return rank + raise LogprobContractError( + f"token_id={token_id} is not covered by any declared vocab shard" + ) + + +@dataclass(frozen=True) +class MaskSpec: + """Active-token mask and ignore index for one logprob invocation. + + Inactive tokens are excluded from drift aggregates and from the + single-owner target gather; their targets may legally hold ``ignore_index``. + """ + + num_tokens: int + active_mask: tuple[bool, ...] + ignore_index: int = -100 + _active_token_count: int = field(init=False, repr=False, compare=False) + _active_mask_sha256: str = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + num_tokens = _positive_int(self.num_tokens, "num_tokens") + _plain_int(self.ignore_index, "ignore_index") + try: + active_mask = tuple(self.active_mask) + except TypeError as exc: + raise LogprobContractError("active_mask must be an iterable of booleans") from exc + for index, value in enumerate(active_mask): + if not isinstance(value, bool): + raise LogprobContractError(f"active_mask[{index}] must be a bool; got {value!r}") + if len(active_mask) != num_tokens: + raise LogprobContractError( + "active_mask must contain exactly one entry per token; " + f"got {len(active_mask)} entries for num_tokens={num_tokens}" + ) + object.__setattr__(self, "active_mask", active_mask) + object.__setattr__(self, "_active_token_count", sum(active_mask)) + object.__setattr__( + self, "_active_mask_sha256", hashlib.sha256(bytes(active_mask)).hexdigest() + ) + + @property + def active_token_count(self) -> int: + return self._active_token_count + + @property + def active_mask_sha256(self) -> str: + """Compact mask identity for provenance and cross-rank agreement.""" + return self._active_mask_sha256 + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic TP-vocab ``(max, sumexp)`` merge semantics. + + Every rank first masks local columns whose global id lies in + ``[real_vocab_size, padded_vocab_size)`` to ``-inf`` (padding never + contributes to the logsumexp), then computes ``m_l = max(local_logits)`` + and ``s_l = sum(exp(local_logits - m_l))`` in fp32. Partials travel by + all-gather -- collectives are transport only, never a numerical + reduction -- and every rank merges in fixed global vocab-shard index + order:: + + M = max_l(m_l) + S = sum_l(s_l * exp(m_l - M)) + LSE = M + log(S) + selected_logp = target_logit - LSE + + The selected target logit comes from a masked single-owner gather; + downcast happens only at the final write. The identity partial for a + padding-only shard, or a row whose local columns are all ``-inf`` after + masking, is ``(m_l, s_l) = (-inf, 0)``: a partial with ``s_l = 0`` + contributes nothing to the merge regardless of its ``m_l``, and + implementations must use this identity directly rather than evaluate + ``exp(-inf - (-inf))``, which would poison the merge with NaN. Averaging + per-rank logsumexp values, or letting a collective reduce numerically, is + never conformant at either determinism scope. + """ + + merge: LogprobMerge = LogprobMerge.MAX_SUMEXP + merge_axis: MergeAxis = MergeAxis.TP_VOCAB + acc_dtype: LogprobDType = LogprobDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_VOCAB_SHARD_INDEX + transport: ReductionTransport = ReductionTransport.ALL_GATHER + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + determinism_scope: DeterminismScope = DeterminismScope.CROSS_TP_BITWISE + + def __post_init__(self) -> None: + object.__setattr__( + self, + "determinism_scope", + _enum_value(DeterminismScope, self.determinism_scope, "determinism_scope"), + ) + object.__setattr__(self, "merge", _enum_value(LogprobMerge, self.merge, "merge")) + object.__setattr__( + self, "merge_axis", _enum_value(MergeAxis, self.merge_axis, "merge_axis") + ) + object.__setattr__( + self, "acc_dtype", _enum_value(LogprobDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "transport", _enum_value(ReductionTransport, self.transport, "transport") + ) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"TP logprob accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class LogprobOutputSpec: + """Output surface every conforming backend must produce: fp32 selected + logprob and fp32 vocab-domain LSE, replicated across the TP group.""" + + selected_logp_dtype: LogprobDType = LogprobDType.FP32 + lse_dtype: LogprobDType = LogprobDType.FP32 + tp_placement: TPPlacement = TPPlacement.REPLICATED + + def __post_init__(self) -> None: + object.__setattr__( + self, + "selected_logp_dtype", + _enum_value(LogprobDType, self.selected_logp_dtype, "selected_logp_dtype"), + ) + object.__setattr__( + self, "lse_dtype", _enum_value(LogprobDType, self.lse_dtype, "lse_dtype") + ) + object.__setattr__( + self, "tp_placement", _enum_value(TPPlacement, self.tp_placement, "tp_placement") + ) + if self.selected_logp_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"selected logprob output must be fp32; got {self.selected_logp_dtype.value}" + ) + if self.lse_dtype is not LogprobDType.FP32: + raise LogprobContractError(f"vocab LSE output must be fp32; got {self.lse_dtype.value}") + + +@dataclass(frozen=True) +class LogprobContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: LogprobRole + dtype: LogprobDType + mask: MaskSpec + sharding: ShardingSpec + reduction: ReductionSpec + output: LogprobOutputSpec = field(default_factory=LogprobOutputSpec) + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(LogprobRole, self.role, "role")) + object.__setattr__(self, "dtype", _enum_value(LogprobDType, self.dtype, "dtype")) + if not isinstance(self.mask, MaskSpec): + raise LogprobContractError("mask must be a MaskSpec") + if not isinstance(self.sharding, ShardingSpec): + raise LogprobContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise LogprobContractError("reduction must be a ReductionSpec") + if not isinstance(self.output, LogprobOutputSpec): + raise LogprobContractError("output must be a LogprobOutputSpec") + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise LogprobContractError( + "export_lse must be True for the WS2 vocab-domain LSE drift contract" + ) + if 0 <= self.mask.ignore_index < self.sharding.real_vocab_size: + raise LogprobContractError( + f"ignore_index={self.mask.ignore_index} must not collide with the real " + f"vocabulary [0, {self.sharding.real_vocab_size})" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "vocab_shard_bounds": [list(pair) for pair in self.sharding.vocab_shard_bounds], + "real_vocab_size": self.sharding.real_vocab_size, + "padded_vocab_size": self.sharding.padded_vocab_size, + "local_vocab_start": self.sharding.local_vocab_start, + "local_vocab_end": self.sharding.local_vocab_end, + } + reduction = { + "merge": self.reduction.merge.value, + "merge_axis": self.reduction.merge_axis.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "transport": self.reduction.transport.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + "determinism_scope": self.reduction.determinism_scope.value, + "cp_is_merge_axis": False, + } + # The digest stands in for the raw per-token mask, which would + # dominate the provenance size. + mask = { + "num_tokens": self.mask.num_tokens, + "active_token_count": self.mask.active_token_count, + "active_mask_sha256": self.mask.active_mask_sha256, + "ignore_index": self.mask.ignore_index, + } + output = { + "selected_logp_dtype": self.output.selected_logp_dtype.value, + "lse_dtype": self.output.lse_dtype.value, + "tp_placement": self.output.tp_placement.value, + } + return { + "semantic_operator": "selected_token_logprob", + "role": self.role.value, + "dtype": self.dtype.value, + "export_lse": self.export_lse, + "lse_domain": "vocab", + "mask": mask, + "sharding": sharding, + "reduction": reduction, + "output": output, + } + + def cross_rank_fingerprint(self) -> str: + """Rank-independent identity for preflight agreement across ranks. + + Excludes ``tp_rank``/``cp_rank`` (and their derived local bounds) so + every rank of one logical invocation computes the same value. + All-gathering this fingerprint together with the resolved backend id + and aborting on mismatch is the documented preflight for distributed + dispatch; ``requested_backend="auto"`` is not distributed-safe + without it. + """ + + payload = self.to_dict() + payload["sharding"] = { + key: value + for key, value in payload["sharding"].items() + if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} + } + # Note: Any future extensions to this payload MUST maintain strict JSON + # serialization determinism across environments to prevent cross-rank + # hashing mismatches. + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class LogprobBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[LogprobRole] + dtypes: frozenset[LogprobDType] + tp_world_sizes: tuple[int, ...] | None = None + cp_world_sizes: tuple[int, ...] | None = None + supports_vocab_padding: bool = False + mask_modes: frozenset[MaskMode] = frozenset() + exports_vocab_lse: bool = False + determinism_scopes: frozenset[DeterminismScope] = frozenset() + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise LogprobContractError("backend_id must be a non-empty string") + if self.backend_id.strip().lower() in RESERVED_DISPATCH_POLICIES: + raise LogprobContractError( + f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" + ) + object.__setattr__(self, "backend_id", self.backend_id.strip()) + try: + roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) + dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) + except TypeError as exc: + raise LogprobContractError("roles and dtypes must be iterables of enum values") from exc + if not roles or not dtypes: + raise LogprobContractError("backend roles and dtypes must not be empty") + tp_world_sizes = self._validated_world_sizes(self.tp_world_sizes, "tp_world_sizes") + cp_world_sizes = self._validated_world_sizes(self.cp_world_sizes, "cp_world_sizes") + try: + mask_modes = frozenset( + _enum_value(MaskMode, value, "mask_modes") for value in self.mask_modes + ) + determinism_scopes = frozenset( + _enum_value(DeterminismScope, value, "determinism_scopes") + for value in self.determinism_scopes + ) + except TypeError as exc: + raise LogprobContractError( + "mask_modes and determinism_scopes must be iterables of enum values" + ) from exc + for flag_name in ("supports_vocab_padding", "exports_vocab_lse"): + if not isinstance(getattr(self, flag_name), bool): + raise LogprobContractError(f"{flag_name} must be a bool") + if self.implementation_kind not in IMPLEMENTATION_KINDS: + raise LogprobContractError( + f"implementation_kind must be one of: {', '.join(sorted(IMPLEMENTATION_KINDS))}" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "mask_modes", mask_modes) + object.__setattr__(self, "determinism_scopes", determinism_scopes) + + @staticmethod + def _validated_world_sizes( + values: tuple[int, ...] | None, field: str + ) -> tuple[int, ...] | None: + if values is None: + return None + try: + sizes = tuple(values) + except TypeError as exc: + raise LogprobContractError(f"{field} must be an iterable of integers") from exc + if not sizes: + raise LogprobContractError(f"{field} must not be empty; use None for unrestricted") + for value in sizes: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must contain positive values; got {value!r}") + if len(set(sizes)) != len(sizes): + raise LogprobContractError(f"{field} must not contain duplicates") + return sizes + + def incompatibilities(self, contract: LogprobContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if self.cp_world_sizes is not None and cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if ( + contract.sharding.padded_vocab_size != contract.sharding.real_vocab_size + and not self.supports_vocab_padding + ): + reasons.append("padded-vs-real vocab masking is unsupported") + if ( + contract.mask.active_token_count != contract.mask.num_tokens + and MaskMode.EXPLICIT_ACTIVE_MASK not in self.mask_modes + ): + # Inactive targets need not hold ignore_index (see MaskMode). + reasons.append("explicit active-token masking is unsupported") + if contract.export_lse and not self.exports_vocab_lse: + reasons.append("vocab-domain LSE export is unsupported") + if contract.reduction.determinism_scope not in self.determinism_scopes: + reasons.append( + f"determinism_scope={contract.reduction.determinism_scope.value} is unsupported" + ) + return tuple(reasons) + + def supports(self, contract: LogprobContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes) if self.cp_world_sizes else None, + "supports_vocab_padding": self.supports_vocab_padding, + "mask_modes": sorted(mode.value for mode in self.mask_modes), + "exports_vocab_lse": self.exports_vocab_lse, + "determinism_scopes": sorted(scope.value for scope in self.determinism_scopes), + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class LogprobDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: LogprobBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "IMPLEMENTATION_KINDS", + "RESERVED_DISPATCH_POLICIES", + "DeterminismScope", + "DowncastPoint", + "LogprobBackendCapability", + "LogprobContract", + "LogprobContractError", + "LogprobDType", + "LogprobDispatchResult", + "LogprobMerge", + "LogprobOutputSpec", + "LogprobRole", + "MaskMode", + "MaskSpec", + "MergeAxis", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "ReductionTransport", + "ShardingSpec", + "TPPlacement", +] diff --git a/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py index a68781d9..aa62b6ef 100644 --- a/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py @@ -161,3 +161,48 @@ def apply( ) return _BatchInvariantLogpSM90Function.apply(logits, target_ids, ignore_index) + + def forward_with_lse( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the exact SM90 path and return its direct FP32 logprob/LSE outputs. + + Unlike the production ``apply`` method, this diagnostic entry point never + falls back to Triton or PyTorch, so comparison provenance stays truthful. + """ + if logits.dim() < 2: + raise ValueError( + f"logits must be at least 2-D ([*lead, V]), got shape {tuple(logits.shape)}" + ) + if logits.shape[:-1] != target_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + if not _sm90_supported(logits): + raise RuntimeError( + "exact cuda-sm90 logprob diagnostics require Hopper, CUDA BF16/FP32 logits, " + "and a 16-byte-aligned vocab row stride; fallback is disabled" + ) + if validate: + vocab_size = logits.size(-1) + valid_targets = target_ids.reshape(-1) + valid_targets = valid_targets[valid_targets != ignore_index] + if valid_targets.numel() and ( + (valid_targets < 0).any() or (valid_targets >= vocab_size).any() + ): + bad = valid_targets[(valid_targets < 0) | (valid_targets >= vocab_size)] + raise ValueError( + f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}" + ) + + lead_shape = logits.shape[:-1] + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() + logp, lse = _C.batch_invariant_logp_sm90(logits_2d, target_1d, int(ignore_index)) + return logp.reshape(lead_shape), lse.reshape(lead_shape) diff --git a/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py index 4ac8bd37..80e08b5f 100644 --- a/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py @@ -45,24 +45,42 @@ def apply( logits_2d = logits.reshape(-1, vocab_size).float() target_1d = target_ids.reshape(-1).to(logits.device, dtype=torch.long) - selected_logp = self._row_wise_selected_logprob( + selected_logp, _ = self._row_wise_selected_logprob_with_lse( logits_2d, target_1d, ignore_index=ignore_index, validate=validate ) return selected_logp.reshape(lead_shape) + def forward_with_lse( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return selected logprob and the FP32 vocab-domain LSE for diagnostics.""" + self._validate_shapes(logits, target_ids) + lead_shape = logits.shape[:-1] + logits_2d = logits.reshape(-1, logits.size(-1)).float() + target_1d = target_ids.reshape(-1).to(logits.device, dtype=torch.long) + logp, lse = self._row_wise_selected_logprob_with_lse( + logits_2d, target_1d, ignore_index=ignore_index, validate=validate + ) + return logp.reshape(lead_shape), lse.reshape(lead_shape) + # ---------------------------------------------------------------------- # # Core Computation # ---------------------------------------------------------------------- # @staticmethod - def _row_wise_selected_logprob( + def _row_wise_selected_logprob_with_lse( logits_2d: torch.Tensor, target_1d: torch.Tensor, *, ignore_index: int, validate: bool = True, - ) -> torch.Tensor: - """Per-row selected logprob with locked reduction order. + ) -> tuple[torch.Tensor, torch.Tensor]: + """Per-row selected logprob and LSE with locked reduction order. The three reduction steps (max, sum-exp, gather) operate on each row independently. PyTorch's ``max(dim=-1)`` and ``sum(dim=-1)`` iterate @@ -104,7 +122,7 @@ def _row_wise_selected_logprob( selected_logp = selected_logp.where(valid_mask, torch.zeros_like(selected_logp)) - return selected_logp + return selected_logp, log_sum_exp # ---------------------------------------------------------------------- # # Helper diff --git a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py new file mode 100644 index 00000000..06eded1c --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic vocab-parallel TP selected-token logprob reference (issue #241 PR3). + +Implements the WS2 contract in ``rl_engine.kernels.logprob_contract`` with a +TP-independent vocab tile decomposition: the padded vocabulary is split into +``num_vocab_tiles`` fixed tiles, every tile's fp32 ``(max, sumexp)`` partial is +computed from a contiguous ``[n, tile]`` tensor, all tile partials travel by +all-gather (transport only), and every rank merges them in global tile-index +order over a fixed ``[n, num_vocab_tiles]`` shape. The TP degree only decides +which rank computes which tiles and never changes any floating-point grouping, +so outputs and gradients are bitwise-identical across TP degrees +(``DeterminismScope.CROSS_TP_BITWISE``) as long as ``num_vocab_tiles`` is held +fixed. A fixed per-shard merge order alone cannot provide this property: +shard boundaries would regroup the combines differently at each degree. + +Consequences of the tile structure: + +- ``num_vocab_tiles`` is part of the numerical identity. It must be pinned + across ranks (enforced by the preflight) and across the TP degrees being + compared; it is never derived from the shard layout. +- Every shard boundary must be tile-aligned; misalignment fails loudly. +- At TP=1 the result matches the WS1 ``NativeBatchInvariantLogpOp`` only + within the #108 logprob tolerance, not bitwise — the WS1 op reduces the + whole ``[n, V]`` row at once, which groups the sums differently. + +Preconditions: logits over the real vocabulary must be finite. A row whose +real-vocab logits are all ``-inf`` has no finite logsumexp; with +``validate=True`` such a row fails loudly if it is active. + +The selected logprob is zero-filled at inactive rows (``MaskSpec.active_mask`` +is the sole authority; with validation enabled an active row can never legally +hold ``ignore_index``). The vocab-domain LSE is returned for every row and is +differentiable everywhere, including inactive rows. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.kernels.logprob_contract import LogprobContract, LogprobContractError, LogprobDType + +BACKEND_ID = "pytorch-vocab-parallel-logp-ws2" +DEFAULT_NUM_VOCAB_TILES = 64 + +_TORCH_TO_CONTRACT_DTYPE = { + torch.bfloat16: LogprobDType.BF16, + torch.float16: LogprobDType.FP16, + torch.float32: LogprobDType.FP32, +} + + +def _require_distributed_initialized(): + import torch.distributed as dist + + if not dist.is_available(): + raise LogprobContractError("vocab-parallel logprob requires torch.distributed.") + if not dist.is_initialized(): + raise LogprobContractError( + "vocab-parallel logprob requires an initialized process group when " + "the contract declares tp_world_size > 1." + ) + return dist + + +def _tile_size(contract: LogprobContract, num_vocab_tiles: int) -> int: + if isinstance(num_vocab_tiles, bool) or not isinstance(num_vocab_tiles, int): + raise LogprobContractError( + f"num_vocab_tiles must be a positive integer; got {num_vocab_tiles!r}" + ) + if num_vocab_tiles <= 0: + raise LogprobContractError( + f"num_vocab_tiles must be a positive integer; got {num_vocab_tiles}" + ) + padded = contract.sharding.padded_vocab_size + if padded % num_vocab_tiles != 0: + raise LogprobContractError( + f"num_vocab_tiles={num_vocab_tiles} must divide " f"padded_vocab_size={padded} exactly" + ) + tile = padded // num_vocab_tiles + for rank, (start, end) in enumerate(contract.sharding.vocab_shard_bounds): + if start % tile != 0 or end % tile != 0: + raise LogprobContractError( + f"vocab_shard_bounds[{rank}]=[{start}, {end}) is not aligned to the " + f"vocab tile size {tile} (num_vocab_tiles={num_vocab_tiles}); " + "cross-TP bitwise determinism requires tile-aligned shard bounds" + ) + return tile + + +def _validate_invocation( + local_logits: torch.Tensor, + target_ids: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> None: + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if local_logits.dim() != 2: + raise LogprobContractError( + f"local_logits must be 2-D [num_tokens, local_vocab]; got {local_logits.dim()}-D" + ) + if target_ids.dim() != 1 or target_ids.shape[0] != local_logits.shape[0]: + raise LogprobContractError( + f"target_ids must be 1-D with one entry per token; got shape " + f"{tuple(target_ids.shape)} for {local_logits.shape[0]} tokens" + ) + sharding = contract.sharding + if local_logits.shape[1] != sharding.local_vocab_size: + raise LogprobContractError( + f"local_logits has {local_logits.shape[1]} vocab columns but the contract " + f"declares local shard [{sharding.local_vocab_start}, " + f"{sharding.local_vocab_end}) of size {sharding.local_vocab_size}" + ) + if local_logits.shape[0] != contract.mask.num_tokens: + raise LogprobContractError( + f"local_logits has {local_logits.shape[0]} tokens but MaskSpec declares " + f"num_tokens={contract.mask.num_tokens}" + ) + declared = _TORCH_TO_CONTRACT_DTYPE.get(local_logits.dtype) + if declared is not contract.dtype: + raise LogprobContractError( + f"local_logits dtype {local_logits.dtype} does not match the contract " + f"dtype {contract.dtype.value}" + ) + if sharding.tp_world_size > 1: + dist = _require_distributed_initialized() + group_rank = dist.get_rank(group=tp_group) + group_world = dist.get_world_size(group=tp_group) + if group_world != sharding.tp_world_size: + raise LogprobContractError( + f"tp_group world size {group_world} does not match the contract " + f"tp_world_size={sharding.tp_world_size}; pass the TP subgroup, " + "not the global group" + ) + if group_rank != sharding.tp_rank: + raise LogprobContractError( + f"tp_group rank {group_rank} does not match the contract " + f"tp_rank={sharding.tp_rank}" + ) + + +def _validate_active_targets( + target_1d: torch.Tensor, active_mask: torch.Tensor, real_vocab_size: int +) -> None: + bad = active_mask & ((target_1d < 0) | (target_1d >= real_vocab_size)) + if bool(bad.any().item()): + bad_values = target_1d[bad] + raise LogprobContractError( + "active target_ids must lie in the real vocabulary " + f"[0, {real_vocab_size}); got values in " + f"[{int(bad_values.min().item())}, {int(bad_values.max().item())}] " + "on active rows" + ) + + +def _preflight_cross_rank_agreement( + contract: LogprobContract, tp_group: Any, num_vocab_tiles: int +) -> None: + """All-gather (fingerprint, backend id, tile count) and abort on mismatch.""" + + dist = _require_distributed_initialized() + payload = (contract.cross_rank_fingerprint(), BACKEND_ID, int(num_vocab_tiles)) + world = dist.get_world_size(group=tp_group) + gathered: list[Any] = [None] * world + dist.all_gather_object(gathered, payload, group=tp_group) + mismatched = [(rank, other) for rank, other in enumerate(gathered) if other != payload] + if mismatched: + rank, other = mismatched[0] + raise LogprobContractError( + "cross-rank preflight failed: rank " + f"{contract.sharding.tp_rank} has {payload} but rank {rank} has {other}; " + "all TP ranks must agree on the contract fingerprint, backend id, and " + "num_vocab_tiles before any collective" + ) + + +def _local_tile_stats(z_masked: torch.Tensor, tile: int) -> tuple[torch.Tensor, torch.Tensor]: + """fp32 per-tile ``(max, sumexp)`` partials for this rank's shard. + + Each tile is reduced as a contiguous ``[n, tile]`` tensor so the reduction + shape and layout are identical no matter which rank computes the tile or + what the local shard size is. An all-``-inf`` (padding-only) tile yields + the identity partial ``(-inf, 0)`` without evaluating ``exp(-inf - (-inf))``. + """ + + n, local_vocab = z_masked.shape + m_parts: list[torch.Tensor] = [] + s_parts: list[torch.Tensor] = [] + for tile_index in range(local_vocab // tile): + block = z_masked[:, tile_index * tile : (tile_index + 1) * tile].contiguous() + m_t = block.max(dim=-1).values + finite = m_t > float("-inf") + m_safe = torch.where(finite, m_t, torch.zeros_like(m_t)) + s_t = (block - m_safe.unsqueeze(-1)).exp().sum(dim=-1) + s_t = torch.where(finite, s_t, torch.zeros_like(s_t)) + m_parts.append(m_t) + s_parts.append(s_t) + return torch.stack(m_parts, dim=1), torch.stack(s_parts, dim=1) + + +def _gather_tile_stats( + local_m: torch.Tensor, + local_s: torch.Tensor, + contract: LogprobContract, + tp_group: Any, + tile: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Assemble all ``num_vocab_tiles`` partials in global tile order.""" + + sharding = contract.sharding + tile_counts = [(end - start) // tile for start, end in sharding.vocab_shard_bounds] + if sharding.tp_world_size == 1: + return local_m.contiguous(), local_s.contiguous() + + dist = _require_distributed_initialized() + n = local_m.shape[0] + max_tiles = max(tile_counts) + packed = local_m.new_zeros((n, max_tiles, 2)) + packed[:, : local_m.shape[1], 0] = local_m + packed[:, : local_s.shape[1], 1] = local_s + packed = packed.contiguous() + gathered = [torch.empty_like(packed) for _ in range(sharding.tp_world_size)] + dist.all_gather(gathered, packed, group=tp_group) + + m_parts = [gathered[rank][:, : tile_counts[rank], 0] for rank in range(len(tile_counts))] + s_parts = [gathered[rank][:, : tile_counts[rank], 1] for rank in range(len(tile_counts))] + return torch.cat(m_parts, dim=1).contiguous(), torch.cat(s_parts, dim=1).contiguous() + + +def _gather_target_logit( + z_masked: torch.Tensor, + safe_target: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> torch.Tensor: + """Exact selected-target logit via a select-by-owner copy.""" + + sharding = contract.sharding + n = z_masked.shape[0] + start = sharding.local_vocab_start + local_vocab = sharding.local_vocab_size + local_idx = (safe_target - start).clamp(0, max(local_vocab - 1, 0)) + owns = (safe_target >= start) & (safe_target < sharding.local_vocab_end) + rows = torch.arange(n, device=z_masked.device) + local_contrib = torch.where( + owns, z_masked[rows, local_idx], torch.zeros_like(safe_target, dtype=z_masked.dtype) + ).contiguous() + + if sharding.tp_world_size == 1: + stacked = local_contrib.unsqueeze(0) + else: + dist = _require_distributed_initialized() + gathered = [torch.empty_like(local_contrib) for _ in range(sharding.tp_world_size)] + dist.all_gather(gathered, local_contrib, group=tp_group) + stacked = torch.stack(gathered, dim=0) + + starts = torch.tensor( + [bound_start for bound_start, _ in sharding.vocab_shard_bounds], + device=safe_target.device, + dtype=torch.long, + ) + owner = torch.bucketize(safe_target, starts, right=True) - 1 + return stacked[owner, rows] + + +def _merge_tile_partials(m_all: torch.Tensor, s_all: torch.Tensor) -> torch.Tensor: + """Fixed-order (max, sumexp) merge over [n, num_vocab_tiles].""" + + M = m_all.max(dim=1).values + finite = M > float("-inf") + M_safe = torch.where(finite, M, torch.zeros_like(M)) + terms = s_all * (m_all - M_safe.unsqueeze(1)).exp() + S = terms.sum(dim=1) + return M + S.log() + + +class _VocabParallelLogprobFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, local_logits, target_1d, active_mask, contract, tp_group, tile): + z_masked = local_logits.float() + sharding = contract.sharding + global_ids = torch.arange( + sharding.local_vocab_start, sharding.local_vocab_end, device=z_masked.device + ) + padding_cols = global_ids >= sharding.real_vocab_size + if bool(padding_cols.any()): + z_masked = z_masked.masked_fill(padding_cols.unsqueeze(0), float("-inf")) + + safe_target = torch.where(active_mask, target_1d, torch.zeros_like(target_1d)) + + local_m, local_s = _local_tile_stats(z_masked, tile) + m_all, s_all = _gather_tile_stats(local_m, local_s, contract, tp_group, tile) + target_logit = _gather_target_logit(z_masked, safe_target, contract, tp_group) + lse = _merge_tile_partials(m_all, s_all) + + selected_logp = torch.where(active_mask, target_logit - lse, torch.zeros_like(lse)) + + ctx.save_for_backward(z_masked, lse, safe_target, active_mask, padding_cols) + ctx.local_vocab_start = sharding.local_vocab_start + ctx.local_vocab_size = sharding.local_vocab_size + ctx.input_dtype = local_logits.dtype + ctx.set_materialize_grads(False) + return selected_logp, lse + + @staticmethod + def backward(ctx, grad_logp, grad_lse): + if not ctx.needs_input_grad[0] or (grad_logp is None and grad_lse is None): + return None, None, None, None, None, None + + z_masked, lse, safe_target, active_mask, padding_cols = ctx.saved_tensors + n, local_vocab = z_masked.shape + finite_row = torch.isfinite(lse) + lse_safe = torch.where(finite_row, lse, torch.zeros_like(lse)) + p = (z_masked - lse_safe.unsqueeze(1)).exp() + p = torch.where(finite_row.unsqueeze(1), p, torch.zeros_like(p)) + + grad = torch.zeros_like(z_masked) + if grad_logp is not None: + local_idx = (safe_target - ctx.local_vocab_start).clamp(0, max(local_vocab - 1, 0)) + owns = (safe_target >= ctx.local_vocab_start) & ( + safe_target < ctx.local_vocab_start + local_vocab + ) + onehot = torch.zeros_like(z_masked) + hit = owns & active_mask + rows = torch.arange(n, device=z_masked.device)[hit] + onehot[rows, local_idx[hit]] = 1.0 + g_logp = torch.where(active_mask, grad_logp, torch.zeros_like(grad_logp)) + grad = grad + g_logp.unsqueeze(1) * (onehot - p) + if grad_lse is not None: + grad = grad + grad_lse.unsqueeze(1) * p + if bool(padding_cols.any()): + grad = grad.masked_fill(padding_cols.unsqueeze(0), 0.0) + return grad.to(ctx.input_dtype), None, None, None, None, None + + +class VocabParallelLogprobOp: + """Deterministic vocab-parallel selected-token logprob (WS2 reference).""" + + op_class = "logprob" + is_batch_invariant = True + + def __init__(self) -> None: + pass + + def __call__( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + ) + + def apply( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + tile = _tile_size(contract, num_vocab_tiles) + _validate_invocation(local_logits, target_ids, contract, tp_group) + + target_1d = target_ids.reshape(-1).to(device=local_logits.device, dtype=torch.long) + active_mask = torch.tensor( + contract.mask.active_mask, dtype=torch.bool, device=local_logits.device + ) + if validate: + _validate_active_targets(target_1d, active_mask, contract.sharding.real_vocab_size) + if contract.sharding.tp_world_size > 1: + _preflight_cross_rank_agreement(contract, tp_group, num_vocab_tiles) + + selected_logp, lse = _VocabParallelLogprobFunction.apply( + local_logits, target_1d, active_mask, contract, tp_group, tile + ) + + if validate and bool((~torch.isfinite(lse) & active_mask).any().item()): + raise LogprobContractError( + "non-finite logsumexp on an active row: logits over the real " + "vocabulary must be finite for every active token" + ) + return selected_logp, lse + + +__all__ = [ + "BACKEND_ID", + "DEFAULT_NUM_VOCAB_TILES", + "VocabParallelLogprobOp", +] diff --git a/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py index 66b99757..804341f3 100644 --- a/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py @@ -10,6 +10,27 @@ _BLOCK_V: int = 1024 +def _launch_batch_invariant_logp( + logits_2d: torch.Tensor, target_1d: torch.Tensor, ignore_index: int +) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = logits_2d.shape[0] + vocab_size = logits_2d.shape[1] + output = torch.empty(num_tokens, device=logits_2d.device, dtype=torch.float32) + lse = torch.empty(num_tokens, device=logits_2d.device, dtype=torch.float32) + _batch_invariant_logp_kernel[(num_tokens,)]( + logits_2d, + target_1d, + output, + lse, + num_tokens, + vocab_size, + logits_2d.stride(0), + ignore_index=ignore_index, + BLOCK_V=_BLOCK_V, + ) + return output, lse + + @triton.jit def _batch_invariant_logp_kernel( logits_ptr, # logits [N, V] @@ -126,22 +147,7 @@ def forward(ctx, logits, target_ids, ignore_index): logits_2d = logits.reshape(-1, vocab_size).contiguous() target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() - num_tokens = logits_2d.shape[0] - output = torch.empty(num_tokens, device=logits.device, dtype=torch.float32) - lse = torch.empty(num_tokens, device=logits.device, dtype=torch.float32) - - grid = (num_tokens,) - _batch_invariant_logp_kernel[grid]( - logits_2d, - target_1d, - output, - lse, - num_tokens, - vocab_size, - logits_2d.stride(0), - ignore_index=ignore_index, - BLOCK_V=_BLOCK_V, - ) + output, lse = _launch_batch_invariant_logp(logits_2d, target_1d, ignore_index) ctx.save_for_backward(logits_2d, target_1d, lse) ctx.ignore_index = ignore_index @@ -237,3 +243,53 @@ def apply( ) return _BatchInvariantLogpFunction.apply(logits, target_ids, ignore_index) + + def forward_with_lse( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return direct FP32 logprob/LSE outputs without an autograd wrapper.""" + self._validate_inputs(logits, target_ids, ignore_index=ignore_index, validate=validate) + lead_shape = logits.shape[:-1] + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() + logp, lse = _launch_batch_invariant_logp(logits_2d, target_1d, ignore_index) + return logp.reshape(lead_shape), lse.reshape(lead_shape) + + @staticmethod + def _validate_inputs( + logits: torch.Tensor, + target_ids: torch.Tensor, + *, + ignore_index: int, + validate: bool, + ) -> None: + if logits.device.type not in ("cuda", "xpu", "hip"): + raise RuntimeError( + "TritonBatchInvariantLogpOp requires a GPU tensor " + f"(CUDA / ROCm / XPU), got device '{logits.device}'." + ) + if logits.dim() < 2: + raise ValueError( + f"logits must be at least 2-D ([*lead, V]), got shape {tuple(logits.shape)}" + ) + if logits.shape[:-1] != target_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + if validate: + vocab_size = logits.size(-1) + valid_targets = target_ids.reshape(-1) + valid_targets = valid_targets[valid_targets != ignore_index] + if valid_targets.numel() and ( + (valid_targets < 0).any() or (valid_targets >= vocab_size).any() + ): + bad = valid_targets[(valid_targets < 0) | (valid_targets >= vocab_size)] + raise ValueError( + f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}" + ) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index efde5c25..bd7ba45c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -8,6 +8,17 @@ import torch +from rl_engine.kernels.logprob_contract import ( + IMPLEMENTATION_KINDS, + DeterminismScope, + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDispatchResult, + LogprobDType, + LogprobRole, + MaskMode, +) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -74,6 +85,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_BATCH_INVARIANT_LOGP_SM90 = ( "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op" ) + # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) + PYTORCH_VOCAB_PARALLEL_LOGP = ( + "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" + ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -168,6 +183,47 @@ def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() + # Truthful descriptors for the existing WS1 batch-invariant logp + # implementations: single-shard (TP=1), ignore-index masking only, no + # vocab-shard metadata, no vocab-domain LSE export. + common_logprob_roles = frozenset({LogprobRole.TRAIN, LogprobRole.INFER}) + common_logprob_dtypes = frozenset({LogprobDType.BF16, LogprobDType.FP16, LogprobDType.FP32}) + base_logprob_capabilities = { + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="pytorch-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="reference", + ), + OpBackend.TRITON_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="triton-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", + ), + OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90: LogprobBackendCapability( + backend_id="cuda-batch-invariant-logp-sm90-ws1", + roles=common_logprob_roles, + dtypes=frozenset({LogprobDType.BF16, LogprobDType.FP32}), + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", + ), + } + self._priority_map = { "cuda": { "logp": [ @@ -292,6 +348,47 @@ def __init__(self): self._adjust_priority_for_hardware() self._adjust_priority_from_env() + # WS2 dispatch owns its candidate list, seeded from the legacy + # batch_invariant_logp priority but decoupled afterwards: neither + # path's registrations may affect the other. + self._logprob_candidates: Dict[str, list] = { + platform: list(ops.get("batch_invariant_logp", [])) + for platform, ops in self._priority_map.items() + } + # Capabilities are scoped per platform: the same backend enum may + # truthfully declare different support on cuda vs rocm vs cpu. + self._logprob_capabilities: Dict[str, Dict[OpBackend, LogprobBackendCapability]] = { + platform: { + backend: base_logprob_capabilities[backend] + for backend in candidates + if backend in base_logprob_capabilities + } + for platform, candidates in self._logprob_candidates.items() + } + + # deterministic vocab-parallel TP logprob reference. + ws2_tp_logprob_capability = LogprobBackendCapability( + backend_id="pytorch-vocab-parallel-logp-ws2", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=None, + cp_world_sizes=None, + supports_vocab_padding=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), + exports_vocab_lse=True, + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", + ) + for ws2_platform in self._priority_map: + self.register_logprob_backend( + OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP, + ws2_tp_logprob_capability, + platform=ws2_platform, + prepend=True, + ) + def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() if rocm_attn_backend in {"flash_attn", "flash-attn", "flash_attention"}: @@ -381,25 +478,188 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: - if backend.name in self._instance_cache: - return self._instance_cache[backend.name] + op_instance = self._get_or_create_backend(backend) + if op_instance is not None: + return op_instance - if backend.name in self._failed_backends: - continue + raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + + def register_logprob_backend( + self, + backend: OpBackend, + capability: LogprobBackendCapability, + *, + platform: Optional[str] = None, + prepend: bool = False, + ) -> None: + """Register (or replace) a backend for WS2 contract-aware logprob dispatch. + + This is the supported seam for making a new backend selectable by + ``get_logprob_op`` (e.g. the deterministic vocab-parallel TP reference + from issue #241 PR 3) without touching the legacy ``get_op`` priority + lists. Registering the same backend again replaces its capability + without duplicating the candidate entry. + """ - op_class = self._load_backend(backend) - if op_class: - try: - op_instance = op_class() - self._instance_cache[backend.name] = op_instance - return op_instance - except Exception as e: - logger.error(f"Failed to instantiate {backend.name}: {e}") - self._failed_backends.add(backend.name) + if not isinstance(backend, OpBackend): + raise LogprobContractError("backend must be an OpBackend") + if not isinstance(capability, LogprobBackendCapability): + raise LogprobContractError("capability must be a LogprobBackendCapability") + resolved_platform = platform if platform is not None else self._platform() + if resolved_platform not in self._priority_map: + raise LogprobContractError( + f"unsupported platform {resolved_platform!r}; expected one of " + f"{sorted(self._priority_map)}" + ) + candidates = self._logprob_candidates.setdefault(resolved_platform, []) + self._logprob_capabilities.setdefault(resolved_platform, {})[backend] = capability + if backend not in candidates: + if prepend: + candidates.insert(0, backend) else: - self._failed_backends.add(backend.name) + candidates.append(backend) + + def get_logprob_op( + self, + contract: LogprobContract, + *, + requested_backend: str = "auto", + ) -> LogprobDispatchResult: + """Resolve only a backend that explicitly supports the WS2 logprob contract. + + This entry point is intentionally separate from legacy ``get_op`` so + existing callers retain their current behavior while WS2 callers cannot + silently fall back to a backend with different distributed semantics. + + ``requested_backend`` is either a case-insensitive policy keyword + (``auto`` | ``production`` | ``reference`` | ``deterministic``) or an + exact, case-sensitive stable backend id. Strictness comes from the + contract's capability checks, not from this policy string, so the + default is ``auto``. With ``tp_world_size > 1``, ``auto`` is rejected: + per-rank auto resolution can diverge across ranks, so distributed + callers must name a policy or backend id and preflight agreement via + ``LogprobContract.cross_rank_fingerprint``. + """ - raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise LogprobContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip() + if requested_backend.lower() == "deterministic": + raise LogprobContractError( + 'requested_backend="deterministic" is not a dispatch policy; request ' + "determinism through ReductionSpec.determinism_scope and match it against " + "backend determinism_scopes instead" + ) + if requested_backend.lower() == "auto" and contract.sharding.tp_world_size > 1: + raise LogprobContractError( + "Unsafe dispatch: requested_backend='auto' is not permitted when tp_world_size > 1 " + "without explicit cross-rank preflighting." + ) + + platform = self._platform() + candidates = self._logprob_candidates.get(platform, []) + rejected: list[str] = [] + # provenance["fallback"] reports only capability/load rejections of + # otherwise-eligible candidates; skips caused purely by the caller's + # own requested_backend policy filter are not fallbacks. + capability_rejections = 0 + + platform_capabilities = self._logprob_capabilities.get(platform, {}) + for backend in candidates: + capability = platform_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no LogprobBackendCapability declared") + capability_rejections += 1 + continue + policy_mismatch = self._logprob_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + # Excluded by the caller's own policy: never a fallback, even + # if the candidate would also have failed capability checks. + rejected.append(f"{backend.name}: {policy_mismatch}") + continue + capability_incompat = list(capability.incompatibilities(contract)) + if capability_incompat: + rejected.append(f"{backend.name}: " + "; ".join(capability_incompat)) + capability_rejections += 1 + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + capability_rejections += 1 + continue + + provenance = { + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": capability_rejections > 0, + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + } + return LogprobDispatchResult( + op=op, + capability=capability, + provenance=provenance, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No logprob backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, dtype={requested['dtype']}, " + f"TP={contract.sharding.tp_world_size}, CP={contract.sharding.cp_world_size}, " + f"padded_vocab={contract.sharding.padded_vocab_size}, " + f"real_vocab={contract.sharding.real_vocab_size}. Rejections: {details}" + ) + + @staticmethod + def _logprob_policy_mismatch( + requested_backend: str, + capability: LogprobBackendCapability, + ) -> str | None: + policy = requested_backend.lower() + if policy == "auto": + return None + if policy in IMPLEMENTATION_KINDS: + if capability.implementation_kind == policy: + return None + return ( + f"implementation_kind={capability.implementation_kind} does not satisfy " + f"requested_backend={policy}" + ) + if capability.backend_id == requested_backend: + return None + return ( + f"backend_id={capability.backend_id} does not match " + f"requested_backend={requested_backend}" + ) + + def _platform(self) -> str: + return self._platform_for_device(None) + + def _get_or_create_backend(self, backend: OpBackend) -> Any | None: + if backend.name in self._instance_cache: + return self._instance_cache[backend.name] + if backend.name in self._failed_backends: + return None + + op_class = self._load_backend(backend) + if op_class is None: + self._failed_backends.add(backend.name) + return None + try: + op = op_class() + except Exception as exc: + logger.error(f"Failed to instantiate {backend.name}: {exc}") + self._failed_backends.add(backend.name) + return None + self._instance_cache[backend.name] = op + return op def _platform_for_device(self, device: torch.device | str | None) -> str: if device is None: diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 42be8c1b..97b585fb 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -3,6 +3,16 @@ """Testing helpers for RL-shaped kernel validation.""" +from .logprob_comparison import ( + LogprobBackendUnavailable, + LogprobCandidate, + LogprobComparisonInputs, + LogprobComparisonReport, + compare_single_gpu_logprob, + make_logprob_candidate, + route_rl_kernel_logs_to_stderr, +) +from .logprob_drift import LogprobDriftStats, summarize_logprob_drift from .reference_ops import ( active_token_count, compute_policy_ratio, @@ -15,13 +25,22 @@ from .rl_batch import SyntheticRLKernelBatch, make_synthetic_rl_kernel_batch __all__ = [ + "LogprobBackendUnavailable", + "LogprobCandidate", + "LogprobComparisonInputs", + "LogprobComparisonReport", + "LogprobDriftStats", "SyntheticRLKernelBatch", "active_token_count", + "compare_single_gpu_logprob", "compute_policy_ratio", "compute_reference_kl", + "make_logprob_candidate", "make_synthetic_rl_kernel_batch", "masked_mean", "masked_sum", "selected_logprobs_reference", + "route_rl_kernel_logs_to_stderr", + "summarize_logprob_drift", "summarize_kernel_drift", ] diff --git a/rl_engine/testing/distributed_logprob_comparison.py b/rl_engine/testing/distributed_logprob_comparison.py new file mode 100644 index 00000000..3e37c88b --- /dev/null +++ b/rl_engine/testing/distributed_logprob_comparison.py @@ -0,0 +1,843 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Distributed WS2 comparison for the vocab-parallel logprob reference.""" + +from __future__ import annotations + +import argparse +import contextlib +import datetime +import hashlib +import json +import os +import pathlib +import shlex +import sys +from dataclasses import asdict, dataclass +from typing import Any, Sequence + +import torch + +if __package__ in (None, ""): + repo_root = pathlib.Path(__file__).resolve().parents[2] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + +_import_output = ( + contextlib.redirect_stdout(sys.stderr) + if __package__ in (None, "") + else contextlib.nullcontext() +) +with _import_output: + from rl_engine.kernels.gtest.tolerance import load_contract as load_tolerance_contract + from rl_engine.kernels.logprob_contract import ( + LogprobContract, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, + ) + from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + BACKEND_ID, + DEFAULT_NUM_VOCAB_TILES, + ) + from rl_engine.kernels.registry import KernelRegistry + from rl_engine.testing.logprob_comparison import route_rl_kernel_logs_to_stderr + from rl_engine.testing.logprob_drift import LogprobDriftStats, summarize_logprob_drift + +_DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} +_TOLERANCE_DTYPES = { + "bf16": "bfloat16", + "fp16": "float16", + "fp32": "float32", +} +_PROCESS_GROUP_TIMEOUT = datetime.timedelta(minutes=5) +_RELATIVE_ERROR_FLOOR = 1.0e-12 + + +@dataclass(frozen=True) +class DistributedLogprobCase: + tp_world_size: int + cp_world_size: int + dtype: str = "bf16" + requested_backend: str = BACKEND_ID + real_vocab_size: int = 151936 + padded_vocab_size: int = 151936 + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES + batch_size: int = 2 + sequence_length: int = 16 + prompt_tokens: int = 8 + seed: int = 123 + ignore_index: int = -100 + + def __post_init__(self) -> None: + for name in ( + "tp_world_size", + "cp_world_size", + "real_vocab_size", + "padded_vocab_size", + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + if self.dtype not in _DTYPES: + raise ValueError(f"dtype must be one of {sorted(_DTYPES)}") + if not self.requested_backend or self.requested_backend.lower() == "auto": + raise ValueError("distributed cases require an explicit non-auto backend") + if self.padded_vocab_size < self.real_vocab_size: + raise ValueError("padded_vocab_size must be at least real_vocab_size") + if self.num_vocab_tiles < self.tp_world_size: + raise ValueError("num_vocab_tiles must be at least tp_world_size") + if self.padded_vocab_size % self.num_vocab_tiles != 0: + raise ValueError("num_vocab_tiles must divide padded_vocab_size exactly") + if self.batch_size <= 0 or self.sequence_length <= 0: + raise ValueError("batch_size and sequence_length must be positive") + if not 0 <= self.prompt_tokens <= self.sequence_length: + raise ValueError("prompt_tokens must be in [0, sequence_length]") + + @property + def world_size(self) -> int: + return self.tp_world_size * self.cp_world_size + + @property + def num_tokens(self) -> int: + return self.batch_size * self.sequence_length + + @property + def case_id(self) -> str: + encoded = json.dumps(asdict(self), sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest()[:16] + + +@dataclass(frozen=True) +class RankTopology: + global_rank: int + world_size: int + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + tp_group_ranks: tuple[int, ...] + + +@dataclass(frozen=True) +class DriftDetail: + stats: LogprobDriftStats + max_rel: float + worst_global_token: int | None + worst_target_id: int | None + worst_owner_rank: int | None + candidate_value: float | None + reference_value: float | None + atol: float + rtol: float + passed: bool + + +@dataclass(frozen=True) +class RankLogprobReport: + global_rank: int + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + sp_world_size: int + dp_world_size: int + token_start: int + token_end: int + vocab_start: int + vocab_end: int + device: str + requested_backend: str + actual_backend: str + fallback: bool + contract_fingerprint: str + contract: dict[str, Any] + capability: dict[str, Any] + tp_outputs_bitwise_replicated: bool + lse: DriftDetail + dlogp: DriftDetail + passed: bool + + +@dataclass(frozen=True) +class DistributedLogprobReport: + schema_version: int + case_id: str + case: dict[str, Any] + launch_command: str + environment: dict[str, Any] + ranks: tuple[RankLogprobReport, ...] + aggregate: dict[str, DriftDetail] + bitwise_fingerprints: dict[str, Any] + passed: bool + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class _RankPayload: + report: RankLogprobReport + candidate_logp: torch.Tensor + candidate_lse: torch.Tensor + reference_logp: torch.Tensor + reference_lse: torch.Tensor + active_mask: torch.Tensor + target_ids: torch.Tensor + global_positions: torch.Tensor + + +def _strict_report_json(payload: dict[str, Any]) -> str: + return json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + + +def plan_distributed_logprob_cases( + *, + tp_world_sizes: Sequence[int] = (1, 2, 4), + cp_world_sizes: Sequence[int] = (1, 2), + **overrides: Any, +) -> tuple[DistributedLogprobCase, ...]: + """Build the scoped issue #241 topology product in deterministic order.""" + + return tuple( + DistributedLogprobCase(tp_world_size=tp, cp_world_size=cp, **overrides) + for tp in tp_world_sizes + for cp in cp_world_sizes + ) + + +def rank_topology(case: DistributedLogprobCase, global_rank: int) -> RankTopology: + if not 0 <= global_rank < case.world_size: + raise ValueError(f"global_rank must be in [0, {case.world_size})") + cp_rank, tp_rank = divmod(global_rank, case.tp_world_size) + group_start = cp_rank * case.tp_world_size + return RankTopology( + global_rank=global_rank, + world_size=case.world_size, + tp_rank=tp_rank, + tp_world_size=case.tp_world_size, + cp_rank=cp_rank, + cp_world_size=case.cp_world_size, + tp_group_ranks=tuple(range(group_start, group_start + case.tp_world_size)), + ) + + +def token_shard_bounds(num_tokens: int, cp_world_size: int) -> tuple[tuple[int, int], ...]: + """Partition token rows contiguously, allowing a one-row imbalance.""" + + if num_tokens < cp_world_size: + raise ValueError("num_tokens must be at least cp_world_size") + quotient, remainder = divmod(num_tokens, cp_world_size) + bounds = [] + cursor = 0 + for cp_rank in range(cp_world_size): + count = quotient + int(cp_rank < remainder) + bounds.append((cursor, cursor + count)) + cursor += count + return tuple(bounds) + + +def vocab_shard_bounds(case: DistributedLogprobCase) -> tuple[tuple[int, int], ...]: + """Assign complete global vocab tiles to TP ranks.""" + + tile_size = case.padded_vocab_size // case.num_vocab_tiles + quotient, remainder = divmod(case.num_vocab_tiles, case.tp_world_size) + bounds = [] + cursor_tiles = 0 + for tp_rank in range(case.tp_world_size): + tile_count = quotient + int(tp_rank < remainder) + start = cursor_tiles * tile_size + cursor_tiles += tile_count + bounds.append((start, cursor_tiles * tile_size)) + return tuple(bounds) + + +def format_launch_command( + case: DistributedLogprobCase, + *, + output: str | pathlib.Path, + device: str = "cuda", + dist_backend: str | None = None, +) -> str: + backend = dist_backend or ("nccl" if device == "cuda" else "gloo") + arguments = [ + "torchrun", + "--standalone", + f"--nproc-per-node={case.world_size}", + "rl_engine/testing/distributed_logprob_comparison.py", + "--tp", + str(case.tp_world_size), + "--cp", + str(case.cp_world_size), + "--dtype", + case.dtype, + "--backend", + case.requested_backend, + "--real-vocab", + str(case.real_vocab_size), + "--padded-vocab", + str(case.padded_vocab_size), + "--num-vocab-tiles", + str(case.num_vocab_tiles), + "--batch", + str(case.batch_size), + "--seq", + str(case.sequence_length), + "--prompt-tokens", + str(case.prompt_tokens), + "--seed", + str(case.seed), + "--device", + device, + "--dist-backend", + backend, + "--output", + str(output), + ] + return shlex.join(arguments) + + +def _canonical_inputs( + case: DistributedLogprobCase, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(case.seed) + logits = torch.randn( + case.num_tokens, + case.padded_vocab_size, + generator=generator, + dtype=torch.float32, + ) + targets = torch.randint( + 0, + case.real_vocab_size, + (case.num_tokens,), + generator=generator, + dtype=torch.long, + ) + active = torch.ones((case.batch_size, case.sequence_length), dtype=torch.bool) + active[:, : case.prompt_tokens] = False + active = active.reshape(-1) + targets = targets.masked_fill(~active, case.ignore_index) + return logits, targets, active + + +def _make_contract( + case: DistributedLogprobCase, + topology: RankTopology, + active_mask: torch.Tensor, +) -> LogprobContract: + return LogprobContract( + role=LogprobRole.TRAIN, + dtype=LogprobDType(case.dtype), + mask=MaskSpec( + num_tokens=int(active_mask.numel()), + active_mask=tuple(bool(value) for value in active_mask.tolist()), + ignore_index=case.ignore_index, + ), + sharding=ShardingSpec( + tp_rank=topology.tp_rank, + tp_world_size=case.tp_world_size, + vocab_shard_bounds=vocab_shard_bounds(case), + real_vocab_size=case.real_vocab_size, + padded_vocab_size=case.padded_vocab_size, + cp_rank=topology.cp_rank, + cp_world_size=case.cp_world_size, + ), + reduction=ReductionSpec(), + ) + + +def _fp32_oracle( + logits: torch.Tensor, + target_ids: torch.Tensor, + active_mask: torch.Tensor, + real_vocab_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + real_logits = logits[:, :real_vocab_size].float() + lse = torch.logsumexp(real_logits, dim=-1) + safe_targets = target_ids.masked_fill(~active_mask, 0) + selected = real_logits.gather(1, safe_targets.unsqueeze(1)).squeeze(1) + logp = torch.where(active_mask, selected - lse, torch.zeros_like(lse)) + return logp, lse + + +def _resolve_tolerance(dtype: str) -> tuple[float, float]: + entry = load_tolerance_contract()["accuracy"]["default"]["logprob"] + tolerance = entry[_TOLERANCE_DTYPES[dtype]] + return float(tolerance["atol"]), float(tolerance["rtol"]) + + +def _drift_detail( + candidate: torch.Tensor, + reference: torch.Tensor, + *, + target_ids: torch.Tensor, + global_positions: torch.Tensor, + sharding: ShardingSpec, + atol: float, + rtol: float, + mask: torch.Tensor | None = None, +) -> DriftDetail: + stats = summarize_logprob_drift(candidate, reference, mask=mask) + diff = (candidate.float() - reference.float()).abs() + selected = torch.ones_like(diff, dtype=torch.bool) if mask is None else mask.to(diff.device) + if not bool(selected.any().item()): + return DriftDetail(stats, 0.0, None, None, None, None, None, atol, rtol, True) + + selected_diff = diff[selected] + selected_ref = reference.float()[selected] + relative = selected_diff.double() / selected_ref.double().abs().clamp_min(_RELATIVE_ERROR_FLOOR) + selected_indices = torch.arange(diff.numel(), device=diff.device)[selected] + worst_selected = int(selected_diff.argmax().item()) + worst_local = int(selected_indices[worst_selected].item()) + target_id = int(target_ids[worst_local].item()) + close = selected_diff <= atol + rtol * selected_ref.abs() + return DriftDetail( + stats=stats, + max_rel=float(relative.max().item()), + worst_global_token=int(global_positions[worst_local].item()), + worst_target_id=target_id, + worst_owner_rank=sharding.owner_rank(target_id) if target_id >= 0 else None, + candidate_value=float(candidate[worst_local].float().item()), + reference_value=float(reference[worst_local].float().item()), + atol=atol, + rtol=rtol, + passed=bool(close.all().item()), + ) + + +def _tp_outputs_replicated( + logp: torch.Tensor, + lse: torch.Tensor, + *, + tp_group: Any, + tp_world_size: int, +) -> bool: + if tp_world_size == 1: + return True + import torch.distributed as dist + + gathered_logp = [torch.empty_like(logp) for _ in range(tp_world_size)] + gathered_lse = [torch.empty_like(lse) for _ in range(tp_world_size)] + dist.all_gather(gathered_logp, logp.contiguous(), group=tp_group) + dist.all_gather(gathered_lse, lse.contiguous(), group=tp_group) + return all(torch.equal(logp, value) for value in gathered_logp) and all( + torch.equal(lse, value) for value in gathered_lse + ) + + +def _execute_rank( + case: DistributedLogprobCase, + topology: RankTopology, + *, + device: torch.device, + tp_group: Any, +) -> _RankPayload: + full_logits, full_targets, full_active = _canonical_inputs(case) + token_start, token_end = token_shard_bounds(case.num_tokens, case.cp_world_size)[ + topology.cp_rank + ] + vocab_start, vocab_end = vocab_shard_bounds(case)[topology.tp_rank] + token_slice = slice(token_start, token_end) + local_active = full_active[token_slice].to(device=device) + local_targets = full_targets[token_slice].to(device=device) + local_fp32 = full_logits[token_slice].to(device=device) + local_logits = local_fp32[:, vocab_start:vocab_end].to(_DTYPES[case.dtype]).contiguous() + positions = torch.arange(token_start, token_end, device=device, dtype=torch.long) + + contract = _make_contract(case, topology, local_active.cpu()) + dispatch = KernelRegistry().get_logprob_op( + contract, + requested_backend=case.requested_backend, + ) + fallback = bool(dispatch.provenance["fallback"]) + if fallback: + raise RuntimeError("distributed logprob dispatch materialized through a fallback") + requested_policy = case.requested_backend.lower() + if requested_policy not in {"reference", "production"} and ( + case.requested_backend != dispatch.capability.backend_id + ): + raise RuntimeError( + f"requested backend {case.requested_backend!r} materialized as " + f"{dispatch.capability.backend_id!r}" + ) + + candidate_logp, candidate_lse = dispatch.op( + local_logits, + local_targets, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=case.num_vocab_tiles, + validate=True, + ) + reference_logp, reference_lse = _fp32_oracle( + local_fp32, + local_targets, + local_active, + case.real_vocab_size, + ) + replicated = _tp_outputs_replicated( + candidate_logp, + candidate_lse, + tp_group=tp_group, + tp_world_size=case.tp_world_size, + ) + atol, rtol = _resolve_tolerance(case.dtype) + lse_drift = _drift_detail( + candidate_lse, + reference_lse, + target_ids=local_targets, + global_positions=positions, + sharding=contract.sharding, + atol=atol, + rtol=rtol, + ) + dlogp_drift = _drift_detail( + candidate_logp, + reference_logp, + target_ids=local_targets, + global_positions=positions, + sharding=contract.sharding, + atol=atol, + rtol=rtol, + mask=local_active, + ) + rank_report = RankLogprobReport( + global_rank=topology.global_rank, + tp_rank=topology.tp_rank, + tp_world_size=topology.tp_world_size, + cp_rank=topology.cp_rank, + cp_world_size=topology.cp_world_size, + sp_world_size=1, + dp_world_size=1, + token_start=token_start, + token_end=token_end, + vocab_start=vocab_start, + vocab_end=vocab_end, + device=str(device), + requested_backend=case.requested_backend, + actual_backend=dispatch.capability.backend_id, + fallback=fallback, + contract_fingerprint=contract.cross_rank_fingerprint(), + contract=contract.to_dict(), + capability=dispatch.capability.to_dict(), + tp_outputs_bitwise_replicated=replicated, + lse=lse_drift, + dlogp=dlogp_drift, + passed=replicated and lse_drift.passed and dlogp_drift.passed, + ) + return _RankPayload( + report=rank_report, + candidate_logp=candidate_logp.detach().cpu(), + candidate_lse=candidate_lse.detach().cpu(), + reference_logp=reference_logp.detach().cpu(), + reference_lse=reference_lse.detach().cpu(), + active_mask=local_active.cpu(), + target_ids=local_targets.cpu(), + global_positions=positions.cpu(), + ) + + +def _aggregate_payloads( + case: DistributedLogprobCase, + payloads: Sequence[_RankPayload], +) -> dict[str, DriftDetail]: + representatives = sorted( + (payload for payload in payloads if payload.report.tp_rank == 0), + key=lambda payload: payload.report.cp_rank, + ) + if len(representatives) != case.cp_world_size: + raise RuntimeError("missing one or more CP representatives in rank reports") + candidate_logp = torch.cat([payload.candidate_logp for payload in representatives]) + candidate_lse = torch.cat([payload.candidate_lse for payload in representatives]) + reference_logp = torch.cat([payload.reference_logp for payload in representatives]) + reference_lse = torch.cat([payload.reference_lse for payload in representatives]) + active_mask = torch.cat([payload.active_mask for payload in representatives]) + target_ids = torch.cat([payload.target_ids for payload in representatives]) + positions = torch.cat([payload.global_positions for payload in representatives]) + sharding = _make_contract( + case, + rank_topology(case, 0), + active_mask, + ).sharding + atol, rtol = _resolve_tolerance(case.dtype) + return { + "lse": _drift_detail( + candidate_lse, + reference_lse, + target_ids=target_ids, + global_positions=positions, + sharding=sharding, + atol=atol, + rtol=rtol, + ), + "dlogp": _drift_detail( + candidate_logp, + reference_logp, + target_ids=target_ids, + global_positions=positions, + sharding=sharding, + atol=atol, + rtol=rtol, + mask=active_mask, + ), + } + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + """Hash the exact CPU tensor bytes for cross-topology comparisons.""" + + data = tensor.detach().cpu().contiguous().numpy().tobytes() + return hashlib.sha256(data).hexdigest() + + +def _aggregate_bitwise_fingerprints( + case: DistributedLogprobCase, + payloads: Sequence[_RankPayload], +) -> dict[str, Any]: + representatives = sorted( + (payload for payload in payloads if payload.report.tp_rank == 0), + key=lambda payload: payload.report.cp_rank, + ) + if len(representatives) != case.cp_world_size: + raise RuntimeError("missing one or more CP representatives in rank reports") + candidate_logp = torch.cat([payload.candidate_logp for payload in representatives]) + candidate_lse = torch.cat([payload.candidate_lse for payload in representatives]) + return { + "candidate_logp_sha256": _tensor_sha256(candidate_logp), + "candidate_lse_sha256": _tensor_sha256(candidate_lse), + "dtype": str(candidate_logp.dtype).replace("torch.", ""), + "shape": list(candidate_logp.shape), + } + + +def _create_tp_group(case: DistributedLogprobCase, topology: RankTopology) -> Any: + if case.world_size == 1: + return None + import torch.distributed as dist + + selected = None + for cp_rank in range(case.cp_world_size): + start = cp_rank * case.tp_world_size + ranks = list(range(start, start + case.tp_world_size)) + group = dist.new_group(ranks=ranks) + if topology.global_rank in ranks: + selected = group + return selected + + +def run_distributed_logprob_case( + case: DistributedLogprobCase, + *, + device_name: str = "cuda", + dist_backend: str | None = None, + output: str | pathlib.Path, +) -> DistributedLogprobReport | None: + """Run one materialized topology; only global rank zero returns the report.""" + + import torch.distributed as dist + + backend = dist_backend or ("nccl" if device_name == "cuda" else "gloo") + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if world_size != case.world_size: + raise RuntimeError( + f"WORLD_SIZE={world_size} does not match TP*CP={case.world_size}; " + "launch exactly the topology declared by the case" + ) + owns_process_group = world_size > 1 and not dist.is_initialized() + try: + if device_name == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable") + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + elif device_name == "cpu": + device = torch.device("cpu") + else: + raise ValueError("device must be cuda or cpu") + + if owns_process_group: + dist.init_process_group(backend=backend, timeout=_PROCESS_GROUP_TIMEOUT) + if dist.is_initialized(): + if dist.get_world_size() != world_size or dist.get_rank() != rank: + raise RuntimeError("initialized process group does not match RANK/WORLD_SIZE") + + topology = rank_topology(case, rank) + tp_group = _create_tp_group(case, topology) + payload = _execute_rank(case, topology, device=device, tp_group=tp_group) + if world_size == 1: + payloads = [payload] + else: + gathered: list[Any] = [None] * world_size + dist.all_gather_object(gathered, payload) + payloads = gathered + + report = None + if rank == 0: + aggregate = _aggregate_payloads(case, payloads) + bitwise_fingerprints = _aggregate_bitwise_fingerprints(case, payloads) + rank_reports = tuple( + payload.report + for payload in sorted(payloads, key=lambda item: item.report.global_rank) + ) + actual_backends = sorted({rank_report.actual_backend for rank_report in rank_reports}) + reduction_specs = { + json.dumps(rank_report.contract["reduction"], sort_keys=True) + for rank_report in rank_reports + } + materialization_consistent = len(actual_backends) == 1 and len(reduction_specs) == 1 + launch_command = format_launch_command( + case, + output=output, + device=device_name, + dist_backend=backend, + ) + report = DistributedLogprobReport( + schema_version=1, + case_id=case.case_id, + case=asdict(case), + launch_command=launch_command, + environment={ + "python": sys.version.split()[0], + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "dist_backend": backend, + "world_size": world_size, + "sp_world_size": 1, + "dp_world_size": 1, + "materialization": { + "actual_backends": actual_backends, + "consistent": materialization_consistent, + }, + "communication": { + "logprob_merge_axis": "tp_vocab", + "cp_is_merge_axis": False, + "report_collection": ("all_gather_object" if world_size > 1 else "none"), + }, + }, + ranks=rank_reports, + aggregate=aggregate, + bitwise_fingerprints=bitwise_fingerprints, + passed=( + materialization_consistent + and all(rank_report.passed for rank_report in rank_reports) + and all(detail.passed for detail in aggregate.values()) + ), + ) + output_path = pathlib.Path(output) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + _strict_report_json(report.to_dict()) + "\n", + encoding="utf-8", + ) + if world_size > 1: + dist.barrier() + return report + finally: + if owns_process_group and dist.is_initialized(): + dist.destroy_process_group() + + +def _case_from_args(args: argparse.Namespace) -> DistributedLogprobCase: + return DistributedLogprobCase( + tp_world_size=args.tp, + cp_world_size=args.cp, + dtype=args.dtype, + requested_backend=args.backend, + real_vocab_size=args.real_vocab, + padded_vocab_size=args.padded_vocab, + num_vocab_tiles=args.num_vocab_tiles, + batch_size=args.batch, + sequence_length=args.seq, + prompt_tokens=args.prompt_tokens, + seed=args.seed, + ) + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the WS2 distributed logprob drift report.") + parser.add_argument("--plan", action="store_true", help="Print the six scoped launch commands.") + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--cp", type=int, default=1) + parser.add_argument("--dtype", choices=tuple(_DTYPES), default="bf16") + parser.add_argument("--backend", default=BACKEND_ID) + parser.add_argument("--real-vocab", type=int, default=151936) + parser.add_argument("--padded-vocab", type=int, default=151936) + parser.add_argument("--num-vocab-tiles", type=int, default=DEFAULT_NUM_VOCAB_TILES) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq", type=int, default=16) + parser.add_argument("--prompt-tokens", type=int, default=8) + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda") + parser.add_argument("--dist-backend", choices=("nccl", "gloo"), default=None) + parser.add_argument("--output", default="artifacts/ws2-logprob/report.json") + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> None: + route_rl_kernel_logs_to_stderr() + args = _parse_args(argv) + if args.plan: + cases = plan_distributed_logprob_cases( + dtype=args.dtype, + requested_backend=args.backend, + real_vocab_size=args.real_vocab, + padded_vocab_size=args.padded_vocab, + num_vocab_tiles=args.num_vocab_tiles, + batch_size=args.batch, + sequence_length=args.seq, + prompt_tokens=args.prompt_tokens, + seed=args.seed, + ) + commands = [ + format_launch_command( + case, + output=pathlib.Path(args.output).parent + / f"tp{case.tp_world_size}-cp{case.cp_world_size}.json", + device=args.device, + dist_backend=args.dist_backend, + ) + for case in cases + ] + print(json.dumps({"commands": commands}, indent=2)) + return + + case = _case_from_args(args) + report = run_distributed_logprob_case( + case, + device_name=args.device, + dist_backend=args.dist_backend, + output=args.output, + ) + if report is not None: + print(_strict_report_json(report.to_dict())) + if not report.passed: + raise SystemExit(1) + + +__all__ = [ + "DistributedLogprobCase", + "DistributedLogprobReport", + "DriftDetail", + "RankLogprobReport", + "RankTopology", + "format_launch_command", + "plan_distributed_logprob_cases", + "rank_topology", + "run_distributed_logprob_case", + "token_shard_bounds", + "vocab_shard_bounds", +] + + +if __name__ == "__main__": + main() diff --git a/rl_engine/testing/logprob_comparison.py b/rl_engine/testing/logprob_comparison.py new file mode 100644 index 00000000..1e9d4e63 --- /dev/null +++ b/rl_engine/testing/logprob_comparison.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Single-GPU selected-logprob comparison.""" + +from __future__ import annotations + +import argparse +import json +import logging +import pathlib +import sys +from collections.abc import Callable, Sequence +from dataclasses import asdict, dataclass, field +from typing import Any + +import torch + +if __package__ in (None, ""): + repo_root = pathlib.Path(__file__).resolve().parents[2] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + from logprob_drift import LogprobDriftStats, summarize_logprob_drift +else: + from .logprob_drift import LogprobDriftStats, summarize_logprob_drift + + +class LogprobBackendUnavailable(RuntimeError): + pass + + +@dataclass(frozen=True) +class LogprobComparisonInputs: + logits: torch.Tensor + target_ids: torch.Tensor + active_token_mask: torch.Tensor | None = None + ignore_index: int = -100 + + +@dataclass(frozen=True) +class LogprobCandidate: + name: str + requested_backend: str + actual_backend: str + fn: Callable[[torch.Tensor, torch.Tensor, int], tuple[torch.Tensor, torch.Tensor]] + provenance: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class _LogprobPathDrift: + candidate_name: str + lse: LogprobDriftStats + dlogp: LogprobDriftStats + bitwise_logp: bool + provenance: dict[str, Any] + + +@dataclass(frozen=True) +class LogprobComparisonReport: + reference_name: str + drifts: tuple[_LogprobPathDrift, ...] + input_provenance: dict[str, Any] + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def make_logprob_candidate(backend: str) -> LogprobCandidate: + normalized = backend.strip().lower().replace("_", "-") + op: Any + if normalized in {"pytorch", "native"}: + from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( + NativeBatchInvariantLogpOp, + ) + + op = NativeBatchInvariantLogpOp() + actual = "pytorch" + elif normalized == "triton": + try: + from rl_engine.kernels.ops.triton.loss.batch_invariant_logp import ( + TritonBatchInvariantLogpOp, + ) + + op = TritonBatchInvariantLogpOp() + except Exception as exc: + raise LogprobBackendUnavailable(f"triton backend is unavailable: {exc}") from exc + actual = "triton" + elif normalized in {"cuda-sm90", "sm90"}: + try: + from rl_engine.kernels.ops.cuda.loss.batch_invariant_logp import ( + BatchInvariantLogpSM90Op, + ) + + op = BatchInvariantLogpSM90Op() + except Exception as exc: + raise LogprobBackendUnavailable(f"cuda-sm90 backend is unavailable: {exc}") from exc + actual = "cuda-sm90" + else: + raise ValueError( + f"unsupported logprob comparison backend {backend!r}; " + "expected pytorch, triton, or cuda-sm90" + ) + + diagnostic = getattr(op, "forward_with_lse", None) + if not callable(diagnostic): + raise LogprobBackendUnavailable( + f"backend {normalized!r} does not expose the required direct LSE diagnostic" + ) + + def run( + logits: torch.Tensor, target_ids: torch.Tensor, ignore_index: int + ) -> tuple[torch.Tensor, torch.Tensor]: + try: + return diagnostic(logits, target_ids, ignore_index=ignore_index, validate=True) + except (RuntimeError, NotImplementedError, OSError) as exc: + raise LogprobBackendUnavailable( + f"exact backend {normalized!r} cannot execute this input: {exc}" + ) from exc + + return LogprobCandidate( + name=f"{actual}-batch-invariant-logp", + requested_backend=actual, + actual_backend=actual, + fn=run, + provenance={ + "requested_alias": normalized, + "implementation": f"{type(op).__module__}.{type(op).__qualname__}", + }, + ) + + +def compare_single_gpu_logprob( + inputs: LogprobComparisonInputs, + *, + candidates: Sequence[str | LogprobCandidate] = ("pytorch",), +) -> LogprobComparisonReport: + active_mask, effective_targets = _validate_inputs(inputs) + reference_logp, reference_lse = _run_ws1_reference( + inputs.logits, effective_targets, inputs.ignore_index + ) + + drifts = [] + for candidate in candidates: + if isinstance(candidate, str): + candidate = make_logprob_candidate(candidate) + logp, lse = _run_candidate( + candidate, + inputs.logits, + effective_targets, + inputs.ignore_index, + ) + drifts.append( + _LogprobPathDrift( + candidate_name=candidate.name, + lse=summarize_logprob_drift(lse, reference_lse), + dlogp=summarize_logprob_drift(logp, reference_logp, mask=active_mask), + bitwise_logp=torch.equal(logp, reference_logp), + provenance=_candidate_provenance(candidate), + ) + ) + + return LogprobComparisonReport( + reference_name="pytorch-batch-invariant-logp", + drifts=tuple(drifts), + input_provenance={ + "device": str(inputs.logits.device), + "input_dtype": str(inputs.logits.dtype), + "output_dtype": str(reference_logp.dtype), + "shape": list(inputs.logits.shape), + "ignore_index": inputs.ignore_index, + "active_token_count": int(active_mask.sum().item()), + "tp_world": 1, + "communication": "none", + }, + ) + + +def _run_ws1_reference( + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int, +) -> tuple[torch.Tensor, torch.Tensor]: + from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp + + op = NativeBatchInvariantLogpOp() + logp = op(logits, target_ids, ignore_index=ignore_index, validate=True) + _, lse = op.forward_with_lse(logits, target_ids, ignore_index=ignore_index, validate=True) + return logp.detach(), lse.detach() + + +def _run_candidate( + candidate: LogprobCandidate, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if candidate.requested_backend != candidate.actual_backend: + raise LogprobBackendUnavailable( + f"requested backend {candidate.requested_backend!r} materialized as " + f"{candidate.actual_backend!r}; silent fallback is forbidden" + ) + logp, lse = candidate.fn(logits, target_ids, ignore_index) + expected_shape = logits.shape[:-1] + for name, value in (("logp", logp), ("lse", lse)): + if not isinstance(value, torch.Tensor): + raise TypeError(f"candidate {candidate.name!r} {name} must be a tensor") + if value.shape != expected_shape: + raise ValueError( + f"candidate {candidate.name!r} {name} shape {tuple(value.shape)} " + f"does not match {tuple(expected_shape)}" + ) + if value.dtype != torch.float32: + raise ValueError(f"candidate {candidate.name!r} {name} must be FP32") + return logp.detach(), lse.detach() + + +def _candidate_provenance(candidate: LogprobCandidate) -> dict[str, Any]: + return { + **candidate.provenance, + "requested_backend": candidate.requested_backend, + "actual_backend": candidate.actual_backend, + "tp_world": 1, + "communication": "none", + "lse_source": "direct", + } + + +def _validate_inputs( + inputs: LogprobComparisonInputs, +) -> tuple[torch.Tensor, torch.Tensor]: + if inputs.logits.dim() < 2: + raise ValueError("logits must be at least 2-D [*lead, vocab]") + if inputs.logits.shape[:-1] != inputs.target_ids.shape: + raise ValueError("target_ids shape must match logits leading shape") + if not inputs.logits.is_floating_point(): + raise ValueError("logits must be floating point") + + if inputs.active_token_mask is None: + active = inputs.target_ids != inputs.ignore_index + else: + if inputs.active_token_mask.shape != inputs.target_ids.shape: + raise ValueError("active_token_mask shape must match target_ids") + if inputs.active_token_mask.dtype != torch.bool: + raise ValueError("active_token_mask must be bool") + active = inputs.active_token_mask.to(device=inputs.target_ids.device) + if bool(((inputs.target_ids == inputs.ignore_index) & active).any().item()): + raise ValueError("active target_ids cannot equal ignore_index") + + effective = inputs.target_ids.to(device=inputs.logits.device, dtype=torch.long).clone() + active = active.to(device=inputs.logits.device, dtype=torch.bool) + effective.masked_fill_(~active, inputs.ignore_index) + valid = effective[active] + vocab_size = inputs.logits.size(-1) + if valid.numel() and ((valid < 0).any() or (valid >= vocab_size).any()): + raise ValueError(f"active target_ids must be in [0, {vocab_size})") + return active, effective + + +def _dtype(name: str) -> torch.dtype: + return { + "fp32": torch.float32, + "bf16": torch.bfloat16, + "fp16": torch.float16, + }[name] + + +def _device(name: str) -> torch.device: + if name == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device(name) + + +def route_rl_kernel_logs_to_stderr() -> None: + from rl_engine.utils.logger import logger + + for handler in logger.handlers: + if isinstance(handler, logging.StreamHandler): + handler.setStream(sys.stderr) + + +def _route_rl_kernel_logs_to_stderr() -> None: + route_rl_kernel_logs_to_stderr() + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run the WS2 TP=1 selected-logprob/LSE comparison harness." + ) + parser.add_argument( + "--candidate", + action="append", + choices=("pytorch", "triton", "cuda-sm90"), + help="Exact backend to compare. Repeat for multiple backends; defaults to pytorch.", + ) + parser.add_argument("--device", default="auto") + parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq", type=int, default=16) + parser.add_argument("--vocab", type=int, default=257) + parser.add_argument("--prompt-tokens", type=int, default=8) + parser.add_argument("--seed", type=int, default=123) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> None: + route_rl_kernel_logs_to_stderr() + args = _parse_args(argv) + device = _device(args.device) + if args.batch < 1 or args.seq < 1 or args.vocab < 1: + raise ValueError("batch, seq, and vocab must be positive") + if not 0 <= args.prompt_tokens <= args.seq: + raise ValueError("prompt-tokens must be in [0, seq]") + + generator = torch.Generator(device=device).manual_seed(args.seed) + logits = torch.randn( + args.batch, + args.seq, + args.vocab, + generator=generator, + device=device, + dtype=_dtype(args.dtype), + ) + target_ids = torch.randint( + 0, + args.vocab, + (args.batch, args.seq), + generator=generator, + device=device, + ) + active_mask = torch.ones((args.batch, args.seq), device=device, dtype=torch.bool) + active_mask[:, : args.prompt_tokens] = False + report = compare_single_gpu_logprob( + LogprobComparisonInputs( + logits=logits, + target_ids=target_ids, + active_token_mask=active_mask, + ), + candidates=tuple(args.candidate or ("pytorch",)), + ) + print(json.dumps(report.to_dict(), indent=2, sort_keys=True)) + + +__all__ = [ + "LogprobBackendUnavailable", + "LogprobCandidate", + "LogprobComparisonInputs", + "LogprobComparisonReport", + "compare_single_gpu_logprob", + "make_logprob_candidate", + "route_rl_kernel_logs_to_stderr", +] + + +if __name__ == "__main__": + main() diff --git a/rl_engine/testing/logprob_drift.py b/rl_engine/testing/logprob_drift.py new file mode 100644 index 00000000..4996dfbc --- /dev/null +++ b/rl_engine/testing/logprob_drift.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared selected-logprob drift summaries.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class LogprobDriftStats: + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + +def summarize_logprob_drift( + candidate: torch.Tensor, + reference: torch.Tensor, + *, + mask: torch.Tensor | None = None, +) -> LogprobDriftStats: + """Summarize absolute drift, optionally over active rows only.""" + + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match reference shape " + f"{tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs() + if mask is None: + values = diff.reshape(-1) + else: + if mask.shape != diff.shape: + raise ValueError("mask shape must match candidate and reference") + if mask.dtype != torch.bool: + raise ValueError("mask must be bool") + values = diff[mask.to(device=diff.device)] + + count = int(values.numel()) + if count == 0: + return LogprobDriftStats(0.0, 0.0, 0.0, 0.0, 0) + return LogprobDriftStats( + max_abs=float(values.max().item()), + mean_abs=float(values.mean().item()), + p95_abs=float(torch.quantile(values, 0.95).item()), + p99_abs=float(torch.quantile(values, 0.99).item()), + active_count=count, + ) + + +__all__ = ["LogprobDriftStats", "summarize_logprob_drift"] diff --git a/tests/test_distributed_logprob_comparison.py b/tests/test_distributed_logprob_comparison.py new file mode 100644 index 00000000..8bb9d4ad --- /dev/null +++ b/tests/test_distributed_logprob_comparison.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +import math +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist + +import rl_engine.testing.distributed_logprob_comparison as distributed_comparison +from rl_engine.kernels.logprob_contract import ShardingSpec +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import BACKEND_ID +from rl_engine.testing.distributed_logprob_comparison import ( + DistributedLogprobCase, + _drift_detail, + _strict_report_json, + format_launch_command, + plan_distributed_logprob_cases, + rank_topology, + run_distributed_logprob_case, + token_shard_bounds, + vocab_shard_bounds, +) +from rl_engine.testing.logprob_drift import summarize_logprob_drift + + +def _small_case(*, tp: int = 1, cp: int = 1) -> DistributedLogprobCase: + return DistributedLogprobCase( + tp_world_size=tp, + cp_world_size=cp, + real_vocab_size=13, + padded_vocab_size=16, + num_vocab_tiles=8, + batch_size=1, + sequence_length=4, + prompt_tokens=1, + seed=7, + ) + + +def test_planner_builds_the_scoped_topology_product(): + cases = plan_distributed_logprob_cases( + real_vocab_size=13, + padded_vocab_size=16, + num_vocab_tiles=8, + ) + + assert [(case.tp_world_size, case.cp_world_size) for case in cases] == [ + (1, 1), + (1, 2), + (2, 1), + (2, 2), + (4, 1), + (4, 2), + ] + assert [case.world_size for case in cases] == [1, 2, 2, 4, 4, 8] + + +def test_rank_mapping_keeps_cp_out_of_the_tp_merge_axis(): + case = _small_case(tp=2, cp=2) + + assert rank_topology(case, 0).tp_group_ranks == (0, 1) + assert rank_topology(case, 1).tp_group_ranks == (0, 1) + assert rank_topology(case, 2).tp_group_ranks == (2, 3) + assert rank_topology(case, 3).tp_group_ranks == (2, 3) + assert [ + (rank_topology(case, rank).cp_rank, rank_topology(case, rank).tp_rank) for rank in range(4) + ] == [ + (0, 0), + (0, 1), + (1, 0), + (1, 1), + ] + + +def test_token_and_vocab_bounds_cover_each_axis_once(): + case = _small_case(tp=4, cp=2) + + assert token_shard_bounds(case.num_tokens, case.cp_world_size) == ((0, 2), (2, 4)) + assert vocab_shard_bounds(case) == ((0, 4), (4, 8), (8, 12), (12, 16)) + + +def test_case_rejects_implicit_backend_and_non_tileable_vocab(): + with pytest.raises(ValueError, match="explicit non-auto backend"): + DistributedLogprobCase(tp_world_size=2, cp_world_size=1, requested_backend="auto") + with pytest.raises(ValueError, match="must divide"): + DistributedLogprobCase( + tp_world_size=2, + cp_world_size=1, + padded_vocab_size=15, + real_vocab_size=13, + num_vocab_tiles=8, + ) + + +def test_launch_command_records_the_materialized_case(tmp_path): + case = _small_case(tp=2, cp=2) + command = format_launch_command(case, output=tmp_path / "report.json") + + assert "--nproc-per-node=4" in command + assert "--tp 2 --cp 2" in command + assert f"--backend {BACKEND_ID}" in command + assert "--real-vocab 13 --padded-vocab 16" in command + + +def test_shared_pr2_drift_summary_preserves_active_mask_semantics(): + candidate = torch.tensor([100.0, 1.0, 3.0]) + reference = torch.tensor([0.0, 2.0, 1.0]) + mask = torch.tensor([False, True, True]) + + stats = summarize_logprob_drift(candidate, reference, mask=mask) + + assert stats.active_count == 2 + assert stats.max_abs == 2.0 + assert stats.mean_abs == 1.5 + + +def test_relative_drift_near_zero_stays_finite(): + sharding = ShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, 16),), + real_vocab_size=13, + padded_vocab_size=16, + ) + detail = _drift_detail( + torch.tensor([1.0]), + torch.tensor([0.0]), + target_ids=torch.tensor([1]), + global_positions=torch.tensor([3]), + sharding=sharding, + atol=0.0, + rtol=0.0, + ) + + assert math.isfinite(detail.max_rel) + assert detail.max_rel == pytest.approx(1.0e12) + assert json.loads(_strict_report_json({"max_rel": detail.max_rel}))["max_rel"] == pytest.approx( + 1.0e12 + ) + with pytest.raises(ValueError, match="Out of range float values"): + _strict_report_json({"max_rel": float("nan")}) + + +def test_tp1_cpu_case_writes_116_compatible_artifact(tmp_path, monkeypatch): + monkeypatch.delenv("RANK", raising=False) + monkeypatch.delenv("LOCAL_RANK", raising=False) + monkeypatch.delenv("WORLD_SIZE", raising=False) + output = tmp_path / "tp1-cp1.json" + + report = run_distributed_logprob_case( + _small_case(), + device_name="cpu", + dist_backend="gloo", + output=output, + ) + + assert report is not None and report.passed + assert report.aggregate["lse"].stats.active_count == 4 + assert report.aggregate["dlogp"].stats.active_count == 3 + assert report.ranks[0].actual_backend == BACKEND_ID + assert report.ranks[0].fallback is False + assert report.ranks[0].tp_outputs_bitwise_replicated + payload = json.loads(output.read_text(encoding="utf-8")) + assert payload["schema_version"] == 1 + assert payload["ranks"][0]["contract"]["reduction"]["cp_is_merge_axis"] is False + assert payload["ranks"][0]["sp_world_size"] == 1 + assert payload["ranks"][0]["dp_world_size"] == 1 + assert payload["environment"]["materialization"]["consistent"] is True + assert payload["aggregate"]["dlogp"]["worst_target_id"] is not None + fingerprints = payload["bitwise_fingerprints"] + assert len(fingerprints["candidate_logp_sha256"]) == 64 + assert len(fingerprints["candidate_lse_sha256"]) == 64 + assert fingerprints["dtype"] == "float32" + assert fingerprints["shape"] == [4] + assert payload["launch_command"].startswith("torchrun --standalone") + + +def test_world_size_mismatch_fails_before_process_group_init(tmp_path, monkeypatch): + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + + with pytest.raises(RuntimeError, match=r"does not match TP\*CP"): + run_distributed_logprob_case( + _small_case(), + device_name="cpu", + dist_backend="gloo", + output=tmp_path / "unused.json", + ) + + +def test_setup_failure_destroys_owned_process_group(tmp_path, monkeypatch): + state = {"initialized": False, "destroyed": False} + + def init_process_group(*, backend, timeout): + state["initialized"] = True + + def destroy_process_group(): + state["destroyed"] = True + state["initialized"] = False + + def fail_group_setup(case, topology): + raise RuntimeError("group setup failed") + + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(dist, "is_initialized", lambda: state["initialized"]) + monkeypatch.setattr(dist, "init_process_group", init_process_group) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + monkeypatch.setattr(dist, "get_rank", lambda: 0) + monkeypatch.setattr(dist, "destroy_process_group", destroy_process_group) + monkeypatch.setattr(distributed_comparison, "_create_tp_group", fail_group_setup) + + with pytest.raises(RuntimeError, match="group setup failed"): + run_distributed_logprob_case( + _small_case(tp=2), + device_name="cpu", + dist_backend="gloo", + output=tmp_path / "unused.json", + ) + + assert state == {"initialized": False, "destroyed": True} + + +def test_partial_initialization_failure_destroys_owned_process_group(tmp_path, monkeypatch): + state = {"initialized": False, "destroyed": False} + + def fail_initialization(*, backend, timeout): + state["initialized"] = True + raise RuntimeError("initialization failed") + + def destroy_process_group(): + state["destroyed"] = True + state["initialized"] = False + + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(dist, "is_initialized", lambda: state["initialized"]) + monkeypatch.setattr(dist, "init_process_group", fail_initialization) + monkeypatch.setattr(dist, "destroy_process_group", destroy_process_group) + + with pytest.raises(RuntimeError, match="initialization failed"): + run_distributed_logprob_case( + _small_case(tp=2), + device_name="cpu", + dist_backend="gloo", + output=tmp_path / "unused.json", + ) + + assert state == {"initialized": False, "destroyed": True} + + +@pytest.mark.skipif(not torch.distributed.is_available(), reason="torch.distributed required") +def test_tp2_cp2_gloo_cli_emits_per_rank_report(tmp_path): + script = ( + Path(__file__).resolve().parents[1] + / "rl_engine" + / "testing" + / "distributed_logprob_comparison.py" + ) + output = tmp_path / "tp2-cp2.json" + environment = os.environ.copy() + environment.setdefault("OMP_NUM_THREADS", "1") + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc-per-node=4", + str(script), + "--tp", + "2", + "--cp", + "2", + "--device", + "cpu", + "--dist-backend", + "gloo", + "--real-vocab", + "13", + "--padded-vocab", + "16", + "--num-vocab-tiles", + "8", + "--batch", + "1", + "--seq", + "4", + "--prompt-tokens", + "1", + "--output", + str(output), + ], + check=True, + capture_output=True, + text=True, + timeout=120, + env=environment, + ) + + payload = json.loads(output.read_text(encoding="utf-8")) + assert payload["passed"] + assert len(payload["ranks"]) == 4 + assert all(rank["tp_outputs_bitwise_replicated"] for rank in payload["ranks"]) + assert {rank["actual_backend"] for rank in payload["ranks"]} == {BACKEND_ID} + assert {rank["cp_rank"] for rank in payload["ranks"]} == {0, 1} + assert payload["aggregate"]["dlogp"]["stats"]["active_count"] == 3 + assert json.loads(result.stdout)["case"]["tp_world_size"] == 2 diff --git a/tests/test_logprob_comparison.py b/tests/test_logprob_comparison.py new file mode 100644 index 00000000..4fc62a13 --- /dev/null +++ b/tests/test_logprob_comparison.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import argparse +import io +import json +import logging +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from rl_engine.kernels.gtest import run_operator_suite +from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case +from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp +from rl_engine.testing.logprob_comparison import ( + LogprobBackendUnavailable, + LogprobCandidate, + LogprobComparisonInputs, + _device, + _route_rl_kernel_logs_to_stderr, + compare_single_gpu_logprob, + make_logprob_candidate, +) +from rl_engine.utils.logger import logger + + +def _inputs() -> LogprobComparisonInputs: + generator = torch.Generator().manual_seed(17) + logits = torch.randn(2, 4, 257, generator=generator, dtype=torch.float32) + target_ids = torch.tensor([[3, 5, 7, 11], [13, 17, 19, 23]]) + active = torch.tensor([[False, False, True, True], [False, True, True, True]]) + return LogprobComparisonInputs(logits, target_ids, active_token_mask=active) + + +def test_single_gpu_pytorch_path_is_bitwise_regression_guard(): + report = compare_single_gpu_logprob(_inputs(), candidates=("pytorch",)) + + assert report.reference_name == "pytorch-batch-invariant-logp" + assert len(report.drifts) == 1 + drift = report.drifts[0] + assert drift.bitwise_logp + assert drift.lse.max_abs == 0.0 + assert drift.dlogp.max_abs == 0.0 + assert drift.dlogp.active_count == 5 + assert drift.provenance["requested_backend"] == "pytorch" + assert drift.provenance["actual_backend"] == "pytorch" + assert drift.provenance["lse_source"] == "direct" + assert report.input_provenance["tp_world"] == 1 + assert report.input_provenance["communication"] == "none" + + +def test_report_serializes_lse_and_active_token_percentiles(): + inputs = _inputs() + reference = make_logprob_candidate("pytorch") + + def shifted(logits, target_ids, ignore_index): + logp, lse = reference.fn(logits, target_ids, ignore_index) + logp = logp.clone() + logp[0, 0] += 100.0 # inactive and therefore excluded from dlogp + logp[0, 2] += 1.0 + lse = lse + torch.arange(lse.numel(), dtype=lse.dtype).reshape_as(lse) * 0.1 + return logp, lse + + candidate = LogprobCandidate( + name="shifted", + requested_backend="shifted", + actual_backend="shifted", + fn=shifted, + ) + report = compare_single_gpu_logprob(inputs, candidates=(candidate,)) + payload = report.to_dict() + drift = payload["drifts"][0] + + assert drift["dlogp"]["active_count"] == 5 + assert drift["dlogp"]["max_abs"] == pytest.approx(1.0) + assert drift["dlogp"]["p95_abs"] == pytest.approx(0.8) + assert drift["dlogp"]["p99_abs"] == pytest.approx(0.96) + assert drift["lse"]["active_count"] == 8 + assert drift["lse"]["p99_abs"] == pytest.approx(0.693, abs=1e-5) + + +def test_canonical_provenance_cannot_be_overridden(): + native = make_logprob_candidate("pytorch") + candidate = LogprobCandidate( + name="custom", + requested_backend="pytorch", + actual_backend="pytorch", + fn=native.fn, + provenance={ + "actual_backend": "fallback", + "tp_world": 8, + "communication": "all-gather", + "lse_source": "reconstructed", + "implementation": "custom", + }, + ) + + provenance = compare_single_gpu_logprob(_inputs(), candidates=(candidate,)).drifts[0].provenance + + assert provenance["actual_backend"] == "pytorch" + assert provenance["tp_world"] == 1 + assert provenance["communication"] == "none" + assert provenance["lse_source"] == "direct" + assert provenance["implementation"] == "custom" + + +def test_all_inactive_tokens_produce_zero_dlogp_statistics(): + inputs = _inputs() + inputs = LogprobComparisonInputs( + inputs.logits, + inputs.target_ids, + active_token_mask=torch.zeros_like(inputs.target_ids, dtype=torch.bool), + ) + drift = compare_single_gpu_logprob(inputs).drifts[0] + + assert drift.dlogp.active_count == 0 + assert drift.dlogp.max_abs == 0.0 + assert drift.dlogp.p95_abs == 0.0 + assert drift.lse.active_count == inputs.target_ids.numel() + + +def test_explicit_backend_mismatch_fails_closed(): + native = make_logprob_candidate("pytorch") + disguised = LogprobCandidate( + name="fallback", + requested_backend="cuda-sm90", + actual_backend="pytorch", + fn=native.fn, + ) + + with pytest.raises(LogprobBackendUnavailable, match="silent fallback is forbidden"): + compare_single_gpu_logprob(_inputs(), candidates=(disguised,)) + + +def test_active_ignore_index_is_rejected(): + inputs = _inputs() + targets = inputs.target_ids.clone() + targets[0, 2] = -100 + + with pytest.raises(ValueError, match="active target_ids cannot equal ignore_index"): + compare_single_gpu_logprob( + LogprobComparisonInputs( + inputs.logits, + targets, + active_token_mask=inputs.active_token_mask, + ) + ) + + +def test_native_diagnostic_lse_satisfies_selected_logit_identity(): + inputs = _inputs() + candidate = make_logprob_candidate("pytorch") + effective = inputs.target_ids.masked_fill(~inputs.active_token_mask, -100) + logp, lse = candidate.fn(inputs.logits, effective, -100) + production_logp = NativeBatchInvariantLogpOp()( + inputs.logits, effective, ignore_index=-100, validate=True + ) + safe_targets = effective.masked_fill(~inputs.active_token_mask, 0) + selected = torch.gather(inputs.logits, -1, safe_targets.unsqueeze(-1)).squeeze(-1) + + assert torch.equal(logp, production_logp) + assert torch.equal(logp[inputs.active_token_mask], (selected - lse)[inputs.active_token_mask]) + + +def test_unsupported_backend_name_is_rejected(): + with pytest.raises(ValueError, match="unsupported logprob comparison backend"): + make_logprob_candidate("unknown") + + +def test_cli_auto_device_resolves_without_constructing_auto(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + assert _device("auto") == torch.device("cpu") + + +def test_cli_routes_rl_kernel_logs_to_stderr_for_machine_readable_stdout(monkeypatch): + stdout = io.StringIO() + stderr = io.StringIO() + original_streams = [ + (handler, handler.stream) + for handler in logger.handlers + if isinstance(handler, logging.StreamHandler) + ] + monkeypatch.setattr(sys, "stdout", stdout) + monkeypatch.setattr(sys, "stderr", stderr) + + try: + _route_rl_kernel_logs_to_stderr() + logger.info("test backend diagnostic") + print(json.dumps({"ok": True})) + finally: + for handler, stream in original_streams: + handler.setStream(stream) + + assert json.loads(stdout.getvalue()) == {"ok": True} + assert "test backend diagnostic" in stderr.getvalue() + + +def test_cli_runs_directly_from_testing_module(): + script = Path(__file__).resolve().parents[1] / "rl_engine" / "testing" / "logprob_comparison.py" + result = subprocess.run( + [ + sys.executable, + str(script), + "--candidate", + "pytorch", + "--device", + "cpu", + "--batch", + "1", + "--seq", + "2", + "--vocab", + "17", + "--prompt-tokens", + "1", + ], + check=True, + capture_output=True, + text=True, + ) + + payload = json.loads(result.stdout) + assert payload["drifts"][0]["provenance"]["actual_backend"] == "pytorch" + assert payload["input_provenance"]["communication"] == "none" + + +def test_operator_comparison_specs_register_batch_invariant_logp(): + args = argparse.Namespace( + op="batch_invariant_logp", + candidate="pytorch", + arch_key=None, + batch=2, + seq=4, + vocab=17, + seed=7, + input_mode="random", + constant_value=0.5, + token_value=3, + normalized_dim=128, + k_dim=16, + n_dim=32, + theta=1.0e6, + eps=1.0e-6, + ) + + case = make_operator_case(args, torch.float32, torch.device("cpu")) + candidate = make_candidate(args) + report = run_operator_suite("batch_invariant_logp", candidates=[candidate], cases=[case]) + + assert report.passed + assert report.candidates[0].cases[0].op_class == "logprob" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_triton_diagnostic_path_reports_direct_lse(): + try: + candidate = make_logprob_candidate("triton") + except LogprobBackendUnavailable as exc: + pytest.skip(str(exc)) + logits = torch.randn(4, 1024, device="cuda", dtype=torch.bfloat16) + targets = torch.tensor([0, 17, 511, 1023], device="cuda") + try: + report = compare_single_gpu_logprob( + LogprobComparisonInputs(logits, targets), candidates=(candidate,) + ) + except LogprobBackendUnavailable as exc: + if isinstance(exc.__cause__, PermissionError): + pytest.skip(str(exc)) + raise + + assert report.drifts[0].provenance["actual_backend"] == "triton" + assert report.drifts[0].provenance["lse_source"] == "direct" + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9, + reason="Hopper CUDA device required", +) +def test_sm90_diagnostic_path_reports_direct_lse_without_fallback(): + try: + candidate = make_logprob_candidate("cuda-sm90") + except LogprobBackendUnavailable as exc: + pytest.skip(str(exc)) + logits = torch.randn(4, 1024, device="cuda", dtype=torch.bfloat16) + targets = torch.tensor([0, 17, 511, 1023], device="cuda") + report = compare_single_gpu_logprob( + LogprobComparisonInputs(logits, targets), candidates=(candidate,) + ) + + assert report.drifts[0].provenance["actual_backend"] == "cuda-sm90" + assert report.drifts[0].provenance["lse_source"] == "direct" diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py new file mode 100644 index 00000000..a74f3b4c --- /dev/null +++ b/tests/test_logprob_contract.py @@ -0,0 +1,636 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 TP-aware logprob contract and contract-aware dispatch tests (issue #241).""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import pytest + +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDType, + LogprobOutputSpec, + LogprobRole, + MaskMode, + MaskSpec, + ReductionSpec, + ShardingSpec, + TPPlacement, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +QWEN3_REAL_VOCAB = 151936 +QWEN3_PADDED_VOCAB = 152064 + + +def _even_bounds(padded_vocab: int, tp_world_size: int) -> tuple[tuple[int, int], ...]: + shard = padded_vocab // tp_world_size + return tuple( + (rank * shard, padded_vocab if rank == tp_world_size - 1 else (rank + 1) * shard) + for rank in range(tp_world_size) + ) + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 2, + cp_rank: int = 0, + cp_world_size: int = 2, + real_vocab_size: int = QWEN3_REAL_VOCAB, + padded_vocab_size: int = QWEN3_PADDED_VOCAB, + vocab_shard_bounds: tuple[tuple[int, int], ...] | None = None, +) -> ShardingSpec: + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=( + vocab_shard_bounds + if vocab_shard_bounds is not None + else _even_bounds(padded_vocab_size, tp_world_size) + ), + real_vocab_size=real_vocab_size, + padded_vocab_size=padded_vocab_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ) + + +def _mask( + *, + num_tokens: int = 8, + active_mask: tuple[bool, ...] | None = None, + ignore_index: int = -100, +) -> MaskSpec: + return MaskSpec( + num_tokens=num_tokens, + active_mask=( + active_mask + if active_mask is not None + else (False, False, True, True, True, True, True, False) + ), + ignore_index=ignore_index, + ) + + +def _contract( + *, + role: str = "train", + dtype: str = "bf16", + mask: MaskSpec | None = None, + sharding: ShardingSpec | None = None, + reduction: ReductionSpec | None = None, +) -> LogprobContract: + return LogprobContract( + role=role, + dtype=dtype, + mask=mask if mask is not None else _mask(), + sharding=sharding if sharding is not None else _sharding(), + reduction=reduction if reduction is not None else ReductionSpec(), + ) + + +def _declared_tp_backend() -> LogprobBackendCapability: + return LogprobBackendCapability( + backend_id="test-deterministic-tp-logprob", + roles=frozenset({LogprobRole.TRAIN, LogprobRole.INFER}), + dtypes=frozenset({LogprobDType.BF16}), + tp_world_sizes=(1, 2, 4), + cp_world_sizes=None, + supports_vocab_padding=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), + exports_vocab_lse=True, + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", + ) + + +def test_qwen3_tp2_bf16_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.tp_world_size == 2 + assert contract.sharding.cp_world_size == 2 + assert contract.sharding.local_vocab_start == 0 + assert contract.sharding.local_vocab_end == QWEN3_PADDED_VOCAB // 2 + assert contract.sharding.local_vocab_size == QWEN3_PADDED_VOCAB // 2 + assert contract.mask.active_token_count == 5 + assert contract.reduction.acc_dtype is LogprobDType.FP32 + assert contract.to_dict()["reduction"] == { + "merge": "max_sumexp", + "merge_axis": "tp_vocab", + "acc_dtype": "fp32", + "order": "global_vocab_shard_index", + "transport": "all_gather", + "downcast_at": "final_write", + "engine": "in_op_reference", + "determinism_scope": "cross_tp_bitwise", + "cp_is_merge_axis": False, + } + json.dumps(contract.to_dict()) + + +@pytest.mark.parametrize("tp_world_size", [1, 2, 4]) +def test_pr4_sweep_tp_degrees_are_representable(tp_world_size): + sharding = _sharding(tp_world_size=tp_world_size, cp_world_size=1) + + assert len(sharding.vocab_shard_bounds) == tp_world_size + assert sharding.vocab_shard_bounds[-1][1] == QWEN3_PADDED_VOCAB + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == tp_world_size - 1 + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 2, "tp_rank=2"), + ("cp_rank", 2, "cp_rank=2"), + ("real_vocab_size", 0, "positive integer"), + ("padded_vocab_size", QWEN3_REAL_VOCAB - 1, "must not be smaller"), + ], +) +def test_invalid_rank_and_vocab_metadata_fail_loudly(field, value, message): + values = { + "tp_rank": 0, + "tp_world_size": 2, + "cp_rank": 0, + "cp_world_size": 2, + "real_vocab_size": QWEN3_REAL_VOCAB, + "padded_vocab_size": QWEN3_PADDED_VOCAB, + "vocab_shard_bounds": _even_bounds(QWEN3_PADDED_VOCAB, 2), + } + values[field] = value + + with pytest.raises(LogprobContractError, match=message): + ShardingSpec(**values) + + +@pytest.mark.parametrize( + ("bounds", "message"), + [ + ((), "one \\(start, end\\) pair per TP rank"), + (((0, 76032),), "one \\(start, end\\) pair per TP rank"), + (((0, 76032), (76032, 76032)), "end > start"), + (((0, 76000), (76032, 152064)), "contiguous"), + (((0, 76064), (76032, 152064)), "contiguous"), + (((0, 76032), (76032, 152000)), "cover padded_vocab_size exactly"), + ], +) +def test_incomplete_or_overlapping_vocab_shard_bounds_fail_loudly(bounds, message): + with pytest.raises(LogprobContractError, match=message): + _sharding(vocab_shard_bounds=bounds) + + +def test_owner_rank_is_unique_and_rejects_out_of_real_vocab_targets(): + sharding = _sharding() + + assert sharding.owner_rank(0) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2 - 1) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2) == 1 + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 1 + + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(-1) + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(QWEN3_REAL_VOCAB) + + +def test_active_token_mask_metadata_is_validated(): + with pytest.raises(LogprobContractError, match="one entry per token"): + _mask(num_tokens=4) + + with pytest.raises(LogprobContractError, match="must be a bool"): + MaskSpec(num_tokens=2, active_mask=(True, 1)) + + all_inactive = _mask(num_tokens=3, active_mask=(False, False, False)) + assert all_inactive.active_token_count == 0 + + +def test_reduction_requires_fp32_accumulation_and_known_semantics(): + with pytest.raises(LogprobContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + with pytest.raises(LogprobContractError, match="merge must be one of"): + ReductionSpec(merge="lse_average") + + with pytest.raises(LogprobContractError, match="transport must be one of"): + ReductionSpec(transport="all_reduce") + + +def test_contract_component_types_and_lse_export_are_enforced(): + with pytest.raises(LogprobContractError, match="mask must be a MaskSpec"): + LogprobContract( + role="train", + dtype="bf16", + mask=None, + sharding=_sharding(), + reduction=ReductionSpec(), + ) + + with pytest.raises(LogprobContractError, match="export_lse must be True"): + replace(_contract(), export_lse=False) + + +def test_ignore_index_must_not_collide_with_the_real_vocabulary(): + with pytest.raises(LogprobContractError, match="must not collide"): + _contract(mask=_mask(ignore_index=5)) + + padding_column = QWEN3_REAL_VOCAB + 1 + contract = _contract(mask=_mask(ignore_index=padding_column)) + assert contract.mask.ignore_index == padding_column + + +def _restrict_to_ws1_candidates(registry: KernelRegistry) -> None: + """Drop the #241 PR3 vocab-parallel reference so only WS1 backends remain.""" + + platform = registry._platform() + registry._logprob_candidates[platform] = [ + backend + for backend in registry._logprob_candidates[platform] + if backend is not OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP + ] + + +def test_current_ws1_backend_rejects_strict_tp_contract_without_fallback(): + registry = KernelRegistry() + _restrict_to_ws1_candidates(registry) + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(_contract(), requested_backend="reference") + + message = str(exc_info.value) + assert "TP=2 is unsupported" in message + assert "vocab-domain LSE export is unsupported" in message + assert "determinism_scope=cross_tp_bitwise is unsupported" in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_current_ws1_backend_rejects_padded_vocab_even_at_tp1(): + registry = KernelRegistry() + _restrict_to_ws1_candidates(registry) + contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(contract) + + message = str(exc_info.value) + assert "TP=1 is unsupported" not in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_ws1_rejections_recorded_when_vocab_parallel_reference_resolves(): + """The WS1 backends still reject strict contracts; they are skipped with + recorded reasons while dispatch resolves the #241 PR3 reference.""" + + registry = KernelRegistry() + platform = registry._platform() + # Order the WS1 backends ahead of the reference so their rejections are + # exercised on the way to a successful resolution. + candidates = registry._logprob_candidates[platform] + candidates.remove(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) + candidates.append(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + assert result.capability.backend_id == "pytorch-vocab-parallel-logp-ws2" + assert result.provenance["fallback"] is True + rejections = " | ".join(result.provenance["prior_rejections"]) + assert "vocab-domain LSE export is unsupported" in rejections + + +def test_undeclared_backend_capability_is_never_selected(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [OpBackend.PYTORCH_NATIVE] + + tp1_contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + with pytest.raises(RuntimeError, match="no LogprobBackendCapability declared"): + registry.get_logprob_op(tp1_contract) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert result.provenance["requested_backend"] == "reference" + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["real_vocab_size"] == QWEN3_REAL_VOCAB + assert result.provenance["contract"]["reduction"]["cp_is_merge_axis"] is False + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): + registry.get_logprob_op(_contract(), requested_backend="another-backend") + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + + +def test_cp_is_a_non_merge_axis_and_cp_agnostic_backends_accept_any_cp_degree(): + capability = _declared_tp_backend() + cp2_contract = _contract(sharding=_sharding(cp_world_size=2, cp_rank=1)) + + assert capability.incompatibilities(cp2_contract) == () + + cp_restricted = replace(capability, cp_world_sizes=(1,)) + assert cp_restricted.incompatibilities(cp2_contract) == ("CP=2 is unsupported",) + + +def test_inactive_tokens_require_explicit_active_mask_support(): + capability = replace(_declared_tp_backend(), mask_modes=frozenset({MaskMode.IGNORE_INDEX})) + contract = _contract() + + assert "explicit active-token masking is unsupported" in ( + capability.incompatibilities(contract) + ) + + fully_active = _contract(mask=_mask(num_tokens=3, active_mask=(True, True, True))) + assert capability.incompatibilities(fully_active) == () + + +def test_backend_id_must_not_shadow_a_reserved_policy_keyword(): + with pytest.raises(LogprobContractError, match="reserved dispatch policy keyword"): + replace(_declared_tp_backend(), backend_id="Deterministic") + + +def test_default_auto_policy_resolves_any_compatible_implementation_kind(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), implementation_kind="reference"), + platform=platform, + ) + + tp1_contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + result = registry.get_logprob_op(tp1_contract) + + assert result.provenance["requested_backend"] == "auto" + assert result.capability.implementation_kind == "reference" + + +def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="REFERENCE") + assert result.capability.backend_id == "test-deterministic-tp-logprob" + + with pytest.raises(RuntimeError, match="does not match requested_backend"): + registry.get_logprob_op(_contract(), requested_backend="Test-Deterministic-TP-Logprob") + + +def test_policy_only_skips_are_not_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="other-compatible-backend"), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_capability_rejections_are_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + + assert result.provenance["fallback"] is True + assert "TP=2 is unsupported" in result.provenance["prior_rejections"][0] + + +def test_ws2_candidate_list_is_decoupled_from_the_legacy_priority_map(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform].insert(0, OpBackend.PYTORCH_NATIVE) + + legacy = registry._priority_map[platform]["batch_invariant_logp"] + assert OpBackend.PYTORCH_NATIVE not in legacy + + legacy.insert(0, OpBackend.PYTORCH_GEMM) + assert OpBackend.PYTORCH_GEMM not in registry._logprob_candidates[platform] + + +def test_register_logprob_backend_is_the_public_registration_seam(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + capability = _declared_tp_backend() + + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, capability, platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(capability, backend_id="replacement-backend"), + platform=platform, + ) + + assert registry._logprob_candidates[platform] == [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + result = registry.get_logprob_op(_contract(), requested_backend="reference") + assert result.capability.backend_id == "replacement-backend" + + with pytest.raises(LogprobContractError, match="capability must be"): + registry.register_logprob_backend(OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, None) + + +def test_backend_id_whitespace_is_normalized_for_dispatch(): + capability = replace(_declared_tp_backend(), backend_id=" padded-id ") + assert capability.backend_id == "padded-id" + + +def test_capabilities_are_scoped_per_platform(): + registry = KernelRegistry() + platform = registry._platform() + other = "rocm" if platform != "rocm" else "cpu" + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="other-platform-backend"), + platform=other, + ) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert ( + registry._logprob_capabilities[other][OpBackend.PYTORCH_BATCH_INVARIANT_LOGP].backend_id + == "other-platform-backend" + ) + + +def test_register_logprob_backend_rejects_unknown_platform(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="unsupported platform"): + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + _declared_tp_backend(), + platform="cuda-typo", + ) + + +def test_non_iterable_roles_and_dtypes_raise_contract_errors(): + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), roles=None) + + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), dtypes=42) + + +def test_requested_deterministic_policy_is_a_loud_error(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="determinism_scope"): + registry.get_logprob_op(_contract(), requested_backend="deterministic") + + +def test_auto_policy_is_rejected_for_tp_sharded_contracts(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + with pytest.raises(LogprobContractError, match="Unsafe dispatch"): + registry.get_logprob_op(_contract()) + + with pytest.raises(LogprobContractError, match="Unsafe dispatch"): + registry.get_logprob_op(_contract(), requested_backend=" AUTO ") + + tp1_contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + assert registry.get_logprob_op(tp1_contract).provenance["requested_backend"] == "auto" + + +def test_determinism_scope_is_part_of_the_typed_contract(): + fixed_only = replace( + _declared_tp_backend(), + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + ) + + assert "determinism_scope=cross_tp_bitwise is unsupported" in ( + fixed_only.incompatibilities(_contract()) + ) + + relaxed = _contract(reduction=ReductionSpec(determinism_scope="fixed_topology")) + assert fixed_only.incompatibilities(relaxed) == () + + +def test_policy_filtered_candidates_never_count_toward_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_output_spec_is_pinned_to_fp32_replicated(): + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(selected_logp_dtype="bf16") + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(lse_dtype="bf16") + + assert LogprobOutputSpec().tp_placement is TPPlacement.REPLICATED + assert _contract().to_dict()["output"] == { + "selected_logp_dtype": "fp32", + "lse_dtype": "fp32", + "tp_placement": "replicated", + } + + +def test_cross_rank_fingerprint_is_rank_independent_and_content_sensitive(): + rank0 = _contract(sharding=_sharding(tp_rank=0)) + rank1 = _contract(sharding=_sharding(tp_rank=1, cp_rank=1)) + + assert rank0.cross_rank_fingerprint() == rank1.cross_rank_fingerprint() + + different_mask = _contract( + mask=_mask(active_mask=(True, True, True, True, True, True, True, False)) + ) + assert rank0.cross_rank_fingerprint() != different_mask.cross_rank_fingerprint() + + +def test_provenance_records_the_active_mask_digest(): + provenance_mask = _contract().to_dict()["mask"] + + assert provenance_mask["active_mask_sha256"] == _mask().active_mask_sha256 + assert len(provenance_mask["active_mask_sha256"]) == 64 + + same_count_different_mask = _mask( + active_mask=(True, True, True, True, True, False, False, False) + ) + assert same_count_different_mask.active_token_count == _mask().active_token_count + assert same_count_different_mask.active_mask_sha256 != _mask().active_mask_sha256 + + +def test_padding_only_shard_is_constructible_for_the_identity_partial(): + sharding = _sharding( + vocab_shard_bounds=((0, QWEN3_REAL_VOCAB), (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB)), + ) + + assert sharding.local_vocab_start == 0 + assert sharding.vocab_shard_bounds[1] == (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB) + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 0 diff --git a/tests/test_vocab_parallel_logp.py b/tests/test_vocab_parallel_logp.py new file mode 100644 index 00000000..87da2ef0 --- /dev/null +++ b/tests/test_vocab_parallel_logp.py @@ -0,0 +1,537 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic vocab-parallel TP logprob reference tests""" + +from __future__ import annotations + +import queue +import tempfile +import traceback +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.kernels.gtest.tolerance import load_contract +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobContract, + LogprobContractError, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + BACKEND_ID, + VocabParallelLogprobOp, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +REAL_VOCAB = 27 +PADDED_VOCAB = 32 +NUM_TILES = 8 +NUM_TOKENS = 6 +ACTIVE = (True, True, True, True, True, False) + + +def _even_bounds(padded: int, world: int) -> tuple[tuple[int, int], ...]: + shard = padded // world + return tuple( + (rank * shard, padded if rank == world - 1 else (rank + 1) * shard) for rank in range(world) + ) + + +def _contract( + *, + tp_rank: int = 0, + tp_world_size: int = 1, + bounds: tuple[tuple[int, int], ...] | None = None, + real_vocab: int = REAL_VOCAB, + padded_vocab: int = PADDED_VOCAB, + num_tokens: int = NUM_TOKENS, + active: tuple[bool, ...] = ACTIVE, + dtype: str = "fp32", +) -> LogprobContract: + return LogprobContract( + role="train", + dtype=dtype, + mask=MaskSpec(num_tokens=num_tokens, active_mask=active), + sharding=ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=( + bounds if bounds is not None else _even_bounds(padded_vocab, tp_world_size) + ), + real_vocab_size=real_vocab, + padded_vocab_size=padded_vocab, + ), + reduction=ReductionSpec(), + ) + + +def _inputs(dtype=torch.float32, seed: int = 2026): + torch.manual_seed(seed) + logits = torch.randn(NUM_TOKENS, PADDED_VOCAB, dtype=torch.float32).to(dtype) + targets = torch.tensor([1, 5, REAL_VOCAB - 1, 0, 13, -100]) + return logits, targets + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + view_dtype = { + torch.float32: torch.int32, + torch.bfloat16: torch.int16, + torch.float16: torch.int16, + }[tensor.dtype] + return tensor.contiguous().view(view_dtype) + + +def _bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and bool((_bits(a) == _bits(b)).all()) + + +def _case_shard_size_mismatch(): + logits, targets = _inputs() + return logits, targets, _contract(tp_rank=0, tp_world_size=2), NUM_TILES, "vocab columns" + + +def _case_mask_length_mismatch(): + logits, targets = _inputs() + contract = _contract(num_tokens=NUM_TOKENS + 1, active=ACTIVE + (True,)) + return logits, targets, contract, NUM_TILES, "num_tokens" + + +def _case_dtype_mismatch(): + logits, targets = _inputs() + return logits, targets, _contract(dtype="bf16"), NUM_TILES, "dtype" + + +def _case_tile_misaligned_bounds(): + # Tile size is 32/8 = 4; a boundary at 6 is misaligned. + logits, targets = _inputs() + contract = _contract(tp_world_size=2, bounds=((0, 6), (6, 32))) + return logits[:, :6], targets, contract, NUM_TILES, "tile" + + +def _case_bad_num_vocab_tiles(): + logits, targets = _inputs() + return logits, targets, _contract(), 7, "num_vocab_tiles" + + +def _case_active_target_out_of_real_vocab(): + logits, targets = _inputs() + bad_targets = targets.clone() + bad_targets[0] = REAL_VOCAB # padding column, active row + return logits, bad_targets, _contract(), NUM_TILES, "real vocabulary" + + +def _case_all_inf_active_row(): + logits, targets = _inputs() + poisoned = logits.clone() + poisoned[0, :] = float("-inf") + return poisoned, targets, _contract(), NUM_TILES, "non-finite" + + +@pytest.mark.parametrize( + "case", + [ + _case_shard_size_mismatch, + _case_mask_length_mismatch, + _case_dtype_mismatch, + _case_tile_misaligned_bounds, + _case_bad_num_vocab_tiles, + _case_active_target_out_of_real_vocab, + _case_all_inf_active_row, + ], + ids=lambda fn: fn.__name__.removeprefix("_case_"), +) +def test_invalid_invocations_fail_loudly(case): + logits, targets, contract, num_tiles, match = case() + with pytest.raises(LogprobContractError, match=match): + VocabParallelLogprobOp()(logits, targets, contract=contract, num_vocab_tiles=num_tiles) + + +class TestSingleRank: + def test_repeated_runs_are_bitwise_identical(self): + contract = _contract() + logits, targets = _inputs() + op = VocabParallelLogprobOp() + logp_a, lse_a = op(logits, targets, contract=contract, num_vocab_tiles=NUM_TILES) + logp_b, lse_b = op(logits, targets, contract=contract, num_vocab_tiles=NUM_TILES) + assert _bitwise_equal(logp_a, logp_b) + assert _bitwise_equal(lse_a, lse_b) + + def test_batch_invariance_same_row_any_context(self): + contract_full = _contract() + logits, targets = _inputs() + op = VocabParallelLogprobOp() + logp_full, lse_full = op(logits, targets, contract=contract_full, num_vocab_tiles=NUM_TILES) + + contract_single = _contract(num_tokens=1, active=(True,)) + logp_one, lse_one = op( + logits[2:3], targets[2:3], contract=contract_single, num_vocab_tiles=NUM_TILES + ) + assert _bitwise_equal(logp_full[2:3], logp_one) + assert _bitwise_equal(lse_full[2:3], lse_one) + + def test_matches_ws1_batch_invariant_logp_within_contract_tolerance(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + contract = _contract(padded_vocab=REAL_VOCAB + 5) + # Use a real==padded contract so the WS1 op sees identical logits. + contract = _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB) + logits, targets = _inputs() + logp, _ = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + ws1 = NativeBatchInvariantLogpOp().apply(logits, targets) + active = torch.tensor(ACTIVE) + assert torch.allclose( + logp[active], ws1[active], atol=tolerance["atol"], rtol=tolerance["rtol"] + ) + + def test_padding_columns_are_excluded_and_finite(self): + contract = _contract() + logits, targets = _inputs() + boosted = logits.clone() + boosted[:, REAL_VOCAB:] = 1e4 # huge padding logits must not leak into LSE + logp, lse = VocabParallelLogprobOp()( + boosted, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + ref_lse = torch.logsumexp(boosted[:, :REAL_VOCAB].float(), dim=-1) + assert torch.isfinite(logp).all() and torch.isfinite(lse).all() + assert torch.allclose(lse, ref_lse, atol=1e-5) + + def test_inactive_rows_zero_filled_lse_still_exported(self): + contract = _contract() + logits, targets = _inputs() + logp, lse = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + assert logp[-1].item() == 0.0 + assert torch.isfinite(lse[-1]) + + +class TestBackward: + def test_grads_match_autograd_oracle(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + contract = _contract() + logits, targets = _inputs() + x = logits.clone().requires_grad_(True) + logp, lse = VocabParallelLogprobOp()( + x, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + (logp.sum() + 0.5 * lse.sum()).backward() + + y = logits.clone().requires_grad_(True) + ref_lse = torch.logsumexp(y[:, :REAL_VOCAB].float(), dim=-1) + safe = targets.clamp(0, REAL_VOCAB - 1) + ref_logp = y[torch.arange(NUM_TOKENS), safe].float() - ref_lse + ref_logp = torch.where(torch.tensor(ACTIVE), ref_logp, torch.zeros_like(ref_logp)) + (ref_logp.sum() + 0.5 * ref_lse.sum()).backward() + + assert torch.allclose(x.grad, y.grad, atol=tolerance["atol"], rtol=tolerance["rtol"]) + assert bool((x.grad[:, REAL_VOCAB:] == 0).all()) + + # No grad requested -> outputs detached from autograd entirely. + logp_ng, lse_ng = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + assert not logp_ng.requires_grad and not lse_ng.requires_grad + + def test_inactive_rows_grad_asymmetry(self): + """The logp term is zeroed on inactive rows; the lse term still flows — + lse is a row property exported (and differentiable) for every row.""" + + contract = _contract() + logits, targets = _inputs() + + x = logits.clone().requires_grad_(True) + _, lse = VocabParallelLogprobOp()(x, targets, contract=contract, num_vocab_tiles=NUM_TILES) + lse.sum().backward() + assert bool((x.grad[-1, :REAL_VOCAB].abs() > 0).any()) + + z = logits.clone().requires_grad_(True) + logp, _ = VocabParallelLogprobOp()(z, targets, contract=contract, num_vocab_tiles=NUM_TILES) + logp.sum().backward() + assert bool((z.grad[-1] == 0).all()) + + +def test_dispatch_resolves_reference_and_leaves_legacy_untouched(): + registry = KernelRegistry() + contract = _contract() + + result = registry.get_logprob_op(contract) + assert result.capability.backend_id == BACKEND_ID + assert result.provenance["fallback"] is False + assert isinstance(result.op, VocabParallelLogprobOp) + assert ( + result.provenance["contract"]["reduction"]["determinism_scope"] + == DeterminismScope.CROSS_TP_BITWISE.value + ) + + by_id = registry.get_logprob_op(contract, requested_backend=BACKEND_ID) + assert by_id.capability.backend_id == BACKEND_ID + by_kind = registry.get_logprob_op(contract, requested_backend="reference") + assert by_kind.capability.backend_id == BACKEND_ID + + for ops in registry._priority_map.values(): + for candidates in ops.values(): + assert OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP not in candidates + + +# Cross-TP bitwise determinism on real ranks (NCCL, one CUDA device per rank) +TP_REAL_VOCAB = 1000 +TP_PADDED_VOCAB = 1024 +TP_NUM_TILES = 32 # tile = 32 columns +TP_TILE = TP_PADDED_VOCAB // TP_NUM_TILES +TP_NUM_TOKENS = 48 +TP_ACTIVE = tuple(index % 7 != 5 for index in range(TP_NUM_TOKENS)) +TP_DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16} +_SPAWN_TIMEOUT_S = 300 + + +def _cuda_device_count() -> int: + return torch.cuda.device_count() if torch.cuda.is_available() else 0 + + +def _requires_gpus(count: int): + return pytest.mark.skipif( + _cuda_device_count() < count, + reason=f"cross-TP determinism needs {count} CUDA devices to place one rank per device", + ) + + +def _tile_counts(world_size: int, uneven: bool) -> list[int]: + """Tiles per rank; bounds are built from whole tiles so they stay tile-aligned.""" + + counts = [TP_NUM_TILES // world_size for _ in range(world_size)] + counts[-1] += TP_NUM_TILES % world_size + if uneven: + for rank in range(world_size - 1): + if counts[rank] > 1: + counts[rank] -= 1 + counts[-1] += 1 + return counts + + +def _tp_bounds(world_size: int, uneven: bool) -> tuple[tuple[int, int], ...]: + bounds, cursor = [], 0 + for count in _tile_counts(world_size, uneven): + bounds.append((cursor, cursor + count * TP_TILE)) + cursor += count * TP_TILE + return tuple(bounds) + + +def _tp_contract(tp_rank: int, tp_world_size: int, bounds, dtype_name: str) -> LogprobContract: + return _contract( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + bounds=bounds, + real_vocab=TP_REAL_VOCAB, + padded_vocab=TP_PADDED_VOCAB, + num_tokens=TP_NUM_TOKENS, + active=TP_ACTIVE, + dtype=dtype_name, + ) + + +def _tp_inputs(device, dtype, seed: int = 2026): + """Identical logits and targets on every rank, seeded on CPU.""" + + gen = torch.Generator(device="cpu").manual_seed(seed) + logits = torch.randn(TP_NUM_TOKENS, TP_PADDED_VOCAB, generator=gen, dtype=torch.float32) + targets = torch.randint(0, TP_REAL_VOCAB, (TP_NUM_TOKENS,), generator=gen) + active = torch.tensor(TP_ACTIVE) + # Inactive rows carry ignore_index; active_mask stays the sole authority. + targets = torch.where(active, targets, torch.full_like(targets, -100)) + return logits.to(device=device, dtype=dtype), targets.to(device) + + +def _nccl_worker(rank, world_size, init_method, result_queue, scenario, uneven, dtype_name): + import torch.distributed as dist + + try: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + backend="nccl", init_method=init_method, rank=rank, world_size=world_size + ) + dtype = TP_DTYPES[dtype_name] + op = VocabParallelLogprobOp() + bounds = _tp_bounds(world_size, uneven) + logits, targets = _tp_inputs(device, dtype) + tiles = TP_NUM_TILES + + if scenario in {"preflight", "misaligned"}: + if scenario == "preflight": + if rank == 0: + tiles = TP_NUM_TILES * 2 + else: + # Nudge the first boundary off the tile grid, on every rank. + split = bounds[0][1] + TP_TILE // 4 + bounds = ((0, split), (split, bounds[1][1])) + bounds[2:] + + start, end = bounds[rank] + try: + op( + logits[:, start:end].contiguous().clone(), + targets, + contract=_tp_contract(rank, world_size, bounds, dtype_name), + tp_group=dist.group.WORLD, + num_vocab_tiles=tiles, + ) + result_queue.put({"ok": False, "rank": rank, "traceback": "no error raised"}) + except LogprobContractError as exc: + result_queue.put({"ok": True, "rank": rank, "message": str(exc)}) + return + + start, end = bounds[rank] + shard = logits[:, start:end].contiguous().clone().requires_grad_(True) + tp_contract = _tp_contract(rank, world_size, bounds, dtype_name) + logp_tp, lse_tp = op( + shard, + targets, + contract=tp_contract, + tp_group=dist.group.WORLD, + num_vocab_tiles=TP_NUM_TILES, + ) + (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() + + # Same ranks, same inputs, run again: the collectives must not perturb bits. + rerun = logits[:, start:end].contiguous().clone() + logp_re, lse_re = op( + rerun, + targets, + contract=tp_contract, + tp_group=dist.group.WORLD, + num_vocab_tiles=TP_NUM_TILES, + ) + + # In-process TP=1 run on the full logits: the cross-TP claim is that a + # TP=n result equals the TP=1 result, bit for bit. + full = logits.clone().requires_grad_(True) + logp_one, lse_one = op( + full, + targets, + contract=_tp_contract(0, 1, ((0, TP_PADDED_VOCAB),), dtype_name), + num_vocab_tiles=TP_NUM_TILES, + ) + (logp_one.sum() + 0.5 * lse_one.sum()).backward() + + result_queue.put( + { + "ok": True, + "rank": rank, + "logp_bits_match": _bitwise_equal(logp_tp, logp_one), + "lse_bits_match": _bitwise_equal(lse_tp, lse_one), + "grad_bits_match": _bitwise_equal(shard.grad, full.grad[:, start:end]), + "rerun_bits_match": ( + _bitwise_equal(logp_re, logp_tp) and _bitwise_equal(lse_re, lse_tp) + ), + "logp_bit_pattern": _bits(logp_tp.detach().float().cpu()).tolist(), + "lse_bit_pattern": _bits(lse_tp.detach().float().cpu()).tolist(), + } + ) + except Exception: # pragma: no cover - forwarded to the parent process + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + import torch.distributed as dist + + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_nccl_scenario(world_size, scenario="correctness", uneven=False, dtype_name="fp32"): + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmpdir: + init_method = (Path(tmpdir) / "nccl_init").as_uri() + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=_nccl_worker, + args=(rank, world_size, init_method, result_queue, scenario, uneven, dtype_name), + ) + for rank in range(world_size) + ] + results = [] + try: + for process in processes: + process.start() + for _ in range(world_size): + try: + results.append(result_queue.get(timeout=_SPAWN_TIMEOUT_S)) + except queue.Empty: + for process in processes: + if process.is_alive(): + process.terminate() + pytest.fail(f"timed out waiting for NCCL workers (scenario={scenario})") + finally: + for process in processes: + process.join(timeout=30) + if process.is_alive(): + process.terminate() + results.sort(key=lambda item: item["rank"]) + for result in results: + assert result["ok"], result.get("traceback") + for process in processes: + assert process.exitcode == 0 + return results + + +class TestCrossTPBitwise: + """TP=n output == TP=1 output, bit for bit, on real NCCL ranks.""" + + @_requires_gpus(2) + @pytest.mark.parametrize("dtype_name", ["fp32", "bf16"]) + @pytest.mark.parametrize("uneven", [False, True], ids=["even", "uneven"]) + def test_tp2_bitwise_identical_to_tp1(self, uneven, dtype_name): + self._assert_matches_tp1(_run_nccl_scenario(2, uneven=uneven, dtype_name=dtype_name)) + + @_requires_gpus(4) + @pytest.mark.parametrize("dtype_name", ["fp32", "bf16"]) + @pytest.mark.parametrize("uneven", [False, True], ids=["even", "uneven"]) + def test_tp4_bitwise_identical_to_tp1(self, uneven, dtype_name): + self._assert_matches_tp1(_run_nccl_scenario(4, uneven=uneven, dtype_name=dtype_name)) + + @staticmethod + def _assert_matches_tp1(results): + for result in results: + rank = result["rank"] + assert result["logp_bits_match"], f"rank {rank} logp bits differ from TP=1" + assert result["lse_bits_match"], f"rank {rank} lse bits differ from TP=1" + assert result["grad_bits_match"], f"rank {rank} grad bits differ from TP=1" + assert result["rerun_bits_match"], f"rank {rank} bits changed between identical runs" + # Outputs are replicated: every rank must hold identical bits. + for other in results[1:]: + assert results[0]["logp_bit_pattern"] == other["logp_bit_pattern"] + assert results[0]["lse_bit_pattern"] == other["lse_bit_pattern"] + + @_requires_gpus(2) + def test_tp2_and_tp4_agree_with_each_other(self): + """The claim is over TP degrees, so pin TP=2 against TP=4 directly.""" + + if _cuda_device_count() < 4: + pytest.skip("needs 4 CUDA devices to compare TP=2 against TP=4") + tp2 = _run_nccl_scenario(2) + tp4 = _run_nccl_scenario(4) + assert tp2[0]["logp_bit_pattern"] == tp4[0]["logp_bit_pattern"] + assert tp2[0]["lse_bit_pattern"] == tp4[0]["lse_bit_pattern"] + + +class TestCrossTPGuards: + """A disagreement must abort loudly on every rank, not strand ranks in a collective.""" + + @_requires_gpus(2) + def test_preflight_rejects_mismatched_num_vocab_tiles(self): + results = _run_nccl_scenario(2, scenario="preflight") + for result in results: + assert "cross-rank preflight failed" in result["message"] + + @_requires_gpus(2) + def test_misaligned_shard_bounds_rejected(self): + results = _run_nccl_scenario(2, scenario="misaligned") + for result in results: + assert "not aligned to the vocab tile size" in result["message"]