From b7ffb3ad496535909215e43b0bb8dfea446aa0a5 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Sun, 2 Aug 2026 20:03:13 +0800 Subject: [PATCH 01/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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 934bc5be978b2291b374766228edf2c9ccc67079 Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 4 Aug 2026 17:04:05 +0800 Subject: [PATCH 07/25] feat: add single-gpu logprob comparison harness --- docs/design/ws2-logprob-single-gpu-harness.md | 94 +++++ .../ops/cuda/loss/batch_invariant_logp.py | 45 +++ .../ops/pytorch/loss/batch_invariant_logp.py | 28 +- .../ops/triton/loss/batch_invariant_logp.py | 88 ++++- rl_engine/testing/__init__.py | 20 + rl_engine/testing/logprob_comparison.py | 361 ++++++++++++++++++ scripts/compare_logprob.py | 96 +++++ tests/test_logprob_comparison.py | 214 +++++++++++ 8 files changed, 925 insertions(+), 21 deletions(-) create mode 100644 docs/design/ws2-logprob-single-gpu-harness.md create mode 100644 rl_engine/testing/logprob_comparison.py create mode 100644 scripts/compare_logprob.py create mode 100644 tests/test_logprob_comparison.py diff --git a/docs/design/ws2-logprob-single-gpu-harness.md b/docs/design/ws2-logprob-single-gpu-harness.md new file mode 100644 index 00000000..7440d465 --- /dev/null +++ b/docs/design/ws2-logprob-single-gpu-harness.md @@ -0,0 +1,94 @@ +# WS2 Single-GPU Logprob Comparison Harness + +This harness is the TP=1 registration and regression guard for issue #241. It compares +selected-token logprob implementations before any distributed communication is introduced. + +## Contract + +For each logical token row, every backend returns direct FP32 values: + +```text +LSE = logsumexp(logits[..., vocab]) +logp = selected_logit - LSE +``` + +The harness uses the merged WS1 batch-invariant PyTorch implementation as its reference. +Reference logp is obtained through the unchanged production call, while reference LSE is +obtained through the diagnostic entry point. The TP=1 PyTorch candidate follows the same +core computation and must be bitwise equal. This is a regression guard, not new +distributed mathematics. + +LSE drift is reported over every logical token row. Selected-token dlogp drift is reported +only over active response/action tokens. Both reports contain max, mean, p95, p99, and the +number of compared values. + +## Exact Backend Selection + +Supported backend names are: + +- `pytorch` +- `triton` +- `cuda-sm90` + +The comparison path does not use registry fallback. An explicitly requested backend must +run exactly or raise `LogprobBackendUnavailable`. In particular, `cuda-sm90` requires a +compiled SM90 extension, Hopper hardware, BF16/FP32 logits, and a compatible vocab row +stride. The production operator may retain its normal fallback behavior outside the +harness. + +Each backend exposes a diagnostic-only `forward_with_lse` method. Existing production +calls remain unchanged: + +```text +op(logits, target_ids) -> logp +op.forward_with_lse(logits, target_ids) -> (logp, lse) +``` + +## Usage + +CPU TP=1 regression guard: + +```bash +python scripts/compare_logprob.py \ + --candidate pytorch \ + --device cpu \ + --dtype fp32 \ + --batch 2 \ + --seq 16 \ + --vocab 257 +``` + +GPU comparison: + +```bash +python scripts/compare_logprob.py \ + --candidate triton \ + --candidate cuda-sm90 \ + --device cuda \ + --dtype bf16 \ + --batch 2 \ + --seq 16 \ + --vocab 151936 +``` + +The command prints a structured JSON report containing input dtype/shape, active-token +count, TP world size, communication mode, requested and actual backends, direct-LSE +provenance, bitwise logp status, and LSE/dlogp drift statistics. + +## Scope Boundary + +This harness is intentionally single-GPU and records `tp_world=1` and +`communication=none`. It does not implement vocab-shard metadata, all-gather transport, +fixed-order cross-rank LSE merging, CP reconstruction, or distributed artifacts. Those +belong to the later PR3 and PR4 work in issue #241. + +## Tests + +```bash +python -m pytest tests/test_logprob_comparison.py -q +``` + +The focused tests cover bitwise TP=1 regression, direct LSE identity, active-token-only +percentiles, zero active tokens, invalid ignore-index usage, structured serialization, +generic operator-harness registration, exact GPU backend diagnostics, and fail-closed +backend provenance. diff --git a/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py index a68781d9..aa62b6ef 100644 --- a/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py @@ -161,3 +161,48 @@ def apply( ) return _BatchInvariantLogpSM90Function.apply(logits, target_ids, ignore_index) + + def forward_with_lse( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the exact SM90 path and return its direct FP32 logprob/LSE outputs. + + Unlike the production ``apply`` method, this diagnostic entry point never + falls back to Triton or PyTorch, so comparison provenance stays truthful. + """ + if logits.dim() < 2: + raise ValueError( + f"logits must be at least 2-D ([*lead, V]), got shape {tuple(logits.shape)}" + ) + if logits.shape[:-1] != target_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + if not _sm90_supported(logits): + raise RuntimeError( + "exact cuda-sm90 logprob diagnostics require Hopper, CUDA BF16/FP32 logits, " + "and a 16-byte-aligned vocab row stride; fallback is disabled" + ) + if validate: + vocab_size = logits.size(-1) + valid_targets = target_ids.reshape(-1) + valid_targets = valid_targets[valid_targets != ignore_index] + if valid_targets.numel() and ( + (valid_targets < 0).any() or (valid_targets >= vocab_size).any() + ): + bad = valid_targets[(valid_targets < 0) | (valid_targets >= vocab_size)] + raise ValueError( + f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}" + ) + + lead_shape = logits.shape[:-1] + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() + logp, lse = _C.batch_invariant_logp_sm90(logits_2d, target_1d, int(ignore_index)) + return logp.reshape(lead_shape), lse.reshape(lead_shape) diff --git a/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py index 4ac8bd37..80e08b5f 100644 --- a/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py @@ -45,24 +45,42 @@ def apply( logits_2d = logits.reshape(-1, vocab_size).float() target_1d = target_ids.reshape(-1).to(logits.device, dtype=torch.long) - selected_logp = self._row_wise_selected_logprob( + selected_logp, _ = self._row_wise_selected_logprob_with_lse( logits_2d, target_1d, ignore_index=ignore_index, validate=validate ) return selected_logp.reshape(lead_shape) + def forward_with_lse( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return selected logprob and the FP32 vocab-domain LSE for diagnostics.""" + self._validate_shapes(logits, target_ids) + lead_shape = logits.shape[:-1] + logits_2d = logits.reshape(-1, logits.size(-1)).float() + target_1d = target_ids.reshape(-1).to(logits.device, dtype=torch.long) + logp, lse = self._row_wise_selected_logprob_with_lse( + logits_2d, target_1d, ignore_index=ignore_index, validate=validate + ) + return logp.reshape(lead_shape), lse.reshape(lead_shape) + # ---------------------------------------------------------------------- # # Core Computation # ---------------------------------------------------------------------- # @staticmethod - def _row_wise_selected_logprob( + def _row_wise_selected_logprob_with_lse( logits_2d: torch.Tensor, target_1d: torch.Tensor, *, ignore_index: int, validate: bool = True, - ) -> torch.Tensor: - """Per-row selected logprob with locked reduction order. + ) -> tuple[torch.Tensor, torch.Tensor]: + """Per-row selected logprob and LSE with locked reduction order. The three reduction steps (max, sum-exp, gather) operate on each row independently. PyTorch's ``max(dim=-1)`` and ``sum(dim=-1)`` iterate @@ -104,7 +122,7 @@ def _row_wise_selected_logprob( selected_logp = selected_logp.where(valid_mask, torch.zeros_like(selected_logp)) - return selected_logp + return selected_logp, log_sum_exp # ---------------------------------------------------------------------- # # Helper diff --git a/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py index 66b99757..804341f3 100644 --- a/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py @@ -10,6 +10,27 @@ _BLOCK_V: int = 1024 +def _launch_batch_invariant_logp( + logits_2d: torch.Tensor, target_1d: torch.Tensor, ignore_index: int +) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = logits_2d.shape[0] + vocab_size = logits_2d.shape[1] + output = torch.empty(num_tokens, device=logits_2d.device, dtype=torch.float32) + lse = torch.empty(num_tokens, device=logits_2d.device, dtype=torch.float32) + _batch_invariant_logp_kernel[(num_tokens,)]( + logits_2d, + target_1d, + output, + lse, + num_tokens, + vocab_size, + logits_2d.stride(0), + ignore_index=ignore_index, + BLOCK_V=_BLOCK_V, + ) + return output, lse + + @triton.jit def _batch_invariant_logp_kernel( logits_ptr, # logits [N, V] @@ -126,22 +147,7 @@ def forward(ctx, logits, target_ids, ignore_index): logits_2d = logits.reshape(-1, vocab_size).contiguous() target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() - num_tokens = logits_2d.shape[0] - output = torch.empty(num_tokens, device=logits.device, dtype=torch.float32) - lse = torch.empty(num_tokens, device=logits.device, dtype=torch.float32) - - grid = (num_tokens,) - _batch_invariant_logp_kernel[grid]( - logits_2d, - target_1d, - output, - lse, - num_tokens, - vocab_size, - logits_2d.stride(0), - ignore_index=ignore_index, - BLOCK_V=_BLOCK_V, - ) + output, lse = _launch_batch_invariant_logp(logits_2d, target_1d, ignore_index) ctx.save_for_backward(logits_2d, target_1d, lse) ctx.ignore_index = ignore_index @@ -237,3 +243,53 @@ def apply( ) return _BatchInvariantLogpFunction.apply(logits, target_ids, ignore_index) + + def forward_with_lse( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return direct FP32 logprob/LSE outputs without an autograd wrapper.""" + self._validate_inputs(logits, target_ids, ignore_index=ignore_index, validate=validate) + lead_shape = logits.shape[:-1] + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() + logp, lse = _launch_batch_invariant_logp(logits_2d, target_1d, ignore_index) + return logp.reshape(lead_shape), lse.reshape(lead_shape) + + @staticmethod + def _validate_inputs( + logits: torch.Tensor, + target_ids: torch.Tensor, + *, + ignore_index: int, + validate: bool, + ) -> None: + if logits.device.type not in ("cuda", "xpu", "hip"): + raise RuntimeError( + "TritonBatchInvariantLogpOp requires a GPU tensor " + f"(CUDA / ROCm / XPU), got device '{logits.device}'." + ) + if logits.dim() < 2: + raise ValueError( + f"logits must be at least 2-D ([*lead, V]), got shape {tuple(logits.shape)}" + ) + if logits.shape[:-1] != target_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + if validate: + vocab_size = logits.size(-1) + valid_targets = target_ids.reshape(-1) + valid_targets = valid_targets[valid_targets != ignore_index] + if valid_targets.numel() and ( + (valid_targets < 0).any() or (valid_targets >= vocab_size).any() + ): + bad = valid_targets[(valid_targets < 0) | (valid_targets >= vocab_size)] + raise ValueError( + f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}" + ) diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 42be8c1b..cc11e625 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -3,6 +3,17 @@ """Testing helpers for RL-shaped kernel validation.""" +from .logprob_comparison import ( + DriftStats, + LogprobBackendUnavailable, + LogprobCandidate, + LogprobComparisonInputs, + LogprobComparisonReport, + LogprobPathDrift, + LogprobPathResult, + compare_single_gpu_logprob, + make_logprob_candidate, +) from .reference_ops import ( active_token_count, compute_policy_ratio, @@ -15,10 +26,19 @@ from .rl_batch import SyntheticRLKernelBatch, make_synthetic_rl_kernel_batch __all__ = [ + "DriftStats", + "LogprobBackendUnavailable", + "LogprobCandidate", + "LogprobComparisonInputs", + "LogprobComparisonReport", + "LogprobPathDrift", + "LogprobPathResult", "SyntheticRLKernelBatch", "active_token_count", + "compare_single_gpu_logprob", "compute_policy_ratio", "compute_reference_kl", + "make_logprob_candidate", "make_synthetic_rl_kernel_batch", "masked_mean", "masked_sum", diff --git a/rl_engine/testing/logprob_comparison.py b/rl_engine/testing/logprob_comparison.py new file mode 100644 index 00000000..601dbced --- /dev/null +++ b/rl_engine/testing/logprob_comparison.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Single-GPU WS2 selected-logprob cross-implementation harness.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Sequence + +import torch + + +class LogprobBackendUnavailable(RuntimeError): + """Raised when an explicitly requested comparison backend cannot run exactly.""" + + +@dataclass(frozen=True) +class LogprobComparisonInputs: + """Logical TP=1 inputs shared by every comparison path.""" + + logits: torch.Tensor + target_ids: torch.Tensor + active_token_mask: torch.Tensor | None = None + ignore_index: int = -100 + + +@dataclass(frozen=True) +class LogprobPathResult: + """Direct selected-logprob and vocab-LSE outputs from one backend.""" + + name: str + logp: torch.Tensor + lse: torch.Tensor + provenance: dict[str, Any] + + +@dataclass(frozen=True) +class LogprobCandidate: + """One exact backend materialization used by the harness.""" + + name: str + requested_backend: str + actual_backend: str + fn: Callable[[torch.Tensor, torch.Tensor, int], tuple[torch.Tensor, torch.Tensor]] + provenance: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class DriftStats: + """Absolute drift statistics over a declared comparison population.""" + + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + def to_dict(self) -> dict[str, Any]: + return { + "max_abs": self.max_abs, + "mean_abs": self.mean_abs, + "p95_abs": self.p95_abs, + "p99_abs": self.p99_abs, + "active_count": self.active_count, + } + + +@dataclass(frozen=True) +class LogprobPathDrift: + """Candidate-vs-reference LSE and active-token dlogp drift.""" + + candidate_name: str + lse: DriftStats + dlogp: DriftStats + bitwise_logp: bool + provenance: dict[str, Any] + + def to_dict(self) -> dict[str, Any]: + return { + "candidate_name": self.candidate_name, + "lse": self.lse.to_dict(), + "dlogp": self.dlogp.to_dict(), + "bitwise_logp": self.bitwise_logp, + "provenance": self.provenance, + } + + +@dataclass(frozen=True) +class LogprobComparisonReport: + """Structured single-GPU report consumed by later WS2 integration.""" + + reference_name: str + drifts: tuple[LogprobPathDrift, ...] + input_provenance: dict[str, Any] + + def to_dict(self) -> dict[str, Any]: + return { + "reference_name": self.reference_name, + "drifts": [drift.to_dict() for drift in self.drifts], + "input_provenance": self.input_provenance, + } + + +def make_logprob_candidate(backend: str) -> LogprobCandidate: + """Materialize an exact built-in backend without registry fallback.""" + + normalized = backend.strip().lower().replace("_", "-") + if normalized in {"pytorch", "native"}: + from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( + NativeBatchInvariantLogpOp, + ) + + op = NativeBatchInvariantLogpOp() + actual = "pytorch" + elif normalized == "triton": + try: + from rl_engine.kernels.ops.triton.loss.batch_invariant_logp import ( + TritonBatchInvariantLogpOp, + ) + + op = TritonBatchInvariantLogpOp() + except Exception as exc: + raise LogprobBackendUnavailable(f"triton backend is unavailable: {exc}") from exc + actual = "triton" + elif normalized in {"cuda-sm90", "sm90"}: + try: + from rl_engine.kernels.ops.cuda.loss.batch_invariant_logp import ( + BatchInvariantLogpSM90Op, + ) + + op = BatchInvariantLogpSM90Op() + except Exception as exc: + raise LogprobBackendUnavailable(f"cuda-sm90 backend is unavailable: {exc}") from exc + actual = "cuda-sm90" + else: + raise ValueError( + f"unsupported logprob comparison backend {backend!r}; " + "expected pytorch, triton, or cuda-sm90" + ) + + diagnostic = getattr(op, "forward_with_lse", None) + if not callable(diagnostic): + raise LogprobBackendUnavailable( + f"backend {normalized!r} does not expose the required direct LSE diagnostic" + ) + + def run( + logits: torch.Tensor, target_ids: torch.Tensor, ignore_index: int + ) -> tuple[torch.Tensor, torch.Tensor]: + try: + return diagnostic(logits, target_ids, ignore_index=ignore_index, validate=True) + except (RuntimeError, NotImplementedError, OSError) as exc: + raise LogprobBackendUnavailable( + f"exact backend {normalized!r} cannot execute this input: {exc}" + ) from exc + + return LogprobCandidate( + name=f"{actual}-batch-invariant-logp", + requested_backend=actual, + actual_backend=actual, + fn=run, + provenance={ + "requested_alias": normalized, + "implementation": f"{type(op).__module__}.{type(op).__qualname__}", + }, + ) + + +def compare_single_gpu_logprob( + inputs: LogprobComparisonInputs, + *, + candidates: Sequence[str | LogprobCandidate] = ("pytorch",), +) -> LogprobComparisonReport: + """Compare exact TP=1 implementations against the WS1 deterministic path.""" + + active_mask, effective_targets = _validate_inputs(inputs) + reference = _run_ws1_reference(inputs.logits, effective_targets, inputs.ignore_index) + + drifts = tuple( + _compare_path( + _run_candidate( + ( + candidate + if isinstance(candidate, LogprobCandidate) + else make_logprob_candidate(candidate) + ), + inputs.logits, + effective_targets, + inputs.ignore_index, + ), + reference, + active_mask, + ) + for candidate in candidates + ) + return LogprobComparisonReport( + reference_name=reference.name, + drifts=drifts, + input_provenance={ + "device": str(inputs.logits.device), + "input_dtype": str(inputs.logits.dtype), + "output_dtype": str(reference.logp.dtype), + "shape": list(inputs.logits.shape), + "ignore_index": inputs.ignore_index, + "active_token_count": int(active_mask.sum().item()), + "tp_world": 1, + "communication": "none", + }, + ) + + +def _run_ws1_reference( + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int, +) -> LogprobPathResult: + """Run the existing deterministic logp path and its direct-LSE diagnostic.""" + from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( + NativeBatchInvariantLogpOp, + ) + + op = NativeBatchInvariantLogpOp() + logp = op(logits, target_ids, ignore_index=ignore_index, validate=True) + _, lse = op.forward_with_lse( + logits, target_ids, ignore_index=ignore_index, validate=True + ) + return LogprobPathResult( + name="pytorch-batch-invariant-logp", + logp=logp.detach(), + lse=lse.detach(), + provenance={ + "requested_backend": "pytorch", + "actual_backend": "pytorch", + "tp_world": 1, + "communication": "none", + "logp_source": "production", + "lse_source": "direct", + }, + ) + + +def _run_candidate( + candidate: LogprobCandidate, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int, +) -> LogprobPathResult: + if candidate.requested_backend != candidate.actual_backend: + raise LogprobBackendUnavailable( + f"requested backend {candidate.requested_backend!r} materialized as " + f"{candidate.actual_backend!r}; silent fallback is forbidden" + ) + logp, lse = candidate.fn(logits, target_ids, ignore_index) + expected_shape = logits.shape[:-1] + for name, value in (("logp", logp), ("lse", lse)): + if not isinstance(value, torch.Tensor): + raise TypeError(f"candidate {candidate.name!r} {name} must be a tensor") + if value.shape != expected_shape: + raise ValueError( + f"candidate {candidate.name!r} {name} shape {tuple(value.shape)} " + f"does not match {tuple(expected_shape)}" + ) + if value.dtype != torch.float32: + raise ValueError(f"candidate {candidate.name!r} {name} must be FP32") + return LogprobPathResult( + name=candidate.name, + logp=logp.detach(), + lse=lse.detach(), + provenance={ + "requested_backend": candidate.requested_backend, + "actual_backend": candidate.actual_backend, + "tp_world": 1, + "communication": "none", + "lse_source": "direct", + **candidate.provenance, + }, + ) + + +def _compare_path( + candidate: LogprobPathResult, + reference: LogprobPathResult, + active_mask: torch.Tensor, +) -> LogprobPathDrift: + return LogprobPathDrift( + candidate_name=candidate.name, + lse=_drift_stats(candidate.lse, reference.lse), + dlogp=_drift_stats(candidate.logp, reference.logp, mask=active_mask), + bitwise_logp=torch.equal(candidate.logp, reference.logp), + provenance=candidate.provenance, + ) + + +def _drift_stats( + candidate: torch.Tensor, + reference: torch.Tensor, + *, + mask: torch.Tensor | None = None, +) -> DriftStats: + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match reference shape " + f"{tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs() + values = diff.reshape(-1) if mask is None else diff[mask.to(device=diff.device)] + count = int(values.numel()) + if count == 0: + return DriftStats(0.0, 0.0, 0.0, 0.0, 0) + return DriftStats( + max_abs=float(values.max().item()), + mean_abs=float(values.mean().item()), + p95_abs=float(torch.quantile(values, 0.95).item()), + p99_abs=float(torch.quantile(values, 0.99).item()), + active_count=count, + ) + + +def _validate_inputs( + inputs: LogprobComparisonInputs, +) -> tuple[torch.Tensor, torch.Tensor]: + if inputs.logits.dim() < 2: + raise ValueError("logits must be at least 2-D [*lead, vocab]") + if inputs.logits.shape[:-1] != inputs.target_ids.shape: + raise ValueError("target_ids shape must match logits leading shape") + if not inputs.logits.is_floating_point(): + raise ValueError("logits must be floating point") + + if inputs.active_token_mask is None: + active = inputs.target_ids != inputs.ignore_index + else: + if inputs.active_token_mask.shape != inputs.target_ids.shape: + raise ValueError("active_token_mask shape must match target_ids") + if inputs.active_token_mask.dtype != torch.bool: + raise ValueError("active_token_mask must be bool") + active = inputs.active_token_mask.to(device=inputs.target_ids.device) + if bool(((inputs.target_ids == inputs.ignore_index) & active).any().item()): + raise ValueError("active target_ids cannot equal ignore_index") + + effective = inputs.target_ids.to(device=inputs.logits.device, dtype=torch.long).clone() + active = active.to(device=inputs.logits.device, dtype=torch.bool) + effective.masked_fill_(~active, inputs.ignore_index) + valid = effective[active] + vocab_size = inputs.logits.size(-1) + if valid.numel() and ((valid < 0).any() or (valid >= vocab_size).any()): + raise ValueError(f"active target_ids must be in [0, {vocab_size})") + return active, effective + + +__all__ = [ + "DriftStats", + "LogprobBackendUnavailable", + "LogprobCandidate", + "LogprobComparisonInputs", + "LogprobComparisonReport", + "LogprobPathDrift", + "LogprobPathResult", + "compare_single_gpu_logprob", + "make_logprob_candidate", +] diff --git a/scripts/compare_logprob.py b/scripts/compare_logprob.py new file mode 100644 index 00000000..a5feb570 --- /dev/null +++ b/scripts/compare_logprob.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys + +import torch + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.testing import ( # noqa: E402 + LogprobComparisonInputs, + compare_single_gpu_logprob, +) + + +def _dtype(name: str) -> torch.dtype: + return { + "fp32": torch.float32, + "bf16": torch.bfloat16, + "fp16": torch.float16, + }[name] + + +def _device(name: str) -> torch.device: + if name == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device(name) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run the WS2 TP=1 selected-logprob/LSE comparison harness." + ) + parser.add_argument( + "--candidate", + action="append", + choices=("pytorch", "triton", "cuda-sm90"), + help="Exact backend to compare. Repeat for multiple backends; defaults to pytorch.", + ) + parser.add_argument("--device", default="auto") + parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq", type=int, default=16) + parser.add_argument("--vocab", type=int, default=257) + parser.add_argument("--prompt-tokens", type=int, default=8) + parser.add_argument("--seed", type=int, default=123) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + device = _device(args.device) + if args.batch < 1 or args.seq < 1 or args.vocab < 1: + raise ValueError("batch, seq, and vocab must be positive") + if not 0 <= args.prompt_tokens <= args.seq: + raise ValueError("prompt-tokens must be in [0, seq]") + + generator = torch.Generator(device=device).manual_seed(args.seed) + logits = torch.randn( + args.batch, + args.seq, + args.vocab, + generator=generator, + device=device, + dtype=_dtype(args.dtype), + ) + target_ids = torch.randint( + 0, + args.vocab, + (args.batch, args.seq), + generator=generator, + device=device, + ) + active_mask = torch.ones((args.batch, args.seq), device=device, dtype=torch.bool) + active_mask[:, : args.prompt_tokens] = False + report = compare_single_gpu_logprob( + LogprobComparisonInputs( + logits=logits, + target_ids=target_ids, + active_token_mask=active_mask, + ), + candidates=tuple(args.candidate or ("pytorch",)), + ) + print(json.dumps(report.to_dict(), indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tests/test_logprob_comparison.py b/tests/test_logprob_comparison.py new file mode 100644 index 00000000..f5051f5d --- /dev/null +++ b/tests/test_logprob_comparison.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import argparse + +import pytest +import torch + +from rl_engine.kernels.gtest import run_operator_suite +from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case +from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( + NativeBatchInvariantLogpOp, +) +from rl_engine.testing.logprob_comparison import ( + LogprobBackendUnavailable, + LogprobCandidate, + LogprobComparisonInputs, + compare_single_gpu_logprob, + make_logprob_candidate, +) +from scripts.compare_logprob import _device + + +def _inputs() -> LogprobComparisonInputs: + generator = torch.Generator().manual_seed(17) + logits = torch.randn(2, 4, 257, generator=generator, dtype=torch.float32) + target_ids = torch.tensor([[3, 5, 7, 11], [13, 17, 19, 23]]) + active = torch.tensor([[False, False, True, True], [False, True, True, True]]) + return LogprobComparisonInputs(logits, target_ids, active_token_mask=active) + + +def test_single_gpu_pytorch_path_is_bitwise_regression_guard(): + report = compare_single_gpu_logprob(_inputs(), candidates=("pytorch",)) + + assert report.reference_name == "pytorch-batch-invariant-logp" + assert len(report.drifts) == 1 + drift = report.drifts[0] + assert drift.bitwise_logp + assert drift.lse.max_abs == 0.0 + assert drift.dlogp.max_abs == 0.0 + assert drift.dlogp.active_count == 5 + assert drift.provenance["requested_backend"] == "pytorch" + assert drift.provenance["actual_backend"] == "pytorch" + assert drift.provenance["lse_source"] == "direct" + assert report.input_provenance["tp_world"] == 1 + assert report.input_provenance["communication"] == "none" + + +def test_report_serializes_lse_and_active_token_percentiles(): + inputs = _inputs() + reference = make_logprob_candidate("pytorch") + + def shifted(logits, target_ids, ignore_index): + logp, lse = reference.fn(logits, target_ids, ignore_index) + logp = logp.clone() + logp[0, 0] += 100.0 # inactive and therefore excluded from dlogp + logp[0, 2] += 1.0 + lse = lse + torch.arange(lse.numel(), dtype=lse.dtype).reshape_as(lse) * 0.1 + return logp, lse + + candidate = LogprobCandidate( + name="shifted", + requested_backend="shifted", + actual_backend="shifted", + fn=shifted, + ) + report = compare_single_gpu_logprob(inputs, candidates=(candidate,)) + payload = report.to_dict() + drift = payload["drifts"][0] + + assert drift["dlogp"]["active_count"] == 5 + assert drift["dlogp"]["max_abs"] == pytest.approx(1.0) + assert drift["dlogp"]["p95_abs"] == pytest.approx(0.8) + assert drift["dlogp"]["p99_abs"] == pytest.approx(0.96) + assert drift["lse"]["active_count"] == 8 + assert drift["lse"]["p99_abs"] == pytest.approx(0.693, abs=1e-5) + + +def test_all_inactive_tokens_produce_zero_dlogp_statistics(): + inputs = _inputs() + inputs = LogprobComparisonInputs( + inputs.logits, + inputs.target_ids, + active_token_mask=torch.zeros_like(inputs.target_ids, dtype=torch.bool), + ) + drift = compare_single_gpu_logprob(inputs).drifts[0] + + assert drift.dlogp.active_count == 0 + assert drift.dlogp.max_abs == 0.0 + assert drift.dlogp.p95_abs == 0.0 + assert drift.lse.active_count == inputs.target_ids.numel() + + +def test_explicit_backend_mismatch_fails_closed(): + native = make_logprob_candidate("pytorch") + disguised = LogprobCandidate( + name="fallback", + requested_backend="cuda-sm90", + actual_backend="pytorch", + fn=native.fn, + ) + + with pytest.raises(LogprobBackendUnavailable, match="silent fallback is forbidden"): + compare_single_gpu_logprob(_inputs(), candidates=(disguised,)) + + +def test_active_ignore_index_is_rejected(): + inputs = _inputs() + targets = inputs.target_ids.clone() + targets[0, 2] = -100 + + with pytest.raises(ValueError, match="active target_ids cannot equal ignore_index"): + compare_single_gpu_logprob( + LogprobComparisonInputs( + inputs.logits, + targets, + active_token_mask=inputs.active_token_mask, + ) + ) + + +def test_native_diagnostic_lse_satisfies_selected_logit_identity(): + inputs = _inputs() + candidate = make_logprob_candidate("pytorch") + effective = inputs.target_ids.masked_fill(~inputs.active_token_mask, -100) + logp, lse = candidate.fn(inputs.logits, effective, -100) + production_logp = NativeBatchInvariantLogpOp()( + inputs.logits, effective, ignore_index=-100, validate=True + ) + safe_targets = effective.masked_fill(~inputs.active_token_mask, 0) + selected = torch.gather(inputs.logits, -1, safe_targets.unsqueeze(-1)).squeeze(-1) + + assert torch.equal(logp, production_logp) + assert torch.equal(logp[inputs.active_token_mask], (selected - lse)[inputs.active_token_mask]) + + +def test_unsupported_backend_name_is_rejected(): + with pytest.raises(ValueError, match="unsupported logprob comparison backend"): + make_logprob_candidate("unknown") + + +def test_cli_auto_device_resolves_without_constructing_auto(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + assert _device("auto") == torch.device("cpu") + + +def test_operator_comparison_specs_register_batch_invariant_logp(): + args = argparse.Namespace( + op="batch_invariant_logp", + candidate="pytorch", + arch_key=None, + batch=2, + seq=4, + vocab=17, + seed=7, + input_mode="random", + constant_value=0.5, + token_value=3, + normalized_dim=128, + k_dim=16, + n_dim=32, + theta=1.0e6, + eps=1.0e-6, + ) + + case = make_operator_case(args, torch.float32, torch.device("cpu")) + candidate = make_candidate(args) + report = run_operator_suite( + "batch_invariant_logp", candidates=[candidate], cases=[case] + ) + + assert report.passed + assert report.candidates[0].cases[0].op_class == "logprob" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_triton_diagnostic_path_reports_direct_lse(): + try: + candidate = make_logprob_candidate("triton") + except LogprobBackendUnavailable as exc: + pytest.skip(str(exc)) + logits = torch.randn(4, 1024, device="cuda", dtype=torch.bfloat16) + targets = torch.tensor([0, 17, 511, 1023], device="cuda") + try: + report = compare_single_gpu_logprob( + LogprobComparisonInputs(logits, targets), candidates=(candidate,) + ) + except LogprobBackendUnavailable as exc: + if isinstance(exc.__cause__, PermissionError): + pytest.skip(str(exc)) + raise + + assert report.drifts[0].provenance["actual_backend"] == "triton" + assert report.drifts[0].provenance["lse_source"] == "direct" + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9, + reason="Hopper CUDA device required", +) +def test_sm90_diagnostic_path_reports_direct_lse_without_fallback(): + try: + candidate = make_logprob_candidate("cuda-sm90") + except LogprobBackendUnavailable as exc: + pytest.skip(str(exc)) + logits = torch.randn(4, 1024, device="cuda", dtype=torch.bfloat16) + targets = torch.tensor([0, 17, 511, 1023], device="cuda") + report = compare_single_gpu_logprob( + LogprobComparisonInputs(logits, targets), candidates=(candidate,) + ) + + assert report.drifts[0].provenance["actual_backend"] == "cuda-sm90" + assert report.drifts[0].provenance["lse_source"] == "direct" From b69426d28bcbdc002a33baed9522e18bde3311e9 Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 4 Aug 2026 20:22:12 +0800 Subject: [PATCH 08/25] fix: keep logprob CLI stdout machine readable --- docs/design/ws2-logprob-single-gpu-harness.md | 7 +++-- scripts/compare_logprob.py | 10 +++++++ tests/test_logprob_comparison.py | 30 ++++++++++++++++++- 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/docs/design/ws2-logprob-single-gpu-harness.md b/docs/design/ws2-logprob-single-gpu-harness.md index 7440d465..c5b869b0 100644 --- a/docs/design/ws2-logprob-single-gpu-harness.md +++ b/docs/design/ws2-logprob-single-gpu-harness.md @@ -71,9 +71,10 @@ python scripts/compare_logprob.py \ --vocab 151936 ``` -The command prints a structured JSON report containing input dtype/shape, active-token -count, TP world size, communication mode, requested and actual backends, direct-LSE -provenance, bitwise logp status, and LSE/dlogp drift statistics. +The command prints a structured JSON report to stdout containing input dtype/shape, +active-token count, TP world size, communication mode, requested and actual backends, +direct-LSE provenance, bitwise logp status, and LSE/dlogp drift statistics. Backend +diagnostic logs are routed to stderr so redirected stdout remains valid JSON. ## Scope Boundary diff --git a/scripts/compare_logprob.py b/scripts/compare_logprob.py index a5feb570..6b42d4ea 100644 --- a/scripts/compare_logprob.py +++ b/scripts/compare_logprob.py @@ -6,6 +6,7 @@ import argparse import json +import logging import pathlib import sys @@ -19,6 +20,7 @@ LogprobComparisonInputs, compare_single_gpu_logprob, ) +from rl_engine.utils.logger import logger # noqa: E402 def _dtype(name: str) -> torch.dtype: @@ -35,6 +37,13 @@ def _device(name: str) -> torch.device: return torch.device(name) +def _route_rl_kernel_logs_to_stderr() -> None: + """Keep stdout machine-readable while preserving backend diagnostics.""" + for handler in logger.handlers: + if isinstance(handler, logging.StreamHandler): + handler.setStream(sys.stderr) + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Run the WS2 TP=1 selected-logprob/LSE comparison harness." @@ -56,6 +65,7 @@ def parse_args() -> argparse.Namespace: def main() -> None: + _route_rl_kernel_logs_to_stderr() args = parse_args() device = _device(args.device) if args.batch < 1 or args.seq < 1 or args.vocab < 1: diff --git a/tests/test_logprob_comparison.py b/tests/test_logprob_comparison.py index f5051f5d..67c1d422 100644 --- a/tests/test_logprob_comparison.py +++ b/tests/test_logprob_comparison.py @@ -2,6 +2,10 @@ # Copyright (c) 2026 RL-Kernel Contributors import argparse +import io +import json +import logging +import sys import pytest import torch @@ -18,7 +22,8 @@ compare_single_gpu_logprob, make_logprob_candidate, ) -from scripts.compare_logprob import _device +from rl_engine.utils.logger import logger +from scripts.compare_logprob import _device, _route_rl_kernel_logs_to_stderr def _inputs() -> LogprobComparisonInputs: @@ -145,6 +150,29 @@ def test_cli_auto_device_resolves_without_constructing_auto(monkeypatch): assert _device("auto") == torch.device("cpu") +def test_cli_routes_rl_kernel_logs_to_stderr_for_machine_readable_stdout(monkeypatch): + stdout = io.StringIO() + stderr = io.StringIO() + original_streams = [ + (handler, handler.stream) + for handler in logger.handlers + if isinstance(handler, logging.StreamHandler) + ] + monkeypatch.setattr(sys, "stdout", stdout) + monkeypatch.setattr(sys, "stderr", stderr) + + try: + _route_rl_kernel_logs_to_stderr() + logger.info("test backend diagnostic") + print(json.dumps({"ok": True})) + finally: + for handler, stream in original_streams: + handler.setStream(stream) + + assert json.loads(stdout.getvalue()) == {"ok": True} + assert "test backend diagnostic" in stderr.getvalue() + + def test_operator_comparison_specs_register_batch_invariant_logp(): args = argparse.Namespace( op="batch_invariant_logp", From 0efcfe12883482247d7544af50fbb1c0adf2a049 Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 4 Aug 2026 20:39:54 +0800 Subject: [PATCH 09/25] refactor: simplify logprob comparison harness --- rl_engine/testing/__init__.py | 6 - rl_engine/testing/logprob_comparison.py | 177 +++++++----------------- 2 files changed, 53 insertions(+), 130 deletions(-) diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index cc11e625..51759abd 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -4,13 +4,10 @@ """Testing helpers for RL-shaped kernel validation.""" from .logprob_comparison import ( - DriftStats, LogprobBackendUnavailable, LogprobCandidate, LogprobComparisonInputs, LogprobComparisonReport, - LogprobPathDrift, - LogprobPathResult, compare_single_gpu_logprob, make_logprob_candidate, ) @@ -26,13 +23,10 @@ from .rl_batch import SyntheticRLKernelBatch, make_synthetic_rl_kernel_batch __all__ = [ - "DriftStats", "LogprobBackendUnavailable", "LogprobCandidate", "LogprobComparisonInputs", "LogprobComparisonReport", - "LogprobPathDrift", - "LogprobPathResult", "SyntheticRLKernelBatch", "active_token_count", "compare_single_gpu_logprob", diff --git a/rl_engine/testing/logprob_comparison.py b/rl_engine/testing/logprob_comparison.py index 601dbced..ad2ef76c 100644 --- a/rl_engine/testing/logprob_comparison.py +++ b/rl_engine/testing/logprob_comparison.py @@ -1,44 +1,31 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Single-GPU WS2 selected-logprob cross-implementation harness.""" +"""Single-GPU selected-logprob comparison.""" from __future__ import annotations -from dataclasses import dataclass, field -from typing import Any, Callable, Sequence +from collections.abc import Callable, Sequence +from dataclasses import asdict, dataclass, field +from typing import Any import torch class LogprobBackendUnavailable(RuntimeError): - """Raised when an explicitly requested comparison backend cannot run exactly.""" + pass @dataclass(frozen=True) class LogprobComparisonInputs: - """Logical TP=1 inputs shared by every comparison path.""" - logits: torch.Tensor target_ids: torch.Tensor active_token_mask: torch.Tensor | None = None ignore_index: int = -100 -@dataclass(frozen=True) -class LogprobPathResult: - """Direct selected-logprob and vocab-LSE outputs from one backend.""" - - name: str - logp: torch.Tensor - lse: torch.Tensor - provenance: dict[str, Any] - - @dataclass(frozen=True) class LogprobCandidate: - """One exact backend materialization used by the harness.""" - name: str requested_backend: str actual_backend: str @@ -47,64 +34,34 @@ class LogprobCandidate: @dataclass(frozen=True) -class DriftStats: - """Absolute drift statistics over a declared comparison population.""" - +class _DriftStats: max_abs: float mean_abs: float p95_abs: float p99_abs: float active_count: int - def to_dict(self) -> dict[str, Any]: - return { - "max_abs": self.max_abs, - "mean_abs": self.mean_abs, - "p95_abs": self.p95_abs, - "p99_abs": self.p99_abs, - "active_count": self.active_count, - } - @dataclass(frozen=True) -class LogprobPathDrift: - """Candidate-vs-reference LSE and active-token dlogp drift.""" - +class _LogprobPathDrift: candidate_name: str - lse: DriftStats - dlogp: DriftStats + lse: _DriftStats + dlogp: _DriftStats bitwise_logp: bool provenance: dict[str, Any] - def to_dict(self) -> dict[str, Any]: - return { - "candidate_name": self.candidate_name, - "lse": self.lse.to_dict(), - "dlogp": self.dlogp.to_dict(), - "bitwise_logp": self.bitwise_logp, - "provenance": self.provenance, - } - @dataclass(frozen=True) class LogprobComparisonReport: - """Structured single-GPU report consumed by later WS2 integration.""" - reference_name: str - drifts: tuple[LogprobPathDrift, ...] + drifts: tuple[_LogprobPathDrift, ...] input_provenance: dict[str, Any] def to_dict(self) -> dict[str, Any]: - return { - "reference_name": self.reference_name, - "drifts": [drift.to_dict() for drift in self.drifts], - "input_provenance": self.input_provenance, - } + return asdict(self) def make_logprob_candidate(backend: str) -> LogprobCandidate: - """Materialize an exact built-in backend without registry fallback.""" - normalized = backend.strip().lower().replace("_", "-") if normalized in {"pytorch", "native"}: from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( @@ -172,35 +129,38 @@ def compare_single_gpu_logprob( *, candidates: Sequence[str | LogprobCandidate] = ("pytorch",), ) -> LogprobComparisonReport: - """Compare exact TP=1 implementations against the WS1 deterministic path.""" - active_mask, effective_targets = _validate_inputs(inputs) - reference = _run_ws1_reference(inputs.logits, effective_targets, inputs.ignore_index) - - drifts = tuple( - _compare_path( - _run_candidate( - ( - candidate - if isinstance(candidate, LogprobCandidate) - else make_logprob_candidate(candidate) - ), - inputs.logits, - effective_targets, - inputs.ignore_index, - ), - reference, - active_mask, - ) - for candidate in candidates + reference_logp, reference_lse = _run_ws1_reference( + inputs.logits, effective_targets, inputs.ignore_index ) + + drifts = [] + for candidate in candidates: + if isinstance(candidate, str): + candidate = make_logprob_candidate(candidate) + logp, lse = _run_candidate( + candidate, + inputs.logits, + effective_targets, + inputs.ignore_index, + ) + drifts.append( + _LogprobPathDrift( + candidate_name=candidate.name, + lse=_drift_stats(lse, reference_lse), + dlogp=_drift_stats(logp, reference_logp, mask=active_mask), + bitwise_logp=torch.equal(logp, reference_logp), + provenance=_candidate_provenance(candidate), + ) + ) + return LogprobComparisonReport( - reference_name=reference.name, - drifts=drifts, + reference_name="pytorch-batch-invariant-logp", + drifts=tuple(drifts), input_provenance={ "device": str(inputs.logits.device), "input_dtype": str(inputs.logits.dtype), - "output_dtype": str(reference.logp.dtype), + "output_dtype": str(reference_logp.dtype), "shape": list(inputs.logits.shape), "ignore_index": inputs.ignore_index, "active_token_count": int(active_mask.sum().item()), @@ -214,8 +174,7 @@ def _run_ws1_reference( logits: torch.Tensor, target_ids: torch.Tensor, ignore_index: int, -) -> LogprobPathResult: - """Run the existing deterministic logp path and its direct-LSE diagnostic.""" +) -> tuple[torch.Tensor, torch.Tensor]: from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( NativeBatchInvariantLogpOp, ) @@ -225,19 +184,7 @@ def _run_ws1_reference( _, lse = op.forward_with_lse( logits, target_ids, ignore_index=ignore_index, validate=True ) - return LogprobPathResult( - name="pytorch-batch-invariant-logp", - logp=logp.detach(), - lse=lse.detach(), - provenance={ - "requested_backend": "pytorch", - "actual_backend": "pytorch", - "tp_world": 1, - "communication": "none", - "logp_source": "production", - "lse_source": "direct", - }, - ) + return logp.detach(), lse.detach() def _run_candidate( @@ -245,7 +192,7 @@ def _run_candidate( logits: torch.Tensor, target_ids: torch.Tensor, ignore_index: int, -) -> LogprobPathResult: +) -> tuple[torch.Tensor, torch.Tensor]: if candidate.requested_backend != candidate.actual_backend: raise LogprobBackendUnavailable( f"requested backend {candidate.requested_backend!r} materialized as " @@ -263,33 +210,18 @@ def _run_candidate( ) if value.dtype != torch.float32: raise ValueError(f"candidate {candidate.name!r} {name} must be FP32") - return LogprobPathResult( - name=candidate.name, - logp=logp.detach(), - lse=lse.detach(), - provenance={ - "requested_backend": candidate.requested_backend, - "actual_backend": candidate.actual_backend, - "tp_world": 1, - "communication": "none", - "lse_source": "direct", - **candidate.provenance, - }, - ) + return logp.detach(), lse.detach() -def _compare_path( - candidate: LogprobPathResult, - reference: LogprobPathResult, - active_mask: torch.Tensor, -) -> LogprobPathDrift: - return LogprobPathDrift( - candidate_name=candidate.name, - lse=_drift_stats(candidate.lse, reference.lse), - dlogp=_drift_stats(candidate.logp, reference.logp, mask=active_mask), - bitwise_logp=torch.equal(candidate.logp, reference.logp), - provenance=candidate.provenance, - ) +def _candidate_provenance(candidate: LogprobCandidate) -> dict[str, Any]: + return { + "requested_backend": candidate.requested_backend, + "actual_backend": candidate.actual_backend, + "tp_world": 1, + "communication": "none", + "lse_source": "direct", + **candidate.provenance, + } def _drift_stats( @@ -297,7 +229,7 @@ def _drift_stats( reference: torch.Tensor, *, mask: torch.Tensor | None = None, -) -> DriftStats: +) -> _DriftStats: if candidate.shape != reference.shape: raise ValueError( f"candidate shape {tuple(candidate.shape)} must match reference shape " @@ -307,8 +239,8 @@ def _drift_stats( values = diff.reshape(-1) if mask is None else diff[mask.to(device=diff.device)] count = int(values.numel()) if count == 0: - return DriftStats(0.0, 0.0, 0.0, 0.0, 0) - return DriftStats( + return _DriftStats(0.0, 0.0, 0.0, 0.0, 0) + return _DriftStats( max_abs=float(values.max().item()), mean_abs=float(values.mean().item()), p95_abs=float(torch.quantile(values, 0.95).item()), @@ -349,13 +281,10 @@ def _validate_inputs( __all__ = [ - "DriftStats", "LogprobBackendUnavailable", "LogprobCandidate", "LogprobComparisonInputs", "LogprobComparisonReport", - "LogprobPathDrift", - "LogprobPathResult", "compare_single_gpu_logprob", "make_logprob_candidate", ] From 115d86c7a1ac9054ee2990ddb82a61869c4d8449 Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 4 Aug 2026 21:30:27 +0800 Subject: [PATCH 10/25] docs: document SM90 logprob validation --- docs/design/ws2-logprob-sm90-validation.md | 134 +++++++++++++++++++++ 1 file changed, 134 insertions(+) create mode 100644 docs/design/ws2-logprob-sm90-validation.md diff --git a/docs/design/ws2-logprob-sm90-validation.md b/docs/design/ws2-logprob-sm90-validation.md new file mode 100644 index 00000000..e94cb52d --- /dev/null +++ b/docs/design/ws2-logprob-sm90-validation.md @@ -0,0 +1,134 @@ +# WS2 Logprob PR2 SM90 Validation + +This document records the Hopper SM90 validation procedure for the PR2 single-GPU +logprob comparison harness from issue #241. It is a validation note for maintainers; +the cloud setup wrapper used during development is intentionally kept outside the +repository. + +## Prerequisites + +The validation host must provide: + +- Python 3.10 or newer; +- CUDA-enabled PyTorch; +- an NVIDIA Hopper GPU with compute capability 9.0, such as H100, H800, or H200; +- `nvidia-smi` and `nvcc`; +- a CUDA development environment capable of compiling the RL-Kernel extension. + +The CUDA version reported by `nvcc` must match `torch.version.cuda`. A runtime-only +image is insufficient because it normally does not include the CUDA compiler. + +## Build + +Activate an environment containing the repository dependencies and a CUDA-enabled +PyTorch installation, then build the editable extension with SM90 enabled: + +```bash +export FORCE_CUDA=1 +export KERNEL_ALIGN_FORCE_SM90=1 +export TORCH_CUDA_ARCH_LIST="9.0+PTX" +export MAX_JOBS=2 + +python -m pip install --no-build-isolation --no-deps -e . +``` + +Verify the extension and SM90 symbol after the build. Import PyTorch first so its +runtime libraries are available to the extension loader: + +```bash +python - <<'PY' +import torch +from rl_engine import _C + +print("torch:", torch.__version__) +print("torch CUDA:", torch.version.cuda) +print("extension:", _C.__file__) +print("SM90 symbol:", hasattr(_C, "batch_invariant_logp_sm90")) +PY +``` + +The final line must report `SM90 symbol: True`. + +## Validation commands + +Run the focused PR2 tests and the complete batch-invariant logprob suite: + +```bash +python -m pytest \ + tests/test_logprob_comparison.py \ + tests/test_operator_inputs.py \ + tests/test_op_checks.py -q + +python -m pytest tests/test_batch_invariant_logp.py -q +``` + +Run the two explicit SM90 comparisons: + +```bash +python scripts/compare_logprob.py \ + --candidate cuda-sm90 \ + --device cuda \ + --dtype bf16 \ + --batch 2 \ + --seq 8 \ + --vocab 1024 \ + --prompt-tokens 3 \ + --seed 7 + +python scripts/compare_logprob.py \ + --candidate cuda-sm90 \ + --device cuda \ + --dtype bf16 \ + --batch 2 \ + --seq 16 \ + --vocab 151936 \ + --prompt-tokens 8 \ + --seed 241 +``` + +The comparison command writes the JSON report to stdout. Diagnostic log messages are +written to stderr so stdout can be redirected directly to a `.json` file. + +## Expected report + +The report must identify the requested and actual backend as `cuda-sm90`, use the +`BatchInvariantLogpSM90Op` implementation, and record: + +```text +tp_world=1 +communication=none +lse_source=direct +``` + +LSE drift is measured over all logical token rows. Selected-logprob drift is measured +only over active response/action tokens. Each drift section includes maximum, mean, +p95, p99, and active-count values. + +## Validation result + +The procedure was validated on: + +```text +GPU: NVIDIA H800 PCIe +Compute capability: 9.0 +Python: 3.11.15 +PyTorch: 2.11.0+cu128 +CUDA toolkit / nvcc: 12.8 +Triton: 3.6.0 +``` + +Results: + +```text +PR2 focused tests: 41 passed +Complete batch-invariant logprob suite: 67 passed +``` + +Observed BF16 SM90 drift against the PyTorch reference: + +| Shape | LSE max abs | dlogp max abs | +| --- | ---: | ---: | +| `[2, 8, 1024]` | `4.76837158203125e-07` | `4.76837158203125e-07` | +| `[2, 16, 151936]` | `9.5367431640625e-07` | `9.5367431640625e-07` | + +Both runs used TP=1, no communication, and the explicit SM90 backend without fallback. From c028b5b71b4d29c3539dd059d36a65008e7bd40e Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 4 Aug 2026 21:50:52 +0800 Subject: [PATCH 11/25] fix: address logprob harness lint and provenance --- rl_engine/testing/logprob_comparison.py | 10 +++----- scripts/compare_logprob.py | 5 +--- tests/test_logprob_comparison.py | 33 ++++++++++++++++++++----- 3 files changed, 31 insertions(+), 17 deletions(-) diff --git a/rl_engine/testing/logprob_comparison.py b/rl_engine/testing/logprob_comparison.py index ad2ef76c..0be7fba4 100644 --- a/rl_engine/testing/logprob_comparison.py +++ b/rl_engine/testing/logprob_comparison.py @@ -175,15 +175,11 @@ def _run_ws1_reference( target_ids: torch.Tensor, ignore_index: int, ) -> tuple[torch.Tensor, torch.Tensor]: - from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( - NativeBatchInvariantLogpOp, - ) + from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp op = NativeBatchInvariantLogpOp() logp = op(logits, target_ids, ignore_index=ignore_index, validate=True) - _, lse = op.forward_with_lse( - logits, target_ids, ignore_index=ignore_index, validate=True - ) + _, lse = op.forward_with_lse(logits, target_ids, ignore_index=ignore_index, validate=True) return logp.detach(), lse.detach() @@ -215,12 +211,12 @@ def _run_candidate( def _candidate_provenance(candidate: LogprobCandidate) -> dict[str, Any]: return { + **candidate.provenance, "requested_backend": candidate.requested_backend, "actual_backend": candidate.actual_backend, "tp_world": 1, "communication": "none", "lse_source": "direct", - **candidate.provenance, } diff --git a/scripts/compare_logprob.py b/scripts/compare_logprob.py index 6b42d4ea..b4736db7 100644 --- a/scripts/compare_logprob.py +++ b/scripts/compare_logprob.py @@ -16,10 +16,7 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from rl_engine.testing import ( # noqa: E402 - LogprobComparisonInputs, - compare_single_gpu_logprob, -) +from rl_engine.testing import LogprobComparisonInputs, compare_single_gpu_logprob # noqa: E402 from rl_engine.utils.logger import logger # noqa: E402 diff --git a/tests/test_logprob_comparison.py b/tests/test_logprob_comparison.py index 67c1d422..d4fece0a 100644 --- a/tests/test_logprob_comparison.py +++ b/tests/test_logprob_comparison.py @@ -12,9 +12,7 @@ from rl_engine.kernels.gtest import run_operator_suite from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case -from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( - NativeBatchInvariantLogpOp, -) +from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp from rl_engine.testing.logprob_comparison import ( LogprobBackendUnavailable, LogprobCandidate, @@ -81,6 +79,31 @@ def shifted(logits, target_ids, ignore_index): assert drift["lse"]["p99_abs"] == pytest.approx(0.693, abs=1e-5) +def test_canonical_provenance_cannot_be_overridden(): + native = make_logprob_candidate("pytorch") + candidate = LogprobCandidate( + name="custom", + requested_backend="pytorch", + actual_backend="pytorch", + fn=native.fn, + provenance={ + "actual_backend": "fallback", + "tp_world": 8, + "communication": "all-gather", + "lse_source": "reconstructed", + "implementation": "custom", + }, + ) + + provenance = compare_single_gpu_logprob(_inputs(), candidates=(candidate,)).drifts[0].provenance + + assert provenance["actual_backend"] == "pytorch" + assert provenance["tp_world"] == 1 + assert provenance["communication"] == "none" + assert provenance["lse_source"] == "direct" + assert provenance["implementation"] == "custom" + + def test_all_inactive_tokens_produce_zero_dlogp_statistics(): inputs = _inputs() inputs = LogprobComparisonInputs( @@ -194,9 +217,7 @@ def test_operator_comparison_specs_register_batch_invariant_logp(): case = make_operator_case(args, torch.float32, torch.device("cpu")) candidate = make_candidate(args) - report = run_operator_suite( - "batch_invariant_logp", candidates=[candidate], cases=[case] - ) + report = run_operator_suite("batch_invariant_logp", candidates=[candidate], cases=[case]) assert report.passed assert report.candidates[0].cases[0].op_class == "logprob" From 7ba09b523b031b5e21de172a1d4d0766fc254e2f Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 4 Aug 2026 22:03:37 +0800 Subject: [PATCH 12/25] fix: type heterogeneous logprob backends --- rl_engine/testing/logprob_comparison.py | 1 + 1 file changed, 1 insertion(+) diff --git a/rl_engine/testing/logprob_comparison.py b/rl_engine/testing/logprob_comparison.py index 0be7fba4..52ca82d7 100644 --- a/rl_engine/testing/logprob_comparison.py +++ b/rl_engine/testing/logprob_comparison.py @@ -63,6 +63,7 @@ def to_dict(self) -> dict[str, Any]: def make_logprob_candidate(backend: str) -> LogprobCandidate: normalized = backend.strip().lower().replace("_", "-") + op: Any if normalized in {"pytorch", "native"}: from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( NativeBatchInvariantLogpOp, From 4231625fbd3c436affc7cf701d8f8c1459cf8ec4 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Thu, 6 Aug 2026 09:31:02 +0800 Subject: [PATCH 13/25] fix comment --- rl_engine/kernels/logprob_contract.py | 3 +++ rl_engine/kernels/registry.py | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index b0bcd7ab..1267740a 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -486,6 +486,9 @@ def cross_rank_fingerprint(self) -> str: for key, value in payload["sharding"].items() if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} } + # Note: Any future extensions to this payload MUST maintain strict JSON + # serialization determinism across environments to prevent cross-rank + # hashing mismatches. encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) return hashlib.sha256(encoded.encode("utf-8")).hexdigest() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 7093b320..60f92527 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -510,6 +510,11 @@ def get_logprob_op( "determinism through ReductionSpec.determinism_scope and match it against " "backend determinism_scopes instead" ) + if requested_backend.strip().lower() == "auto" and contract.sharding.tp_world_size > 1: + raise LogprobContractError( + "Unsafe dispatch: requested_backend='auto' is not permitted when tp_world_size > 1 " + "without explicit cross-rank preflighting." + ) platform = self._platform() candidates = self._logprob_candidates.get(platform, []) From 2360a7106de7b6e42a4aee89f0ae4a7ac7709ad0 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Fri, 7 Aug 2026 20:31:56 +0800 Subject: [PATCH 14/25] test(ws2): align dispatch tests with the auto+TP>1 unsafe-dispatch guard The guard added per review rejects requested_backend="auto" whenever tp_world_size > 1, so TP-sharded dispatch tests now name an explicit policy and auto-policy tests use TP=1 contracts. Add coverage for the guard itself and document the restriction in get_logprob_op. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_012PyjQEqDJwy9Cos4Sb9QBK --- rl_engine/kernels/registry.py | 7 +++++-- tests/test_logprob_contract.py | 32 ++++++++++++++++++++++++++------ 2 files changed, 31 insertions(+), 8 deletions(-) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 60f92527..32cfdc37 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -496,7 +496,10 @@ def get_logprob_op( (``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``. + default is ``auto``. With ``tp_world_size > 1``, ``auto`` is rejected: + per-rank auto resolution can diverge across ranks, so distributed + callers must name a policy or backend id and preflight agreement via + ``LogprobContract.cross_rank_fingerprint``. """ if not isinstance(contract, LogprobContract): @@ -510,7 +513,7 @@ def get_logprob_op( "determinism through ReductionSpec.determinism_scope and match it against " "backend determinism_scopes instead" ) - if requested_backend.strip().lower() == "auto" and contract.sharding.tp_world_size > 1: + if requested_backend.lower() == "auto" and contract.sharding.tp_world_size > 1: raise LogprobContractError( "Unsafe dispatch: requested_backend='auto' is not permitted when tp_world_size > 1 " "without explicit cross-rank preflighting." diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py index 40cea4e2..42dc6dd1 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -251,7 +251,7 @@ 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()) + registry.get_logprob_op(_contract(), requested_backend="reference") message = str(exc_info.value) assert "TP=2 is unsupported" in message @@ -277,8 +277,9 @@ def test_undeclared_backend_capability_is_never_selected(): platform = registry._platform() registry._logprob_candidates[platform] = [OpBackend.PYTORCH_NATIVE] + tp1_contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) with pytest.raises(RuntimeError, match="no LogprobBackendCapability declared"): - registry.get_logprob_op(_contract()) + registry.get_logprob_op(tp1_contract) def test_declared_compatible_backend_resolves_and_records_provenance(): @@ -354,7 +355,8 @@ def test_default_auto_policy_resolves_any_compatible_implementation_kind(): platform=platform, ) - result = registry.get_logprob_op(_contract()) + tp1_contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + result = registry.get_logprob_op(tp1_contract) assert result.provenance["requested_backend"] == "auto" assert result.capability.implementation_kind == "reference" @@ -407,7 +409,7 @@ def test_capability_rejections_are_reported_as_fallback(): OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform ) - result = registry.get_logprob_op(_contract()) + result = registry.get_logprob_op(_contract(), requested_backend="reference") assert result.provenance["fallback"] is True assert "TP=2 is unsupported" in result.provenance["prior_rejections"][0] @@ -441,7 +443,7 @@ def test_register_logprob_backend_is_the_public_registration_seam(): ) assert registry._logprob_candidates[platform] == [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] - result = registry.get_logprob_op(_contract()) + result = registry.get_logprob_op(_contract(), requested_backend="reference") assert result.capability.backend_id == "replacement-backend" with pytest.raises(LogprobContractError, match="capability must be"): @@ -467,7 +469,7 @@ def test_capabilities_are_scoped_per_platform(): platform=other, ) - result = registry.get_logprob_op(_contract()) + result = registry.get_logprob_op(_contract(), requested_backend="reference") assert result.capability.backend_id == "test-deterministic-tp-logprob" assert ( @@ -502,6 +504,24 @@ def test_requested_deterministic_policy_is_a_loud_error(): registry.get_logprob_op(_contract(), requested_backend="deterministic") +def test_auto_policy_is_rejected_for_tp_sharded_contracts(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + with pytest.raises(LogprobContractError, match="Unsafe dispatch"): + registry.get_logprob_op(_contract()) + + with pytest.raises(LogprobContractError, match="Unsafe dispatch"): + registry.get_logprob_op(_contract(), requested_backend=" AUTO ") + + tp1_contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + assert registry.get_logprob_op(tp1_contract).provenance["requested_backend"] == "auto" + + def test_determinism_scope_is_part_of_the_typed_contract(): fixed_only = replace( _declared_tp_backend(), From 4eebb3bd1dd6be2f8c1ef2f2ae54ce80973132ce Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Sat, 8 Aug 2026 17:37:29 +0800 Subject: [PATCH 15/25] refactor: colocate logprob harness tooling and docs --- docs/design/ws2-logprob-single-gpu-harness.md | 95 ------------- docs/design/ws2-logprob-sm90-validation.md | 134 ------------------ docs/operators/batch-invariant-logp.md | 103 +++++++++++++- rl_engine/testing/logprob_comparison.py | 94 ++++++++++++ scripts/compare_logprob.py | 103 -------------- tests/test_logprob_comparison.py | 32 ++++- 6 files changed, 226 insertions(+), 335 deletions(-) delete mode 100644 docs/design/ws2-logprob-single-gpu-harness.md delete mode 100644 docs/design/ws2-logprob-sm90-validation.md delete mode 100644 scripts/compare_logprob.py diff --git a/docs/design/ws2-logprob-single-gpu-harness.md b/docs/design/ws2-logprob-single-gpu-harness.md deleted file mode 100644 index c5b869b0..00000000 --- a/docs/design/ws2-logprob-single-gpu-harness.md +++ /dev/null @@ -1,95 +0,0 @@ -# WS2 Single-GPU Logprob Comparison Harness - -This harness is the TP=1 registration and regression guard for issue #241. It compares -selected-token logprob implementations before any distributed communication is introduced. - -## Contract - -For each logical token row, every backend returns direct FP32 values: - -```text -LSE = logsumexp(logits[..., vocab]) -logp = selected_logit - LSE -``` - -The harness uses the merged WS1 batch-invariant PyTorch implementation as its reference. -Reference logp is obtained through the unchanged production call, while reference LSE is -obtained through the diagnostic entry point. The TP=1 PyTorch candidate follows the same -core computation and must be bitwise equal. This is a regression guard, not new -distributed mathematics. - -LSE drift is reported over every logical token row. Selected-token dlogp drift is reported -only over active response/action tokens. Both reports contain max, mean, p95, p99, and the -number of compared values. - -## Exact Backend Selection - -Supported backend names are: - -- `pytorch` -- `triton` -- `cuda-sm90` - -The comparison path does not use registry fallback. An explicitly requested backend must -run exactly or raise `LogprobBackendUnavailable`. In particular, `cuda-sm90` requires a -compiled SM90 extension, Hopper hardware, BF16/FP32 logits, and a compatible vocab row -stride. The production operator may retain its normal fallback behavior outside the -harness. - -Each backend exposes a diagnostic-only `forward_with_lse` method. Existing production -calls remain unchanged: - -```text -op(logits, target_ids) -> logp -op.forward_with_lse(logits, target_ids) -> (logp, lse) -``` - -## Usage - -CPU TP=1 regression guard: - -```bash -python scripts/compare_logprob.py \ - --candidate pytorch \ - --device cpu \ - --dtype fp32 \ - --batch 2 \ - --seq 16 \ - --vocab 257 -``` - -GPU comparison: - -```bash -python scripts/compare_logprob.py \ - --candidate triton \ - --candidate cuda-sm90 \ - --device cuda \ - --dtype bf16 \ - --batch 2 \ - --seq 16 \ - --vocab 151936 -``` - -The command prints a structured JSON report to stdout containing input dtype/shape, -active-token count, TP world size, communication mode, requested and actual backends, -direct-LSE provenance, bitwise logp status, and LSE/dlogp drift statistics. Backend -diagnostic logs are routed to stderr so redirected stdout remains valid JSON. - -## Scope Boundary - -This harness is intentionally single-GPU and records `tp_world=1` and -`communication=none`. It does not implement vocab-shard metadata, all-gather transport, -fixed-order cross-rank LSE merging, CP reconstruction, or distributed artifacts. Those -belong to the later PR3 and PR4 work in issue #241. - -## Tests - -```bash -python -m pytest tests/test_logprob_comparison.py -q -``` - -The focused tests cover bitwise TP=1 regression, direct LSE identity, active-token-only -percentiles, zero active tokens, invalid ignore-index usage, structured serialization, -generic operator-harness registration, exact GPU backend diagnostics, and fail-closed -backend provenance. diff --git a/docs/design/ws2-logprob-sm90-validation.md b/docs/design/ws2-logprob-sm90-validation.md deleted file mode 100644 index e94cb52d..00000000 --- a/docs/design/ws2-logprob-sm90-validation.md +++ /dev/null @@ -1,134 +0,0 @@ -# WS2 Logprob PR2 SM90 Validation - -This document records the Hopper SM90 validation procedure for the PR2 single-GPU -logprob comparison harness from issue #241. It is a validation note for maintainers; -the cloud setup wrapper used during development is intentionally kept outside the -repository. - -## Prerequisites - -The validation host must provide: - -- Python 3.10 or newer; -- CUDA-enabled PyTorch; -- an NVIDIA Hopper GPU with compute capability 9.0, such as H100, H800, or H200; -- `nvidia-smi` and `nvcc`; -- a CUDA development environment capable of compiling the RL-Kernel extension. - -The CUDA version reported by `nvcc` must match `torch.version.cuda`. A runtime-only -image is insufficient because it normally does not include the CUDA compiler. - -## Build - -Activate an environment containing the repository dependencies and a CUDA-enabled -PyTorch installation, then build the editable extension with SM90 enabled: - -```bash -export FORCE_CUDA=1 -export KERNEL_ALIGN_FORCE_SM90=1 -export TORCH_CUDA_ARCH_LIST="9.0+PTX" -export MAX_JOBS=2 - -python -m pip install --no-build-isolation --no-deps -e . -``` - -Verify the extension and SM90 symbol after the build. Import PyTorch first so its -runtime libraries are available to the extension loader: - -```bash -python - <<'PY' -import torch -from rl_engine import _C - -print("torch:", torch.__version__) -print("torch CUDA:", torch.version.cuda) -print("extension:", _C.__file__) -print("SM90 symbol:", hasattr(_C, "batch_invariant_logp_sm90")) -PY -``` - -The final line must report `SM90 symbol: True`. - -## Validation commands - -Run the focused PR2 tests and the complete batch-invariant logprob suite: - -```bash -python -m pytest \ - tests/test_logprob_comparison.py \ - tests/test_operator_inputs.py \ - tests/test_op_checks.py -q - -python -m pytest tests/test_batch_invariant_logp.py -q -``` - -Run the two explicit SM90 comparisons: - -```bash -python scripts/compare_logprob.py \ - --candidate cuda-sm90 \ - --device cuda \ - --dtype bf16 \ - --batch 2 \ - --seq 8 \ - --vocab 1024 \ - --prompt-tokens 3 \ - --seed 7 - -python scripts/compare_logprob.py \ - --candidate cuda-sm90 \ - --device cuda \ - --dtype bf16 \ - --batch 2 \ - --seq 16 \ - --vocab 151936 \ - --prompt-tokens 8 \ - --seed 241 -``` - -The comparison command writes the JSON report to stdout. Diagnostic log messages are -written to stderr so stdout can be redirected directly to a `.json` file. - -## Expected report - -The report must identify the requested and actual backend as `cuda-sm90`, use the -`BatchInvariantLogpSM90Op` implementation, and record: - -```text -tp_world=1 -communication=none -lse_source=direct -``` - -LSE drift is measured over all logical token rows. Selected-logprob drift is measured -only over active response/action tokens. Each drift section includes maximum, mean, -p95, p99, and active-count values. - -## Validation result - -The procedure was validated on: - -```text -GPU: NVIDIA H800 PCIe -Compute capability: 9.0 -Python: 3.11.15 -PyTorch: 2.11.0+cu128 -CUDA toolkit / nvcc: 12.8 -Triton: 3.6.0 -``` - -Results: - -```text -PR2 focused tests: 41 passed -Complete batch-invariant logprob suite: 67 passed -``` - -Observed BF16 SM90 drift against the PyTorch reference: - -| Shape | LSE max abs | dlogp max abs | -| --- | ---: | ---: | -| `[2, 8, 1024]` | `4.76837158203125e-07` | `4.76837158203125e-07` | -| `[2, 16, 151936]` | `9.5367431640625e-07` | `9.5367431640625e-07` | - -Both runs used TP=1, no communication, and the explicit SM90 backend without fallback. diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index fbc0e9f1..d8e95616 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -179,6 +179,101 @@ fp16/bf16 backward: checked against fp32 reference with relaxed tolerance CPU-vs-CUDA comparisons use tolerance-based checks; batch-invariance checks within the same backend use exact equality where appropriate. +## TP=1 Comparison Harness + +The single-GPU comparison harness is the TP=1 registration and regression guard +for issue #241. It uses the batch-invariant PyTorch implementation as the +reference and compares exact `pytorch`, `triton`, or `cuda-sm90` backends before +distributed communication is introduced. + +Each backend exposes a diagnostic-only entry point while the production contract +remains unchanged: + +```text +op(logits, target_ids) -> logp +op.forward_with_lse(logits, target_ids) -> (logp, lse) +``` + +The harness reports LSE drift over every logical token row and selected-logprob +drift over active response/action tokens only. Drift summaries contain max, +mean, p95, p99, and the number of compared values. Reports also record requested +and actual backends, implementation, direct-LSE provenance, input shape and +dtype, `tp_world=1`, and `communication=none`. + +Backend selection is exact and does not use registry fallback. In particular, +an explicit `cuda-sm90` comparison fails unless the compiled SM90 extension, +Hopper hardware, input dtype, and vocab row stride satisfy the kernel contract. + +Run the PyTorch TP=1 guard directly from the kernel-specific testing module: + +```bash +python rl_engine/testing/logprob_comparison.py \ + --candidate pytorch \ + --device cpu \ + --dtype fp32 \ + --batch 2 \ + --seq 16 \ + --vocab 257 +``` + +On a GPU, repeat `--candidate` to compare multiple exact backends: + +```bash +python rl_engine/testing/logprob_comparison.py \ + --candidate triton \ + --candidate cuda-sm90 \ + --device cuda \ + --dtype bf16 \ + --batch 2 \ + --seq 16 \ + --vocab 151936 +``` + +The command writes structured JSON to stdout and routes backend diagnostics to +stderr. The harness does not implement vocab sharding, collective communication, +cross-rank LSE merging, or CP reconstruction. + +### SM90 validation + +SM90 validation requires a Hopper GPU, CUDA-enabled PyTorch, and an `nvcc` +toolkit matching `torch.version.cuda`. Build the extension with: + +```bash +export FORCE_CUDA=1 +export KERNEL_ALIGN_FORCE_SM90=1 +export TORCH_CUDA_ARCH_LIST="9.0+PTX" + +python -m pip install --no-build-isolation --no-deps -e . +``` + +Run the focused harness tests, the complete operator suite, and an explicit +SM90 comparison: + +```bash +python -m pytest \ + tests/test_logprob_comparison.py \ + tests/test_operator_inputs.py \ + tests/test_op_checks.py -q + +python -m pytest tests/test_batch_invariant_logp.py -q + +python rl_engine/testing/logprob_comparison.py \ + --candidate cuda-sm90 \ + --device cuda \ + --dtype bf16 \ + --batch 2 \ + --seq 16 \ + --vocab 151936 \ + --prompt-tokens 8 \ + --seed 241 +``` + +The PR2 path was validated on an NVIDIA H800 PCIe with PyTorch 2.11.0+cu128, +CUDA 12.8, and Triton 3.6.0. The focused tests passed 41 cases and the complete +batch-invariant suite passed 67 cases. For BF16 shape `[2, 16, 151936]`, both +LSE and active-token dlogp had maximum absolute drift +`9.5367431640625e-07` against the PyTorch reference, with no backend fallback. + ## Minimal Example ```python @@ -206,11 +301,14 @@ out.sum().backward() python -m pytest tests/test_batch_invariant_logp.py -q -rs ``` -All backends (Native, Triton) are tested in a single file. Coverage includes: +All production backends are tested in a single file. Coverage includes correctness, leading-shape preservation, batch-invariance (bitwise), validation, ignore-index behavior, backward correctness, CUDA smoke cases, registry dispatch, and Triton-specific fp32/fp16/bf16 correctness, large vocab, backward -gradient batch-invariance, and ignored-row zero gradients. +gradient batch-invariance, and ignored-row zero gradients. The focused +`tests/test_logprob_comparison.py` suite covers TP=1 bitwise regression, direct +LSE identity, active-token drift statistics, structured serialization, exact +backend diagnostics, and fail-closed provenance. Triton tests skip when Triton or CUDA is unavailable. On Windows, run via WSL/Linux with CUDA. @@ -223,4 +321,5 @@ WSL/Linux with CUDA. - `csrc/cuda/batch_invariant_logp_kernel_sm90.cu` - `rl_engine/kernels/registry.py` - `tests/test_batch_invariant_logp.py` +- `tests/test_logprob_comparison.py` - `benchmarks/benchmark_batch_invariant_logp.py` diff --git a/rl_engine/testing/logprob_comparison.py b/rl_engine/testing/logprob_comparison.py index 52ca82d7..d207f75b 100644 --- a/rl_engine/testing/logprob_comparison.py +++ b/rl_engine/testing/logprob_comparison.py @@ -5,12 +5,22 @@ from __future__ import annotations +import argparse +import json +import logging +import pathlib +import sys from collections.abc import Callable, Sequence from dataclasses import asdict, dataclass, field from typing import Any import torch +if __package__ in (None, ""): + repo_root = pathlib.Path(__file__).resolve().parents[2] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + class LogprobBackendUnavailable(RuntimeError): pass @@ -277,6 +287,86 @@ def _validate_inputs( return active, effective +def _dtype(name: str) -> torch.dtype: + return { + "fp32": torch.float32, + "bf16": torch.bfloat16, + "fp16": torch.float16, + }[name] + + +def _device(name: str) -> torch.device: + if name == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device(name) + + +def _route_rl_kernel_logs_to_stderr() -> None: + from rl_engine.utils.logger import logger + + for handler in logger.handlers: + if isinstance(handler, logging.StreamHandler): + handler.setStream(sys.stderr) + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run the WS2 TP=1 selected-logprob/LSE comparison harness." + ) + parser.add_argument( + "--candidate", + action="append", + choices=("pytorch", "triton", "cuda-sm90"), + help="Exact backend to compare. Repeat for multiple backends; defaults to pytorch.", + ) + parser.add_argument("--device", default="auto") + parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq", type=int, default=16) + parser.add_argument("--vocab", type=int, default=257) + parser.add_argument("--prompt-tokens", type=int, default=8) + parser.add_argument("--seed", type=int, default=123) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> None: + _route_rl_kernel_logs_to_stderr() + args = _parse_args(argv) + device = _device(args.device) + if args.batch < 1 or args.seq < 1 or args.vocab < 1: + raise ValueError("batch, seq, and vocab must be positive") + if not 0 <= args.prompt_tokens <= args.seq: + raise ValueError("prompt-tokens must be in [0, seq]") + + generator = torch.Generator(device=device).manual_seed(args.seed) + logits = torch.randn( + args.batch, + args.seq, + args.vocab, + generator=generator, + device=device, + dtype=_dtype(args.dtype), + ) + target_ids = torch.randint( + 0, + args.vocab, + (args.batch, args.seq), + generator=generator, + device=device, + ) + active_mask = torch.ones((args.batch, args.seq), device=device, dtype=torch.bool) + active_mask[:, : args.prompt_tokens] = False + report = compare_single_gpu_logprob( + LogprobComparisonInputs( + logits=logits, + target_ids=target_ids, + active_token_mask=active_mask, + ), + candidates=tuple(args.candidate or ("pytorch",)), + ) + print(json.dumps(report.to_dict(), indent=2, sort_keys=True)) + + __all__ = [ "LogprobBackendUnavailable", "LogprobCandidate", @@ -285,3 +375,7 @@ def _validate_inputs( "compare_single_gpu_logprob", "make_logprob_candidate", ] + + +if __name__ == "__main__": + main() diff --git a/scripts/compare_logprob.py b/scripts/compare_logprob.py deleted file mode 100644 index b4736db7..00000000 --- a/scripts/compare_logprob.py +++ /dev/null @@ -1,103 +0,0 @@ -#!/usr/bin/env python -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -from __future__ import annotations - -import argparse -import json -import logging -import pathlib -import sys - -import torch - -REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from rl_engine.testing import LogprobComparisonInputs, compare_single_gpu_logprob # noqa: E402 -from rl_engine.utils.logger import logger # noqa: E402 - - -def _dtype(name: str) -> torch.dtype: - return { - "fp32": torch.float32, - "bf16": torch.bfloat16, - "fp16": torch.float16, - }[name] - - -def _device(name: str) -> torch.device: - if name == "auto": - return torch.device("cuda" if torch.cuda.is_available() else "cpu") - return torch.device(name) - - -def _route_rl_kernel_logs_to_stderr() -> None: - """Keep stdout machine-readable while preserving backend diagnostics.""" - for handler in logger.handlers: - if isinstance(handler, logging.StreamHandler): - handler.setStream(sys.stderr) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Run the WS2 TP=1 selected-logprob/LSE comparison harness." - ) - parser.add_argument( - "--candidate", - action="append", - choices=("pytorch", "triton", "cuda-sm90"), - help="Exact backend to compare. Repeat for multiple backends; defaults to pytorch.", - ) - parser.add_argument("--device", default="auto") - parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") - parser.add_argument("--batch", type=int, default=2) - parser.add_argument("--seq", type=int, default=16) - parser.add_argument("--vocab", type=int, default=257) - parser.add_argument("--prompt-tokens", type=int, default=8) - parser.add_argument("--seed", type=int, default=123) - return parser.parse_args() - - -def main() -> None: - _route_rl_kernel_logs_to_stderr() - args = parse_args() - device = _device(args.device) - if args.batch < 1 or args.seq < 1 or args.vocab < 1: - raise ValueError("batch, seq, and vocab must be positive") - if not 0 <= args.prompt_tokens <= args.seq: - raise ValueError("prompt-tokens must be in [0, seq]") - - generator = torch.Generator(device=device).manual_seed(args.seed) - logits = torch.randn( - args.batch, - args.seq, - args.vocab, - generator=generator, - device=device, - dtype=_dtype(args.dtype), - ) - target_ids = torch.randint( - 0, - args.vocab, - (args.batch, args.seq), - generator=generator, - device=device, - ) - active_mask = torch.ones((args.batch, args.seq), device=device, dtype=torch.bool) - active_mask[:, : args.prompt_tokens] = False - report = compare_single_gpu_logprob( - LogprobComparisonInputs( - logits=logits, - target_ids=target_ids, - active_token_mask=active_mask, - ), - candidates=tuple(args.candidate or ("pytorch",)), - ) - print(json.dumps(report.to_dict(), indent=2, sort_keys=True)) - - -if __name__ == "__main__": - main() diff --git a/tests/test_logprob_comparison.py b/tests/test_logprob_comparison.py index d4fece0a..98f976f9 100644 --- a/tests/test_logprob_comparison.py +++ b/tests/test_logprob_comparison.py @@ -5,6 +5,7 @@ import io import json import logging +import subprocess import sys import pytest @@ -17,11 +18,12 @@ LogprobBackendUnavailable, LogprobCandidate, LogprobComparisonInputs, + _device, + _route_rl_kernel_logs_to_stderr, compare_single_gpu_logprob, make_logprob_candidate, ) from rl_engine.utils.logger import logger -from scripts.compare_logprob import _device, _route_rl_kernel_logs_to_stderr def _inputs() -> LogprobComparisonInputs: @@ -196,6 +198,34 @@ def test_cli_routes_rl_kernel_logs_to_stderr_for_machine_readable_stdout(monkeyp assert "test backend diagnostic" in stderr.getvalue() +def test_cli_runs_directly_from_testing_module(): + result = subprocess.run( + [ + sys.executable, + "rl_engine/testing/logprob_comparison.py", + "--candidate", + "pytorch", + "--device", + "cpu", + "--batch", + "1", + "--seq", + "2", + "--vocab", + "17", + "--prompt-tokens", + "1", + ], + check=True, + capture_output=True, + text=True, + ) + + payload = json.loads(result.stdout) + assert payload["drifts"][0]["provenance"]["actual_backend"] == "pytorch" + assert payload["input_provenance"]["communication"] == "none" + + def test_operator_comparison_specs_register_batch_invariant_logp(): args = argparse.Namespace( op="batch_invariant_logp", From 19488cc7baf127be2b5ce7820223a51c5a556fe5 Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Sat, 8 Aug 2026 17:51:35 +0800 Subject: [PATCH 16/25] test: resolve logprob CLI path reliably --- tests/test_logprob_comparison.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_logprob_comparison.py b/tests/test_logprob_comparison.py index 98f976f9..4fc62a13 100644 --- a/tests/test_logprob_comparison.py +++ b/tests/test_logprob_comparison.py @@ -7,6 +7,7 @@ import logging import subprocess import sys +from pathlib import Path import pytest import torch @@ -199,10 +200,11 @@ def test_cli_routes_rl_kernel_logs_to_stderr_for_machine_readable_stdout(monkeyp def test_cli_runs_directly_from_testing_module(): + script = Path(__file__).resolve().parents[1] / "rl_engine" / "testing" / "logprob_comparison.py" result = subprocess.run( [ sys.executable, - "rl_engine/testing/logprob_comparison.py", + str(script), "--candidate", "pytorch", "--device", From b7d9d895580bc60558bdacadd097baa8e7b676d4 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Wed, 5 Aug 2026 16:05:57 +0800 Subject: [PATCH 17/25] 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 c4d553c5..3c2fdfc8 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 @@ -336,3 +349,6 @@ WSL/Linux with CUDA. - `tests/test_batch_invariant_logp.py` - `tests/test_logprob_comparison.py` - `benchmarks/benchmark_batch_invariant_logp.py` +- `rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py` +- `rl_engine/kernels/logprob_contract.py` +- `tests/test_vocab_parallel_logp.py` diff --git a/rl_engine/kernels/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 32cfdc37..4b2c5306 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 42dc6dd1..307acfdf 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(), requested_backend="reference") @@ -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 3866d3c18a76668d9eff00804cdd262c63b56f72 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Wed, 5 Aug 2026 16:10:26 +0800 Subject: [PATCH 18/25] 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 65f3c6f05456cb6e2223c9e267ebb0d86705fcbd Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Wed, 5 Aug 2026 21:13:04 +0800 Subject: [PATCH 19/25] 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 05d19eb20bac1fc165c38b6094a658af83b07145 Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Sat, 8 Aug 2026 22:07:36 +0800 Subject: [PATCH 20/25] test: align PR3 dispatch with latest PR1 guard --- tests/test_logprob_contract.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py index 307acfdf..a74f3b4c 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -297,7 +297,7 @@ def test_ws1_rejections_recorded_when_vocab_parallel_reference_resolves(): candidates.remove(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) candidates.append(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) - result = registry.get_logprob_op(_contract()) + result = registry.get_logprob_op(_contract(), requested_backend="reference") assert result.capability.backend_id == "pytorch-vocab-parallel-logp-ws2" assert result.provenance["fallback"] is True rejections = " | ".join(result.provenance["prior_rejections"]) From f36a63d217623c4b396ed336b8f965a19582dd7a Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 11 Aug 2026 20:18:05 +0800 Subject: [PATCH 21/25] feat(ws2): add distributed logprob drift runner --- docs/operators/batch-invariant-logp.md | 65 ++ rl_engine/testing/__init__.py | 5 + .../testing/distributed_logprob_comparison.py | 807 ++++++++++++++++++ rl_engine/testing/logprob_comparison.py | 54 +- rl_engine/testing/logprob_drift.py | 57 ++ tests/test_distributed_logprob_comparison.py | 217 +++++ 6 files changed, 1165 insertions(+), 40 deletions(-) create mode 100644 rl_engine/testing/distributed_logprob_comparison.py create mode 100644 rl_engine/testing/logprob_drift.py create mode 100644 tests/test_distributed_logprob_comparison.py diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index 3c2fdfc8..8d1b6487 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -300,6 +300,65 @@ batch-invariant suite passed 67 cases. For BF16 shape `[2, 16, 151936]`, both LSE and active-token dlogp had maximum absolute drift `9.5367431640625e-07` against the PyTorch reference, with no backend fallback. +## Distributed WS2 Drift Report + +The issue #241 PR4 runner materializes one TP/CP topology per `torchrun` +invocation. TP partitions the vocabulary and is the only numerical merge axis; +CP partitions token rows and is recorded in provenance without participating in +the vocab-domain LSE merge. For global rank `r`: + +```text +tp_rank = r % tp_world_size +cp_rank = r // tp_world_size +``` + +Every case generates the same seeded FP32 logical logits, targets, and active +mask. The candidate receives a BF16 token/vocab shard through the explicit +`pytorch-vocab-parallel-logp-ws2` backend, while the independent oracle computes +`torch.logsumexp` over the complete real-vocab FP32 token slice. Distributed +dispatch rejects `auto`, capability fallback, topology mismatches, non-tileable +vocabularies, and incomplete materialization. + +Reports follow the issue #116 fields and contain per-rank and aggregate LSE and +active-token dlogp summaries: max/mean/p95/p99 absolute drift, max relative +drift, worst global token position, target id, target owner rank, #108 tolerance, +and pass/fail. Provenance includes TP/CP topology, dtype, shard bounds, backend +capability, contract fingerprint, reduction spec, merge order, transport, and +the exact launch command. Replicated TP outputs are checked bitwise before one +representative per CP shard is included in aggregate statistics. + +Print the scoped TP=1/2/4 x CP=1/2 launch matrix without starting workers: + +```bash +python rl_engine/testing/distributed_logprob_comparison.py \ + --plan \ + --device cuda \ + --dtype bf16 \ + --output artifacts/ws2-logprob/report.json +``` + +Run one TP=2, CP=2 Qwen3-vocab case on four local GPUs: + +```bash +torchrun --standalone --nproc-per-node=4 \ + rl_engine/testing/distributed_logprob_comparison.py \ + --tp 2 \ + --cp 2 \ + --dtype bf16 \ + --backend pytorch-vocab-parallel-logp-ws2 \ + --real-vocab 151936 \ + --padded-vocab 151936 \ + --num-vocab-tiles 64 \ + --batch 2 \ + --seq 16 \ + --prompt-tokens 8 \ + --output artifacts/ws2-logprob/tp2-cp2.json +``` + +The full matrix requires up to eight ranks for TP=4, CP=2. CPU/Gloo cases are +available for topology and artifact validation; the scoped numerical gate is +BF16 on CUDA/NCCL. + ## Minimal Example ```python @@ -335,6 +394,9 @@ gradient batch-invariance, and ignored-row zero gradients. The focused `tests/test_logprob_comparison.py` suite covers TP=1 bitwise regression, direct LSE identity, active-token drift statistics, structured serialization, exact backend diagnostics, and fail-closed provenance. +`tests/test_distributed_logprob_comparison.py` covers topology planning, TP/CP +rank mapping, token/vocab sharding, explicit backend materialization, #116 JSON +artifacts, and a real four-process TP=2, CP=2 Gloo smoke case. Triton tests skip when Triton or CUDA is unavailable. On Windows, run via WSL/Linux with CUDA. @@ -348,6 +410,9 @@ WSL/Linux with CUDA. - `rl_engine/kernels/registry.py` - `tests/test_batch_invariant_logp.py` - `tests/test_logprob_comparison.py` +- `rl_engine/testing/logprob_drift.py` +- `rl_engine/testing/distributed_logprob_comparison.py` +- `tests/test_distributed_logprob_comparison.py` - `benchmarks/benchmark_batch_invariant_logp.py` - `rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py` - `rl_engine/kernels/logprob_contract.py` diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 51759abd..97b585fb 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -10,7 +10,9 @@ LogprobComparisonReport, compare_single_gpu_logprob, make_logprob_candidate, + route_rl_kernel_logs_to_stderr, ) +from .logprob_drift import LogprobDriftStats, summarize_logprob_drift from .reference_ops import ( active_token_count, compute_policy_ratio, @@ -27,6 +29,7 @@ "LogprobCandidate", "LogprobComparisonInputs", "LogprobComparisonReport", + "LogprobDriftStats", "SyntheticRLKernelBatch", "active_token_count", "compare_single_gpu_logprob", @@ -37,5 +40,7 @@ "masked_mean", "masked_sum", "selected_logprobs_reference", + "route_rl_kernel_logs_to_stderr", + "summarize_logprob_drift", "summarize_kernel_drift", ] diff --git a/rl_engine/testing/distributed_logprob_comparison.py b/rl_engine/testing/distributed_logprob_comparison.py new file mode 100644 index 00000000..b5b8e283 --- /dev/null +++ b/rl_engine/testing/distributed_logprob_comparison.py @@ -0,0 +1,807 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Distributed WS2 comparison for the vocab-parallel logprob reference.""" + +from __future__ import annotations + +import argparse +import contextlib +import hashlib +import json +import os +import pathlib +import shlex +import sys +from dataclasses import asdict, dataclass +from typing import Any, Sequence + +import torch + +if __package__ in (None, ""): + repo_root = pathlib.Path(__file__).resolve().parents[2] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + +_import_output = ( + contextlib.redirect_stdout(sys.stderr) + if __package__ in (None, "") + else contextlib.nullcontext() +) +with _import_output: + from rl_engine.kernels.gtest.tolerance import load_contract as load_tolerance_contract + from rl_engine.kernels.logprob_contract import ( + LogprobContract, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, + ) + from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + BACKEND_ID, + DEFAULT_NUM_VOCAB_TILES, + ) + from rl_engine.kernels.registry import KernelRegistry + from rl_engine.testing.logprob_comparison import route_rl_kernel_logs_to_stderr + from rl_engine.testing.logprob_drift import LogprobDriftStats, summarize_logprob_drift + +_DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} +_TOLERANCE_DTYPES = { + "bf16": "bfloat16", + "fp16": "float16", + "fp32": "float32", +} + + +@dataclass(frozen=True) +class DistributedLogprobCase: + tp_world_size: int + cp_world_size: int + dtype: str = "bf16" + requested_backend: str = BACKEND_ID + real_vocab_size: int = 151936 + padded_vocab_size: int = 151936 + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES + batch_size: int = 2 + sequence_length: int = 16 + prompt_tokens: int = 8 + seed: int = 123 + ignore_index: int = -100 + + def __post_init__(self) -> None: + for name in ( + "tp_world_size", + "cp_world_size", + "real_vocab_size", + "padded_vocab_size", + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + if self.dtype not in _DTYPES: + raise ValueError(f"dtype must be one of {sorted(_DTYPES)}") + if not self.requested_backend or self.requested_backend.lower() == "auto": + raise ValueError("distributed cases require an explicit non-auto backend") + if self.padded_vocab_size < self.real_vocab_size: + raise ValueError("padded_vocab_size must be at least real_vocab_size") + if self.num_vocab_tiles < self.tp_world_size: + raise ValueError("num_vocab_tiles must be at least tp_world_size") + if self.padded_vocab_size % self.num_vocab_tiles != 0: + raise ValueError("num_vocab_tiles must divide padded_vocab_size exactly") + if self.batch_size <= 0 or self.sequence_length <= 0: + raise ValueError("batch_size and sequence_length must be positive") + if not 0 <= self.prompt_tokens <= self.sequence_length: + raise ValueError("prompt_tokens must be in [0, sequence_length]") + + @property + def world_size(self) -> int: + return self.tp_world_size * self.cp_world_size + + @property + def num_tokens(self) -> int: + return self.batch_size * self.sequence_length + + @property + def case_id(self) -> str: + encoded = json.dumps(asdict(self), sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest()[:16] + + +@dataclass(frozen=True) +class RankTopology: + global_rank: int + world_size: int + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + tp_group_ranks: tuple[int, ...] + + +@dataclass(frozen=True) +class DriftDetail: + stats: LogprobDriftStats + max_rel: float + worst_global_token: int | None + worst_target_id: int | None + worst_owner_rank: int | None + candidate_value: float | None + reference_value: float | None + atol: float + rtol: float + passed: bool + + +@dataclass(frozen=True) +class RankLogprobReport: + global_rank: int + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + sp_world_size: int + dp_world_size: int + token_start: int + token_end: int + vocab_start: int + vocab_end: int + device: str + requested_backend: str + actual_backend: str + fallback: bool + contract_fingerprint: str + contract: dict[str, Any] + capability: dict[str, Any] + tp_outputs_bitwise_replicated: bool + lse: DriftDetail + dlogp: DriftDetail + passed: bool + + +@dataclass(frozen=True) +class DistributedLogprobReport: + schema_version: int + case_id: str + case: dict[str, Any] + launch_command: str + environment: dict[str, Any] + ranks: tuple[RankLogprobReport, ...] + aggregate: dict[str, DriftDetail] + passed: bool + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class _RankPayload: + report: RankLogprobReport + candidate_logp: torch.Tensor + candidate_lse: torch.Tensor + reference_logp: torch.Tensor + reference_lse: torch.Tensor + active_mask: torch.Tensor + target_ids: torch.Tensor + global_positions: torch.Tensor + + +def plan_distributed_logprob_cases( + *, + tp_world_sizes: Sequence[int] = (1, 2, 4), + cp_world_sizes: Sequence[int] = (1, 2), + **overrides: Any, +) -> tuple[DistributedLogprobCase, ...]: + """Build the scoped issue #241 topology product in deterministic order.""" + + return tuple( + DistributedLogprobCase(tp_world_size=tp, cp_world_size=cp, **overrides) + for tp in tp_world_sizes + for cp in cp_world_sizes + ) + + +def rank_topology(case: DistributedLogprobCase, global_rank: int) -> RankTopology: + if not 0 <= global_rank < case.world_size: + raise ValueError(f"global_rank must be in [0, {case.world_size})") + cp_rank, tp_rank = divmod(global_rank, case.tp_world_size) + group_start = cp_rank * case.tp_world_size + return RankTopology( + global_rank=global_rank, + world_size=case.world_size, + tp_rank=tp_rank, + tp_world_size=case.tp_world_size, + cp_rank=cp_rank, + cp_world_size=case.cp_world_size, + tp_group_ranks=tuple(range(group_start, group_start + case.tp_world_size)), + ) + + +def token_shard_bounds(num_tokens: int, cp_world_size: int) -> tuple[tuple[int, int], ...]: + """Partition token rows contiguously, allowing a one-row imbalance.""" + + if num_tokens < cp_world_size: + raise ValueError("num_tokens must be at least cp_world_size") + quotient, remainder = divmod(num_tokens, cp_world_size) + bounds = [] + cursor = 0 + for cp_rank in range(cp_world_size): + count = quotient + int(cp_rank < remainder) + bounds.append((cursor, cursor + count)) + cursor += count + return tuple(bounds) + + +def vocab_shard_bounds(case: DistributedLogprobCase) -> tuple[tuple[int, int], ...]: + """Assign complete global vocab tiles to TP ranks.""" + + tile_size = case.padded_vocab_size // case.num_vocab_tiles + quotient, remainder = divmod(case.num_vocab_tiles, case.tp_world_size) + bounds = [] + cursor_tiles = 0 + for tp_rank in range(case.tp_world_size): + tile_count = quotient + int(tp_rank < remainder) + start = cursor_tiles * tile_size + cursor_tiles += tile_count + bounds.append((start, cursor_tiles * tile_size)) + return tuple(bounds) + + +def format_launch_command( + case: DistributedLogprobCase, + *, + output: str | pathlib.Path, + device: str = "cuda", + dist_backend: str | None = None, +) -> str: + backend = dist_backend or ("nccl" if device == "cuda" else "gloo") + arguments = [ + "torchrun", + "--standalone", + f"--nproc-per-node={case.world_size}", + "rl_engine/testing/distributed_logprob_comparison.py", + "--tp", + str(case.tp_world_size), + "--cp", + str(case.cp_world_size), + "--dtype", + case.dtype, + "--backend", + case.requested_backend, + "--real-vocab", + str(case.real_vocab_size), + "--padded-vocab", + str(case.padded_vocab_size), + "--num-vocab-tiles", + str(case.num_vocab_tiles), + "--batch", + str(case.batch_size), + "--seq", + str(case.sequence_length), + "--prompt-tokens", + str(case.prompt_tokens), + "--seed", + str(case.seed), + "--device", + device, + "--dist-backend", + backend, + "--output", + str(output), + ] + return shlex.join(arguments) + + +def _canonical_inputs( + case: DistributedLogprobCase, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(case.seed) + logits = torch.randn( + case.num_tokens, + case.padded_vocab_size, + generator=generator, + dtype=torch.float32, + ) + targets = torch.randint( + 0, + case.real_vocab_size, + (case.num_tokens,), + generator=generator, + dtype=torch.long, + ) + active = torch.ones((case.batch_size, case.sequence_length), dtype=torch.bool) + active[:, : case.prompt_tokens] = False + active = active.reshape(-1) + targets = targets.masked_fill(~active, case.ignore_index) + return logits, targets, active + + +def _make_contract( + case: DistributedLogprobCase, + topology: RankTopology, + active_mask: torch.Tensor, +) -> LogprobContract: + return LogprobContract( + role=LogprobRole.TRAIN, + dtype=LogprobDType(case.dtype), + mask=MaskSpec( + num_tokens=int(active_mask.numel()), + active_mask=tuple(bool(value) for value in active_mask.tolist()), + ignore_index=case.ignore_index, + ), + sharding=ShardingSpec( + tp_rank=topology.tp_rank, + tp_world_size=case.tp_world_size, + vocab_shard_bounds=vocab_shard_bounds(case), + real_vocab_size=case.real_vocab_size, + padded_vocab_size=case.padded_vocab_size, + cp_rank=topology.cp_rank, + cp_world_size=case.cp_world_size, + ), + reduction=ReductionSpec(), + ) + + +def _fp32_oracle( + logits: torch.Tensor, + target_ids: torch.Tensor, + active_mask: torch.Tensor, + real_vocab_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + real_logits = logits[:, :real_vocab_size].float() + lse = torch.logsumexp(real_logits, dim=-1) + safe_targets = target_ids.masked_fill(~active_mask, 0) + selected = real_logits.gather(1, safe_targets.unsqueeze(1)).squeeze(1) + logp = torch.where(active_mask, selected - lse, torch.zeros_like(lse)) + return logp, lse + + +def _resolve_tolerance(dtype: str) -> tuple[float, float]: + entry = load_tolerance_contract()["accuracy"]["default"]["logprob"] + tolerance = entry[_TOLERANCE_DTYPES[dtype]] + return float(tolerance["atol"]), float(tolerance["rtol"]) + + +def _drift_detail( + candidate: torch.Tensor, + reference: torch.Tensor, + *, + target_ids: torch.Tensor, + global_positions: torch.Tensor, + sharding: ShardingSpec, + atol: float, + rtol: float, + mask: torch.Tensor | None = None, +) -> DriftDetail: + stats = summarize_logprob_drift(candidate, reference, mask=mask) + diff = (candidate.float() - reference.float()).abs() + selected = torch.ones_like(diff, dtype=torch.bool) if mask is None else mask.to(diff.device) + if not bool(selected.any().item()): + return DriftDetail(stats, 0.0, None, None, None, None, None, atol, rtol, True) + + selected_diff = diff[selected] + selected_ref = reference.float()[selected] + relative = selected_diff / selected_ref.abs().clamp_min(torch.finfo(torch.float32).tiny) + selected_indices = torch.arange(diff.numel(), device=diff.device)[selected] + worst_selected = int(selected_diff.argmax().item()) + worst_local = int(selected_indices[worst_selected].item()) + target_id = int(target_ids[worst_local].item()) + close = selected_diff <= atol + rtol * selected_ref.abs() + return DriftDetail( + stats=stats, + max_rel=float(relative.max().item()), + worst_global_token=int(global_positions[worst_local].item()), + worst_target_id=target_id, + worst_owner_rank=sharding.owner_rank(target_id) if target_id >= 0 else None, + candidate_value=float(candidate[worst_local].float().item()), + reference_value=float(reference[worst_local].float().item()), + atol=atol, + rtol=rtol, + passed=bool(close.all().item()), + ) + + +def _tp_outputs_replicated( + logp: torch.Tensor, + lse: torch.Tensor, + *, + tp_group: Any, + tp_world_size: int, +) -> bool: + if tp_world_size == 1: + return True + import torch.distributed as dist + + gathered_logp = [torch.empty_like(logp) for _ in range(tp_world_size)] + gathered_lse = [torch.empty_like(lse) for _ in range(tp_world_size)] + dist.all_gather(gathered_logp, logp.contiguous(), group=tp_group) + dist.all_gather(gathered_lse, lse.contiguous(), group=tp_group) + return all(torch.equal(logp, value) for value in gathered_logp) and all( + torch.equal(lse, value) for value in gathered_lse + ) + + +def _execute_rank( + case: DistributedLogprobCase, + topology: RankTopology, + *, + device: torch.device, + tp_group: Any, +) -> _RankPayload: + full_logits, full_targets, full_active = _canonical_inputs(case) + token_start, token_end = token_shard_bounds(case.num_tokens, case.cp_world_size)[ + topology.cp_rank + ] + vocab_start, vocab_end = vocab_shard_bounds(case)[topology.tp_rank] + token_slice = slice(token_start, token_end) + local_active = full_active[token_slice].to(device=device) + local_targets = full_targets[token_slice].to(device=device) + local_fp32 = full_logits[token_slice].to(device=device) + local_logits = local_fp32[:, vocab_start:vocab_end].to(_DTYPES[case.dtype]).contiguous() + positions = torch.arange(token_start, token_end, device=device, dtype=torch.long) + + contract = _make_contract(case, topology, local_active.cpu()) + dispatch = KernelRegistry().get_logprob_op( + contract, + requested_backend=case.requested_backend, + ) + fallback = bool(dispatch.provenance["fallback"]) + if fallback: + raise RuntimeError("distributed logprob dispatch materialized through a fallback") + requested_policy = case.requested_backend.lower() + if requested_policy not in {"reference", "production"} and ( + case.requested_backend != dispatch.capability.backend_id + ): + raise RuntimeError( + f"requested backend {case.requested_backend!r} materialized as " + f"{dispatch.capability.backend_id!r}" + ) + + candidate_logp, candidate_lse = dispatch.op( + local_logits, + local_targets, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=case.num_vocab_tiles, + validate=True, + ) + reference_logp, reference_lse = _fp32_oracle( + local_fp32, + local_targets, + local_active, + case.real_vocab_size, + ) + replicated = _tp_outputs_replicated( + candidate_logp, + candidate_lse, + tp_group=tp_group, + tp_world_size=case.tp_world_size, + ) + atol, rtol = _resolve_tolerance(case.dtype) + lse_drift = _drift_detail( + candidate_lse, + reference_lse, + target_ids=local_targets, + global_positions=positions, + sharding=contract.sharding, + atol=atol, + rtol=rtol, + ) + dlogp_drift = _drift_detail( + candidate_logp, + reference_logp, + target_ids=local_targets, + global_positions=positions, + sharding=contract.sharding, + atol=atol, + rtol=rtol, + mask=local_active, + ) + rank_report = RankLogprobReport( + global_rank=topology.global_rank, + tp_rank=topology.tp_rank, + tp_world_size=topology.tp_world_size, + cp_rank=topology.cp_rank, + cp_world_size=topology.cp_world_size, + sp_world_size=1, + dp_world_size=1, + token_start=token_start, + token_end=token_end, + vocab_start=vocab_start, + vocab_end=vocab_end, + device=str(device), + requested_backend=case.requested_backend, + actual_backend=dispatch.capability.backend_id, + fallback=fallback, + contract_fingerprint=contract.cross_rank_fingerprint(), + contract=contract.to_dict(), + capability=dispatch.capability.to_dict(), + tp_outputs_bitwise_replicated=replicated, + lse=lse_drift, + dlogp=dlogp_drift, + passed=replicated and lse_drift.passed and dlogp_drift.passed, + ) + return _RankPayload( + report=rank_report, + candidate_logp=candidate_logp.detach().cpu(), + candidate_lse=candidate_lse.detach().cpu(), + reference_logp=reference_logp.detach().cpu(), + reference_lse=reference_lse.detach().cpu(), + active_mask=local_active.cpu(), + target_ids=local_targets.cpu(), + global_positions=positions.cpu(), + ) + + +def _aggregate_payloads( + case: DistributedLogprobCase, + payloads: Sequence[_RankPayload], +) -> dict[str, DriftDetail]: + representatives = sorted( + (payload for payload in payloads if payload.report.tp_rank == 0), + key=lambda payload: payload.report.cp_rank, + ) + if len(representatives) != case.cp_world_size: + raise RuntimeError("missing one or more CP representatives in rank reports") + candidate_logp = torch.cat([payload.candidate_logp for payload in representatives]) + candidate_lse = torch.cat([payload.candidate_lse for payload in representatives]) + reference_logp = torch.cat([payload.reference_logp for payload in representatives]) + reference_lse = torch.cat([payload.reference_lse for payload in representatives]) + active_mask = torch.cat([payload.active_mask for payload in representatives]) + target_ids = torch.cat([payload.target_ids for payload in representatives]) + positions = torch.cat([payload.global_positions for payload in representatives]) + sharding = _make_contract( + case, + rank_topology(case, 0), + active_mask, + ).sharding + atol, rtol = _resolve_tolerance(case.dtype) + return { + "lse": _drift_detail( + candidate_lse, + reference_lse, + target_ids=target_ids, + global_positions=positions, + sharding=sharding, + atol=atol, + rtol=rtol, + ), + "dlogp": _drift_detail( + candidate_logp, + reference_logp, + target_ids=target_ids, + global_positions=positions, + sharding=sharding, + atol=atol, + rtol=rtol, + mask=active_mask, + ), + } + + +def _create_tp_group(case: DistributedLogprobCase, topology: RankTopology) -> Any: + if case.world_size == 1: + return None + import torch.distributed as dist + + selected = None + for cp_rank in range(case.cp_world_size): + start = cp_rank * case.tp_world_size + ranks = list(range(start, start + case.tp_world_size)) + group = dist.new_group(ranks=ranks) + if topology.global_rank in ranks: + selected = group + return selected + + +def run_distributed_logprob_case( + case: DistributedLogprobCase, + *, + device_name: str = "cuda", + dist_backend: str | None = None, + output: str | pathlib.Path, +) -> DistributedLogprobReport | None: + """Run one materialized topology; only global rank zero returns the report.""" + + import torch.distributed as dist + + backend = dist_backend or ("nccl" if device_name == "cuda" else "gloo") + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if world_size != case.world_size: + raise RuntimeError( + f"WORLD_SIZE={world_size} does not match TP*CP={case.world_size}; " + "launch exactly the topology declared by the case" + ) + initialized_here = False + if world_size > 1 and not dist.is_initialized(): + dist.init_process_group(backend=backend) + initialized_here = True + if dist.is_initialized(): + if dist.get_world_size() != world_size or dist.get_rank() != rank: + raise RuntimeError("initialized process group does not match RANK/WORLD_SIZE") + + if device_name == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable") + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + elif device_name == "cpu": + device = torch.device("cpu") + else: + raise ValueError("device must be cuda or cpu") + + topology = rank_topology(case, rank) + tp_group = _create_tp_group(case, topology) + try: + payload = _execute_rank(case, topology, device=device, tp_group=tp_group) + if world_size == 1: + payloads = [payload] + else: + gathered: list[Any] = [None] * world_size + dist.all_gather_object(gathered, payload) + payloads = gathered + + report = None + if rank == 0: + aggregate = _aggregate_payloads(case, payloads) + rank_reports = tuple( + payload.report + for payload in sorted(payloads, key=lambda item: item.report.global_rank) + ) + actual_backends = sorted({rank_report.actual_backend for rank_report in rank_reports}) + reduction_specs = { + json.dumps(rank_report.contract["reduction"], sort_keys=True) + for rank_report in rank_reports + } + materialization_consistent = len(actual_backends) == 1 and len(reduction_specs) == 1 + launch_command = format_launch_command( + case, + output=output, + device=device_name, + dist_backend=backend, + ) + report = DistributedLogprobReport( + schema_version=1, + case_id=case.case_id, + case=asdict(case), + launch_command=launch_command, + environment={ + "python": sys.version.split()[0], + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "dist_backend": backend, + "world_size": world_size, + "sp_world_size": 1, + "dp_world_size": 1, + "materialization": { + "actual_backends": actual_backends, + "consistent": materialization_consistent, + }, + "communication": { + "logprob_merge_axis": "tp_vocab", + "cp_is_merge_axis": False, + "report_collection": ("all_gather_object" if world_size > 1 else "none"), + }, + }, + ranks=rank_reports, + aggregate=aggregate, + passed=( + materialization_consistent + and all(rank_report.passed for rank_report in rank_reports) + and all(detail.passed for detail in aggregate.values()) + ), + ) + output_path = pathlib.Path(output) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(report.to_dict(), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + if world_size > 1: + dist.barrier() + return report + finally: + if initialized_here and dist.is_initialized(): + dist.destroy_process_group() + + +def _case_from_args(args: argparse.Namespace) -> DistributedLogprobCase: + return DistributedLogprobCase( + tp_world_size=args.tp, + cp_world_size=args.cp, + dtype=args.dtype, + requested_backend=args.backend, + real_vocab_size=args.real_vocab, + padded_vocab_size=args.padded_vocab, + num_vocab_tiles=args.num_vocab_tiles, + batch_size=args.batch, + sequence_length=args.seq, + prompt_tokens=args.prompt_tokens, + seed=args.seed, + ) + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the WS2 distributed logprob drift report.") + parser.add_argument("--plan", action="store_true", help="Print the six scoped launch commands.") + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--cp", type=int, default=1) + parser.add_argument("--dtype", choices=tuple(_DTYPES), default="bf16") + parser.add_argument("--backend", default=BACKEND_ID) + parser.add_argument("--real-vocab", type=int, default=151936) + parser.add_argument("--padded-vocab", type=int, default=151936) + parser.add_argument("--num-vocab-tiles", type=int, default=DEFAULT_NUM_VOCAB_TILES) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq", type=int, default=16) + parser.add_argument("--prompt-tokens", type=int, default=8) + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda") + parser.add_argument("--dist-backend", choices=("nccl", "gloo"), default=None) + parser.add_argument("--output", default="artifacts/ws2-logprob/report.json") + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> None: + route_rl_kernel_logs_to_stderr() + args = _parse_args(argv) + if args.plan: + cases = plan_distributed_logprob_cases( + dtype=args.dtype, + requested_backend=args.backend, + real_vocab_size=args.real_vocab, + padded_vocab_size=args.padded_vocab, + num_vocab_tiles=args.num_vocab_tiles, + batch_size=args.batch, + sequence_length=args.seq, + prompt_tokens=args.prompt_tokens, + seed=args.seed, + ) + commands = [ + format_launch_command( + case, + output=pathlib.Path(args.output).parent + / f"tp{case.tp_world_size}-cp{case.cp_world_size}.json", + device=args.device, + dist_backend=args.dist_backend, + ) + for case in cases + ] + print(json.dumps({"commands": commands}, indent=2)) + return + + case = _case_from_args(args) + report = run_distributed_logprob_case( + case, + device_name=args.device, + dist_backend=args.dist_backend, + output=args.output, + ) + if report is not None: + print(json.dumps(report.to_dict(), indent=2, sort_keys=True)) + if not report.passed: + raise SystemExit(1) + + +__all__ = [ + "DistributedLogprobCase", + "DistributedLogprobReport", + "DriftDetail", + "RankLogprobReport", + "RankTopology", + "format_launch_command", + "plan_distributed_logprob_cases", + "rank_topology", + "run_distributed_logprob_case", + "token_shard_bounds", + "vocab_shard_bounds", +] + + +if __name__ == "__main__": + main() diff --git a/rl_engine/testing/logprob_comparison.py b/rl_engine/testing/logprob_comparison.py index d207f75b..1e9d4e63 100644 --- a/rl_engine/testing/logprob_comparison.py +++ b/rl_engine/testing/logprob_comparison.py @@ -20,6 +20,9 @@ repo_root = pathlib.Path(__file__).resolve().parents[2] if str(repo_root) not in sys.path: sys.path.insert(0, str(repo_root)) + from logprob_drift import LogprobDriftStats, summarize_logprob_drift +else: + from .logprob_drift import LogprobDriftStats, summarize_logprob_drift class LogprobBackendUnavailable(RuntimeError): @@ -43,20 +46,11 @@ class LogprobCandidate: provenance: dict[str, Any] = field(default_factory=dict) -@dataclass(frozen=True) -class _DriftStats: - max_abs: float - mean_abs: float - p95_abs: float - p99_abs: float - active_count: int - - @dataclass(frozen=True) class _LogprobPathDrift: candidate_name: str - lse: _DriftStats - dlogp: _DriftStats + lse: LogprobDriftStats + dlogp: LogprobDriftStats bitwise_logp: bool provenance: dict[str, Any] @@ -158,8 +152,8 @@ def compare_single_gpu_logprob( drifts.append( _LogprobPathDrift( candidate_name=candidate.name, - lse=_drift_stats(lse, reference_lse), - dlogp=_drift_stats(logp, reference_logp, mask=active_mask), + lse=summarize_logprob_drift(lse, reference_lse), + dlogp=summarize_logprob_drift(logp, reference_logp, mask=active_mask), bitwise_logp=torch.equal(logp, reference_logp), provenance=_candidate_provenance(candidate), ) @@ -231,31 +225,6 @@ def _candidate_provenance(candidate: LogprobCandidate) -> dict[str, Any]: } -def _drift_stats( - candidate: torch.Tensor, - reference: torch.Tensor, - *, - mask: torch.Tensor | None = None, -) -> _DriftStats: - if candidate.shape != reference.shape: - raise ValueError( - f"candidate shape {tuple(candidate.shape)} must match reference shape " - f"{tuple(reference.shape)}" - ) - diff = (candidate.float() - reference.float()).abs() - values = diff.reshape(-1) if mask is None else diff[mask.to(device=diff.device)] - count = int(values.numel()) - if count == 0: - return _DriftStats(0.0, 0.0, 0.0, 0.0, 0) - return _DriftStats( - max_abs=float(values.max().item()), - mean_abs=float(values.mean().item()), - p95_abs=float(torch.quantile(values, 0.95).item()), - p99_abs=float(torch.quantile(values, 0.99).item()), - active_count=count, - ) - - def _validate_inputs( inputs: LogprobComparisonInputs, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -301,7 +270,7 @@ def _device(name: str) -> torch.device: return torch.device(name) -def _route_rl_kernel_logs_to_stderr() -> None: +def route_rl_kernel_logs_to_stderr() -> None: from rl_engine.utils.logger import logger for handler in logger.handlers: @@ -309,6 +278,10 @@ def _route_rl_kernel_logs_to_stderr() -> None: handler.setStream(sys.stderr) +def _route_rl_kernel_logs_to_stderr() -> None: + route_rl_kernel_logs_to_stderr() + + def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser( description="Run the WS2 TP=1 selected-logprob/LSE comparison harness." @@ -330,7 +303,7 @@ def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: def main(argv: Sequence[str] | None = None) -> None: - _route_rl_kernel_logs_to_stderr() + route_rl_kernel_logs_to_stderr() args = _parse_args(argv) device = _device(args.device) if args.batch < 1 or args.seq < 1 or args.vocab < 1: @@ -374,6 +347,7 @@ def main(argv: Sequence[str] | None = None) -> None: "LogprobComparisonReport", "compare_single_gpu_logprob", "make_logprob_candidate", + "route_rl_kernel_logs_to_stderr", ] diff --git a/rl_engine/testing/logprob_drift.py b/rl_engine/testing/logprob_drift.py new file mode 100644 index 00000000..4996dfbc --- /dev/null +++ b/rl_engine/testing/logprob_drift.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared selected-logprob drift summaries.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class LogprobDriftStats: + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + +def summarize_logprob_drift( + candidate: torch.Tensor, + reference: torch.Tensor, + *, + mask: torch.Tensor | None = None, +) -> LogprobDriftStats: + """Summarize absolute drift, optionally over active rows only.""" + + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match reference shape " + f"{tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs() + if mask is None: + values = diff.reshape(-1) + else: + if mask.shape != diff.shape: + raise ValueError("mask shape must match candidate and reference") + if mask.dtype != torch.bool: + raise ValueError("mask must be bool") + values = diff[mask.to(device=diff.device)] + + count = int(values.numel()) + if count == 0: + return LogprobDriftStats(0.0, 0.0, 0.0, 0.0, 0) + return LogprobDriftStats( + max_abs=float(values.max().item()), + mean_abs=float(values.mean().item()), + p95_abs=float(torch.quantile(values, 0.95).item()), + p99_abs=float(torch.quantile(values, 0.99).item()), + active_count=count, + ) + + +__all__ = ["LogprobDriftStats", "summarize_logprob_drift"] diff --git a/tests/test_distributed_logprob_comparison.py b/tests/test_distributed_logprob_comparison.py new file mode 100644 index 00000000..022ccc16 --- /dev/null +++ b/tests/test_distributed_logprob_comparison.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import BACKEND_ID +from rl_engine.testing.distributed_logprob_comparison import ( + DistributedLogprobCase, + format_launch_command, + plan_distributed_logprob_cases, + rank_topology, + run_distributed_logprob_case, + token_shard_bounds, + vocab_shard_bounds, +) +from rl_engine.testing.logprob_drift import summarize_logprob_drift + + +def _small_case(*, tp: int = 1, cp: int = 1) -> DistributedLogprobCase: + return DistributedLogprobCase( + tp_world_size=tp, + cp_world_size=cp, + real_vocab_size=13, + padded_vocab_size=16, + num_vocab_tiles=8, + batch_size=1, + sequence_length=4, + prompt_tokens=1, + seed=7, + ) + + +def test_planner_builds_the_scoped_topology_product(): + cases = plan_distributed_logprob_cases( + real_vocab_size=13, + padded_vocab_size=16, + num_vocab_tiles=8, + ) + + assert [(case.tp_world_size, case.cp_world_size) for case in cases] == [ + (1, 1), + (1, 2), + (2, 1), + (2, 2), + (4, 1), + (4, 2), + ] + assert [case.world_size for case in cases] == [1, 2, 2, 4, 4, 8] + + +def test_rank_mapping_keeps_cp_out_of_the_tp_merge_axis(): + case = _small_case(tp=2, cp=2) + + assert rank_topology(case, 0).tp_group_ranks == (0, 1) + assert rank_topology(case, 1).tp_group_ranks == (0, 1) + assert rank_topology(case, 2).tp_group_ranks == (2, 3) + assert rank_topology(case, 3).tp_group_ranks == (2, 3) + assert [ + (rank_topology(case, rank).cp_rank, rank_topology(case, rank).tp_rank) for rank in range(4) + ] == [ + (0, 0), + (0, 1), + (1, 0), + (1, 1), + ] + + +def test_token_and_vocab_bounds_cover_each_axis_once(): + case = _small_case(tp=4, cp=2) + + assert token_shard_bounds(case.num_tokens, case.cp_world_size) == ((0, 2), (2, 4)) + assert vocab_shard_bounds(case) == ((0, 4), (4, 8), (8, 12), (12, 16)) + + +def test_case_rejects_implicit_backend_and_non_tileable_vocab(): + with pytest.raises(ValueError, match="explicit non-auto backend"): + DistributedLogprobCase(tp_world_size=2, cp_world_size=1, requested_backend="auto") + with pytest.raises(ValueError, match="must divide"): + DistributedLogprobCase( + tp_world_size=2, + cp_world_size=1, + padded_vocab_size=15, + real_vocab_size=13, + num_vocab_tiles=8, + ) + + +def test_launch_command_records_the_materialized_case(tmp_path): + case = _small_case(tp=2, cp=2) + command = format_launch_command(case, output=tmp_path / "report.json") + + assert "--nproc-per-node=4" in command + assert "--tp 2 --cp 2" in command + assert f"--backend {BACKEND_ID}" in command + assert "--real-vocab 13 --padded-vocab 16" in command + + +def test_shared_pr2_drift_summary_preserves_active_mask_semantics(): + candidate = torch.tensor([100.0, 1.0, 3.0]) + reference = torch.tensor([0.0, 2.0, 1.0]) + mask = torch.tensor([False, True, True]) + + stats = summarize_logprob_drift(candidate, reference, mask=mask) + + assert stats.active_count == 2 + assert stats.max_abs == 2.0 + assert stats.mean_abs == 1.5 + + +def test_tp1_cpu_case_writes_116_compatible_artifact(tmp_path, monkeypatch): + monkeypatch.delenv("RANK", raising=False) + monkeypatch.delenv("LOCAL_RANK", raising=False) + monkeypatch.delenv("WORLD_SIZE", raising=False) + output = tmp_path / "tp1-cp1.json" + + report = run_distributed_logprob_case( + _small_case(), + device_name="cpu", + dist_backend="gloo", + output=output, + ) + + assert report is not None and report.passed + assert report.aggregate["lse"].stats.active_count == 4 + assert report.aggregate["dlogp"].stats.active_count == 3 + assert report.ranks[0].actual_backend == BACKEND_ID + assert report.ranks[0].fallback is False + assert report.ranks[0].tp_outputs_bitwise_replicated + payload = json.loads(output.read_text(encoding="utf-8")) + assert payload["schema_version"] == 1 + assert payload["ranks"][0]["contract"]["reduction"]["cp_is_merge_axis"] is False + assert payload["ranks"][0]["sp_world_size"] == 1 + assert payload["ranks"][0]["dp_world_size"] == 1 + assert payload["environment"]["materialization"]["consistent"] is True + assert payload["aggregate"]["dlogp"]["worst_target_id"] is not None + assert payload["launch_command"].startswith("torchrun --standalone") + + +def test_world_size_mismatch_fails_before_process_group_init(tmp_path, monkeypatch): + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + + with pytest.raises(RuntimeError, match=r"does not match TP\*CP"): + run_distributed_logprob_case( + _small_case(), + device_name="cpu", + dist_backend="gloo", + output=tmp_path / "unused.json", + ) + + +@pytest.mark.skipif(not torch.distributed.is_available(), reason="torch.distributed required") +def test_tp2_cp2_gloo_cli_emits_per_rank_report(tmp_path): + script = ( + Path(__file__).resolve().parents[1] + / "rl_engine" + / "testing" + / "distributed_logprob_comparison.py" + ) + output = tmp_path / "tp2-cp2.json" + environment = os.environ.copy() + environment.setdefault("OMP_NUM_THREADS", "1") + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc-per-node=4", + str(script), + "--tp", + "2", + "--cp", + "2", + "--device", + "cpu", + "--dist-backend", + "gloo", + "--real-vocab", + "13", + "--padded-vocab", + "16", + "--num-vocab-tiles", + "8", + "--batch", + "1", + "--seq", + "4", + "--prompt-tokens", + "1", + "--output", + str(output), + ], + check=True, + capture_output=True, + text=True, + timeout=120, + env=environment, + ) + + payload = json.loads(output.read_text(encoding="utf-8")) + assert payload["passed"] + assert len(payload["ranks"]) == 4 + assert all(rank["tp_outputs_bitwise_replicated"] for rank in payload["ranks"]) + assert {rank["actual_backend"] for rank in payload["ranks"]} == {BACKEND_ID} + assert {rank["cp_rank"] for rank in payload["ranks"]} == {0, 1} + assert payload["aggregate"]["dlogp"]["stats"]["active_count"] == 3 + assert json.loads(result.stdout)["case"]["tp_world_size"] == 2 From 1d9bac1a258d857485ae0b4bd00fa0e768f126cc Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 11 Aug 2026 22:24:25 +0800 Subject: [PATCH 22/25] ci(ws2): run logprob comparison tests --- .github/workflows/ci.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2d9ba91e..993ccf52 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -82,6 +82,11 @@ jobs: - name: Run WS2 Vocab-Parallel Logprob Tests (CPU-safe) run: python -m pytest tests/test_vocab_parallel_logp.py -v + - name: Run WS2 Logprob Comparison Tests (CPU-safe) + run: | + python -m pytest tests/test_logprob_comparison.py -v + python -m pytest tests/test_distributed_logprob_comparison.py -v + docs: runs-on: ubuntu-latest steps: From f6b5a07798996ed235aae3c3c360e2bb0fe1256c Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 11 Aug 2026 23:02:17 +0800 Subject: [PATCH 23/25] fix(ws2): harden distributed drift reporting --- docs/operators/batch-invariant-logp.md | 5 ++- .../testing/distributed_logprob_comparison.py | 15 ++++++--- tests/test_distributed_logprob_comparison.py | 31 +++++++++++++++++++ 3 files changed, 46 insertions(+), 5 deletions(-) diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index 8d1b6487..b060c10f 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -58,7 +58,10 @@ or device, dispatch is unchanged (Triton -> PyTorch). `VocabParallelLogprobOp` (`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) -**TP=1, TP=2, and TP=4 produce bit-identical results.** +defines a cross-TP bitwise contract for TP=1, TP=2, and TP=4 when +`num_vocab_tiles` is fixed and every vocabulary-shard boundary is tile-aligned. +The complete BF16 CUDA/NCCL validation matrix for this contract is tracked by +issue #241 PR4. 1. Split the padded vocabulary into `num_vocab_tiles` fixed tiles. 2. Each rank computes fp32 `(max, sumexp)` for the tiles it owns. Every tile diff --git a/rl_engine/testing/distributed_logprob_comparison.py b/rl_engine/testing/distributed_logprob_comparison.py index b5b8e283..ad3bb05d 100644 --- a/rl_engine/testing/distributed_logprob_comparison.py +++ b/rl_engine/testing/distributed_logprob_comparison.py @@ -7,6 +7,7 @@ import argparse import contextlib +import datetime import hashlib import json import os @@ -56,6 +57,8 @@ "fp16": "float16", "fp32": "float32", } +_PROCESS_GROUP_TIMEOUT = datetime.timedelta(minutes=5) +_RELATIVE_ERROR_FLOOR = 1.0e-12 @dataclass(frozen=True) @@ -190,6 +193,10 @@ class _RankPayload: global_positions: torch.Tensor +def _strict_report_json(payload: dict[str, Any]) -> str: + return json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + + def plan_distributed_logprob_cases( *, tp_world_sizes: Sequence[int] = (1, 2, 4), @@ -385,7 +392,7 @@ def _drift_detail( selected_diff = diff[selected] selected_ref = reference.float()[selected] - relative = selected_diff / selected_ref.abs().clamp_min(torch.finfo(torch.float32).tiny) + relative = selected_diff.double() / selected_ref.double().abs().clamp_min(_RELATIVE_ERROR_FLOOR) selected_indices = torch.arange(diff.numel(), device=diff.device)[selected] worst_selected = int(selected_diff.argmax().item()) worst_local = int(selected_indices[worst_selected].item()) @@ -620,7 +627,7 @@ def run_distributed_logprob_case( ) initialized_here = False if world_size > 1 and not dist.is_initialized(): - dist.init_process_group(backend=backend) + dist.init_process_group(backend=backend, timeout=_PROCESS_GROUP_TIMEOUT) initialized_here = True if dist.is_initialized(): if dist.get_world_size() != world_size or dist.get_rank() != rank: @@ -700,7 +707,7 @@ def run_distributed_logprob_case( output_path = pathlib.Path(output) output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text( - json.dumps(report.to_dict(), indent=2, sort_keys=True) + "\n", + _strict_report_json(report.to_dict()) + "\n", encoding="utf-8", ) if world_size > 1: @@ -783,7 +790,7 @@ def main(argv: Sequence[str] | None = None) -> None: output=args.output, ) if report is not None: - print(json.dumps(report.to_dict(), indent=2, sort_keys=True)) + print(_strict_report_json(report.to_dict())) if not report.passed: raise SystemExit(1) diff --git a/tests/test_distributed_logprob_comparison.py b/tests/test_distributed_logprob_comparison.py index 022ccc16..a41113e6 100644 --- a/tests/test_distributed_logprob_comparison.py +++ b/tests/test_distributed_logprob_comparison.py @@ -4,6 +4,7 @@ from __future__ import annotations import json +import math import os import subprocess import sys @@ -12,9 +13,12 @@ import pytest import torch +from rl_engine.kernels.logprob_contract import ShardingSpec from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import BACKEND_ID from rl_engine.testing.distributed_logprob_comparison import ( DistributedLogprobCase, + _drift_detail, + _strict_report_json, format_launch_command, plan_distributed_logprob_cases, rank_topology, @@ -116,6 +120,33 @@ def test_shared_pr2_drift_summary_preserves_active_mask_semantics(): assert stats.mean_abs == 1.5 +def test_relative_drift_near_zero_stays_finite(): + sharding = ShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, 16),), + real_vocab_size=13, + padded_vocab_size=16, + ) + detail = _drift_detail( + torch.tensor([1.0]), + torch.tensor([0.0]), + target_ids=torch.tensor([1]), + global_positions=torch.tensor([3]), + sharding=sharding, + atol=0.0, + rtol=0.0, + ) + + assert math.isfinite(detail.max_rel) + assert detail.max_rel == pytest.approx(1.0e12) + assert json.loads(_strict_report_json({"max_rel": detail.max_rel}))["max_rel"] == pytest.approx( + 1.0e12 + ) + with pytest.raises(ValueError, match="Out of range float values"): + _strict_report_json({"max_rel": float("nan")}) + + def test_tp1_cpu_case_writes_116_compatible_artifact(tmp_path, monkeypatch): monkeypatch.delenv("RANK", raising=False) monkeypatch.delenv("LOCAL_RANK", raising=False) From 8a4f4bef5029cbc1711273e68c11979594160943 Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Tue, 11 Aug 2026 23:18:33 +0800 Subject: [PATCH 24/25] fix(ws2): clean up process groups on setup failure --- .../testing/distributed_logprob_comparison.py | 41 ++++++------ tests/test_distributed_logprob_comparison.py | 65 +++++++++++++++++++ 2 files changed, 85 insertions(+), 21 deletions(-) diff --git a/rl_engine/testing/distributed_logprob_comparison.py b/rl_engine/testing/distributed_logprob_comparison.py index ad3bb05d..d3884194 100644 --- a/rl_engine/testing/distributed_logprob_comparison.py +++ b/rl_engine/testing/distributed_logprob_comparison.py @@ -625,27 +625,26 @@ def run_distributed_logprob_case( f"WORLD_SIZE={world_size} does not match TP*CP={case.world_size}; " "launch exactly the topology declared by the case" ) - initialized_here = False - if world_size > 1 and not dist.is_initialized(): - dist.init_process_group(backend=backend, timeout=_PROCESS_GROUP_TIMEOUT) - initialized_here = True - if dist.is_initialized(): - if dist.get_world_size() != world_size or dist.get_rank() != rank: - raise RuntimeError("initialized process group does not match RANK/WORLD_SIZE") - - if device_name == "cuda": - if not torch.cuda.is_available(): - raise RuntimeError("CUDA was requested but is unavailable") - device = torch.device("cuda", local_rank) - torch.cuda.set_device(device) - elif device_name == "cpu": - device = torch.device("cpu") - else: - raise ValueError("device must be cuda or cpu") - - topology = rank_topology(case, rank) - tp_group = _create_tp_group(case, topology) + owns_process_group = world_size > 1 and not dist.is_initialized() try: + if device_name == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable") + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + elif device_name == "cpu": + device = torch.device("cpu") + else: + raise ValueError("device must be cuda or cpu") + + if owns_process_group: + dist.init_process_group(backend=backend, timeout=_PROCESS_GROUP_TIMEOUT) + if dist.is_initialized(): + if dist.get_world_size() != world_size or dist.get_rank() != rank: + raise RuntimeError("initialized process group does not match RANK/WORLD_SIZE") + + topology = rank_topology(case, rank) + tp_group = _create_tp_group(case, topology) payload = _execute_rank(case, topology, device=device, tp_group=tp_group) if world_size == 1: payloads = [payload] @@ -714,7 +713,7 @@ def run_distributed_logprob_case( dist.barrier() return report finally: - if initialized_here and dist.is_initialized(): + if owns_process_group and dist.is_initialized(): dist.destroy_process_group() diff --git a/tests/test_distributed_logprob_comparison.py b/tests/test_distributed_logprob_comparison.py index a41113e6..8309e469 100644 --- a/tests/test_distributed_logprob_comparison.py +++ b/tests/test_distributed_logprob_comparison.py @@ -12,7 +12,9 @@ import pytest import torch +import torch.distributed as dist +import rl_engine.testing.distributed_logprob_comparison as distributed_comparison from rl_engine.kernels.logprob_contract import ShardingSpec from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import BACKEND_ID from rl_engine.testing.distributed_logprob_comparison import ( @@ -189,6 +191,69 @@ def test_world_size_mismatch_fails_before_process_group_init(tmp_path, monkeypat ) +def test_setup_failure_destroys_owned_process_group(tmp_path, monkeypatch): + state = {"initialized": False, "destroyed": False} + + def init_process_group(*, backend, timeout): + state["initialized"] = True + + def destroy_process_group(): + state["destroyed"] = True + state["initialized"] = False + + def fail_group_setup(case, topology): + raise RuntimeError("group setup failed") + + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(dist, "is_initialized", lambda: state["initialized"]) + monkeypatch.setattr(dist, "init_process_group", init_process_group) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + monkeypatch.setattr(dist, "get_rank", lambda: 0) + monkeypatch.setattr(dist, "destroy_process_group", destroy_process_group) + monkeypatch.setattr(distributed_comparison, "_create_tp_group", fail_group_setup) + + with pytest.raises(RuntimeError, match="group setup failed"): + run_distributed_logprob_case( + _small_case(tp=2), + device_name="cpu", + dist_backend="gloo", + output=tmp_path / "unused.json", + ) + + assert state == {"initialized": False, "destroyed": True} + + +def test_partial_initialization_failure_destroys_owned_process_group(tmp_path, monkeypatch): + state = {"initialized": False, "destroyed": False} + + def fail_initialization(*, backend, timeout): + state["initialized"] = True + raise RuntimeError("initialization failed") + + def destroy_process_group(): + state["destroyed"] = True + state["initialized"] = False + + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(dist, "is_initialized", lambda: state["initialized"]) + monkeypatch.setattr(dist, "init_process_group", fail_initialization) + monkeypatch.setattr(dist, "destroy_process_group", destroy_process_group) + + with pytest.raises(RuntimeError, match="initialization failed"): + run_distributed_logprob_case( + _small_case(tp=2), + device_name="cpu", + dist_backend="gloo", + output=tmp_path / "unused.json", + ) + + assert state == {"initialized": False, "destroyed": True} + + @pytest.mark.skipif(not torch.distributed.is_available(), reason="torch.distributed required") def test_tp2_cp2_gloo_cli_emits_per_rank_report(tmp_path): script = ( From c2c99c1ab9a63ce324b88608b0dd6b9e29b7216c Mon Sep 17 00:00:00 2001 From: hihaluemen <1596916766@qq.com> Date: Thu, 20 Aug 2026 00:35:35 +0800 Subject: [PATCH 25/25] feat(ws2): record cross-topology logprob fingerprints --- .../testing/distributed_logprob_comparison.py | 30 +++++++++++++++++++ tests/test_distributed_logprob_comparison.py | 5 ++++ 2 files changed, 35 insertions(+) diff --git a/rl_engine/testing/distributed_logprob_comparison.py b/rl_engine/testing/distributed_logprob_comparison.py index d3884194..3e37c88b 100644 --- a/rl_engine/testing/distributed_logprob_comparison.py +++ b/rl_engine/testing/distributed_logprob_comparison.py @@ -175,6 +175,7 @@ class DistributedLogprobReport: environment: dict[str, Any] ranks: tuple[RankLogprobReport, ...] aggregate: dict[str, DriftDetail] + bitwise_fingerprints: dict[str, Any] passed: bool def to_dict(self) -> dict[str, Any]: @@ -590,6 +591,33 @@ def _aggregate_payloads( } +def _tensor_sha256(tensor: torch.Tensor) -> str: + """Hash the exact CPU tensor bytes for cross-topology comparisons.""" + + data = tensor.detach().cpu().contiguous().numpy().tobytes() + return hashlib.sha256(data).hexdigest() + + +def _aggregate_bitwise_fingerprints( + case: DistributedLogprobCase, + payloads: Sequence[_RankPayload], +) -> dict[str, Any]: + representatives = sorted( + (payload for payload in payloads if payload.report.tp_rank == 0), + key=lambda payload: payload.report.cp_rank, + ) + if len(representatives) != case.cp_world_size: + raise RuntimeError("missing one or more CP representatives in rank reports") + candidate_logp = torch.cat([payload.candidate_logp for payload in representatives]) + candidate_lse = torch.cat([payload.candidate_lse for payload in representatives]) + return { + "candidate_logp_sha256": _tensor_sha256(candidate_logp), + "candidate_lse_sha256": _tensor_sha256(candidate_lse), + "dtype": str(candidate_logp.dtype).replace("torch.", ""), + "shape": list(candidate_logp.shape), + } + + def _create_tp_group(case: DistributedLogprobCase, topology: RankTopology) -> Any: if case.world_size == 1: return None @@ -656,6 +684,7 @@ def run_distributed_logprob_case( report = None if rank == 0: aggregate = _aggregate_payloads(case, payloads) + bitwise_fingerprints = _aggregate_bitwise_fingerprints(case, payloads) rank_reports = tuple( payload.report for payload in sorted(payloads, key=lambda item: item.report.global_rank) @@ -697,6 +726,7 @@ def run_distributed_logprob_case( }, ranks=rank_reports, aggregate=aggregate, + bitwise_fingerprints=bitwise_fingerprints, passed=( materialization_consistent and all(rank_report.passed for rank_report in rank_reports) diff --git a/tests/test_distributed_logprob_comparison.py b/tests/test_distributed_logprob_comparison.py index 8309e469..8bb9d4ad 100644 --- a/tests/test_distributed_logprob_comparison.py +++ b/tests/test_distributed_logprob_comparison.py @@ -175,6 +175,11 @@ def test_tp1_cpu_case_writes_116_compatible_artifact(tmp_path, monkeypatch): assert payload["ranks"][0]["dp_world_size"] == 1 assert payload["environment"]["materialization"]["consistent"] is True assert payload["aggregate"]["dlogp"]["worst_target_id"] is not None + fingerprints = payload["bitwise_fingerprints"] + assert len(fingerprints["candidate_logp_sha256"]) == 64 + assert len(fingerprints["candidate_lse_sha256"]) == 64 + assert fingerprints["dtype"] == "float32" + assert fingerprints["shape"] == [4] assert payload["launch_command"].startswith("torchrun --standalone")