From b7ffb3ad496535909215e43b0bb8dfea446aa0a5 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Sun, 2 Aug 2026 20:03:13 +0800 Subject: [PATCH 01/10] feat(ws2): add TP-aware logprob contract and dispatch metadata Implements PR 1 of issue #241: a typed contract for vocab-parallel selected-token logprob, mirroring the WS2 attention contract pattern. - rl_engine/kernels/logprob_contract.py: LogprobContract, ShardingSpec (per-rank vocab shard bounds, padded-vs-real vocab, TP/CP rank metadata, owner_rank resolution), MaskSpec (active-token mask, ignore_index), ReductionSpec (fp32 (max, sumexp) merge in fixed global vocab-shard index order, all-gather transport, CP declared a non-merge axis), and LogprobBackendCapability. - KernelRegistry.get_logprob_op(contract): contract-aware dispatch that only selects backends with a declared capability; incompatible or undeclared candidates are rejected with explicit reasons and never used as a silent fallback. Existing WS1 batch-invariant logp backends are declared truthfully as single-shard references, so strict WS2 requests fail loudly until the deterministic vocab-parallel TP reference (PR 3) lands. Legacy get_op() behavior is unchanged. - Design doc, runtime-dispatch and operator doc updates, and CPU-safe contract/dispatch tests covering the Qwen3-8B TP=2 BF16 target and the TP=1/2/4 sweep shapes. Tolerance values remain owned by #108. --- .github/workflows/ci.yml | 3 + docs/design/runtime-dispatch.md | 7 + docs/design/ws2-tp-logprob-contract.md | 196 ++++++++++ docs/operators/batch-invariant-logp.md | 13 + rl_engine/kernels/logprob_contract.py | 501 +++++++++++++++++++++++++ rl_engine/kernels/registry.py | 192 ++++++++++ tests/test_logprob_contract.py | 402 ++++++++++++++++++++ 7 files changed, 1314 insertions(+) create mode 100644 docs/design/ws2-tp-logprob-contract.md create mode 100644 rl_engine/kernels/logprob_contract.py create mode 100644 tests/test_logprob_contract.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..28cdb58d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,6 +76,9 @@ 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 + docs: runs-on: ubuntu-latest steps: diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index bedf3475..84f29ec3 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -11,6 +11,13 @@ 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. +See [WS2 TP-aware logprob contract](ws2-tp-logprob-contract.md). + ## LogP Priority | Platform | Priority | diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md new file mode 100644 index 00000000..a392393b --- /dev/null +++ b/docs/design/ws2-tp-logprob-contract.md @@ -0,0 +1,196 @@ +# WS2 TP-Aware Logprob Contract + +Status: PR1 contract and dispatch metadata + +Tracking and shared contracts: + +- [#241: TP-aware deterministic logprob](https://github.com/RL-Align/RL-Kernel/issues/241) +- [#83: WS2 roadmap](https://github.com/RL-Align/RL-Kernel/issues/83) +- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) +- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) +- [#116: WS2 tolerance and drift-report format](https://github.com/RL-Align/RL-Kernel/issues/116) +- [Cross-config logprob drift contract](ws2_cross_config_logprob_drift_contract.md) + +## Scope + +This contract describes the logical inputs and deterministic reduction semantics for +selected-token log-probability under vocab-parallel tensor parallelism (TP): + +```text +selected_logp[t] = logits[t, target[t]] - logsumexp_vocab(logits[t, :]) +``` + +Under vocab-parallel TP each rank holds one vocabulary shard, so the vocabulary-wide +`logsumexp` requires a cross-rank reduction. This contract lets runtime dispatch reject a +backend whose numerical semantics do not match the requested layout. + +This PR1 layer does not shard tensors, launch a collective, merge `(max, sumexp)` partial +states, or implement a kernel. The single-GPU harness registration, the deterministic +vocab-parallel TP reference, and the cross-config integration belong to later PRs in #241. + +Context parallelism (CP) 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. + +## Contract Objects + +`rl_engine.kernels.logprob_contract` defines: + +- `LogprobContract`: role, logits dtype, mask, sharding, reduction, and LSE export; +- `ShardingSpec`: per-rank vocab-shard bounds, padded-vs-real vocabulary, TP/CP rank + metadata, and target-token ownership; +- `MaskSpec`: active-token mask and ignore index; +- `ReductionSpec`: fixed `(max, sumexp)` merge semantics; +- `LogprobBackendCapability`: the layouts and semantics a backend explicitly supports. + +Construction performs validation immediately. A structurally valid contract means that the +request is complete and internally consistent; it does not mean that an installed backend can +materialize it. + +`ShardingSpec.vocab_shard_bounds` lists every TP rank's half-open `[start, end)` vocab range +indexed by TP rank. The full table is required on every rank: it defines target ownership +and the fixed merge order without any collective, and it makes an incomplete or overlapping +partition a loud construction-time error instead of a silent runtime divergence. +`ShardingSpec.owner_rank(token_id)` resolves the unique owning rank for a real-vocab token +and rejects everything else. + +`padded_vocab_size` is the shard-covered (weight) vocabulary; `real_vocab_size` is the +tokenizer vocabulary. Padding columns occupy `[real_vocab_size, padded_vocab_size)` and must +be excluded from the logsumexp by any conforming implementation. The two sizes are equal +when the vocabulary is unpadded. + +Inactive tokens (prompt, padding, masked-out response positions) are excluded from every +drift aggregate and are exempt from the exactly-one-owner target gather; their targets may +legally hold `ignore_index`. `ignore_index` must not collide with the real vocabulary. + +## Qwen3-8B TP=2 BF16 Example + +```python +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + MaskSpec, + ReductionSpec, + ShardingSpec, +) + +sharding = ShardingSpec( + tp_rank=0, + tp_world_size=2, + vocab_shard_bounds=((0, 76032), (76032, 152064)), + real_vocab_size=151936, + padded_vocab_size=152064, + cp_rank=0, + cp_world_size=2, +) + +contract = LogprobContract( + role="train", + dtype="bf16", + mask=MaskSpec( + num_tokens=8, + active_mask=(False, False, True, True, True, True, True, False), + ignore_index=-100, + ), + sharding=sharding, + reduction=ReductionSpec(), +) +``` + +Each rank owns one contiguous vocab shard; the 128 padding columns at the end of rank 1's +shard are outside the real vocabulary and never contribute to the logsumexp. The two leading +prompt tokens and the trailing padding token are inactive. + +## Reduction Semantics + +The only PR1 reduction contract is: + +```text +partial state: (local_max, local_sumexp), fp32 +merge: max_sumexp +merge_axis: tp_vocab +order: global_vocab_shard_index +transport: all_gather +downcast_at: final_write +engine: in_op_reference +``` + +Every rank computes `m_l = max(local_logits)` and `s_l = sum(exp(local_logits - m_l))` in +fp32, the partials travel by all-gather (collectives are transport only, never a numerical +reduction), and every rank merges in fixed global vocab-shard index order: + +```text +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: exactly one rank holds +each active token's target column. Downcast happens only at the final write. Because the +merge order is fixed by shard index, TP=2 is bitwise-equal to TP=1 by construction; averaging +per-rank logsumexp values or letting a collective reduce numerically is not conformant. + +The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift +reports follow the #116 format. This contract does not introduce another tolerance table. +The selected-token metric remains the cross-config convention: + +```text +dlogp = training-side recomputed logp - rollout-side old logp +``` + +computed over active response tokens only. + +## Contract-Aware Dispatch + +Legacy callers continue to use `KernelRegistry.get_op()`. WS2 callers use: + +```python +result = kernel_registry.get_logprob_op(contract) +op = result.op +provenance = result.provenance +``` + +Dispatch considers only backends with a `LogprobBackendCapability`. It checks role, dtype, +TP/CP degree, padded-vs-real vocab masking, inactive-token support, vocab-domain LSE export, +and deterministic TP merge. An undeclared or incompatible backend is skipped with an +explicit rejection reason; there is no silent fallback. + +`requested_backend` accepts a case-insensitive policy keyword (`auto` | `production` | +`reference` | `deterministic`; default `auto`) or an exact, case-sensitive stable backend +id. Strictness comes from the contract's capability checks, not from the policy string. A +backend id may never shadow a policy keyword; capability construction rejects that. The +provenance `fallback` flag reports only capability or load rejections of otherwise-eligible +candidates — skips caused purely by the caller's own policy filter are not fallbacks. + +WS2 dispatch resolves from its own candidate list, seeded from but decoupled from the legacy +`batch_invariant_logp` priority list: registering a TP-vocab backend for WS2 dispatch does +not change what legacy `get_op("batch_invariant_logp")` returns to WS1 callers. + +The current WS1 batch-invariant logp implementations are single-shard (TP=1) references: +they accept full-vocabulary logits with ignore-index masking but carry no vocab-shard +metadata, no padded-vs-real vocab distinction, and no public vocab-domain LSE export. Strict +WS2 requests therefore fail clearly today. The later deterministic vocab-parallel reference +becomes selectable by registering a capability that truthfully declares those features; no +controller branch or silent fallback is required. + +Successful dispatch provenance records: + +- requested and actual backend ids; +- platform and fallback status; +- prior candidate rejection reasons; +- the complete requested contract, including shard bounds, padded and real vocab sizes, + merge semantics, and the explicit `cp_is_merge_axis: false` declaration; +- the selected backend capability descriptor. + +## Validation + +Contract and dispatch behavior are covered by: + +```bash +python -m pytest tests/test_logprob_contract.py -q +``` + +The tests include Qwen3-8B TP=2 BF16 construction with padded vocab, the TP=1/2/4 sweep +shapes, incomplete/overlapping shard-bound rejection, owner-rank resolution, active-mask and +ignore-index validation, fp32-accumulation and merge-semantics enforcement, undeclared +backend rejection, no incompatible fallback, and JSON-compatible provenance. diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index fbc0e9f1..b0671c40 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -54,6 +54,19 @@ 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). +### WS2 TP-aware dispatch + +WS2 distributed callers use a separate contract-aware entry point, +`kernel_registry.get_logprob_op(contract)`. It validates explicit vocab-shard ownership, +padded-vs-real vocab metadata, active-token masking, and fixed `(max, sumexp)` merge +semantics before selecting a backend. Legacy `get_op("batch_invariant_logp")` behavior +remains unchanged. + +The backends above are single-shard (TP=1) references and do not yet export vocab-domain +LSE or carry vocab-shard metadata, so they are declared incompatible with strict WS2 +requests instead of being selected as a silent fallback. See +[WS2 TP-aware logprob contract](../design/ws2-tp-logprob-contract.md). + ## Benchmarks `benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py new file mode 100644 index 00000000..8f2667e1 --- /dev/null +++ b/rl_engine/kernels/logprob_contract.py @@ -0,0 +1,501 @@ +# 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 + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + +# Dispatch policy keywords accepted by KernelRegistry.get_logprob_op; a stable +# backend id must never shadow one of these, or it becomes unselectable by id. +RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) + + +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 partial states. + + Every rank contributes ``(local_max, local_sumexp)`` computed in the + accumulation dtype; the merged result is + ``M = max(m_l)``, ``S = sum(s_l * exp(m_l - M))``, ``LSE = M + log(S)``. + """ + + 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" + + +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 TP rank. The full table is required on every rank: + it defines target-token ownership and the fixed global-shard-index merge + order without any collective, and makes an incomplete partition a loud + construction-time error instead of a silent runtime divergence. + + ``padded_vocab_size`` is the shard-covered (weight) vocabulary; + ``real_vocab_size`` is the tokenizer vocabulary. Padding columns occupy + ``[real_vocab_size, padded_vocab_size)`` and must be excluded from the + logsumexp by any conforming implementation. + """ + + 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 ownership for one logprob invocation. + + Inactive tokens are excluded from every drift aggregate and are exempt + from the exactly-one-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) + + 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)) + + @property + def active_token_count(self) -> int: + return self._active_token_count + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic TP-vocab ``(max, sumexp)`` merge semantics.""" + + 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 + + def __post_init__(self) -> None: + 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 LogprobContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: LogprobRole + dtype: LogprobDType + mask: MaskSpec + sharding: ShardingSpec + reduction: ReductionSpec + 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.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, + "cp_is_merge_axis": False, + } + mask = { + "num_tokens": self.mask.num_tokens, + "active_token_count": self.mask.active_token_count, + "active_mask": list(self.mask.active_mask), + "ignore_index": self.mask.ignore_index, + } + 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, + } + + +@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 + supports_inactive_tokens: bool = False + exports_vocab_lse: bool = False + deterministic_tp_merge: 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 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" + ) + roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) + dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) + 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") + for flag_name in ( + "supports_vocab_padding", + "supports_inactive_tokens", + "exports_vocab_lse", + "deterministic_tp_merge", + ): + if not isinstance(getattr(self, flag_name), bool): + raise LogprobContractError(f"{flag_name} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise LogprobContractError( + "implementation_kind must be production, reference, or deterministic" + ) + 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) + + @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 not self.supports_inactive_tokens + ): + reasons.append("inactive-token (ignore_index) masking is unsupported") + if contract.export_lse and not self.exports_vocab_lse: + reasons.append("vocab-domain LSE export is unsupported") + if tp_size > 1 and not self.deterministic_tp_merge: + reasons.append("deterministic TP (max, sumexp) merge 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, + "supports_inactive_tokens": self.supports_inactive_tokens, + "exports_vocab_lse": self.exports_vocab_lse, + "deterministic_tp_merge": self.deterministic_tp_merge, + "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__ = [ + "DowncastPoint", + "LogprobBackendCapability", + "LogprobContract", + "LogprobContractError", + "LogprobDispatchResult", + "LogprobDType", + "LogprobMerge", + "LogprobRole", + "MaskSpec", + "MergeAxis", + "RESERVED_DISPATCH_POLICIES", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "ReductionTransport", + "ShardingSpec", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index fb2feb6f..35ae951b 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -8,6 +8,14 @@ import torch +from rl_engine.kernels.logprob_contract import ( + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDispatchResult, + LogprobDType, + LogprobRole, +) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -164,6 +172,51 @@ def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() + # These descriptors report what the existing WS1 batch-invariant logp + # implementations actually support: single-shard (TP=1) logits with + # ignore-index masking, no vocab-shard metadata, no padded-vs-real + # vocab distinction, and no public vocab-domain LSE export. A strict + # WS2 request is rejected with explicit reasons until the deterministic + # vocab-parallel TP reference backend lands (issue #241 PR 3) instead + # of silently selecting an incompatible fallback. + common_logprob_roles = frozenset({LogprobRole.TRAIN, LogprobRole.INFER}) + common_logprob_dtypes = frozenset({LogprobDType.BF16, LogprobDType.FP16, LogprobDType.FP32}) + self._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, + supports_inactive_tokens=True, + exports_vocab_lse=False, + deterministic_tp_merge=False, + 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, + supports_inactive_tokens=True, + exports_vocab_lse=False, + deterministic_tp_merge=False, + implementation_kind="deterministic", + ), + 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, + supports_inactive_tokens=True, + exports_vocab_lse=False, + deterministic_tp_merge=False, + implementation_kind="deterministic", + ), + } + self._priority_map = { "cuda": { "logp": [ @@ -280,6 +333,17 @@ def __init__(self): self._adjust_priority_for_hardware() self._adjust_priority_from_env() + # WS2 contract-aware dispatch owns its candidate list. It is seeded + # from the legacy batch_invariant_logp priority (after hardware/env + # adjustments) but deliberately decoupled afterwards: registering a + # TP-vocab backend for WS2 dispatch must not change what legacy + # get_op("batch_invariant_logp") returns to WS1 callers, and vice + # versa. + self._logprob_candidates: Dict[str, list] = { + platform: list(ops.get("batch_invariant_logp", [])) + for platform, ops in self._priority_map.items() + } + 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"}: @@ -389,6 +453,134 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + 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() + + 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 + + for backend in candidates: + capability = self._logprob_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no LogprobBackendCapability declared") + capability_rejections += 1 + continue + capability_incompat = list(capability.incompatibilities(contract)) + policy_mismatch = self._logprob_policy_mismatch(requested_backend, capability) + reasons = capability_incompat + ([policy_mismatch] if policy_mismatch else []) + if reasons: + rejected.append(f"{backend.name}: " + "; ".join(reasons)) + if capability_incompat: + capability_rejections += 1 + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + capability_rejections += 1 + continue + + provenance = { + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": capability_rejections > 0, + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + } + return LogprobDispatchResult( + op=op, + capability=capability, + provenance=provenance, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No logprob backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, dtype={requested['dtype']}, " + f"TP={contract.sharding.tp_world_size}, CP={contract.sharding.cp_world_size}, " + f"padded_vocab={contract.sharding.padded_vocab_size}, " + f"real_vocab={contract.sharding.real_vocab_size}. Rejections: {details}" + ) + + @staticmethod + def _logprob_policy_mismatch( + requested_backend: str, + capability: LogprobBackendCapability, + ) -> str | None: + policy = requested_backend.lower() + if policy == "auto": + return None + if policy in {"production", "reference", "deterministic"}: + 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: + if device_ctx.is_rocm: + return "rocm" + if device_ctx.device_type == "cuda": + return "cuda" + return "cpu" + + 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: if device_ctx.is_rocm: diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py new file mode 100644 index 00000000..d18083f4 --- /dev/null +++ b/tests/test_logprob_contract.py @@ -0,0 +1,402 @@ +# 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 ( + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +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, (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, + supports_inactive_tokens=True, + exports_vocab_lse=True, + deterministic_tp_merge=True, + implementation_kind="deterministic", + ) + + +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", + "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 test_current_ws1_backend_rejects_strict_tp_contract_without_fallback(): + registry = KernelRegistry() + + 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 "deterministic TP (max, sumexp) merge 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() + 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_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] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + result = registry.get_logprob_op(_contract(), requested_backend="deterministic") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert result.provenance["requested_backend"] == "deterministic" + 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] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + 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_declared_backend_support(): + capability = replace(_declared_tp_backend(), supports_inactive_tokens=False) + contract = _contract() + + assert "inactive-token (ignore_index) 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] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = replace( + _declared_tp_backend(), implementation_kind="reference" + ) + + 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] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + result = registry.get_logprob_op(_contract(), requested_backend="DETERMINISTIC") + 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] = [ + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + ] + registry._logprob_capabilities[OpBackend.TRITON_BATCH_INVARIANT_LOGP] = replace( + _declared_tp_backend(), backend_id="other-compatible-backend" + ) + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + 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] = [ + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + ] + registry._logprob_capabilities[OpBackend.TRITON_BATCH_INVARIANT_LOGP] = replace( + _declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,) + ) + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + 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 From cdc11ba17621ef52e94f90fec8a477d0b9990075 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Sun, 2 Aug 2026 22:50:48 +0800 Subject: [PATCH 02/10] fix(ws2): address CodeRabbit review on logprob contract PR MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - docs: correct the TP-invariance claim — fixed merge order gives determinism per TP degree; cross-degree bitwise equality additionally requires a TP-degree-independent local tile decomposition (PR 3 obligation), otherwise #108 tolerances apply - contract: store backend_id stripped so id-based dispatch matches; summarize the active mask in to_dict() provenance instead of copying every per-token boolean; sort __all__ per RUF022 - registry: add public register_logprob_backend() seam for PR 3 and tests; delegate _platform() to _platform_for_device(None); reuse _get_or_create_backend() in get_op so WS2 and legacy dispatch share one cache/blacklist code path - tests: use the registration seam instead of poking private state, pin _even_bounds' last bound for non-divisible vocabularies, assert candidate-list decoupling in both directions, cover registration replace semantics and backend_id normalization --- docs/design/ws2-tp-logprob-contract.md | 16 +++-- rl_engine/kernels/logprob_contract.py | 8 ++- rl_engine/kernels/registry.py | 54 ++++++++++------ tests/test_logprob_contract.py | 88 +++++++++++++++++++------- 4 files changed, 116 insertions(+), 50 deletions(-) diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md index a392393b..3b1cf2fd 100644 --- a/docs/design/ws2-tp-logprob-contract.md +++ b/docs/design/ws2-tp-logprob-contract.md @@ -127,8 +127,15 @@ selected_logp = target_logit - LSE The selected target logit comes from a masked single-owner gather: exactly one rank holds each active token's target column. Downcast happens only at the final write. Because the -merge order is fixed by shard index, TP=2 is bitwise-equal to TP=1 by construction; averaging -per-rank logsumexp values or letting a collective reduce numerically is not conformant. +merge order is fixed by shard index, the result is deterministic and reproducible at every +TP degree by construction. Cross-degree bitwise equality (TP=2 equal to TP=1, the #241 PR 3 +acceptance target) requires one further condition: the local per-shard reduction must use a +TP-degree-independent tile decomposition, so that the same partial sums are formed in the +same order regardless of how the vocabulary is sharded. Providing that decomposition is an +obligation of the deterministic reference implementation; a backend without it is still +deterministic per degree, and its cross-degree drift is judged against the #108 tolerance +table instead. Averaging per-rank logsumexp values or letting a collective reduce +numerically is not conformant in either case. The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift reports follow the #116 format. This contract does not introduce another tolerance table. @@ -170,8 +177,9 @@ The current WS1 batch-invariant logp implementations are single-shard (TP=1) ref they accept full-vocabulary logits with ignore-index masking but carry no vocab-shard metadata, no padded-vs-real vocab distinction, and no public vocab-domain LSE export. Strict WS2 requests therefore fail clearly today. The later deterministic vocab-parallel reference -becomes selectable by registering a capability that truthfully declares those features; no -controller branch or silent fallback is required. +becomes selectable through `KernelRegistry.register_logprob_backend(backend, capability)` +by declaring a capability that truthfully describes those features; no controller branch or +silent fallback is required. Successful dispatch provenance records: diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 8f2667e1..815d232f 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -342,10 +342,11 @@ def to_dict(self) -> dict[str, Any]: "engine": self.reduction.engine.value, "cp_is_merge_axis": False, } + # The per-token mask is deliberately summarized: provenance exists for + # logging/serialization and the raw mask would dominate its size. mask = { "num_tokens": self.mask.num_tokens, "active_token_count": self.mask.active_token_count, - "active_mask": list(self.mask.active_mask), "ignore_index": self.mask.ignore_index, } return { @@ -382,6 +383,7 @@ def __post_init__(self) -> None: raise LogprobContractError( f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" ) + object.__setattr__(self, "backend_id", self.backend_id.strip()) roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) if not roles or not dtypes: @@ -482,17 +484,17 @@ class LogprobDispatchResult: __all__ = [ + "RESERVED_DISPATCH_POLICIES", "DowncastPoint", "LogprobBackendCapability", "LogprobContract", "LogprobContractError", - "LogprobDispatchResult", "LogprobDType", + "LogprobDispatchResult", "LogprobMerge", "LogprobRole", "MaskSpec", "MergeAxis", - "RESERVED_DISPATCH_POLICIES", "ReductionEngine", "ReductionOrder", "ReductionSpec", diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 35ae951b..d9b741bf 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -433,25 +433,41 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: - if backend.name in self._instance_cache: - return self._instance_cache[backend.name] + op_instance = self._get_or_create_backend(backend) + if op_instance is not None: + return op_instance - if backend.name in self._failed_backends: - continue + raise RuntimeError(f"No functional backend found for {op_type} on {platform}") - 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) - else: - self._failed_backends.add(backend.name) + 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. + """ - raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + 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() + candidates = self._logprob_candidates.setdefault(resolved_platform, []) + self._logprob_capabilities[backend] = capability + if backend not in candidates: + if prepend: + candidates.insert(0, backend) + else: + candidates.append(backend) def get_logprob_op( self, @@ -556,11 +572,7 @@ def _logprob_policy_mismatch( ) def _platform(self) -> str: - if device_ctx.is_rocm: - return "rocm" - if device_ctx.device_type == "cuda": - return "cuda" - return "cpu" + return self._platform_for_device(None) def _get_or_create_backend(self, backend: OpBackend) -> Any | None: if backend.name in self._instance_cache: diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py index d18083f4..8a7a1ec5 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -28,7 +28,10 @@ def _even_bounds(padded_vocab: int, tp_world_size: int) -> tuple[tuple[int, int], ...]: shard = padded_vocab // tp_world_size - return tuple((rank * shard, (rank + 1) * shard) for rank in range(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( @@ -274,8 +277,10 @@ def test_undeclared_backend_capability_is_never_selected(): def test_declared_compatible_backend_resolves_and_records_provenance(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + 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="deterministic") @@ -293,8 +298,10 @@ def test_declared_compatible_backend_resolves_and_records_provenance(): def test_requested_stable_backend_id_is_enforced(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + 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") @@ -333,9 +340,11 @@ def test_backend_id_must_not_shadow_a_reserved_policy_keyword(): def test_default_auto_policy_resolves_any_compatible_implementation_kind(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = replace( - _declared_tp_backend(), implementation_kind="reference" + 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()) @@ -347,8 +356,10 @@ def test_default_auto_policy_resolves_any_compatible_implementation_kind(): def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + 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="DETERMINISTIC") assert result.capability.backend_id == "test-deterministic-tp-logprob" @@ -360,14 +371,15 @@ def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): def test_policy_only_skips_are_not_reported_as_fallback(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [ + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( OpBackend.TRITON_BATCH_INVARIANT_LOGP, - OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, - ] - registry._logprob_capabilities[OpBackend.TRITON_BATCH_INVARIANT_LOGP] = replace( - _declared_tp_backend(), backend_id="other-compatible-backend" + 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 ) - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") @@ -378,14 +390,15 @@ def test_policy_only_skips_are_not_reported_as_fallback(): def test_capability_rejections_are_reported_as_fallback(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [ + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( OpBackend.TRITON_BATCH_INVARIANT_LOGP, - OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, - ] - registry._logprob_capabilities[OpBackend.TRITON_BATCH_INVARIANT_LOGP] = replace( - _declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,) + 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 ) - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() result = registry.get_logprob_op(_contract()) @@ -400,3 +413,34 @@ def test_ws2_candidate_list_is_decoupled_from_the_legacy_priority_map(): 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" From 6455715eb09cfd4189f53a3aceb0a0df342a4dd8 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Sun, 2 Aug 2026 23:06:28 +0800 Subject: [PATCH 03/10] fix(ws2): address second CodeRabbit round on logprob contract - docs: state that cross-TP bitwise equality needs a global tile-level merge structure independent of TP partitioning (per-shard tiles alone leave different grouping at shard boundaries), and that padded columns are masked to -inf before the local (max, sumexp) partials - registry: scope logprob capabilities per platform so the same backend enum can declare different support on cuda/rocm/cpu; validate the platform argument of register_logprob_backend against known platforms - contract: derive IMPLEMENTATION_KINDS from RESERVED_DISPATCH_POLICIES and use it for the kind check; wrap non-iterable roles/dtypes in LogprobContractError for consistent error handling - tests: cover per-platform capability scoping, unknown-platform rejection, and non-iterable roles/dtypes --- docs/design/ws2-tp-logprob-contract.md | 24 +++++++++------ rl_engine/kernels/logprob_contract.py | 15 ++++++--- rl_engine/kernels/registry.py | 22 ++++++++++++-- tests/test_logprob_contract.py | 42 ++++++++++++++++++++++++++ 4 files changed, 87 insertions(+), 16 deletions(-) diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md index 3b1cf2fd..28685bb4 100644 --- a/docs/design/ws2-tp-logprob-contract.md +++ b/docs/design/ws2-tp-logprob-contract.md @@ -114,9 +114,11 @@ downcast_at: final_write engine: in_op_reference ``` -Every rank computes `m_l = max(local_logits)` and `s_l = sum(exp(local_logits - m_l))` in -fp32, the partials travel by all-gather (collectives are transport only, never a numerical -reduction), and every rank merges in fixed global vocab-shard index order: +Every rank first masks every local column whose global id lies in +`[real_vocab_size, padded_vocab_size)` to `-inf`, so padding never contributes to the +logsumexp, then computes `m_l = max(local_logits)` and `s_l = sum(exp(local_logits - m_l))` +in fp32. The partials travel by all-gather (collectives are transport only, never a +numerical reduction), and every rank merges in fixed global vocab-shard index order: ```text M = max_l(m_l) @@ -129,12 +131,16 @@ The selected target logit comes from a masked single-owner gather: exactly one r each active token's target column. Downcast happens only at the final write. Because the merge order is fixed by shard index, the result is deterministic and reproducible at every TP degree by construction. Cross-degree bitwise equality (TP=2 equal to TP=1, the #241 PR 3 -acceptance target) requires one further condition: the local per-shard reduction must use a -TP-degree-independent tile decomposition, so that the same partial sums are formed in the -same order regardless of how the vocabulary is sharded. Providing that decomposition is an -obligation of the deterministic reference implementation; a backend without it is still -deterministic per degree, and its cross-degree drift is judged against the #108 tolerance -table instead. Averaging per-rank logsumexp values or letting a collective reduce +acceptance target) requires one further condition: the entire reduction must 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 that the TP degree only selects which rank computes +which tiles and never changes the floating-point grouping. A TP-degree-independent +decomposition inside each shard is not sufficient on its own, because shard boundaries +would still group the combines differently across degrees. Providing that global structure +is an obligation of the deterministic reference implementation; a backend without it is +still deterministic per degree, and its cross-degree drift is judged against the #108 +tolerance table instead. Averaging per-rank logsumexp values or letting a collective reduce numerically is not conformant in either case. The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 815d232f..2cb22a25 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -30,6 +30,9 @@ # Dispatch policy keywords accepted by KernelRegistry.get_logprob_op; a stable # backend id must never shadow one of these, or it becomes unselectable by id. RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) +# Policies a backend can declare as its implementation kind; "auto" is a +# selection strategy, not an implementation kind. +IMPLEMENTATION_KINDS = RESERVED_DISPATCH_POLICIES - {"auto"} class LogprobContractError(ValueError): @@ -384,8 +387,11 @@ def __post_init__(self) -> None: f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" ) object.__setattr__(self, "backend_id", self.backend_id.strip()) - roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) - dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) + 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") @@ -398,9 +404,9 @@ def __post_init__(self) -> None: ): if not isinstance(getattr(self, flag_name), bool): raise LogprobContractError(f"{flag_name} must be a bool") - if self.implementation_kind not in {"production", "reference", "deterministic"}: + if self.implementation_kind not in IMPLEMENTATION_KINDS: raise LogprobContractError( - "implementation_kind must be production, reference, or deterministic" + f"implementation_kind must be one of: {', '.join(sorted(IMPLEMENTATION_KINDS))}" ) object.__setattr__(self, "roles", roles) object.__setattr__(self, "dtypes", dtypes) @@ -484,6 +490,7 @@ class LogprobDispatchResult: __all__ = [ + "IMPLEMENTATION_KINDS", "RESERVED_DISPATCH_POLICIES", "DowncastPoint", "LogprobBackendCapability", diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index d9b741bf..0f4f2875 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -181,7 +181,7 @@ def __init__(self): # of silently selecting an incompatible fallback. common_logprob_roles = frozenset({LogprobRole.TRAIN, LogprobRole.INFER}) common_logprob_dtypes = frozenset({LogprobDType.BF16, LogprobDType.FP16, LogprobDType.FP32}) - self._logprob_capabilities = { + base_logprob_capabilities = { OpBackend.PYTORCH_BATCH_INVARIANT_LOGP: LogprobBackendCapability( backend_id="pytorch-batch-invariant-logp-ws1", roles=common_logprob_roles, @@ -343,6 +343,16 @@ def __init__(self): 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() + } def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() @@ -461,8 +471,13 @@ def register_logprob_backend( 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[backend] = capability + self._logprob_capabilities.setdefault(resolved_platform, {})[backend] = capability if backend not in candidates: if prepend: candidates.insert(0, backend) @@ -502,8 +517,9 @@ def get_logprob_op( # own requested_backend policy filter are not fallbacks. capability_rejections = 0 + platform_capabilities = self._logprob_capabilities.get(platform, {}) for backend in candidates: - capability = self._logprob_capabilities.get(backend) + capability = platform_capabilities.get(backend) if capability is None: rejected.append(f"{backend.name}: no LogprobBackendCapability declared") capability_rejections += 1 diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py index 8a7a1ec5..670225f0 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -444,3 +444,45 @@ def test_register_logprob_backend_is_the_public_registration_seam(): 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) From 3b4eaef3dd7ccb3645653e79166b03d0492355d1 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Sun, 2 Aug 2026 23:48:14 +0800 Subject: [PATCH 04/10] feat(ws2): make determinism scope and invocation surface part of the typed contract Address external review: the cross-TP bitwise guarantee lived only in prose, so a fixed-topology-deterministic backend could pass dispatch as fully conformant. - DeterminismScope (fixed_topology | cross_tp_bitwise): requested via ReductionSpec (default cross_tp_bitwise, the #241 PR 3 target), declared per backend via determinism_scopes, enforced by dispatch; replaces the deterministic_tp_merge bool - MaskMode (explicit_active_mask | ignore_index) replaces supports_inactive_tokens: the contract permits inactive targets that do not hold ignore_index, so ignore-index-only backends are rejected for contracts with inactive tokens - LogprobOutputSpec pins the output surface: fp32 selected logprob and fp32 vocab LSE, replicated across the TP group - implementation_kind is now a tier (reference | production); determinism is no longer conflated with it, and requesting "deterministic" as a policy raises a loud error pointing at determinism_scope - fallback provenance: policy evaluation now precedes capability checks, so a candidate excluded by the caller's own policy never counts as a fallback even when it also lacks capabilities - docs: define the (-inf, 0) identity partial for padding-only or all--inf shards; document that requested_backend="auto" is not distributed-safe and specify the preflight fingerprint agreement - LogprobContract.cross_rank_fingerprint(): rank-independent identity for that preflight; provenance now records active_mask_sha256 so masks with equal active counts remain distinguishable --- docs/design/ws2-tp-logprob-contract.md | 68 +++++++--- rl_engine/kernels/logprob_contract.py | 175 ++++++++++++++++++++++--- rl_engine/kernels/registry.py | 42 +++--- tests/test_logprob_contract.py | 116 ++++++++++++++-- 4 files changed, 343 insertions(+), 58 deletions(-) diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md index 28685bb4..cf9179b8 100644 --- a/docs/design/ws2-tp-logprob-contract.md +++ b/docs/design/ws2-tp-logprob-contract.md @@ -36,12 +36,17 @@ for provenance and must never widen the merge. `rl_engine.kernels.logprob_contract` defines: -- `LogprobContract`: role, logits dtype, mask, sharding, reduction, and LSE export; +- `LogprobContract`: role, logits dtype, mask, sharding, reduction, output surface, and + LSE export, plus a rank-independent `cross_rank_fingerprint()`; - `ShardingSpec`: per-rank vocab-shard bounds, padded-vs-real vocabulary, TP/CP rank metadata, and target-token ownership; - `MaskSpec`: active-token mask and ignore index; -- `ReductionSpec`: fixed `(max, sumexp)` merge semantics; -- `LogprobBackendCapability`: the layouts and semantics a backend explicitly supports. +- `ReductionSpec`: fixed `(max, sumexp)` merge semantics and the requested determinism + scope; +- `LogprobOutputSpec`: the output surface — fp32 selected logprob and fp32 vocab-domain + LSE, replicated across the TP group; +- `LogprobBackendCapability`: the layouts and semantics a backend explicitly supports, + including its mask modes and determinism scopes. Construction performs validation immediately. A structurally valid contract means that the request is complete and internally consistent; it does not mean that an installed backend can @@ -140,8 +145,21 @@ decomposition inside each shard is not sufficient on its own, because shard boun would still group the combines differently across degrees. Providing that global structure is an obligation of the deterministic reference implementation; a backend without it is still deterministic per degree, and its cross-degree drift is judged against the #108 -tolerance table instead. Averaging per-rank logsumexp values or letting a collective reduce -numerically is not conformant in either case. +tolerance table instead. The contract expresses this distinction as +`ReductionSpec.determinism_scope`: `cross_tp_bitwise` (the #241 target and the default) +versus `fixed_topology`. A backend declares the scopes it honors in +`LogprobBackendCapability.determinism_scopes`, and dispatch rejects a backend that cannot +honor the requested scope — prose obligations are not enough; the guarantee is part of the +typed contract. + +A shard may lie entirely inside the padded region, and a row's local columns may all be +`-inf` after masking. The identity partial for these cases is defined as +`(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 +evaluating `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. The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift reports follow the #116 format. This contract does not introduce another tolerance table. @@ -164,16 +182,20 @@ provenance = result.provenance ``` Dispatch considers only backends with a `LogprobBackendCapability`. It checks role, dtype, -TP/CP degree, padded-vs-real vocab masking, inactive-token support, vocab-domain LSE export, -and deterministic TP merge. An undeclared or incompatible backend is skipped with an -explicit rejection reason; there is no silent fallback. +TP/CP degree, padded-vs-real vocab masking, explicit active-mask support, vocab-domain LSE +export, and the requested determinism scope. An undeclared or incompatible backend is +skipped with an explicit rejection reason; there is no silent fallback. `requested_backend` accepts a case-insensitive policy keyword (`auto` | `production` | -`reference` | `deterministic`; default `auto`) or an exact, case-sensitive stable backend -id. Strictness comes from the contract's capability checks, not from the policy string. A -backend id may never shadow a policy keyword; capability construction rejects that. The -provenance `fallback` flag reports only capability or load rejections of otherwise-eligible -candidates — skips caused purely by the caller's own policy filter are not fallbacks. +`reference`; default `auto`) or an exact, case-sensitive stable backend id. The keywords +select an implementation tier; determinism is not a tier — it is requested through +`ReductionSpec.determinism_scope`, so `requested_backend="deterministic"` raises a loud +error instead of silently matching nothing. Strictness comes from the contract's +capability checks, not from the policy string. A backend id may never shadow a reserved +keyword; capability construction rejects that. The provenance `fallback` flag reports only +capability or load rejections of policy-eligible candidates — a candidate excluded by the +caller's own policy never counts as a fallback, even if it would also have failed +capability checks. WS2 dispatch resolves from its own candidate list, seeded from but decoupled from the legacy `batch_invariant_logp` priority list: registering a TP-vocab backend for WS2 dispatch does @@ -192,10 +214,26 @@ Successful dispatch provenance records: - requested and actual backend ids; - platform and fallback status; - prior candidate rejection reasons; -- the complete requested contract, including shard bounds, padded and real vocab sizes, - merge semantics, and the explicit `cp_is_merge_axis: false` declaration; +- the complete dispatch-relevant contract, including shard bounds, padded and real vocab + sizes, merge and output semantics, the explicit `cp_is_merge_axis: false` declaration, + and the active-mask digest (`active_mask_sha256`) — the mask's identity without its + per-token payload; - the selected backend capability descriptor. +### Distributed dispatch safety + +`get_logprob_op` resolves locally on each rank, so `requested_backend="auto"` is not +distributed-safe on its own: a load failure on one rank can resolve a different backend +than its peers, which for a collective-bearing implementation means divergent numerical +schedules or a deadlock. For `tp_world_size > 1` a caller must either request an exact +backend id or run a preflight agreement before any collective: all-gather the resolved +backend id together with `LogprobContract.cross_rank_fingerprint()` — a rank-independent +hash covering the shard-bounds table, vocab sizes, reduction/output semantics, and the +active-mask digest, excluding rank-local fields — and abort on any mismatch. Implementing +this preflight is an obligation of the #241 PR 3/PR 4 work; the backend invocation +protocol (how the contract and mask reach the implementation) is likewise defined there, +against this contract. + ## Validation Contract and dispatch behavior are covered by: diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 2cb22a25..6941c78e 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -21,6 +21,8 @@ from __future__ import annotations +import hashlib +import json from dataclasses import dataclass, field from enum import Enum from typing import Any, TypeVar @@ -29,10 +31,15 @@ # Dispatch policy keywords accepted by KernelRegistry.get_logprob_op; a stable # backend id must never shadow one of these, or it becomes unselectable by id. +# "deterministic" stays reserved even though it is no longer a policy: +# determinism is expressed through DeterminismScope, and requesting it as a +# policy is a loud error rather than a silent id mismatch. RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) -# Policies a backend can declare as its implementation kind; "auto" is a -# selection strategy, not an implementation kind. -IMPLEMENTATION_KINDS = RESERVED_DISPATCH_POLICIES - {"auto"} +# Implementation tiers a backend can declare. Determinism is deliberately a +# separate axis (DeterminismScope): a backend can be a deterministic reference, +# a deterministic production implementation, or a non-deterministic production +# implementation. +IMPLEMENTATION_KINDS = frozenset({"production", "reference"}) class LogprobContractError(ValueError): @@ -85,6 +92,40 @@ 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 (see the design doc); fixed shard-order + merging alone is not sufficient. + """ + + FIXED_TOPOLOGY = "fixed_topology" + CROSS_TP_BITWISE = "cross_tp_bitwise" + + +class MaskMode(str, Enum): + """How a backend consumes inactive-token information. + + ``explicit_active_mask``: the backend honors an arbitrary active-token + mask. ``ignore_index``: the backend only recognizes inactive tokens whose + target id equals ``ignore_index``. 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) @@ -233,6 +274,7 @@ class MaskSpec: 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") @@ -251,11 +293,19 @@ def __post_init__(self) -> None: ) 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: @@ -268,8 +318,14 @@ class ReductionSpec: 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") @@ -291,6 +347,39 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True) +class LogprobOutputSpec: + """Output surface every conforming backend must produce. + + Selected logprob and vocab-domain LSE are fp32 and replicated across the + TP group; the ``downcast_at: final_write`` rule applies to any consumer + downcast after these outputs, never inside the reduction. + """ + + 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.""" @@ -300,6 +389,7 @@ class LogprobContract: mask: MaskSpec sharding: ShardingSpec reduction: ReductionSpec + output: LogprobOutputSpec = field(default_factory=LogprobOutputSpec) export_lse: bool = True def __post_init__(self) -> None: @@ -311,6 +401,8 @@ def __post_init__(self) -> None: 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" @@ -343,15 +435,24 @@ def to_dict(self) -> dict[str, Any]: "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 per-token mask is deliberately summarized: provenance exists for - # logging/serialization and the raw mask would dominate its size. + # logging/serialization and the raw mask would dominate its size. The + # digest keeps the mask *identity* observable, so two masks with the + # same active count still produce distinguishable provenance. 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, @@ -361,7 +462,28 @@ def to_dict(self) -> dict[str, Any]: "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) @@ -374,9 +496,9 @@ class LogprobBackendCapability: tp_world_sizes: tuple[int, ...] | None = None cp_world_sizes: tuple[int, ...] | None = None supports_vocab_padding: bool = False - supports_inactive_tokens: bool = False + mask_modes: frozenset[MaskMode] = frozenset() exports_vocab_lse: bool = False - deterministic_tp_merge: bool = False + determinism_scopes: frozenset[DeterminismScope] = frozenset() implementation_kind: str = "production" def __post_init__(self) -> None: @@ -396,12 +518,19 @@ def __post_init__(self) -> None: 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") - for flag_name in ( - "supports_vocab_padding", - "supports_inactive_tokens", - "exports_vocab_lse", - "deterministic_tp_merge", - ): + 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: @@ -412,6 +541,8 @@ def __post_init__(self) -> None: 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( @@ -453,13 +584,17 @@ def incompatibilities(self, contract: LogprobContract) -> tuple[str, ...]: reasons.append("padded-vs-real vocab masking is unsupported") if ( contract.mask.active_token_count != contract.mask.num_tokens - and not self.supports_inactive_tokens + and MaskMode.EXPLICIT_ACTIVE_MASK not in self.mask_modes ): - reasons.append("inactive-token (ignore_index) masking is unsupported") + # The contract does not require inactive targets to hold + # ignore_index, so ignore-index-only masking is insufficient. + 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 tp_size > 1 and not self.deterministic_tp_merge: - reasons.append("deterministic TP (max, sumexp) merge 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: @@ -473,9 +608,9 @@ def to_dict(self) -> dict[str, Any]: "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, - "supports_inactive_tokens": self.supports_inactive_tokens, + "mask_modes": sorted(mode.value for mode in self.mask_modes), "exports_vocab_lse": self.exports_vocab_lse, - "deterministic_tp_merge": self.deterministic_tp_merge, + "determinism_scopes": sorted(scope.value for scope in self.determinism_scopes), "implementation_kind": self.implementation_kind, } @@ -492,6 +627,7 @@ class LogprobDispatchResult: __all__ = [ "IMPLEMENTATION_KINDS", "RESERVED_DISPATCH_POLICIES", + "DeterminismScope", "DowncastPoint", "LogprobBackendCapability", "LogprobContract", @@ -499,7 +635,9 @@ class LogprobDispatchResult: "LogprobDType", "LogprobDispatchResult", "LogprobMerge", + "LogprobOutputSpec", "LogprobRole", + "MaskMode", "MaskSpec", "MergeAxis", "ReductionEngine", @@ -507,4 +645,5 @@ class LogprobDispatchResult: "ReductionSpec", "ReductionTransport", "ShardingSpec", + "TPPlacement", ] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 0f4f2875..b17acf3e 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -9,12 +9,15 @@ import torch from rl_engine.kernels.logprob_contract import ( + IMPLEMENTATION_KINDS, + DeterminismScope, LogprobBackendCapability, LogprobContract, LogprobContractError, LogprobDispatchResult, LogprobDType, LogprobRole, + MaskMode, ) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -188,9 +191,9 @@ def __init__(self): dtypes=common_logprob_dtypes, tp_world_sizes=(1,), supports_vocab_padding=False, - supports_inactive_tokens=True, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), exports_vocab_lse=False, - deterministic_tp_merge=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), implementation_kind="reference", ), OpBackend.TRITON_BATCH_INVARIANT_LOGP: LogprobBackendCapability( @@ -199,10 +202,10 @@ def __init__(self): dtypes=common_logprob_dtypes, tp_world_sizes=(1,), supports_vocab_padding=False, - supports_inactive_tokens=True, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), exports_vocab_lse=False, - deterministic_tp_merge=False, - implementation_kind="deterministic", + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", ), OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90: LogprobBackendCapability( backend_id="cuda-batch-invariant-logp-sm90-ws1", @@ -210,10 +213,10 @@ def __init__(self): dtypes=frozenset({LogprobDType.BF16, LogprobDType.FP32}), tp_world_sizes=(1,), supports_vocab_padding=False, - supports_inactive_tokens=True, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), exports_vocab_lse=False, - deterministic_tp_merge=False, - implementation_kind="deterministic", + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", ), } @@ -508,6 +511,12 @@ def get_logprob_op( 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, []) @@ -524,13 +533,16 @@ def get_logprob_op( rejected.append(f"{backend.name}: no LogprobBackendCapability declared") capability_rejections += 1 continue - capability_incompat = list(capability.incompatibilities(contract)) policy_mismatch = self._logprob_policy_mismatch(requested_backend, capability) - reasons = capability_incompat + ([policy_mismatch] if policy_mismatch else []) - if reasons: - rejected.append(f"{backend.name}: " + "; ".join(reasons)) - if capability_incompat: - capability_rejections += 1 + if policy_mismatch is not None: + # Excluded by the caller's own policy: never a fallback, even + # if the candidate would also have failed capability checks. + rejected.append(f"{backend.name}: {policy_mismatch}") + continue + capability_incompat = list(capability.incompatibilities(contract)) + if capability_incompat: + rejected.append(f"{backend.name}: " + "; ".join(capability_incompat)) + capability_rejections += 1 continue op = self._get_or_create_backend(backend) @@ -573,7 +585,7 @@ def _logprob_policy_mismatch( policy = requested_backend.lower() if policy == "auto": return None - if policy in {"production", "reference", "deterministic"}: + if policy in IMPLEMENTATION_KINDS: if capability.implementation_kind == policy: return None return ( diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py index 670225f0..40cea4e2 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -11,14 +11,18 @@ 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 @@ -101,10 +105,12 @@ def _declared_tp_backend() -> LogprobBackendCapability: tp_world_sizes=(1, 2, 4), cp_world_sizes=None, supports_vocab_padding=True, - supports_inactive_tokens=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), exports_vocab_lse=True, - deterministic_tp_merge=True, - implementation_kind="deterministic", + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", ) @@ -126,6 +132,7 @@ def test_qwen3_tp2_bf16_contract_is_representable_and_serializable(): "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()) @@ -249,7 +256,7 @@ def test_current_ws1_backend_rejects_strict_tp_contract_without_fallback(): message = str(exc_info.value) assert "TP=2 is unsupported" in message assert "vocab-domain LSE export is unsupported" in message - assert "deterministic TP (max, sumexp) merge is unsupported" in message + assert "determinism_scope=cross_tp_bitwise is unsupported" in message assert "padded-vs-real vocab masking is unsupported" in message @@ -282,11 +289,11 @@ def test_declared_compatible_backend_resolves_and_records_provenance(): OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform ) - result = registry.get_logprob_op(_contract(), requested_backend="deterministic") + 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"] == "deterministic" + 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 @@ -320,11 +327,11 @@ def test_cp_is_a_non_merge_axis_and_cp_agnostic_backends_accept_any_cp_degree(): assert cp_restricted.incompatibilities(cp2_contract) == ("CP=2 is unsupported",) -def test_inactive_tokens_require_declared_backend_support(): - capability = replace(_declared_tp_backend(), supports_inactive_tokens=False) +def test_inactive_tokens_require_explicit_active_mask_support(): + capability = replace(_declared_tp_backend(), mask_modes=frozenset({MaskMode.IGNORE_INDEX})) contract = _contract() - assert "inactive-token (ignore_index) masking is unsupported" in ( + assert "explicit active-token masking is unsupported" in ( capability.incompatibilities(contract) ) @@ -361,7 +368,7 @@ def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform ) - result = registry.get_logprob_op(_contract(), requested_backend="DETERMINISTIC") + 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"): @@ -486,3 +493,92 @@ 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(), 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 From e6dbeefaba51f23176ddd72e70f257b81f070e1a Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Mon, 3 Aug 2026 06:09:14 +0800 Subject: [PATCH 05/10] docs(ws2): drop standalone design doc per review Fold the normative reduction semantics (padded-column masking, fp32 (max, sumexp) merge formulas, the (-inf, 0) identity partial, and the cross-TP tile-structure requirement) into the ReductionSpec and DeterminismScope docstrings, and repoint the runtime-dispatch and batch-invariant-logp doc references at the module. The contract summary moves to the PR description. --- docs/design/runtime-dispatch.md | 3 +- docs/design/ws2-tp-logprob-contract.md | 248 ------------------------- docs/operators/batch-invariant-logp.md | 4 +- rl_engine/kernels/logprob_contract.py | 34 +++- 4 files changed, 35 insertions(+), 254 deletions(-) delete mode 100644 docs/design/ws2-tp-logprob-contract.md diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index 84f29ec3..23c1586a 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -16,7 +16,8 @@ addition to platform priority, this path requires a backend capability descripto 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. -See [WS2 TP-aware logprob contract](ws2-tp-logprob-contract.md). +The contract objects and their normative reduction semantics are documented in +`rl_engine.kernels.logprob_contract`. ## LogP Priority diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md deleted file mode 100644 index cf9179b8..00000000 --- a/docs/design/ws2-tp-logprob-contract.md +++ /dev/null @@ -1,248 +0,0 @@ -# WS2 TP-Aware Logprob Contract - -Status: PR1 contract and dispatch metadata - -Tracking and shared contracts: - -- [#241: TP-aware deterministic logprob](https://github.com/RL-Align/RL-Kernel/issues/241) -- [#83: WS2 roadmap](https://github.com/RL-Align/RL-Kernel/issues/83) -- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) -- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) -- [#116: WS2 tolerance and drift-report format](https://github.com/RL-Align/RL-Kernel/issues/116) -- [Cross-config logprob drift contract](ws2_cross_config_logprob_drift_contract.md) - -## Scope - -This contract describes the logical inputs and deterministic reduction semantics for -selected-token log-probability under vocab-parallel tensor parallelism (TP): - -```text -selected_logp[t] = logits[t, target[t]] - logsumexp_vocab(logits[t, :]) -``` - -Under vocab-parallel TP each rank holds one vocabulary shard, so the vocabulary-wide -`logsumexp` requires a cross-rank reduction. This contract lets runtime dispatch reject a -backend whose numerical semantics do not match the requested layout. - -This PR1 layer does not shard tensors, launch a collective, merge `(max, sumexp)` partial -states, or implement a kernel. The single-GPU harness registration, the deterministic -vocab-parallel TP reference, and the cross-config integration belong to later PRs in #241. - -Context parallelism (CP) 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. - -## Contract Objects - -`rl_engine.kernels.logprob_contract` defines: - -- `LogprobContract`: role, logits dtype, mask, sharding, reduction, output surface, and - LSE export, plus a rank-independent `cross_rank_fingerprint()`; -- `ShardingSpec`: per-rank vocab-shard bounds, padded-vs-real vocabulary, TP/CP rank - metadata, and target-token ownership; -- `MaskSpec`: active-token mask and ignore index; -- `ReductionSpec`: fixed `(max, sumexp)` merge semantics and the requested determinism - scope; -- `LogprobOutputSpec`: the output surface — fp32 selected logprob and fp32 vocab-domain - LSE, replicated across the TP group; -- `LogprobBackendCapability`: the layouts and semantics a backend explicitly supports, - including its mask modes and determinism scopes. - -Construction performs validation immediately. A structurally valid contract means that the -request is complete and internally consistent; it does not mean that an installed backend can -materialize it. - -`ShardingSpec.vocab_shard_bounds` lists every TP rank's half-open `[start, end)` vocab range -indexed by TP rank. The full table is required on every rank: it defines target ownership -and the fixed merge order without any collective, and it makes an incomplete or overlapping -partition a loud construction-time error instead of a silent runtime divergence. -`ShardingSpec.owner_rank(token_id)` resolves the unique owning rank for a real-vocab token -and rejects everything else. - -`padded_vocab_size` is the shard-covered (weight) vocabulary; `real_vocab_size` is the -tokenizer vocabulary. Padding columns occupy `[real_vocab_size, padded_vocab_size)` and must -be excluded from the logsumexp by any conforming implementation. The two sizes are equal -when the vocabulary is unpadded. - -Inactive tokens (prompt, padding, masked-out response positions) are excluded from every -drift aggregate and are exempt from the exactly-one-owner target gather; their targets may -legally hold `ignore_index`. `ignore_index` must not collide with the real vocabulary. - -## Qwen3-8B TP=2 BF16 Example - -```python -from rl_engine.kernels.logprob_contract import ( - LogprobContract, - MaskSpec, - ReductionSpec, - ShardingSpec, -) - -sharding = ShardingSpec( - tp_rank=0, - tp_world_size=2, - vocab_shard_bounds=((0, 76032), (76032, 152064)), - real_vocab_size=151936, - padded_vocab_size=152064, - cp_rank=0, - cp_world_size=2, -) - -contract = LogprobContract( - role="train", - dtype="bf16", - mask=MaskSpec( - num_tokens=8, - active_mask=(False, False, True, True, True, True, True, False), - ignore_index=-100, - ), - sharding=sharding, - reduction=ReductionSpec(), -) -``` - -Each rank owns one contiguous vocab shard; the 128 padding columns at the end of rank 1's -shard are outside the real vocabulary and never contribute to the logsumexp. The two leading -prompt tokens and the trailing padding token are inactive. - -## Reduction Semantics - -The only PR1 reduction contract is: - -```text -partial state: (local_max, local_sumexp), fp32 -merge: max_sumexp -merge_axis: tp_vocab -order: global_vocab_shard_index -transport: all_gather -downcast_at: final_write -engine: in_op_reference -``` - -Every rank first masks every local column whose global id lies in -`[real_vocab_size, padded_vocab_size)` to `-inf`, so padding never contributes to the -logsumexp, then computes `m_l = max(local_logits)` and `s_l = sum(exp(local_logits - m_l))` -in fp32. The partials travel by all-gather (collectives are transport only, never a -numerical reduction), and every rank merges in fixed global vocab-shard index order: - -```text -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: exactly one rank holds -each active token's target column. Downcast happens only at the final write. Because the -merge order is fixed by shard index, the result is deterministic and reproducible at every -TP degree by construction. Cross-degree bitwise equality (TP=2 equal to TP=1, the #241 PR 3 -acceptance target) requires one further condition: the entire reduction must 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 that the TP degree only selects which rank computes -which tiles and never changes the floating-point grouping. A TP-degree-independent -decomposition inside each shard is not sufficient on its own, because shard boundaries -would still group the combines differently across degrees. Providing that global structure -is an obligation of the deterministic reference implementation; a backend without it is -still deterministic per degree, and its cross-degree drift is judged against the #108 -tolerance table instead. The contract expresses this distinction as -`ReductionSpec.determinism_scope`: `cross_tp_bitwise` (the #241 target and the default) -versus `fixed_topology`. A backend declares the scopes it honors in -`LogprobBackendCapability.determinism_scopes`, and dispatch rejects a backend that cannot -honor the requested scope — prose obligations are not enough; the guarantee is part of the -typed contract. - -A shard may lie entirely inside the padded region, and a row's local columns may all be -`-inf` after masking. The identity partial for these cases is defined as -`(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 -evaluating `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. - -The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift -reports follow the #116 format. This contract does not introduce another tolerance table. -The selected-token metric remains the cross-config convention: - -```text -dlogp = training-side recomputed logp - rollout-side old logp -``` - -computed over active response tokens only. - -## Contract-Aware Dispatch - -Legacy callers continue to use `KernelRegistry.get_op()`. WS2 callers use: - -```python -result = kernel_registry.get_logprob_op(contract) -op = result.op -provenance = result.provenance -``` - -Dispatch considers only backends with a `LogprobBackendCapability`. It checks role, dtype, -TP/CP degree, padded-vs-real vocab masking, explicit active-mask support, vocab-domain LSE -export, and the requested determinism scope. An undeclared or incompatible backend is -skipped with an explicit rejection reason; there is no silent fallback. - -`requested_backend` accepts a case-insensitive policy keyword (`auto` | `production` | -`reference`; default `auto`) or an exact, case-sensitive stable backend id. The keywords -select an implementation tier; determinism is not a tier — it is requested through -`ReductionSpec.determinism_scope`, so `requested_backend="deterministic"` raises a loud -error instead of silently matching nothing. Strictness comes from the contract's -capability checks, not from the policy string. A backend id may never shadow a reserved -keyword; capability construction rejects that. The provenance `fallback` flag reports only -capability or load rejections of policy-eligible candidates — a candidate excluded by the -caller's own policy never counts as a fallback, even if it would also have failed -capability checks. - -WS2 dispatch resolves from its own candidate list, seeded from but decoupled from the legacy -`batch_invariant_logp` priority list: registering a TP-vocab backend for WS2 dispatch does -not change what legacy `get_op("batch_invariant_logp")` returns to WS1 callers. - -The current WS1 batch-invariant logp implementations are single-shard (TP=1) references: -they accept full-vocabulary logits with ignore-index masking but carry no vocab-shard -metadata, no padded-vs-real vocab distinction, and no public vocab-domain LSE export. Strict -WS2 requests therefore fail clearly today. The later deterministic vocab-parallel reference -becomes selectable through `KernelRegistry.register_logprob_backend(backend, capability)` -by declaring a capability that truthfully describes those features; no controller branch or -silent fallback is required. - -Successful dispatch provenance records: - -- requested and actual backend ids; -- platform and fallback status; -- prior candidate rejection reasons; -- the complete dispatch-relevant contract, including shard bounds, padded and real vocab - sizes, merge and output semantics, the explicit `cp_is_merge_axis: false` declaration, - and the active-mask digest (`active_mask_sha256`) — the mask's identity without its - per-token payload; -- the selected backend capability descriptor. - -### Distributed dispatch safety - -`get_logprob_op` resolves locally on each rank, so `requested_backend="auto"` is not -distributed-safe on its own: a load failure on one rank can resolve a different backend -than its peers, which for a collective-bearing implementation means divergent numerical -schedules or a deadlock. For `tp_world_size > 1` a caller must either request an exact -backend id or run a preflight agreement before any collective: all-gather the resolved -backend id together with `LogprobContract.cross_rank_fingerprint()` — a rank-independent -hash covering the shard-bounds table, vocab sizes, reduction/output semantics, and the -active-mask digest, excluding rank-local fields — and abort on any mismatch. Implementing -this preflight is an obligation of the #241 PR 3/PR 4 work; the backend invocation -protocol (how the contract and mask reach the implementation) is likewise defined there, -against this contract. - -## Validation - -Contract and dispatch behavior are covered by: - -```bash -python -m pytest tests/test_logprob_contract.py -q -``` - -The tests include Qwen3-8B TP=2 BF16 construction with padded vocab, the TP=1/2/4 sweep -shapes, incomplete/overlapping shard-bound rejection, owner-rank resolution, active-mask and -ignore-index validation, fp32-accumulation and merge-semantics enforcement, undeclared -backend rejection, no incompatible fallback, and JSON-compatible provenance. diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index b0671c40..4bca1f8c 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -64,8 +64,8 @@ remains unchanged. The backends above are single-shard (TP=1) references and do not yet export vocab-domain LSE or carry vocab-shard metadata, so they are declared incompatible with strict WS2 -requests instead of being selected as a silent fallback. See -[WS2 TP-aware logprob contract](../design/ws2-tp-logprob-contract.md). +requests instead of being selected as a silent fallback. The contract objects are +documented in `rl_engine.kernels.logprob_contract`. ## Benchmarks diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 6941c78e..9f6d72ba 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -100,8 +100,12 @@ class DeterminismScope(str, Enum): ``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 (see the design doc); fixed shard-order - merging alone is not sufficient. + 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" @@ -309,7 +313,31 @@ def active_mask_sha256(self) -> str: @dataclass(frozen=True) class ReductionSpec: - """Deterministic TP-vocab ``(max, sumexp)`` merge semantics.""" + """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 From 878ba88a44e5f9ac24d1a38b4536a60c02721602 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Mon, 3 Aug 2026 06:21:30 +0800 Subject: [PATCH 06/10] style(ws2): align comment density with sibling kernel modules Shrink class docstrings toward the attention-contract one-liner style and cut design-rationale comments; the normative reduction semantics stay in the ReductionSpec and DeterminismScope docstrings. --- rl_engine/kernels/logprob_contract.py | 65 +++++++++------------------ rl_engine/kernels/registry.py | 19 +++----- 2 files changed, 26 insertions(+), 58 deletions(-) diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 9f6d72ba..b0bcd7ab 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -29,16 +29,10 @@ _EnumT = TypeVar("_EnumT", bound=Enum) -# Dispatch policy keywords accepted by KernelRegistry.get_logprob_op; a stable -# backend id must never shadow one of these, or it becomes unselectable by id. -# "deterministic" stays reserved even though it is no longer a policy: -# determinism is expressed through DeterminismScope, and requesting it as a -# policy is a loud error rather than a silent id mismatch. +# 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"}) -# Implementation tiers a backend can declare. Determinism is deliberately a -# separate axis (DeterminismScope): a backend can be a deterministic reference, -# a deterministic production implementation, or a non-deterministic production -# implementation. +# Backend tiers; determinism is a separate axis (DeterminismScope). IMPLEMENTATION_KINDS = frozenset({"production", "reference"}) @@ -58,12 +52,7 @@ class LogprobDType(str, Enum): class LogprobMerge(str, Enum): - """Merge primitive for per-shard partial states. - - Every rank contributes ``(local_max, local_sumexp)`` computed in the - accumulation dtype; the merged result is - ``M = max(m_l)``, ``S = sum(s_l * exp(m_l - M))``, ``LSE = M + log(S)``. - """ + """Merge primitive for per-shard ``(local_max, local_sumexp)`` partials.""" MAX_SUMEXP = "max_sumexp" @@ -115,11 +104,9 @@ class DeterminismScope(str, Enum): class MaskMode(str, Enum): """How a backend consumes inactive-token information. - ``explicit_active_mask``: the backend honors an arbitrary active-token - mask. ``ignore_index``: the backend only recognizes inactive tokens whose - target id equals ``ignore_index``. The contract permits inactive targets - that do NOT hold ``ignore_index``, so an ignore-index-only backend cannot - serve a contract with inactive tokens. + 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" @@ -161,15 +148,11 @@ 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 TP rank. The full table is required on every rank: - it defines target-token ownership and the fixed global-shard-index merge - order without any collective, and makes an incomplete partition a loud - construction-time error instead of a silent runtime divergence. - - ``padded_vocab_size`` is the shard-covered (weight) vocabulary; - ``real_vocab_size`` is the tokenizer vocabulary. Padding columns occupy - ``[real_vocab_size, padded_vocab_size)`` and must be excluded from the - logsumexp by any conforming implementation. + 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 @@ -267,11 +250,10 @@ def owner_rank(self, token_id: int) -> int: @dataclass(frozen=True) class MaskSpec: - """Active-token ownership for one logprob invocation. + """Active-token mask and ignore index for one logprob invocation. - Inactive tokens are excluded from every drift aggregate and are exempt - from the exactly-one-owner target gather; their targets may legally hold - ``ignore_index``. + Inactive tokens are excluded from drift aggregates and from the + single-owner target gather; their targets may legally hold ``ignore_index``. """ num_tokens: int @@ -377,12 +359,8 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class LogprobOutputSpec: - """Output surface every conforming backend must produce. - - Selected logprob and vocab-domain LSE are fp32 and replicated across the - TP group; the ``downcast_at: final_write`` rule applies to any consumer - downcast after these outputs, never inside the reduction. - """ + """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 @@ -466,10 +444,8 @@ def to_dict(self) -> dict[str, Any]: "determinism_scope": self.reduction.determinism_scope.value, "cp_is_merge_axis": False, } - # The per-token mask is deliberately summarized: provenance exists for - # logging/serialization and the raw mask would dominate its size. The - # digest keeps the mask *identity* observable, so two masks with the - # same active count still produce distinguishable provenance. + # 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, @@ -614,8 +590,7 @@ def incompatibilities(self, contract: LogprobContract) -> tuple[str, ...]: contract.mask.active_token_count != contract.mask.num_tokens and MaskMode.EXPLICIT_ACTIVE_MASK not in self.mask_modes ): - # The contract does not require inactive targets to hold - # ignore_index, so ignore-index-only masking is insufficient. + # 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") diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index b17acf3e..7093b320 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -175,13 +175,9 @@ def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() - # These descriptors report what the existing WS1 batch-invariant logp - # implementations actually support: single-shard (TP=1) logits with - # ignore-index masking, no vocab-shard metadata, no padded-vs-real - # vocab distinction, and no public vocab-domain LSE export. A strict - # WS2 request is rejected with explicit reasons until the deterministic - # vocab-parallel TP reference backend lands (issue #241 PR 3) instead - # of silently selecting an incompatible fallback. + # 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 = { @@ -336,12 +332,9 @@ def __init__(self): self._adjust_priority_for_hardware() self._adjust_priority_from_env() - # WS2 contract-aware dispatch owns its candidate list. It is seeded - # from the legacy batch_invariant_logp priority (after hardware/env - # adjustments) but deliberately decoupled afterwards: registering a - # TP-vocab backend for WS2 dispatch must not change what legacy - # get_op("batch_invariant_logp") returns to WS1 callers, and vice - # versa. + # 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() From a63bea2dc0269cba82fca44e36cf3941de5aae33 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Wed, 5 Aug 2026 16:05:57 +0800 Subject: [PATCH 07/10] init vocab parallel logp --- .github/workflows/ci.yml | 3 + docs/operators/batch-invariant-logp.md | 36 +- .../ops/pytorch/loss/vocab_parallel_logp.py | 408 +++++++++++++++++ rl_engine/kernels/registry.py | 27 ++ tests/test_logprob_contract.py | 32 ++ tests/test_operator_inputs.py | 26 ++ tests/test_vocab_parallel_logp.py | 426 ++++++++++++++++++ 7 files changed, 948 insertions(+), 10 deletions(-) create mode 100644 rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py create mode 100644 tests/test_vocab_parallel_logp.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 28cdb58d..2d9ba91e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -79,6 +79,9 @@ jobs: - 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/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index 4bca1f8c..5068c33f 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -54,18 +54,31 @@ 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). -### WS2 TP-aware dispatch +## Tensor Parallel -WS2 distributed callers use a separate contract-aware entry point, -`kernel_registry.get_logprob_op(contract)`. It validates explicit vocab-shard ownership, -padded-vs-real vocab metadata, active-token masking, and fixed `(max, sumexp)` merge -semantics before selecting a backend. Legacy `get_op("batch_invariant_logp")` behavior -remains unchanged. +`VocabParallelLogprobOp` +(`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) +**TP=1, TP=2, and TP=4 produce bit-identical results.** -The backends above are single-shard (TP=1) references and do not yet export vocab-domain -LSE or carry vocab-shard metadata, so they are declared incompatible with strict WS2 -requests instead of being selected as a silent fallback. The contract objects are -documented in `rl_engine.kernels.logprob_contract`. +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 @@ -237,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/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 7093b320..213aa7a2 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -85,6 +85,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_BATCH_INVARIANT_LOGP_SM90 = ( "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op" ) + # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) + PYTORCH_VOCAB_PARALLEL_LOGP = ( + "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" + ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -350,6 +354,29 @@ def __init__(self): for platform, candidates in self._logprob_candidates.items() } + # deterministic vocab-parallel TP logprob reference. + ws2_tp_logprob_capability = LogprobBackendCapability( + backend_id="pytorch-vocab-parallel-logp-ws2", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=None, + cp_world_sizes=None, + supports_vocab_padding=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), + exports_vocab_lse=True, + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", + ) + for ws2_platform in self._priority_map: + self.register_logprob_backend( + OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP, + ws2_tp_logprob_capability, + platform=ws2_platform, + prepend=True, + ) + def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() if rocm_attn_backend in {"flash_attn", "flash-attn", "flash_attention"}: diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py index 40cea4e2..a961f312 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -247,8 +247,20 @@ def test_ignore_index_must_not_collide_with_the_real_vocabulary(): 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()) @@ -262,6 +274,7 @@ def test_current_ws1_backend_rejects_strict_tp_contract_without_fallback(): 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: @@ -272,6 +285,25 @@ def test_current_ws1_backend_rejects_padded_vocab_even_at_tp1(): 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() diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 4f742734..2080f6fd 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -40,6 +40,7 @@ def _args(**overrides): "logp", "linear_logp", "batch_invariant_logp", + "vocab_parallel_logp", "rope", "silu", "swiglu", @@ -72,6 +73,31 @@ def test_constant_batch_invariant_logp_inputs_match_operator_contract(): assert torch.equal(inputs["target_ids"], torch.full((1, 2), 3, dtype=torch.long)) +def test_constant_vocab_parallel_logp_inputs_match_operator_contract(): + args = _args(input_mode="constant", constant_value=0.5, token_value=3) + inputs = make_operator_inputs("vocab_parallel_logp", args, torch.float32, torch.device("cpu")) + + # vocab=17 rounds up to padded=20 with 4 tiles; tokens flatten to batch*seq. + assert torch.equal(inputs["local_logits"], torch.full((2, 20), 0.5)) + assert torch.equal(inputs["target_ids"], torch.full((2,), 3, dtype=torch.long)) + assert inputs["contract"].sharding.real_vocab_size == 17 + assert inputs["contract"].sharding.padded_vocab_size == 20 + assert inputs["num_vocab_tiles"] == 4 + assert operator_shape_name("vocab_parallel_logp", args) == "2x20" + + +def test_vocab_parallel_logp_inputs_run_through_the_operator(): + from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import VocabParallelLogprobOp + + args = _args(input_mode="random", seed=7) + inputs = make_operator_inputs("vocab_parallel_logp", args, torch.float32, torch.device("cpu")) + + logp, lse = VocabParallelLogprobOp()(**inputs) + assert logp.shape == inputs["target_ids"].shape + assert logp.dtype == torch.float32 and lse.dtype == torch.float32 + assert torch.isfinite(logp).all() and torch.isfinite(lse).all() + + def test_random_logp_inputs_are_seeded(): args = _args(input_mode="random", seed=7) first = make_operator_inputs("logp", args, torch.float32, torch.device("cpu")) diff --git a/tests/test_vocab_parallel_logp.py b/tests/test_vocab_parallel_logp.py new file mode 100644 index 00000000..510aec70 --- /dev/null +++ b/tests/test_vocab_parallel_logp.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic vocab-parallel TP logprob reference tests (issue #241 PR3). + +Bit-level determinism assertions compare raw bit patterns via +``tensor.view(torch.int32)`` rather than ``torch.equal``: value equality +treats ``-0.0 == 0.0`` as equal and ``NaN != NaN`` as different, neither of +which is what a bitwise claim means. +""" + +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}[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 + + +# --------------------------------------------------------------------------- +# Multi-rank gloo tests (spawn pattern from tests/test_linear_logp.py) +# --------------------------------------------------------------------------- + + +def _gloo_available() -> bool: + return torch.distributed.is_available() and torch.distributed.is_gloo_available() + + +requires_gloo = pytest.mark.skipif( + not _gloo_available(), reason="requires torch.distributed with the gloo backend" +) + +_WORLD_SIZE = 4 +_UNEVEN_BOUNDS = ((0, 4), (4, 16), (16, 24), (24, 32)) # tile-aligned (tile=4) + + +def _tp_worker(rank, world_size, init_method, result_queue, scenario): + import torch.distributed as dist + + torch.set_num_threads(1) + try: + dist.init_process_group( + backend="gloo", init_method=init_method, rank=rank, world_size=world_size + ) + dtype = torch.bfloat16 if scenario == "bf16" else torch.float32 + dtype_name = "bf16" if scenario == "bf16" else "fp32" + bounds = _UNEVEN_BOUNDS if scenario == "uneven" else _even_bounds(PADDED_VOCAB, world_size) + logits, targets = _inputs(dtype=dtype) + start, end = bounds[rank] + + op = VocabParallelLogprobOp() + tiles = 16 if scenario == "preflight" and rank == 0 else NUM_TILES + contract_tp = _contract( + tp_rank=rank, tp_world_size=world_size, bounds=bounds, dtype=dtype_name + ) + + if scenario == "preflight": + try: + op( + logits[:, start:end].contiguous().clone(), + targets, + contract=contract_tp, + tp_group=dist.group.WORLD, + num_vocab_tiles=tiles, + ) + result_queue.put({"ok": False, "rank": rank, "traceback": "no error raised"}) + except LogprobContractError: + result_queue.put({"ok": True, "rank": rank}) + return + + shard = logits[:, start:end].contiguous().clone().requires_grad_(True) + logp_tp, lse_tp = op( + shard, + targets, + contract=contract_tp, + tp_group=dist.group.WORLD, + num_vocab_tiles=NUM_TILES, + ) + (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() + + # In-process TP=1 run of the same op on the full logits: the cross-TP + # bitwise claim is TP=n output == TP=1 output, bit for bit. + full = logits.clone().requires_grad_(True) + contract_tp1 = _contract(dtype=dtype_name) + logp_one, lse_one = op(full, targets, contract=contract_tp1, num_vocab_tiles=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]), + "logp": logp_tp.detach().float(), + "lse": lse_tp.detach().float(), + } + ) + except Exception: + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _run_gloo_scenario(scenario): + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmpdir: + init_method = (Path(tmpdir) / "gloo_init").as_uri() + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=_tp_worker, + args=(rank, _WORLD_SIZE, init_method, result_queue, scenario), + ) + 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=60)) + except queue.Empty: + for process in processes: + if process.is_alive(): + process.terminate() + pytest.fail("timed out waiting for vocab-parallel gloo workers") + finally: + for process in processes: + process.join(timeout=10) + 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 + + +@requires_gloo +@pytest.mark.parametrize("scenario", ["even", "uneven", "bf16"]) +def test_tp4_bitwise_identical_to_tp1(scenario): + results = _run_gloo_scenario(scenario) + for result in results: + assert result["logp_bits_match"], f"rank {result['rank']} logp bits differ from TP=1" + assert result["lse_bits_match"], f"rank {result['rank']} lse bits differ from TP=1" + assert result["grad_bits_match"], f"rank {result['rank']} grad bits differ from TP=1" + # Outputs are replicated: every rank must hold identical bits. + for other in results[1:]: + assert _bitwise_equal(results[0]["logp"], other["logp"]) + assert _bitwise_equal(results[0]["lse"], other["lse"]) + + +@requires_gloo +def test_preflight_rejects_mismatched_num_vocab_tiles(): + _run_gloo_scenario("preflight") From 4fcdc30cded0b25498874a966b38ac75fe31cb13 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Wed, 5 Aug 2026 16:10:26 +0800 Subject: [PATCH 08/10] init vocab parallel logp --- tests/test_operator_inputs.py | 26 -------------------------- 1 file changed, 26 deletions(-) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 2080f6fd..4f742734 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -40,7 +40,6 @@ def _args(**overrides): "logp", "linear_logp", "batch_invariant_logp", - "vocab_parallel_logp", "rope", "silu", "swiglu", @@ -73,31 +72,6 @@ def test_constant_batch_invariant_logp_inputs_match_operator_contract(): assert torch.equal(inputs["target_ids"], torch.full((1, 2), 3, dtype=torch.long)) -def test_constant_vocab_parallel_logp_inputs_match_operator_contract(): - args = _args(input_mode="constant", constant_value=0.5, token_value=3) - inputs = make_operator_inputs("vocab_parallel_logp", args, torch.float32, torch.device("cpu")) - - # vocab=17 rounds up to padded=20 with 4 tiles; tokens flatten to batch*seq. - assert torch.equal(inputs["local_logits"], torch.full((2, 20), 0.5)) - assert torch.equal(inputs["target_ids"], torch.full((2,), 3, dtype=torch.long)) - assert inputs["contract"].sharding.real_vocab_size == 17 - assert inputs["contract"].sharding.padded_vocab_size == 20 - assert inputs["num_vocab_tiles"] == 4 - assert operator_shape_name("vocab_parallel_logp", args) == "2x20" - - -def test_vocab_parallel_logp_inputs_run_through_the_operator(): - from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import VocabParallelLogprobOp - - args = _args(input_mode="random", seed=7) - inputs = make_operator_inputs("vocab_parallel_logp", args, torch.float32, torch.device("cpu")) - - logp, lse = VocabParallelLogprobOp()(**inputs) - assert logp.shape == inputs["target_ids"].shape - assert logp.dtype == torch.float32 and lse.dtype == torch.float32 - assert torch.isfinite(logp).all() and torch.isfinite(lse).all() - - def test_random_logp_inputs_are_seeded(): args = _args(input_mode="random", seed=7) first = make_operator_inputs("logp", args, torch.float32, torch.device("cpu")) From 6ffade40974e1fba5a1c5fbbd7ad9e6d6c9fe9ad Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Wed, 5 Aug 2026 21:13:04 +0800 Subject: [PATCH 09/10] adding cross tp testing --- tests/test_vocab_parallel_logp.py | 257 +++++++++++++++++++++--------- 1 file changed, 184 insertions(+), 73 deletions(-) diff --git a/tests/test_vocab_parallel_logp.py b/tests/test_vocab_parallel_logp.py index 510aec70..87da2ef0 100644 --- a/tests/test_vocab_parallel_logp.py +++ b/tests/test_vocab_parallel_logp.py @@ -1,13 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Deterministic vocab-parallel TP logprob reference tests (issue #241 PR3). - -Bit-level determinism assertions compare raw bit patterns via -``tensor.view(torch.int32)`` rather than ``torch.equal``: value equality -treats ``-0.0 == 0.0`` as equal and ``NaN != NaN`` as different, neither of -which is what a bitwise claim means. -""" +"""Deterministic vocab-parallel TP logprob reference tests""" from __future__ import annotations @@ -86,7 +80,11 @@ def _inputs(dtype=torch.float32, seed: int = 2026): def _bits(tensor: torch.Tensor) -> torch.Tensor: - view_dtype = {torch.float32: torch.int32, torch.bfloat16: torch.int16}[tensor.dtype] + view_dtype = { + torch.float32: torch.int32, + torch.bfloat16: torch.int16, + torch.float16: torch.int16, + }[tensor.dtype] return tensor.contiguous().view(view_dtype) @@ -283,72 +281,143 @@ def test_dispatch_resolves_reference_and_leaves_legacy_untouched(): assert OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP not in candidates -# --------------------------------------------------------------------------- -# Multi-rank gloo tests (spawn pattern from tests/test_linear_logp.py) -# --------------------------------------------------------------------------- +# 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 _gloo_available() -> bool: - return torch.distributed.is_available() and torch.distributed.is_gloo_available() +def _cuda_device_count() -> int: + return torch.cuda.device_count() if torch.cuda.is_available() else 0 -requires_gloo = pytest.mark.skipif( - not _gloo_available(), reason="requires torch.distributed with the gloo backend" -) +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, + ) -_WORLD_SIZE = 4 -_UNEVEN_BOUNDS = ((0, 4), (4, 16), (16, 24), (24, 32)) # tile-aligned (tile=4) +def _tp_inputs(device, dtype, seed: int = 2026): + """Identical logits and targets on every rank, seeded on CPU.""" -def _tp_worker(rank, world_size, init_method, result_queue, scenario): + 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 - torch.set_num_threads(1) try: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) dist.init_process_group( - backend="gloo", init_method=init_method, rank=rank, world_size=world_size + backend="nccl", init_method=init_method, rank=rank, world_size=world_size ) - dtype = torch.bfloat16 if scenario == "bf16" else torch.float32 - dtype_name = "bf16" if scenario == "bf16" else "fp32" - bounds = _UNEVEN_BOUNDS if scenario == "uneven" else _even_bounds(PADDED_VOCAB, world_size) - logits, targets = _inputs(dtype=dtype) - start, end = bounds[rank] - + dtype = TP_DTYPES[dtype_name] op = VocabParallelLogprobOp() - tiles = 16 if scenario == "preflight" and rank == 0 else NUM_TILES - contract_tp = _contract( - tp_rank=rank, tp_world_size=world_size, bounds=bounds, dtype=dtype_name - ) - - if scenario == "preflight": + 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=contract_tp, + 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: - result_queue.put({"ok": True, "rank": rank}) + 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=contract_tp, + contract=tp_contract, tp_group=dist.group.WORLD, - num_vocab_tiles=NUM_TILES, + num_vocab_tiles=TP_NUM_TILES, ) (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() - # In-process TP=1 run of the same op on the full logits: the cross-TP - # bitwise claim is TP=n output == TP=1 output, bit for bit. + # 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) - contract_tp1 = _contract(dtype=dtype_name) - logp_one, lse_one = op(full, targets, contract=contract_tp1, num_vocab_tiles=NUM_TILES) + 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( @@ -358,45 +427,50 @@ def _tp_worker(rank, world_size, init_method, result_queue, scenario): "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]), - "logp": logp_tp.detach().float(), - "lse": lse_tp.detach().float(), + "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: + except Exception: # pragma: no cover - forwarded to the parent process result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) raise finally: - if torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() + import torch.distributed as dist + if dist.is_initialized(): + dist.destroy_process_group() -def _run_gloo_scenario(scenario): + +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) / "gloo_init").as_uri() + init_method = (Path(tmpdir) / "nccl_init").as_uri() result_queue = ctx.Queue() processes = [ ctx.Process( - target=_tp_worker, - args=(rank, _WORLD_SIZE, init_method, result_queue, scenario), + target=_nccl_worker, + args=(rank, world_size, init_method, result_queue, scenario, uneven, dtype_name), ) - for rank in range(_WORLD_SIZE) + for rank in range(world_size) ] results = [] try: for process in processes: process.start() - for _ in range(_WORLD_SIZE): + for _ in range(world_size): try: - results.append(result_queue.get(timeout=60)) + results.append(result_queue.get(timeout=_SPAWN_TIMEOUT_S)) except queue.Empty: for process in processes: if process.is_alive(): process.terminate() - pytest.fail("timed out waiting for vocab-parallel gloo workers") + pytest.fail(f"timed out waiting for NCCL workers (scenario={scenario})") finally: for process in processes: - process.join(timeout=10) + process.join(timeout=30) if process.is_alive(): process.terminate() results.sort(key=lambda item: item["rank"]) @@ -407,20 +481,57 @@ def _run_gloo_scenario(scenario): return results -@requires_gloo -@pytest.mark.parametrize("scenario", ["even", "uneven", "bf16"]) -def test_tp4_bitwise_identical_to_tp1(scenario): - results = _run_gloo_scenario(scenario) - for result in results: - assert result["logp_bits_match"], f"rank {result['rank']} logp bits differ from TP=1" - assert result["lse_bits_match"], f"rank {result['rank']} lse bits differ from TP=1" - assert result["grad_bits_match"], f"rank {result['rank']} grad bits differ from TP=1" - # Outputs are replicated: every rank must hold identical bits. - for other in results[1:]: - assert _bitwise_equal(results[0]["logp"], other["logp"]) - assert _bitwise_equal(results[0]["lse"], other["lse"]) - - -@requires_gloo -def test_preflight_rejects_mismatched_num_vocab_tiles(): - _run_gloo_scenario("preflight") +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"] From 3627fcde7c805deb10238109eaf5d21e8b935b1c Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Wed, 12 Aug 2026 11:34:36 +0800 Subject: [PATCH 10/10] ws2 deterministic grpo loss PR5 --- docs/operators/grpo-loss.md | 146 ++- rl_engine/kernels/loss_contract.py | 784 ++++++++++++++++ .../ops/pytorch/loss/distributed_grpo_loss.py | 447 +++++++++ rl_engine/kernels/registry.py | 173 ++++ tests/test_distributed_grpo_loss.py | 866 ++++++++++++++++++ tests/test_grpo_loss_contract.py | 392 ++++++++ 6 files changed, 2802 insertions(+), 6 deletions(-) create mode 100644 rl_engine/kernels/loss_contract.py create mode 100644 rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py create mode 100644 tests/test_distributed_grpo_loss.py create mode 100644 tests/test_grpo_loss_contract.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/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/registry.py b/rl_engine/kernels/registry.py index 213aa7a2..8391dbbb 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -19,6 +19,15 @@ 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 @@ -89,6 +98,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): 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" @@ -377,6 +390,34 @@ def __init__(self): 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"}: @@ -597,6 +638,138 @@ def get_logprob_op( 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: + 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. + """ + + 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, 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"