diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..2d9ba91e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,6 +76,12 @@ 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 + 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..5068c33f 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -54,6 +54,32 @@ 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`) +**TP=1, TP=2, and TP=4 produce bit-identical results.** + +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 @@ -224,3 +250,6 @@ WSL/Linux with CUDA. - `rl_engine/kernels/registry.py` - `tests/test_batch_invariant_logp.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/docs/operators/grpo-loss.md b/docs/operators/grpo-loss.md index 07fcac16..fbfe5751 100644 --- a/docs/operators/grpo-loss.md +++ b/docs/operators/grpo-loss.md @@ -15,6 +15,11 @@ fused ratio/KL kernel (logits → ratio/KL via online softmax), and the group-no logits --[ratio_kl op]--> (ratio, kl) --[group adv + clipped surrogate]--> loss ``` +The backends above consume dense `[B, T, V]` logits and reduce with a plain masked +mean, so they are single-shard only. For vocab-parallel TP, or when the reduction +must be bitwise reproducible across parallel degrees, see [Tensor and Data +Parallel](#tensor-and-data-parallel) below. + ## Entry Point ```python from rl_engine.kernels.registry import kernel_registry @@ -74,6 +79,117 @@ op mirrors this using `NativeRatioKLOp`. Gradients flow into `policy_logits` only (`ref_logits` is frozen; `old_logps` is cached). +## Tensor and Data Parallel + +`DistributedGRPOLossOp` +(`rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py`) +**Every TP × DP degree produces bit-identical loss, per-sequence totals, and +gradients.** It is a reference backend on top of the deterministic +[vocab-parallel logprob](batch-invariant-logp.md#tensor-parallel); the backends +above stay the default single-GPU path. + +The objective is elementwise — ratio, clipping, reference KL and the group-relative +advantage all act per token or per sequence — so the only place the parallel layout +can change the answer is the final sum over tokens. + +1. Selected logprobs come from `VocabParallelLogprobOp`, so ratio and KL inherit + cross-TP bitwise equality for free. TP needs no further handling here: by the + time the objective sees a logprob, the vocabulary has already been reduced. +2. A DP rank owns each of its sequences **whole**. `sequence_shard_bounds` is a + contiguous `[0, num_sequences)` partition in DP-rank order, exactly like + `ShardingSpec.vocab_shard_bounds` for the vocabulary. +3. Two nested reductions, each with a contract-fixed extent: `padded_seq_len` + token slots → one sequence total (entirely local, since the rank owns the + whole sequence); `num_sequences` totals → the scalar numerator. +4. Only per-sequence totals cross a rank boundary. They travel by `all_gather` + and are concatenated in DP-rank order into a `[num_sequences]` vector. The + collective moves bytes and placement is an exact copy, so every degree + performs identical arithmetic on identical inputs. `all_reduce` is excluded on + purpose: its combine order follows the collective's topology, not the declared + sequence order. +5. Advantages are replicated, not merged — rewards are one scalar per sequence, + so every rank normalizes every group over the identical `[num_sequences]` + tensor and keeps its own slice. This is what lets an advantage group straddle + DP ranks with no extra machinery. +6. The normalizer divides by a **global** active-token count, gathered as + integers so it is exact at every degree. + +Step 3 is why the determinism argument is short: because no sequence's token sum +is ever split across ranks, there is no partial-sequence state to merge and no +alignment rule to get wrong. + +```python +sequence_shard_bounds = ((0, 8),) # DP=1 +sequence_shard_bounds = ((0, 4), (4, 8)) # DP=2 +``` + +### Context parallelism is out of scope + +`cp_world_size` must be 1; anything else raises `LossContractError` rather than +silently reducing over a partial batch. CP is an attention-level concern — attention +is the only op with a cross-token dependency — and this operator consumes logits, by +which point CP has already been resolved upstream. Supporting it here would mean +splitting a sequence's token sum across ranks and reducing over the very axis the +logprob contract declares a *non-merge* axis. In a CP job that reduction belongs to +the caller. `cp_rank`/`cp_world_size` are carried for provenance only, mirroring +`ShardingSpec`. + +### Normalizer semantics + +`TokenNormalizer` makes the GRPO normalizer ambiguity explicit. The modes differ by +more than a scale factor once sequence lengths vary, so the choice is part of the +numerical identity and travels in the contract fingerprint. + +| Mode | Denominator | Notes | +| --- | --- | --- | +| `global_active_tokens` (default) | active tokens in the global batch | Matches `NativeGRPOLossOp`'s masked mean at DP=1. Long sequences weigh more. | +| `per_sequence_then_mean` | per-sequence count, then mean over live sequences | Original GRPO form; sequences weigh equally. | +| `fixed_constant` | declared constant | Dr.GRPO form; independent of the mask. | + +Usage goes through the contract-aware entry point: + +```python +from rl_engine.kernels.registry import kernel_registry + +dispatched = kernel_registry.get_loss_op(contract) # GRPOLossContract from +result = dispatched.op.apply( # rl_engine.kernels.loss_contract + policy_local_logits, # [n, local_vocab] differentiable + action_ids, # [n] + old_logps, # [n] + rewards, # [local_num_sequences] + contract=contract, + ref_local_logits=ref_local_logits, # required when beta > 0 + tp_group=tp_group, # vocab-parallel subgroup + dp_group=dp_group, # data-parallel subgroup +) +result.loss.backward() # gradients flow into policy_local_logits only + +loss, policy_loss, kl = result # unpacks like the single-GPU op +``` + +A preflight `all_gather_object` runs on **both** the DP and TP axes before any other +collective. Neither alone suffices: the loss is replicated across TP, so two TP +siblings disagreeing on `beta` would compute different losses for one sharded model, +and the logprob path's own preflight cannot see that because `beta` is not part of +the logprob contract. Other loud failures, with no silent fallback: sequence bounds +that are non-contiguous or leave a gap; a nested logprob contract whose token count +disagrees with the owned sequences; a determinism scope stronger or weaker than the +logprob path's; and population-std advantages over a singleton group. + +### Comparing configurations + +Compare `per_sequence_policy` / `per_sequence_kl`, not the scalar loss. Measured on +this operator's test inputs, regrouping the token sum — a real change of the +summation tree — moves the per-sequence vector in 12 of 12 seeds but the scalar loss +in only 5 of 12: averaging `num_sequences` totals into one fp32 number rounds most +reorderings away. A drift report that compares only the scalar will under-report +reduction differences. `GRPOLossResult` exposes both, plus `advantages`, +`per_sequence_active_tokens`, and a `provenance` dict. + +Bitwise here means *across parallel degrees on one PyTorch build and GPU model*. It +rests on PyTorch's reduction kernels being deterministic for a fixed shape on a fixed +device; cross-version and cross-architecture equality is neither tested nor claimed. + ## Accuracy Reference semantics (`NativeGRPOLossOp`): @@ -92,6 +208,11 @@ loss = masked_mean(policy, completion_mask) + beta * masked_mean(kl, completion_ The Triton op matches the native reference (forward and backward) to `atol=1e-4`. +For `DistributedGRPOLossOp`, the reference-equals-policy identity is exact rather +than approximate: with `ref_logits is policy_logits` and `old_logps == logp_policy` +the ratio is `exp(0) = 1` bitwise, so the result is invariant to the clip epsilon, +and the KL is exactly `0.0`. + ## Performance Notes The cost is dominated by the [`ratio_kl`](ratio-kl.md) stage (the vocab-dimension work); @@ -118,18 +239,31 @@ online — the forward peak is independent of `V`. ## Tests ```bash -python -m pytest tests/test_grpo_loss.py -v +python -m pytest tests/test_grpo_loss.py -v # single-GPU backends +python -m pytest tests/test_grpo_loss_contract.py -v # TP/DP/CP contract, CPU only +python -m pytest tests/test_distributed_grpo_loss.py -v ``` -Covers the native reference (group advantages + loss from logits), Triton forward/backward -vs native, masked-token invariance, an SGD loss step, and registry dispatch. Triton tests -skip without CUDA + Triton. +`test_grpo_loss.py` covers the native reference (group advantages + loss from logits), +Triton forward/backward vs native, masked-token invariance, an SGD loss step, and +registry dispatch. Triton tests skip without CUDA + Triton. + +`test_distributed_grpo_loss.py` covers every `(TP, DP)` combination reachable with +four ranks — `tp2`, `tp4`, `dp2`, `dp4`, `tp2xdp2` — each compared bitwise against a +single-rank GPU baseline, plus the KL=0 identity, run-to-run stability, negative +controls, and the DP- and TP-axis preflight guards. Larger degrees run unchanged on a +bigger node. Multi-rank tests need one GPU per rank and skip otherwise; they are +deliberately small (1000-token vocabulary, 8 sequences of 32 slots) and each worker +caps itself with `torch.cuda.set_per_process_memory_fraction`, so the suite can share +a node with a running training job. ## Implementation Files - `rl_engine/kernels/ops/pytorch/loss/grpo_loss.py` - `rl_engine/kernels/ops/triton/loss/grpo_loss.py` - `rl_engine/kernels/ops/triton/loss/ratio_kl.py`, `rl_engine/kernels/ops/pytorch/loss/ratio_kl.py` -- `rl_engine/kernels/registry.py` -- `tests/test_grpo_loss.py` +- `rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py` +- `rl_engine/kernels/loss_contract.py` +- `rl_engine/kernels/registry.py` (`register_loss_backend`, `get_loss_op`) +- `tests/test_grpo_loss.py`, `tests/test_grpo_loss_contract.py`, `tests/test_distributed_grpo_loss.py` - `benchmarks/benchmark_ratio_kl.py` diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py new file mode 100644 index 00000000..b0bcd7ab --- /dev/null +++ b/rl_engine/kernels/logprob_contract.py @@ -0,0 +1,652 @@ +# 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"} + } + 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/loss_contract.py b/rl_engine/kernels/loss_contract.py new file mode 100644 index 00000000..78f2e280 --- /dev/null +++ b/rl_engine/kernels/loss_contract.py @@ -0,0 +1,784 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for deterministic GRPO loss on the TP-aware logprob path. + +The GRPO objective consumes selected-token log-probabilities and reduces them +to a scalar:: + + ratio_t = exp(logp_policy_t - old_logp_t) + surrogate = -min(ratio_t * adv_t, clip(ratio_t) * adv_t) + loss = normalize(sum_t surrogate_t) + beta * normalize(sum_t kl_t) +""" + +from __future__ import annotations + +import hashlib +import json +import math +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, TypeVar + +from rl_engine.kernels.logprob_contract import ( + IMPLEMENTATION_KINDS, + RESERVED_DISPATCH_POLICIES, + DeterminismScope, + DowncastPoint, + LogprobContract, + LogprobDType, +) + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +class LossContractError(ValueError): + """Raised when loss metadata does not describe a valid GRPO invocation.""" + + +class TokenNormalizer(str, Enum): + """Denominator applied to the summed per-token loss terms. + + ``global_active_tokens``: divide by the number of active tokens in the + *global* batch, gathered across every DP rank. Long sequences therefore + contribute proportionally more. This matches the existing single-GPU + ``NativeGRPOLossOp`` masked mean at DP=1. + + ``per_sequence_then_mean``: divide each sequence's sum by that sequence's + own active-token count, then average over sequences that hold at least one + active token. Sequences are weighted equally regardless of length. + + ``fixed_constant``: divide by a declared constant, independent of the mask. + Requires ``LossReductionSpec.fixed_normalizer_constant``. + + The three differ by more than a scale factor once sequence lengths vary, so + the choice is part of the numerical identity and travels in the fingerprint. + """ + + GLOBAL_ACTIVE_TOKENS = "global_active_tokens" + PER_SEQUENCE_THEN_MEAN = "per_sequence_then_mean" + FIXED_CONSTANT = "fixed_constant" + + +class SummationOrder(str, Enum): + """Fixed combine order for per-token partials. + + ``sequence_major_fixed``: within one sequence, tokens combine over the full + ``padded_seq_len`` extent on the single rank that owns the sequence; + sequences then combine in ascending global sequence index. Both extents are + contract-fixed, so the floating-point grouping is identical at every DP + degree. Only *which rank* computes a sequence changes. + """ + + SEQUENCE_MAJOR_FIXED = "sequence_major_fixed" + + +class LossTransport(str, Enum): + """Collectives move partial sums only; they never reduce numerically. + + ``all_reduce`` is excluded on purpose: NCCL's reduction order depends on + world size and topology, so it would silently regroup the per-sequence + combines and break the cross-DP guarantee ``SummationOrder`` provides. + """ + + ALL_GATHER = "all_gather" + + +class KLEstimator(str, Enum): + """Per-token reference-KL estimator. + + ``k3_unbiased``: ``exp(logp_ref - logp_policy) - (logp_ref - logp_policy) - 1``, + the non-negative low-variance estimator used by the existing ratio/KL op. + + ``k1_log_ratio``: ``logp_policy - logp_ref``, the plain log-ratio. + """ + + K3_UNBIASED = "k3_unbiased" + K1_LOG_RATIO = "k1_log_ratio" + + +class ClipMode(str, Enum): + MIN_OF_UNCLIPPED_AND_CLIPPED = "min_of_unclipped_and_clipped" + + +class AdvantageNormalizer(str, Enum): + """Group-relative reward normalization. + + ``mean_std_population``: ``(r - mean) / std`` with the population (biased) + standard deviation, the original GRPO form. + + ``mean_only``: ``r - mean``, the Dr.GRPO form that drops the std divisor to + avoid its length/difficulty bias. + """ + + MEAN_STD_POPULATION = "mean_std_population" + MEAN_ONLY = "mean_only" + + +class VarianceFormula(str, Enum): + """Only the two-pass form is conformant. + + ``E[x^2] - E[x]^2`` cancels catastrophically once rewards share a large + offset, and its error depends on group size, so it cannot support a bitwise + claim. The two-pass form subtracts the mean before squaring. + """ + + TWO_PASS = "two_pass" + + +class GroupReplication(str, Enum): + """How advantage groups are evaluated when they span DP ranks. + + ``replicated_all_gather``: per-sequence rewards are all-gathered and *every* + rank normalizes *every* group identically, then keeps its own slice. + Rewards are one scalar per sequence, so replicating the whole computation is + cheaper than making a partial-statistic merge bitwise-reproducible. + """ + + REPLICATED_ALL_GATHER = "replicated_all_gather" + + +class LossPlacement(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 LossContractError(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 LossContractError(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 LossContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _non_negative_float(value: Any, field: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise LossContractError(f"{field} must be a real number; got {value!r}") + value = float(value) + if not math.isfinite(value) or value < 0.0: + raise LossContractError(f"{field} must be finite and non-negative; got {value!r}") + return value + + +@dataclass(frozen=True) +class ClipSpec: + """Asymmetric PPO-style ratio clipping bounds. + + Separate low/high epsilons cover the "clip-higher" variants; passing the + same value twice recovers the symmetric ``[1-eps, 1+eps]`` form. + """ + + clip_eps_low: float = 0.2 + clip_eps_high: float = 0.2 + mode: ClipMode = ClipMode.MIN_OF_UNCLIPPED_AND_CLIPPED + + def __post_init__(self) -> None: + object.__setattr__(self, "mode", _enum_value(ClipMode, self.mode, "clip.mode")) + low = _non_negative_float(self.clip_eps_low, "clip_eps_low") + high = _non_negative_float(self.clip_eps_high, "clip_eps_high") + if low >= 1.0: + raise LossContractError( + f"clip_eps_low={low} must be smaller than 1.0; the lower clip bound " + "1 - clip_eps_low must stay positive" + ) + object.__setattr__(self, "clip_eps_low", low) + object.__setattr__(self, "clip_eps_high", high) + + @property + def lower_bound(self) -> float: + return 1.0 - self.clip_eps_low + + @property + def upper_bound(self) -> float: + return 1.0 + self.clip_eps_high + + +@dataclass(frozen=True) +class AdvantageSpec: + """Group-relative advantage normalization semantics.""" + + normalizer: AdvantageNormalizer = AdvantageNormalizer.MEAN_STD_POPULATION + variance: VarianceFormula = VarianceFormula.TWO_PASS + std_eps: float = 1e-6 + replication: GroupReplication = GroupReplication.REPLICATED_ALL_GATHER + + def __post_init__(self) -> None: + object.__setattr__( + self, + "normalizer", + _enum_value(AdvantageNormalizer, self.normalizer, "advantage.normalizer"), + ) + object.__setattr__( + self, "variance", _enum_value(VarianceFormula, self.variance, "advantage.variance") + ) + object.__setattr__( + self, + "replication", + _enum_value(GroupReplication, self.replication, "advantage.replication"), + ) + std_eps = _non_negative_float(self.std_eps, "advantage.std_eps") + if std_eps <= 0.0: + raise LossContractError( + f"advantage.std_eps={std_eps} must be strictly positive; it is the floor " + "that keeps a zero-variance group from dividing by zero" + ) + object.__setattr__(self, "std_eps", std_eps) + + +@dataclass(frozen=True) +class ObjectiveSpec: + """The GRPO objective itself: clipping, reference KL, advantage shaping.""" + + clip: ClipSpec = field(default_factory=ClipSpec) + advantage: AdvantageSpec = field(default_factory=AdvantageSpec) + kl_estimator: KLEstimator = KLEstimator.K3_UNBIASED + beta: float = 0.0 + + def __post_init__(self) -> None: + if not isinstance(self.clip, ClipSpec): + raise LossContractError("objective.clip must be a ClipSpec") + if not isinstance(self.advantage, AdvantageSpec): + raise LossContractError("objective.advantage must be an AdvantageSpec") + object.__setattr__( + self, + "kl_estimator", + _enum_value(KLEstimator, self.kl_estimator, "objective.kl_estimator"), + ) + object.__setattr__(self, "beta", _non_negative_float(self.beta, "objective.beta")) + + @property + def uses_reference_model(self) -> bool: + """Whether reference logits are required at all. + + ``beta == 0`` drops the KL term from the loss, so a backend may skip the + reference forward entirely. The KL is still *reported*, so a caller + that wants the diagnostic must supply reference logits regardless. + """ + + return self.beta > 0.0 + + +@dataclass(frozen=True) +class LossReductionSpec: + """Deterministic token/sequence summation and normalizer semantics. + + Per-token loss terms are accumulated in fp32, combined in the order given by + ``summation_order``, moved between ranks by ``transport`` (never reduced by + it), and divided by the denominator selected by ``token_normalizer``. The + scalar is downcast, if at all, only at ``downcast_at``. + + ``determinism_scope`` reuses the logprob scale. ``cross_tp_bitwise`` here + means the scalar loss and its gradient are bitwise-identical across TP and + DP degrees, given a fixed vocab tile count. + """ + + token_normalizer: TokenNormalizer = TokenNormalizer.GLOBAL_ACTIVE_TOKENS + summation_order: SummationOrder = SummationOrder.SEQUENCE_MAJOR_FIXED + acc_dtype: LogprobDType = LogprobDType.FP32 + transport: LossTransport = LossTransport.ALL_GATHER + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + determinism_scope: DeterminismScope = DeterminismScope.CROSS_TP_BITWISE + fixed_normalizer_constant: int | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "token_normalizer", + _enum_value(TokenNormalizer, self.token_normalizer, "token_normalizer"), + ) + object.__setattr__( + self, + "summation_order", + _enum_value(SummationOrder, self.summation_order, "summation_order"), + ) + object.__setattr__( + self, "acc_dtype", _enum_value(LogprobDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__( + self, "transport", _enum_value(LossTransport, self.transport, "transport") + ) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__( + self, + "determinism_scope", + _enum_value(DeterminismScope, self.determinism_scope, "determinism_scope"), + ) + if self.acc_dtype is not LogprobDType.FP32: + raise LossContractError(f"loss accumulation must be fp32; got {self.acc_dtype.value}") + + needs_constant = self.token_normalizer is TokenNormalizer.FIXED_CONSTANT + if needs_constant: + object.__setattr__( + self, + "fixed_normalizer_constant", + _positive_int(self.fixed_normalizer_constant, "fixed_normalizer_constant"), + ) + elif self.fixed_normalizer_constant is not None: + raise LossContractError( + "fixed_normalizer_constant is only meaningful for " + f"token_normalizer={TokenNormalizer.FIXED_CONSTANT.value}; got " + f"{self.token_normalizer.value}" + ) + + +@dataclass(frozen=True) +class LossOutputSpec: + """Output surface: fp32 scalars replicated across every DP and TP rank.""" + + loss_dtype: LogprobDType = LogprobDType.FP32 + placement: LossPlacement = LossPlacement.REPLICATED + + def __post_init__(self) -> None: + object.__setattr__( + self, "loss_dtype", _enum_value(LogprobDType, self.loss_dtype, "loss_dtype") + ) + object.__setattr__( + self, "placement", _enum_value(LossPlacement, self.placement, "placement") + ) + if self.loss_dtype is not LogprobDType.FP32: + raise LossContractError(f"loss output must be fp32; got {self.loss_dtype.value}") + + +@dataclass(frozen=True) +class LossShardingSpec: + """Which sequences of the global batch this DP rank owns. + + The global batch is ``num_sequences`` sequences of ``padded_seq_len`` token + slots each. ``sequence_shard_bounds`` lists every DP rank's half-open + ``[start, end)`` sequence range, indexed by rank; the full table is required + on every rank and must form a contiguous ``[0, num_sequences)`` partition, + exactly as ``ShardingSpec.vocab_shard_bounds`` does for the vocabulary. + """ + + dp_rank: int + dp_world_size: int + num_sequences: int + padded_seq_len: int + sequence_shard_bounds: tuple[tuple[int, int], ...] + group_boundaries: tuple[int, ...] + cp_rank: int = 0 + cp_world_size: int = 1 + + def __post_init__(self) -> None: + dp_world_size = _positive_int(self.dp_world_size, "dp_world_size") + dp_rank = _non_negative_int(self.dp_rank, "dp_rank") + if dp_rank >= dp_world_size: + raise LossContractError( + f"dp_rank={dp_rank} must be smaller than dp_world_size={dp_world_size}" + ) + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + _non_negative_int(self.cp_rank, "cp_rank") + if cp_world_size != 1: + raise LossContractError( + f"cp_world_size={cp_world_size} is unsupported: context parallelism splits a " + "sequence's tokens across ranks, so the loss reduction would span an axis this " + "contract does not model. Reduce across CP outside this operator." + ) + if self.cp_rank != 0: + raise LossContractError(f"cp_rank must be 0 when cp_world_size=1; got {self.cp_rank}") + + _positive_int(self.num_sequences, "num_sequences") + _positive_int(self.padded_seq_len, "padded_seq_len") + object.__setattr__( + self, + "sequence_shard_bounds", + self._validated_bounds(self.sequence_shard_bounds, dp_world_size, self.num_sequences), + ) + object.__setattr__( + self, + "group_boundaries", + self._validated_groups(self.group_boundaries, self.num_sequences), + ) + + @staticmethod + def _validated_bounds( + raw: Any, dp_world_size: int, num_sequences: int + ) -> tuple[tuple[int, int], ...]: + try: + bounds = tuple((pair[0], pair[1]) for pair in raw) + except (TypeError, IndexError, KeyError) as exc: + raise LossContractError( + "sequence_shard_bounds must be an iterable of (start, end) integer pairs" + ) from exc + if len(bounds) != dp_world_size: + raise LossContractError( + "sequence_shard_bounds must declare exactly one (start, end) pair per DP rank; " + f"got {len(bounds)} pairs for dp_world_size={dp_world_size}" + ) + expected_start = 0 + for rank, (start, end) in enumerate(bounds): + for name, value in ( + (f"sequence_shard_bounds[{rank}][0]", start), + (f"sequence_shard_bounds[{rank}][1]", end), + ): + if isinstance(value, bool) or not isinstance(value, int): + raise LossContractError(f"{name} must be an integer; got {value!r}") + if end <= start: + raise LossContractError( + f"sequence_shard_bounds[{rank}] must satisfy end > start; got [{start}, {end})" + ) + if start != expected_start: + raise LossContractError( + "sequence_shard_bounds must form a contiguous [0, num_sequences) partition " + f"in DP-rank order; rank {rank} starts at {start}, expected {expected_start}" + ) + expected_start = end + if expected_start != num_sequences: + raise LossContractError( + "sequence_shard_bounds must cover num_sequences exactly; covered " + f"{expected_start}, declared {num_sequences}" + ) + return bounds + + @staticmethod + def _validated_groups(raw: Any, num_sequences: int) -> tuple[int, ...]: + try: + offsets = tuple(raw) + except TypeError as exc: + raise LossContractError("group_boundaries must be an iterable of integers") from exc + if len(offsets) < 2: + raise LossContractError( + "group_boundaries must hold num_groups + 1 offsets, so at least 2 entries" + ) + for index, value in enumerate(offsets): + if isinstance(value, bool) or not isinstance(value, int): + raise LossContractError(f"group_boundaries[{index}] must be an integer") + if offsets[0] != 0 or offsets[-1] != num_sequences: + raise LossContractError( + "group_boundaries must start at 0 and end at " + f"num_sequences={num_sequences}; got [{offsets[0]}, ..., {offsets[-1]}]" + ) + for index in range(1, len(offsets)): + if offsets[index] <= offsets[index - 1]: + raise LossContractError( + "group_boundaries must be strictly increasing; " + f"offset {index} is {offsets[index]} after {offsets[index - 1]}" + ) + return offsets + + @property + def num_groups(self) -> int: + return len(self.group_boundaries) - 1 + + @property + def group_sizes(self) -> tuple[int, ...]: + return tuple( + self.group_boundaries[i + 1] - self.group_boundaries[i] for i in range(self.num_groups) + ) + + @property + def local_sequence_start(self) -> int: + return self.sequence_shard_bounds[self.dp_rank][0] + + @property + def local_sequence_end(self) -> int: + return self.sequence_shard_bounds[self.dp_rank][1] + + @property + def local_num_sequences(self) -> int: + start, end = self.sequence_shard_bounds[self.dp_rank] + return end - start + + @property + def local_num_token_slots(self) -> int: + """Token slots this rank holds -- the row count of its logprob call.""" + + return self.local_num_sequences * self.padded_seq_len + + +@dataclass(frozen=True) +class GRPOLossContract: + """Complete semantic request for one deterministic GRPO loss invocation.""" + + logprob: LogprobContract + sharding: LossShardingSpec + objective: ObjectiveSpec = field(default_factory=ObjectiveSpec) + reduction: LossReductionSpec = field(default_factory=LossReductionSpec) + output: LossOutputSpec = field(default_factory=LossOutputSpec) + + def __post_init__(self) -> None: + if not isinstance(self.logprob, LogprobContract): + raise LossContractError("logprob must be a LogprobContract") + if not isinstance(self.sharding, LossShardingSpec): + raise LossContractError("sharding must be a LossShardingSpec") + if not isinstance(self.objective, ObjectiveSpec): + raise LossContractError("objective must be an ObjectiveSpec") + if not isinstance(self.reduction, LossReductionSpec): + raise LossContractError("reduction must be a LossReductionSpec") + if not isinstance(self.output, LossOutputSpec): + raise LossContractError("output must be a LossOutputSpec") + + # The nested logprob contract describes this rank's own rows, so its + # token count must match the cells this rank owns. Catching the + # mismatch here turns a silent shape error deep inside the reduction + # into a contract failure at construction. + expected_tokens = self.sharding.local_num_token_slots + if self.logprob.mask.num_tokens != expected_tokens: + raise LossContractError( + f"logprob.mask.num_tokens={self.logprob.mask.num_tokens} must equal the " + f"{expected_tokens} token slots this rank owns " + f"({self.sharding.local_num_sequences} sequences x " + f"{self.sharding.padded_seq_len} slots per sequence)" + ) + if self.logprob.reduction.determinism_scope is not self.reduction.determinism_scope: + raise LossContractError( + "the loss cannot claim a stronger or weaker determinism scope than the " + f"logprob path it consumes; loss={self.reduction.determinism_scope.value}, " + f"logprob={self.logprob.reduction.determinism_scope.value}" + ) + if self.logprob.sharding.cp_world_size != self.sharding.cp_world_size: + raise LossContractError( + f"cp_world_size disagrees between the logprob contract " + f"({self.logprob.sharding.cp_world_size}) and the loss sharding " + f"({self.sharding.cp_world_size})" + ) + if self.objective.advantage.normalizer is AdvantageNormalizer.MEAN_STD_POPULATION: + # A singleton group has zero population variance, so its advantage + # would collapse to 0 and the sequence would contribute nothing. + # Reject it rather than silently training on a dead group. + small = [index for index, size in enumerate(self.sharding.group_sizes) if size < 2] + if small: + raise LossContractError( + f"advantage normalizer {AdvantageNormalizer.MEAN_STD_POPULATION.value} " + f"needs at least 2 sequences per group; groups {small} are smaller" + ) + + @property + def global_token_slots(self) -> int: + return self.sharding.num_sequences * self.sharding.padded_seq_len + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "dp_rank": self.sharding.dp_rank, + "dp_world_size": self.sharding.dp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "num_sequences": self.sharding.num_sequences, + "padded_seq_len": self.sharding.padded_seq_len, + "sequence_shard_bounds": [list(pair) for pair in self.sharding.sequence_shard_bounds], + "group_boundaries": list(self.sharding.group_boundaries), + "local_sequence_start": self.sharding.local_sequence_start, + "local_sequence_end": self.sharding.local_sequence_end, + } + objective = { + "clip_eps_low": self.objective.clip.clip_eps_low, + "clip_eps_high": self.objective.clip.clip_eps_high, + "clip_mode": self.objective.clip.mode.value, + "kl_estimator": self.objective.kl_estimator.value, + "beta": self.objective.beta, + "advantage_normalizer": self.objective.advantage.normalizer.value, + "advantage_variance": self.objective.advantage.variance.value, + "advantage_std_eps": self.objective.advantage.std_eps, + "advantage_replication": self.objective.advantage.replication.value, + } + reduction = { + "token_normalizer": self.reduction.token_normalizer.value, + "summation_order": self.reduction.summation_order.value, + "acc_dtype": self.reduction.acc_dtype.value, + "transport": self.reduction.transport.value, + "downcast_at": self.reduction.downcast_at.value, + "determinism_scope": self.reduction.determinism_scope.value, + "fixed_normalizer_constant": self.reduction.fixed_normalizer_constant, + "cp_is_merge_axis": False, + "dp_is_merge_axis": True, + } + return { + "semantic_operator": "grpo_loss", + "logprob": self.logprob.to_dict(), + "sharding": sharding, + "objective": objective, + "reduction": reduction, + "output": { + "loss_dtype": self.output.loss_dtype.value, + "placement": self.output.placement.value, + }, + } + + def cross_rank_fingerprint(self) -> str: + """Rank-independent identity for preflight agreement across all ranks. + + Drops every rank-local field so all ``dp_world_size x tp_world_size`` + ranks of one logical invocation agree. That includes the nested + logprob contract's mask: each DP rank holds a different slice of + sequences, so its ``num_tokens`` and mask digest legitimately differ + even though the invocation is the same one. The global token geometry + is still pinned, by ``num_sequences``/``padded_seq_len`` and by the full + ``sequence_shard_bounds`` table, which every rank declares identically. + """ + + payload = self.to_dict() + logprob = payload["logprob"] + logprob.pop("mask", None) + logprob["sharding"] = { + key: value + for key, value in logprob["sharding"].items() + if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} + } + payload["sharding"] = { + key: value + for key, value in payload["sharding"].items() + if key not in {"dp_rank", "cp_rank", "local_sequence_start", "local_sequence_end"} + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class LossBackendCapability: + """Capabilities a concrete loss backend declares to contract-aware dispatch.""" + + backend_id: str + token_normalizers: frozenset[TokenNormalizer] + kl_estimators: frozenset[KLEstimator] + advantage_normalizers: frozenset[AdvantageNormalizer] + determinism_scopes: frozenset[DeterminismScope] + dp_world_sizes: tuple[int, ...] | None = None + supports_variable_group_sizes: bool = False + supports_asymmetric_clip: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise LossContractError("backend_id must be a non-empty string") + if self.backend_id.strip().lower() in RESERVED_DISPATCH_POLICIES: + raise LossContractError( + f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" + ) + object.__setattr__(self, "backend_id", self.backend_id.strip()) + for name, enum_type in ( + ("token_normalizers", TokenNormalizer), + ("kl_estimators", KLEstimator), + ("advantage_normalizers", AdvantageNormalizer), + ("determinism_scopes", DeterminismScope), + ): + try: + values = frozenset( + _enum_value(enum_type, value, name) for value in getattr(self, name) + ) + except TypeError as exc: + raise LossContractError(f"{name} must be an iterable of enum values") from exc + if not values: + raise LossContractError(f"{name} must not be empty") + object.__setattr__(self, name, values) + object.__setattr__( + self, + "dp_world_sizes", + self._validated_world_sizes(self.dp_world_sizes, "dp_world_sizes"), + ) + for flag_name in ("supports_variable_group_sizes", "supports_asymmetric_clip"): + if not isinstance(getattr(self, flag_name), bool): + raise LossContractError(f"{flag_name} must be a bool") + if self.implementation_kind not in IMPLEMENTATION_KINDS: + raise LossContractError( + f"implementation_kind must be one of: {', '.join(sorted(IMPLEMENTATION_KINDS))}" + ) + + @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 LossContractError(f"{field} must be an iterable of integers") from exc + if not sizes: + raise LossContractError(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 LossContractError(f"{field} must contain positive values; got {value!r}") + if len(set(sizes)) != len(sizes): + raise LossContractError(f"{field} must not contain duplicates") + return sizes + + def incompatibilities(self, contract: GRPOLossContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + reduction = contract.reduction + objective = contract.objective + sharding = contract.sharding + if reduction.token_normalizer not in self.token_normalizers: + reasons.append(f"token_normalizer={reduction.token_normalizer.value} is unsupported") + if objective.kl_estimator not in self.kl_estimators: + reasons.append(f"kl_estimator={objective.kl_estimator.value} is unsupported") + if objective.advantage.normalizer not in self.advantage_normalizers: + reasons.append( + f"advantage normalizer={objective.advantage.normalizer.value} is unsupported" + ) + if reduction.determinism_scope not in self.determinism_scopes: + reasons.append(f"determinism_scope={reduction.determinism_scope.value} is unsupported") + if self.dp_world_sizes is not None and sharding.dp_world_size not in self.dp_world_sizes: + reasons.append(f"DP={sharding.dp_world_size} is unsupported") + if len(set(sharding.group_sizes)) > 1 and not self.supports_variable_group_sizes: + reasons.append("variable advantage group sizes are unsupported") + if ( + objective.clip.clip_eps_low != objective.clip.clip_eps_high + and not self.supports_asymmetric_clip + ): + reasons.append("asymmetric ratio clipping is unsupported") + return tuple(reasons) + + def supports(self, contract: GRPOLossContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "token_normalizers": sorted(item.value for item in self.token_normalizers), + "kl_estimators": sorted(item.value for item in self.kl_estimators), + "advantage_normalizers": sorted(item.value for item in self.advantage_normalizers), + "determinism_scopes": sorted(item.value for item in self.determinism_scopes), + "dp_world_sizes": list(self.dp_world_sizes) if self.dp_world_sizes else None, + "supports_variable_group_sizes": self.supports_variable_group_sizes, + "supports_asymmetric_clip": self.supports_asymmetric_clip, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class LossDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: LossBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "AdvantageNormalizer", + "AdvantageSpec", + "ClipMode", + "ClipSpec", + "GRPOLossContract", + "GroupReplication", + "KLEstimator", + "LossBackendCapability", + "LossContractError", + "LossDispatchResult", + "LossOutputSpec", + "LossPlacement", + "LossReductionSpec", + "LossShardingSpec", + "LossTransport", + "ObjectiveSpec", + "SummationOrder", + "TokenNormalizer", + "VarianceFormula", +] diff --git a/rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py b/rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py new file mode 100644 index 00000000..bbec450d --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic DP-aware GRPO loss.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Iterator + +import torch + +from rl_engine.kernels.loss_contract import ( + AdvantageNormalizer, + GRPOLossContract, + KLEstimator, + LossContractError, + TokenNormalizer, +) +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + DEFAULT_NUM_VOCAB_TILES, + VocabParallelLogprobOp, +) + +BACKEND_ID = "pytorch-distributed-grpo-loss-ws2" + +# Channel layout of the packed per-sequence tensor moved by the single fp32 +# all-gather; counts travel separately as integers so they stay exact. +_CH_POLICY = 0 +_CH_KL = 1 + + +@dataclass(frozen=True) +class GRPOLossResult: + """Scalar loss terms, per-sequence diagnostics, and bound provenance. + + Unpacks as ``loss, policy_loss, kl`` so it can stand in for the tuple the + single-GPU GRPO ops return. + + The ``per_sequence_*`` vectors are the reduction's last mesh-independent + intermediate, detached and replicated on every rank. They are the right + surface for comparing two configurations: the scalar loss averages + ``num_sequences`` totals into 24 mantissa bits and routinely rounds a real + reordering away, whereas the per-sequence vector preserves it. They also + give a drift report somewhere to point when one sequence is responsible. + """ + + loss: torch.Tensor + policy_loss: torch.Tensor + kl: torch.Tensor + advantages: torch.Tensor + per_sequence_policy: torch.Tensor + per_sequence_kl: torch.Tensor + per_sequence_active_tokens: torch.Tensor + provenance: dict[str, Any] = field(default_factory=dict) + + def __iter__(self) -> Iterator[torch.Tensor]: + yield from (self.loss, self.policy_loss, self.kl) + + +def _require_distributed_initialized(): + import torch.distributed as dist + + if not dist.is_available(): + raise LossContractError("distributed GRPO loss requires torch.distributed.") + if not dist.is_initialized(): + raise LossContractError( + "distributed GRPO loss requires an initialized process group when " + "the contract declares dp_world_size > 1." + ) + return dist + + +def _validate_invocation( + policy_local_logits: torch.Tensor, + ref_local_logits: torch.Tensor | None, + action_ids: torch.Tensor, + old_logps: torch.Tensor, + rewards: torch.Tensor, + contract: GRPOLossContract, + dp_group: Any, +) -> None: + sharding = contract.sharding + num_rows = sharding.local_num_token_slots + if policy_local_logits.dim() != 2: + raise LossContractError( + "policy_local_logits must be 2-D [num_tokens, local_vocab]; got " + f"{policy_local_logits.dim()}-D" + ) + if policy_local_logits.shape[0] != num_rows: + raise LossContractError( + f"policy_local_logits has {policy_local_logits.shape[0]} rows but this rank owns " + f"{num_rows} token slots" + ) + if ref_local_logits is not None and ref_local_logits.shape != policy_local_logits.shape: + raise LossContractError( + f"ref_local_logits shape {tuple(ref_local_logits.shape)} must match " + f"policy_local_logits shape {tuple(policy_local_logits.shape)}" + ) + if ref_local_logits is None and contract.objective.uses_reference_model: + raise LossContractError( + f"objective.beta={contract.objective.beta} puts the reference KL in the loss, " + "so ref_local_logits is required" + ) + for name, tensor in (("action_ids", action_ids), ("old_logps", old_logps)): + if tensor.dim() != 1 or tensor.shape[0] != num_rows: + raise LossContractError( + f"{name} must be 1-D with one entry per owned token slot; got shape " + f"{tuple(tensor.shape)} for {num_rows} slots" + ) + if rewards.dim() != 1 or rewards.shape[0] != sharding.local_num_sequences: + raise LossContractError( + "rewards must be 1-D with one entry per sequence this rank owns; got shape " + f"{tuple(rewards.shape)} for {sharding.local_num_sequences} sequences" + ) + + if sharding.dp_world_size > 1: + dist = _require_distributed_initialized() + group_world = dist.get_world_size(group=dp_group) + group_rank = dist.get_rank(group=dp_group) + if group_world != sharding.dp_world_size: + raise LossContractError( + f"dp_group world size {group_world} does not match the contract " + f"dp_world_size={sharding.dp_world_size}; pass the DP subgroup, " + "not the global group" + ) + if group_rank != sharding.dp_rank: + raise LossContractError( + f"dp_group rank {group_rank} does not match the contract dp_rank={sharding.dp_rank}" + ) + + +def _preflight_cross_rank_agreement( + contract: GRPOLossContract, dp_group: Any, tp_group: Any, num_vocab_tiles: int +) -> None: + """All-gather (fingerprint, backend id, vocab tile count) and abort on mismatch. + + Checked over the DP group *and* the TP group. Neither alone is sufficient: + the loss is replicated across TP, so two TP siblings that disagree on, say, + ``beta`` would compute different losses and produce inconsistent gradients + for one sharded model -- and the logprob path's own preflight cannot catch + that, because ``beta`` is not part of the logprob contract. Agreement + within both groups implies agreement across the whole DP x TP grid by + transitivity. + + Runs before any other collective, including the logprob path's own TP + preflight, so a rank that joined the wrong logical invocation fails here + rather than deadlocking a later reduction. + """ + + checks = [ + (axis, group, world_size) + for axis, group, world_size in ( + ("DP", dp_group, contract.sharding.dp_world_size), + ("TP", tp_group, contract.logprob.sharding.tp_world_size), + ) + if world_size > 1 + ] + if not checks: + return + + dist = _require_distributed_initialized() + payload = (contract.cross_rank_fingerprint(), BACKEND_ID, int(num_vocab_tiles)) + for axis, group, _ in checks: + gathered: list[Any] = [None] * dist.get_world_size(group=group) + dist.all_gather_object(gathered, payload, group=group) + mismatched = [(rank, other) for rank, other in enumerate(gathered) if other != payload] + if mismatched: + rank, other = mismatched[0] + raise LossContractError( + f"cross-rank preflight failed on the {axis} axis: this rank has {payload} " + f"but {axis} rank {rank} has {other}; every rank must agree on the contract " + "fingerprint, backend id and num_vocab_tiles before any collective" + ) + + +def _gather_global_rewards( + rewards: torch.Tensor, contract: GRPOLossContract, dp_group: Any +) -> torch.Tensor: + """Assemble the global ``[num_sequences]`` reward vector on every rank. + + ``sequence_shard_bounds`` is a contiguous partition in DP-rank order, so + concatenating the gathered slices in rank order reproduces the global vector + exactly -- no ownership arbitration is needed. + """ + + sharding = contract.sharding + local = rewards.float() + if sharding.dp_world_size == 1: + return local.contiguous() + + dist = _require_distributed_initialized() + max_local = max(end - start for start, end in sharding.sequence_shard_bounds) + padded = local.new_zeros(max_local) + padded[: local.shape[0]] = local + gathered = [torch.empty_like(padded) for _ in range(sharding.dp_world_size)] + dist.all_gather(gathered, padded.contiguous(), group=dp_group) + return torch.cat( + [ + gathered[rank][: end - start] + for rank, (start, end) in enumerate(sharding.sequence_shard_bounds) + ], + dim=0, + ) + + +def _group_advantages(global_rewards: torch.Tensor, contract: GRPOLossContract) -> torch.Tensor: + """Group-relative advantages over the global reward vector. + + Evaluated identically on every rank from an identically shaped input, so no + merge is involved and the result is bitwise-equal mesh-wide. The variance + is two-pass: centring before squaring keeps the result meaningful when the + rewards share a large offset, which ``E[x^2] - E[x]^2`` does not. + """ + + advantage = contract.objective.advantage + boundaries = contract.sharding.group_boundaries + parts: list[torch.Tensor] = [] + for index in range(len(boundaries) - 1): + start, end = boundaries[index], boundaries[index + 1] + group = global_rewards[start:end] + count = float(end - start) + centered = group - group.sum() / count + if advantage.normalizer is AdvantageNormalizer.MEAN_ONLY: + parts.append(centered) + continue + variance = (centered * centered).sum() / count + parts.append(centered / variance.clamp_min(advantage.std_eps**2).sqrt()) + return torch.cat(parts, dim=0) + + +def _sequence_totals(values: torch.Tensor, contract: GRPOLossContract) -> torch.Tensor: + """Reduce this rank's per-token values to one total per owned sequence. + + The reduced extent is ``padded_seq_len``, which the contract fixes, so this + sum is byte-for-byte the same work at every DP degree. + """ + + sharding = contract.sharding + view = values.reshape(sharding.local_num_sequences, sharding.padded_seq_len) + return view.sum(dim=1) + + +def _assemble_global_vector( + local_totals: torch.Tensor, contract: GRPOLossContract, dp_group: Any +) -> torch.Tensor: + """Place every rank's per-sequence totals into the fixed global vector. + + Returns ``[num_sequences, ...]``. The length is a property of the contract, + never of the DP degree, and filling it is pure placement -- no arithmetic + touches the gathered values -- so the downstream reduction sees identical + inputs at every degree. + """ + + sharding = contract.sharding + if sharding.dp_world_size == 1: + return local_totals + + dist = _require_distributed_initialized() + trailing = local_totals.shape[1:] + max_local = max(end - start for start, end in sharding.sequence_shard_bounds) + padded = local_totals.new_zeros((max_local, *trailing)) + padded[: local_totals.shape[0]] = local_totals.detach() + gathered = [torch.empty_like(padded) for _ in range(sharding.dp_world_size)] + dist.all_gather(gathered, padded.contiguous(), group=dp_group) + + # This rank's own slice comes from the live tensor: all_gather severs the + # graph, and the other ranks' slices are constants here anyway. + pieces = [] + for rank, (start, end) in enumerate(sharding.sequence_shard_bounds): + pieces.append(local_totals if rank == sharding.dp_rank else gathered[rank][: end - start]) + return torch.cat(pieces, dim=0) + + +def _normalized( + per_sequence_totals: torch.Tensor, + per_sequence_counts: torch.Tensor, + contract: GRPOLossContract, +) -> torch.Tensor: + """Apply the declared token normalizer to fixed-order sequence totals.""" + + reduction = contract.reduction + normalizer = reduction.token_normalizer + if normalizer is TokenNormalizer.FIXED_CONSTANT: + return per_sequence_totals.sum() / float(reduction.fixed_normalizer_constant) + if normalizer is TokenNormalizer.GLOBAL_ACTIVE_TOKENS: + return per_sequence_totals.sum() / per_sequence_counts.sum().to(per_sequence_totals.dtype) + + # PER_SEQUENCE_THEN_MEAN: sequences with no active token contribute nothing + # and are excluded from the outer denominator rather than counted as zero. + live = per_sequence_counts > 0 + denominators = per_sequence_counts.to(per_sequence_totals.dtype).clamp_min(1.0) + per_sequence_means = torch.where( + live, per_sequence_totals / denominators, torch.zeros_like(per_sequence_totals) + ) + return per_sequence_means.sum() / live.sum().to(per_sequence_totals.dtype) + + +class DistributedGRPOLossOp: + """Deterministic GRPO loss over TP-sharded logits with a DP-invariant reduction. + + The WS2 reference (issue #241 PR5). ``policy_local_logits`` is this rank's + ``[n, local_vocab]`` vocabulary shard, not a dense ``[n, vocab]`` tensor: + tensor parallelism is delegated to the vocab-parallel logprob path, and this + operator owns the sum over tokens and sequences. + """ + + op_class = "grpo_loss" + is_batch_invariant = True + + def __init__(self) -> None: + self._logprob = VocabParallelLogprobOp() + + def __call__(self, *args: Any, **kwargs: Any) -> GRPOLossResult: + return self.apply(*args, **kwargs) + + def apply( + self, + policy_local_logits: torch.Tensor, + action_ids: torch.Tensor, + old_logps: torch.Tensor, + rewards: torch.Tensor, + *, + contract: GRPOLossContract, + ref_local_logits: torch.Tensor | None = None, + tp_group: Any = None, + dp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + ) -> GRPOLossResult: + if not isinstance(contract, GRPOLossContract): + raise LossContractError("contract must be a GRPOLossContract") + _validate_invocation( + policy_local_logits, + ref_local_logits, + action_ids, + old_logps, + rewards, + contract, + dp_group, + ) + sharding = contract.sharding + objective = contract.objective + if validate: + _preflight_cross_rank_agreement(contract, dp_group, tp_group, num_vocab_tiles) + + logp_policy, _ = self._logprob.apply( + policy_local_logits, + action_ids, + contract=contract.logprob, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + ) + active = torch.tensor( + contract.logprob.mask.active_mask, + dtype=torch.bool, + device=policy_local_logits.device, + ) + + delta = (logp_policy - old_logps.float()).masked_fill(~active, 0.0) + ratio = delta.exp() + + if ref_local_logits is None: + kl_terms = torch.zeros_like(logp_policy) + else: + with torch.no_grad(): + logp_ref, _ = self._logprob.apply( + ref_local_logits, + action_ids, + contract=contract.logprob, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=False, + ) + diff = (logp_ref - logp_policy).masked_fill(~active, 0.0) + if objective.kl_estimator is KLEstimator.K3_UNBIASED: + kl_terms = diff.exp() - diff - 1.0 + else: + kl_terms = -diff + kl_terms = kl_terms.masked_fill(~active, 0.0) + + global_rewards = _gather_global_rewards(rewards, contract, dp_group) + advantages = _group_advantages(global_rewards, contract) + adv_tokens = ( + advantages[sharding.local_sequence_start : sharding.local_sequence_end] + .reshape(-1, 1) + .expand(sharding.local_num_sequences, sharding.padded_seq_len) + .reshape(-1) + ) + + clip = objective.clip + unclipped = ratio * adv_tokens + clipped = ratio.clamp(clip.lower_bound, clip.upper_bound) * adv_tokens + policy_terms = (-torch.minimum(unclipped, clipped)).masked_fill(~active, 0.0) + + packed = torch.stack( + (_sequence_totals(policy_terms, contract), _sequence_totals(kl_terms, contract)), + dim=1, + ) + # Fixed-extent reductions: [local_seqs, padded_seq_len] -> [local_seqs], + # gathered into [num_sequences] -> scalar, at every DP degree. + totals = _assemble_global_vector(packed, contract, dp_group) + per_sequence_policy = totals[:, _CH_POLICY] + per_sequence_kl = totals[:, _CH_KL] + per_sequence_counts = _assemble_global_vector( + _sequence_totals(active.to(torch.long), contract), contract, dp_group + ) + + total_active = int(per_sequence_counts.sum().item()) + if validate and total_active == 0: + raise LossContractError( + "the global batch holds no active tokens; the loss normalizer would divide by zero" + ) + + policy_loss = _normalized(per_sequence_policy, per_sequence_counts, contract) + kl = _normalized(per_sequence_kl, per_sequence_counts, contract) + loss = policy_loss + objective.beta * kl + + provenance = { + "backend_id": BACKEND_ID, + "implementation_kind": "reference", + "num_vocab_tiles": int(num_vocab_tiles), + "padded_seq_len": int(sharding.padded_seq_len), + "num_sequences": int(sharding.num_sequences), + "global_active_tokens": total_active, + "reference_model_used": ref_local_logits is not None, + "requested_contract": contract.to_dict(), + "cross_rank_fingerprint": contract.cross_rank_fingerprint(), + } + return GRPOLossResult( + loss=loss, + policy_loss=policy_loss, + kl=kl, + advantages=advantages, + per_sequence_policy=per_sequence_policy.detach(), + per_sequence_kl=per_sequence_kl.detach(), + per_sequence_active_tokens=per_sequence_counts.detach(), + provenance=provenance, + ) + + +__all__ = [ + "BACKEND_ID", + "GRPOLossResult", + "DistributedGRPOLossOp", +] 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/registry.py b/rl_engine/kernels/registry.py index fb2feb6f..8391dbbb 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -8,6 +8,26 @@ import torch +from rl_engine.kernels.logprob_contract import ( + IMPLEMENTATION_KINDS, + DeterminismScope, + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDispatchResult, + LogprobDType, + LogprobRole, + MaskMode, +) +from rl_engine.kernels.loss_contract import ( + AdvantageNormalizer, + GRPOLossContract, + KLEstimator, + LossBackendCapability, + LossContractError, + LossDispatchResult, + TokenNormalizer, +) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -74,6 +94,14 @@ 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" + ) + # Deterministic GRPO loss on the TP-aware logprob path (WS2 #241 PR5) + PYTORCH_DISTRIBUTED_GRPO_LOSS = ( + "rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss.DistributedGRPOLossOp" + ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -164,6 +192,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": [ @@ -280,6 +349,75 @@ 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, + ) + + # WS2 loss dispatch starts empty: the existing single-GPU GRPO ops + # declare no LossBackendCapability, so they are unreachable from + # get_loss_op and cannot be picked up as a silent fallback for a + # contract that asks for deterministic mesh-wide reduction. + self._loss_candidates: Dict[str, list] = {platform: [] for platform in self._priority_map} + self._loss_capabilities: Dict[str, Dict[OpBackend, LossBackendCapability]] = { + platform: {} for platform in self._priority_map + } + ws2_grpo_loss_capability = LossBackendCapability( + backend_id="pytorch-distributed-grpo-loss-ws2", + token_normalizers=frozenset(TokenNormalizer), + kl_estimators=frozenset(KLEstimator), + advantage_normalizers=frozenset(AdvantageNormalizer), + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + dp_world_sizes=None, + supports_variable_group_sizes=True, + supports_asymmetric_clip=True, + implementation_kind="reference", + ) + for ws2_platform in self._priority_map: + self.register_loss_backend( + OpBackend.PYTORCH_DISTRIBUTED_GRPO_LOSS, + ws2_grpo_loss_capability, + platform=ws2_platform, + ) + 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"}: @@ -369,25 +507,312 @@ 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 + + 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. + """ + + 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: + 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``. + """ + + 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" + ) + + 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 - if backend.name in self._failed_backends: + 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 - 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) + 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}" + ) + + def register_loss_backend( + self, + backend: OpBackend, + capability: LossBackendCapability, + *, + platform: Optional[str] = None, + prepend: bool = False, + ) -> None: + """Register (or replace) a backend for WS2 contract-aware loss dispatch. + + The seam that makes a loss backend selectable by ``get_loss_op``. It is + deliberately separate from the logprob registry: a backend may serve one + contract and not the other, and the deterministic GRPO loss consumes a + logprob backend rather than being one. + """ + + if not isinstance(backend, OpBackend): + raise LossContractError("backend must be an OpBackend") + if not isinstance(capability, LossBackendCapability): + raise LossContractError("capability must be a LossBackendCapability") + resolved_platform = platform if platform is not None else self._platform() + if resolved_platform not in self._priority_map: + raise LossContractError( + f"unsupported platform {resolved_platform!r}; expected one of " + f"{sorted(self._priority_map)}" + ) + candidates = self._loss_candidates.setdefault(resolved_platform, []) + self._loss_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_loss_op( + self, + contract: GRPOLossContract, + *, + requested_backend: str = "auto", + ) -> LossDispatchResult: + """Resolve only a backend that explicitly supports the WS2 loss contract. + + Mirrors ``get_logprob_op``: strictness comes from the contract's + capability checks rather than from the policy string, and a contract no + registered backend can serve fails loudly instead of falling back to an + op with different reduction semantics. + """ - raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + if not isinstance(contract, GRPOLossContract): + raise LossContractError("contract must be a GRPOLossContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise LossContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip() + if requested_backend.lower() == "deterministic": + raise LossContractError( + 'requested_backend="deterministic" is not a dispatch policy; request ' + "determinism through LossReductionSpec.determinism_scope and match it " + "against backend determinism_scopes instead" + ) + + platform = self._platform() + candidates = self._loss_candidates.get(platform, []) + rejected: list[str] = [] + capability_rejections = 0 + + platform_capabilities = self._loss_capabilities.get(platform, {}) + for backend in candidates: + capability = platform_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no LossBackendCapability declared") + capability_rejections += 1 + continue + policy_mismatch = self._loss_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + 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 LossDispatchResult(op=op, capability=capability, provenance=provenance) + + details = " | ".join(rejected) if rejected else "no candidates registered" + raise RuntimeError( + "No loss backend supports the requested WS2 GRPO contract on " + f"{platform}: normalizer={contract.reduction.token_normalizer.value}, " + f"kl={contract.objective.kl_estimator.value}, " + f"TP={contract.logprob.sharding.tp_world_size}, " + f"DP={contract.sharding.dp_world_size}, CP={contract.sharding.cp_world_size}. " + f"Rejections: {details}" + ) + + @staticmethod + def _loss_policy_mismatch( + requested_backend: str, + capability: LossBackendCapability, + ) -> 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}" + ) + + @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/tests/test_distributed_grpo_loss.py b/tests/test_distributed_grpo_loss.py new file mode 100644 index 00000000..b0a2321e --- /dev/null +++ b/tests/test_distributed_grpo_loss.py @@ -0,0 +1,866 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic GRPO loss on the TP-aware logprob path. + +The headline claim under test is that the scalar loss and its gradient are +bitwise-identical across every TP x DP degree, given a fixed vocab tile count. +``TestMeshBitwise`` is the file's centre of gravity: it runs the reachable +degrees on real NCCL ranks and compares raw bit patterns against the single-rank +baseline. + +The remaining classes support that claim rather than duplicate it: the single- +rank tests pin the objective's algebraic identities (ratio exactly 1, KL exactly +0), the negative controls prove the bitwise comparisons are not vacuous, and the +guard tests check that a rank which disagrees about the invocation aborts +instead of corrupting the merge. + +Comparisons are made on ``per_sequence_policy``/``per_sequence_kl`` rather than +on the scalar loss alone. That is not a stylistic choice: measured over this +file's own inputs, regrouping the token sum moves the per-sequence vector in 12 +of 12 seeds but the scalar loss in only 5 of 12, because averaging +``NUM_SEQUENCES`` totals into one fp32 number rounds most reorderings away. +Asserting on the scalar alone would let a wrong reduction pass most of the time. + +Context parallelism is out of scope (see ``LossShardingSpec``); the contract +rejects ``cp_world_size > 1`` and ``TestGuards`` covers that. + +Multi-rank tests need one GPU per rank and skip otherwise. They stay small on +purpose -- a 1000-token vocabulary and 8 sequences of 32 slots -- so they can +share a node with a running training job; each worker additionally caps itself +with ``set_per_process_memory_fraction`` so a regression here cannot starve a +co-tenant. +""" + +from __future__ import annotations + +import math +import os +import queue +import tempfile +import traceback +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.loss_contract import ( + AdvantageNormalizer, + AdvantageSpec, + ClipSpec, + GRPOLossContract, + KLEstimator, + LossContractError, + LossReductionSpec, + LossShardingSpec, + ObjectiveSpec, + TokenNormalizer, +) +from rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss import ( + BACKEND_ID, + DistributedGRPOLossOp, +) + +# Global batch geometry, shared by every configuration so the comparisons are +# between meshes rather than between problems. The degrees below are chosen so +# that DP and CP each reach 4 while every shard stays tile-aligned. +NUM_SEQUENCES = 8 +PADDED_SEQ_LEN = 32 +NUM_TOKEN_SLOTS = NUM_SEQUENCES * PADDED_SEQ_LEN +# Deliberately straddles the DP=2 split at 4 and the DP=4 splits at 2/4/6, so the +# replicated-advantage path is exercised by groups that no single rank owns. +GROUP_BOUNDARIES = (0, 3, NUM_SEQUENCES) +REAL_VOCAB = 1000 +PADDED_VOCAB = 1024 +NUM_VOCAB_TILES = 32 +BETA = 0.04 +SEED = 20260811 + +_SPAWN_TIMEOUT_S = 600 +# Fraction of each card the workers may allocate. The test tensors need a few +# megabytes; the cap exists so a bug cannot balloon into a co-tenant job. +_MEMORY_FRACTION = 0.02 + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + """Raw bit pattern, so -0.0 vs 0.0 and NaN vs NaN compare honestly.""" + + view_dtype = { + torch.float32: torch.int32, + torch.bfloat16: torch.int16, + torch.float16: torch.int16, + }[tensor.dtype] + return tensor.contiguous().view(view_dtype) + + +def _scalar_bits(tensor: torch.Tensor) -> int: + return int(_bits(tensor.detach().float().cpu()).item()) + + +def _vector_bits(tensor: torch.Tensor) -> list[int]: + return _bits(tensor.detach().float().cpu()).tolist() + + +def _reduction_fingerprint(result) -> tuple: + """Everything a change of reduction order can move, at full resolution. + + The per-sequence vectors come first because they are the sensitive part; + the scalars are carried along so a normalizer bug is caught too. + """ + + return ( + _vector_bits(result.per_sequence_policy), + _vector_bits(result.per_sequence_kl), + _scalar_bits(result.loss), + _scalar_bits(result.policy_loss), + _scalar_bits(result.kl), + ) + + +def _cuda_device_count() -> int: + try: + return torch.cuda.device_count() + except Exception: # pragma: no cover - driver-level failures + return 0 + + +def _requires_gpus(count: int): + return pytest.mark.skipif( + _cuda_device_count() < count, + reason=f"needs {count} CUDA devices for a real {count}-rank mesh", + ) + + +# --------------------------------------------------------------------------- # +# Global problem definition +# --------------------------------------------------------------------------- # +def _global_active_mask() -> tuple[bool, ...]: + """Right-padded sequences of strictly decreasing length. + + The lengths are staggered so that no two DP shards hold the same number of + active tokens; equal counts would let a rank-local normalizer accidentally + agree with the global one and pass a test it should fail. + """ + + mask: list[bool] = [] + for seq in range(NUM_SEQUENCES): + real_len = PADDED_SEQ_LEN - seq * 3 + mask.extend(slot < real_len for slot in range(PADDED_SEQ_LEN)) + return tuple(mask) + + +GLOBAL_ACTIVE_MASK = _global_active_mask() + + +def _global_inputs(seed: int = SEED) -> dict[str, torch.Tensor]: + """Deterministic global tensors every rank slices its own view out of. + + ``old_logps`` is centred on ``-log(REAL_VOCAB)``, the scale of a selected + logprob under near-uniform logits, so the importance ratios land around 1 + with real spread. Leaving it centred on 0 would make every ratio ~1e-3 and + every per-token loss term nearly identical, and a sum of near-identical + values is almost invariant to how it is grouped -- which would quietly + drain the power out of every bitwise comparison in this file. + """ + + gen = torch.Generator().manual_seed(seed) + return { + "policy": torch.randn(NUM_TOKEN_SLOTS, PADDED_VOCAB, generator=gen), + "ref": torch.randn(NUM_TOKEN_SLOTS, PADDED_VOCAB, generator=gen), + "action_ids": torch.randint(0, REAL_VOCAB, (NUM_TOKEN_SLOTS,), generator=gen), + "old_logps": torch.randn(NUM_TOKEN_SLOTS, generator=gen) * 0.5 - math.log(REAL_VOCAB), + "rewards": torch.randn(NUM_SEQUENCES, generator=gen), + } + + +def _dp_bounds(dp: int) -> tuple[tuple[int, int], ...]: + """Contiguous sequence partition in DP-rank order.""" + + seqs = NUM_SEQUENCES // dp + return tuple((d * seqs, (d + 1) * seqs) for d in range(dp)) + + +def _owned_rows(bounds: tuple[int, int]) -> list[int]: + """Global token-slot indices this shard owns, in canonical local row order.""" + + start, end = bounds + return [ + seq * PADDED_SEQ_LEN + slot for seq in range(start, end) for slot in range(PADDED_SEQ_LEN) + ] + + +def _build_contract( + *, + tp_rank: int, + tp: int, + dp_rank: int, + dp: int, + objective: ObjectiveSpec | None = None, + reduction: LossReductionSpec | None = None, +) -> GRPOLossContract: + bounds = _dp_bounds(dp) + rows = _owned_rows(bounds[dp_rank]) + shard = PADDED_VOCAB // tp + logprob = LogprobContract( + role=LogprobRole.TRAIN, + dtype=LogprobDType.FP32, + mask=MaskSpec( + num_tokens=len(rows), + active_mask=tuple(GLOBAL_ACTIVE_MASK[row] for row in rows), + ), + sharding=ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp, + vocab_shard_bounds=tuple((r * shard, (r + 1) * shard) for r in range(tp)), + real_vocab_size=REAL_VOCAB, + padded_vocab_size=PADDED_VOCAB, + ), + reduction=ReductionSpec(), + ) + sharding = LossShardingSpec( + dp_rank=dp_rank, + dp_world_size=dp, + num_sequences=NUM_SEQUENCES, + padded_seq_len=PADDED_SEQ_LEN, + sequence_shard_bounds=bounds, + group_boundaries=GROUP_BOUNDARIES, + ) + return GRPOLossContract( + logprob=logprob, + sharding=sharding, + objective=objective if objective is not None else ObjectiveSpec(beta=BETA), + reduction=reduction if reduction is not None else LossReductionSpec(), + ) + + +def _rank_inputs( + globals_: dict[str, torch.Tensor], + *, + tp_rank: int, + tp: int, + bounds: tuple[int, int], + device: torch.device, +) -> dict[str, torch.Tensor]: + rows = _owned_rows(bounds) + shard = PADDED_VOCAB // tp + cols = slice(tp_rank * shard, (tp_rank + 1) * shard) + start, end = bounds + return { + "policy": globals_["policy"][rows, cols].clone().to(device).requires_grad_(True), + "ref": globals_["ref"][rows, cols].clone().to(device), + "action_ids": globals_["action_ids"][rows].clone().to(device), + "old_logps": globals_["old_logps"][rows].clone().to(device), + "rewards": globals_["rewards"][start:end].clone().to(device), + } + + +def _single_rank_setup( + *, + objective: ObjectiveSpec | None = None, + reduction: LossReductionSpec | None = None, + device: str = "cpu", + globals_: dict[str, torch.Tensor] | None = None, +) -> tuple[GRPOLossContract, dict[str, torch.Tensor]]: + """Contract and tensors for one rank owning the whole batch.""" + + contract = _build_contract( + tp_rank=0, tp=1, dp_rank=0, dp=1, objective=objective, reduction=reduction + ) + tensors = _rank_inputs( + globals_ if globals_ is not None else _global_inputs(), + tp_rank=0, + tp=1, + bounds=(0, NUM_SEQUENCES), + device=torch.device(device), + ) + return contract, tensors + + +def _run_single_rank( + *, + num_vocab_tiles: int = NUM_VOCAB_TILES, + with_reference: bool = True, + **setup, +): + """Baseline invocation: one rank, no collectives.""" + + contract, tensors = _single_rank_setup(**setup) + result = DistributedGRPOLossOp().apply( + tensors["policy"], + tensors["action_ids"], + tensors["old_logps"], + tensors["rewards"], + contract=contract, + ref_local_logits=tensors["ref"] if with_reference else None, + num_vocab_tiles=num_vocab_tiles, + ) + return result, tensors, contract + + +# --------------------------------------------------------------------------- # +# Multi-rank harness +# --------------------------------------------------------------------------- # +def _mesh_worker(rank, world_size, init_method, result_queue, tp, dp, scenario): + """One NCCL rank of a TP x DP mesh. + + Only plain Python values go back on the queue. Tensors sent over a + multiprocessing queue travel by shared memory and are lost when the sender + exits before the parent maps them, which surfaces as an empty queue rather + than an error. + """ + + payload = {"rank": rank} + try: + import torch.distributed as dist + + device_index = rank % torch.cuda.device_count() + torch.cuda.set_device(device_index) + torch.cuda.set_per_process_memory_fraction(_MEMORY_FRACTION, device_index) + device = torch.device("cuda", device_index) + dist.init_process_group( + backend="nccl", init_method=init_method, world_size=world_size, rank=rank + ) + + # Rank layout puts TP fastest: rank = dp_rank * tp + tp_rank. + tp_rank = rank % tp + dp_rank = rank // tp + # new_group is collective, so every rank builds every subgroup in the + # same order even though it only keeps one of each. + tp_group = None + dp_group = None + if tp > 1: + tp_groups = [ + dist.new_group(ranks=list(range(base * tp, (base + 1) * tp))) for base in range(dp) + ] + tp_group = tp_groups[dp_rank] + if dp > 1: + dp_groups = [ + dist.new_group(ranks=list(range(offset, world_size, tp))) for offset in range(tp) + ] + dp_group = dp_groups[tp_rank] + + objective = ObjectiveSpec(beta=BETA) + if scenario == "preflight_dp_mismatch" and dp_rank == 1: + # Perturb a pure fingerprint field. Changing the batch geometry + # instead would change this rank's local shapes, so it would fail + # while constructing its contract -- before the preflight -- and + # strand the other ranks inside the all-gather. + objective = ObjectiveSpec(beta=BETA * 2) + if scenario == "preflight_tp_mismatch" and tp_rank == 1: + # beta is invisible to the logprob contract, so only the loss + # preflight's TP-axis check can catch this one. + objective = ObjectiveSpec(beta=BETA * 2) + + contract = _build_contract( + tp_rank=tp_rank, tp=tp, dp_rank=dp_rank, dp=dp, objective=objective + ) + tensors = _rank_inputs( + _global_inputs(), + tp_rank=tp_rank, + tp=tp, + bounds=contract.sharding.sequence_shard_bounds[dp_rank], + device=device, + ) + op = DistributedGRPOLossOp() + result = op.apply( + tensors["policy"], + tensors["action_ids"], + tensors["old_logps"], + tensors["rewards"], + contract=contract, + ref_local_logits=tensors["ref"], + tp_group=tp_group, + dp_group=dp_group, + num_vocab_tiles=NUM_VOCAB_TILES, + ) + result.loss.backward() + + grad = tensors["policy"].grad + payload.update( + { + "ok": True, + "loss_bits": _scalar_bits(result.loss), + "policy_bits": _scalar_bits(result.policy_loss), + "kl_bits": _scalar_bits(result.kl), + "per_sequence_policy_bits": _vector_bits(result.per_sequence_policy), + "per_sequence_kl_bits": _vector_bits(result.per_sequence_kl), + "per_sequence_counts": result.per_sequence_active_tokens.cpu().tolist(), + "advantage_bits": _vector_bits(result.advantages), + "global_active": result.provenance["global_active_tokens"], + "backend_id": result.provenance["backend_id"], + "grad_nonzero": bool(grad.abs().sum().item() > 0.0), + # Gradient of one fixed global token slot, keyed by vocab shard + # so the parent can reassemble the full row across TP ranks. + "grad_row_bits": _vector_bits(grad[0]) if dp_rank == 0 else None, + "vocab_start": contract.logprob.sharding.local_vocab_start, + } + ) + except BaseException as exc: # the parent re-raises the text + payload.update( + { + "ok": False, + "error": f"{type(exc).__name__}: {exc}", + "tb": traceback.format_exc(), + } + ) + finally: + try: + import torch.distributed as dist + + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + except Exception: # pragma: no cover - teardown best effort + pass + try: + result_queue.put(payload) + except Exception: # pragma: no cover - queue already closed + pass + + +def _run_mesh(tp: int, dp: int, *, scenario: str = "correctness") -> list[dict]: + """Spawn a ``tp * dp`` NCCL mesh and collect one payload per rank.""" + + world_size = tp * dp + os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + ctx = mp.get_context("spawn") + result_queue = ctx.Queue() + with tempfile.TemporaryDirectory() as tmp: + init_method = f"file://{Path(tmp) / 'store'}" + spawned = mp.spawn( + _mesh_worker, + args=(world_size, init_method, result_queue, tp, dp, scenario), + nprocs=world_size, + join=False, + ) + payloads: list[dict] = [] + try: + for _ in range(world_size): + payloads.append(result_queue.get(timeout=_SPAWN_TIMEOUT_S)) + except queue.Empty: # pragma: no cover - only on a genuine hang + pytest.fail( + f"TP={tp} DP={dp}: only {len(payloads)}/{world_size} ranks reported " + f"within {_SPAWN_TIMEOUT_S}s" + ) + finally: + spawned.join(timeout=_SPAWN_TIMEOUT_S) + return sorted(payloads, key=lambda item: item["rank"]) + + +def _require_all_ok(payloads: list[dict]) -> None: + failures = [item for item in payloads if not item.get("ok")] + if failures: + head = failures[0] + pytest.fail(f"rank {head['rank']} failed: {head['error']}\n{head['tb']}") + + +def _assemble_grad_row(payloads: list[dict]) -> list[int]: + """Reassemble the fixed token slot's gradient row across TP vocab shards.""" + + contributions = [item for item in payloads if item["grad_row_bits"] is not None] + contributions.sort(key=lambda item: item["vocab_start"]) + row: list[int] = [] + for item in contributions: + row.extend(item["grad_row_bits"]) + return row + + +def _consensus(payloads: list[dict]) -> dict: + """Collapse the mesh's per-rank payloads into the one replicated answer.""" + + _require_all_ok(payloads) + replicated = ( + "loss_bits", + "policy_bits", + "kl_bits", + "per_sequence_policy_bits", + "per_sequence_kl_bits", + "per_sequence_counts", + "global_active", + "advantage_bits", + ) + for key in replicated: + values = {repr(item[key]) for item in payloads} + assert len(values) == 1, f"ranks disagree on {key}: {values}" + assert all(item["grad_nonzero"] for item in payloads), ( + "some rank produced a zero gradient; the all-gather likely severed the " + "graph for that rank's own block" + ) + head = payloads[0] + consensus = {key: head[key] for key in replicated} + consensus["grad_row"] = _assemble_grad_row(payloads) + return consensus + + +# --------------------------------------------------------------------------- # +# Single-rank behaviour +# --------------------------------------------------------------------------- # +class TestSingleRank: + def test_forward_and_backward_run(self): + result, tensors, contract = _run_single_rank() + result.loss.backward() + assert result.loss.dtype is torch.float32 + assert result.loss.shape == () + assert tensors["policy"].grad.abs().sum().item() > 0.0 + assert result.provenance["backend_id"] == BACKEND_ID + assert result.provenance["global_active_tokens"] == sum(GLOBAL_ACTIVE_MASK) + + def test_unpacks_as_the_legacy_triple(self): + result, _, _ = _run_single_rank() + loss, policy_loss, kl = result + assert _scalar_bits(loss) == _scalar_bits(result.loss) + assert _scalar_bits(policy_loss) == _scalar_bits(result.policy_loss) + assert _scalar_bits(kl) == _scalar_bits(result.kl) + + def test_run_to_run_bitwise_stability(self): + first, _, _ = _run_single_rank() + second, _, _ = _run_single_rank() + assert _scalar_bits(first.loss) == _scalar_bits(second.loss) + assert _scalar_bits(first.policy_loss) == _scalar_bits(second.policy_loss) + assert _scalar_bits(first.kl) == _scalar_bits(second.kl) + + def test_advantages_are_group_centred(self): + result, _, _ = _run_single_rank() + advantages = result.advantages + for start, end in zip(GROUP_BOUNDARIES[:-1], GROUP_BOUNDARIES[1:], strict=True): + assert advantages[start:end].sum().item() == pytest.approx(0.0, abs=1e-5) + + def test_inactive_tokens_do_not_affect_the_loss(self): + # Inactive slots carry real logits here, so a backend that forgot to + # mask them would shift the loss rather than merely change a padding. + baseline, _, _ = _run_single_rank() + perturbed_globals = _global_inputs() + inactive = [i for i, flag in enumerate(GLOBAL_ACTIVE_MASK) if not flag] + perturbed_globals["policy"][inactive] += 7.5 + perturbed_globals["old_logps"][inactive] -= 3.25 + perturbed, _, _ = _run_single_rank(globals_=perturbed_globals) + assert _scalar_bits(perturbed.loss) == _scalar_bits(baseline.loss) + + def test_dispatch_resolves_this_backend(self): + from rl_engine.kernels.registry import kernel_registry + + _, _, contract = _run_single_rank() + dispatched = kernel_registry.get_loss_op(contract) + assert dispatched.capability.backend_id == BACKEND_ID + assert isinstance(dispatched.op, DistributedGRPOLossOp) + assert dispatched.provenance["fallback"] is False + + +class TestKLZeroIdentity: + """Acceptance criterion: the reference-equals-policy identity is exact.""" + + def _identity_run(self, **overrides): + # old_logps set to the operator's own selected logprob, so the ratio is + # exp(0) = 1 exactly rather than approximately. + contract, tensors = _single_rank_setup(**overrides) + op = DistributedGRPOLossOp() + with torch.no_grad(): + logp, _ = op._logprob.apply( + tensors["policy"], + tensors["action_ids"], + contract=contract.logprob, + num_vocab_tiles=NUM_VOCAB_TILES, + ) + result = op.apply( + tensors["policy"], + tensors["action_ids"], + logp, + tensors["rewards"], + contract=contract, + ref_local_logits=tensors["policy"].detach(), + num_vocab_tiles=NUM_VOCAB_TILES, + ) + return result, tensors + + def test_kl_is_exactly_zero(self): + result, _ = self._identity_run() + assert _scalar_bits(result.kl) == _scalar_bits(torch.zeros(())) + + def test_loss_reduces_to_the_policy_term(self): + result, _ = self._identity_run() + assert _scalar_bits(result.loss) == _scalar_bits(result.policy_loss) + + def test_ratio_is_exactly_one_so_clipping_cannot_bind(self): + # A ratio that is only approximately 1 would land outside a sufficiently + # tight clip range and change the answer; an exact 1 cannot. + tight, _ = self._identity_run( + objective=ObjectiveSpec(beta=BETA, clip=ClipSpec(clip_eps_low=1e-7, clip_eps_high=1e-7)) + ) + loose, _ = self._identity_run( + objective=ObjectiveSpec(beta=BETA, clip=ClipSpec(clip_eps_low=0.9, clip_eps_high=0.9)) + ) + assert _scalar_bits(tight.loss) == _scalar_bits(loose.loss) + + def test_policy_term_matches_the_masked_advantage_mean(self): + result, _ = self._identity_run() + advantages = result.advantages + mask = torch.tensor(GLOBAL_ACTIVE_MASK) + per_token = advantages.reshape(-1, 1).expand(NUM_SEQUENCES, PADDED_SEQ_LEN).reshape(-1) + expected = -per_token.masked_fill(~mask, 0.0).sum() / mask.sum() + # Not a bitwise comparison: this reference sums the flat [N] vector, + # while the operator sums through its fixed tile structure. + assert result.policy_loss.item() == pytest.approx(expected.item(), abs=1e-6) + + +class TestNormalizers: + def test_global_active_tokens_matches_a_flat_masked_mean(self): + result, _, _ = _run_single_rank() + assert result.provenance["global_active_tokens"] == sum(GLOBAL_ACTIVE_MASK) + + def test_normalizers_disagree_on_unequal_sequence_lengths(self): + # If these ever agreed, the normalizer choice would be untestable and + # the sequence lengths in this file would have stopped being staggered. + token_mean, _, _ = _run_single_rank() + seq_mean, _, _ = _run_single_rank( + reduction=LossReductionSpec(token_normalizer=TokenNormalizer.PER_SEQUENCE_THEN_MEAN) + ) + fixed, _, _ = _run_single_rank( + reduction=LossReductionSpec( + token_normalizer=TokenNormalizer.FIXED_CONSTANT, + fixed_normalizer_constant=NUM_TOKEN_SLOTS, + ) + ) + values = { + _scalar_bits(token_mean.loss), + _scalar_bits(seq_mean.loss), + _scalar_bits(fixed.loss), + } + assert len(values) == 3 + + def test_fixed_constant_normalizer_scales_the_token_sum(self): + active = sum(GLOBAL_ACTIVE_MASK) + token_mean, _, _ = _run_single_rank() + fixed, _, _ = _run_single_rank( + reduction=LossReductionSpec( + token_normalizer=TokenNormalizer.FIXED_CONSTANT, + fixed_normalizer_constant=NUM_TOKEN_SLOTS, + ) + ) + assert fixed.policy_loss.item() == pytest.approx( + token_mean.policy_loss.item() * active / NUM_TOKEN_SLOTS, rel=1e-6 + ) + + def test_kl_estimators_differ(self): + k3, _, _ = _run_single_rank() + k1, _, _ = _run_single_rank( + objective=ObjectiveSpec(beta=BETA, kl_estimator=KLEstimator.K1_LOG_RATIO) + ) + assert _scalar_bits(k3.kl) != _scalar_bits(k1.kl) + # k3 is non-negative by construction; the plain log-ratio is not. + assert k3.kl.item() >= 0.0 + + def test_mean_only_advantage_skips_the_std_divisor(self): + std_normalized, _, _ = _run_single_rank() + mean_only, _, _ = _run_single_rank( + objective=ObjectiveSpec( + beta=BETA, + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY), + ) + ) + assert _scalar_bits(std_normalized.loss) != _scalar_bits(mean_only.loss) + + +class TestNegativeControls: + """Prove the bitwise assertions elsewhere are not comparing constants. + + Each control changes something that genuinely regroups or rescales the + reduction and asserts the compared surface notices. If one of these ever + starts passing trivially, the corresponding positive assertion has stopped + meaning anything. + """ + + def test_regrouping_the_token_sum_moves_the_per_sequence_vector(self): + # Split each sequence's token sum in half before adding, changing the + # summation tree without changing a single input value. The scalar loss + # absorbs that most of the time; the per-sequence vector does not, which + # is why it is the comparison surface everywhere else in this file. + import rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss as module + + baseline, _, _ = _run_single_rank() + original = module._sequence_totals + + def halved(values, contract): + view = values.reshape( + contract.sharding.local_num_sequences, 2, contract.sharding.padded_seq_len // 2 + ) + return view.sum(dim=2).sum(dim=1) + + module._sequence_totals = halved + try: + regrouped, _, _ = _run_single_rank() + finally: + module._sequence_totals = original + assert _vector_bits(baseline.per_sequence_policy) != _vector_bits( + regrouped.per_sequence_policy + ) + + def test_vocab_tile_count_perturbs_the_logprob_it_consumes(self): + baseline, _, _ = _run_single_rank() + retiled, _, _ = _run_single_rank(num_vocab_tiles=NUM_VOCAB_TILES * 2) + assert _reduction_fingerprint(baseline) != _reduction_fingerprint(retiled) + + def test_sequence_order_matters_to_the_scalar(self): + # An all_reduce would combine sequence totals in topology order rather + # than in global sequence order. Permuting the assembled grid is a + # stand-in for that mistake, and the loss must notice. + import rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss as module + + baseline, _, _ = _run_single_rank() + permutation = torch.tensor([5, 2, 7, 0, 3, 6, 1, 4]) + original = module._assemble_global_vector + module._assemble_global_vector = lambda totals, contract, group: original( + totals, contract, group + )[permutation] + try: + permuted, _, _ = _run_single_rank() + finally: + module._assemble_global_vector = original + assert _vector_bits(baseline.per_sequence_policy) != _vector_bits( + permuted.per_sequence_policy + ) + + def test_clip_epsilon_perturbs_the_loss(self): + baseline, _, _ = _run_single_rank() + clipped, _, _ = _run_single_rank( + objective=ObjectiveSpec(beta=BETA, clip=ClipSpec(clip_eps_low=0.01, clip_eps_high=0.01)) + ) + assert _reduction_fingerprint(baseline) != _reduction_fingerprint(clipped) + + +class TestGuards: + def test_reference_logits_required_when_beta_is_positive(self): + contract, tensors = _single_rank_setup() + with pytest.raises(LossContractError, match="ref_local_logits is required"): + DistributedGRPOLossOp().apply( + tensors["policy"], + tensors["action_ids"], + tensors["old_logps"], + tensors["rewards"], + contract=contract, + num_vocab_tiles=NUM_VOCAB_TILES, + ) + + def test_reference_optional_when_beta_is_zero(self): + result, _, _ = _run_single_rank(objective=ObjectiveSpec(beta=0.0), with_reference=False) + assert _scalar_bits(result.loss) == _scalar_bits(result.policy_loss) + assert _scalar_bits(result.kl) == _scalar_bits(torch.zeros(())) + assert result.provenance["reference_model_used"] is False + + def test_row_count_must_match_owned_slots(self): + contract, tensors = _single_rank_setup() + with pytest.raises(LossContractError, match="token slots"): + DistributedGRPOLossOp().apply( + tensors["policy"][:-1], + tensors["action_ids"][:-1], + tensors["old_logps"][:-1], + tensors["rewards"], + contract=contract, + ref_local_logits=tensors["ref"][:-1], + num_vocab_tiles=NUM_VOCAB_TILES, + ) + + def test_reward_count_must_match_owned_sequences(self): + contract, tensors = _single_rank_setup() + with pytest.raises(LossContractError, match="one entry per sequence"): + DistributedGRPOLossOp().apply( + tensors["policy"], + tensors["action_ids"], + tensors["old_logps"], + tensors["rewards"][:-1], + contract=contract, + ref_local_logits=tensors["ref"], + num_vocab_tiles=NUM_VOCAB_TILES, + ) + + +# --------------------------------------------------------------------------- # +# The claim: bitwise equality across the mesh +# --------------------------------------------------------------------------- # +# Every reachable (TP, DP) combination with at most 4 ranks. Larger degrees +# need a bigger node and run unchanged there via the same helper. +MESH_CONFIGS = [ + pytest.param(2, 1, id="tp2"), + pytest.param(4, 1, id="tp4"), + pytest.param(1, 2, id="dp2"), + pytest.param(1, 4, id="dp4"), + pytest.param(2, 2, id="tp2xdp2"), +] + + +@pytest.fixture(scope="module") +def gpu_baseline() -> dict: + """The TP=DP=1 answer, computed on one GPU so the mesh comparison is + device-for-device rather than CPU-versus-GPU. Built once for the module.""" + + if _cuda_device_count() < 1: + pytest.skip("needs a CUDA device") + return _consensus(_run_mesh(1, 1)) + + +@_requires_gpus(1) +class TestMeshBitwise: + @pytest.mark.parametrize(("tp", "dp"), MESH_CONFIGS) + def test_mesh_matches_the_single_rank_baseline(self, tp, dp, gpu_baseline): + world = tp * dp + if _cuda_device_count() < world: + pytest.skip(f"needs {world} CUDA devices for TP={tp} DP={dp}") + baseline = gpu_baseline + actual = _consensus(_run_mesh(tp, dp)) + label = f"TP={tp} DP={dp}" + + assert actual["global_active"] == baseline["global_active"] + assert actual["per_sequence_counts"] == baseline["per_sequence_counts"] + assert actual["advantage_bits"] == baseline["advantage_bits"] + # The sensitive comparison: per-sequence totals, before the scalar + # average rounds a reordering away. + assert actual["per_sequence_policy_bits"] == baseline["per_sequence_policy_bits"], ( + f"{label} per-sequence policy totals differ from the single-rank baseline" + ) + assert actual["per_sequence_kl_bits"] == baseline["per_sequence_kl_bits"], ( + f"{label} per-sequence KL totals differ from the single-rank baseline" + ) + assert actual["loss_bits"] == baseline["loss_bits"], f"{label} loss differs" + assert actual["policy_bits"] == baseline["policy_bits"] + assert actual["kl_bits"] == baseline["kl_bits"] + assert actual["grad_row"] == baseline["grad_row"], ( + f"{label} gradient differs from the single-rank baseline" + ) + + @_requires_gpus(4) + def test_pure_axes_agree_with_each_other(self): + # Transitively implied by both matching the baseline; kept because a + # direct TP-vs-DP comparison names the culprit when a shared drift moves + # both away from the baseline at once. + tp4 = _consensus(_run_mesh(4, 1)) + dp4 = _consensus(_run_mesh(1, 4)) + assert tp4["per_sequence_policy_bits"] == dp4["per_sequence_policy_bits"] + assert tp4["loss_bits"] == dp4["loss_bits"] + assert tp4["grad_row"] == dp4["grad_row"] + + +@_requires_gpus(2) +class TestMeshGuards: + def test_preflight_rejects_a_dp_rank_that_disagrees(self): + payloads = _run_mesh(1, 2, scenario="preflight_dp_mismatch") + assert all(not item.get("ok") for item in payloads), ( + "every rank must abort when one of them declares a different objective" + ) + assert any("preflight" in item.get("error", "") for item in payloads) + + def test_preflight_rejects_a_tp_rank_that_disagrees(self): + # beta is not part of the logprob contract, so the logprob path's own TP + # preflight cannot see this; only the loss preflight's TP-axis check can. + # Without it two TP siblings would compute different losses for one + # sharded model and nothing would notice. + payloads = _run_mesh(2, 1, scenario="preflight_tp_mismatch") + assert all(not item.get("ok") for item in payloads) + assert any("TP axis" in item.get("error", "") for item in payloads) diff --git a/tests/test_grpo_loss_contract.py b/tests/test_grpo_loss_contract.py new file mode 100644 index 00000000..2a17bce4 --- /dev/null +++ b/tests/test_grpo_loss_contract.py @@ -0,0 +1,392 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Unit tests for the WS2 deterministic GRPO loss contract. + +These are pure-Python contract checks: no tensors, no collectives, no GPU. +The distributed behaviour they describe is exercised in +``tests/test_distributed_grpo_loss.py``. +""" + +from __future__ import annotations + +import json + +import pytest + +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobContract, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.loss_contract import ( + AdvantageNormalizer, + AdvantageSpec, + ClipSpec, + GRPOLossContract, + KLEstimator, + LossBackendCapability, + LossContractError, + LossReductionSpec, + LossShardingSpec, + ObjectiveSpec, + TokenNormalizer, +) + +# Global batch geometry shared by every fixture below: 8 sequences of 8 token +# slots, in two equal advantage groups. Tests that care about variable or +# straddling groups override group_boundaries explicitly. +NUM_SEQUENCES = 8 +PADDED_SEQ_LEN = 8 +GROUP_BOUNDARIES = (0, 4, NUM_SEQUENCES) +REAL_VOCAB = 30 +PADDED_VOCAB = 32 + + +def _dp_bounds(dp: int) -> tuple[tuple[int, int], ...]: + """Contiguous sequence partition in DP-rank order.""" + + seqs = NUM_SEQUENCES // dp + return tuple((d * seqs, (d + 1) * seqs) for d in range(dp)) + + +def _logprob_contract(num_tokens: int, **overrides) -> LogprobContract: + kwargs = { + "role": LogprobRole.TRAIN, + "dtype": LogprobDType.FP32, + "mask": MaskSpec(num_tokens=num_tokens, active_mask=(True,) * num_tokens), + "sharding": ShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, PADDED_VOCAB),), + real_vocab_size=REAL_VOCAB, + padded_vocab_size=PADDED_VOCAB, + ), + "reduction": ReductionSpec(), + } + kwargs.update(overrides) + return LogprobContract(**kwargs) + + +def _sharding(*, dp_rank: int = 0, dp: int = 1, **overrides) -> LossShardingSpec: + kwargs = { + "dp_rank": dp_rank, + "dp_world_size": dp, + "num_sequences": NUM_SEQUENCES, + "padded_seq_len": PADDED_SEQ_LEN, + "sequence_shard_bounds": _dp_bounds(dp), + "group_boundaries": GROUP_BOUNDARIES, + } + kwargs.update(overrides) + return LossShardingSpec(**kwargs) + + +def _contract(*, dp_rank: int = 0, dp: int = 1, **overrides) -> GRPOLossContract: + sharding = overrides.pop("sharding", _sharding(dp_rank=dp_rank, dp=dp)) + kwargs = { + "logprob": _logprob_contract(sharding.local_num_token_slots), + "sharding": sharding, + "objective": ObjectiveSpec(), + "reduction": LossReductionSpec(), + } + kwargs.update(overrides) + return GRPOLossContract(**kwargs) + + +class TestSequenceOwnership: + def test_single_rank_owns_every_sequence(self): + sharding = _sharding() + assert sharding.local_sequence_start == 0 + assert sharding.local_sequence_end == NUM_SEQUENCES + assert sharding.local_num_token_slots == NUM_SEQUENCES * PADDED_SEQ_LEN + assert sharding.num_groups == 2 + assert sharding.group_sizes == (4, 4) + + @pytest.mark.parametrize("dp", [1, 2, 4, 8]) + def test_dp_shapes_partition_the_batch(self, dp): + total = 0 + for dp_rank in range(dp): + sharding = _sharding(dp_rank=dp_rank, dp=dp) + assert sharding.local_num_sequences == NUM_SEQUENCES // dp + total += sharding.local_num_token_slots + assert total == NUM_SEQUENCES * PADDED_SEQ_LEN + + def test_bounds_must_be_contiguous_in_rank_order(self): + with pytest.raises(LossContractError, match="contiguous"): + _sharding(dp=2, sequence_shard_bounds=((0, 4), (5, 8))) + + def test_bounds_must_cover_every_sequence(self): + with pytest.raises(LossContractError, match="cover num_sequences"): + _sharding(dp=2, sequence_shard_bounds=((0, 3), (3, 6))) + + def test_bounds_count_must_match_dp_world_size(self): + with pytest.raises( + LossContractError, match="exactly one \\(start, end\\) pair per DP rank" + ): + _sharding(dp=2, sequence_shard_bounds=((0, NUM_SEQUENCES),)) + + def test_empty_shard_rejected(self): + with pytest.raises(LossContractError, match="end > start"): + _sharding(dp=2, sequence_shard_bounds=((0, 0), (0, NUM_SEQUENCES))) + + def test_context_parallelism_is_rejected(self): + # CP splits a sequence's tokens across ranks, which this contract does + # not model; it must fail loudly rather than silently drop the rest. + with pytest.raises(LossContractError, match="cp_world_size=2 is unsupported"): + _sharding(cp_world_size=2) + + def test_cp_rank_must_be_zero(self): + with pytest.raises(LossContractError, match="cp_rank must be 0"): + _sharding(cp_rank=1) + + +class TestGroupBoundaries: + @pytest.mark.parametrize( + ("boundaries", "match"), + [ + ((1, NUM_SEQUENCES), "must start at 0"), + ((0, 3), "must start at 0"), + ((0, 2, 2, NUM_SEQUENCES), "strictly increasing"), + ((0,), "at least 2 entries"), + ], + ) + def test_malformed_boundaries_rejected(self, boundaries, match): + with pytest.raises(LossContractError, match=match): + _sharding(group_boundaries=boundaries) + + def test_variable_group_sizes_accepted(self): + sharding = _sharding(group_boundaries=(0, 7, NUM_SEQUENCES)) + assert sharding.group_sizes == (7, 1) + + def test_groups_may_straddle_dp_shards(self): + # A group split at 3 crosses the DP=2 shard boundary at 4, so no rank + # owns that group alone. The contract permits it; that is why the + # operator replicates advantages instead of merging partial statistics. + sharding = _sharding(dp=2, dp_rank=0, group_boundaries=(0, 3, NUM_SEQUENCES)) + assert sharding.sequence_shard_bounds == ((0, 4), (4, NUM_SEQUENCES)) + assert sharding.group_sizes == (3, 5) + + def test_population_std_rejects_singleton_groups(self): + # A one-sequence group has zero population variance, so its advantage + # would silently collapse to zero rather than fail. + with pytest.raises(LossContractError, match="at least 2 sequences per group"): + _contract( + sharding=_sharding(group_boundaries=(0, 7, NUM_SEQUENCES)), + objective=ObjectiveSpec(), + ) + + def test_mean_only_allows_singleton_groups(self): + contract = _contract( + sharding=_sharding(group_boundaries=(0, 7, NUM_SEQUENCES)), + objective=ObjectiveSpec( + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY) + ), + ) + assert contract.sharding.group_sizes == (7, 1) + + +class TestSpecValidation: + def test_accumulation_must_be_fp32(self): + with pytest.raises(LossContractError, match="must be fp32"): + LossReductionSpec(acc_dtype=LogprobDType.BF16) + + def test_fixed_constant_normalizer_requires_its_constant(self): + with pytest.raises(LossContractError, match="fixed_normalizer_constant"): + LossReductionSpec(token_normalizer=TokenNormalizer.FIXED_CONSTANT) + + def test_constant_rejected_for_other_normalizers(self): + with pytest.raises(LossContractError, match="only meaningful for"): + LossReductionSpec(fixed_normalizer_constant=32) + + def test_fixed_constant_normalizer_accepts_its_constant(self): + spec = LossReductionSpec( + token_normalizer=TokenNormalizer.FIXED_CONSTANT, fixed_normalizer_constant=32 + ) + assert spec.fixed_normalizer_constant == 32 + + def test_lower_clip_bound_must_stay_positive(self): + with pytest.raises(LossContractError, match="must be smaller than 1.0"): + ClipSpec(clip_eps_low=1.0) + + def test_asymmetric_clip_bounds(self): + clip = ClipSpec(clip_eps_low=0.2, clip_eps_high=0.28) + assert clip.lower_bound == pytest.approx(0.8) + assert clip.upper_bound == pytest.approx(1.28) + + def test_std_eps_must_be_positive(self): + with pytest.raises(LossContractError, match="strictly positive"): + AdvantageSpec(std_eps=0.0) + + def test_negative_beta_rejected(self): + with pytest.raises(LossContractError, match="non-negative"): + ObjectiveSpec(beta=-0.01) + + def test_uses_reference_model_tracks_beta(self): + assert not ObjectiveSpec(beta=0.0).uses_reference_model + assert ObjectiveSpec(beta=0.04).uses_reference_model + + +class TestContractCoherence: + def test_logprob_token_count_must_match_owned_slots(self): + sharding = _sharding(dp=2, dp_rank=0) + with pytest.raises(LossContractError, match="token slots this rank owns"): + GRPOLossContract( + logprob=_logprob_contract(sharding.local_num_token_slots + 1), + sharding=sharding, + ) + + def test_determinism_scope_must_agree_with_logprob_path(self): + tokens = _sharding().local_num_token_slots + with pytest.raises(LossContractError, match="stronger or weaker determinism scope"): + GRPOLossContract( + logprob=_logprob_contract( + tokens, + reduction=ReductionSpec(determinism_scope=DeterminismScope.FIXED_TOPOLOGY), + ), + sharding=_sharding(), + ) + + def test_global_token_slots(self): + assert _contract().global_token_slots == NUM_SEQUENCES * PADDED_SEQ_LEN + + +class TestFingerprint: + @pytest.mark.parametrize("dp", [2, 4, 8]) + def test_every_dp_rank_agrees(self, dp): + # This is the property the distributed preflight relies on. Each rank + # holds a different slice of sequences, so their nested logprob masks + # genuinely differ -- the fingerprint must still match. + fingerprints = { + _contract(dp_rank=rank, dp=dp).cross_rank_fingerprint() for rank in range(dp) + } + assert len(fingerprints) == 1 + + def test_dp_degree_changes_the_fingerprint(self): + # A different partition is a different logical invocation, so a rank + # that joined the wrong one must be caught rather than merged with. + assert _contract(dp=1).cross_rank_fingerprint() != _contract(dp=2).cross_rank_fingerprint() + + @pytest.mark.parametrize( + "overrides", + [ + { + "reduction": LossReductionSpec( + token_normalizer=TokenNormalizer.PER_SEQUENCE_THEN_MEAN + ) + }, + {"objective": ObjectiveSpec(beta=0.04)}, + {"objective": ObjectiveSpec(clip=ClipSpec(clip_eps_high=0.28))}, + {"objective": ObjectiveSpec(kl_estimator=KLEstimator.K1_LOG_RATIO)}, + { + "objective": ObjectiveSpec( + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY) + ) + }, + {"sharding": _sharding(group_boundaries=(0, 3, NUM_SEQUENCES))}, + ], + ids=["normalizer", "beta", "clip", "kl", "advantage", "groups"], + ) + def test_numerical_identity_changes_the_fingerprint(self, overrides): + assert ( + _contract().cross_rank_fingerprint() != _contract(**overrides).cross_rank_fingerprint() + ) + + def test_to_dict_is_json_serializable_and_stable(self): + contract = _contract() + first = json.dumps(contract.to_dict(), sort_keys=True) + second = json.dumps(contract.to_dict(), sort_keys=True) + assert first == second + payload = contract.to_dict() + assert payload["semantic_operator"] == "grpo_loss" + assert payload["reduction"]["cp_is_merge_axis"] is False + assert payload["reduction"]["dp_is_merge_axis"] is True + assert payload["logprob"]["reduction"]["cp_is_merge_axis"] is False + + +def _capability(**overrides) -> LossBackendCapability: + kwargs = { + "backend_id": "test-loss-backend", + "token_normalizers": frozenset({TokenNormalizer.GLOBAL_ACTIVE_TOKENS}), + "kl_estimators": frozenset({KLEstimator.K3_UNBIASED}), + "advantage_normalizers": frozenset({AdvantageNormalizer.MEAN_STD_POPULATION}), + "determinism_scopes": frozenset({DeterminismScope.CROSS_TP_BITWISE}), + "implementation_kind": "reference", + } + kwargs.update(overrides) + return LossBackendCapability(**kwargs) + + +class TestBackendCapability: + def test_matching_capability_supports_contract(self): + assert _capability().supports(_contract()) + + def test_backend_id_may_not_shadow_a_dispatch_policy(self): + with pytest.raises(LossContractError, match="reserved dispatch policy"): + _capability(backend_id="reference") + + def test_unsupported_normalizer_is_reported(self): + contract = _contract( + reduction=LossReductionSpec(token_normalizer=TokenNormalizer.PER_SEQUENCE_THEN_MEAN) + ) + reasons = _capability().incompatibilities(contract) + assert any("token_normalizer" in reason for reason in reasons) + + def test_unsupported_dp_degree_is_reported(self): + capability = _capability(dp_world_sizes=(1,)) + reasons = capability.incompatibilities(_contract(dp=2)) + assert any("DP=2" in reason for reason in reasons) + + def test_variable_group_sizes_gated(self): + contract = _contract( + sharding=_sharding(group_boundaries=(0, 7, NUM_SEQUENCES)), + objective=ObjectiveSpec( + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY) + ), + ) + capability = _capability(advantage_normalizers=frozenset({AdvantageNormalizer.MEAN_ONLY})) + assert any( + "variable advantage group sizes" in reason + for reason in capability.incompatibilities(contract) + ) + + def test_variable_group_sizes_allowed_when_declared(self): + contract = _contract( + sharding=_sharding(group_boundaries=(0, 7, NUM_SEQUENCES)), + objective=ObjectiveSpec( + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY) + ), + ) + capability = _capability( + advantage_normalizers=frozenset({AdvantageNormalizer.MEAN_ONLY}), + supports_variable_group_sizes=True, + ) + assert capability.supports(contract) + + def test_asymmetric_clip_gated(self): + contract = _contract(objective=ObjectiveSpec(clip=ClipSpec(clip_eps_high=0.28))) + assert any( + "asymmetric ratio clipping" in reason + for reason in _capability().incompatibilities(contract) + ) + assert _capability(supports_asymmetric_clip=True).supports(contract) + + def test_every_incompatibility_is_reported_at_once(self): + contract = _contract( + dp=2, + reduction=LossReductionSpec(token_normalizer=TokenNormalizer.PER_SEQUENCE_THEN_MEAN), + objective=ObjectiveSpec( + kl_estimator=KLEstimator.K1_LOG_RATIO, + clip=ClipSpec(clip_eps_high=0.28), + ), + ) + reasons = _capability(dp_world_sizes=(1,)).incompatibilities(contract) + assert len(reasons) >= 4 + + def test_to_dict_round_trips_declared_flags(self): + payload = _capability(dp_world_sizes=(1, 2)).to_dict() + assert payload["dp_world_sizes"] == [1, 2] + assert payload["implementation_kind"] == "reference" diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py new file mode 100644 index 00000000..a961f312 --- /dev/null +++ b/tests/test_logprob_contract.py @@ -0,0 +1,616 @@ +# 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()) + + 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()) + 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] + + with pytest.raises(RuntimeError, match="no LogprobBackendCapability declared"): + registry.get_logprob_op(_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, + ) + + result = registry.get_logprob_op(_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()) + + 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()) + 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()) + + 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_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"]