diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..993ccf52 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,6 +76,17 @@ jobs: run: | python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu" + - name: Run WS2 Logprob Contract Tests (CPU-safe) + run: python -m pytest tests/test_logprob_contract.py -v + + - name: Run WS2 Vocab-Parallel Logprob Tests (CPU-safe) + run: python -m pytest tests/test_vocab_parallel_logp.py -v + + - name: Run WS2 Logprob Comparison Tests (CPU-safe) + run: | + python -m pytest tests/test_logprob_comparison.py -v + python -m pytest tests/test_distributed_logprob_comparison.py -v + docs: runs-on: ubuntu-latest steps: diff --git a/README.md b/README.md index f1ab58d8..0be7704e 100644 --- a/README.md +++ b/README.md @@ -117,8 +117,12 @@ RL-Kernel sits between high-level alignment libraries and low-level GPU kernels, git clone https://github.com/RL-Align/RL-Kernel.git cd RL-Kernel -# Install core dependencies (CUDA 12.4+ recommended) -pip install -e . +# CPU-only / pure-Python fallback +python -m pip install -e . + +# Native CUDA or ROCm extension (install a matching PyTorch build first) +RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e . +python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)" ``` ### Contributions diff --git a/csrc/deterministic_logp_kernel.cu b/csrc/deterministic_logp_kernel.cu index 1ca23287..c238e11d 100644 --- a/csrc/deterministic_logp_kernel.cu +++ b/csrc/deterministic_logp_kernel.cu @@ -15,6 +15,15 @@ constexpr int kDeterministicLogpMediumVocabLimit = 4096; constexpr int kDeterministicLogpWarpSize = 32; constexpr float kDeterministicLogpNegInf = -3.4028234663852886e38F; +template +__device__ __forceinline__ T deterministic_logp_shfl_down_32(T value, unsigned int delta) { +#if defined(__HIPCC__) || defined(__HIP_PLATFORM_AMD__) + return __shfl_down(value, delta, kDeterministicLogpWarpSize); +#else + return __shfl_down_sync(0xffffffffu, value, delta, kDeterministicLogpWarpSize); +#endif +} + template struct DeterministicLogpBlockTraits { static_assert( @@ -36,7 +45,7 @@ __device__ __forceinline__ float deterministicBlockReduceMax(float val) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset)); + val = fmaxf(val, deterministic_logp_shfl_down_32(val, offset)); } if (lane == 0) { @@ -50,7 +59,7 @@ __device__ __forceinline__ float deterministicBlockReduceMax(float val) { if (wid == 0) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset)); + val = fmaxf(val, deterministic_logp_shfl_down_32(val, offset)); } } return val; @@ -66,7 +75,7 @@ __device__ __forceinline__ float deterministicBlockReduceSum(float val) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - val += __shfl_down_sync(0xffffffff, val, offset); + val += deterministic_logp_shfl_down_32(val, offset); } if (lane == 0) { @@ -80,7 +89,7 @@ __device__ __forceinline__ float deterministicBlockReduceSum(float val) { if (wid == 0) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - val += __shfl_down_sync(0xffffffff, val, offset); + val += deterministic_logp_shfl_down_32(val, offset); } } return val; diff --git a/csrc/fused_logp_kernel.cu b/csrc/fused_logp_kernel.cu index b620b047..679a6a30 100644 --- a/csrc/fused_logp_kernel.cu +++ b/csrc/fused_logp_kernel.cu @@ -5,6 +5,17 @@ #include #include +constexpr int kFusedLogpLogicalWarpSize = 32; + +template +__device__ __forceinline__ T fused_logp_shfl_down_32(T value, unsigned int delta) { +#if defined(__HIPCC__) || defined(__HIP_PLATFORM_AMD__) + return __shfl_down(value, delta, kFusedLogpLogicalWarpSize); +#else + return __shfl_down_sync(0xffffffffu, value, delta, kFusedLogpLogicalWarpSize); +#endif +} + template __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) { static __shared__ float shared[32]; @@ -14,7 +25,7 @@ __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) { float f_val = static_cast(val); for (int offset = 16; offset > 0; offset /= 2) - f_val = max(f_val, __shfl_down_sync(0xffffffff, f_val, offset)); + f_val = max(f_val, fused_logp_shfl_down_32(f_val, offset)); if (lane == 0) shared[wid] = f_val; __syncthreads(); @@ -22,7 +33,7 @@ __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) { f_val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : -1e20f; if (wid == 0) { for (int offset = 16; offset > 0; offset /= 2) - f_val = max(f_val, __shfl_down_sync(0xffffffff, f_val, offset)); + f_val = max(f_val, fused_logp_shfl_down_32(f_val, offset)); } return static_cast(f_val); } @@ -36,7 +47,7 @@ __device__ __forceinline__ scalar_t blockReduceSum(scalar_t val) { float f_val = static_cast(val); for (int offset = 16; offset > 0; offset /= 2) - f_val += __shfl_down_sync(0xffffffff, f_val, offset); + f_val += fused_logp_shfl_down_32(f_val, offset); if (lane == 0) shared[wid] = f_val; __syncthreads(); @@ -44,7 +55,7 @@ __device__ __forceinline__ scalar_t blockReduceSum(scalar_t val) { f_val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f; if (wid == 0) { for (int offset = 16; offset > 0; offset /= 2) - f_val += __shfl_down_sync(0xffffffff, f_val, offset); + f_val += fused_logp_shfl_down_32(f_val, offset); } return static_cast(f_val); } @@ -82,8 +93,8 @@ __device__ __forceinline__ LogSumExpState blockReduceLogSumExp(LogSumExpState st for (int offset = 16; offset > 0; offset /= 2) { LogSumExpState other{ - __shfl_down_sync(0xffffffff, state.max_val, offset), - __shfl_down_sync(0xffffffff, state.sum_exp, offset)}; + fused_logp_shfl_down_32(state.max_val, offset), + fused_logp_shfl_down_32(state.sum_exp, offset)}; state = merge_logsumexp_state(state, other); } @@ -100,8 +111,8 @@ __device__ __forceinline__ LogSumExpState blockReduceLogSumExp(LogSumExpState st if (wid == 0) { for (int offset = 16; offset > 0; offset /= 2) { LogSumExpState other{ - __shfl_down_sync(0xffffffff, state.max_val, offset), - __shfl_down_sync(0xffffffff, state.sum_exp, offset)}; + fused_logp_shfl_down_32(state.max_val, offset), + fused_logp_shfl_down_32(state.sum_exp, offset)}; state = merge_logsumexp_state(state, other); } } diff --git a/csrc/ops.cpp b/csrc/ops.cpp index eee328a4..97b753d3 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -253,6 +253,7 @@ std::vector deterministic_attention_backward( // Prefix-Shared Attention Declarations & Wrappers +#if !defined(USE_ROCM) void prefix_shared_attention_forward( const __nv_bfloat16 *Q, // [bs, G, len_q, DIM] const __nv_bfloat16 *K, // [bs, len_kv, DIM] @@ -296,6 +297,7 @@ at::Tensor prefix_shared_attention( return O; } #endif +#endif // PyBind11 Module Registration PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -355,8 +357,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out"); m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32"); - // registry Prefix-Shared Attention + // Prefix-shared attention uses NVIDIA PTX and falls back to PyTorch SDPA on ROCm. +#if !defined(USE_ROCM) m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); +#endif // registry Batch-Invariant Deterministic GEMM m.def("det_gemm_fwd", &det_gemm_fwd, "Batch-invariant deterministic GEMM forward (C=A@B)"); diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm new file mode 100644 index 00000000..93ec25f8 --- /dev/null +++ b/docker/Dockerfile.rocm @@ -0,0 +1,28 @@ +# docker/Dockerfile.rocm +# base: docker build -f docker/Dockerfile.rocm_base -t rl-kernel:rocm-dev . +# build: docker build -f docker/Dockerfile.rocm -t /rl-kernel-ci:rocm . +# push: docker push /rl-kernel-ci:rocm + +FROM rl-kernel:rocm-dev + +# Build a portable extension by default. The list follows the multi-architecture +# ROCm profile used by vLLM: MI200 (gfx90a), MI300/MI325 (gfx942), MI350/MI355 +# (gfx950), plus supported RDNA 3/4 targets. Override it at build time with +# --build-arg PYTORCH_ROCM_ARCH=, or at run time with `docker run -e`. +# Newer GPU targets need no setup.py change: pass the gfx target supported by the +# installed PyTorch/ROCm pair. +ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1150;gfx1151;gfx1200;gfx1201 +ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV MAX_JOBS=8 + +USER root +WORKDIR /opt/rl-kernel + +COPY pyproject.toml setup.py* requirements*.txt ./ + +RUN pip install --no-cache-dir -U pip \ + && pip install --no-cache-dir -r requirements.txt \ + && pip install --no-cache-dir pytest + +USER rlkernel +WORKDIR /workspace/RL-Kernel diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base new file mode 100644 index 00000000..ab577175 --- /dev/null +++ b/docker/Dockerfile.rocm_base @@ -0,0 +1,52 @@ +# docker/Dockerfile.rocm_base +# Build: docker build -f docker/Dockerfile.rocm_base -t rl-kernel:rocm-dev . +# +# To build only the targets deployed in a particular image, override the default: +# docker build -f docker/Dockerfile.rocm_base \ +# --build-arg PYTORCH_ROCM_ARCH='gfx942;gfx950' -t rl-kernel:rocm-dev . + +ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete +FROM ${BASE_IMAGE} + +# This is a build-target list, not a hardware allow-list. It covers the target +# families supported by the vLLM ROCm 7.2 reference image: MI200 (gfx90a), +# MI300/MI325 (gfx942), MI350/MI355 (gfx950), and RDNA 3/4. For future GPUs, +# pass the target accepted by the selected ROCm/PyTorch toolchain at build time. +ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1150;gfx1151;gfx1200;gfx1201 +# Keep this wheel index aligned with BASE_IMAGE's ROCm release when overriding it. +ARG PYTORCH_INDEX_URL=https://download.pytorch.org/whl/rocm7.2 +ARG PYTORCH_VERSION=2.12.1 +ARG RL_KERNEL_USER=rlkernel +ARG RL_KERNEL_UID=10001 +ARG RL_KERNEL_GID=10001 + +ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:${PATH} +ENV ROCM_PATH=/opt/rocm +ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:${LD_LIBRARY_PATH} +ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV MAX_JOBS=8 +ENV DEBIAN_FRONTEND=noninteractive + +RUN apt-get update -y \ + && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + cmake \ + git \ + ninja-build \ + pkg-config \ + python3 \ + python3-dev \ + python3-pip \ + python3-venv \ + && python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel \ + && python3 -m pip install --no-cache-dir --index-url "${PYTORCH_INDEX_URL}" "torch==${PYTORCH_VERSION}" \ + && python3 -c "import torch; assert torch.version.hip is not None, torch.__version__" \ + && groupadd --gid "${RL_KERNEL_GID}" "${RL_KERNEL_USER}" \ + && useradd --uid "${RL_KERNEL_UID}" --gid "${RL_KERNEL_GID}" --create-home --shell /bin/bash "${RL_KERNEL_USER}" \ + && install -d --owner "${RL_KERNEL_USER}" --group "${RL_KERNEL_USER}" /workspace/RL-Kernel \ + && rm -rf /var/lib/apt/lists/* + +ENV HOME=/home/${RL_KERNEL_USER} +WORKDIR /workspace/RL-Kernel +USER ${RL_KERNEL_USER} diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index bedf3475..23c1586a 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -11,6 +11,14 @@ logical type, and the registry selects the first available backend for the curre 4. Cache successfully constructed operator instances. 5. Skip backends that already failed in the current process. +WS2 TP-aware logprob uses the stricter `KernelRegistry.get_logprob_op(contract)` path. In +addition to platform priority, this path requires a backend capability descriptor and checks +the requested role, dtype, TP/CP layout, padded-vs-real vocab masking, inactive-token +support, vocab-domain LSE export, and deterministic TP merge semantics. Incompatible +candidates produce explicit rejection reasons and are never used as an undeclared fallback. +The contract objects and their normative reduction semantics are documented in +`rl_engine.kernels.logprob_contract`. + ## LogP Priority | Platform | Priority | diff --git a/docs/getting_started/faq.md b/docs/getting_started/faq.md index 878c8c14..27cdd0c1 100644 --- a/docs/getting_started/faq.md +++ b/docs/getting_started/faq.md @@ -13,8 +13,8 @@ change more often than RL-Kernel's public API. | --- | --- | --- | --- | | Read docs or edit docs | `pip install -r requirements-docs.txt` | No | Use `mkdocs build --strict -f mkdocs.yaml` before opening a PR. | | Run CPU/mock tests | `pip install -e ".[dev]"` | No | Matches the default CI style: fallback and mocked integration coverage. | -| Run CUDA operators | `pip install -e ".[cuda]"` | Yes, NVIDIA | Requires a CUDA-enabled PyTorch wheel and a working CUDA toolchain for source builds. | -| Run ROCm operators | `pip install -e ".[rocm]"` | Yes, AMD | Requires a ROCm-enabled PyTorch wheel and ROCm compiler/runtime environment. | +| Run CUDA operators | `RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[cuda]"` | Yes, NVIDIA | Requires a CUDA-enabled PyTorch wheel and a working CUDA toolchain. | +| Run ROCm operators | `RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[rocm]"` | Yes, AMD | Requires a ROCm-enabled PyTorch wheel and ROCm compiler/runtime environment. | | Run real vLLM rollout | `pip install -e ".[vllm]"` | Runtime-dependent | Core tests do not need vLLM; install this only where real vLLM is used. | Do not install every optional extra by default. Install the smallest environment @@ -84,7 +84,13 @@ The important rule is that PyTorch must match your runtime: ```bash git clone https://github.com/RL-Align/RL-Kernel.git cd RL-Kernel + +# CPU-only / pure-Python fallback pip install -e . + +# Native CUDA or ROCm extension (after installing a matching PyTorch build) +RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e . +python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)" ``` The examples on this page use `python3` for system-level commands. Inside an @@ -93,9 +99,9 @@ activated virtual environment, `python` is also fine. ### Which optional extras exist? ```bash -pip install -e ".[cuda]" -pip install -e ".[rocm]" -pip install -e ".[vllm]" +RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[cuda]" +RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[rocm]" +pip install --no-build-isolation -e ".[vllm]" pip install -e ".[dev]" ``` diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md index 71fe227d..0e8bad46 100644 --- a/docs/getting_started/installation.md +++ b/docs/getting_started/installation.md @@ -14,15 +14,15 @@ git clone https://github.com/RL-Align/RL-Kernel.git cd RL-Kernel # Optional: pin the compile target. If unset, the build targets your GPU's arch. # export TORCH_CUDA_ARCH_LIST="9.0+PTX" # e.g. Hopper; or "8.6+PTX", "12.0+PTX" -pip install --no-build-isolation -e . +RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e . ``` -Without `--no-build-isolation`, PyTorch is invisible to the isolated build -environment, the extension is silently skipped, and the library falls back to the -slower pure-PyTorch kernels. Confirm the compiled extension is present with: +`RL_KERNEL_REQUIRE_EXT=1` makes the build fail if `_C` cannot be compiled. Without +`--no-build-isolation`, PyTorch is invisible to the isolated build environment. +Confirm the compiled extension is present with: ```bash -python -c "from rl_engine import _C; print('compiled extension OK')" +python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)" ``` A CPU-only install (plain `pip install -e .` on a machine with no GPU) remains @@ -34,15 +34,15 @@ The extras add optional dependencies on top of the compiled package, so they use the same `--no-build-isolation` flag as the source build above. ```bash -pip install --no-build-isolation -e ".[cuda]" +RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e ".[cuda]" ``` ```bash -pip install --no-build-isolation -e ".[rocm]" +RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e ".[rocm]" ``` ```bash -pip install --no-build-isolation -e ".[vllm]" +python -m pip install --no-build-isolation -e ".[vllm]" ``` Install the vLLM extra only on rollout or benchmark environments that need the diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index fbc0e9f1..5bbacaa7 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -54,6 +54,58 @@ CUDA priority list when the extension exposes `_C.batch_invariant_logp_sm90` (built with `KERNEL_ALIGN_FORCE_SM90=1`) on an SM90 device. On any other build or device, dispatch is unchanged (Triton -> PyTorch). +## Tensor Parallel + +`VocabParallelLogprobOp` +(`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) +defines a cross-TP bitwise contract for TP=1, TP=2, and TP=4 when +`num_vocab_tiles` is fixed and every vocabulary-shard boundary is tile-aligned. +The complete BF16 CUDA/NCCL validation matrix for this contract is tracked by +issue #241 PR4. + +1. Split the padded vocabulary into `num_vocab_tiles` fixed tiles. +2. Each rank computes fp32 `(max, sumexp)` for the tiles it owns. Every tile + is reduced as the same contiguous `[n, tile]` shape, on any rank. +3. All tile partials are shared with `all_gather`. The collective only moves + bytes; it never does math, so it cannot round anything. +4. Every rank merges all tiles in the same fixed order, over the same + `[n, num_vocab_tiles]` shape. `LSE = M + log(sum(s_t * exp(m_t - M)))`. +5. The target logit is copied from the rank that owns it (never summed). +6. `logp = target_logit - LSE`. Inactive rows become `0.0`. + +Usage goes through the contract-aware entry point: + +```python +from rl_engine.kernels.registry import kernel_registry + +result = kernel_registry.get_logprob_op(contract) # LogprobContract from +op = result.op # rl_engine.kernels.logprob_contract +logp, lse = op(local_logits, target_ids, contract=contract, tp_group=tp_group) +``` + +### Vime CP=2 runtime provider + +The optional Vime adapter is owned by RL-Kernel and can be selected without +patching Megatron or vLLM: + +```text +--selected-logprob-provider rl_engine.integrations.vime.logp.provider +--selected-logprob-provider-mode strict +``` + +Vime passes the local `[T, V_local]` logits, shifted targets, TP subgroup, +and CP row-ownership metadata. The provider builds the same `LogprobContract` +used by the distributed report, dispatches the explicit +`pytorch-vocab-parallel-logp-ws2` backend, and returns selected logp as `[T, 1]`. +When entropy is requested, it uses the same fixed TP-rank order and returns +full-vocabulary entropy for the existing loss surface. CP rank/layout are +recorded in provenance and never participate in the vocabulary LSE merge. + +The provider fails closed for undeclared real/padded vocabulary sizes, TP/CP +metadata mismatches, unsupported top-p replay masks, and backend fallback. +`auto` mode may then use Vime's native path; `strict` mode reports the +configuration error. This adapter does not import Vime. + ## Benchmarks `benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the @@ -179,6 +231,160 @@ fp16/bf16 backward: checked against fp32 reference with relaxed tolerance CPU-vs-CUDA comparisons use tolerance-based checks; batch-invariance checks within the same backend use exact equality where appropriate. +## TP=1 Comparison Harness + +The single-GPU comparison harness is the TP=1 registration and regression guard +for issue #241. It uses the batch-invariant PyTorch implementation as the +reference and compares exact `pytorch`, `triton`, or `cuda-sm90` backends before +distributed communication is introduced. + +Each backend exposes a diagnostic-only entry point while the production contract +remains unchanged: + +```text +op(logits, target_ids) -> logp +op.forward_with_lse(logits, target_ids) -> (logp, lse) +``` + +The harness reports LSE drift over every logical token row and selected-logprob +drift over active response/action tokens only. Drift summaries contain max, +mean, p95, p99, and the number of compared values. Reports also record requested +and actual backends, implementation, direct-LSE provenance, input shape and +dtype, `tp_world=1`, and `communication=none`. + +Backend selection is exact and does not use registry fallback. In particular, +an explicit `cuda-sm90` comparison fails unless the compiled SM90 extension, +Hopper hardware, input dtype, and vocab row stride satisfy the kernel contract. + +Run the PyTorch TP=1 guard directly from the kernel-specific testing module: + +```bash +python rl_engine/testing/logprob_comparison.py \ + --candidate pytorch \ + --device cpu \ + --dtype fp32 \ + --batch 2 \ + --seq 16 \ + --vocab 257 +``` + +On a GPU, repeat `--candidate` to compare multiple exact backends: + +```bash +python rl_engine/testing/logprob_comparison.py \ + --candidate triton \ + --candidate cuda-sm90 \ + --device cuda \ + --dtype bf16 \ + --batch 2 \ + --seq 16 \ + --vocab 151936 +``` + +The command writes structured JSON to stdout and routes backend diagnostics to +stderr. The harness does not implement vocab sharding, collective communication, +cross-rank LSE merging, or CP reconstruction. + +### SM90 validation + +SM90 validation requires a Hopper GPU, CUDA-enabled PyTorch, and an `nvcc` +toolkit matching `torch.version.cuda`. Build the extension with: + +```bash +export FORCE_CUDA=1 +export KERNEL_ALIGN_FORCE_SM90=1 +export TORCH_CUDA_ARCH_LIST="9.0+PTX" + +python -m pip install --no-build-isolation --no-deps -e . +``` + +Run the focused harness tests, the complete operator suite, and an explicit +SM90 comparison: + +```bash +python -m pytest \ + tests/test_logprob_comparison.py \ + tests/test_operator_inputs.py \ + tests/test_op_checks.py -q + +python -m pytest tests/test_batch_invariant_logp.py -q + +python rl_engine/testing/logprob_comparison.py \ + --candidate cuda-sm90 \ + --device cuda \ + --dtype bf16 \ + --batch 2 \ + --seq 16 \ + --vocab 151936 \ + --prompt-tokens 8 \ + --seed 241 +``` + +The PR2 path was validated on an NVIDIA H800 PCIe with PyTorch 2.11.0+cu128, +CUDA 12.8, and Triton 3.6.0. The focused tests passed 41 cases and the complete +batch-invariant suite passed 67 cases. For BF16 shape `[2, 16, 151936]`, both +LSE and active-token dlogp had maximum absolute drift +`9.5367431640625e-07` against the PyTorch reference, with no backend fallback. + +## Distributed WS2 Drift Report + +The issue #241 PR4 runner materializes one TP/CP topology per `torchrun` +invocation. TP partitions the vocabulary and is the only numerical merge axis; +CP partitions token rows and is recorded in provenance without participating in +the vocab-domain LSE merge. For global rank `r`: + +```text +tp_rank = r % tp_world_size +cp_rank = r // tp_world_size +``` + +Every case generates the same seeded FP32 logical logits, targets, and active +mask. The candidate receives a BF16 token/vocab shard through the explicit +`pytorch-vocab-parallel-logp-ws2` backend, while the independent oracle computes +`torch.logsumexp` over the complete real-vocab FP32 token slice. Distributed +dispatch rejects `auto`, capability fallback, topology mismatches, non-tileable +vocabularies, and incomplete materialization. + +Reports follow the issue #116 fields and contain per-rank and aggregate LSE and +active-token dlogp summaries: max/mean/p95/p99 absolute drift, max relative +drift, worst global token position, target id, target owner rank, #108 tolerance, +and pass/fail. Provenance includes TP/CP topology, dtype, shard bounds, backend +capability, contract fingerprint, reduction spec, merge order, transport, and +the exact launch command. Replicated TP outputs are checked bitwise before one +representative per CP shard is included in aggregate statistics. + +Print the scoped TP=1/2/4 x CP=1/2 launch matrix without starting workers: + +```bash +python rl_engine/testing/distributed_logprob_comparison.py \ + --plan \ + --device cuda \ + --dtype bf16 \ + --output artifacts/ws2-logprob/report.json +``` + +Run one TP=2, CP=2 Qwen3-vocab case on four local GPUs: + +```bash +torchrun --standalone --nproc-per-node=4 \ + rl_engine/testing/distributed_logprob_comparison.py \ + --tp 2 \ + --cp 2 \ + --dtype bf16 \ + --backend pytorch-vocab-parallel-logp-ws2 \ + --real-vocab 151936 \ + --padded-vocab 151936 \ + --num-vocab-tiles 64 \ + --batch 2 \ + --seq 16 \ + --prompt-tokens 8 \ + --output artifacts/ws2-logprob/tp2-cp2.json +``` + +The full matrix requires up to eight ranks for TP=4, CP=2. CPU/Gloo cases are +available for topology and artifact validation; the scoped numerical gate is +BF16 on CUDA/NCCL. + ## Minimal Example ```python @@ -206,11 +412,17 @@ out.sum().backward() python -m pytest tests/test_batch_invariant_logp.py -q -rs ``` -All backends (Native, Triton) are tested in a single file. Coverage includes: +All production backends are tested in a single file. Coverage includes correctness, leading-shape preservation, batch-invariance (bitwise), validation, ignore-index behavior, backward correctness, CUDA smoke cases, registry dispatch, and Triton-specific fp32/fp16/bf16 correctness, large vocab, backward -gradient batch-invariance, and ignored-row zero gradients. +gradient batch-invariance, and ignored-row zero gradients. The focused +`tests/test_logprob_comparison.py` suite covers TP=1 bitwise regression, direct +LSE identity, active-token drift statistics, structured serialization, exact +backend diagnostics, and fail-closed provenance. +`tests/test_distributed_logprob_comparison.py` covers topology planning, TP/CP +rank mapping, token/vocab sharding, explicit backend materialization, #116 JSON +artifacts, and a real four-process TP=2, CP=2 Gloo smoke case. Triton tests skip when Triton or CUDA is unavailable. On Windows, run via WSL/Linux with CUDA. @@ -223,4 +435,11 @@ WSL/Linux with CUDA. - `csrc/cuda/batch_invariant_logp_kernel_sm90.cu` - `rl_engine/kernels/registry.py` - `tests/test_batch_invariant_logp.py` +- `tests/test_logprob_comparison.py` +- `rl_engine/testing/logprob_drift.py` +- `rl_engine/testing/distributed_logprob_comparison.py` +- `tests/test_distributed_logprob_comparison.py` - `benchmarks/benchmark_batch_invariant_logp.py` +- `rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py` +- `rl_engine/kernels/logprob_contract.py` +- `tests/test_vocab_parallel_logp.py` diff --git a/envs.py b/envs.py index 34aded7c..03c65ef2 100644 --- a/envs.py +++ b/envs.py @@ -26,3 +26,4 @@ def env_flag(name: str, default: bool = False) -> bool: KERNEL_ALIGN_NCU_LINEINFO = "KERNEL_ALIGN_NCU_LINEINFO" KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC = "KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC" KERNEL_ALIGN_FORCE_SM90 = "KERNEL_ALIGN_FORCE_SM90" +RL_KERNEL_REQUIRE_EXT = "RL_KERNEL_REQUIRE_EXT" diff --git a/examples/vime_qwen3_8b_tp2_cp2/README.md b/examples/vime_qwen3_8b_tp2_cp2/README.md new file mode 100644 index 00000000..30d3734a --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/README.md @@ -0,0 +1,69 @@ +# Vime Qwen3-8B TP=2 CP=2 validation + +This example is the recommended reproducible entry point for the Vime-side +selected-logprob integration. It keeps framework glue in Vime and keeps the +numerical provider, contract, provenance, and report in RL-Kernel. + +The example is deliberately strict: + +- Megatron training uses `TP=2`, `CP=2`, `PP=1`, and four actor ranks. +- vLLM rollout uses processed logprobs and `top_p=1.0`. +- Vime must load `rl_engine.integrations.vime.logp.provider` in `strict` mode. +- A native fallback or a missing provider marker is not reported as a pass. +- Attention and FFN are not declared consistent from configuration alone. They + require executed Megatron and vLLM readbacks, so the report marks them + `unclaimed` until those artifacts are supplied. The readback must use + `rlkernel.operator_runtime_evidence.v1` and report exact-zero comparison + metrics for both sides. + +The Vime companion must be installed or checked out separately. This example +does not modify `vllm-project/vime`. + +## Dry run + +```bash +python examples/vime_qwen3_8b_tp2_cp2/run.py \ + --vime-root /path/to/RL-Align/vime \ + --rl-kernel-root /path/to/RL-Kernel \ + --output reports/qwen3_8b_tp2_cp2.validation.json +``` + +## Execute + +The Vime script expects model/checkpoint/data paths through environment +variables. Override them before adding `--run`: + +```bash +export MODEL_ROOT=/models/Qwen3-8B +export TORCH_DIST_ROOT=/models/Qwen3-8B_torch_dist +export PROMPT_DATA=/data/dapo-math-17k.jsonl +export RL_KERNEL_ROOT=/path/to/RL-Kernel + +python examples/vime_qwen3_8b_tp2_cp2/run.py \ + --vime-root /path/to/RL-Align/vime \ + --rl-kernel-root "$RL_KERNEL_ROOT" \ + --output reports/qwen3_8b_tp2_cp2.validation.json \ + --run +``` + +When the Megatron/vLLM launch also emits the operator readback artifact, pass +it explicitly: + +```bash +python examples/vime_qwen3_8b_tp2_cp2/run.py \ + --vime-root /path/to/RL-Align/vime \ + --rl-kernel-root "$RL_KERNEL_ROOT" \ + --runtime-evidence reports/qwen3_8b_tp2_cp2.runtime-evidence.json \ + --output reports/qwen3_8b_tp2_cp2.validation.json \ + --run +``` + +The evidence file is intentionally post-execution. It must include training +and rollout identities for `attention` and `ffn`, plus `passed: true` and +exact-zero `out`, backward, and (for attention) `LSE` comparison metrics. A +configured backend without this readback remains `unclaimed`. + +The runner writes a JSON report and a sibling combined log. The report records +the exact command, both repository revisions, provider backend identity, strict +fallback status, and the claim boundary. It does not fabricate numerical drift +when the GPU run was not executed. diff --git a/examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json b/examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json new file mode 100644 index 00000000..93332cbb --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json @@ -0,0 +1,43 @@ +{ + "schema_version": "rlkernel.vime_qwen3_8b_tp2_cp2.v1", + "model": "Qwen/Qwen3-8B", + "training": { + "framework": "megatron", + "tensor_model_parallel_size": 2, + "context_parallel_size": 2, + "pipeline_model_parallel_size": 1, + "world_size": 4, + "dtype": "bf16" + }, + "rollout": { + "framework": "vllm", + "top_p": 1.0, + "logprobs_mode": "processed_logprobs" + }, + "selected_logprob_provider": { + "path": "rl_engine.integrations.vime.logp.provider", + "mode": "strict", + "backend_id": "pytorch-vocab-parallel-logp-ws2", + "real_vocab_size": 151936, + "padded_vocab_size": 152064, + "num_vocab_tiles": 64 + }, + "operator_evidence": { + "logp": { + "training": "rl-kernel-provider", + "rollout": "vllm-native-processed-logprobs", + "required_runtime_marker": "Selected-logprob provider active" + }, + "attention": { + "training": "runtime-readback-required", + "rollout": "runtime-readback-required", + "status": "not_claimed_without_Megatron_and_vLLM_readback" + }, + "ffn": { + "training": "runtime-readback-required", + "rollout": "runtime-readback-required", + "status": "not_claimed_without_Megatron_and_vLLM_readback" + } + }, + "vime_script": "scripts/run-qwen3-8B-rlkernel-tp2-cp2.sh" +} diff --git a/examples/vime_qwen3_8b_tp2_cp2/run.py b/examples/vime_qwen3_8b_tp2_cp2/run.py new file mode 100644 index 00000000..4823dd58 --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/run.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Run and archive the Vime Qwen3-8B TP=2/CP=2 validation entry point. + +This is an integration example, not a synthetic pass generator. A dry run +only records the exact launch contract. ``--run`` executes Vime and records +whether the strict RL-Kernel provider was actually observed in the log. The +report deliberately leaves attention/FFN unclaimed until both framework +readbacks are supplied. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Mapping + + +DEFAULT_CONFIG = Path(__file__).with_name("qwen3_8b_tp2_cp2.json") +PROVIDER_MARKER = "Selected-logprob provider active" +FALLBACK_MARKERS = ("using native path", "fallback=True", "fallback=true") +RUNTIME_EVIDENCE_SCHEMA = "rlkernel.operator_runtime_evidence.v1" +_OPERATOR_METRICS = { + "attention": ("out_max_abs", "lse_max_abs", "dq_max_abs", "dk_max_abs", "dv_max_abs"), + "ffn": ("out_max_abs", "dx_max_abs", "dw_max_abs"), +} + + +def load_config(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict): + raise ValueError("example config must contain a JSON object") + return value + + +def validate_config(config: Mapping[str, Any]) -> None: + training = config.get("training") + rollout = config.get("rollout") + provider = config.get("selected_logprob_provider") + if not isinstance(training, Mapping) or not isinstance(rollout, Mapping) or not isinstance(provider, Mapping): + raise ValueError("training, rollout, and selected_logprob_provider sections are required") + expected = { + "tensor_model_parallel_size": 2, + "context_parallel_size": 2, + "pipeline_model_parallel_size": 1, + "world_size": 4, + } + for name, value in expected.items(): + if training.get(name) != value: + raise ValueError(f"training.{name} must be {value!r}") + if rollout.get("top_p") != 1.0: + raise ValueError("rollout.top_p must remain 1.0 for the strict provider contract") + if provider.get("mode") != "strict": + raise ValueError("selected_logprob_provider.mode must be strict") + if provider.get("path") != "rl_engine.integrations.vime.logp.provider": + raise ValueError("example must use the RL-Kernel Vime provider") + if provider.get("backend_id") != "pytorch-vocab-parallel-logp-ws2": + raise ValueError("example must pin the WS2 vocab-parallel backend") + + +def load_runtime_evidence(path: Path | None) -> dict[str, Any] | None: + """Load post-execution readback without treating configuration as evidence.""" + + if path is None: + return None + with path.open(encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict) or value.get("schema_version") != RUNTIME_EVIDENCE_SCHEMA: + raise ValueError(f"runtime evidence must use schema {RUNTIME_EVIDENCE_SCHEMA!r}") + return value + + +def _operator_evidence_status(evidence: Mapping[str, Any] | None, operator: str) -> str: + if evidence is None: + return "unclaimed" + operators = evidence.get("operators") + item = operators.get(operator) if isinstance(operators, Mapping) else None + if not isinstance(item, Mapping): + return "unclaimed" + training = item.get("training") + rollout = item.get("rollout") + comparison = item.get("comparison") + if not isinstance(training, Mapping) or not isinstance(rollout, Mapping): + return "unclaimed" + if not isinstance(comparison, Mapping) or comparison.get("passed") is not True: + return "failed" + required_identity = ("implementation_id", "backend_id", "contract_id") + if any(not training.get(name) or not rollout.get(name) for name in required_identity): + return "failed" + if training["implementation_id"] != rollout["implementation_id"]: + return "failed" + for metric in _OPERATOR_METRICS[operator]: + value = comparison.get(metric) + if not isinstance(value, (int, float)) or isinstance(value, bool) or value != 0.0: + return "failed" + return "passed" + + +def validate_runtime_evidence(evidence: Mapping[str, Any] | None) -> None: + """Reject malformed evidence before it can affect a report.""" + + if evidence is None: + return + for operator in _OPERATOR_METRICS: + status = _operator_evidence_status(evidence, operator) + if status == "failed": + raise ValueError(f"runtime evidence for {operator} is incomplete or non-zero") + + +def _revision(path: Path) -> str | None: + try: + result = subprocess.run( + ["git", "-C", str(path), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return None + return result.stdout.strip() or None + + +def build_environment(vime_root: Path, rl_kernel_root: Path) -> dict[str, str]: + env = dict(os.environ) + existing = [str(vime_root), str(rl_kernel_root), "/root/Megatron-LM"] + if env.get("PYTHONPATH"): + existing.append(env["PYTHONPATH"]) + env["PYTHONPATH"] = os.pathsep.join(existing) + env["RL_KERNEL_ROOT"] = str(rl_kernel_root) + env["TP_SIZE"] = "2" + env["CP_SIZE"] = "2" + env["ROLLOUT_TOP_P"] = "1.0" + return env + + +def build_command(config: Mapping[str, Any], vime_root: Path) -> list[str]: + script = vime_root / str(config.get("vime_script", "")) + if not script.is_file(): + raise FileNotFoundError(f"Vime entry script does not exist: {script}") + return ["bash", str(script)] + + +def build_report( + config: Mapping[str, Any], + *, + vime_root: Path, + rl_kernel_root: Path, + command: list[str], + status: str, + returncode: int | None, + log_text: str, + log_path: Path | None, + runtime_evidence: Mapping[str, Any] | None = None, + runtime_evidence_path: Path | None = None, +) -> dict[str, Any]: + provider_active = PROVIDER_MARKER in log_text + fallback_observed = any(marker in log_text for marker in FALLBACK_MARKERS) + strict_provider_passed = status == "passed" and provider_active and not fallback_observed + effective_status = "passed" if strict_provider_passed else ("failed" if status == "passed" else status) + attention_status = _operator_evidence_status(runtime_evidence, "attention") + ffn_status = _operator_evidence_status(runtime_evidence, "ffn") + return { + "schema_version": "rlkernel.vime_validation_report.v1", + "created_at": datetime.now(timezone.utc).isoformat(), + "status": effective_status, + "claim_boundary": { + "qwen3_8b_tp2_cp2_vime_training": strict_provider_passed, + "attention_train_infer_consistency": attention_status, + "ffn_train_infer_consistency": ffn_status, + "reason": ( + "attention and FFN require executed Megatron/vLLM runtime readbacks; " + "the evidence contract accepts only exact-zero comparison metrics" + ), + }, + "config": dict(config), + "topology": config["training"], + "provider": { + "configured_path": config["selected_logprob_provider"]["path"], + "configured_mode": config["selected_logprob_provider"]["mode"], + "backend_id": config["selected_logprob_provider"]["backend_id"], + "active_observed": provider_active, + "fallback_observed": fallback_observed, + }, + "command": command, + "returncode": returncode, + "artifacts": { + "log": None if log_path is None else str(log_path), + "runtime_evidence": None if runtime_evidence_path is None else str(runtime_evidence_path), + }, + "runtime_evidence": None if runtime_evidence is None else dict(runtime_evidence), + "revisions": { + "vime": _revision(vime_root), + "rl_kernel": _revision(rl_kernel_root), + }, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG) + parser.add_argument("--vime-root", type=Path, default=Path(os.environ.get("VIME_ROOT", "."))) + parser.add_argument("--rl-kernel-root", type=Path, default=Path(os.environ.get("RL_KERNEL_ROOT", "."))) + parser.add_argument("--output", type=Path, default=Path("qwen3_8b_tp2_cp2.validation.json")) + parser.add_argument( + "--runtime-evidence", + type=Path, + default=None, + help="post-execution Megatron/vLLM operator readback JSON (strict exact-zero contract)", + ) + parser.add_argument("--run", action="store_true", help="execute the Vime script") + args = parser.parse_args(argv) + + config = load_config(args.config) + validate_config(config) + runtime_evidence = load_runtime_evidence(args.runtime_evidence) + validate_runtime_evidence(runtime_evidence) + vime_root = args.vime_root.resolve() + rl_kernel_root = args.rl_kernel_root.resolve() + command = build_command(config, vime_root) + + status = "not_run" + returncode: int | None = None + log_text = "" + log_path: Path | None = None + if args.run: + args.output.parent.mkdir(parents=True, exist_ok=True) + log_path = args.output.with_suffix(".log") + env = build_environment(vime_root, rl_kernel_root) + with log_path.open("w", encoding="utf-8") as log_handle: + process = subprocess.run(command, cwd=vime_root, env=env, stdout=log_handle, stderr=subprocess.STDOUT) + returncode = process.returncode + log_text = log_path.read_text(encoding="utf-8", errors="replace") + status = "passed" if returncode == 0 else "failed" + + report = build_report( + config, + vime_root=vime_root, + rl_kernel_root=rl_kernel_root, + command=command, + status=status, + returncode=returncode, + log_text=log_text, + log_path=log_path, + runtime_evidence=runtime_evidence, + runtime_evidence_path=args.runtime_evidence, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if report["status"] in {"passed", "not_run"} else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/rl_engine/__init__.py b/rl_engine/__init__.py index e69de29b..306d6813 100644 --- a/rl_engine/__init__.py +++ b/rl_engine/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import torch # noqa: F401 # Load torch shared libraries before importing rl_engine._C. diff --git a/rl_engine/integrations/__init__.py b/rl_engine/integrations/__init__.py new file mode 100644 index 00000000..29bcc68f --- /dev/null +++ b/rl_engine/integrations/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Optional framework adapters owned by RL-Kernel.""" + diff --git a/rl_engine/integrations/vime/__init__.py b/rl_engine/integrations/vime/__init__.py new file mode 100644 index 00000000..44470314 --- /dev/null +++ b/rl_engine/integrations/vime/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Vime adapter entry points without a Vime runtime dependency.""" + +from .logp import ProviderResult, SelectedLogprobProviderUnavailable, provider + +__all__ = ["ProviderResult", "SelectedLogprobProviderUnavailable", "provider"] + diff --git a/rl_engine/integrations/vime/logp.py b/rl_engine/integrations/vime/logp.py new file mode 100644 index 00000000..db436d69 --- /dev/null +++ b/rl_engine/integrations/vime/logp.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime selected-logprob provider for Vime's Megatron backend. + +The adapter intentionally accepts and returns structural objects: RL-Kernel +never imports Vime. Vime remains responsible for constructing locally owned +CP token rows and response masks; this provider owns only the TP-vocabulary +reduction. CP rank and layout are recorded and validated as row ownership +metadata, never passed to the numerical merge. +""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch + +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 kernel_registry + + +class SelectedLogprobProviderUnavailable(RuntimeError): + """Request Vime's native provider fallback in ``auto`` mode. + + Vime recognizes the marker instead of importing this class, which keeps + the dependency direction from Vime to RL-Kernel at runtime only. + """ + + selected_logprob_provider_unavailable = True + + +@dataclass(frozen=True) +class ProviderResult: + """Structural result understood by the Vime provider boundary.""" + + selected_logprobs: torch.Tensor + entropy: torch.Tensor | None + backend_id: str + contract_id: str + provenance: Mapping[str, Any] + + +def _as_positive_int(value: Any, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise SelectedLogprobProviderUnavailable(f"{name} must be a positive integer; got {value!r}") + return value + + +def _metadata(request: Any) -> Mapping[str, Any]: + value = getattr(request, "metadata", None) + if not isinstance(value, Mapping): + raise SelectedLogprobProviderUnavailable("request.metadata must provide vocab-parallel metadata") + return value + + +def _request_tensor(request: Any, name: str) -> torch.Tensor: + value = getattr(request, name, None) + if not isinstance(value, torch.Tensor): + raise SelectedLogprobProviderUnavailable(f"request.{name} must be a torch.Tensor") + return value + + +def _tp_coordinates(tp_group: Any) -> tuple[int, int]: + if tp_group is not None and hasattr(tp_group, "rank") and hasattr(tp_group, "size"): + return int(tp_group.rank()), int(tp_group.size()) + + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank(group=tp_group), dist.get_world_size(group=tp_group) + return 0, 1 + + +def _tile_count(metadata: Mapping[str, Any], padded_vocab_size: int) -> int: + configured = metadata.get("num_vocab_tiles", os.getenv("RL_KERNEL_LOGPROB_NUM_VOCAB_TILES")) + if configured is None or configured == "": + configured = DEFAULT_NUM_VOCAB_TILES + try: + tiles = int(configured) + except (TypeError, ValueError) as exc: + raise SelectedLogprobProviderUnavailable( + f"num_vocab_tiles must be an integer; got {configured!r}" + ) from exc + if tiles <= 0 or padded_vocab_size % tiles: + raise SelectedLogprobProviderUnavailable( + f"num_vocab_tiles={tiles} must divide padded_vocab_size={padded_vocab_size}" + ) + return tiles + + +def _contract_for_request(request: Any) -> tuple[LogprobContract, int]: + logits = _request_tensor(request, "logits") + targets = _request_tensor(request, "target_ids") + metadata = _metadata(request) + if logits.ndim != 2 or targets.shape != (logits.shape[0],): + raise SelectedLogprobProviderUnavailable("request must contain local [T, V] logits and aligned [T] targets") + if logits.dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise SelectedLogprobProviderUnavailable(f"unsupported logit dtype {logits.dtype}") + if targets.device != logits.device: + raise SelectedLogprobProviderUnavailable("target_ids must share the local logits device") + + cp = getattr(request, "context_parallel", None) + cp_world_size = _as_positive_int(getattr(cp, "world_size", None), "context_parallel.world_size") + cp_rank = getattr(cp, "rank", None) + if isinstance(cp_rank, bool) or not isinstance(cp_rank, int) or not 0 <= cp_rank < cp_world_size: + raise SelectedLogprobProviderUnavailable( + f"context_parallel.rank={cp_rank!r} is invalid for CP={cp_world_size}" + ) + if getattr(cp, "layout", None) not in ({"single"} if cp_world_size == 1 else {"zigzag", "allgather"}): + raise SelectedLogprobProviderUnavailable("context_parallel layout does not describe local CP token ownership") + + tp_rank, tp_world_size = _tp_coordinates(getattr(request, "tensor_parallel_group", None)) + declared_tp_rank = metadata.get("tp_rank") + declared_tp_world_size = metadata.get("tp_world_size") + if declared_tp_rank is not None and declared_tp_rank != tp_rank: + raise SelectedLogprobProviderUnavailable( + f"metadata tp_rank={declared_tp_rank} disagrees with TP group rank={tp_rank}" + ) + if declared_tp_world_size is not None and declared_tp_world_size != tp_world_size: + raise SelectedLogprobProviderUnavailable( + f"metadata tp_world_size={declared_tp_world_size} disagrees with TP group size={tp_world_size}" + ) + + real_vocab_size = _as_positive_int(metadata.get("real_vocab_size"), "real_vocab_size") + padded_vocab_size = _as_positive_int(metadata.get("padded_vocab_size"), "padded_vocab_size") + if logits.shape[1] * tp_world_size != padded_vocab_size: + raise SelectedLogprobProviderUnavailable( + "local vocab width and TP group do not cover padded_vocab_size exactly: " + f"{logits.shape[1]} * {tp_world_size} != {padded_vocab_size}" + ) + if real_vocab_size > padded_vocab_size: + raise SelectedLogprobProviderUnavailable("real_vocab_size must not exceed padded_vocab_size") + + bounds = tuple( + (rank * logits.shape[1], (rank + 1) * logits.shape[1]) for rank in range(tp_world_size) + ) + active_mask = (True,) * logits.shape[0] + contract = LogprobContract( + role=LogprobRole.TRAIN, + dtype={ + torch.bfloat16: LogprobDType.BF16, + torch.float16: LogprobDType.FP16, + torch.float32: LogprobDType.FP32, + }[logits.dtype], + mask=MaskSpec(num_tokens=logits.shape[0], active_mask=active_mask), + sharding=ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=bounds, + real_vocab_size=real_vocab_size, + padded_vocab_size=padded_vocab_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ), + reduction=ReductionSpec(), + ) + return contract, _tile_count(metadata, padded_vocab_size) + + +def provider(request: Any) -> ProviderResult: + """Compute Vime selected logprobs on the explicit WS2 TP/CP contract. + + Top-p replay is deliberately unavailable until it has a separately + validated fixed-order mask contract. In Vime ``auto`` mode this signals + native execution; in ``strict`` mode it fails instead of changing sampled + distribution semantics. + """ + + if getattr(request, "log_prob_keep_mask", None) is not None: + raise SelectedLogprobProviderUnavailable( + "RL-Kernel WS2 logprob does not yet materialize Vime top-p replay masks" + ) + + contract, num_vocab_tiles = _contract_for_request(request) + dispatch = kernel_registry.get_logprob_op(contract, requested_backend=BACKEND_ID) + if dispatch.provenance["actual_backend"] != BACKEND_ID or dispatch.provenance["fallback"]: + raise RuntimeError("explicit WS2 backend dispatch changed during materialization") + if getattr(request, "with_entropy", False): + selected_logp, _lse, entropy = dispatch.op.apply_with_entropy( + request.logits, + request.target_ids, + contract=contract, + tp_group=getattr(request, "tensor_parallel_group", None), + num_vocab_tiles=num_vocab_tiles, + with_entropy_grad=bool(getattr(request, "with_entropy_grad", False)), + ) + else: + selected_logp, _lse = dispatch.op( + request.logits, + request.target_ids, + contract=contract, + tp_group=getattr(request, "tensor_parallel_group", None), + num_vocab_tiles=num_vocab_tiles, + ) + entropy = None + provenance = dict(dispatch.provenance) + provenance["request"] = { + "logits_shape": list(request.logits.shape), + "logits_dtype": str(request.logits.dtype).replace("torch.", ""), + "target_shape": list(request.target_ids.shape), + "target_dtype": str(request.target_ids.dtype).replace("torch.", ""), + "real_vocab_size": contract.sharding.real_vocab_size, + "padded_vocab_size": contract.sharding.padded_vocab_size, + "tp_rank": contract.sharding.tp_rank, + "tp_world_size": contract.sharding.tp_world_size, + "cp_rank": contract.sharding.cp_rank, + "cp_world_size": contract.sharding.cp_world_size, + } + provenance["execution"] = { + "role": "vime_training_selected_logprob", + "strict_backend": True, + "top_p_replay": False, + } + provenance["cp_row_ownership"] = { + "cp_rank": contract.sharding.cp_rank, + "cp_world_size": contract.sharding.cp_world_size, + "layout": getattr(request.context_parallel, "layout"), + "local_token_rows": int(request.logits.shape[0]), + "cp_is_merge_axis": False, + } + provenance["num_vocab_tiles"] = num_vocab_tiles + return ProviderResult( + selected_logprobs=selected_logp.unsqueeze(-1), + entropy=entropy, + backend_id=dispatch.capability.backend_id, + contract_id=contract.cross_rank_fingerprint(), + provenance=provenance, + ) + + +__all__ = ["ProviderResult", "SelectedLogprobProviderUnavailable", "provider"] diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py new file mode 100644 index 00000000..1267740a --- /dev/null +++ b/rl_engine/kernels/logprob_contract.py @@ -0,0 +1,655 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for TP-aware selected-token log-probability. + +The objects in this module describe a vocab-parallel logprob invocation: + +``selected_logp[t] = logits[t, target[t]] - logsumexp_vocab(logits[t, :])`` + +Under vocab-parallel tensor parallelism the vocabulary-wide ``logsumexp`` +requires cross-rank reduction. This module only *describes* that invocation +(shard ownership, merge semantics, mask/ignore-index metadata); it does not +shard tensors, launch collectives, or implement the ``(max, sumexp)`` merge. +Keeping description and materialization separate lets dispatch reject an +incompatible backend before any numerically different path is launched. + +Context parallelism is a declared non-merge axis: CP partitions tokens, never +the vocabulary, so the logprob reduction spans TP vocab shards only. CP rank +metadata is carried for provenance and must never widen the merge. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + +# Policy keywords accepted by KernelRegistry.get_logprob_op; a backend id must +# never shadow one of these, or it becomes unselectable by id. +RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) +# Backend tiers; determinism is a separate axis (DeterminismScope). +IMPLEMENTATION_KINDS = frozenset({"production", "reference"}) + + +class LogprobContractError(ValueError): + """Raised when logprob metadata does not describe a valid invocation.""" + + +class LogprobRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class LogprobDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class LogprobMerge(str, Enum): + """Merge primitive for per-shard ``(local_max, local_sumexp)`` partials.""" + + MAX_SUMEXP = "max_sumexp" + + +class MergeAxis(str, Enum): + """The only reduction axis of this contract; CP is a non-merge axis.""" + + TP_VOCAB = "tp_vocab" + + +class ReductionOrder(str, Enum): + GLOBAL_VOCAB_SHARD_INDEX = "global_vocab_shard_index" + + +class ReductionTransport(str, Enum): + """Collectives move partial states only; they never reduce numerically.""" + + ALL_GATHER = "all_gather" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class DeterminismScope(str, Enum): + """Strength of the reduction's determinism guarantee. + + ``fixed_topology``: bitwise-reproducible for one fixed TP degree; results + at different TP degrees are compared against the #108 tolerance table. + + ``cross_tp_bitwise``: additionally bitwise-equal across TP degrees. This + requires the entire reduction to follow a global tile-level structure that + is independent of TP partitioning: a fixed tile decomposition of the + vocabulary plus a fixed merge order and rescaling tree over those tiles, + identical at every TP degree, so the TP degree only selects which rank + computes which tiles and never changes the floating-point grouping. + Fixed shard-order merging alone is not sufficient, because shard + boundaries would still group the combines differently across degrees. + """ + + FIXED_TOPOLOGY = "fixed_topology" + CROSS_TP_BITWISE = "cross_tp_bitwise" + + +class MaskMode(str, Enum): + """How a backend consumes inactive-token information. + + The contract permits inactive targets that do not hold ``ignore_index``, + so an ``ignore_index``-only backend cannot serve a contract with inactive + tokens. + """ + + EXPLICIT_ACTIVE_MASK = "explicit_active_mask" + IGNORE_INDEX = "ignore_index" + + +class TPPlacement(str, Enum): + REPLICATED = "replicated" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise LogprobContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise LogprobContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _plain_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise LogprobContractError(f"{field} must be an integer; got {value!r}") + return value + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical vocab-parallel TP ownership for one logprob invocation. + + ``vocab_shard_bounds`` lists every TP rank's half-open ``[start, end)`` + vocab range, indexed by rank; the full table is required on every rank and + must form a contiguous ``[0, padded_vocab_size)`` partition. + ``padded_vocab_size`` is the shard-covered (weight) vocabulary, + ``real_vocab_size`` the tokenizer vocabulary; padding columns occupy + ``[real_vocab_size, padded_vocab_size)``. + """ + + tp_rank: int + tp_world_size: int + vocab_shard_bounds: tuple[tuple[int, int], ...] + real_vocab_size: int + padded_vocab_size: int + cp_rank: int = 0 + cp_world_size: int = 1 + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + if tp_rank >= tp_world_size: + raise LogprobContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if cp_rank >= cp_world_size: + raise LogprobContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + real_vocab_size = _positive_int(self.real_vocab_size, "real_vocab_size") + padded_vocab_size = _positive_int(self.padded_vocab_size, "padded_vocab_size") + if padded_vocab_size < real_vocab_size: + raise LogprobContractError( + f"padded_vocab_size={padded_vocab_size} must not be smaller than " + f"real_vocab_size={real_vocab_size}" + ) + + try: + bounds = tuple((pair[0], pair[1]) for pair in self.vocab_shard_bounds) + except (TypeError, IndexError) as exc: + raise LogprobContractError( + "vocab_shard_bounds must be an iterable of (start, end) integer pairs" + ) from exc + if len(bounds) != tp_world_size: + raise LogprobContractError( + "vocab_shard_bounds must declare exactly one (start, end) pair per TP rank; " + f"got {len(bounds)} pairs for tp_world_size={tp_world_size}" + ) + expected_start = 0 + for rank, (start, end) in enumerate(bounds): + start = _plain_int(start, f"vocab_shard_bounds[{rank}][0]") + end = _plain_int(end, f"vocab_shard_bounds[{rank}][1]") + if end <= start: + raise LogprobContractError( + f"vocab_shard_bounds[{rank}] must satisfy end > start; got [{start}, {end})" + ) + if start != expected_start: + raise LogprobContractError( + "vocab_shard_bounds must form a contiguous [0, padded_vocab_size) " + f"partition in TP-rank order; rank {rank} starts at {start}, " + f"expected {expected_start}" + ) + expected_start = end + if expected_start != padded_vocab_size: + raise LogprobContractError( + "vocab_shard_bounds must cover padded_vocab_size exactly; " + f"covered {expected_start}, declared {padded_vocab_size}" + ) + object.__setattr__(self, "vocab_shard_bounds", bounds) + + @property + def local_vocab_start(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][0] + + @property + def local_vocab_end(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][1] + + @property + def local_vocab_size(self) -> int: + start, end = self.vocab_shard_bounds[self.tp_rank] + return end - start + + def owner_rank(self, token_id: int) -> int: + """Return the unique TP rank owning ``token_id``; error outside real vocab.""" + + token_id = _plain_int(token_id, "token_id") + if token_id < 0 or token_id >= self.real_vocab_size: + raise LogprobContractError( + f"token_id={token_id} is outside the real vocabulary " + f"[0, {self.real_vocab_size}); mask it as inactive instead" + ) + for rank, (start, end) in enumerate(self.vocab_shard_bounds): + if start <= token_id < end: + return rank + raise LogprobContractError( + f"token_id={token_id} is not covered by any declared vocab shard" + ) + + +@dataclass(frozen=True) +class MaskSpec: + """Active-token mask and ignore index for one logprob invocation. + + Inactive tokens are excluded from drift aggregates and from the + single-owner target gather; their targets may legally hold ``ignore_index``. + """ + + num_tokens: int + active_mask: tuple[bool, ...] + ignore_index: int = -100 + _active_token_count: int = field(init=False, repr=False, compare=False) + _active_mask_sha256: str = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + num_tokens = _positive_int(self.num_tokens, "num_tokens") + _plain_int(self.ignore_index, "ignore_index") + try: + active_mask = tuple(self.active_mask) + except TypeError as exc: + raise LogprobContractError("active_mask must be an iterable of booleans") from exc + for index, value in enumerate(active_mask): + if not isinstance(value, bool): + raise LogprobContractError(f"active_mask[{index}] must be a bool; got {value!r}") + if len(active_mask) != num_tokens: + raise LogprobContractError( + "active_mask must contain exactly one entry per token; " + f"got {len(active_mask)} entries for num_tokens={num_tokens}" + ) + object.__setattr__(self, "active_mask", active_mask) + object.__setattr__(self, "_active_token_count", sum(active_mask)) + object.__setattr__( + self, "_active_mask_sha256", hashlib.sha256(bytes(active_mask)).hexdigest() + ) + + @property + def active_token_count(self) -> int: + return self._active_token_count + + @property + def active_mask_sha256(self) -> str: + """Compact mask identity for provenance and cross-rank agreement.""" + return self._active_mask_sha256 + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic TP-vocab ``(max, sumexp)`` merge semantics. + + Every rank first masks local columns whose global id lies in + ``[real_vocab_size, padded_vocab_size)`` to ``-inf`` (padding never + contributes to the logsumexp), then computes ``m_l = max(local_logits)`` + and ``s_l = sum(exp(local_logits - m_l))`` in fp32. Partials travel by + all-gather -- collectives are transport only, never a numerical + reduction -- and every rank merges in fixed global vocab-shard index + order:: + + M = max_l(m_l) + S = sum_l(s_l * exp(m_l - M)) + LSE = M + log(S) + selected_logp = target_logit - LSE + + The selected target logit comes from a masked single-owner gather; + downcast happens only at the final write. The identity partial for a + padding-only shard, or a row whose local columns are all ``-inf`` after + masking, is ``(m_l, s_l) = (-inf, 0)``: a partial with ``s_l = 0`` + contributes nothing to the merge regardless of its ``m_l``, and + implementations must use this identity directly rather than evaluate + ``exp(-inf - (-inf))``, which would poison the merge with NaN. Averaging + per-rank logsumexp values, or letting a collective reduce numerically, is + never conformant at either determinism scope. + """ + + merge: LogprobMerge = LogprobMerge.MAX_SUMEXP + merge_axis: MergeAxis = MergeAxis.TP_VOCAB + acc_dtype: LogprobDType = LogprobDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_VOCAB_SHARD_INDEX + transport: ReductionTransport = ReductionTransport.ALL_GATHER + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + determinism_scope: DeterminismScope = DeterminismScope.CROSS_TP_BITWISE + + def __post_init__(self) -> None: + object.__setattr__( + self, + "determinism_scope", + _enum_value(DeterminismScope, self.determinism_scope, "determinism_scope"), + ) + object.__setattr__(self, "merge", _enum_value(LogprobMerge, self.merge, "merge")) + object.__setattr__( + self, "merge_axis", _enum_value(MergeAxis, self.merge_axis, "merge_axis") + ) + object.__setattr__( + self, "acc_dtype", _enum_value(LogprobDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "transport", _enum_value(ReductionTransport, self.transport, "transport") + ) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"TP logprob accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class LogprobOutputSpec: + """Output surface every conforming backend must produce: fp32 selected + logprob and fp32 vocab-domain LSE, replicated across the TP group.""" + + selected_logp_dtype: LogprobDType = LogprobDType.FP32 + lse_dtype: LogprobDType = LogprobDType.FP32 + tp_placement: TPPlacement = TPPlacement.REPLICATED + + def __post_init__(self) -> None: + object.__setattr__( + self, + "selected_logp_dtype", + _enum_value(LogprobDType, self.selected_logp_dtype, "selected_logp_dtype"), + ) + object.__setattr__( + self, "lse_dtype", _enum_value(LogprobDType, self.lse_dtype, "lse_dtype") + ) + object.__setattr__( + self, "tp_placement", _enum_value(TPPlacement, self.tp_placement, "tp_placement") + ) + if self.selected_logp_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"selected logprob output must be fp32; got {self.selected_logp_dtype.value}" + ) + if self.lse_dtype is not LogprobDType.FP32: + raise LogprobContractError(f"vocab LSE output must be fp32; got {self.lse_dtype.value}") + + +@dataclass(frozen=True) +class LogprobContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: LogprobRole + dtype: LogprobDType + mask: MaskSpec + sharding: ShardingSpec + reduction: ReductionSpec + output: LogprobOutputSpec = field(default_factory=LogprobOutputSpec) + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(LogprobRole, self.role, "role")) + object.__setattr__(self, "dtype", _enum_value(LogprobDType, self.dtype, "dtype")) + if not isinstance(self.mask, MaskSpec): + raise LogprobContractError("mask must be a MaskSpec") + if not isinstance(self.sharding, ShardingSpec): + raise LogprobContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise LogprobContractError("reduction must be a ReductionSpec") + if not isinstance(self.output, LogprobOutputSpec): + raise LogprobContractError("output must be a LogprobOutputSpec") + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise LogprobContractError( + "export_lse must be True for the WS2 vocab-domain LSE drift contract" + ) + if 0 <= self.mask.ignore_index < self.sharding.real_vocab_size: + raise LogprobContractError( + f"ignore_index={self.mask.ignore_index} must not collide with the real " + f"vocabulary [0, {self.sharding.real_vocab_size})" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "vocab_shard_bounds": [list(pair) for pair in self.sharding.vocab_shard_bounds], + "real_vocab_size": self.sharding.real_vocab_size, + "padded_vocab_size": self.sharding.padded_vocab_size, + "local_vocab_start": self.sharding.local_vocab_start, + "local_vocab_end": self.sharding.local_vocab_end, + } + reduction = { + "merge": self.reduction.merge.value, + "merge_axis": self.reduction.merge_axis.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "transport": self.reduction.transport.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + "determinism_scope": self.reduction.determinism_scope.value, + "cp_is_merge_axis": False, + } + # The digest stands in for the raw per-token mask, which would + # dominate the provenance size. + mask = { + "num_tokens": self.mask.num_tokens, + "active_token_count": self.mask.active_token_count, + "active_mask_sha256": self.mask.active_mask_sha256, + "ignore_index": self.mask.ignore_index, + } + output = { + "selected_logp_dtype": self.output.selected_logp_dtype.value, + "lse_dtype": self.output.lse_dtype.value, + "tp_placement": self.output.tp_placement.value, + } + return { + "semantic_operator": "selected_token_logprob", + "role": self.role.value, + "dtype": self.dtype.value, + "export_lse": self.export_lse, + "lse_domain": "vocab", + "mask": mask, + "sharding": sharding, + "reduction": reduction, + "output": output, + } + + def cross_rank_fingerprint(self) -> str: + """Rank-independent identity for preflight agreement across ranks. + + Excludes ``tp_rank``/``cp_rank`` (and their derived local bounds) so + every rank of one logical invocation computes the same value. + All-gathering this fingerprint together with the resolved backend id + and aborting on mismatch is the documented preflight for distributed + dispatch; ``requested_backend="auto"`` is not distributed-safe + without it. + """ + + payload = self.to_dict() + payload["sharding"] = { + key: value + for key, value in payload["sharding"].items() + if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} + } + # Note: Any future extensions to this payload MUST maintain strict JSON + # serialization determinism across environments to prevent cross-rank + # hashing mismatches. + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class LogprobBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[LogprobRole] + dtypes: frozenset[LogprobDType] + tp_world_sizes: tuple[int, ...] | None = None + cp_world_sizes: tuple[int, ...] | None = None + supports_vocab_padding: bool = False + mask_modes: frozenset[MaskMode] = frozenset() + exports_vocab_lse: bool = False + determinism_scopes: frozenset[DeterminismScope] = frozenset() + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise LogprobContractError("backend_id must be a non-empty string") + if self.backend_id.strip().lower() in RESERVED_DISPATCH_POLICIES: + raise LogprobContractError( + f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" + ) + object.__setattr__(self, "backend_id", self.backend_id.strip()) + try: + roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) + dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) + except TypeError as exc: + raise LogprobContractError("roles and dtypes must be iterables of enum values") from exc + if not roles or not dtypes: + raise LogprobContractError("backend roles and dtypes must not be empty") + tp_world_sizes = self._validated_world_sizes(self.tp_world_sizes, "tp_world_sizes") + cp_world_sizes = self._validated_world_sizes(self.cp_world_sizes, "cp_world_sizes") + try: + mask_modes = frozenset( + _enum_value(MaskMode, value, "mask_modes") for value in self.mask_modes + ) + determinism_scopes = frozenset( + _enum_value(DeterminismScope, value, "determinism_scopes") + for value in self.determinism_scopes + ) + except TypeError as exc: + raise LogprobContractError( + "mask_modes and determinism_scopes must be iterables of enum values" + ) from exc + for flag_name in ("supports_vocab_padding", "exports_vocab_lse"): + if not isinstance(getattr(self, flag_name), bool): + raise LogprobContractError(f"{flag_name} must be a bool") + if self.implementation_kind not in IMPLEMENTATION_KINDS: + raise LogprobContractError( + f"implementation_kind must be one of: {', '.join(sorted(IMPLEMENTATION_KINDS))}" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "mask_modes", mask_modes) + object.__setattr__(self, "determinism_scopes", determinism_scopes) + + @staticmethod + def _validated_world_sizes( + values: tuple[int, ...] | None, field: str + ) -> tuple[int, ...] | None: + if values is None: + return None + try: + sizes = tuple(values) + except TypeError as exc: + raise LogprobContractError(f"{field} must be an iterable of integers") from exc + if not sizes: + raise LogprobContractError(f"{field} must not be empty; use None for unrestricted") + for value in sizes: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must contain positive values; got {value!r}") + if len(set(sizes)) != len(sizes): + raise LogprobContractError(f"{field} must not contain duplicates") + return sizes + + def incompatibilities(self, contract: LogprobContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if self.cp_world_sizes is not None and cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if ( + contract.sharding.padded_vocab_size != contract.sharding.real_vocab_size + and not self.supports_vocab_padding + ): + reasons.append("padded-vs-real vocab masking is unsupported") + if ( + contract.mask.active_token_count != contract.mask.num_tokens + and MaskMode.EXPLICIT_ACTIVE_MASK not in self.mask_modes + ): + # Inactive targets need not hold ignore_index (see MaskMode). + reasons.append("explicit active-token masking is unsupported") + if contract.export_lse and not self.exports_vocab_lse: + reasons.append("vocab-domain LSE export is unsupported") + if contract.reduction.determinism_scope not in self.determinism_scopes: + reasons.append( + f"determinism_scope={contract.reduction.determinism_scope.value} is unsupported" + ) + return tuple(reasons) + + def supports(self, contract: LogprobContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes) if self.cp_world_sizes else None, + "supports_vocab_padding": self.supports_vocab_padding, + "mask_modes": sorted(mode.value for mode in self.mask_modes), + "exports_vocab_lse": self.exports_vocab_lse, + "determinism_scopes": sorted(scope.value for scope in self.determinism_scopes), + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class LogprobDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: LogprobBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "IMPLEMENTATION_KINDS", + "RESERVED_DISPATCH_POLICIES", + "DeterminismScope", + "DowncastPoint", + "LogprobBackendCapability", + "LogprobContract", + "LogprobContractError", + "LogprobDType", + "LogprobDispatchResult", + "LogprobMerge", + "LogprobOutputSpec", + "LogprobRole", + "MaskMode", + "MaskSpec", + "MergeAxis", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "ReductionTransport", + "ShardingSpec", + "TPPlacement", +] diff --git a/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py index a68781d9..aa62b6ef 100644 --- a/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py @@ -161,3 +161,48 @@ def apply( ) return _BatchInvariantLogpSM90Function.apply(logits, target_ids, ignore_index) + + def forward_with_lse( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the exact SM90 path and return its direct FP32 logprob/LSE outputs. + + Unlike the production ``apply`` method, this diagnostic entry point never + falls back to Triton or PyTorch, so comparison provenance stays truthful. + """ + if logits.dim() < 2: + raise ValueError( + f"logits must be at least 2-D ([*lead, V]), got shape {tuple(logits.shape)}" + ) + if logits.shape[:-1] != target_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + if not _sm90_supported(logits): + raise RuntimeError( + "exact cuda-sm90 logprob diagnostics require Hopper, CUDA BF16/FP32 logits, " + "and a 16-byte-aligned vocab row stride; fallback is disabled" + ) + if validate: + vocab_size = logits.size(-1) + valid_targets = target_ids.reshape(-1) + valid_targets = valid_targets[valid_targets != ignore_index] + if valid_targets.numel() and ( + (valid_targets < 0).any() or (valid_targets >= vocab_size).any() + ): + bad = valid_targets[(valid_targets < 0) | (valid_targets >= vocab_size)] + raise ValueError( + f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}" + ) + + lead_shape = logits.shape[:-1] + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() + logp, lse = _C.batch_invariant_logp_sm90(logits_2d, target_1d, int(ignore_index)) + return logp.reshape(lead_shape), lse.reshape(lead_shape) diff --git a/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py index 4ac8bd37..80e08b5f 100644 --- a/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py @@ -45,24 +45,42 @@ def apply( logits_2d = logits.reshape(-1, vocab_size).float() target_1d = target_ids.reshape(-1).to(logits.device, dtype=torch.long) - selected_logp = self._row_wise_selected_logprob( + selected_logp, _ = self._row_wise_selected_logprob_with_lse( logits_2d, target_1d, ignore_index=ignore_index, validate=validate ) return selected_logp.reshape(lead_shape) + def forward_with_lse( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return selected logprob and the FP32 vocab-domain LSE for diagnostics.""" + self._validate_shapes(logits, target_ids) + lead_shape = logits.shape[:-1] + logits_2d = logits.reshape(-1, logits.size(-1)).float() + target_1d = target_ids.reshape(-1).to(logits.device, dtype=torch.long) + logp, lse = self._row_wise_selected_logprob_with_lse( + logits_2d, target_1d, ignore_index=ignore_index, validate=validate + ) + return logp.reshape(lead_shape), lse.reshape(lead_shape) + # ---------------------------------------------------------------------- # # Core Computation # ---------------------------------------------------------------------- # @staticmethod - def _row_wise_selected_logprob( + def _row_wise_selected_logprob_with_lse( logits_2d: torch.Tensor, target_1d: torch.Tensor, *, ignore_index: int, validate: bool = True, - ) -> torch.Tensor: - """Per-row selected logprob with locked reduction order. + ) -> tuple[torch.Tensor, torch.Tensor]: + """Per-row selected logprob and LSE with locked reduction order. The three reduction steps (max, sum-exp, gather) operate on each row independently. PyTorch's ``max(dim=-1)`` and ``sum(dim=-1)`` iterate @@ -104,7 +122,7 @@ def _row_wise_selected_logprob( selected_logp = selected_logp.where(valid_mask, torch.zeros_like(selected_logp)) - return selected_logp + return selected_logp, log_sum_exp # ---------------------------------------------------------------------- # # Helper diff --git a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py new file mode 100644 index 00000000..0de98102 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py @@ -0,0 +1,527 @@ +# 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 _gather_entropy_partials( + local_entropy: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> torch.Tensor: + """Merge per-shard entropy contributions in TP-rank order. + + Entropy is not part of the WS2 selected-logprob contract, but the Vime + adapter needs it for the existing training loss surface. The collective + transports independent per-shard contributions; every TP rank performs + the same explicit rank-ordered sum afterwards. + """ + + if contract.sharding.tp_world_size == 1: + return local_entropy + + dist = _require_distributed_initialized() + gathered = [torch.empty_like(local_entropy) for _ in range(contract.sharding.tp_world_size)] + dist.all_gather(gathered, local_entropy.contiguous(), group=tp_group) + merged = gathered[0].clone() + for partial in gathered[1:]: + merged = merged + partial + return merged + + +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, + with_entropy, + with_entropy_grad, + ): + 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)) + + if with_entropy: + finite_row = torch.isfinite(lse) + lse_safe = torch.where(finite_row, lse, torch.zeros_like(lse)) + probabilities = (z_masked - lse_safe.unsqueeze(1)).exp() + probabilities = torch.where(finite_row.unsqueeze(1), probabilities, torch.zeros_like(probabilities)) + finite_logits = torch.isfinite(z_masked) + log_gap = torch.where( + finite_logits, + lse_safe.unsqueeze(1) - z_masked, + torch.zeros_like(z_masked), + ) + local_entropy = (probabilities * log_gap).sum(dim=1) + entropy = _gather_entropy_partials(local_entropy, contract, tp_group) + else: + entropy = local_logits.new_empty((0,), dtype=torch.float32) + + ctx.save_for_backward(z_masked, lse, safe_target, active_mask, padding_cols, entropy) + ctx.local_vocab_start = sharding.local_vocab_start + ctx.local_vocab_size = sharding.local_vocab_size + ctx.input_dtype = local_logits.dtype + ctx.with_entropy_grad = bool(with_entropy and with_entropy_grad) + ctx.set_materialize_grads(False) + if with_entropy and not ctx.with_entropy_grad: + ctx.mark_non_differentiable(entropy) + return selected_logp, lse, entropy + + @staticmethod + def backward(ctx, grad_logp, grad_lse, grad_entropy): + if not ctx.needs_input_grad[0] or ( + grad_logp is None and grad_lse is None and grad_entropy is None + ): + return None, None, None, None, None, None, None, None + + z_masked, lse, safe_target, active_mask, padding_cols, entropy = 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 ctx.with_entropy_grad and grad_entropy is not None: + entropy_input = lse_safe.unsqueeze(1) - z_masked - entropy.unsqueeze(1) + entropy_input = torch.where(torch.isfinite(z_masked), entropy_input, torch.zeros_like(entropy_input)) + grad = grad + grad_entropy.unsqueeze(1) * p * entropy_input + 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, 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]: + selected_logp, lse, _ = self._apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + with_entropy=False, + with_entropy_grad=False, + ) + return selected_logp, lse + + def apply_with_entropy( + 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, + with_entropy_grad: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return selected logprob, vocabulary LSE, and full-vocabulary entropy. + + The method is intentionally separate from :meth:`apply` so the WS2 + selected-logprob surface remains unchanged for existing callers. + """ + + return self._apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + with_entropy=True, + with_entropy_grad=with_entropy_grad, + ) + + def _apply( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any, + num_vocab_tiles: int, + validate: bool, + with_entropy: bool, + with_entropy_grad: bool, + ) -> tuple[torch.Tensor, 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, entropy = _VocabParallelLogprobFunction.apply( + local_logits, + target_1d, + active_mask, + contract, + tp_group, + tile, + with_entropy, + with_entropy_grad, + ) + + 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, entropy + + +__all__ = [ + "BACKEND_ID", + "DEFAULT_NUM_VOCAB_TILES", + "VocabParallelLogprobOp", +] diff --git a/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py index 66b99757..804341f3 100644 --- a/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py @@ -10,6 +10,27 @@ _BLOCK_V: int = 1024 +def _launch_batch_invariant_logp( + logits_2d: torch.Tensor, target_1d: torch.Tensor, ignore_index: int +) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = logits_2d.shape[0] + vocab_size = logits_2d.shape[1] + output = torch.empty(num_tokens, device=logits_2d.device, dtype=torch.float32) + lse = torch.empty(num_tokens, device=logits_2d.device, dtype=torch.float32) + _batch_invariant_logp_kernel[(num_tokens,)]( + logits_2d, + target_1d, + output, + lse, + num_tokens, + vocab_size, + logits_2d.stride(0), + ignore_index=ignore_index, + BLOCK_V=_BLOCK_V, + ) + return output, lse + + @triton.jit def _batch_invariant_logp_kernel( logits_ptr, # logits [N, V] @@ -126,22 +147,7 @@ def forward(ctx, logits, target_ids, ignore_index): logits_2d = logits.reshape(-1, vocab_size).contiguous() target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() - num_tokens = logits_2d.shape[0] - output = torch.empty(num_tokens, device=logits.device, dtype=torch.float32) - lse = torch.empty(num_tokens, device=logits.device, dtype=torch.float32) - - grid = (num_tokens,) - _batch_invariant_logp_kernel[grid]( - logits_2d, - target_1d, - output, - lse, - num_tokens, - vocab_size, - logits_2d.stride(0), - ignore_index=ignore_index, - BLOCK_V=_BLOCK_V, - ) + output, lse = _launch_batch_invariant_logp(logits_2d, target_1d, ignore_index) ctx.save_for_backward(logits_2d, target_1d, lse) ctx.ignore_index = ignore_index @@ -237,3 +243,53 @@ def apply( ) return _BatchInvariantLogpFunction.apply(logits, target_ids, ignore_index) + + def forward_with_lse( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return direct FP32 logprob/LSE outputs without an autograd wrapper.""" + self._validate_inputs(logits, target_ids, ignore_index=ignore_index, validate=validate) + lead_shape = logits.shape[:-1] + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() + logp, lse = _launch_batch_invariant_logp(logits_2d, target_1d, ignore_index) + return logp.reshape(lead_shape), lse.reshape(lead_shape) + + @staticmethod + def _validate_inputs( + logits: torch.Tensor, + target_ids: torch.Tensor, + *, + ignore_index: int, + validate: bool, + ) -> None: + if logits.device.type not in ("cuda", "xpu", "hip"): + raise RuntimeError( + "TritonBatchInvariantLogpOp requires a GPU tensor " + f"(CUDA / ROCm / XPU), got device '{logits.device}'." + ) + if logits.dim() < 2: + raise ValueError( + f"logits must be at least 2-D ([*lead, V]), got shape {tuple(logits.shape)}" + ) + if logits.shape[:-1] != target_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + if validate: + vocab_size = logits.size(-1) + valid_targets = target_ids.reshape(-1) + valid_targets = valid_targets[valid_targets != ignore_index] + if valid_targets.numel() and ( + (valid_targets < 0).any() or (valid_targets >= vocab_size).any() + ): + bad = valid_targets[(valid_targets < 0) | (valid_targets >= vocab_size)] + raise ValueError( + f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}" + ) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index efde5c25..bd7ba45c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -8,6 +8,17 @@ import torch +from rl_engine.kernels.logprob_contract import ( + IMPLEMENTATION_KINDS, + DeterminismScope, + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDispatchResult, + LogprobDType, + LogprobRole, + MaskMode, +) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -74,6 +85,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_BATCH_INVARIANT_LOGP_SM90 = ( "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op" ) + # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) + PYTORCH_VOCAB_PARALLEL_LOGP = ( + "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" + ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -168,6 +183,47 @@ def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() + # Truthful descriptors for the existing WS1 batch-invariant logp + # implementations: single-shard (TP=1), ignore-index masking only, no + # vocab-shard metadata, no vocab-domain LSE export. + common_logprob_roles = frozenset({LogprobRole.TRAIN, LogprobRole.INFER}) + common_logprob_dtypes = frozenset({LogprobDType.BF16, LogprobDType.FP16, LogprobDType.FP32}) + base_logprob_capabilities = { + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="pytorch-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="reference", + ), + OpBackend.TRITON_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="triton-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", + ), + OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90: LogprobBackendCapability( + backend_id="cuda-batch-invariant-logp-sm90-ws1", + roles=common_logprob_roles, + dtypes=frozenset({LogprobDType.BF16, LogprobDType.FP32}), + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", + ), + } + self._priority_map = { "cuda": { "logp": [ @@ -292,6 +348,47 @@ def __init__(self): self._adjust_priority_for_hardware() self._adjust_priority_from_env() + # WS2 dispatch owns its candidate list, seeded from the legacy + # batch_invariant_logp priority but decoupled afterwards: neither + # path's registrations may affect the other. + self._logprob_candidates: Dict[str, list] = { + platform: list(ops.get("batch_invariant_logp", [])) + for platform, ops in self._priority_map.items() + } + # Capabilities are scoped per platform: the same backend enum may + # truthfully declare different support on cuda vs rocm vs cpu. + self._logprob_capabilities: Dict[str, Dict[OpBackend, LogprobBackendCapability]] = { + platform: { + backend: base_logprob_capabilities[backend] + for backend in candidates + if backend in base_logprob_capabilities + } + for platform, candidates in self._logprob_candidates.items() + } + + # deterministic vocab-parallel TP logprob reference. + ws2_tp_logprob_capability = LogprobBackendCapability( + backend_id="pytorch-vocab-parallel-logp-ws2", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=None, + cp_world_sizes=None, + supports_vocab_padding=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), + exports_vocab_lse=True, + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", + ) + for ws2_platform in self._priority_map: + self.register_logprob_backend( + OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP, + ws2_tp_logprob_capability, + platform=ws2_platform, + prepend=True, + ) + def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() if rocm_attn_backend in {"flash_attn", "flash-attn", "flash_attention"}: @@ -381,25 +478,188 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: - if backend.name in self._instance_cache: - return self._instance_cache[backend.name] + op_instance = self._get_or_create_backend(backend) + if op_instance is not None: + return op_instance - if backend.name in self._failed_backends: - continue + raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + + def register_logprob_backend( + self, + backend: OpBackend, + capability: LogprobBackendCapability, + *, + platform: Optional[str] = None, + prepend: bool = False, + ) -> None: + """Register (or replace) a backend for WS2 contract-aware logprob dispatch. + + This is the supported seam for making a new backend selectable by + ``get_logprob_op`` (e.g. the deterministic vocab-parallel TP reference + from issue #241 PR 3) without touching the legacy ``get_op`` priority + lists. Registering the same backend again replaces its capability + without duplicating the candidate entry. + """ - op_class = self._load_backend(backend) - if op_class: - try: - op_instance = op_class() - self._instance_cache[backend.name] = op_instance - return op_instance - except Exception as e: - logger.error(f"Failed to instantiate {backend.name}: {e}") - self._failed_backends.add(backend.name) + if not isinstance(backend, OpBackend): + raise LogprobContractError("backend must be an OpBackend") + if not isinstance(capability, LogprobBackendCapability): + raise LogprobContractError("capability must be a LogprobBackendCapability") + resolved_platform = platform if platform is not None else self._platform() + if resolved_platform not in self._priority_map: + raise LogprobContractError( + f"unsupported platform {resolved_platform!r}; expected one of " + f"{sorted(self._priority_map)}" + ) + candidates = self._logprob_candidates.setdefault(resolved_platform, []) + self._logprob_capabilities.setdefault(resolved_platform, {})[backend] = capability + if backend not in candidates: + if prepend: + candidates.insert(0, backend) else: - self._failed_backends.add(backend.name) + candidates.append(backend) + + def get_logprob_op( + self, + contract: LogprobContract, + *, + requested_backend: str = "auto", + ) -> LogprobDispatchResult: + """Resolve only a backend that explicitly supports the WS2 logprob contract. + + This entry point is intentionally separate from legacy ``get_op`` so + existing callers retain their current behavior while WS2 callers cannot + silently fall back to a backend with different distributed semantics. + + ``requested_backend`` is either a case-insensitive policy keyword + (``auto`` | ``production`` | ``reference`` | ``deterministic``) or an + exact, case-sensitive stable backend id. Strictness comes from the + contract's capability checks, not from this policy string, so the + default is ``auto``. With ``tp_world_size > 1``, ``auto`` is rejected: + per-rank auto resolution can diverge across ranks, so distributed + callers must name a policy or backend id and preflight agreement via + ``LogprobContract.cross_rank_fingerprint``. + """ - raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise LogprobContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip() + if requested_backend.lower() == "deterministic": + raise LogprobContractError( + 'requested_backend="deterministic" is not a dispatch policy; request ' + "determinism through ReductionSpec.determinism_scope and match it against " + "backend determinism_scopes instead" + ) + if requested_backend.lower() == "auto" and contract.sharding.tp_world_size > 1: + raise LogprobContractError( + "Unsafe dispatch: requested_backend='auto' is not permitted when tp_world_size > 1 " + "without explicit cross-rank preflighting." + ) + + platform = self._platform() + candidates = self._logprob_candidates.get(platform, []) + rejected: list[str] = [] + # provenance["fallback"] reports only capability/load rejections of + # otherwise-eligible candidates; skips caused purely by the caller's + # own requested_backend policy filter are not fallbacks. + capability_rejections = 0 + + platform_capabilities = self._logprob_capabilities.get(platform, {}) + for backend in candidates: + capability = platform_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no LogprobBackendCapability declared") + capability_rejections += 1 + continue + policy_mismatch = self._logprob_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + # Excluded by the caller's own policy: never a fallback, even + # if the candidate would also have failed capability checks. + rejected.append(f"{backend.name}: {policy_mismatch}") + continue + capability_incompat = list(capability.incompatibilities(contract)) + if capability_incompat: + rejected.append(f"{backend.name}: " + "; ".join(capability_incompat)) + capability_rejections += 1 + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + capability_rejections += 1 + continue + + provenance = { + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": capability_rejections > 0, + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + } + return LogprobDispatchResult( + op=op, + capability=capability, + provenance=provenance, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No logprob backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, dtype={requested['dtype']}, " + f"TP={contract.sharding.tp_world_size}, CP={contract.sharding.cp_world_size}, " + f"padded_vocab={contract.sharding.padded_vocab_size}, " + f"real_vocab={contract.sharding.real_vocab_size}. Rejections: {details}" + ) + + @staticmethod + def _logprob_policy_mismatch( + requested_backend: str, + capability: LogprobBackendCapability, + ) -> str | None: + policy = requested_backend.lower() + if policy == "auto": + return None + if policy in IMPLEMENTATION_KINDS: + if capability.implementation_kind == policy: + return None + return ( + f"implementation_kind={capability.implementation_kind} does not satisfy " + f"requested_backend={policy}" + ) + if capability.backend_id == requested_backend: + return None + return ( + f"backend_id={capability.backend_id} does not match " + f"requested_backend={requested_backend}" + ) + + def _platform(self) -> str: + return self._platform_for_device(None) + + def _get_or_create_backend(self, backend: OpBackend) -> Any | None: + if backend.name in self._instance_cache: + return self._instance_cache[backend.name] + if backend.name in self._failed_backends: + return None + + op_class = self._load_backend(backend) + if op_class is None: + self._failed_backends.add(backend.name) + return None + try: + op = op_class() + except Exception as exc: + logger.error(f"Failed to instantiate {backend.name}: {exc}") + self._failed_backends.add(backend.name) + return None + self._instance_cache[backend.name] = op + return op def _platform_for_device(self, device: torch.device | str | None) -> str: if device is None: diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 42be8c1b..97b585fb 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -3,6 +3,16 @@ """Testing helpers for RL-shaped kernel validation.""" +from .logprob_comparison import ( + LogprobBackendUnavailable, + LogprobCandidate, + LogprobComparisonInputs, + LogprobComparisonReport, + compare_single_gpu_logprob, + make_logprob_candidate, + route_rl_kernel_logs_to_stderr, +) +from .logprob_drift import LogprobDriftStats, summarize_logprob_drift from .reference_ops import ( active_token_count, compute_policy_ratio, @@ -15,13 +25,22 @@ from .rl_batch import SyntheticRLKernelBatch, make_synthetic_rl_kernel_batch __all__ = [ + "LogprobBackendUnavailable", + "LogprobCandidate", + "LogprobComparisonInputs", + "LogprobComparisonReport", + "LogprobDriftStats", "SyntheticRLKernelBatch", "active_token_count", + "compare_single_gpu_logprob", "compute_policy_ratio", "compute_reference_kl", + "make_logprob_candidate", "make_synthetic_rl_kernel_batch", "masked_mean", "masked_sum", "selected_logprobs_reference", + "route_rl_kernel_logs_to_stderr", + "summarize_logprob_drift", "summarize_kernel_drift", ] diff --git a/rl_engine/testing/distributed_logprob_comparison.py b/rl_engine/testing/distributed_logprob_comparison.py new file mode 100644 index 00000000..3e37c88b --- /dev/null +++ b/rl_engine/testing/distributed_logprob_comparison.py @@ -0,0 +1,843 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Distributed WS2 comparison for the vocab-parallel logprob reference.""" + +from __future__ import annotations + +import argparse +import contextlib +import datetime +import hashlib +import json +import os +import pathlib +import shlex +import sys +from dataclasses import asdict, dataclass +from typing import Any, Sequence + +import torch + +if __package__ in (None, ""): + repo_root = pathlib.Path(__file__).resolve().parents[2] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + +_import_output = ( + contextlib.redirect_stdout(sys.stderr) + if __package__ in (None, "") + else contextlib.nullcontext() +) +with _import_output: + from rl_engine.kernels.gtest.tolerance import load_contract as load_tolerance_contract + from rl_engine.kernels.logprob_contract import ( + LogprobContract, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, + ) + from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + BACKEND_ID, + DEFAULT_NUM_VOCAB_TILES, + ) + from rl_engine.kernels.registry import KernelRegistry + from rl_engine.testing.logprob_comparison import route_rl_kernel_logs_to_stderr + from rl_engine.testing.logprob_drift import LogprobDriftStats, summarize_logprob_drift + +_DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} +_TOLERANCE_DTYPES = { + "bf16": "bfloat16", + "fp16": "float16", + "fp32": "float32", +} +_PROCESS_GROUP_TIMEOUT = datetime.timedelta(minutes=5) +_RELATIVE_ERROR_FLOOR = 1.0e-12 + + +@dataclass(frozen=True) +class DistributedLogprobCase: + tp_world_size: int + cp_world_size: int + dtype: str = "bf16" + requested_backend: str = BACKEND_ID + real_vocab_size: int = 151936 + padded_vocab_size: int = 151936 + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES + batch_size: int = 2 + sequence_length: int = 16 + prompt_tokens: int = 8 + seed: int = 123 + ignore_index: int = -100 + + def __post_init__(self) -> None: + for name in ( + "tp_world_size", + "cp_world_size", + "real_vocab_size", + "padded_vocab_size", + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + if self.dtype not in _DTYPES: + raise ValueError(f"dtype must be one of {sorted(_DTYPES)}") + if not self.requested_backend or self.requested_backend.lower() == "auto": + raise ValueError("distributed cases require an explicit non-auto backend") + if self.padded_vocab_size < self.real_vocab_size: + raise ValueError("padded_vocab_size must be at least real_vocab_size") + if self.num_vocab_tiles < self.tp_world_size: + raise ValueError("num_vocab_tiles must be at least tp_world_size") + if self.padded_vocab_size % self.num_vocab_tiles != 0: + raise ValueError("num_vocab_tiles must divide padded_vocab_size exactly") + if self.batch_size <= 0 or self.sequence_length <= 0: + raise ValueError("batch_size and sequence_length must be positive") + if not 0 <= self.prompt_tokens <= self.sequence_length: + raise ValueError("prompt_tokens must be in [0, sequence_length]") + + @property + def world_size(self) -> int: + return self.tp_world_size * self.cp_world_size + + @property + def num_tokens(self) -> int: + return self.batch_size * self.sequence_length + + @property + def case_id(self) -> str: + encoded = json.dumps(asdict(self), sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest()[:16] + + +@dataclass(frozen=True) +class RankTopology: + global_rank: int + world_size: int + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + tp_group_ranks: tuple[int, ...] + + +@dataclass(frozen=True) +class DriftDetail: + stats: LogprobDriftStats + max_rel: float + worst_global_token: int | None + worst_target_id: int | None + worst_owner_rank: int | None + candidate_value: float | None + reference_value: float | None + atol: float + rtol: float + passed: bool + + +@dataclass(frozen=True) +class RankLogprobReport: + global_rank: int + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + sp_world_size: int + dp_world_size: int + token_start: int + token_end: int + vocab_start: int + vocab_end: int + device: str + requested_backend: str + actual_backend: str + fallback: bool + contract_fingerprint: str + contract: dict[str, Any] + capability: dict[str, Any] + tp_outputs_bitwise_replicated: bool + lse: DriftDetail + dlogp: DriftDetail + passed: bool + + +@dataclass(frozen=True) +class DistributedLogprobReport: + schema_version: int + case_id: str + case: dict[str, Any] + launch_command: str + environment: dict[str, Any] + ranks: tuple[RankLogprobReport, ...] + aggregate: dict[str, DriftDetail] + bitwise_fingerprints: dict[str, Any] + passed: bool + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class _RankPayload: + report: RankLogprobReport + candidate_logp: torch.Tensor + candidate_lse: torch.Tensor + reference_logp: torch.Tensor + reference_lse: torch.Tensor + active_mask: torch.Tensor + target_ids: torch.Tensor + global_positions: torch.Tensor + + +def _strict_report_json(payload: dict[str, Any]) -> str: + return json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + + +def plan_distributed_logprob_cases( + *, + tp_world_sizes: Sequence[int] = (1, 2, 4), + cp_world_sizes: Sequence[int] = (1, 2), + **overrides: Any, +) -> tuple[DistributedLogprobCase, ...]: + """Build the scoped issue #241 topology product in deterministic order.""" + + return tuple( + DistributedLogprobCase(tp_world_size=tp, cp_world_size=cp, **overrides) + for tp in tp_world_sizes + for cp in cp_world_sizes + ) + + +def rank_topology(case: DistributedLogprobCase, global_rank: int) -> RankTopology: + if not 0 <= global_rank < case.world_size: + raise ValueError(f"global_rank must be in [0, {case.world_size})") + cp_rank, tp_rank = divmod(global_rank, case.tp_world_size) + group_start = cp_rank * case.tp_world_size + return RankTopology( + global_rank=global_rank, + world_size=case.world_size, + tp_rank=tp_rank, + tp_world_size=case.tp_world_size, + cp_rank=cp_rank, + cp_world_size=case.cp_world_size, + tp_group_ranks=tuple(range(group_start, group_start + case.tp_world_size)), + ) + + +def token_shard_bounds(num_tokens: int, cp_world_size: int) -> tuple[tuple[int, int], ...]: + """Partition token rows contiguously, allowing a one-row imbalance.""" + + if num_tokens < cp_world_size: + raise ValueError("num_tokens must be at least cp_world_size") + quotient, remainder = divmod(num_tokens, cp_world_size) + bounds = [] + cursor = 0 + for cp_rank in range(cp_world_size): + count = quotient + int(cp_rank < remainder) + bounds.append((cursor, cursor + count)) + cursor += count + return tuple(bounds) + + +def vocab_shard_bounds(case: DistributedLogprobCase) -> tuple[tuple[int, int], ...]: + """Assign complete global vocab tiles to TP ranks.""" + + tile_size = case.padded_vocab_size // case.num_vocab_tiles + quotient, remainder = divmod(case.num_vocab_tiles, case.tp_world_size) + bounds = [] + cursor_tiles = 0 + for tp_rank in range(case.tp_world_size): + tile_count = quotient + int(tp_rank < remainder) + start = cursor_tiles * tile_size + cursor_tiles += tile_count + bounds.append((start, cursor_tiles * tile_size)) + return tuple(bounds) + + +def format_launch_command( + case: DistributedLogprobCase, + *, + output: str | pathlib.Path, + device: str = "cuda", + dist_backend: str | None = None, +) -> str: + backend = dist_backend or ("nccl" if device == "cuda" else "gloo") + arguments = [ + "torchrun", + "--standalone", + f"--nproc-per-node={case.world_size}", + "rl_engine/testing/distributed_logprob_comparison.py", + "--tp", + str(case.tp_world_size), + "--cp", + str(case.cp_world_size), + "--dtype", + case.dtype, + "--backend", + case.requested_backend, + "--real-vocab", + str(case.real_vocab_size), + "--padded-vocab", + str(case.padded_vocab_size), + "--num-vocab-tiles", + str(case.num_vocab_tiles), + "--batch", + str(case.batch_size), + "--seq", + str(case.sequence_length), + "--prompt-tokens", + str(case.prompt_tokens), + "--seed", + str(case.seed), + "--device", + device, + "--dist-backend", + backend, + "--output", + str(output), + ] + return shlex.join(arguments) + + +def _canonical_inputs( + case: DistributedLogprobCase, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(case.seed) + logits = torch.randn( + case.num_tokens, + case.padded_vocab_size, + generator=generator, + dtype=torch.float32, + ) + targets = torch.randint( + 0, + case.real_vocab_size, + (case.num_tokens,), + generator=generator, + dtype=torch.long, + ) + active = torch.ones((case.batch_size, case.sequence_length), dtype=torch.bool) + active[:, : case.prompt_tokens] = False + active = active.reshape(-1) + targets = targets.masked_fill(~active, case.ignore_index) + return logits, targets, active + + +def _make_contract( + case: DistributedLogprobCase, + topology: RankTopology, + active_mask: torch.Tensor, +) -> LogprobContract: + return LogprobContract( + role=LogprobRole.TRAIN, + dtype=LogprobDType(case.dtype), + mask=MaskSpec( + num_tokens=int(active_mask.numel()), + active_mask=tuple(bool(value) for value in active_mask.tolist()), + ignore_index=case.ignore_index, + ), + sharding=ShardingSpec( + tp_rank=topology.tp_rank, + tp_world_size=case.tp_world_size, + vocab_shard_bounds=vocab_shard_bounds(case), + real_vocab_size=case.real_vocab_size, + padded_vocab_size=case.padded_vocab_size, + cp_rank=topology.cp_rank, + cp_world_size=case.cp_world_size, + ), + reduction=ReductionSpec(), + ) + + +def _fp32_oracle( + logits: torch.Tensor, + target_ids: torch.Tensor, + active_mask: torch.Tensor, + real_vocab_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + real_logits = logits[:, :real_vocab_size].float() + lse = torch.logsumexp(real_logits, dim=-1) + safe_targets = target_ids.masked_fill(~active_mask, 0) + selected = real_logits.gather(1, safe_targets.unsqueeze(1)).squeeze(1) + logp = torch.where(active_mask, selected - lse, torch.zeros_like(lse)) + return logp, lse + + +def _resolve_tolerance(dtype: str) -> tuple[float, float]: + entry = load_tolerance_contract()["accuracy"]["default"]["logprob"] + tolerance = entry[_TOLERANCE_DTYPES[dtype]] + return float(tolerance["atol"]), float(tolerance["rtol"]) + + +def _drift_detail( + candidate: torch.Tensor, + reference: torch.Tensor, + *, + target_ids: torch.Tensor, + global_positions: torch.Tensor, + sharding: ShardingSpec, + atol: float, + rtol: float, + mask: torch.Tensor | None = None, +) -> DriftDetail: + stats = summarize_logprob_drift(candidate, reference, mask=mask) + diff = (candidate.float() - reference.float()).abs() + selected = torch.ones_like(diff, dtype=torch.bool) if mask is None else mask.to(diff.device) + if not bool(selected.any().item()): + return DriftDetail(stats, 0.0, None, None, None, None, None, atol, rtol, True) + + selected_diff = diff[selected] + selected_ref = reference.float()[selected] + relative = selected_diff.double() / selected_ref.double().abs().clamp_min(_RELATIVE_ERROR_FLOOR) + selected_indices = torch.arange(diff.numel(), device=diff.device)[selected] + worst_selected = int(selected_diff.argmax().item()) + worst_local = int(selected_indices[worst_selected].item()) + target_id = int(target_ids[worst_local].item()) + close = selected_diff <= atol + rtol * selected_ref.abs() + return DriftDetail( + stats=stats, + max_rel=float(relative.max().item()), + worst_global_token=int(global_positions[worst_local].item()), + worst_target_id=target_id, + worst_owner_rank=sharding.owner_rank(target_id) if target_id >= 0 else None, + candidate_value=float(candidate[worst_local].float().item()), + reference_value=float(reference[worst_local].float().item()), + atol=atol, + rtol=rtol, + passed=bool(close.all().item()), + ) + + +def _tp_outputs_replicated( + logp: torch.Tensor, + lse: torch.Tensor, + *, + tp_group: Any, + tp_world_size: int, +) -> bool: + if tp_world_size == 1: + return True + import torch.distributed as dist + + gathered_logp = [torch.empty_like(logp) for _ in range(tp_world_size)] + gathered_lse = [torch.empty_like(lse) for _ in range(tp_world_size)] + dist.all_gather(gathered_logp, logp.contiguous(), group=tp_group) + dist.all_gather(gathered_lse, lse.contiguous(), group=tp_group) + return all(torch.equal(logp, value) for value in gathered_logp) and all( + torch.equal(lse, value) for value in gathered_lse + ) + + +def _execute_rank( + case: DistributedLogprobCase, + topology: RankTopology, + *, + device: torch.device, + tp_group: Any, +) -> _RankPayload: + full_logits, full_targets, full_active = _canonical_inputs(case) + token_start, token_end = token_shard_bounds(case.num_tokens, case.cp_world_size)[ + topology.cp_rank + ] + vocab_start, vocab_end = vocab_shard_bounds(case)[topology.tp_rank] + token_slice = slice(token_start, token_end) + local_active = full_active[token_slice].to(device=device) + local_targets = full_targets[token_slice].to(device=device) + local_fp32 = full_logits[token_slice].to(device=device) + local_logits = local_fp32[:, vocab_start:vocab_end].to(_DTYPES[case.dtype]).contiguous() + positions = torch.arange(token_start, token_end, device=device, dtype=torch.long) + + contract = _make_contract(case, topology, local_active.cpu()) + dispatch = KernelRegistry().get_logprob_op( + contract, + requested_backend=case.requested_backend, + ) + fallback = bool(dispatch.provenance["fallback"]) + if fallback: + raise RuntimeError("distributed logprob dispatch materialized through a fallback") + requested_policy = case.requested_backend.lower() + if requested_policy not in {"reference", "production"} and ( + case.requested_backend != dispatch.capability.backend_id + ): + raise RuntimeError( + f"requested backend {case.requested_backend!r} materialized as " + f"{dispatch.capability.backend_id!r}" + ) + + candidate_logp, candidate_lse = dispatch.op( + local_logits, + local_targets, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=case.num_vocab_tiles, + validate=True, + ) + reference_logp, reference_lse = _fp32_oracle( + local_fp32, + local_targets, + local_active, + case.real_vocab_size, + ) + replicated = _tp_outputs_replicated( + candidate_logp, + candidate_lse, + tp_group=tp_group, + tp_world_size=case.tp_world_size, + ) + atol, rtol = _resolve_tolerance(case.dtype) + lse_drift = _drift_detail( + candidate_lse, + reference_lse, + target_ids=local_targets, + global_positions=positions, + sharding=contract.sharding, + atol=atol, + rtol=rtol, + ) + dlogp_drift = _drift_detail( + candidate_logp, + reference_logp, + target_ids=local_targets, + global_positions=positions, + sharding=contract.sharding, + atol=atol, + rtol=rtol, + mask=local_active, + ) + rank_report = RankLogprobReport( + global_rank=topology.global_rank, + tp_rank=topology.tp_rank, + tp_world_size=topology.tp_world_size, + cp_rank=topology.cp_rank, + cp_world_size=topology.cp_world_size, + sp_world_size=1, + dp_world_size=1, + token_start=token_start, + token_end=token_end, + vocab_start=vocab_start, + vocab_end=vocab_end, + device=str(device), + requested_backend=case.requested_backend, + actual_backend=dispatch.capability.backend_id, + fallback=fallback, + contract_fingerprint=contract.cross_rank_fingerprint(), + contract=contract.to_dict(), + capability=dispatch.capability.to_dict(), + tp_outputs_bitwise_replicated=replicated, + lse=lse_drift, + dlogp=dlogp_drift, + passed=replicated and lse_drift.passed and dlogp_drift.passed, + ) + return _RankPayload( + report=rank_report, + candidate_logp=candidate_logp.detach().cpu(), + candidate_lse=candidate_lse.detach().cpu(), + reference_logp=reference_logp.detach().cpu(), + reference_lse=reference_lse.detach().cpu(), + active_mask=local_active.cpu(), + target_ids=local_targets.cpu(), + global_positions=positions.cpu(), + ) + + +def _aggregate_payloads( + case: DistributedLogprobCase, + payloads: Sequence[_RankPayload], +) -> dict[str, DriftDetail]: + representatives = sorted( + (payload for payload in payloads if payload.report.tp_rank == 0), + key=lambda payload: payload.report.cp_rank, + ) + if len(representatives) != case.cp_world_size: + raise RuntimeError("missing one or more CP representatives in rank reports") + candidate_logp = torch.cat([payload.candidate_logp for payload in representatives]) + candidate_lse = torch.cat([payload.candidate_lse for payload in representatives]) + reference_logp = torch.cat([payload.reference_logp for payload in representatives]) + reference_lse = torch.cat([payload.reference_lse for payload in representatives]) + active_mask = torch.cat([payload.active_mask for payload in representatives]) + target_ids = torch.cat([payload.target_ids for payload in representatives]) + positions = torch.cat([payload.global_positions for payload in representatives]) + sharding = _make_contract( + case, + rank_topology(case, 0), + active_mask, + ).sharding + atol, rtol = _resolve_tolerance(case.dtype) + return { + "lse": _drift_detail( + candidate_lse, + reference_lse, + target_ids=target_ids, + global_positions=positions, + sharding=sharding, + atol=atol, + rtol=rtol, + ), + "dlogp": _drift_detail( + candidate_logp, + reference_logp, + target_ids=target_ids, + global_positions=positions, + sharding=sharding, + atol=atol, + rtol=rtol, + mask=active_mask, + ), + } + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + """Hash the exact CPU tensor bytes for cross-topology comparisons.""" + + data = tensor.detach().cpu().contiguous().numpy().tobytes() + return hashlib.sha256(data).hexdigest() + + +def _aggregate_bitwise_fingerprints( + case: DistributedLogprobCase, + payloads: Sequence[_RankPayload], +) -> dict[str, Any]: + representatives = sorted( + (payload for payload in payloads if payload.report.tp_rank == 0), + key=lambda payload: payload.report.cp_rank, + ) + if len(representatives) != case.cp_world_size: + raise RuntimeError("missing one or more CP representatives in rank reports") + candidate_logp = torch.cat([payload.candidate_logp for payload in representatives]) + candidate_lse = torch.cat([payload.candidate_lse for payload in representatives]) + return { + "candidate_logp_sha256": _tensor_sha256(candidate_logp), + "candidate_lse_sha256": _tensor_sha256(candidate_lse), + "dtype": str(candidate_logp.dtype).replace("torch.", ""), + "shape": list(candidate_logp.shape), + } + + +def _create_tp_group(case: DistributedLogprobCase, topology: RankTopology) -> Any: + if case.world_size == 1: + return None + import torch.distributed as dist + + selected = None + for cp_rank in range(case.cp_world_size): + start = cp_rank * case.tp_world_size + ranks = list(range(start, start + case.tp_world_size)) + group = dist.new_group(ranks=ranks) + if topology.global_rank in ranks: + selected = group + return selected + + +def run_distributed_logprob_case( + case: DistributedLogprobCase, + *, + device_name: str = "cuda", + dist_backend: str | None = None, + output: str | pathlib.Path, +) -> DistributedLogprobReport | None: + """Run one materialized topology; only global rank zero returns the report.""" + + import torch.distributed as dist + + backend = dist_backend or ("nccl" if device_name == "cuda" else "gloo") + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if world_size != case.world_size: + raise RuntimeError( + f"WORLD_SIZE={world_size} does not match TP*CP={case.world_size}; " + "launch exactly the topology declared by the case" + ) + owns_process_group = world_size > 1 and not dist.is_initialized() + try: + if device_name == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable") + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + elif device_name == "cpu": + device = torch.device("cpu") + else: + raise ValueError("device must be cuda or cpu") + + if owns_process_group: + dist.init_process_group(backend=backend, timeout=_PROCESS_GROUP_TIMEOUT) + if dist.is_initialized(): + if dist.get_world_size() != world_size or dist.get_rank() != rank: + raise RuntimeError("initialized process group does not match RANK/WORLD_SIZE") + + topology = rank_topology(case, rank) + tp_group = _create_tp_group(case, topology) + payload = _execute_rank(case, topology, device=device, tp_group=tp_group) + if world_size == 1: + payloads = [payload] + else: + gathered: list[Any] = [None] * world_size + dist.all_gather_object(gathered, payload) + payloads = gathered + + report = None + if rank == 0: + aggregate = _aggregate_payloads(case, payloads) + bitwise_fingerprints = _aggregate_bitwise_fingerprints(case, payloads) + rank_reports = tuple( + payload.report + for payload in sorted(payloads, key=lambda item: item.report.global_rank) + ) + actual_backends = sorted({rank_report.actual_backend for rank_report in rank_reports}) + reduction_specs = { + json.dumps(rank_report.contract["reduction"], sort_keys=True) + for rank_report in rank_reports + } + materialization_consistent = len(actual_backends) == 1 and len(reduction_specs) == 1 + launch_command = format_launch_command( + case, + output=output, + device=device_name, + dist_backend=backend, + ) + report = DistributedLogprobReport( + schema_version=1, + case_id=case.case_id, + case=asdict(case), + launch_command=launch_command, + environment={ + "python": sys.version.split()[0], + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "dist_backend": backend, + "world_size": world_size, + "sp_world_size": 1, + "dp_world_size": 1, + "materialization": { + "actual_backends": actual_backends, + "consistent": materialization_consistent, + }, + "communication": { + "logprob_merge_axis": "tp_vocab", + "cp_is_merge_axis": False, + "report_collection": ("all_gather_object" if world_size > 1 else "none"), + }, + }, + ranks=rank_reports, + aggregate=aggregate, + bitwise_fingerprints=bitwise_fingerprints, + passed=( + materialization_consistent + and all(rank_report.passed for rank_report in rank_reports) + and all(detail.passed for detail in aggregate.values()) + ), + ) + output_path = pathlib.Path(output) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + _strict_report_json(report.to_dict()) + "\n", + encoding="utf-8", + ) + if world_size > 1: + dist.barrier() + return report + finally: + if owns_process_group and dist.is_initialized(): + dist.destroy_process_group() + + +def _case_from_args(args: argparse.Namespace) -> DistributedLogprobCase: + return DistributedLogprobCase( + tp_world_size=args.tp, + cp_world_size=args.cp, + dtype=args.dtype, + requested_backend=args.backend, + real_vocab_size=args.real_vocab, + padded_vocab_size=args.padded_vocab, + num_vocab_tiles=args.num_vocab_tiles, + batch_size=args.batch, + sequence_length=args.seq, + prompt_tokens=args.prompt_tokens, + seed=args.seed, + ) + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the WS2 distributed logprob drift report.") + parser.add_argument("--plan", action="store_true", help="Print the six scoped launch commands.") + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--cp", type=int, default=1) + parser.add_argument("--dtype", choices=tuple(_DTYPES), default="bf16") + parser.add_argument("--backend", default=BACKEND_ID) + parser.add_argument("--real-vocab", type=int, default=151936) + parser.add_argument("--padded-vocab", type=int, default=151936) + parser.add_argument("--num-vocab-tiles", type=int, default=DEFAULT_NUM_VOCAB_TILES) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq", type=int, default=16) + parser.add_argument("--prompt-tokens", type=int, default=8) + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda") + parser.add_argument("--dist-backend", choices=("nccl", "gloo"), default=None) + parser.add_argument("--output", default="artifacts/ws2-logprob/report.json") + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> None: + route_rl_kernel_logs_to_stderr() + args = _parse_args(argv) + if args.plan: + cases = plan_distributed_logprob_cases( + dtype=args.dtype, + requested_backend=args.backend, + real_vocab_size=args.real_vocab, + padded_vocab_size=args.padded_vocab, + num_vocab_tiles=args.num_vocab_tiles, + batch_size=args.batch, + sequence_length=args.seq, + prompt_tokens=args.prompt_tokens, + seed=args.seed, + ) + commands = [ + format_launch_command( + case, + output=pathlib.Path(args.output).parent + / f"tp{case.tp_world_size}-cp{case.cp_world_size}.json", + device=args.device, + dist_backend=args.dist_backend, + ) + for case in cases + ] + print(json.dumps({"commands": commands}, indent=2)) + return + + case = _case_from_args(args) + report = run_distributed_logprob_case( + case, + device_name=args.device, + dist_backend=args.dist_backend, + output=args.output, + ) + if report is not None: + print(_strict_report_json(report.to_dict())) + if not report.passed: + raise SystemExit(1) + + +__all__ = [ + "DistributedLogprobCase", + "DistributedLogprobReport", + "DriftDetail", + "RankLogprobReport", + "RankTopology", + "format_launch_command", + "plan_distributed_logprob_cases", + "rank_topology", + "run_distributed_logprob_case", + "token_shard_bounds", + "vocab_shard_bounds", +] + + +if __name__ == "__main__": + main() diff --git a/rl_engine/testing/logprob_comparison.py b/rl_engine/testing/logprob_comparison.py new file mode 100644 index 00000000..1e9d4e63 --- /dev/null +++ b/rl_engine/testing/logprob_comparison.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Single-GPU selected-logprob comparison.""" + +from __future__ import annotations + +import argparse +import json +import logging +import pathlib +import sys +from collections.abc import Callable, Sequence +from dataclasses import asdict, dataclass, field +from typing import Any + +import torch + +if __package__ in (None, ""): + repo_root = pathlib.Path(__file__).resolve().parents[2] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + from logprob_drift import LogprobDriftStats, summarize_logprob_drift +else: + from .logprob_drift import LogprobDriftStats, summarize_logprob_drift + + +class LogprobBackendUnavailable(RuntimeError): + pass + + +@dataclass(frozen=True) +class LogprobComparisonInputs: + logits: torch.Tensor + target_ids: torch.Tensor + active_token_mask: torch.Tensor | None = None + ignore_index: int = -100 + + +@dataclass(frozen=True) +class LogprobCandidate: + name: str + requested_backend: str + actual_backend: str + fn: Callable[[torch.Tensor, torch.Tensor, int], tuple[torch.Tensor, torch.Tensor]] + provenance: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class _LogprobPathDrift: + candidate_name: str + lse: LogprobDriftStats + dlogp: LogprobDriftStats + bitwise_logp: bool + provenance: dict[str, Any] + + +@dataclass(frozen=True) +class LogprobComparisonReport: + reference_name: str + drifts: tuple[_LogprobPathDrift, ...] + input_provenance: dict[str, Any] + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def make_logprob_candidate(backend: str) -> LogprobCandidate: + normalized = backend.strip().lower().replace("_", "-") + op: Any + if normalized in {"pytorch", "native"}: + from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import ( + NativeBatchInvariantLogpOp, + ) + + op = NativeBatchInvariantLogpOp() + actual = "pytorch" + elif normalized == "triton": + try: + from rl_engine.kernels.ops.triton.loss.batch_invariant_logp import ( + TritonBatchInvariantLogpOp, + ) + + op = TritonBatchInvariantLogpOp() + except Exception as exc: + raise LogprobBackendUnavailable(f"triton backend is unavailable: {exc}") from exc + actual = "triton" + elif normalized in {"cuda-sm90", "sm90"}: + try: + from rl_engine.kernels.ops.cuda.loss.batch_invariant_logp import ( + BatchInvariantLogpSM90Op, + ) + + op = BatchInvariantLogpSM90Op() + except Exception as exc: + raise LogprobBackendUnavailable(f"cuda-sm90 backend is unavailable: {exc}") from exc + actual = "cuda-sm90" + else: + raise ValueError( + f"unsupported logprob comparison backend {backend!r}; " + "expected pytorch, triton, or cuda-sm90" + ) + + diagnostic = getattr(op, "forward_with_lse", None) + if not callable(diagnostic): + raise LogprobBackendUnavailable( + f"backend {normalized!r} does not expose the required direct LSE diagnostic" + ) + + def run( + logits: torch.Tensor, target_ids: torch.Tensor, ignore_index: int + ) -> tuple[torch.Tensor, torch.Tensor]: + try: + return diagnostic(logits, target_ids, ignore_index=ignore_index, validate=True) + except (RuntimeError, NotImplementedError, OSError) as exc: + raise LogprobBackendUnavailable( + f"exact backend {normalized!r} cannot execute this input: {exc}" + ) from exc + + return LogprobCandidate( + name=f"{actual}-batch-invariant-logp", + requested_backend=actual, + actual_backend=actual, + fn=run, + provenance={ + "requested_alias": normalized, + "implementation": f"{type(op).__module__}.{type(op).__qualname__}", + }, + ) + + +def compare_single_gpu_logprob( + inputs: LogprobComparisonInputs, + *, + candidates: Sequence[str | LogprobCandidate] = ("pytorch",), +) -> LogprobComparisonReport: + active_mask, effective_targets = _validate_inputs(inputs) + reference_logp, reference_lse = _run_ws1_reference( + inputs.logits, effective_targets, inputs.ignore_index + ) + + drifts = [] + for candidate in candidates: + if isinstance(candidate, str): + candidate = make_logprob_candidate(candidate) + logp, lse = _run_candidate( + candidate, + inputs.logits, + effective_targets, + inputs.ignore_index, + ) + drifts.append( + _LogprobPathDrift( + candidate_name=candidate.name, + lse=summarize_logprob_drift(lse, reference_lse), + dlogp=summarize_logprob_drift(logp, reference_logp, mask=active_mask), + bitwise_logp=torch.equal(logp, reference_logp), + provenance=_candidate_provenance(candidate), + ) + ) + + return LogprobComparisonReport( + reference_name="pytorch-batch-invariant-logp", + drifts=tuple(drifts), + input_provenance={ + "device": str(inputs.logits.device), + "input_dtype": str(inputs.logits.dtype), + "output_dtype": str(reference_logp.dtype), + "shape": list(inputs.logits.shape), + "ignore_index": inputs.ignore_index, + "active_token_count": int(active_mask.sum().item()), + "tp_world": 1, + "communication": "none", + }, + ) + + +def _run_ws1_reference( + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int, +) -> tuple[torch.Tensor, torch.Tensor]: + from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp + + op = NativeBatchInvariantLogpOp() + logp = op(logits, target_ids, ignore_index=ignore_index, validate=True) + _, lse = op.forward_with_lse(logits, target_ids, ignore_index=ignore_index, validate=True) + return logp.detach(), lse.detach() + + +def _run_candidate( + candidate: LogprobCandidate, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if candidate.requested_backend != candidate.actual_backend: + raise LogprobBackendUnavailable( + f"requested backend {candidate.requested_backend!r} materialized as " + f"{candidate.actual_backend!r}; silent fallback is forbidden" + ) + logp, lse = candidate.fn(logits, target_ids, ignore_index) + expected_shape = logits.shape[:-1] + for name, value in (("logp", logp), ("lse", lse)): + if not isinstance(value, torch.Tensor): + raise TypeError(f"candidate {candidate.name!r} {name} must be a tensor") + if value.shape != expected_shape: + raise ValueError( + f"candidate {candidate.name!r} {name} shape {tuple(value.shape)} " + f"does not match {tuple(expected_shape)}" + ) + if value.dtype != torch.float32: + raise ValueError(f"candidate {candidate.name!r} {name} must be FP32") + return logp.detach(), lse.detach() + + +def _candidate_provenance(candidate: LogprobCandidate) -> dict[str, Any]: + return { + **candidate.provenance, + "requested_backend": candidate.requested_backend, + "actual_backend": candidate.actual_backend, + "tp_world": 1, + "communication": "none", + "lse_source": "direct", + } + + +def _validate_inputs( + inputs: LogprobComparisonInputs, +) -> tuple[torch.Tensor, torch.Tensor]: + if inputs.logits.dim() < 2: + raise ValueError("logits must be at least 2-D [*lead, vocab]") + if inputs.logits.shape[:-1] != inputs.target_ids.shape: + raise ValueError("target_ids shape must match logits leading shape") + if not inputs.logits.is_floating_point(): + raise ValueError("logits must be floating point") + + if inputs.active_token_mask is None: + active = inputs.target_ids != inputs.ignore_index + else: + if inputs.active_token_mask.shape != inputs.target_ids.shape: + raise ValueError("active_token_mask shape must match target_ids") + if inputs.active_token_mask.dtype != torch.bool: + raise ValueError("active_token_mask must be bool") + active = inputs.active_token_mask.to(device=inputs.target_ids.device) + if bool(((inputs.target_ids == inputs.ignore_index) & active).any().item()): + raise ValueError("active target_ids cannot equal ignore_index") + + effective = inputs.target_ids.to(device=inputs.logits.device, dtype=torch.long).clone() + active = active.to(device=inputs.logits.device, dtype=torch.bool) + effective.masked_fill_(~active, inputs.ignore_index) + valid = effective[active] + vocab_size = inputs.logits.size(-1) + if valid.numel() and ((valid < 0).any() or (valid >= vocab_size).any()): + raise ValueError(f"active target_ids must be in [0, {vocab_size})") + return active, effective + + +def _dtype(name: str) -> torch.dtype: + return { + "fp32": torch.float32, + "bf16": torch.bfloat16, + "fp16": torch.float16, + }[name] + + +def _device(name: str) -> torch.device: + if name == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device(name) + + +def route_rl_kernel_logs_to_stderr() -> None: + from rl_engine.utils.logger import logger + + for handler in logger.handlers: + if isinstance(handler, logging.StreamHandler): + handler.setStream(sys.stderr) + + +def _route_rl_kernel_logs_to_stderr() -> None: + route_rl_kernel_logs_to_stderr() + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run the WS2 TP=1 selected-logprob/LSE comparison harness." + ) + parser.add_argument( + "--candidate", + action="append", + choices=("pytorch", "triton", "cuda-sm90"), + help="Exact backend to compare. Repeat for multiple backends; defaults to pytorch.", + ) + parser.add_argument("--device", default="auto") + parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq", type=int, default=16) + parser.add_argument("--vocab", type=int, default=257) + parser.add_argument("--prompt-tokens", type=int, default=8) + parser.add_argument("--seed", type=int, default=123) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> None: + route_rl_kernel_logs_to_stderr() + args = _parse_args(argv) + device = _device(args.device) + if args.batch < 1 or args.seq < 1 or args.vocab < 1: + raise ValueError("batch, seq, and vocab must be positive") + if not 0 <= args.prompt_tokens <= args.seq: + raise ValueError("prompt-tokens must be in [0, seq]") + + generator = torch.Generator(device=device).manual_seed(args.seed) + logits = torch.randn( + args.batch, + args.seq, + args.vocab, + generator=generator, + device=device, + dtype=_dtype(args.dtype), + ) + target_ids = torch.randint( + 0, + args.vocab, + (args.batch, args.seq), + generator=generator, + device=device, + ) + active_mask = torch.ones((args.batch, args.seq), device=device, dtype=torch.bool) + active_mask[:, : args.prompt_tokens] = False + report = compare_single_gpu_logprob( + LogprobComparisonInputs( + logits=logits, + target_ids=target_ids, + active_token_mask=active_mask, + ), + candidates=tuple(args.candidate or ("pytorch",)), + ) + print(json.dumps(report.to_dict(), indent=2, sort_keys=True)) + + +__all__ = [ + "LogprobBackendUnavailable", + "LogprobCandidate", + "LogprobComparisonInputs", + "LogprobComparisonReport", + "compare_single_gpu_logprob", + "make_logprob_candidate", + "route_rl_kernel_logs_to_stderr", +] + + +if __name__ == "__main__": + main() diff --git a/rl_engine/testing/logprob_drift.py b/rl_engine/testing/logprob_drift.py new file mode 100644 index 00000000..4996dfbc --- /dev/null +++ b/rl_engine/testing/logprob_drift.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared selected-logprob drift summaries.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class LogprobDriftStats: + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + +def summarize_logprob_drift( + candidate: torch.Tensor, + reference: torch.Tensor, + *, + mask: torch.Tensor | None = None, +) -> LogprobDriftStats: + """Summarize absolute drift, optionally over active rows only.""" + + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match reference shape " + f"{tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs() + if mask is None: + values = diff.reshape(-1) + else: + if mask.shape != diff.shape: + raise ValueError("mask shape must match candidate and reference") + if mask.dtype != torch.bool: + raise ValueError("mask must be bool") + values = diff[mask.to(device=diff.device)] + + count = int(values.numel()) + if count == 0: + return LogprobDriftStats(0.0, 0.0, 0.0, 0.0, 0) + return LogprobDriftStats( + max_abs=float(values.max().item()), + mean_abs=float(values.mean().item()), + p95_abs=float(torch.quantile(values, 0.95).item()), + p99_abs=float(torch.quantile(values, 0.99).item()), + active_count=count, + ) + + +__all__ = ["LogprobDriftStats", "summarize_logprob_drift"] diff --git a/setup.py b/setup.py index a17ecb40..6a35bcb3 100644 --- a/setup.py +++ b/setup.py @@ -3,6 +3,7 @@ import importlib.util import os +import warnings from pathlib import Path from setuptools import find_packages, setup @@ -24,16 +25,27 @@ def _load_envs_module(): def _load_torch_extension_tools(): try: import torch - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - except ImportError: - return None, None, None, None + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + return None, None, None + + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension - try: - from torch.utils.cpp_extension import ROCMExtension - except ImportError: - ROCMExtension = None - return torch, BuildExtension, CUDAExtension, ROCMExtension +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) def _cuda_define_from_env(name: str, macro: str) -> list[str]: @@ -46,9 +58,56 @@ def _cuda_define_from_env(name: str, macro: str) -> list[str]: return [f"-D{macro}={parsed}"] +_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( + "-Xfatbin", + "-compress-all", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + def get_extensions(): - torch, _, CUDAExtension, ROCMExtension = _load_torch_extension_tools() + torch, _, CUDAExtension = _load_torch_extension_tools() if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) return [] extensions = [] @@ -56,47 +115,49 @@ def get_extensions(): torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = torch.version.hip is not None + is_rocm = getattr(torch.version, "hip", None) is not None - if is_rocm and ROCMExtension is not None: - extensions.append( - ROCMExtension( - name="rl_engine._C", - sources=[ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cpp", - ], - extra_compile_args={ - "cxx": ["-O3", "-std=c++17"], - "hipcc": ["-O3", "--use_fast_math", "-Xhipcc", "-compress-all"], - }, - extra_link_args=list(torch_rpath), - ) + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." ) - elif torch.cuda.is_available(): + + if is_rocm or torch.cuda.is_available(): cuda_sources = [ "csrc/ops.cpp", "csrc/fused_logp_kernel.cu", "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/attention/prefix_shared_attention.cu", "csrc/cuda/gemm/det_gemm_kernel.cu", "csrc/cuda/rmsnorm.cu", "csrc/cuda/activation.cu", "csrc/cuda/attention/deterministic_attention.cu", ] + if not is_rocm: + # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). + # The ROCm dispatcher falls back to PyTorch SDPA for this operator. + cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): nvcc_flags.append("--use_fast_math") - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") nvcc_flags.extend( _cuda_define_from_env( "FUSED_LOGP_TWOPASS_BLOCK_SIZE", @@ -139,44 +200,53 @@ def get_extensions(): "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", ) ) - if envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): nvcc_flags.append("-lineinfo") - if os.name == "nt" and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC): + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): nvcc_flags.append("-allow-unsupported-compiler") nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] extra_link_args = list(torch_rpath) - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - "csrc/cuda/embedding_lm_head_sm90.cu", # single-card batch-invariant embedding/lm-head - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - if enable_sm90 and present_sm90: - tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant - cuda_sources.extend(present_sm90) - nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + if enable_sm90 and present_sm90: + tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant + cuda_sources.extend(present_sm90) + nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) extensions.append( CUDAExtension( @@ -190,11 +260,19 @@ def get_extensions(): extra_link_args=extra_link_args, ) ) + + if _native_extension_required() and not extensions: + raise RuntimeError( + "rl_engine._C was requested but no CUDA/ROCm build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + return extensions def get_cmdclass(): - _, BuildExtension, _, _ = _load_torch_extension_tools() + _, BuildExtension, _ = _load_torch_extension_tools() if BuildExtension is None: return {} return {"build_ext": BuildExtension} diff --git a/tests/test_distributed_logprob_comparison.py b/tests/test_distributed_logprob_comparison.py new file mode 100644 index 00000000..8bb9d4ad --- /dev/null +++ b/tests/test_distributed_logprob_comparison.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +import math +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist + +import rl_engine.testing.distributed_logprob_comparison as distributed_comparison +from rl_engine.kernels.logprob_contract import ShardingSpec +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import BACKEND_ID +from rl_engine.testing.distributed_logprob_comparison import ( + DistributedLogprobCase, + _drift_detail, + _strict_report_json, + format_launch_command, + plan_distributed_logprob_cases, + rank_topology, + run_distributed_logprob_case, + token_shard_bounds, + vocab_shard_bounds, +) +from rl_engine.testing.logprob_drift import summarize_logprob_drift + + +def _small_case(*, tp: int = 1, cp: int = 1) -> DistributedLogprobCase: + return DistributedLogprobCase( + tp_world_size=tp, + cp_world_size=cp, + real_vocab_size=13, + padded_vocab_size=16, + num_vocab_tiles=8, + batch_size=1, + sequence_length=4, + prompt_tokens=1, + seed=7, + ) + + +def test_planner_builds_the_scoped_topology_product(): + cases = plan_distributed_logprob_cases( + real_vocab_size=13, + padded_vocab_size=16, + num_vocab_tiles=8, + ) + + assert [(case.tp_world_size, case.cp_world_size) for case in cases] == [ + (1, 1), + (1, 2), + (2, 1), + (2, 2), + (4, 1), + (4, 2), + ] + assert [case.world_size for case in cases] == [1, 2, 2, 4, 4, 8] + + +def test_rank_mapping_keeps_cp_out_of_the_tp_merge_axis(): + case = _small_case(tp=2, cp=2) + + assert rank_topology(case, 0).tp_group_ranks == (0, 1) + assert rank_topology(case, 1).tp_group_ranks == (0, 1) + assert rank_topology(case, 2).tp_group_ranks == (2, 3) + assert rank_topology(case, 3).tp_group_ranks == (2, 3) + assert [ + (rank_topology(case, rank).cp_rank, rank_topology(case, rank).tp_rank) for rank in range(4) + ] == [ + (0, 0), + (0, 1), + (1, 0), + (1, 1), + ] + + +def test_token_and_vocab_bounds_cover_each_axis_once(): + case = _small_case(tp=4, cp=2) + + assert token_shard_bounds(case.num_tokens, case.cp_world_size) == ((0, 2), (2, 4)) + assert vocab_shard_bounds(case) == ((0, 4), (4, 8), (8, 12), (12, 16)) + + +def test_case_rejects_implicit_backend_and_non_tileable_vocab(): + with pytest.raises(ValueError, match="explicit non-auto backend"): + DistributedLogprobCase(tp_world_size=2, cp_world_size=1, requested_backend="auto") + with pytest.raises(ValueError, match="must divide"): + DistributedLogprobCase( + tp_world_size=2, + cp_world_size=1, + padded_vocab_size=15, + real_vocab_size=13, + num_vocab_tiles=8, + ) + + +def test_launch_command_records_the_materialized_case(tmp_path): + case = _small_case(tp=2, cp=2) + command = format_launch_command(case, output=tmp_path / "report.json") + + assert "--nproc-per-node=4" in command + assert "--tp 2 --cp 2" in command + assert f"--backend {BACKEND_ID}" in command + assert "--real-vocab 13 --padded-vocab 16" in command + + +def test_shared_pr2_drift_summary_preserves_active_mask_semantics(): + candidate = torch.tensor([100.0, 1.0, 3.0]) + reference = torch.tensor([0.0, 2.0, 1.0]) + mask = torch.tensor([False, True, True]) + + stats = summarize_logprob_drift(candidate, reference, mask=mask) + + assert stats.active_count == 2 + assert stats.max_abs == 2.0 + assert stats.mean_abs == 1.5 + + +def test_relative_drift_near_zero_stays_finite(): + sharding = ShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, 16),), + real_vocab_size=13, + padded_vocab_size=16, + ) + detail = _drift_detail( + torch.tensor([1.0]), + torch.tensor([0.0]), + target_ids=torch.tensor([1]), + global_positions=torch.tensor([3]), + sharding=sharding, + atol=0.0, + rtol=0.0, + ) + + assert math.isfinite(detail.max_rel) + assert detail.max_rel == pytest.approx(1.0e12) + assert json.loads(_strict_report_json({"max_rel": detail.max_rel}))["max_rel"] == pytest.approx( + 1.0e12 + ) + with pytest.raises(ValueError, match="Out of range float values"): + _strict_report_json({"max_rel": float("nan")}) + + +def test_tp1_cpu_case_writes_116_compatible_artifact(tmp_path, monkeypatch): + monkeypatch.delenv("RANK", raising=False) + monkeypatch.delenv("LOCAL_RANK", raising=False) + monkeypatch.delenv("WORLD_SIZE", raising=False) + output = tmp_path / "tp1-cp1.json" + + report = run_distributed_logprob_case( + _small_case(), + device_name="cpu", + dist_backend="gloo", + output=output, + ) + + assert report is not None and report.passed + assert report.aggregate["lse"].stats.active_count == 4 + assert report.aggregate["dlogp"].stats.active_count == 3 + assert report.ranks[0].actual_backend == BACKEND_ID + assert report.ranks[0].fallback is False + assert report.ranks[0].tp_outputs_bitwise_replicated + payload = json.loads(output.read_text(encoding="utf-8")) + assert payload["schema_version"] == 1 + assert payload["ranks"][0]["contract"]["reduction"]["cp_is_merge_axis"] is False + assert payload["ranks"][0]["sp_world_size"] == 1 + assert payload["ranks"][0]["dp_world_size"] == 1 + assert payload["environment"]["materialization"]["consistent"] is True + assert payload["aggregate"]["dlogp"]["worst_target_id"] is not None + fingerprints = payload["bitwise_fingerprints"] + assert len(fingerprints["candidate_logp_sha256"]) == 64 + assert len(fingerprints["candidate_lse_sha256"]) == 64 + assert fingerprints["dtype"] == "float32" + assert fingerprints["shape"] == [4] + assert payload["launch_command"].startswith("torchrun --standalone") + + +def test_world_size_mismatch_fails_before_process_group_init(tmp_path, monkeypatch): + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + + with pytest.raises(RuntimeError, match=r"does not match TP\*CP"): + run_distributed_logprob_case( + _small_case(), + device_name="cpu", + dist_backend="gloo", + output=tmp_path / "unused.json", + ) + + +def test_setup_failure_destroys_owned_process_group(tmp_path, monkeypatch): + state = {"initialized": False, "destroyed": False} + + def init_process_group(*, backend, timeout): + state["initialized"] = True + + def destroy_process_group(): + state["destroyed"] = True + state["initialized"] = False + + def fail_group_setup(case, topology): + raise RuntimeError("group setup failed") + + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(dist, "is_initialized", lambda: state["initialized"]) + monkeypatch.setattr(dist, "init_process_group", init_process_group) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + monkeypatch.setattr(dist, "get_rank", lambda: 0) + monkeypatch.setattr(dist, "destroy_process_group", destroy_process_group) + monkeypatch.setattr(distributed_comparison, "_create_tp_group", fail_group_setup) + + with pytest.raises(RuntimeError, match="group setup failed"): + run_distributed_logprob_case( + _small_case(tp=2), + device_name="cpu", + dist_backend="gloo", + output=tmp_path / "unused.json", + ) + + assert state == {"initialized": False, "destroyed": True} + + +def test_partial_initialization_failure_destroys_owned_process_group(tmp_path, monkeypatch): + state = {"initialized": False, "destroyed": False} + + def fail_initialization(*, backend, timeout): + state["initialized"] = True + raise RuntimeError("initialization failed") + + def destroy_process_group(): + state["destroyed"] = True + state["initialized"] = False + + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(dist, "is_initialized", lambda: state["initialized"]) + monkeypatch.setattr(dist, "init_process_group", fail_initialization) + monkeypatch.setattr(dist, "destroy_process_group", destroy_process_group) + + with pytest.raises(RuntimeError, match="initialization failed"): + run_distributed_logprob_case( + _small_case(tp=2), + device_name="cpu", + dist_backend="gloo", + output=tmp_path / "unused.json", + ) + + assert state == {"initialized": False, "destroyed": True} + + +@pytest.mark.skipif(not torch.distributed.is_available(), reason="torch.distributed required") +def test_tp2_cp2_gloo_cli_emits_per_rank_report(tmp_path): + script = ( + Path(__file__).resolve().parents[1] + / "rl_engine" + / "testing" + / "distributed_logprob_comparison.py" + ) + output = tmp_path / "tp2-cp2.json" + environment = os.environ.copy() + environment.setdefault("OMP_NUM_THREADS", "1") + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc-per-node=4", + str(script), + "--tp", + "2", + "--cp", + "2", + "--device", + "cpu", + "--dist-backend", + "gloo", + "--real-vocab", + "13", + "--padded-vocab", + "16", + "--num-vocab-tiles", + "8", + "--batch", + "1", + "--seq", + "4", + "--prompt-tokens", + "1", + "--output", + str(output), + ], + check=True, + capture_output=True, + text=True, + timeout=120, + env=environment, + ) + + payload = json.loads(output.read_text(encoding="utf-8")) + assert payload["passed"] + assert len(payload["ranks"]) == 4 + assert all(rank["tp_outputs_bitwise_replicated"] for rank in payload["ranks"]) + assert {rank["actual_backend"] for rank in payload["ranks"]} == {BACKEND_ID} + assert {rank["cp_rank"] for rank in payload["ranks"]} == {0, 1} + assert payload["aggregate"]["dlogp"]["stats"]["active_count"] == 3 + assert json.loads(result.stdout)["case"]["tp_world_size"] == 2 diff --git a/tests/test_logprob_comparison.py b/tests/test_logprob_comparison.py new file mode 100644 index 00000000..4fc62a13 --- /dev/null +++ b/tests/test_logprob_comparison.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import argparse +import io +import json +import logging +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from rl_engine.kernels.gtest import run_operator_suite +from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case +from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp +from rl_engine.testing.logprob_comparison import ( + LogprobBackendUnavailable, + LogprobCandidate, + LogprobComparisonInputs, + _device, + _route_rl_kernel_logs_to_stderr, + compare_single_gpu_logprob, + make_logprob_candidate, +) +from rl_engine.utils.logger import logger + + +def _inputs() -> LogprobComparisonInputs: + generator = torch.Generator().manual_seed(17) + logits = torch.randn(2, 4, 257, generator=generator, dtype=torch.float32) + target_ids = torch.tensor([[3, 5, 7, 11], [13, 17, 19, 23]]) + active = torch.tensor([[False, False, True, True], [False, True, True, True]]) + return LogprobComparisonInputs(logits, target_ids, active_token_mask=active) + + +def test_single_gpu_pytorch_path_is_bitwise_regression_guard(): + report = compare_single_gpu_logprob(_inputs(), candidates=("pytorch",)) + + assert report.reference_name == "pytorch-batch-invariant-logp" + assert len(report.drifts) == 1 + drift = report.drifts[0] + assert drift.bitwise_logp + assert drift.lse.max_abs == 0.0 + assert drift.dlogp.max_abs == 0.0 + assert drift.dlogp.active_count == 5 + assert drift.provenance["requested_backend"] == "pytorch" + assert drift.provenance["actual_backend"] == "pytorch" + assert drift.provenance["lse_source"] == "direct" + assert report.input_provenance["tp_world"] == 1 + assert report.input_provenance["communication"] == "none" + + +def test_report_serializes_lse_and_active_token_percentiles(): + inputs = _inputs() + reference = make_logprob_candidate("pytorch") + + def shifted(logits, target_ids, ignore_index): + logp, lse = reference.fn(logits, target_ids, ignore_index) + logp = logp.clone() + logp[0, 0] += 100.0 # inactive and therefore excluded from dlogp + logp[0, 2] += 1.0 + lse = lse + torch.arange(lse.numel(), dtype=lse.dtype).reshape_as(lse) * 0.1 + return logp, lse + + candidate = LogprobCandidate( + name="shifted", + requested_backend="shifted", + actual_backend="shifted", + fn=shifted, + ) + report = compare_single_gpu_logprob(inputs, candidates=(candidate,)) + payload = report.to_dict() + drift = payload["drifts"][0] + + assert drift["dlogp"]["active_count"] == 5 + assert drift["dlogp"]["max_abs"] == pytest.approx(1.0) + assert drift["dlogp"]["p95_abs"] == pytest.approx(0.8) + assert drift["dlogp"]["p99_abs"] == pytest.approx(0.96) + assert drift["lse"]["active_count"] == 8 + assert drift["lse"]["p99_abs"] == pytest.approx(0.693, abs=1e-5) + + +def test_canonical_provenance_cannot_be_overridden(): + native = make_logprob_candidate("pytorch") + candidate = LogprobCandidate( + name="custom", + requested_backend="pytorch", + actual_backend="pytorch", + fn=native.fn, + provenance={ + "actual_backend": "fallback", + "tp_world": 8, + "communication": "all-gather", + "lse_source": "reconstructed", + "implementation": "custom", + }, + ) + + provenance = compare_single_gpu_logprob(_inputs(), candidates=(candidate,)).drifts[0].provenance + + assert provenance["actual_backend"] == "pytorch" + assert provenance["tp_world"] == 1 + assert provenance["communication"] == "none" + assert provenance["lse_source"] == "direct" + assert provenance["implementation"] == "custom" + + +def test_all_inactive_tokens_produce_zero_dlogp_statistics(): + inputs = _inputs() + inputs = LogprobComparisonInputs( + inputs.logits, + inputs.target_ids, + active_token_mask=torch.zeros_like(inputs.target_ids, dtype=torch.bool), + ) + drift = compare_single_gpu_logprob(inputs).drifts[0] + + assert drift.dlogp.active_count == 0 + assert drift.dlogp.max_abs == 0.0 + assert drift.dlogp.p95_abs == 0.0 + assert drift.lse.active_count == inputs.target_ids.numel() + + +def test_explicit_backend_mismatch_fails_closed(): + native = make_logprob_candidate("pytorch") + disguised = LogprobCandidate( + name="fallback", + requested_backend="cuda-sm90", + actual_backend="pytorch", + fn=native.fn, + ) + + with pytest.raises(LogprobBackendUnavailable, match="silent fallback is forbidden"): + compare_single_gpu_logprob(_inputs(), candidates=(disguised,)) + + +def test_active_ignore_index_is_rejected(): + inputs = _inputs() + targets = inputs.target_ids.clone() + targets[0, 2] = -100 + + with pytest.raises(ValueError, match="active target_ids cannot equal ignore_index"): + compare_single_gpu_logprob( + LogprobComparisonInputs( + inputs.logits, + targets, + active_token_mask=inputs.active_token_mask, + ) + ) + + +def test_native_diagnostic_lse_satisfies_selected_logit_identity(): + inputs = _inputs() + candidate = make_logprob_candidate("pytorch") + effective = inputs.target_ids.masked_fill(~inputs.active_token_mask, -100) + logp, lse = candidate.fn(inputs.logits, effective, -100) + production_logp = NativeBatchInvariantLogpOp()( + inputs.logits, effective, ignore_index=-100, validate=True + ) + safe_targets = effective.masked_fill(~inputs.active_token_mask, 0) + selected = torch.gather(inputs.logits, -1, safe_targets.unsqueeze(-1)).squeeze(-1) + + assert torch.equal(logp, production_logp) + assert torch.equal(logp[inputs.active_token_mask], (selected - lse)[inputs.active_token_mask]) + + +def test_unsupported_backend_name_is_rejected(): + with pytest.raises(ValueError, match="unsupported logprob comparison backend"): + make_logprob_candidate("unknown") + + +def test_cli_auto_device_resolves_without_constructing_auto(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + assert _device("auto") == torch.device("cpu") + + +def test_cli_routes_rl_kernel_logs_to_stderr_for_machine_readable_stdout(monkeypatch): + stdout = io.StringIO() + stderr = io.StringIO() + original_streams = [ + (handler, handler.stream) + for handler in logger.handlers + if isinstance(handler, logging.StreamHandler) + ] + monkeypatch.setattr(sys, "stdout", stdout) + monkeypatch.setattr(sys, "stderr", stderr) + + try: + _route_rl_kernel_logs_to_stderr() + logger.info("test backend diagnostic") + print(json.dumps({"ok": True})) + finally: + for handler, stream in original_streams: + handler.setStream(stream) + + assert json.loads(stdout.getvalue()) == {"ok": True} + assert "test backend diagnostic" in stderr.getvalue() + + +def test_cli_runs_directly_from_testing_module(): + script = Path(__file__).resolve().parents[1] / "rl_engine" / "testing" / "logprob_comparison.py" + result = subprocess.run( + [ + sys.executable, + str(script), + "--candidate", + "pytorch", + "--device", + "cpu", + "--batch", + "1", + "--seq", + "2", + "--vocab", + "17", + "--prompt-tokens", + "1", + ], + check=True, + capture_output=True, + text=True, + ) + + payload = json.loads(result.stdout) + assert payload["drifts"][0]["provenance"]["actual_backend"] == "pytorch" + assert payload["input_provenance"]["communication"] == "none" + + +def test_operator_comparison_specs_register_batch_invariant_logp(): + args = argparse.Namespace( + op="batch_invariant_logp", + candidate="pytorch", + arch_key=None, + batch=2, + seq=4, + vocab=17, + seed=7, + input_mode="random", + constant_value=0.5, + token_value=3, + normalized_dim=128, + k_dim=16, + n_dim=32, + theta=1.0e6, + eps=1.0e-6, + ) + + case = make_operator_case(args, torch.float32, torch.device("cpu")) + candidate = make_candidate(args) + report = run_operator_suite("batch_invariant_logp", candidates=[candidate], cases=[case]) + + assert report.passed + assert report.candidates[0].cases[0].op_class == "logprob" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_triton_diagnostic_path_reports_direct_lse(): + try: + candidate = make_logprob_candidate("triton") + except LogprobBackendUnavailable as exc: + pytest.skip(str(exc)) + logits = torch.randn(4, 1024, device="cuda", dtype=torch.bfloat16) + targets = torch.tensor([0, 17, 511, 1023], device="cuda") + try: + report = compare_single_gpu_logprob( + LogprobComparisonInputs(logits, targets), candidates=(candidate,) + ) + except LogprobBackendUnavailable as exc: + if isinstance(exc.__cause__, PermissionError): + pytest.skip(str(exc)) + raise + + assert report.drifts[0].provenance["actual_backend"] == "triton" + assert report.drifts[0].provenance["lse_source"] == "direct" + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9, + reason="Hopper CUDA device required", +) +def test_sm90_diagnostic_path_reports_direct_lse_without_fallback(): + try: + candidate = make_logprob_candidate("cuda-sm90") + except LogprobBackendUnavailable as exc: + pytest.skip(str(exc)) + logits = torch.randn(4, 1024, device="cuda", dtype=torch.bfloat16) + targets = torch.tensor([0, 17, 511, 1023], device="cuda") + report = compare_single_gpu_logprob( + LogprobComparisonInputs(logits, targets), candidates=(candidate,) + ) + + assert report.drifts[0].provenance["actual_backend"] == "cuda-sm90" + assert report.drifts[0].provenance["lse_source"] == "direct" diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py new file mode 100644 index 00000000..a74f3b4c --- /dev/null +++ b/tests/test_logprob_contract.py @@ -0,0 +1,636 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 TP-aware logprob contract and contract-aware dispatch tests (issue #241).""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import pytest + +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDType, + LogprobOutputSpec, + LogprobRole, + MaskMode, + MaskSpec, + ReductionSpec, + ShardingSpec, + TPPlacement, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +QWEN3_REAL_VOCAB = 151936 +QWEN3_PADDED_VOCAB = 152064 + + +def _even_bounds(padded_vocab: int, tp_world_size: int) -> tuple[tuple[int, int], ...]: + shard = padded_vocab // tp_world_size + return tuple( + (rank * shard, padded_vocab if rank == tp_world_size - 1 else (rank + 1) * shard) + for rank in range(tp_world_size) + ) + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 2, + cp_rank: int = 0, + cp_world_size: int = 2, + real_vocab_size: int = QWEN3_REAL_VOCAB, + padded_vocab_size: int = QWEN3_PADDED_VOCAB, + vocab_shard_bounds: tuple[tuple[int, int], ...] | None = None, +) -> ShardingSpec: + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=( + vocab_shard_bounds + if vocab_shard_bounds is not None + else _even_bounds(padded_vocab_size, tp_world_size) + ), + real_vocab_size=real_vocab_size, + padded_vocab_size=padded_vocab_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ) + + +def _mask( + *, + num_tokens: int = 8, + active_mask: tuple[bool, ...] | None = None, + ignore_index: int = -100, +) -> MaskSpec: + return MaskSpec( + num_tokens=num_tokens, + active_mask=( + active_mask + if active_mask is not None + else (False, False, True, True, True, True, True, False) + ), + ignore_index=ignore_index, + ) + + +def _contract( + *, + role: str = "train", + dtype: str = "bf16", + mask: MaskSpec | None = None, + sharding: ShardingSpec | None = None, + reduction: ReductionSpec | None = None, +) -> LogprobContract: + return LogprobContract( + role=role, + dtype=dtype, + mask=mask if mask is not None else _mask(), + sharding=sharding if sharding is not None else _sharding(), + reduction=reduction if reduction is not None else ReductionSpec(), + ) + + +def _declared_tp_backend() -> LogprobBackendCapability: + return LogprobBackendCapability( + backend_id="test-deterministic-tp-logprob", + roles=frozenset({LogprobRole.TRAIN, LogprobRole.INFER}), + dtypes=frozenset({LogprobDType.BF16}), + tp_world_sizes=(1, 2, 4), + cp_world_sizes=None, + supports_vocab_padding=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), + exports_vocab_lse=True, + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", + ) + + +def test_qwen3_tp2_bf16_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.tp_world_size == 2 + assert contract.sharding.cp_world_size == 2 + assert contract.sharding.local_vocab_start == 0 + assert contract.sharding.local_vocab_end == QWEN3_PADDED_VOCAB // 2 + assert contract.sharding.local_vocab_size == QWEN3_PADDED_VOCAB // 2 + assert contract.mask.active_token_count == 5 + assert contract.reduction.acc_dtype is LogprobDType.FP32 + assert contract.to_dict()["reduction"] == { + "merge": "max_sumexp", + "merge_axis": "tp_vocab", + "acc_dtype": "fp32", + "order": "global_vocab_shard_index", + "transport": "all_gather", + "downcast_at": "final_write", + "engine": "in_op_reference", + "determinism_scope": "cross_tp_bitwise", + "cp_is_merge_axis": False, + } + json.dumps(contract.to_dict()) + + +@pytest.mark.parametrize("tp_world_size", [1, 2, 4]) +def test_pr4_sweep_tp_degrees_are_representable(tp_world_size): + sharding = _sharding(tp_world_size=tp_world_size, cp_world_size=1) + + assert len(sharding.vocab_shard_bounds) == tp_world_size + assert sharding.vocab_shard_bounds[-1][1] == QWEN3_PADDED_VOCAB + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == tp_world_size - 1 + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 2, "tp_rank=2"), + ("cp_rank", 2, "cp_rank=2"), + ("real_vocab_size", 0, "positive integer"), + ("padded_vocab_size", QWEN3_REAL_VOCAB - 1, "must not be smaller"), + ], +) +def test_invalid_rank_and_vocab_metadata_fail_loudly(field, value, message): + values = { + "tp_rank": 0, + "tp_world_size": 2, + "cp_rank": 0, + "cp_world_size": 2, + "real_vocab_size": QWEN3_REAL_VOCAB, + "padded_vocab_size": QWEN3_PADDED_VOCAB, + "vocab_shard_bounds": _even_bounds(QWEN3_PADDED_VOCAB, 2), + } + values[field] = value + + with pytest.raises(LogprobContractError, match=message): + ShardingSpec(**values) + + +@pytest.mark.parametrize( + ("bounds", "message"), + [ + ((), "one \\(start, end\\) pair per TP rank"), + (((0, 76032),), "one \\(start, end\\) pair per TP rank"), + (((0, 76032), (76032, 76032)), "end > start"), + (((0, 76000), (76032, 152064)), "contiguous"), + (((0, 76064), (76032, 152064)), "contiguous"), + (((0, 76032), (76032, 152000)), "cover padded_vocab_size exactly"), + ], +) +def test_incomplete_or_overlapping_vocab_shard_bounds_fail_loudly(bounds, message): + with pytest.raises(LogprobContractError, match=message): + _sharding(vocab_shard_bounds=bounds) + + +def test_owner_rank_is_unique_and_rejects_out_of_real_vocab_targets(): + sharding = _sharding() + + assert sharding.owner_rank(0) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2 - 1) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2) == 1 + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 1 + + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(-1) + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(QWEN3_REAL_VOCAB) + + +def test_active_token_mask_metadata_is_validated(): + with pytest.raises(LogprobContractError, match="one entry per token"): + _mask(num_tokens=4) + + with pytest.raises(LogprobContractError, match="must be a bool"): + MaskSpec(num_tokens=2, active_mask=(True, 1)) + + all_inactive = _mask(num_tokens=3, active_mask=(False, False, False)) + assert all_inactive.active_token_count == 0 + + +def test_reduction_requires_fp32_accumulation_and_known_semantics(): + with pytest.raises(LogprobContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + with pytest.raises(LogprobContractError, match="merge must be one of"): + ReductionSpec(merge="lse_average") + + with pytest.raises(LogprobContractError, match="transport must be one of"): + ReductionSpec(transport="all_reduce") + + +def test_contract_component_types_and_lse_export_are_enforced(): + with pytest.raises(LogprobContractError, match="mask must be a MaskSpec"): + LogprobContract( + role="train", + dtype="bf16", + mask=None, + sharding=_sharding(), + reduction=ReductionSpec(), + ) + + with pytest.raises(LogprobContractError, match="export_lse must be True"): + replace(_contract(), export_lse=False) + + +def test_ignore_index_must_not_collide_with_the_real_vocabulary(): + with pytest.raises(LogprobContractError, match="must not collide"): + _contract(mask=_mask(ignore_index=5)) + + padding_column = QWEN3_REAL_VOCAB + 1 + contract = _contract(mask=_mask(ignore_index=padding_column)) + assert contract.mask.ignore_index == padding_column + + +def _restrict_to_ws1_candidates(registry: KernelRegistry) -> None: + """Drop the #241 PR3 vocab-parallel reference so only WS1 backends remain.""" + + platform = registry._platform() + registry._logprob_candidates[platform] = [ + backend + for backend in registry._logprob_candidates[platform] + if backend is not OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP + ] + + +def test_current_ws1_backend_rejects_strict_tp_contract_without_fallback(): + registry = KernelRegistry() + _restrict_to_ws1_candidates(registry) + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(_contract(), requested_backend="reference") + + message = str(exc_info.value) + assert "TP=2 is unsupported" in message + assert "vocab-domain LSE export is unsupported" in message + assert "determinism_scope=cross_tp_bitwise is unsupported" in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_current_ws1_backend_rejects_padded_vocab_even_at_tp1(): + registry = KernelRegistry() + _restrict_to_ws1_candidates(registry) + contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(contract) + + message = str(exc_info.value) + assert "TP=1 is unsupported" not in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_ws1_rejections_recorded_when_vocab_parallel_reference_resolves(): + """The WS1 backends still reject strict contracts; they are skipped with + recorded reasons while dispatch resolves the #241 PR3 reference.""" + + registry = KernelRegistry() + platform = registry._platform() + # Order the WS1 backends ahead of the reference so their rejections are + # exercised on the way to a successful resolution. + candidates = registry._logprob_candidates[platform] + candidates.remove(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) + candidates.append(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + assert result.capability.backend_id == "pytorch-vocab-parallel-logp-ws2" + assert result.provenance["fallback"] is True + rejections = " | ".join(result.provenance["prior_rejections"]) + assert "vocab-domain LSE export is unsupported" in rejections + + +def test_undeclared_backend_capability_is_never_selected(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [OpBackend.PYTORCH_NATIVE] + + tp1_contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + with pytest.raises(RuntimeError, match="no LogprobBackendCapability declared"): + registry.get_logprob_op(tp1_contract) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert result.provenance["requested_backend"] == "reference" + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["real_vocab_size"] == QWEN3_REAL_VOCAB + assert result.provenance["contract"]["reduction"]["cp_is_merge_axis"] is False + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): + registry.get_logprob_op(_contract(), requested_backend="another-backend") + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + + +def test_cp_is_a_non_merge_axis_and_cp_agnostic_backends_accept_any_cp_degree(): + capability = _declared_tp_backend() + cp2_contract = _contract(sharding=_sharding(cp_world_size=2, cp_rank=1)) + + assert capability.incompatibilities(cp2_contract) == () + + cp_restricted = replace(capability, cp_world_sizes=(1,)) + assert cp_restricted.incompatibilities(cp2_contract) == ("CP=2 is unsupported",) + + +def test_inactive_tokens_require_explicit_active_mask_support(): + capability = replace(_declared_tp_backend(), mask_modes=frozenset({MaskMode.IGNORE_INDEX})) + contract = _contract() + + assert "explicit active-token masking is unsupported" in ( + capability.incompatibilities(contract) + ) + + fully_active = _contract(mask=_mask(num_tokens=3, active_mask=(True, True, True))) + assert capability.incompatibilities(fully_active) == () + + +def test_backend_id_must_not_shadow_a_reserved_policy_keyword(): + with pytest.raises(LogprobContractError, match="reserved dispatch policy keyword"): + replace(_declared_tp_backend(), backend_id="Deterministic") + + +def test_default_auto_policy_resolves_any_compatible_implementation_kind(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), implementation_kind="reference"), + platform=platform, + ) + + tp1_contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + result = registry.get_logprob_op(tp1_contract) + + assert result.provenance["requested_backend"] == "auto" + assert result.capability.implementation_kind == "reference" + + +def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="REFERENCE") + assert result.capability.backend_id == "test-deterministic-tp-logprob" + + with pytest.raises(RuntimeError, match="does not match requested_backend"): + registry.get_logprob_op(_contract(), requested_backend="Test-Deterministic-TP-Logprob") + + +def test_policy_only_skips_are_not_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="other-compatible-backend"), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_capability_rejections_are_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + + assert result.provenance["fallback"] is True + assert "TP=2 is unsupported" in result.provenance["prior_rejections"][0] + + +def test_ws2_candidate_list_is_decoupled_from_the_legacy_priority_map(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform].insert(0, OpBackend.PYTORCH_NATIVE) + + legacy = registry._priority_map[platform]["batch_invariant_logp"] + assert OpBackend.PYTORCH_NATIVE not in legacy + + legacy.insert(0, OpBackend.PYTORCH_GEMM) + assert OpBackend.PYTORCH_GEMM not in registry._logprob_candidates[platform] + + +def test_register_logprob_backend_is_the_public_registration_seam(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + capability = _declared_tp_backend() + + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, capability, platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(capability, backend_id="replacement-backend"), + platform=platform, + ) + + assert registry._logprob_candidates[platform] == [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + result = registry.get_logprob_op(_contract(), requested_backend="reference") + assert result.capability.backend_id == "replacement-backend" + + with pytest.raises(LogprobContractError, match="capability must be"): + registry.register_logprob_backend(OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, None) + + +def test_backend_id_whitespace_is_normalized_for_dispatch(): + capability = replace(_declared_tp_backend(), backend_id=" padded-id ") + assert capability.backend_id == "padded-id" + + +def test_capabilities_are_scoped_per_platform(): + registry = KernelRegistry() + platform = registry._platform() + other = "rocm" if platform != "rocm" else "cpu" + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="other-platform-backend"), + platform=other, + ) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert ( + registry._logprob_capabilities[other][OpBackend.PYTORCH_BATCH_INVARIANT_LOGP].backend_id + == "other-platform-backend" + ) + + +def test_register_logprob_backend_rejects_unknown_platform(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="unsupported platform"): + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + _declared_tp_backend(), + platform="cuda-typo", + ) + + +def test_non_iterable_roles_and_dtypes_raise_contract_errors(): + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), roles=None) + + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), dtypes=42) + + +def test_requested_deterministic_policy_is_a_loud_error(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="determinism_scope"): + registry.get_logprob_op(_contract(), requested_backend="deterministic") + + +def test_auto_policy_is_rejected_for_tp_sharded_contracts(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + with pytest.raises(LogprobContractError, match="Unsafe dispatch"): + registry.get_logprob_op(_contract()) + + with pytest.raises(LogprobContractError, match="Unsafe dispatch"): + registry.get_logprob_op(_contract(), requested_backend=" AUTO ") + + tp1_contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + assert registry.get_logprob_op(tp1_contract).provenance["requested_backend"] == "auto" + + +def test_determinism_scope_is_part_of_the_typed_contract(): + fixed_only = replace( + _declared_tp_backend(), + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + ) + + assert "determinism_scope=cross_tp_bitwise is unsupported" in ( + fixed_only.incompatibilities(_contract()) + ) + + relaxed = _contract(reduction=ReductionSpec(determinism_scope="fixed_topology")) + assert fixed_only.incompatibilities(relaxed) == () + + +def test_policy_filtered_candidates_never_count_toward_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_output_spec_is_pinned_to_fp32_replicated(): + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(selected_logp_dtype="bf16") + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(lse_dtype="bf16") + + assert LogprobOutputSpec().tp_placement is TPPlacement.REPLICATED + assert _contract().to_dict()["output"] == { + "selected_logp_dtype": "fp32", + "lse_dtype": "fp32", + "tp_placement": "replicated", + } + + +def test_cross_rank_fingerprint_is_rank_independent_and_content_sensitive(): + rank0 = _contract(sharding=_sharding(tp_rank=0)) + rank1 = _contract(sharding=_sharding(tp_rank=1, cp_rank=1)) + + assert rank0.cross_rank_fingerprint() == rank1.cross_rank_fingerprint() + + different_mask = _contract( + mask=_mask(active_mask=(True, True, True, True, True, True, True, False)) + ) + assert rank0.cross_rank_fingerprint() != different_mask.cross_rank_fingerprint() + + +def test_provenance_records_the_active_mask_digest(): + provenance_mask = _contract().to_dict()["mask"] + + assert provenance_mask["active_mask_sha256"] == _mask().active_mask_sha256 + assert len(provenance_mask["active_mask_sha256"]) == 64 + + same_count_different_mask = _mask( + active_mask=(True, True, True, True, True, False, False, False) + ) + assert same_count_different_mask.active_token_count == _mask().active_token_count + assert same_count_different_mask.active_mask_sha256 != _mask().active_mask_sha256 + + +def test_padding_only_shard_is_constructible_for_the_identity_partial(): + sharding = _sharding( + vocab_shard_bounds=((0, QWEN3_REAL_VOCAB), (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB)), + ) + + assert sharding.local_vocab_start == 0 + assert sharding.vocab_shard_bounds[1] == (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB) + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 0 diff --git a/tests/test_vime_logprob_provider.py b/tests/test_vime_logprob_provider.py new file mode 100644 index 00000000..021f4e9f --- /dev/null +++ b/tests/test_vime_logprob_provider.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU coverage for the optional Vime WS2 selected-logprob adapter.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.integrations.vime.logp import SelectedLogprobProviderUnavailable, provider + + +def _request(*, cp_rank: int = 0, with_entropy: bool = False, keep_mask=None): + logits = torch.tensor( + [[0.25, -0.5, 1.0, 0.1, -0.3, 0.6, -0.7, 0.4] for _ in range(3)], + dtype=torch.float32, + requires_grad=True, + ) + return SimpleNamespace( + logits=logits, + target_ids=torch.tensor([2, 5, 0]), + tensor_parallel_group=None, + context_parallel=SimpleNamespace( + world_size=2, + rank=cp_rank, + layout="zigzag", + ), + with_entropy=with_entropy, + with_entropy_grad=with_entropy, + log_prob_keep_mask=keep_mask, + metadata={ + "real_vocab_size": 7, + "padded_vocab_size": 8, + "tp_rank": 0, + "tp_world_size": 1, + "num_vocab_tiles": 4, + }, + ) + + +def test_provider_runs_locally_with_cp2_row_metadata(): + request = _request(cp_rank=1) + + result = provider(request) + reference = torch.log_softmax(request.logits[:, :7], dim=-1)[ + torch.arange(request.logits.size(0)), request.target_ids + ] + + assert result.selected_logprobs.shape == (3, 1) + torch.testing.assert_close(result.selected_logprobs.squeeze(-1), reference) + assert result.backend_id == "pytorch-vocab-parallel-logp-ws2" + assert result.provenance["cp_row_ownership"] == { + "cp_rank": 1, + "cp_world_size": 2, + "layout": "zigzag", + "local_token_rows": 3, + "cp_is_merge_axis": False, + } + + +def test_provider_entropy_preserves_vime_semantics_and_autograd(): + request = _request(with_entropy=True) + + result = provider(request) + reference_logits = request.logits.detach().clone().requires_grad_(True) + log_probs = torch.log_softmax(reference_logits[:, :7], dim=-1) + reference_logp = log_probs[torch.arange(reference_logits.size(0)), request.target_ids] + reference_entropy = -(log_probs.exp() * log_probs).sum(dim=-1) + + torch.testing.assert_close(result.selected_logprobs.squeeze(-1), reference_logp) + torch.testing.assert_close(result.entropy, reference_entropy) + (result.selected_logprobs.sum() + result.entropy.sum()).backward() + (reference_logp.sum() + reference_entropy.sum()).backward() + torch.testing.assert_close(request.logits.grad[:, :7], reference_logits.grad[:, :7]) + assert bool((request.logits.grad[:, 7] == 0).all()) + + +def test_provider_rejects_top_p_replay_without_changing_its_semantics(): + request = _request(keep_mask=torch.ones((3, 8), dtype=torch.bool)) + + with pytest.raises(SelectedLogprobProviderUnavailable, match="top-p replay"): + provider(request) + + +def test_provider_rejects_local_vocab_metadata_that_cannot_describe_tp_ownership(): + request = _request() + request.metadata["padded_vocab_size"] = 16 + + with pytest.raises(SelectedLogprobProviderUnavailable, match="cover padded_vocab_size"): + provider(request) diff --git a/tests/test_vime_qwen3_example.py b/tests/test_vime_qwen3_example.py new file mode 100644 index 00000000..007d4ba1 --- /dev/null +++ b/tests/test_vime_qwen3_example.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from examples.vime_qwen3_8b_tp2_cp2.run import ( + build_report, + load_config, + validate_config, + validate_runtime_evidence, +) + + +ROOT = Path(__file__).parents[1] +CONFIG = ROOT / "examples" / "vime_qwen3_8b_tp2_cp2" / "qwen3_8b_tp2_cp2.json" + + +def test_qwen3_example_config_is_strict_and_explicit(): + config = load_config(CONFIG) + validate_config(config) + assert config["training"]["tensor_model_parallel_size"] == 2 + assert config["training"]["context_parallel_size"] == 2 + assert config["selected_logprob_provider"]["mode"] == "strict" + + +def test_qwen3_example_report_does_not_claim_unread_back_attention_or_ffn(tmp_path): + config = load_config(CONFIG) + report = build_report( + config, + vime_root=tmp_path / "vime", + rl_kernel_root=tmp_path / "rl-kernel", + command=["bash", "run.sh"], + status="passed", + returncode=0, + log_text="Selected-logprob provider active: backend_id=pytorch-vocab-parallel-logp-ws2", + log_path=tmp_path / "run.log", + ) + assert report["status"] == "passed" + assert report["claim_boundary"]["qwen3_8b_tp2_cp2_vime_training"] is True + assert report["claim_boundary"]["attention_train_infer_consistency"] == "unclaimed" + assert report["claim_boundary"]["ffn_train_infer_consistency"] == "unclaimed" + assert report["provider"]["fallback_observed"] is False + + +def test_qwen3_example_fails_closed_when_provider_marker_is_missing(tmp_path): + config = load_config(CONFIG) + report = build_report( + config, + vime_root=tmp_path / "vime", + rl_kernel_root=tmp_path / "rl-kernel", + command=["bash", "run.sh"], + status="passed", + returncode=0, + log_text="training completed without provider provenance", + log_path=None, + ) + assert report["status"] == "failed" + assert report["claim_boundary"]["qwen3_8b_tp2_cp2_vime_training"] is False + + +def _runtime_evidence(): + return { + "schema_version": "rlkernel.operator_runtime_evidence.v1", + "operators": { + "attention": { + "training": {"implementation_id": "rlk.attn", "backend_id": "rlk", "contract_id": "a"}, + "rollout": {"implementation_id": "rlk.attn", "backend_id": "rlk", "contract_id": "a"}, + "comparison": { + "passed": True, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "dq_max_abs": 0.0, + "dk_max_abs": 0.0, + "dv_max_abs": 0.0, + }, + }, + "ffn": { + "training": {"implementation_id": "rlk.ffn", "backend_id": "rlk", "contract_id": "f"}, + "rollout": {"implementation_id": "rlk.ffn", "backend_id": "rlk", "contract_id": "f"}, + "comparison": { + "passed": True, + "out_max_abs": 0.0, + "dx_max_abs": 0.0, + "dw_max_abs": 0.0, + }, + }, + }, + } + + +def test_qwen3_example_accepts_only_exact_zero_runtime_evidence(tmp_path): + evidence = _runtime_evidence() + validate_runtime_evidence(evidence) + config = load_config(CONFIG) + report = build_report( + config, + vime_root=tmp_path / "vime", + rl_kernel_root=tmp_path / "rl-kernel", + command=["bash", "run.sh"], + status="passed", + returncode=0, + log_text="Selected-logprob provider active: backend_id=pytorch-vocab-parallel-logp-ws2", + log_path=None, + runtime_evidence=evidence, + ) + assert report["claim_boundary"]["attention_train_infer_consistency"] == "passed" + assert report["claim_boundary"]["ffn_train_infer_consistency"] == "passed" + + +def test_qwen3_example_rejects_nonzero_runtime_evidence(): + evidence = _runtime_evidence() + evidence["operators"]["attention"]["comparison"]["out_max_abs"] = 1e-6 + with pytest.raises(ValueError, match="attention"): + validate_runtime_evidence(evidence) + + +@pytest.mark.parametrize("bad_path", ["", "other.provider"]) +def test_qwen3_example_rejects_non_rlkernel_provider(bad_path): + config = load_config(CONFIG) + config["selected_logprob_provider"]["path"] = bad_path + with pytest.raises(ValueError, match="RL-Kernel Vime provider"): + validate_config(config) diff --git a/tests/test_vocab_parallel_logp.py b/tests/test_vocab_parallel_logp.py new file mode 100644 index 00000000..a0dfe5d8 --- /dev/null +++ b/tests/test_vocab_parallel_logp.py @@ -0,0 +1,560 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic vocab-parallel TP logprob reference tests""" + +from __future__ import annotations + +import queue +import tempfile +import traceback +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.kernels.gtest.tolerance import load_contract +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobContract, + LogprobContractError, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + BACKEND_ID, + VocabParallelLogprobOp, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +REAL_VOCAB = 27 +PADDED_VOCAB = 32 +NUM_TILES = 8 +NUM_TOKENS = 6 +ACTIVE = (True, True, True, True, True, False) + + +def _even_bounds(padded: int, world: int) -> tuple[tuple[int, int], ...]: + shard = padded // world + return tuple( + (rank * shard, padded if rank == world - 1 else (rank + 1) * shard) for rank in range(world) + ) + + +def _contract( + *, + tp_rank: int = 0, + tp_world_size: int = 1, + bounds: tuple[tuple[int, int], ...] | None = None, + real_vocab: int = REAL_VOCAB, + padded_vocab: int = PADDED_VOCAB, + num_tokens: int = NUM_TOKENS, + active: tuple[bool, ...] = ACTIVE, + dtype: str = "fp32", +) -> LogprobContract: + return LogprobContract( + role="train", + dtype=dtype, + mask=MaskSpec(num_tokens=num_tokens, active_mask=active), + sharding=ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=( + bounds if bounds is not None else _even_bounds(padded_vocab, tp_world_size) + ), + real_vocab_size=real_vocab, + padded_vocab_size=padded_vocab, + ), + reduction=ReductionSpec(), + ) + + +def _inputs(dtype=torch.float32, seed: int = 2026): + torch.manual_seed(seed) + logits = torch.randn(NUM_TOKENS, PADDED_VOCAB, dtype=torch.float32).to(dtype) + targets = torch.tensor([1, 5, REAL_VOCAB - 1, 0, 13, -100]) + return logits, targets + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + view_dtype = { + torch.float32: torch.int32, + torch.bfloat16: torch.int16, + torch.float16: torch.int16, + }[tensor.dtype] + return tensor.contiguous().view(view_dtype) + + +def _bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and bool((_bits(a) == _bits(b)).all()) + + +def _case_shard_size_mismatch(): + logits, targets = _inputs() + return logits, targets, _contract(tp_rank=0, tp_world_size=2), NUM_TILES, "vocab columns" + + +def _case_mask_length_mismatch(): + logits, targets = _inputs() + contract = _contract(num_tokens=NUM_TOKENS + 1, active=ACTIVE + (True,)) + return logits, targets, contract, NUM_TILES, "num_tokens" + + +def _case_dtype_mismatch(): + logits, targets = _inputs() + return logits, targets, _contract(dtype="bf16"), NUM_TILES, "dtype" + + +def _case_tile_misaligned_bounds(): + # Tile size is 32/8 = 4; a boundary at 6 is misaligned. + logits, targets = _inputs() + contract = _contract(tp_world_size=2, bounds=((0, 6), (6, 32))) + return logits[:, :6], targets, contract, NUM_TILES, "tile" + + +def _case_bad_num_vocab_tiles(): + logits, targets = _inputs() + return logits, targets, _contract(), 7, "num_vocab_tiles" + + +def _case_active_target_out_of_real_vocab(): + logits, targets = _inputs() + bad_targets = targets.clone() + bad_targets[0] = REAL_VOCAB # padding column, active row + return logits, bad_targets, _contract(), NUM_TILES, "real vocabulary" + + +def _case_all_inf_active_row(): + logits, targets = _inputs() + poisoned = logits.clone() + poisoned[0, :] = float("-inf") + return poisoned, targets, _contract(), NUM_TILES, "non-finite" + + +@pytest.mark.parametrize( + "case", + [ + _case_shard_size_mismatch, + _case_mask_length_mismatch, + _case_dtype_mismatch, + _case_tile_misaligned_bounds, + _case_bad_num_vocab_tiles, + _case_active_target_out_of_real_vocab, + _case_all_inf_active_row, + ], + ids=lambda fn: fn.__name__.removeprefix("_case_"), +) +def test_invalid_invocations_fail_loudly(case): + logits, targets, contract, num_tiles, match = case() + with pytest.raises(LogprobContractError, match=match): + VocabParallelLogprobOp()(logits, targets, contract=contract, num_vocab_tiles=num_tiles) + + +class TestSingleRank: + def test_repeated_runs_are_bitwise_identical(self): + contract = _contract() + logits, targets = _inputs() + op = VocabParallelLogprobOp() + logp_a, lse_a = op(logits, targets, contract=contract, num_vocab_tiles=NUM_TILES) + logp_b, lse_b = op(logits, targets, contract=contract, num_vocab_tiles=NUM_TILES) + assert _bitwise_equal(logp_a, logp_b) + assert _bitwise_equal(lse_a, lse_b) + + def test_batch_invariance_same_row_any_context(self): + contract_full = _contract() + logits, targets = _inputs() + op = VocabParallelLogprobOp() + logp_full, lse_full = op(logits, targets, contract=contract_full, num_vocab_tiles=NUM_TILES) + + contract_single = _contract(num_tokens=1, active=(True,)) + logp_one, lse_one = op( + logits[2:3], targets[2:3], contract=contract_single, num_vocab_tiles=NUM_TILES + ) + assert _bitwise_equal(logp_full[2:3], logp_one) + assert _bitwise_equal(lse_full[2:3], lse_one) + + def test_matches_ws1_batch_invariant_logp_within_contract_tolerance(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + contract = _contract(padded_vocab=REAL_VOCAB + 5) + # Use a real==padded contract so the WS1 op sees identical logits. + contract = _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB) + logits, targets = _inputs() + logp, _ = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + ws1 = NativeBatchInvariantLogpOp().apply(logits, targets) + active = torch.tensor(ACTIVE) + assert torch.allclose( + logp[active], ws1[active], atol=tolerance["atol"], rtol=tolerance["rtol"] + ) + + def test_padding_columns_are_excluded_and_finite(self): + contract = _contract() + logits, targets = _inputs() + boosted = logits.clone() + boosted[:, REAL_VOCAB:] = 1e4 # huge padding logits must not leak into LSE + logp, lse = VocabParallelLogprobOp()( + boosted, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + ref_lse = torch.logsumexp(boosted[:, :REAL_VOCAB].float(), dim=-1) + assert torch.isfinite(logp).all() and torch.isfinite(lse).all() + assert torch.allclose(lse, ref_lse, atol=1e-5) + + def test_inactive_rows_zero_filled_lse_still_exported(self): + contract = _contract() + logits, targets = _inputs() + logp, lse = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + assert logp[-1].item() == 0.0 + assert torch.isfinite(lse[-1]) + + +class TestBackward: + def test_grads_match_autograd_oracle(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + contract = _contract() + logits, targets = _inputs() + x = logits.clone().requires_grad_(True) + logp, lse = VocabParallelLogprobOp()( + x, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + (logp.sum() + 0.5 * lse.sum()).backward() + + y = logits.clone().requires_grad_(True) + ref_lse = torch.logsumexp(y[:, :REAL_VOCAB].float(), dim=-1) + safe = targets.clamp(0, REAL_VOCAB - 1) + ref_logp = y[torch.arange(NUM_TOKENS), safe].float() - ref_lse + ref_logp = torch.where(torch.tensor(ACTIVE), ref_logp, torch.zeros_like(ref_logp)) + (ref_logp.sum() + 0.5 * ref_lse.sum()).backward() + + assert torch.allclose(x.grad, y.grad, atol=tolerance["atol"], rtol=tolerance["rtol"]) + assert bool((x.grad[:, REAL_VOCAB:] == 0).all()) + + # No grad requested -> outputs detached from autograd entirely. + logp_ng, lse_ng = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + assert not logp_ng.requires_grad and not lse_ng.requires_grad + + def test_inactive_rows_grad_asymmetry(self): + """The logp term is zeroed on inactive rows; the lse term still flows — + lse is a row property exported (and differentiable) for every row.""" + + contract = _contract() + logits, targets = _inputs() + + x = logits.clone().requires_grad_(True) + _, lse = VocabParallelLogprobOp()(x, targets, contract=contract, num_vocab_tiles=NUM_TILES) + lse.sum().backward() + assert bool((x.grad[-1, :REAL_VOCAB].abs() > 0).any()) + + z = logits.clone().requires_grad_(True) + logp, _ = VocabParallelLogprobOp()(z, targets, contract=contract, num_vocab_tiles=NUM_TILES) + logp.sum().backward() + assert bool((z.grad[-1] == 0).all()) + + def test_entropy_matches_full_vocab_oracle_and_backpropagates(self): + contract = _contract() + logits, targets = _inputs() + x = logits.clone().requires_grad_(True) + logp, _lse, entropy = VocabParallelLogprobOp().apply_with_entropy( + x, + targets, + contract=contract, + num_vocab_tiles=NUM_TILES, + ) + + y = logits.clone().requires_grad_(True) + ref_log_probs = torch.log_softmax(y[:, :REAL_VOCAB].float(), dim=-1) + safe = targets.clamp(0, REAL_VOCAB - 1) + ref_logp = ref_log_probs[torch.arange(NUM_TOKENS), safe] + ref_logp = torch.where(torch.tensor(ACTIVE), ref_logp, torch.zeros_like(ref_logp)) + ref_entropy = -(ref_log_probs.exp() * ref_log_probs).sum(dim=-1) + + torch.testing.assert_close(entropy, ref_entropy) + (logp.sum() + entropy.sum()).backward() + (ref_logp.sum() + ref_entropy.sum()).backward() + torch.testing.assert_close(x.grad, y.grad, atol=2e-5, rtol=2e-5) + + +def test_dispatch_resolves_reference_and_leaves_legacy_untouched(): + registry = KernelRegistry() + contract = _contract() + + result = registry.get_logprob_op(contract) + assert result.capability.backend_id == BACKEND_ID + assert result.provenance["fallback"] is False + assert isinstance(result.op, VocabParallelLogprobOp) + assert ( + result.provenance["contract"]["reduction"]["determinism_scope"] + == DeterminismScope.CROSS_TP_BITWISE.value + ) + + by_id = registry.get_logprob_op(contract, requested_backend=BACKEND_ID) + assert by_id.capability.backend_id == BACKEND_ID + by_kind = registry.get_logprob_op(contract, requested_backend="reference") + assert by_kind.capability.backend_id == BACKEND_ID + + for ops in registry._priority_map.values(): + for candidates in ops.values(): + assert OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP not in candidates + + +# Cross-TP bitwise determinism on real ranks (NCCL, one CUDA device per rank) +TP_REAL_VOCAB = 1000 +TP_PADDED_VOCAB = 1024 +TP_NUM_TILES = 32 # tile = 32 columns +TP_TILE = TP_PADDED_VOCAB // TP_NUM_TILES +TP_NUM_TOKENS = 48 +TP_ACTIVE = tuple(index % 7 != 5 for index in range(TP_NUM_TOKENS)) +TP_DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16} +_SPAWN_TIMEOUT_S = 300 + + +def _cuda_device_count() -> int: + return torch.cuda.device_count() if torch.cuda.is_available() else 0 + + +def _requires_gpus(count: int): + return pytest.mark.skipif( + _cuda_device_count() < count, + reason=f"cross-TP determinism needs {count} CUDA devices to place one rank per device", + ) + + +def _tile_counts(world_size: int, uneven: bool) -> list[int]: + """Tiles per rank; bounds are built from whole tiles so they stay tile-aligned.""" + + counts = [TP_NUM_TILES // world_size for _ in range(world_size)] + counts[-1] += TP_NUM_TILES % world_size + if uneven: + for rank in range(world_size - 1): + if counts[rank] > 1: + counts[rank] -= 1 + counts[-1] += 1 + return counts + + +def _tp_bounds(world_size: int, uneven: bool) -> tuple[tuple[int, int], ...]: + bounds, cursor = [], 0 + for count in _tile_counts(world_size, uneven): + bounds.append((cursor, cursor + count * TP_TILE)) + cursor += count * TP_TILE + return tuple(bounds) + + +def _tp_contract(tp_rank: int, tp_world_size: int, bounds, dtype_name: str) -> LogprobContract: + return _contract( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + bounds=bounds, + real_vocab=TP_REAL_VOCAB, + padded_vocab=TP_PADDED_VOCAB, + num_tokens=TP_NUM_TOKENS, + active=TP_ACTIVE, + dtype=dtype_name, + ) + + +def _tp_inputs(device, dtype, seed: int = 2026): + """Identical logits and targets on every rank, seeded on CPU.""" + + gen = torch.Generator(device="cpu").manual_seed(seed) + logits = torch.randn(TP_NUM_TOKENS, TP_PADDED_VOCAB, generator=gen, dtype=torch.float32) + targets = torch.randint(0, TP_REAL_VOCAB, (TP_NUM_TOKENS,), generator=gen) + active = torch.tensor(TP_ACTIVE) + # Inactive rows carry ignore_index; active_mask stays the sole authority. + targets = torch.where(active, targets, torch.full_like(targets, -100)) + return logits.to(device=device, dtype=dtype), targets.to(device) + + +def _nccl_worker(rank, world_size, init_method, result_queue, scenario, uneven, dtype_name): + import torch.distributed as dist + + try: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + backend="nccl", init_method=init_method, rank=rank, world_size=world_size + ) + dtype = TP_DTYPES[dtype_name] + op = VocabParallelLogprobOp() + bounds = _tp_bounds(world_size, uneven) + logits, targets = _tp_inputs(device, dtype) + tiles = TP_NUM_TILES + + if scenario in {"preflight", "misaligned"}: + if scenario == "preflight": + if rank == 0: + tiles = TP_NUM_TILES * 2 + else: + # Nudge the first boundary off the tile grid, on every rank. + split = bounds[0][1] + TP_TILE // 4 + bounds = ((0, split), (split, bounds[1][1])) + bounds[2:] + + start, end = bounds[rank] + try: + op( + logits[:, start:end].contiguous().clone(), + targets, + contract=_tp_contract(rank, world_size, bounds, dtype_name), + tp_group=dist.group.WORLD, + num_vocab_tiles=tiles, + ) + result_queue.put({"ok": False, "rank": rank, "traceback": "no error raised"}) + except LogprobContractError as exc: + result_queue.put({"ok": True, "rank": rank, "message": str(exc)}) + return + + start, end = bounds[rank] + shard = logits[:, start:end].contiguous().clone().requires_grad_(True) + tp_contract = _tp_contract(rank, world_size, bounds, dtype_name) + logp_tp, lse_tp = op( + shard, + targets, + contract=tp_contract, + tp_group=dist.group.WORLD, + num_vocab_tiles=TP_NUM_TILES, + ) + (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() + + # Same ranks, same inputs, run again: the collectives must not perturb bits. + rerun = logits[:, start:end].contiguous().clone() + logp_re, lse_re = op( + rerun, + targets, + contract=tp_contract, + tp_group=dist.group.WORLD, + num_vocab_tiles=TP_NUM_TILES, + ) + + # In-process TP=1 run on the full logits: the cross-TP claim is that a + # TP=n result equals the TP=1 result, bit for bit. + full = logits.clone().requires_grad_(True) + logp_one, lse_one = op( + full, + targets, + contract=_tp_contract(0, 1, ((0, TP_PADDED_VOCAB),), dtype_name), + num_vocab_tiles=TP_NUM_TILES, + ) + (logp_one.sum() + 0.5 * lse_one.sum()).backward() + + result_queue.put( + { + "ok": True, + "rank": rank, + "logp_bits_match": _bitwise_equal(logp_tp, logp_one), + "lse_bits_match": _bitwise_equal(lse_tp, lse_one), + "grad_bits_match": _bitwise_equal(shard.grad, full.grad[:, start:end]), + "rerun_bits_match": ( + _bitwise_equal(logp_re, logp_tp) and _bitwise_equal(lse_re, lse_tp) + ), + "logp_bit_pattern": _bits(logp_tp.detach().float().cpu()).tolist(), + "lse_bit_pattern": _bits(lse_tp.detach().float().cpu()).tolist(), + } + ) + except Exception: # pragma: no cover - forwarded to the parent process + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + import torch.distributed as dist + + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_nccl_scenario(world_size, scenario="correctness", uneven=False, dtype_name="fp32"): + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmpdir: + init_method = (Path(tmpdir) / "nccl_init").as_uri() + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=_nccl_worker, + args=(rank, world_size, init_method, result_queue, scenario, uneven, dtype_name), + ) + for rank in range(world_size) + ] + results = [] + try: + for process in processes: + process.start() + for _ in range(world_size): + try: + results.append(result_queue.get(timeout=_SPAWN_TIMEOUT_S)) + except queue.Empty: + for process in processes: + if process.is_alive(): + process.terminate() + pytest.fail(f"timed out waiting for NCCL workers (scenario={scenario})") + finally: + for process in processes: + process.join(timeout=30) + if process.is_alive(): + process.terminate() + results.sort(key=lambda item: item["rank"]) + for result in results: + assert result["ok"], result.get("traceback") + for process in processes: + assert process.exitcode == 0 + return results + + +class TestCrossTPBitwise: + """TP=n output == TP=1 output, bit for bit, on real NCCL ranks.""" + + @_requires_gpus(2) + @pytest.mark.parametrize("dtype_name", ["fp32", "bf16"]) + @pytest.mark.parametrize("uneven", [False, True], ids=["even", "uneven"]) + def test_tp2_bitwise_identical_to_tp1(self, uneven, dtype_name): + self._assert_matches_tp1(_run_nccl_scenario(2, uneven=uneven, dtype_name=dtype_name)) + + @_requires_gpus(4) + @pytest.mark.parametrize("dtype_name", ["fp32", "bf16"]) + @pytest.mark.parametrize("uneven", [False, True], ids=["even", "uneven"]) + def test_tp4_bitwise_identical_to_tp1(self, uneven, dtype_name): + self._assert_matches_tp1(_run_nccl_scenario(4, uneven=uneven, dtype_name=dtype_name)) + + @staticmethod + def _assert_matches_tp1(results): + for result in results: + rank = result["rank"] + assert result["logp_bits_match"], f"rank {rank} logp bits differ from TP=1" + assert result["lse_bits_match"], f"rank {rank} lse bits differ from TP=1" + assert result["grad_bits_match"], f"rank {rank} grad bits differ from TP=1" + assert result["rerun_bits_match"], f"rank {rank} bits changed between identical runs" + # Outputs are replicated: every rank must hold identical bits. + for other in results[1:]: + assert results[0]["logp_bit_pattern"] == other["logp_bit_pattern"] + assert results[0]["lse_bit_pattern"] == other["lse_bit_pattern"] + + @_requires_gpus(2) + def test_tp2_and_tp4_agree_with_each_other(self): + """The claim is over TP degrees, so pin TP=2 against TP=4 directly.""" + + if _cuda_device_count() < 4: + pytest.skip("needs 4 CUDA devices to compare TP=2 against TP=4") + tp2 = _run_nccl_scenario(2) + tp4 = _run_nccl_scenario(4) + assert tp2[0]["logp_bit_pattern"] == tp4[0]["logp_bit_pattern"] + assert tp2[0]["lse_bit_pattern"] == tp4[0]["lse_bit_pattern"] + + +class TestCrossTPGuards: + """A disagreement must abort loudly on every rank, not strand ranks in a collective.""" + + @_requires_gpus(2) + def test_preflight_rejects_mismatched_num_vocab_tiles(self): + results = _run_nccl_scenario(2, scenario="preflight") + for result in results: + assert "cross-rank preflight failed" in result["message"] + + @_requires_gpus(2) + def test_misaligned_shard_bounds_rejected(self): + results = _run_nccl_scenario(2, scenario="misaligned") + for result in results: + assert "not aligned to the vocab tile size" in result["message"]