diff --git a/.gitignore b/.gitignore index bab026fa..2fa0a916 100644 --- a/.gitignore +++ b/.gitignore @@ -216,3 +216,10 @@ _dev_notes/ # Local C8 execute dumps; default output is under TMPDIR. ws1-c8-ci.json + +# hipify outputs generated during ROCm builds (sources live in csrc/*.cu, csrc/cuda/, csrc/hip/hip_*.hip) +csrc/*.hip +csrc/*_hip.cpp +csrc/hip/** +!csrc/hip/hip_*.hip +csrc/hip/*_hip.hip diff --git a/_dev_notes/rocm_logprob_implementation_summary.md b/_dev_notes/rocm_logprob_implementation_summary.md new file mode 100644 index 00000000..d019b230 --- /dev/null +++ b/_dev_notes/rocm_logprob_implementation_summary.md @@ -0,0 +1,174 @@ +# ROCm Logprob 实现说明 + +日期：2026-08-22 + +这次工作是在临时分支 `work/ws2-logprob-rocm` 上完成的。它基于 issue #241 的 PR1-PR5 integration，再合入当前 PR4 的最新更新。 + +## 一句话概括 + +新增了一个 ROCm 版 WS2 vocab-parallel logprob 后端：每个 GPU 用 HIP/CUDA 可移植 kernel 计算本地 vocab tile 的 FP32 `(max, sumexp)`，TP 之间仍使用 issue #241 已确定的 all-gather + 固定顺序 merge。这样不会为了追求 ROCm 速度而改变原来的数值契约。 + +## 改了什么 + +### 1. 增加本地 tile partial kernel + +文件：`csrc/deterministic_logp_kernel.cu` + +新增 `deterministic_logp_tile_stats`： + +- 输入一个 TP rank 的 local vocab logits； +- 每个 vocab tile 输出 FP32 `max` 和 `sumexp`； +- 过滤真实 vocab 之外的 padding 列； +- 使用固定 block reduction，不使用 atomic 或不确定的全局归约； +- 同一份 `.cu` 源码可由 CUDA 或 HIP 编译。 + +这个 kernel 只负责 rank-local 计算，不负责 RCCL/NCCL，也不负责全局 LSE 合并。跨 TP 的合并顺序仍由已有 WS2 Python 实现控制。 + +### 2. 增加 ROCm backend 注册 + +文件： + +- `rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py` +- `rl_engine/kernels/ops/rocm/loss/__init__.py` +- `rl_engine/kernels/registry.py` + +新增 backend： + +```text +rocm-vocab-parallel-logp-ws2 +``` + +它只在 ROCm platform 注册，放在 PyTorch reference 前面。backend 继承现有 `VocabParallelLogprobOp`，所以以下逻辑仍然共用一份实现： + +- TP contract 和 preflight； +- tile-aligned shard 检查； +- all-gather transport； +- global tile order merge； +- selected-token owner gather； +- entropy 的显式 rank-order merge； +- backward 和 active-mask 语义。 + +只有在 ROCm native extension 已经加载、且提供新符号时，registry 才会注册这个 production backend。native 不可用时，registry 不会把它显示成可用的 ROCm production 实现，而是明确使用已有的 PyTorch reference backend。这样 backend provenance 是准确的，不会出现“界面显示用了 ROCm，实际却悄悄跑了另一条路径”的情况。 + +如果调用方显式要求 native ROCm backend，但扩展或符号缺失，会直接抛出清晰的 `RuntimeError`，不会在 wrapper 内部静默 fallback。PyTorch reference 始终使用纯 PyTorch tile 统计；ROCm 子类只有在 native 能力已由 registry 检查通过时才启用 HIP tile kernel。 + +### 3. 增加 native Python binding 和类型声明 + +文件： + +- `csrc/ops.cpp` +- `rl_engine/_C.pyi` + +注册了 `deterministic_logp_tile_stats` 的 Python binding 和类型签名。 + +同时修复了一个 ROCm 构建隐患：`deterministic_collective_*` 使用 CUDA IPC，但 ROCm 构建不编译对应源文件。现在这些符号的声明和 pybind 注册在 HIP 编译时都会被排除，避免 ROCm 链接阶段出现 unresolved symbol。 + +Prefix-Shared Attention 的 NVIDIA PTX 注册也改成 HIP 编译时排除，与 PR4 的 source gating 保持一致。 + +### 4. 增加测试 + +文件：`tests/test_rocm_logprob_backend.py` + +覆盖： + +- ROCm backend 继承并保持 WS2 operator surface； +- backend 只在 ROCm registry 中注册； +- native extension 可用时才注册 ROCm production backend；不可用时首选 PyTorch reference； +- native tile kernel 和 binding 存在； +- CPU-only 环境可以导入 ROCm wrapper，不要求 native extension。 +- reference/native 两条执行路径不会互相静默切换。 + +另外补上了和 PR #319/#325 思路一致的验证入口： + +- ROCm native 路径复用 TP2/TP4 的 TP1 对照、重复执行、forward/backward 和 bitwise 检查； +- ROCm 多卡 `TP2 x CP2` CLI 测试要求所有 rank 的实际 backend 都是 + `rocm-vocab-parallel-logp-ws2`，且 provenance 不得标记 fallback； +- 显式请求 native backend 但扩展缺失时，测试要求直接 fail fast。 + +这些测试在非 ROCm 或 GPU 数量不足的环境会跳过；真正执行需要带 ROCm native extension 的多卡机器。 + +## 没有改什么 + +- 没有改 issue #241 的 TP/CP 语义；CP 仍不是 logprob 的 merge axis。 +- 没有用 RCCL `all_reduce` 取代固定顺序的 all-gather + local merge。 +- 没有把 AITER/Composable Kernel 强行接进 strict path。 +- 没有把 SM90/TMA/WGMMA 源文件加入 ROCm build。 +- 没有修改 Vime provider 的 contract 或 entropy 语义。 +- 没有删除任何仓库文件；只是在 ROCm 编译时排除了不适用的 CUDA IPC/PTX 注册。 + +## 验证结果 + +在当前 Windows/CUDA 开发环境中： + +```text +70 passed, 11 skipped +``` + +通过的测试包括： + +- ROCm backend dispatch 测试； +- issue #241 logprob contract 测试； +- vocab-parallel logprob reference 测试； +- Vime selected-logprob provider 测试。 + +CUDA extension build 没有完成，原因是当前机器缺少 Microsoft Visual C++ `cl.exe`；同时本地 `nvcc` 是 CUDA 12.6，而 PyTorch 是 CUDA 12.8。当前机器也没有 `hipcc` 和 ROCm runtime，因此以下项目尚未在本地验证： + +- gfx942/gfx950 HIP 编译； +- MI300X RCCL TP2/TP4； +- ROCm native tile kernel 与 PyTorch reference 的真实数值/bitwise 对比。 + +建议在 ROCm 机器上运行： + +```bash +PYTORCH_ROCM_ARCH=gfx942 RL_KERNEL_REQUIRE_EXT=1 MAX_JOBS=16 \ + python setup.py build_ext --inplace + +PYTHONHASHSEED=0 pytest -q -ra \ + tests/test_rocm_logprob_backend.py \ + tests/test_logprob_contract.py \ + tests/test_vocab_parallel_logp.py \ + tests/test_distributed_logprob_comparison.py +``` + +真实多 GPU 验证还需要补充 RCCL TP2/TP4 的运行命令和结果；在拿到 gfx942 结果前，不应把这个 backend 宣称为已完成性能优化，只能称为 strict semantics-preserving ROCm implementation。 + +## 当前提交 + +实现尚未推送远程；代码位于当前工作分支 `work/ws2-logprob-rocm`。合并基线提交是 `b811abc`，本次实现文件仍在工作区，待 ROCm 环境验证后再拆分成正式 PR commits。 + +## 2026-08-23 补充：ROCm 性能调优 + +在 MI300X 上做完基准测试（`benchmarks/benchmark_rocm_logp.py`，结果见 +`benchmarks/results/pr328_rocm_mi300x/report.md`）后，对 ROCm backend 做了两处调整， +TP contract、tile 顺序 merge、selected-target 传输和 active-mask 语义都没有变： + +1. 新增 ROCm 专用文件 `csrc/hip/hip_deterministic_logp_kernel.hip`（只在 ROCm 构建时编译， + 共享的 `csrc/deterministic_logp_kernel.cu` 保持 SM90 调优版本不动）。其中 + `hip_deterministic_logp_tile_stats` 直接读取 BF16/FP16/FP32 shard（kernel 内部逐元素精确 + 转成 FP32，并自行过滤 padding 列），不再先做一份 FP32 拷贝；每个线程固定处理 8 个连续 + 元素（向量化 load），累加顺序只由 `(BlockSize, Vec)` 决定，与 rank、shard 偏移和存储 + dtype 无关，所以 TP=n 与 TP=1 仍然 bitwise 一致，BF16 直读与 FP32 上转的 partial 也 + bitwise 一致。全 padding 的 tile 现在返回 `(-inf, 0)` identity partial（原来是 `-FLT_MAX`）。 +2. 新增 `hip_deterministic_logp_backward`：从保存的输入 shard 一次 fused pass 生成 + `grad_logits`（`g_logp * (onehot - p) + g_lse * p`，padding 列为 0，非有限行 `p = 0`）， + 替代共享 Python autograd 里约 9 次 `[tokens, vocab]` FP32 elementwise pass。 + `RocmVocabParallelLogprobOp.apply` 走这条 HIP autograd 路径；`apply_with_entropy` + 仍沿用共享路径（带 HIP tile kernel），因为 entropy 梯度本来就需要完整的概率张量。 + +构建期可调参数：`DETERMINISTIC_LOGP_TILE_BLOCK_SIZE`（默认 128）、 +`DETERMINISTIC_LOGP_TILE_VECTOR_ELEMENTS`（默认 8）、`DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE` +（默认 256），通过 `setup.py` 的环境变量注入。 + +## 2026-08-23 补充：Triton vocab-parallel backend + +`apply` 所用的 fused autograd 路径被提成共享实现 +（`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py` 里的 +`apply_with_kernels` / `VocabParallelLogprobKernels`）：backend 只需提供 +`tile_stats` 和 `backward` 两个 kernel，TP 传输、固定 tile 顺序 merge、 +target ownership、mask 语义全部复用。基于这个接口新增了 +`rl_engine/kernels/ops/triton/loss/vocab_parallel_logp.py` +（`triton-vocab-parallel-logp-ws2`）：两个 Triton kernel，按 `BLOCK_V=1024` +从 tile 起点分块归约，masked lane 贡献 identity，所以归约顺序只由 `BLOCK_V` +决定，TP=n 与 TP=1 仍 bitwise 一致；同一份源码可在 CUDA 和 ROCm 上运行。 +registry 在 `cuda`/`rocm` 平台都注册它，排在 PyTorch reference 之前；ROCm 上 +HIP backend 仍然排第一。 diff --git a/benchmarks/benchmark_rocm_logp.py b/benchmarks/benchmark_rocm_logp.py new file mode 100644 index 00000000..1ffd737e --- /dev/null +++ b/benchmarks/benchmark_rocm_logp.py @@ -0,0 +1,2281 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""ROCm benchmark for the WS2 vocab-parallel logprob path (PR #328). + +Operator-only: seeded logits, no checkpoint, tokenizer, or model server. + +Single GPU (TP=1), Qwen3 vocabulary ``V=151936`` (64 tiles of 2374 columns): + +- ``native``: ``torch.logsumexp`` + ``gather`` on FP32 logits (plain PyTorch, + not batch-invariant by contract). +- ``ws1-pytorch`` / ``ws1-triton``: existing single-shard batch-invariant ops. +- ``ws2-reference``: ``pytorch-vocab-parallel-logp-ws2`` (PyTorch tile loop). +- ``ws2-rocm``: ``rocm-vocab-parallel-logp-ws2`` (HIP tile-stats kernel). + +Plus a component table for the HIP ``deterministic_logp_tile_stats`` kernel +against the PyTorch ``_local_tile_stats`` loop it replaces. + +Distributed (one process per GPU, RCCL via ProcessGroupNCCL): + +- ``native``: Megatron-style vocab-parallel logprob (all-reduce MAX, all-reduce + SUM of exp, all-reduce SUM of the owned target logit). +- ``ws2-reference`` and ``ws2-rocm``: the contract-aware WS2 operator with the + fixed tile-order merge; CP ranks shard tokens and never join the merge. + +Every path reports latency (GPU events on a single GPU; synchronized wall +clock and slowest rank per sample when distributed), peak device memory, +FP64 accuracy, repeat bitwise stability, and batch invariance. + +Usage: + python benchmarks/benchmark_rocm_logp.py \ + --warmup 5 --samples 20 --training-samples 10 \ + --output-dir benchmarks/results/pr328_rocm_mi300x +""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import os +import platform +import statistics +import tempfile +import threading +import time +import traceback +from pathlib import Path +from typing import Any, Callable + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.loss.vocab_parallel_logp import ( + CudaVocabParallelLogprobOp, + native_tile_stats_available, +) +from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + VocabParallelLogprobOp, + _local_tile_stats, +) +from rl_engine.kernels.ops.rocm.loss.vocab_parallel_logp import RocmVocabParallelLogprobOp +from rl_engine.kernels.ops.triton.loss.vocab_parallel_logp import TritonVocabParallelLogprobOp + +REAL_VOCAB = 151936 # Qwen3 tokenizer/lm_head width; 151936 = 64 * 2374, so no padding. +NUM_TILES = 64 +IGNORE_INDEX = -100 +LOGIT_SCALE = 2.0 +SINGLE_TOKENS = (1, 8, 32, 128, 512, 2048) +SINGLE_DTYPES = ("bf16", "fp32") +DISTRIBUTED_TOKENS = (256, 2048) +TOPOLOGIES = ( + ("tp2", 2, 1), + ("tp4", 4, 1), + ("tp8", 8, 1), + ("tp2_cp2", 2, 2), + ("tp4_cp2", 4, 2), + ("tp2_cp4", 2, 4), +) +DISTRIBUTED_PATHS = ("native", "ws2-reference", "ws2-triton", "ws2-cuda", "ws2-rocm") +WS2_KERNEL_PATHS = ("ws2-triton", "ws2-cuda", "ws2-rocm") +_DTYPES = {"bf16": torch.bfloat16, "fp32": torch.float32} +_SPAWN_TIMEOUT_S = 1800 + + +# --------------------------------------------------------------------------- helpers + + +def _percentile(values: list[float], percentile: float) -> float: + ordered = sorted(values) + if not ordered: + return float("nan") + position = (len(ordered) - 1) * percentile + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def _summary_ms(values: list[float]) -> dict[str, float]: + return { + "median_ms": statistics.median(values), + "p95_ms": _percentile(values, 0.95), + "min_ms": min(values), + "max_ms": max(values), + } + + +def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float: + actual_float = actual.detach().double() + expected_float = expected.detach().double() + denominator = torch.linalg.vector_norm(expected_float) + if denominator.item() == 0.0: + return float(torch.linalg.vector_norm(actual_float - expected_float).item()) + return float((torch.linalg.vector_norm(actual_float - expected_float) / denominator).item()) + + +def _accuracy(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + difference = actual.detach().double() - expected.detach().double() + return { + "max_abs": float(difference.abs().max().item()) if difference.numel() else 0.0, + "relative_l2": _relative_l2(actual, expected), + } + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + return tensor.detach().float().contiguous().view(torch.int32) + + +def _bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and bool(torch.equal(_bits(a), _bits(b))) + + +def _mismatch_count(a: torch.Tensor, b: torch.Tensor) -> int: + return int((_bits(a) != _bits(b)).sum().item()) + + +def _gpu_event_samples(function: Callable[[], Any], *, warmup: int, samples: int) -> list[float]: + for _ in range(warmup): + function() + torch.cuda.synchronize() + events = [] + for _ in range(samples): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + function() + end.record() + events.append((start, end)) + torch.cuda.synchronize() + return [float(start.elapsed_time(end)) for start, end in events] + + +def _host_wall_samples(function: Callable[[], Any], *, warmup: int, samples: int) -> list[float]: + """Wall-clock timing for host execution, where CUDA events do not apply.""" + for _ in range(warmup): + function() + timings = [] + for _ in range(samples): + start = time.perf_counter() + function() + timings.append((time.perf_counter() - start) * 1000.0) + return timings + + +def _timed_samples( + function: Callable[[], Any], *, warmup: int, samples: int, device: torch.device +) -> list[float]: + if device.type == "cuda": + return _gpu_event_samples(function, warmup=warmup, samples=samples) + return _host_wall_samples(function, warmup=warmup, samples=samples) + + +def _device_synchronize(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize() + + +def _device_empty_cache(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.empty_cache() + else: + gc.collect() + + +def _rss_mib() -> float: + with open("/proc/self/statm", "r", encoding="ascii") as handle: + resident_pages = int(handle.read().split()[1]) + return resident_pages * os.sysconf("SC_PAGE_SIZE") / (1024.0 * 1024.0) + + +def _host_peak_rss_mib(function: Callable[[], Any]) -> float: + """Peak resident-set increase during one host call, sampled from /proc. + + This is the closest host analogue of ``torch.cuda.max_memory_allocated``, + but it is an RSS high-water delta rather than an allocator statistic: it + includes caching-allocator reuse and page-level granularity, so it is an + approximation and not directly comparable to the device figures. A call + served entirely from already-resident pages can legitimately report ~0. + """ + gc.collect() + baseline = _rss_mib() + peak = baseline + stop = threading.Event() + + def sampler() -> None: + nonlocal peak + while not stop.is_set(): + peak = max(peak, _rss_mib()) + stop.wait(0.001) + + thread = threading.Thread(target=sampler, daemon=True) + thread.start() + try: + function() + finally: + stop.set() + thread.join() + return float(max(peak, _rss_mib()) - baseline) + + +def _peak_memory_mib(function: Callable[[], Any], device: torch.device) -> float: + """Peak memory used by one call, above what was live before it.""" + if device.type != "cuda": + return _host_peak_rss_mib(function) + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + baseline = torch.cuda.memory_allocated() + function() + torch.cuda.synchronize() + return float((torch.cuda.max_memory_allocated() - baseline) / (1024.0 * 1024.0)) + + +def _seeded_logits( + num_tokens: int, vocab: int, *, seed: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Identical FP32 logits, targets, and active mask on every rank. + + Generated on-device from a seeded CUDA generator so multi-GB inputs are not + materialized on the host; the same seed yields identical values on every + MI300X rank. + """ + generator = torch.Generator(device=device).manual_seed(seed) + logits = torch.randn(num_tokens, vocab, generator=generator, device=device) * LOGIT_SCALE + targets = torch.randint(0, vocab, (num_tokens,), generator=generator, device=device) + active = (torch.arange(num_tokens, device=device) % 7) != 5 + return logits, targets, active + + +def _fp64_oracle( + logits_fp32: torch.Tensor, targets: torch.Tensor, active: torch.Tensor, real_vocab: int +) -> tuple[torch.Tensor, torch.Tensor]: + z = logits_fp32[:, :real_vocab].double() + lse = torch.logsumexp(z, dim=-1) + safe = torch.where(active, targets, torch.zeros_like(targets)) + selected = z.gather(1, safe.unsqueeze(1)).squeeze(1) + logp = torch.where(active, selected - lse, torch.zeros_like(lse)) + return logp, lse + + +def _contract( + *, + num_tokens: int, + active: tuple[bool, ...], + tp_rank: int, + tp_world_size: int, + bounds: tuple[tuple[int, int], ...], + real_vocab: int, + padded_vocab: int, + dtype: str, + cp_rank: int = 0, + cp_world_size: int = 1, +) -> 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, + real_vocab_size=real_vocab, + padded_vocab_size=padded_vocab, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ), + reduction=ReductionSpec(), + ) + + +# --------------------------------------------------------------------------- single GPU + + +def _native_logp( + logits: torch.Tensor, targets: torch.Tensor, active: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + z = logits.float() + lse = torch.logsumexp(z, dim=-1) + safe = torch.where(active, targets, torch.zeros_like(targets)) + selected = z.gather(1, safe.unsqueeze(1)).squeeze(1) + logp = torch.where(active, selected - lse, torch.zeros_like(lse)) + return logp, lse + + +def _single_gpu_paths( + device: torch.device, +) -> dict[str, tuple[Callable[..., Any], Callable[..., Any]]]: + ws1_pytorch = NativeBatchInvariantLogpOp() + ws2_reference = VocabParallelLogprobOp() + + def ws1_pytorch_fn(logits, targets, active, contract): + ignore_targets = torch.where(active, targets, torch.full_like(targets, IGNORE_INDEX)) + return ws1_pytorch.forward_with_lse(logits, ignore_targets, IGNORE_INDEX, validate=False) + + def ws1_pytorch_train(logits, targets, active, contract): + ignore_targets = torch.where(active, targets, torch.full_like(targets, IGNORE_INDEX)) + return ws1_pytorch.apply(logits, ignore_targets, IGNORE_INDEX, validate=False) + + def ws2_reference_fn(logits, targets, active, contract): + return ws2_reference.apply( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES, validate=False + ) + + # path -> (forward returning (logp, lse), training forward returning logp with autograd) + paths: dict[str, tuple[Callable[..., Any], Callable[..., Any]]] = { + "native": ( + lambda logits, targets, active, contract: _native_logp(logits, targets, active), + lambda logits, targets, active, contract: _native_logp(logits, targets, active)[0], + ), + "ws1-pytorch": (ws1_pytorch_fn, ws1_pytorch_train), + } + if device.type != "cuda": + # Triton and the native extensions have no host backend; the reference + # tile loop and the plain-PyTorch baseline are the whole CPU story. + paths["ws2-reference"] = (ws2_reference_fn, lambda *a: ws2_reference_fn(*a)[0]) + return paths + try: + from rl_engine.kernels.ops.triton.loss.batch_invariant_logp import ( + TritonBatchInvariantLogpOp, + ) + + ws1_triton = TritonBatchInvariantLogpOp() + probe = torch.randn(2, 256, device=device, dtype=torch.bfloat16) + ws1_triton.forward_with_lse(probe, torch.zeros(2, device=device, dtype=torch.long)) + torch.cuda.synchronize() + + def ws1_triton_fn(logits, targets, active, contract): + ignore_targets = torch.where(active, targets, torch.full_like(targets, IGNORE_INDEX)) + return ws1_triton.forward_with_lse(logits, ignore_targets, IGNORE_INDEX) + + def ws1_triton_train(logits, targets, active, contract): + ignore_targets = torch.where(active, targets, torch.full_like(targets, IGNORE_INDEX)) + return ws1_triton.apply(logits, ignore_targets, IGNORE_INDEX) + + paths["ws1-triton"] = (ws1_triton_fn, ws1_triton_train) + except Exception as exc: # pragma: no cover - environment dependent + print(f"ws1-triton unavailable: {exc}") + paths["ws2-reference"] = ( + ws2_reference_fn, + lambda *a: ws2_reference_fn(*a)[0], + ) + ws2_triton = TritonVocabParallelLogprobOp() + + def ws2_triton_fn(logits, targets, active, contract): + return ws2_triton.apply( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES, validate=False + ) + + paths["ws2-triton"] = (ws2_triton_fn, lambda *a: ws2_triton_fn(*a)[0]) + + if native_tile_stats_available(): + ws2_cuda = CudaVocabParallelLogprobOp() + + def ws2_cuda_fn(logits, targets, active, contract): + return ws2_cuda.apply( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES, validate=False + ) + + paths["ws2-cuda"] = (ws2_cuda_fn, lambda *a: ws2_cuda_fn(*a)[0]) + + if torch.version.hip is not None: + ws2_rocm = RocmVocabParallelLogprobOp() + + def ws2_rocm_fn(logits, targets, active, contract): + return ws2_rocm.apply( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES, validate=False + ) + + paths["ws2-rocm"] = (ws2_rocm_fn, lambda *a: ws2_rocm_fn(*a)[0]) + return paths + + +def _single_gpu_benchmarks( + *, + warmup: int, + samples: int, + training_samples: int, + tokens: tuple[int, ...], + device: torch.device, +) -> dict[str, Any]: + if device.type == "cuda": + torch.cuda.set_device(device) + paths = _single_gpu_paths(device) + bounds = ((0, REAL_VOCAB),) + cases: list[dict[str, Any]] = [] + validate_overhead: list[dict[str, Any]] = [] + # The validate=True overhead is measured on the fastest native backend the + # platform actually has: HIP on ROCm, the CUDA tile-stats kernel on CUDA, + # and the reference tile loop on the host. + if torch.version.hip is not None: + validate_op: Any = RocmVocabParallelLogprobOp() + validate_op_path = "ws2-rocm" + elif device.type == "cuda" and native_tile_stats_available(): + validate_op = CudaVocabParallelLogprobOp() + validate_op_path = "ws2-cuda" + else: + validate_op = VocabParallelLogprobOp() + validate_op_path = "ws2-reference" + + for dtype_name in SINGLE_DTYPES: + dtype = _DTYPES[dtype_name] + for num_tokens in tokens: + logits_fp32, targets, active = _seeded_logits( + num_tokens, REAL_VOCAB, seed=2026 + num_tokens, device=device + ) + logits = logits_fp32.to(dtype).contiguous() + oracle_logp, oracle_lse = _fp64_oracle(logits.float(), targets, active, REAL_VOCAB) + logits_fp32 = None + contract = _contract( + num_tokens=num_tokens, + active=tuple(bool(flag) for flag in active.tolist()), + tp_rank=0, + tp_world_size=1, + bounds=bounds, + real_vocab=REAL_VOCAB, + padded_vocab=REAL_VOCAB, + dtype=dtype_name, + ) + row = min(3, num_tokens - 1) + row_contract = _contract( + num_tokens=1, + active=(bool(active[row].item()),), + tp_rank=0, + tp_world_size=1, + bounds=bounds, + real_vocab=REAL_VOCAB, + padded_vocab=REAL_VOCAB, + dtype=dtype_name, + ) + outputs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + for path_name, (path, train_path) in paths.items(): + + def forward(): + return path(logits, targets, active, contract) + + def train_step(): + leaf = logits.detach().clone().requires_grad_(True) + logp = train_path(leaf, targets, active, contract) + (logp * active).sum().backward() + return leaf.grad + + try: + first_logp, first_lse = forward() + second_logp, second_lse = forward() + row_logp, row_lse = path( + logits[row : row + 1].contiguous(), + targets[row : row + 1], + active[row : row + 1], + row_contract, + ) + train_step() + _device_synchronize(device) + except Exception as exc: + print(f"{path_name} failed for {dtype_name} M={num_tokens}: {exc}") + continue + outputs[path_name] = (first_logp.detach(), first_lse.detach()) + forward_times = _timed_samples( + forward, warmup=warmup, samples=samples, device=device + ) + train_times = _timed_samples( + train_step, + warmup=max(1, warmup // 2), + samples=training_samples, + device=device, + ) + grad = train_step() + _device_synchronize(device) + active_idx = active.nonzero().squeeze(1) + cases.append( + { + "dtype": dtype_name, + "tokens": num_tokens, + "path": path_name, + "forward": _summary_ms(forward_times), + "train_fwd_bwd": _summary_ms(train_times), + "forward_peak_mib": _peak_memory_mib(forward, device), + "train_peak_mib": _peak_memory_mib(train_step, device), + "logp_vs_fp64": _accuracy(first_logp[active_idx], oracle_logp[active_idx]), + "lse_vs_fp64": _accuracy(first_lse, oracle_lse), + "repeat_bitwise": _bitwise_equal(first_logp, second_logp) + and _bitwise_equal(first_lse, second_lse), + "batch_invariant": _bitwise_equal(row_logp[0], first_logp[row]) + and _bitwise_equal(row_lse[0], first_lse[row]), + "grad_finite": bool(torch.isfinite(grad).all().item()), + } + ) + print( + f"single {dtype_name} M={num_tokens:5d} {path_name:14s} " + f"fwd={cases[-1]['forward']['median_ms']:.4f}ms " + f"train={cases[-1]['train_fwd_bwd']['median_ms']:.4f}ms " + f"peak={cases[-1]['train_peak_mib']:.1f}MiB", + flush=True, + ) + if "ws2-reference" in outputs: + ref_logp, ref_lse = outputs["ws2-reference"] + for kernel_path in WS2_KERNEL_PATHS: + if kernel_path not in outputs: + continue + k_logp, k_lse = outputs[kernel_path] + for case in cases: + if ( + case["dtype"] == dtype_name + and case["tokens"] == num_tokens + and case["path"] == kernel_path + ): + case["mismatch_vs_reference"] = _mismatch_count( + k_logp, ref_logp + ) + _mismatch_count(k_lse, ref_lse) + case["rel_l2_vs_reference"] = max( + _relative_l2(k_logp, ref_logp), _relative_l2(k_lse, ref_lse) + ) + # validate=True production entry point overhead (host-side checks + .item() sync) + if dtype_name == "bf16": + + def validated(): + return validate_op.apply( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES, validate=True + ) + + def unvalidated(): + return validate_op.apply( + logits, + targets, + contract=contract, + num_vocab_tiles=NUM_TILES, + validate=False, + ) + + validate_overhead.append( + { + "tokens": num_tokens, + "path": validate_op_path, + "validate_true": _summary_ms( + _timed_samples(validated, warmup=warmup, samples=samples, device=device) + ), + "validate_false": _summary_ms( + _timed_samples( + unvalidated, warmup=warmup, samples=samples, device=device + ) + ), + } + ) + logits = oracle_logp = oracle_lse = outputs = None + _device_empty_cache(device) + + return { + "cases": cases, + "validate_overhead": validate_overhead, + "paths": list(paths), + "device": device.type, + } + + +def _tile_stats_component( + *, warmup: int, samples: int, tokens: tuple[int, ...], device: torch.device +) -> list[dict[str, Any]]: + """Native tile-stats kernel versus the PyTorch tile loop it replaces. + + ``hip_deterministic_logp_tile_stats`` on ROCm, ``deterministic_logp_tile_stats`` + (``csrc/deterministic_logp_kernel.cu``) on CUDA. The two kernels share the + same fixed per-tile reduction contract, so the row means the same thing on + either platform; the ``hip_*`` result keys are kept for schema stability. + """ + if device.type != "cuda": + return [] + rows: list[dict[str, Any]] = [] + tile = REAL_VOCAB // NUM_TILES + for dtype_name in SINGLE_DTYPES: + dtype = _DTYPES[dtype_name] + for num_tokens in tokens: + logits_fp32, _, _ = _seeded_logits( + num_tokens, REAL_VOCAB, seed=99 + num_tokens, device=device + ) + logits = logits_fp32.to(dtype).contiguous() + z32 = logits.float() + logits_fp32 = None + + def pytorch_loop(): + return _local_tile_stats(z32, tile) + + tile_stats = getattr( + _C, "hip_deterministic_logp_tile_stats", _C.deterministic_logp_tile_stats + ) + kernel_symbol = ( + "hip_deterministic_logp_tile_stats" + if hasattr(_C, "hip_deterministic_logp_tile_stats") + else "deterministic_logp_tile_stats" + ) + + def hip_kernel_fp32(): + return tile_stats(z32, 0, REAL_VOCAB, NUM_TILES) + + def hip_kernel_input_dtype(): + return tile_stats(logits, 0, REAL_VOCAB, NUM_TILES) + + ref_m, ref_s = pytorch_loop() + hip_m, hip_s = hip_kernel_fp32() + hip_m2, hip_s2 = hip_kernel_fp32() + rows.append( + { + "dtype": dtype_name, + "tokens": num_tokens, + "kernel_symbol": kernel_symbol, + "pytorch_loop": _summary_ms( + _timed_samples(pytorch_loop, warmup=warmup, samples=samples, device=device) + ), + "hip_fp32_input": _summary_ms( + _timed_samples( + hip_kernel_fp32, warmup=warmup, samples=samples, device=device + ) + ), + "hip_native_dtype_input": _summary_ms( + _timed_samples( + hip_kernel_input_dtype, warmup=warmup, samples=samples, device=device + ) + ), + "pytorch_loop_peak_mib": _peak_memory_mib(pytorch_loop, device), + "hip_peak_mib": _peak_memory_mib(hip_kernel_fp32, device), + "max_bitwise": _bitwise_equal(hip_m, ref_m), + "sumexp_rel_l2": _relative_l2(hip_s, ref_s), + "sumexp_max_rel": float( + ((hip_s - ref_s).abs() / ref_s.abs().clamp_min(1e-30)).max().item() + ), + "repeat_bitwise": _bitwise_equal(hip_m, hip_m2) + and _bitwise_equal(hip_s, hip_s2), + } + ) + print( + f"tile-stats {dtype_name} M={num_tokens:5d} " + f"loop={rows[-1]['pytorch_loop']['median_ms']:.4f}ms " + f"hip={rows[-1]['hip_fp32_input']['median_ms']:.4f}ms", + flush=True, + ) + logits = z32 = None + _device_empty_cache(device) + return rows + + +# --------------------------------------------------------------------------- distributed + + +class _NativeVocabParallelLogp(torch.autograd.Function): + """Megatron-style vocab-parallel logprob with RCCL all-reduce, no fixed order.""" + + @staticmethod + def forward(ctx, local_logits, targets, active, vocab_start, group): + z = local_logits.float() + local_vocab = z.shape[1] + global_max = z.max(dim=-1).values + dist.all_reduce(global_max, op=dist.ReduceOp.MAX, group=group) + sum_exp = (z - global_max.unsqueeze(1)).exp().sum(dim=-1) + dist.all_reduce(sum_exp, op=dist.ReduceOp.SUM, group=group) + lse = global_max + sum_exp.log() + owned = active & (targets >= vocab_start) & (targets < vocab_start + local_vocab) + local_index = torch.where(owned, targets - vocab_start, torch.zeros_like(targets)) + selected = z.gather(1, local_index.unsqueeze(1)).squeeze(1) * owned + dist.all_reduce(selected, op=dist.ReduceOp.SUM, group=group) + logp = torch.where(active, selected - lse, torch.zeros_like(lse)) + ctx.save_for_backward(z, lse, local_index, owned, active) + ctx.input_dtype = local_logits.dtype + ctx.set_materialize_grads(False) + return logp, lse + + @staticmethod + def backward(ctx, grad_logp, grad_lse): + z, lse, local_index, owned, active = ctx.saved_tensors + probabilities = (z - lse.unsqueeze(1)).exp() + scale = torch.zeros_like(lse) + if grad_logp is not None: + scale = scale - grad_logp * active + if grad_lse is not None: + scale = scale + grad_lse + grad = probabilities * scale.unsqueeze(1) + if grad_logp is not None: + grad.scatter_add_( + 1, local_index.unsqueeze(1), (grad_logp * owned).unsqueeze(1).to(grad.dtype) + ) + return grad.to(ctx.input_dtype), None, None, None, None + + +def _tp_bounds(tp_world_size: int) -> tuple[tuple[int, int], ...]: + tile = REAL_VOCAB // NUM_TILES + per_rank = NUM_TILES // tp_world_size + return tuple( + (rank * per_rank * tile, (rank + 1) * per_rank * tile) for rank in range(tp_world_size) + ) + + +def _token_bounds(num_tokens: int, cp_world_size: int) -> tuple[tuple[int, int], ...]: + 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 _distributed_wall_samples( + function: Callable[[], Any], *, warmup: int, samples: int +) -> list[float]: + for _ in range(warmup): + function() + torch.cuda.synchronize() + dist.barrier() + timings = [] + for _ in range(samples): + torch.cuda.synchronize() + start = time.perf_counter() + function() + torch.cuda.synchronize() + timings.append((time.perf_counter() - start) * 1000.0) + dist.barrier() + return timings + + +def _slowest_rank_summary(local_timings: list[float]) -> dict[str, float]: + gathered: list[list[float] | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, local_timings) + slowest = [ + max(float(rank_timings[index]) for rank_timings in gathered if rank_timings is not None) + for index in range(len(local_timings)) + ] + return _summary_ms(slowest) + + +def _all_max(value: float) -> float: + gathered: list[float | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, float(value)) + return max(float(item) for item in gathered if item is not None) + + +def _all_all(flag: bool) -> bool: + gathered: list[bool | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, bool(flag)) + return all(bool(item) for item in gathered) + + +def _all_sum(value: int) -> int: + gathered: list[int | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, int(value)) + return sum(int(item) for item in gathered if item is not None) + + +def _tp_replicated( + logp: torch.Tensor, lse: torch.Tensor, tp_group: Any, tp_world_size: int +) -> bool: + if tp_world_size == 1: + return True + payload = torch.stack([_bits(logp), _bits(lse)]) + gathered = [torch.empty_like(payload) for _ in range(tp_world_size)] + dist.all_gather(gathered, payload, group=tp_group) + return all(torch.equal(gathered[0], other) for other in gathered[1:]) + + +def _distributed_worker( + rank: int, + world_size: int, + init_method: str, + topology: tuple[str, int, int], + tokens_list: tuple[int, ...], + config: dict[str, Any], + result_queue: Any, +) -> None: + name, tp_world_size, cp_world_size = topology + 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, + device_id=device, + ) + cp_rank, tp_rank = divmod(rank, tp_world_size) + tp_group = None + for group_cp_rank in range(cp_world_size): + ranks = list(range(group_cp_rank * tp_world_size, (group_cp_rank + 1) * tp_world_size)) + group = dist.new_group(ranks=ranks) + if rank in ranks: + tp_group = group + bounds = _tp_bounds(tp_world_size) + vocab_start, vocab_end = bounds[tp_rank] + ops: dict[str, Any] = { + "ws2-reference": VocabParallelLogprobOp(), + "ws2-triton": TritonVocabParallelLogprobOp(), + } + if native_tile_stats_available(): + ops["ws2-cuda"] = CudaVocabParallelLogprobOp() + if torch.version.hip is not None: + ops["ws2-rocm"] = RocmVocabParallelLogprobOp() + results: list[dict[str, Any]] = [] + + for num_tokens in tokens_list: + full_logits, full_targets, full_active = _seeded_logits( + num_tokens, REAL_VOCAB, seed=4100 + num_tokens, device=device + ) + token_start, token_end = _token_bounds(num_tokens, cp_world_size)[cp_rank] + local_tokens = token_end - token_start + targets = full_targets[token_start:token_end].contiguous() + active = full_active[token_start:token_end].contiguous() + local_fp32 = full_logits[token_start:token_end] + shard = local_fp32[:, vocab_start:vocab_end].to(torch.bfloat16).contiguous() + # The FP64 oracle sees the BF16-rounded logits every path actually consumes. + oracle_logp, oracle_lse = _fp64_oracle( + local_fp32.to(torch.bfloat16).float(), targets, active, REAL_VOCAB + ) + full_logits = local_fp32 = None + torch.cuda.empty_cache() + contract = _contract( + num_tokens=local_tokens, + active=tuple(bool(flag) for flag in active.tolist()), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + bounds=bounds, + real_vocab=REAL_VOCAB, + padded_vocab=REAL_VOCAB, + dtype="bf16", + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ) + outputs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + for path_name in DISTRIBUTED_PATHS: + if path_name != "native" and path_name not in ops: + continue + if path_name == "native": + + def forward(x=shard): + return _NativeVocabParallelLogp.apply( + x, targets, active, vocab_start, tp_group + ) + + else: + op = ops[path_name] + + def forward(x=shard, op=op): + return op.apply( + x, + targets, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=NUM_TILES, + validate=False, + ) + + def train_step(): + leaf = shard.detach().clone().requires_grad_(True) + logp, _ = forward(leaf) + (logp * active).sum().backward() + return leaf.grad + + first_logp, first_lse = forward() + second_logp, second_lse = forward() + torch.cuda.synchronize() + outputs[path_name] = (first_logp.detach(), first_lse.detach()) + forward_times = _distributed_wall_samples( + forward, warmup=config["warmup"], samples=config["samples"] + ) + train_times = _distributed_wall_samples( + train_step, + warmup=max(1, config["warmup"] // 2), + samples=config["training_samples"], + ) + grad = train_step() + torch.cuda.synchronize() + active_idx = active.nonzero().squeeze(1) + logp_acc = _accuracy(first_logp[active_idx], oracle_logp[active_idx]) + lse_acc = _accuracy(first_lse, oracle_lse) + entry = { + "topology": name, + "tp": tp_world_size, + "cp": cp_world_size, + "tokens": num_tokens, + "tokens_per_cp_rank": local_tokens, + "local_vocab": vocab_end - vocab_start, + "path": path_name, + "forward": _slowest_rank_summary(forward_times), + "train_fwd_bwd": _slowest_rank_summary(train_times), + "forward_peak_mib": _all_max(_peak_memory_mib(forward, device)), + "train_peak_mib": _all_max(_peak_memory_mib(train_step, device)), + "logp_vs_fp64_max_abs": _all_max(logp_acc["max_abs"]), + "logp_vs_fp64_rel_l2": _all_max(logp_acc["relative_l2"]), + "lse_vs_fp64_max_abs": _all_max(lse_acc["max_abs"]), + "tp_replicated": _all_all( + _tp_replicated(first_logp, first_lse, tp_group, tp_world_size) + ), + "repeat_bitwise": _all_all( + _bitwise_equal(first_logp, second_logp) + and _bitwise_equal(first_lse, second_lse) + ), + "grad_finite": _all_all(bool(torch.isfinite(grad).all().item())), + } + results.append(entry) + if rank == 0: + print( + f"dist {name} M={num_tokens:5d} {path_name:14s} " + f"fwd={entry['forward']['median_ms']:.4f}ms " + f"train={entry['train_fwd_bwd']['median_ms']:.4f}ms " + f"peak={entry['train_peak_mib']:.1f}MiB", + flush=True, + ) + ref_logp, ref_lse = outputs["ws2-reference"] + for other_path in ("native",) + WS2_KERNEL_PATHS: + if other_path not in outputs: + continue + o_logp, o_lse = outputs[other_path] + mismatch = _mismatch_count(o_logp, ref_logp) + _mismatch_count(o_lse, ref_lse) + rel = max(_relative_l2(o_logp, ref_logp), _relative_l2(o_lse, ref_lse)) + # Count once per TP group (outputs are replicated inside the group). + mismatch_total = _all_sum(mismatch if tp_rank == 0 else 0) + rel_max = _all_max(rel) + for entry in results: + if entry["tokens"] == num_tokens and entry["path"] == other_path: + entry["mismatch_vs_reference"] = mismatch_total + entry["rel_l2_vs_reference"] = rel_max + shard = oracle_logp = oracle_lse = outputs = None + torch.cuda.empty_cache() + dist.barrier() + if rank == 0: + result_queue.put({"ok": True, "results": results}) + except Exception: + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_world( + topology: tuple[str, int, int], tokens_list: tuple[int, ...], config: dict[str, Any] +) -> list[dict[str, Any]]: + name, tp_world_size, cp_world_size = topology + world_size = tp_world_size * cp_world_size + context = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmpdir: + init_method = (Path(tmpdir) / "rccl_init").as_uri() + result_queue = context.Queue() + processes = [ + context.Process( + target=_distributed_worker, + args=(rank, world_size, init_method, topology, tokens_list, config, result_queue), + ) + for rank in range(world_size) + ] + for process in processes: + process.start() + try: + payload = result_queue.get(timeout=_SPAWN_TIMEOUT_S) + finally: + for process in processes: + process.join(timeout=120) + if process.is_alive(): + process.terminate() + if not payload.get("ok"): + raise RuntimeError(f"{name} rank {payload.get('rank')} failed:\n{payload.get('traceback')}") + return payload["results"] + + +# --------------------------------------------------------------------------- report + + +PATH_DESCRIPTIONS = { + "native": ( + "`torch.logsumexp` + `gather` on FP32 logits (plain PyTorch, not batch-invariant " + "by contract)" + ), + "ws1-pytorch": ( + "`pytorch-batch-invariant-logp-ws1`, the single-shard batch-invariant PyTorch op" + ), + "ws1-triton": ( + "`triton-batch-invariant-logp-ws1`, the single-shard Triton online-softmax op; it has " + "no vocab-parallel (TP) path, so it appears only in the single-GPU results" + ), + "ws2-reference": ( + "`pytorch-vocab-parallel-logp-ws2`, the WS2 vocab-parallel reference operator: a PyTorch " + "tile loop for the per-tile FP32 `(max, sumexp)` partials, all-gather of the partials, " + "fixed global tile-order merge, and a PyTorch autograd backward" + ), + "ws2-triton": ( + "`triton-vocab-parallel-logp-ws2`, the same contract, transport, and merge with two " + "Triton kernels (tile statistics read from the stored shard, fused backward); one " + "source for CUDA and ROCm" + ), + "ws2-cuda": ( + "`cuda-vocab-parallel-logp-ws2`, the same contract, transport, and merge with two " + "CUDA kernels from `csrc/deterministic_logp_kernel.cu`: " + "`deterministic_logp_tile_stats` reads the stored BF16/FP16/FP32 shard directly " + "(16-byte vector loads, no FP32 copy) and `deterministic_logp_backward` produces the " + "gradient in one fused pass" + ), + "ws2-rocm": ( + "`rocm-vocab-parallel-logp-ws2`, the same contract, transport, and merge with two HIP " + "kernels: `hip_deterministic_logp_tile_stats` reads the stored BF16/FP16/FP32 shard " + "directly (8-element vector loads, no FP32 copy) and `hip_deterministic_logp_backward` " + "produces the gradient in one fused pass" + ), +} +DISTRIBUTED_DESCRIPTIONS = { + "native": ( + "`native` is a Megatron-style vocab-parallel logprob using RCCL all-reduce (MAX, SUM of " + "exp, SUM of the owned target logit) through ProcessGroupNCCL" + ), + "ws2-reference": ( + "the WS2 operators all-gather per-tile `(max, sumexp)` partials over RCCL and merge them " + "in fixed global tile order; CP ranks shard tokens and never enter the merge" + ), + "ws2-triton": "", + "ws2-cuda": "", + "ws2-rocm": "", +} + + +class ReportStyle: + """Which measured paths a report shows, how it names them, and what it compares against.""" + + def __init__( + self, + *, + paths: tuple[str, ...], + names: dict[str, str], + baseline: str, + table_tokens: tuple[int, ...] | None, + show_tuning_baseline: bool, + command: str, + ) -> None: + if baseline not in paths: + raise ValueError( + f"report baseline {baseline!r} is not among the reported paths {paths}" + ) + self.paths = paths + self.names = names + self.baseline = baseline + self.table_tokens = table_tokens + self.show_tuning_baseline = show_tuning_baseline + self.command = command + + def name(self, path: str) -> str: + return self.names.get(path, path) + + def keep_tokens(self, tokens: int) -> bool: + return self.table_tokens is None or tokens in self.table_tokens + + +def _fmt_ratio(numerator: float, denominator: float) -> str: + if denominator <= 0: + return "n/a" + return f"{numerator / denominator:.2f}×" + + +def _median_ratio( + numerator: dict[str, Any] | None, denominator: dict[str, Any] | None, key: str +) -> str: + if numerator is None or denominator is None: + return "n/a" + return _fmt_ratio(numerator[key]["median_ms"], denominator[key]["median_ms"]) + + +def _field_ratio( + numerator: dict[str, Any] | None, denominator: dict[str, Any] | None, key: str +) -> str: + if numerator is None or denominator is None: + return "n/a" + return _fmt_ratio(float(numerator[key]), float(denominator[key])) + + +def _range_text(values: list[float], fmt: str = "{:.1f}") -> str: + if not values: + return "n/a" + low, high = fmt.format(min(values)), fmt.format(max(values)) + if low == high: + return low + return f"{low}-{high}" + + +def _lookup(cases: list[dict[str, Any]], **match: Any) -> dict[str, Any] | None: + for case in cases: + if all(case.get(key) == value for key, value in match.items()): + return case + return None + + +def _yes(flag: Any) -> str: + return "yes" if flag else "no" + + +def _default_compare_label(payload: dict[str, Any]) -> str: + env = payload.get("environment", {}) + if env.get("device") == "cpu": + return "cpu" + gpu = str(env.get("gpu", "device")) + for token in ("MI300X", "MI250X", "MI325X", "H100", "H200", "A100", "B200"): + if token.lower() in gpu.lower(): + return token.lower() + return gpu.split()[0].lower() if gpu else "device" + + +def _report_platforms( + payload: dict[str, Any], comparisons: list[str], label: str | None +) -> list[tuple[str, dict[str, Any]]]: + """This run first, then every ``--compare-with`` platform, in the given order. + + Each entry feeds the same tables, so one report can carry several devices. + """ + + platforms: list[tuple[str, dict[str, Any]]] = [ + (label or _default_compare_label(payload), payload) + ] + for spec in comparisons: + if "=" not in spec: + raise ValueError(f"--compare-with must be LABEL=path/to/results.json, got {spec!r}") + other_label, _, other_path = spec.partition("=") + platforms.append( + (other_label.strip(), json.loads(Path(other_path).read_text(encoding="utf-8"))) + ) + return platforms + + +def _write_report( + payload: dict[str, Any], + output_directory: Path, + style: ReportStyle, + platforms: list[tuple[str, dict[str, Any]]] | None = None, +) -> None: + env = payload["environment"] + cfg = payload["config"] + if platforms is None: + platforms = _report_platforms(payload, [], None) + multi = len(platforms) > 1 + plat_head = "Platform | " if multi else "" + plat_div = "---|" if multi else "" + + def plat_cell(label: str) -> str: + return f"{label} | " if multi else "" + + def tagged(rows: list[dict[str, Any]], label: str) -> list[dict[str, Any]]: + return [dict(row, platform=label) for row in rows] + + single = [ + case + for label, pl in platforms + for case in tagged(pl["single_gpu"]["cases"], label) + if case["path"] in style.paths + ] + component = [ + row for label, pl in platforms for row in tagged(pl.get("tile_stats_component", []), label) + ] + distributed = [ + row + for label, pl in platforms + for row in tagged(pl.get("distributed", []), label) + if row["path"] in style.paths + ] + base = style.baseline + base_name = style.name(base) + others = [path for path in style.paths if path != base] + lines: list[str] = [] + add = lines.append + + device_type = env.get("device", "cuda") + is_host = device_type == "cpu" + platform_label = ( + "host (CPU)" if is_host else ("ROCm" if env.get("hip") not in (None, "None") else "CUDA") + ) + if multi: + add("# PR #328 vocab-parallel logprob performance analysis") + else: + add(f"# PR #328 {platform_label} vocab-parallel logprob performance analysis") + add("") + add("> Operator-only benchmark. No model checkpoint or serving engine was used.") + if multi: + add("") + add( + "> Every platform below ran this same harness, so the seeded logits, the " + f"`V={REAL_VOCAB}` / {NUM_TILES}-tile split, and the FP64 oracle are identical and " + "only the device and backend differ. Backends are not available everywhere: " + "`ws2-rocm` is ROCm-only, `ws2-cuda` is CUDA-only, `ws2-triton` compiles from one " + "source on both, and the PyTorch paths are the only ones that also run on the host. " + "A missing row means the backend cannot exist on that platform, not that it failed." + ) + add("") + add("## Environment") + add("") + if multi: + keys = sorted({key for _, pl in platforms for key in pl["environment"]}) + add("| Item | " + " | ".join(label for label, _ in platforms) + " |") + add("|---|" + "---|" * len(platforms)) + for key in keys: + values = " | ".join(str(pl["environment"].get(key, "n/a")) for _, pl in platforms) + add(f"| {key} | {values} |") + else: + add("| Item | Value |") + add("|---|---|") + for key in sorted(env): + add(f"| {key} | {env[key]} |") + add("") + add("## Methodology") + add("") + add( + f"- Qwen3 vocabulary `V={REAL_VOCAB}` split into {NUM_TILES} tiles of " + f"{REAL_VOCAB // NUM_TILES} columns; seeded logits (`randn * {LOGIT_SCALE}`), " + "random targets, every seventh token inactive." + ) + add("- Measured paths:") + for path in style.paths: + add(f" - `{style.name(path)}`: {PATH_DESCRIPTIONS[path]}.") + dist_notes = [ + DISTRIBUTED_DESCRIPTIONS.get(p, "") for p in style.paths if DISTRIBUTED_DESCRIPTIONS.get(p) + ] + if dist_notes and distributed: + collective = "/".join( + sorted( + { + "RCCL" if pl["environment"].get("hip") not in (None, "None") else "NCCL" + for _, pl in platforms + if pl.get("distributed") + } + ) + ) + add( + "- Distributed: one process per GPU; " + + "; ".join(note.replace("RCCL", collective) for note in dist_notes) + + "." + ) + add( + "- Forward returns the selected-token logprob and the vocabulary LSE; forward+backward " + "computes `grad_logits` for `sum(active * logp)`. The WS2 operators run with " + "`validate=False`; the `validate=True` production entry point is measured separately." + ) + if is_host: + add( + "- Timing: `time.perf_counter` wall clock, median and p95. Peak memory is the " + "per-call increase in resident set size sampled from `/proc/self/statm`, which is " + "an allocator-inclusive high-water approximation rather than the exact tensor " + "bytes reported by `torch.cuda.max_memory_allocated` on device runs; treat the " + "host memory column as indicative only." + ) + else: + add( + "- Single-GPU timing: GPU events, median and p95. Distributed timing: synchronized " + "wall clock, slowest rank per sample. Peak memory is the per-call increase in " + "`torch.cuda.max_memory_allocated` (distributed: max over ranks)." + ) + add( + "- Accuracy is against an FP64 `logsumexp` of the same (BF16-rounded) logits. Repeat = " + "two identical calls are bitwise equal; batch-invariant = a row computed alone is " + "bitwise equal to the same row inside the batch; TP-replicated = every TP rank holds " + "identical bits." + ) + add( + f"- {cfg['warmup']} warmups, {cfg['samples']} measured forward samples, " + f"{cfg['training_samples']} measured forward+backward samples. Raw medians, p95, " + "minimum, maximum, and every measured path are in `results.json`." + ) + if style.table_tokens is not None: + add( + "- Tables show " + + ", ".join(f"{t}" for t in style.table_tokens) + + " tokens; the figures cover the full token sweep." + ) + add("") + add("Reproduce this report from the repository root:") + add("") + add("```bash") + add(style.command) + add("```") + add("") + + # ---- key findings + all_single, all_distributed, all_component = single, distributed, component + add("## Key findings") + add("") + for _plat, _payload in platforms: + prefix = f"{_plat} " if multi else "" + single = [c for c in all_single if c["platform"] == _plat] + distributed = [d for d in all_distributed if d["platform"] == _plat] + component = [r for r in all_component if r["platform"] == _plat] + _env = _payload["environment"] + single_label = "Host" if _env.get("device") == "cpu" else "Single GPU" + _tile_symbol = ( + "hip_deterministic_logp_tile_stats" + if _env.get("hip") not in (None, "None") + else "deterministic_logp_tile_stats" + ) + if component: + _tile_symbol = component[0].get("kernel_symbol", _tile_symbol) + for path in others: + name = style.name(path) + speed_fwd, speed_train, mem = [], [], [] + for dtype_name in SINGLE_DTYPES: + for case in single: + if case["dtype"] != dtype_name or case["path"] != path: + continue + if not style.keep_tokens(case["tokens"]): + continue + ref = _lookup(single, dtype=dtype_name, tokens=case["tokens"], path=base) + if ref is None: + continue + speed_fwd.append(ref["forward"]["median_ms"] / case["forward"]["median_ms"]) + speed_train.append( + ref["train_fwd_bwd"]["median_ms"] / case["train_fwd_bwd"]["median_ms"] + ) + mem.append(case["train_peak_mib"] / max(ref["train_peak_mib"], 1e-6)) + if speed_fwd: + add( + f"- {prefix}{single_label}: `{name}` is " + f"{_range_text(speed_fwd, '{:.2f}')}x faster than " + f"`{base_name}` in forward and {_range_text(speed_train, '{:.2f}')}x in " + f"forward+backward, with {_range_text(mem, '{:.2f}')}x its peak memory." + ) + d_fwd, d_train, d_mem, d_abs = [], [], [], [] + for d in distributed: + if d["path"] != path or not style.keep_tokens(d["tokens"]): + continue + ref = _lookup(distributed, topology=d["topology"], tokens=d["tokens"], path=base) + if ref is None: + continue + d_fwd.append(ref["forward"]["median_ms"] / d["forward"]["median_ms"]) + d_train.append(ref["train_fwd_bwd"]["median_ms"] / d["train_fwd_bwd"]["median_ms"]) + d_mem.append(d["train_peak_mib"] / max(ref["train_peak_mib"], 1e-6)) + d_abs.append(d["forward"]["median_ms"]) + if d_fwd: + add( + f"- {prefix}Distributed: `{name}` is " + f"{_range_text(d_fwd, '{:.2f}')}x faster than " + f"`{base_name}` in forward and {_range_text(d_train, '{:.2f}')}x in " + f"forward+backward across {len({d['topology'] for d in distributed})} TP/CP " + f"topologies, at {_range_text(d_mem, '{:.2f}')}x the per-rank peak memory " + f"(absolute forward {_range_text(d_abs, '{:.3f}')} ms)." + ) + for kernel_path in WS2_KERNEL_PATHS: + if kernel_path not in style.paths or "ws1-triton" not in style.paths: + continue + ratios_fwd, ratios_train = [], [] + for dtype_name in SINGLE_DTYPES: + for case in single: + if case["path"] != kernel_path or case["dtype"] != dtype_name: + continue + if not style.keep_tokens(case["tokens"]): + continue + tri = _lookup( + single, dtype=dtype_name, tokens=case["tokens"], path="ws1-triton" + ) + if tri is None: + continue + ratios_fwd.append(case["forward"]["median_ms"] / tri["forward"]["median_ms"]) + ratios_train.append( + case["train_fwd_bwd"]["median_ms"] / tri["train_fwd_bwd"]["median_ms"] + ) + if ratios_fwd: + add( + f"- {prefix}`{style.name(kernel_path)}` runs at " + f"{_range_text(ratios_fwd, '{:.2f}')}x the " + f"latency of `{style.name('ws1-triton')}` in forward and " + f"{_range_text(ratios_train, '{:.2f}')}x in forward+backward with the same " + "peak memory, while carrying the vocab-parallel contract (tile partials, " + "all-gather, fixed tile-order merge, vocab-domain LSE export) that the " + "single-shard Triton op does not provide; the gap is the operator's fixed " + "Python/launch floor, not the kernels." + ) + if component: + comp_speed = [ + r["pytorch_loop"]["median_ms"] / r["hip_fp32_input"]["median_ms"] + for r in component + if style.keep_tokens(r["tokens"]) + ] + comp_mem = [ + r["pytorch_loop_peak_mib"] / max(r["hip_peak_mib"], 1e-6) + for r in component + if style.keep_tokens(r["tokens"]) + ] + add( + f"- {prefix}The `{_tile_symbol}` kernel alone is " + f"{_range_text(comp_speed, '{:.1f}')}x faster than the PyTorch tile loop and " + f"allocates {_range_text(comp_mem, '{:.0f}')}x less transient memory (it writes " + f"only the `[tokens, {NUM_TILES}]` FP32 partials)." + ) + for kernel_path in WS2_KERNEL_PATHS: + if kernel_path not in style.paths or "ws2-reference" not in style.paths: + continue + mism = [ + c.get("mismatch_vs_reference") + for c in single + if c["path"] == kernel_path and c.get("mismatch_vs_reference") is not None + ] + relr = [c.get("rel_l2_vs_reference", 0.0) for c in single if c["path"] == kernel_path] + if not mism: + continue + add( + f"- {prefix}`{style.name(kernel_path)}` vs " + f"`{style.name('ws2-reference')}`: tile maxima are " + "bitwise equal; sumexp partials differ only by FP32 summation order, so final " + "outputs differ in " + f"{_range_text([float(m) for m in mism], '{:.0f}')} elements per case with " + f"relative-L2 {_range_text(relr, '{:.1e}')}. Both paths are equally close to FP64." + ) + ws2 = [c for c in single if c["path"].startswith("ws2")] + if ws2: + add( + f"- {prefix}Repeat bitwise: " + f"{_yes(all(c['repeat_bitwise'] for c in ws2))}; batch-invariant: " + f"{_yes(all(c['batch_invariant'] for c in ws2))}; all gradients finite: " + f"{_yes(all(c['grad_finite'] for c in ws2))}." + ) + ws2_dist = [d for d in distributed if d["path"].startswith("ws2")] + if ws2_dist: + add( + f"- {prefix}Distributed: TP-replicated and repeat bitwise on every topology: " + f"{_yes(all(d['tp_replicated'] and d['repeat_bitwise'] for d in ws2_dist))}." + ) + add("") + + single, distributed, component = all_single, all_distributed, all_component + + # ---- single GPU tables + for dtype_name in SINGLE_DTYPES: + rows = [c for c in single if c["dtype"] == dtype_name and style.keep_tokens(c["tokens"])] + if not rows: + continue + add(f"## Single-GPU logprob ({dtype_name.upper()} logits, V={REAL_VOCAB})") + add("") + add("### Forward") + add("") + add( + f"| Tokens | {plat_head}Path | Median (ms) | p95 (ms) | Speedup vs {base_name} | " + "Peak MiB | logp max-abs vs FP64 | LSE max-abs vs FP64 | Repeat | Batch-inv |" + ) + add("|---:|" + plat_div + "---|---:|---:|---:|---:|---:|---:|:---:|:---:|") + for tokens in sorted({c["tokens"] for c in rows}): + for label, _ in platforms: + # Speedups are always against that platform's own baseline. + ref = _lookup(single, dtype=dtype_name, tokens=tokens, path=base, platform=label) + for path in style.paths: + case = _lookup(rows, tokens=tokens, path=path, platform=label) + if case is None: + continue + add( + f"| {tokens} | {plat_cell(label)}{style.name(path)} | " + f"{case['forward']['median_ms']:.4f} | " + f"{case['forward']['p95_ms']:.4f} | " + f"{_median_ratio(ref, case, 'forward')} | " + f"{case['forward_peak_mib']:.1f} | " + f"{case['logp_vs_fp64']['max_abs']:.3e} | " + f"{case['lse_vs_fp64']['max_abs']:.3e} | " + f"{_yes(case['repeat_bitwise'])} | " + f"{_yes(case['batch_invariant'])} |" + ) + add("") + add("### Forward+backward") + add("") + add( + f"| Tokens | {plat_head}Path | Median (ms) | p95 (ms) | Speedup vs {base_name} | " + f"Peak MiB | Memory vs {base_name} | Grad finite |" + ) + add("|---:|" + plat_div + "---|---:|---:|---:|---:|---:|:---:|") + for tokens in sorted({c["tokens"] for c in rows}): + for label, _ in platforms: + ref = _lookup(single, dtype=dtype_name, tokens=tokens, path=base, platform=label) + for path in style.paths: + case = _lookup(rows, tokens=tokens, path=path, platform=label) + if case is None: + continue + add( + f"| {tokens} | {plat_cell(label)}{style.name(path)} | " + f"{case['train_fwd_bwd']['median_ms']:.4f} | " + f"{case['train_fwd_bwd']['p95_ms']:.4f} | " + f"{_median_ratio(ref, case, 'train_fwd_bwd')} | " + f"{case['train_peak_mib']:.1f} | " + f"{_field_ratio(case, ref, 'train_peak_mib')} | " + f"{_yes(case['grad_finite'])} |" + ) + add("") + kernel_paths = [p for p in WS2_KERNEL_PATHS if p in style.paths] + if kernel_paths and "ws2-reference" in style.paths: + add(f"### Numerics versus `{style.name('ws2-reference')}`") + add("") + add(f"| Tokens | {plat_head}Path | Mismatched elements (logp+LSE) | Relative L2 |") + add("|---:|" + plat_div + "---|---:|---:|") + for tokens in sorted({c["tokens"] for c in rows}): + for label, _ in platforms: + for kernel_path in kernel_paths: + case = _lookup(rows, tokens=tokens, path=kernel_path, platform=label) + if case is None: + continue + add( + f"| {tokens} | {plat_cell(label)}{style.name(kernel_path)} | " + f"{case.get('mismatch_vs_reference', 'n/a')} | " + f"{case.get('rel_l2_vs_reference', float('nan')):.3e} |" + ) + add("") + + overhead_rows: list[tuple[str, str, dict[str, Any]]] = [] + for label, pl in platforms: + for row in pl["single_gpu"].get("validate_overhead", []): + if not style.keep_tokens(row["tokens"]): + continue + row_path = row.get("path", "ws2-rocm") + if row_path in style.paths: + overhead_rows.append((label, row_path, row)) + if overhead_rows: + measured = sorted({style.name(path) for _, path, _ in overhead_rows}) + add(f"### `validate=True` production entry point ({', '.join(measured)}, BF16)") + add("") + add(f"| Tokens | {plat_head}Path | validate=False (ms) | validate=True (ms) | Overhead |") + add("|---:|" + plat_div + "---|---:|---:|---:|") + for label, row_path, row in overhead_rows: + overhead_ratio = _fmt_ratio( + row["validate_true"]["median_ms"], row["validate_false"]["median_ms"] + ) + add( + f"| {row['tokens']} | {plat_cell(label)}{style.name(row_path)} | " + f"{row['validate_false']['median_ms']:.4f} | " + f"{row['validate_true']['median_ms']:.4f} | {overhead_ratio} |" + ) + add("") + add( + "`validate=True` adds host-side target-range checks and a non-finite LSE check that " + "synchronizes the stream; the cost is a fixed per-call overhead." + ) + add("") + + # ---- tile-stats component + plat_order = {label: index for index, (label, _) in enumerate(platforms)} + comp_rows = [r for r in component if style.keep_tokens(r["tokens"])] + comp_rows.sort(key=lambda r: (r["dtype"], r["tokens"], plat_order.get(r["platform"], 99))) + if comp_rows: + add("## Tile-stats kernel") + add("") + + def _symbol(row: dict[str, Any]) -> str: + # results.json written before kernel_symbol existed: infer from the platform. + hip_build = dict(platforms)[row["platform"]]["environment"].get("hip") not in ( + None, + "None", + ) + default = ( + "hip_deterministic_logp_tile_stats" + if hip_build + else "deterministic_logp_tile_stats" + ) + return row.get("kernel_symbol", default) + + symbols = sorted({_symbol(row) for row in comp_rows}) + add( + f"{', '.join(f'`{s}`' for s in symbols)} computes the per-row, per-tile FP32 " + "`(max, sumexp)` partials that the operator all-gathers and merges; the PyTorch tile " + f"loop is what `{style.name('ws2-reference')}` uses for the same step. Tile maxima are " + "bitwise equal; sums differ only by FP32 summation order." + ) + add("") + add( + f"| Logits dtype | Tokens | {plat_head}Kernel | PyTorch tile loop (ms) | " + "Kernel on FP32 (ms) | Kernel on stored dtype (ms) | Speedup | Loop peak MiB | " + "Kernel peak MiB | Max bitwise | sumexp max rel | Repeat |" + ) + add("|---|---:|" + plat_div + "---|---:|---:|---:|---:|---:|---:|:---:|---:|:---:|") + for row in comp_rows: + kernel_speedup = _fmt_ratio( + row["pytorch_loop"]["median_ms"], row["hip_fp32_input"]["median_ms"] + ) + add( + f"| {row['dtype']} | {row['tokens']} | {plat_cell(row['platform'])}" + f"`{_symbol(row)}` | {row['pytorch_loop']['median_ms']:.4f} | " + f"{row['hip_fp32_input']['median_ms']:.4f} | " + f"{row['hip_native_dtype_input']['median_ms']:.4f} | " + f"{kernel_speedup} | " + f"{row['pytorch_loop_peak_mib']:.1f} | {row['hip_peak_mib']:.1f} | " + f"{_yes(row['max_bitwise'])} | {row['sumexp_max_rel']:.2e} | " + f"{_yes(row['repeat_bitwise'])} |" + ) + add("") + + # ---- distributed + topo_order = {name: index for index, (name, _, _) in enumerate(TOPOLOGIES)} + path_order = {path: index for index, path in enumerate(style.paths)} + dist_rows = [d for d in distributed if style.keep_tokens(d["tokens"])] + dist_rows.sort( + key=lambda d: ( + topo_order.get(d["topology"], 99), + d["tokens"], + plat_order.get(d["platform"], 99), + path_order.get(d["path"], 99), + ) + ) + if dist_rows: + collectives = sorted( + { + "RCCL" if pl["environment"].get("hip") not in (None, "None") else "NCCL" + for label, pl in platforms + if pl.get("distributed") + } + ) + add(f"## Distributed vocab-parallel logprob (BF16, {'/'.join(collectives)})") + add("") + absent = [ + style.name(path) for path in style.paths if path not in {d["path"] for d in distributed} + ] + if absent: + add( + "Only the vocab-parallel operators take part here. " + + ", ".join(f"`{name}`" for name in absent) + + " is a single-shard op that consumes the full `[tokens, V]` logits on one GPU; " + "it has no TP implementation (no vocab shard input, TP group, or partial merge), " + "so there is no comparable distributed row for it." + ) + add("") + add("### Forward") + add("") + add( + f"| Topology | Tokens | {plat_head}Path | Median (ms) | p95 (ms) | " + f"Speedup vs {base_name} | Peak MiB/rank | logp max-abs vs FP64 | TP-replicated | " + "Repeat |" + ) + add("|---|---:|" + plat_div + "---|---:|---:|---:|---:|---:|:---:|:---:|") + for d in dist_rows: + ref = _lookup( + distributed, + topology=d["topology"], + tokens=d["tokens"], + path=base, + platform=d["platform"], + ) + add( + f"| {d['topology']} | {d['tokens']} | {plat_cell(d['platform'])}" + f"{style.name(d['path'])} | " + f"{d['forward']['median_ms']:.4f} | {d['forward']['p95_ms']:.4f} | " + f"{_median_ratio(ref, d, 'forward')} | {d['forward_peak_mib']:.1f} | " + f"{d['logp_vs_fp64_max_abs']:.3e} | {_yes(d['tp_replicated'])} | " + f"{_yes(d['repeat_bitwise'])} |" + ) + add("") + add("### Forward+backward") + add("") + add( + f"| Topology | Tokens | {plat_head}Path | Median (ms) | p95 (ms) | " + f"Speedup vs {base_name} | Peak MiB/rank | Memory vs {base_name} | Grad finite |" + ) + add("|---|---:|" + plat_div + "---|---:|---:|---:|---:|---:|:---:|") + for d in dist_rows: + ref = _lookup( + distributed, + topology=d["topology"], + tokens=d["tokens"], + path=base, + platform=d["platform"], + ) + add( + f"| {d['topology']} | {d['tokens']} | {plat_cell(d['platform'])}" + f"{style.name(d['path'])} | " + f"{d['train_fwd_bwd']['median_ms']:.4f} | {d['train_fwd_bwd']['p95_ms']:.4f} | " + f"{_median_ratio(ref, d, 'train_fwd_bwd')} | {d['train_peak_mib']:.1f} | " + f"{_field_ratio(d, ref, 'train_peak_mib')} | {_yes(d['grad_finite'])} |" + ) + add("") + kernel_paths = [p for p in WS2_KERNEL_PATHS if p in style.paths] + if kernel_paths and "ws2-reference" in style.paths: + add(f"### Numerics versus `{style.name('ws2-reference')}` (distributed)") + add("") + add( + f"| Topology | Tokens | {plat_head}Path | Mismatched elements (logp+LSE) | " + "Relative L2 |" + ) + add("|---|---:|" + plat_div + "---|---:|---:|") + for d in dist_rows: + if d["path"] not in kernel_paths: + continue + add( + f"| {d['topology']} | {d['tokens']} | {plat_cell(d['platform'])}" + f"{style.name(d['path'])} | " + f"{d.get('mismatch_vs_reference', 'n/a')} | " + f"{d.get('rel_l2_vs_reference', float('nan')):.3e} |" + ) + add("") + + # ---- optional tuning history + baseline = payload.get("baseline") + if baseline and style.show_tuning_baseline and "ws2-rocm" in style.paths: + name = style.name("ws2-rocm") + add("## ROCm tuning: before versus after") + add("") + add( + f"Baseline commit `{baseline.get('git_commit')}` ran the PR's first ROCm backend: the " + "HIP tile kernel on an FP32 copy of the shard, with the shared PyTorch autograd " + f"backward. `{name}` rows only." + ) + add("") + add( + "| dtype | Tokens | Fwd before (ms) | Fwd after (ms) | Speedup | " + "Fwd+bwd before (ms) | Fwd+bwd after (ms) | Speedup | Peak before (MiB) | " + "Peak after (MiB) | Memory ratio |" + ) + add("|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|") + for before in baseline["single_gpu"]: + if not style.keep_tokens(before["tokens"]): + continue + after = _lookup(single, dtype=before["dtype"], tokens=before["tokens"], path="ws2-rocm") + if after is None: + continue + add( + f"| {before['dtype']} | {before['tokens']} | " + f"{before['forward']['median_ms']:.4f} | " + f"{after['forward']['median_ms']:.4f} | " + f"{_median_ratio(before, after, 'forward')} | " + f"{before['train_fwd_bwd']['median_ms']:.4f} | " + f"{after['train_fwd_bwd']['median_ms']:.4f} | " + f"{_median_ratio(before, after, 'train_fwd_bwd')} | " + f"{before['train_peak_mib']:.1f} | {after['train_peak_mib']:.1f} | " + f"{_field_ratio(after, before, 'train_peak_mib')} |" + ) + add("") + + add("## Figures") + add("") + if multi: + add( + "One line per backend and device across the full token sweep. The grid puts " + "latency and peak memory, forward and forward+backward, on one page." + ) + add("") + add("![Single-device latency and memory grid](single_gpu_grid.png)") + add("") + if multi: + add( + "The host and reference paths span three orders of magnitude, which flattens the " + "kernel backends against each other. The second grid drops them and re-scales to " + "the kernel backends alone, where the differences between Triton and the two " + "vendor kernels are legible." + ) + add("") + add("![Kernel backends only](single_gpu_grid_kernels.png)") + add("") + add("![Single-device latency](single_gpu_latency.png)") + add("") + add("![Single-device peak memory](single_gpu_memory.png)") + add("") + if distributed: + add("![Distributed latency](distributed_logp_latency.png)") + add("") + (output_directory / "report.md").write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def _platform_kind(payload: dict[str, Any]) -> str: + env = payload.get("environment", {}) + if env.get("device") == "cpu" or str(env.get("gpu", "")).startswith("n/a"): + return "cpu" + return "rocm" if env.get("hip") not in (None, "None") else "cuda" + + +# Fixed colours so a series looks the same in every figure of the set. +_SERIES_STYLE = { + "cpu": {"color": "#7f7f7f", "marker": "v", "linestyle": ":"}, + "native-torch": {"color": "#000000", "marker": "o", "linestyle": "-"}, + "triton-cuda": {"color": "#1f77b4", "marker": "s", "linestyle": "--"}, + "triton-rocm": {"color": "#d62728", "marker": "^", "linestyle": "--"}, + "cuda": {"color": "#2ca02c", "marker": "D", "linestyle": "-"}, + "hip": {"color": "#ff7f0e", "marker": "P", "linestyle": "-"}, +} +_SERIES_ORDER = ("cpu", "native-torch", "triton-cuda", "triton-rocm", "cuda", "hip") + + +def _figure_series( + platforms: list[tuple[str, dict[str, Any]]], +) -> list[tuple[str, str, str]]: + """``(display label, platform label, path)`` for one line per backend/device. + + The PyTorch reference is drawn once from the primary accelerator as + ``native-torch``; the host run contributes the ``cpu`` line. Each + accelerator adds its Triton line and its vendor-kernel line, so a ROCm + + CUDA + host set yields the six series in ``_SERIES_ORDER``. + """ + + kinds = {label: _platform_kind(pl) for label, pl in platforms} + measured = { + label: {case["path"] for case in pl["single_gpu"]["cases"]} for label, pl in platforms + } + series: list[tuple[str, str, str]] = [] + + for label, kind in kinds.items(): + if kind == "cpu" and "ws2-reference" in measured[label]: + series.append(("cpu", label, "ws2-reference")) + break + + # The reference line is drawn once, from the primary accelerator. On a + # host-only run the "cpu" entry above already is that line, so skip it. + primary = next((label for label, kind in kinds.items() if kind != "cpu"), None) + if primary is not None and "ws2-reference" in measured.get(primary, set()): + series.append(("native-torch", primary, "ws2-reference")) + + for label, kind in kinds.items(): + if kind == "cpu": + continue + if "ws2-triton" in measured[label]: + series.append((f"triton-{kind}", label, "ws2-triton")) + vendor = "ws2-rocm" if kind == "rocm" else "ws2-cuda" + if vendor in measured[label]: + series.append(("hip" if kind == "rocm" else "cuda", label, vendor)) + + rank = {name: index for index, name in enumerate(_SERIES_ORDER)} + series.sort(key=lambda item: rank.get(item[0], len(rank))) + return series + + +def _series_value(case: dict[str, Any] | None, key: str) -> float: + if case is None: + return float("nan") + value = case[key] + return float(value["median_ms"] if isinstance(value, dict) else value) + + +def _write_figures( + payload: dict[str, Any], + output_directory: Path, + style: ReportStyle, + platforms: list[tuple[str, dict[str, Any]]] | None = None, +) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + if platforms is None: + platforms = _report_platforms(payload, [], None) + series = _figure_series(platforms) + if not series: + return + by_label = dict(platforms) + + single: dict[str, list[dict[str, Any]]] = { + label: [c for c in pl["single_gpu"]["cases"] if c["dtype"] == "bf16"] + for label, pl in platforms + } + tokens = sorted({c["tokens"] for cases in single.values() for c in cases}) + + panels = ( + ("forward", "Forward latency", "median ms", True), + ("train_fwd_bwd", "Forward+backward latency", "median ms", True), + ("forward_peak_mib", "Forward peak memory", "peak MiB above live", False), + ("train_peak_mib", "Forward+backward peak memory", "peak MiB above live", False), + ) + + def draw(axis, key, title, ylabel, log_y, chosen_series=None) -> None: + for index, (name, platform_label, path) in enumerate(chosen_series or series): + ys = [ + _series_value(_lookup(single[platform_label], tokens=t, path=path), key) + for t in tokens + ] + # Series routinely coincide (cuda tracks the reference's memory; triton + # and hip share it). Draw later ones thinner so the overlap stays legible. + axis.plot( + tokens, + ys, + label=name, + linewidth=3.2 - 0.4 * index, + markersize=7 - 0.5 * index, + zorder=3 + index, + **_SERIES_STYLE.get(name, {}), + ) + axis.set_xscale("log", base=2) + if log_y: + axis.set_yscale("log") + else: + # symlog keeps the host run's legitimate ~0 MiB readings on the axis; + # memory is never negative, so clip the mirrored half away. + axis.set_yscale("symlog", linthresh=1.0) + axis.set_ylim(bottom=0) + axis.set_xlabel("tokens") + axis.set_ylabel(ylabel) + axis.set_title(f"BF16: {title}", fontsize=11) + axis.grid(True, which="both", alpha=0.3) + axis.legend(fontsize=8) + + # Latency and memory keep their own files; the grid puts all four panels together. + for filename, chosen in ( + ("single_gpu_latency.png", panels[:2]), + ("single_gpu_memory.png", panels[2:]), + ): + figure, axes = plt.subplots(1, 2, figsize=(12, 4.5)) + for axis, (key, title, ylabel, log_y) in zip(axes, chosen): + draw(axis, key, title, ylabel, log_y) + figure.tight_layout() + figure.savefig(output_directory / filename, dpi=180) + plt.close(figure) + + def grid(filename: str, chosen_series, subtitle: str) -> None: + figure, axes = plt.subplots(2, 2, figsize=(13, 9)) + for axis, (key, title, ylabel, log_y) in zip(axes.flat, panels): + draw(axis, key, title, ylabel, log_y, chosen_series) + figure.suptitle( + f"Single-device vocab-parallel logprob, BF16, V={REAL_VOCAB} — {subtitle} (" + + ", ".join(name for name, _, _ in chosen_series) + + ")", + fontsize=12, + ) + figure.tight_layout(rect=(0, 0, 1, 0.96)) + figure.savefig(output_directory / filename, dpi=180) + plt.close(figure) + + grid("single_gpu_grid.png", series, "all paths") + + # The host and reference lines span three orders of magnitude, which flattens + # the kernel backends against each other. Re-draw them on their own scale. + kernel_series = [item for item in series if item[0] not in ("cpu", "native-torch")] + if len(kernel_series) >= 2 and len(kernel_series) != len(series): + grid("single_gpu_grid_kernels.png", kernel_series, "kernel backends only") + + # ---- distributed: same series minus the host run + dist_series = [item for item in series if item[0] != "cpu"] + distributed = { + label: pl.get("distributed", []) or [] for label, pl in platforms if label in by_label + } + dist_series = [item for item in dist_series if distributed.get(item[1])] + if not dist_series: + return + + cells: list[tuple[str, int]] = [] + for _, platform_label, _ in dist_series: + for row in distributed[platform_label]: + key = (row["topology"], row["tokens"]) + if key not in cells: + cells.append(key) + topo_rank = {name: index for index, (name, _, _) in enumerate(TOPOLOGIES)} + cells.sort(key=lambda c: (topo_rank.get(c[0], 99), c[1])) + labels = [f"{topology}\nM={tokens_}" for topology, tokens_ in cells] + + figure, axes = plt.subplots(1, 2, figsize=(max(12, 1.15 * len(labels)), 5.0)) + xs = list(range(len(labels))) + width = 0.8 / max(len(dist_series), 1) + for axis, key, direction in zip( + axes, ("forward", "train_fwd_bwd"), ("Forward", "Forward+backward") + ): + for index, (name, platform_label, path) in enumerate(dist_series): + values = [ + _series_value( + _lookup(distributed[platform_label], path=path, topology=topology, tokens=t), + key, + ) + for topology, t in cells + ] + offset = (index - (len(dist_series) - 1) / 2) * width + axis.bar( + [x + offset for x in xs], + values, + width, + label=name, + color=_SERIES_STYLE.get(name, {}).get("color"), + zorder=3, + ) + axis.set_xticks(xs) + axis.set_xticklabels(labels, fontsize=8) + axis.set_ylabel("slowest-rank median ms") + axis.set_title(f"Distributed vocab-parallel logprob, BF16: {direction}") + axis.grid(True, axis="y", alpha=0.3) + axis.legend(fontsize=8) + figure.tight_layout() + figure.savefig(output_directory / "distributed_logp_latency.png", dpi=180) + plt.close(figure) + + +def _extension_symbols() -> str: + if not _EXT_AVAILABLE or _C is None: + return "none (pure-Python fallback)" + candidates = ( + "deterministic_logp_tile_stats", + "hip_deterministic_logp_tile_stats", + "hip_deterministic_logp_backward", + ) + present = [name for name in candidates if hasattr(_C, name)] + return ", ".join(present) if present else "none" + + +def _environment(device: torch.device) -> dict[str, Any]: + environment: dict[str, Any] = { + "device": device.type, + "torch": torch.__version__, + "cuda": torch.version.cuda, + "hip": torch.version.hip, + "python": os.sys.version.split()[0], + "git_commit": os.popen("git rev-parse HEAD").read().strip(), + "extension_symbols": _extension_symbols(), + } + if device.type == "cuda": + properties = torch.cuda.get_device_properties(0) + environment.update( + { + "gpu": torch.cuda.get_device_name(0), + "gpu_count": torch.cuda.device_count(), + "architecture": ( + getattr(properties, "gcnArchName", "unknown") + if torch.version.hip is not None + else f"sm_{properties.major}{properties.minor}" + ), + "native_collective": ( + "torch.distributed ProcessGroupNCCL" + + (" (RCCL on ROCm)" if torch.version.hip is not None else " (NCCL)") + ), + } + ) + else: + environment.update( + { + "gpu": "n/a (host execution)", + "gpu_count": 0, + "architecture": platform.processor() or platform.machine(), + "cpu_count": os.cpu_count(), + "torch_threads": torch.get_num_threads(), + "native_collective": "n/a (single-process host run)", + } + ) + return environment + + +def _validate_environment(require_distributed: bool, device: torch.device) -> None: + if device.type == "cpu": + if require_distributed: + raise RuntimeError( + "--device cpu cannot run the distributed section; add --skip-distributed" + ) + return + if not torch.cuda.is_available(): + raise RuntimeError("no CUDA/ROCm GPU is visible") + if torch.version.hip is not None: + if not _EXT_AVAILABLE or _C is None or not hasattr(_C, "hip_deterministic_logp_backward"): + raise RuntimeError( + "rl_engine._C with hip_deterministic_logp_* is unavailable; build with " + "PYTORCH_ROCM_ARCH=gfx942 RL_KERNEL_REQUIRE_EXT=1 " + "python setup.py build_ext --inplace" + ) + elif not native_tile_stats_available(): + # Not fatal: ws2-triton and the reference still run, only ws2-cuda drops out. + print( + "warning: rl_engine._C with deterministic_logp_tile_stats is unavailable; " + "the ws2-cuda path will be skipped. Build with " + "TORCH_CUDA_ARCH_LIST=9.0 RL_KERNEL_REQUIRE_EXT=1 " + "python setup.py build_ext --inplace" + ) + if require_distributed and (not dist.is_available() or not dist.is_nccl_available()): + raise RuntimeError("PyTorch NCCL/ProcessGroupNCCL support is unavailable") + + +ALL_PATHS = ( + "native", + "ws1-pytorch", + "ws1-triton", + "ws2-reference", + "ws2-triton", + "ws2-cuda", + "ws2-rocm", +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--output-dir", type=Path, default=Path("benchmarks/results/rocm_logp")) + parser.add_argument( + "--device", + choices=("cuda", "cpu"), + default="cuda", + help="device for the single-device section; 'cpu' implies --skip-distributed", + ) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--samples", type=int, default=10) + parser.add_argument("--training-samples", type=int, default=5) + parser.add_argument("--skip-distributed", action="store_true") + parser.add_argument("--skip-single", action="store_true") + parser.add_argument( + "--baseline", + type=Path, + default=None, + help="earlier results.json; adds a before/after tuning section for the ROCm backend", + ) + parser.add_argument( + "--render-from", + type=Path, + default=None, + help="skip measurement and render report/figures from this results.json", + ) + parser.add_argument( + "--report-paths", + type=str, + default=",".join(ALL_PATHS), + help="comma-separated measured paths to show, in order (default: all)", + ) + parser.add_argument( + "--rename", + type=str, + default="", + help="display names, e.g. 'ws2-reference=native,ws2-rocm=strict-hip'", + ) + parser.add_argument( + "--report-baseline", + type=str, + default=None, + help="path used for speedup/memory ratio columns (default: first reported path)", + ) + parser.add_argument( + "--table-tokens", + type=str, + default="", + help="comma-separated token counts to show in tables (default: all measured)", + ) + parser.add_argument( + "--topologies", + type=str, + default=",".join(name for name, _, _ in TOPOLOGIES), + help="comma-separated subset of " + ",".join(name for name, _, _ in TOPOLOGIES), + ) + parser.add_argument( + "--compare-with", + action="append", + default=[], + metavar="LABEL=PATH", + help=( + "results.json from another platform's run of this harness; adds a cross-platform " + "comparison section to the report. Repeat for each platform, e.g. " + "--compare-with h100=benchmarks/results/pr328_cuda_h100/results.json" + ), + ) + parser.add_argument( + "--compare-label", + type=str, + default=None, + help="label for this run inside the comparison section (default: derived from the GPU)", + ) + parser.add_argument("--tokens", type=str, default=",".join(str(t) for t in SINGLE_TOKENS)) + parser.add_argument( + "--distributed-tokens", type=str, default=",".join(str(t) for t in DISTRIBUTED_TOKENS) + ) + return parser.parse_args() + + +def _report_style( + args: argparse.Namespace, output_directory: Path, config: dict[str, Any] +) -> ReportStyle: + paths = tuple(p.strip() for p in args.report_paths.split(",") if p.strip()) + unknown = [p for p in paths if p not in ALL_PATHS] + if unknown: + raise ValueError(f"unknown report paths {unknown}; choose from {ALL_PATHS}") + names: dict[str, str] = {} + for item in (piece.strip() for piece in args.rename.split(",") if piece.strip()): + key, _, value = item.partition("=") + names[key.strip()] = value.strip() + table_tokens = tuple(int(t) for t in args.table_tokens.split(",") if t.strip()) or None + # Always show the measurement command: rerunning it with the same report flags + # regenerates this report; --render-from only re-renders an existing results.json. + command_parts = [ + "python benchmarks/benchmark_rocm_logp.py \\", + f" --warmup {config['warmup']} \\", + f" --samples {config['samples']} \\", + f" --training-samples {config['training_samples']} \\", + ] + if paths != ALL_PATHS: + command_parts.append(f" --report-paths {','.join(paths)} \\") + if names: + command_parts.append(" --rename " + ",".join(f"{k}={v}" for k, v in names.items()) + " \\") + if args.report_baseline: + command_parts.append(f" --report-baseline {args.report_baseline} \\") + if table_tokens: + command_parts.append(f" --table-tokens {','.join(str(t) for t in table_tokens)} \\") + for spec in getattr(args, "compare_with", []) or []: + command_parts.append(f" --compare-with {spec} \\") + if getattr(args, "compare_label", None): + command_parts.append(f" --compare-label {args.compare_label} \\") + command_parts.append(f" --output-dir {output_directory.as_posix()}") + return ReportStyle( + paths=paths, + names=names, + baseline=args.report_baseline or paths[0], + table_tokens=table_tokens, + show_tuning_baseline=args.baseline is not None, + command="\n".join(command_parts), + ) + + +def main() -> None: + args = parse_args() + output_directory: Path = args.output_dir + output_directory.mkdir(parents=True, exist_ok=True) + + if args.render_from is not None: + payload = json.loads(args.render_from.read_text(encoding="utf-8")) + else: + device = torch.device("cuda", 0) if args.device == "cuda" else torch.device("cpu") + if device.type == "cpu": + # Triton, the native kernels, and NCCL are all device-only. + args.skip_distributed = True + _validate_environment(require_distributed=not args.skip_distributed, device=device) + config = { + "warmup": args.warmup, + "samples": args.samples, + "training_samples": args.training_samples, + } + tokens = tuple(int(t) for t in args.tokens.split(",") if t) + distributed_tokens = tuple(int(t) for t in args.distributed_tokens.split(",") if t) + selected = {name.strip() for name in args.topologies.split(",") if name.strip()} + payload = { + "environment": _environment(device), + "config": config, + "single_gpu": {"cases": [], "validate_overhead": [], "paths": []}, + "tile_stats_component": [], + "distributed": [], + } + if not args.skip_single: + payload["single_gpu"] = _single_gpu_benchmarks( + warmup=args.warmup, + samples=args.samples, + training_samples=args.training_samples, + tokens=tokens, + device=device, + ) + payload["tile_stats_component"] = _tile_stats_component( + warmup=args.warmup, samples=args.samples, tokens=tokens, device=device + ) + if not args.skip_distributed: + device_count = torch.cuda.device_count() + for topology in TOPOLOGIES: + name, tp, cp = topology + if name not in selected: + continue + if tp * cp > device_count: + print(f"skipping {name}: needs {tp * cp} GPUs, {device_count} visible") + continue + payload["distributed"].extend( + _run_distributed_world(topology, distributed_tokens, config) + ) + style = _report_style(args, output_directory, payload["config"]) + if args.baseline is not None: + baseline = json.loads(args.baseline.read_text(encoding="utf-8")) + payload["baseline"] = { + "git_commit": baseline.get("environment", {}).get("git_commit"), + "single_gpu": [c for c in baseline["single_gpu"]["cases"] if c["path"] == "ws2-rocm"], + "tile_stats_component": baseline.get("tile_stats_component", []), + "distributed": [d for d in baseline.get("distributed", []) if d["path"] == "ws2-rocm"], + } + elif args.render_from is not None: + payload.pop("baseline", None) + (output_directory / "results.json").write_text( + json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8" + ) + platforms = _report_platforms(payload, args.compare_with, args.compare_label) + _write_report(payload, output_directory, style, platforms) + _write_figures(payload, output_directory, style, platforms) + print(json.dumps({"output_dir": str(output_directory), "status": "ok"})) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/results/pr328_rocm_mi300x/distributed_logp_latency.png b/benchmarks/results/pr328_rocm_mi300x/distributed_logp_latency.png new file mode 100644 index 00000000..9263fa28 Binary files /dev/null and b/benchmarks/results/pr328_rocm_mi300x/distributed_logp_latency.png differ diff --git a/benchmarks/results/pr328_rocm_mi300x/report.md b/benchmarks/results/pr328_rocm_mi300x/report.md new file mode 100644 index 00000000..e281083e --- /dev/null +++ b/benchmarks/results/pr328_rocm_mi300x/report.md @@ -0,0 +1,297 @@ +# PR #328 vocab-parallel logprob performance analysis + +> Operator-only benchmark. No model checkpoint or serving engine was used. + +> Every platform below ran this same harness, so the seeded logits, the `V=151936` / 64-tile split, and the FP64 oracle are identical and only the device and backend differ. Backends are not available everywhere: `ws2-rocm` is ROCm-only, `ws2-cuda` is CUDA-only, `ws2-triton` compiles from one source on both, and the PyTorch paths are the only ones that also run on the host. A missing row means the backend cannot exist on that platform, not that it failed. + +## Environment + +| Item | mi300x | h100 | cpu | +|---|---|---|---| +| architecture | gfx942:sramecc+:xnack- | sm_90 | x86_64 | +| cpu_count | n/a | n/a | 192 | +| cuda | n/a | 13.0 | 13.0 | +| device | n/a | cuda | cpu | +| extension_symbols | hip_deterministic_logp_tile_stats, hip_deterministic_logp_backward | deterministic_logp_tile_stats | deterministic_logp_tile_stats | +| git_commit | e9f1d2a5c67a283bd987978614b4444a9416c1f0 | dd5fe05be2252e04b0308149b46ac783bd68e42a | dd5fe05be2252e04b0308149b46ac783bd68e42a | +| gpu | AMD Instinct MI300X | NVIDIA H100 80GB HBM3 | n/a (host execution) | +| gpu_count | 8 | 8 | 0 | +| hip | 7.14.60850 | None | None | +| native_collective | torch.distributed ProcessGroupNCCL (RCCL on ROCm) | torch.distributed ProcessGroupNCCL (NCCL) | n/a (single-process host run) | +| python | 3.12.3 | 3.11.15 | 3.11.15 | +| torch | 2.12.0+rocm7.14.0a20260608 | 2.13.0+cu130 | 2.13.0+cu130 | +| torch_threads | n/a | n/a | 96 | + +## Methodology + +- Qwen3 vocabulary `V=151936` split into 64 tiles of 2374 columns; seeded logits (`randn * 2.0`), random targets, every seventh token inactive. +- Measured paths: + - `native`: `pytorch-vocab-parallel-logp-ws2`, the WS2 vocab-parallel reference operator: a PyTorch tile loop for the per-tile FP32 `(max, sumexp)` partials, all-gather of the partials, fixed global tile-order merge, and a PyTorch autograd backward. + - `triton`: `triton-vocab-parallel-logp-ws2`, the same contract, transport, and merge with two Triton kernels (tile statistics read from the stored shard, fused backward); one source for CUDA and ROCm. + - `cuda`: `cuda-vocab-parallel-logp-ws2`, the same contract, transport, and merge with two CUDA kernels from `csrc/deterministic_logp_kernel.cu`: `deterministic_logp_tile_stats` reads the stored BF16/FP16/FP32 shard directly (16-byte vector loads, no FP32 copy) and `deterministic_logp_backward` produces the gradient in one fused pass. + - `hip`: `rocm-vocab-parallel-logp-ws2`, the same contract, transport, and merge with two HIP kernels: `hip_deterministic_logp_tile_stats` reads the stored BF16/FP16/FP32 shard directly (8-element vector loads, no FP32 copy) and `hip_deterministic_logp_backward` produces the gradient in one fused pass. +- Distributed: one process per GPU; the WS2 operators all-gather per-tile `(max, sumexp)` partials over NCCL/RCCL and merge them in fixed global tile order; CP ranks shard tokens and never enter the merge. +- Forward returns the selected-token logprob and the vocabulary LSE; forward+backward computes `grad_logits` for `sum(active * logp)`. The WS2 operators run with `validate=False`; the `validate=True` production entry point is measured separately. +- Single-GPU timing: GPU events, median and p95. Distributed timing: synchronized wall clock, slowest rank per sample. Peak memory is the per-call increase in `torch.cuda.max_memory_allocated` (distributed: max over ranks). +- Accuracy is against an FP64 `logsumexp` of the same (BF16-rounded) logits. Repeat = two identical calls are bitwise equal; batch-invariant = a row computed alone is bitwise equal to the same row inside the batch; TP-replicated = every TP rank holds identical bits. +- 5 warmups, 20 measured forward samples, 10 measured forward+backward samples. Raw medians, p95, minimum, maximum, and every measured path are in `results.json`. +- Tables show 2048 tokens; the figures cover the full token sweep. + +Reproduce this report from the repository root: + +```bash +python benchmarks/benchmark_rocm_logp.py \ + --warmup 5 \ + --samples 20 \ + --training-samples 10 \ + --report-paths ws2-reference,ws2-triton,ws2-cuda,ws2-rocm \ + --rename ws2-reference=native,ws2-triton=triton,ws2-cuda=cuda,ws2-rocm=hip \ + --report-baseline ws2-reference \ + --table-tokens 2048 \ + --compare-with h100=benchmarks/results/pr328_cuda_h100/results.json \ + --compare-with cpu=benchmarks/results/pr328_cpu/results.json \ + --compare-label mi300x \ + --output-dir benchmarks/results/pr328_rocm_mi300x +``` + +## Key findings + +- mi300x Single GPU: `triton` is 5.26-6.45x faster than `native` in forward and 4.29-5.71x in forward+backward, with 0.15-0.33x its peak memory. +- mi300x Distributed: `triton` is 1.62-4.28x faster than `native` in forward and 1.82-4.05x in forward+backward across 6 TP/CP topologies, at 0.15x the per-rank peak memory (absolute forward 0.824-1.115 ms). +- mi300x Single GPU: `hip` is 5.59-6.88x faster than `native` in forward and 4.47-6.07x in forward+backward, with 0.15-0.33x its peak memory. +- mi300x Distributed: `hip` is 1.70-4.67x faster than `native` in forward and 2.05-4.49x in forward+backward across 6 TP/CP topologies, at 0.15x the per-rank peak memory (absolute forward 0.755-1.063 ms). +- mi300x The `hip_deterministic_logp_tile_stats` kernel alone is 9.6-9.9x faster than the PyTorch tile loop and allocates 57x less transient memory (it writes only the `[tokens, 64]` FP32 partials). +- mi300x `triton` vs `native`: tile maxima are bitwise equal; sumexp partials differ only by FP32 summation order, so final outputs differ in 0-169 elements per case with relative-L2 0.0e+00-6.9e-08. Both paths are equally close to FP64. +- mi300x `hip` vs `native`: tile maxima are bitwise equal; sumexp partials differ only by FP32 summation order, so final outputs differ in 0-65 elements per case with relative-L2 0.0e+00-1.3e-08. Both paths are equally close to FP64. +- mi300x Repeat bitwise: yes; batch-invariant: yes; all gradients finite: yes. +- mi300x Distributed: TP-replicated and repeat bitwise on every topology: yes. +- h100 Single GPU: `triton` is 6.80-8.83x faster than `native` in forward and 5.28-8.12x in forward+backward, with 0.15-0.33x its peak memory. +- h100 Distributed: `triton` is 1.74-4.33x faster than `native` in forward and 2.13-5.02x in forward+backward across 6 TP/CP topologies, at 0.15x the per-rank peak memory (absolute forward 0.817-1.077 ms). +- h100 Single GPU: `cuda` is 6.98-9.55x faster than `native` in forward and 5.49-8.46x in forward+backward, with 0.15-0.33x its peak memory. +- h100 Distributed: `cuda` is 1.89-4.75x faster than `native` in forward and 2.41-5.48x in forward+backward across 6 TP/CP topologies, at 0.15x the per-rank peak memory (absolute forward 0.756-0.995 ms). +- h100 The `deterministic_logp_tile_stats` kernel alone is 10.5-10.6x faster than the PyTorch tile loop and allocates 57x less transient memory (it writes only the `[tokens, 64]` FP32 partials). +- h100 `triton` vs `native`: tile maxima are bitwise equal; sumexp partials differ only by FP32 summation order, so final outputs differ in 0-81 elements per case with relative-L2 0.0e+00-6.9e-08. Both paths are equally close to FP64. +- h100 `cuda` vs `native`: tile maxima are bitwise equal; sumexp partials differ only by FP32 summation order, so final outputs differ in 0-73 elements per case with relative-L2 0.0e+00-1.3e-08. Both paths are equally close to FP64. +- h100 Repeat bitwise: yes; batch-invariant: yes; all gradients finite: yes. +- h100 Distributed: TP-replicated and repeat bitwise on every topology: yes. +- cpu Repeat bitwise: yes; batch-invariant: yes; all gradients finite: yes. + +## Single-GPU logprob (BF16 logits, V=151936) + +### Forward + +| Tokens | Platform | Path | Median (ms) | p95 (ms) | Speedup vs native | Peak MiB | logp max-abs vs FP64 | LSE max-abs vs FP64 | Repeat | Batch-inv | +|---:|---|---|---:|---:|---:|---:|---:|---:|:---:|:---:| +| 2048 | mi300x | native | 6.1387 | 6.3072 | 1.00× | 1245.0 | 1.557e-06 | 1.382e-06 | yes | yes | +| 2048 | mi300x | triton | 0.9514 | 0.9801 | 6.45× | 2.0 | 1.557e-06 | 1.318e-06 | yes | yes | +| 2048 | mi300x | hip | 0.8923 | 0.9162 | 6.88× | 2.0 | 1.621e-06 | 1.318e-06 | yes | yes | +| 2048 | h100 | native | 7.0518 | 7.0676 | 1.00× | 1246.0 | 1.225e-06 | 8.285e-07 | yes | yes | +| 2048 | h100 | triton | 0.7990 | 0.8314 | 8.83× | 2.0 | 1.225e-06 | 8.285e-07 | yes | yes | +| 2048 | h100 | cuda | 0.7388 | 0.7532 | 9.55× | 2.0 | 1.225e-06 | 8.285e-07 | yes | yes | +| 2048 | cpu | native | 421.5963 | 430.5902 | 1.00× | 1277.0 | 1.245e-06 | 8.234e-07 | yes | yes | + +### Forward+backward + +| Tokens | Platform | Path | Median (ms) | p95 (ms) | Speedup vs native | Peak MiB | Memory vs native | Grad finite | +|---:|---|---|---:|---:|---:|---:|---:|:---:| +| 2048 | mi300x | native | 12.3477 | 12.4680 | 1.00× | 7715.7 | 1.00× | yes | +| 2048 | mi300x | triton | 2.1619 | 2.1945 | 5.71× | 1187.1 | 0.15× | yes | +| 2048 | mi300x | hip | 2.0359 | 2.0513 | 6.07× | 1187.1 | 0.15× | yes | +| 2048 | h100 | native | 16.6661 | 16.8009 | 1.00× | 7722.7 | 1.00× | yes | +| 2048 | h100 | triton | 2.0525 | 2.1072 | 8.12× | 1189.1 | 0.15× | yes | +| 2048 | h100 | cuda | 1.9692 | 2.0044 | 8.46× | 1188.1 | 0.15× | yes | +| 2048 | cpu | native | 1442.7101 | 1452.4850 | 1.00× | 7714.4 | 1.00× | yes | + +### Numerics versus `native` + +| Tokens | Platform | Path | Mismatched elements (logp+LSE) | Relative L2 | +|---:|---|---|---:|---:| +| 2048 | mi300x | triton | 169 | 1.713e-08 | +| 2048 | mi300x | hip | 65 | 1.098e-08 | +| 2048 | h100 | triton | 81 | 1.147e-08 | +| 2048 | h100 | cuda | 73 | 1.051e-08 | + +## Single-GPU logprob (FP32 logits, V=151936) + +### Forward + +| Tokens | Platform | Path | Median (ms) | p95 (ms) | Speedup vs native | Peak MiB | logp max-abs vs FP64 | LSE max-abs vs FP64 | Repeat | Batch-inv | +|---:|---|---|---:|---:|---:|---:|---:|---:|:---:|:---:| +| 2048 | mi300x | native | 5.5791 | 5.6527 | 1.00× | 58.0 | 1.747e-06 | 1.271e-06 | yes | yes | +| 2048 | mi300x | triton | 1.0602 | 1.0851 | 5.26× | 2.0 | 1.711e-06 | 1.197e-06 | yes | yes | +| 2048 | mi300x | hip | 0.9972 | 1.0351 | 5.59× | 2.0 | 1.747e-06 | 1.271e-06 | yes | yes | +| 2048 | h100 | native | 6.1386 | 6.3147 | 1.00× | 58.3 | 1.662e-06 | 7.926e-07 | yes | yes | +| 2048 | h100 | triton | 0.9026 | 0.9550 | 6.80× | 2.0 | 1.662e-06 | 8.059e-07 | yes | yes | +| 2048 | h100 | cuda | 0.8794 | 0.8956 | 6.98× | 2.0 | 1.662e-06 | 8.059e-07 | yes | yes | +| 2048 | cpu | native | 312.5294 | 319.2598 | 1.00× | 90.0 | 1.645e-06 | 8.017e-07 | yes | yes | + +### Forward+backward + +| Tokens | Platform | Path | Median (ms) | p95 (ms) | Speedup vs native | Peak MiB | Memory vs native | Grad finite | +|---:|---|---|---:|---:|---:|---:|---:|:---:| +| 2048 | mi300x | native | 12.1054 | 12.1808 | 1.00× | 7122.2 | 1.00× | yes | +| 2048 | mi300x | triton | 2.8249 | 2.8799 | 4.29× | 2374.1 | 0.33× | yes | +| 2048 | mi300x | hip | 2.7086 | 3.3877 | 4.47× | 2374.1 | 0.33× | yes | +| 2048 | h100 | native | 15.7429 | 16.3057 | 1.00× | 7128.2 | 1.00× | yes | +| 2048 | h100 | triton | 2.9805 | 3.0255 | 5.28× | 2376.1 | 0.33× | yes | +| 2048 | h100 | cuda | 2.8676 | 2.8790 | 5.49× | 2376.1 | 0.33× | yes | +| 2048 | cpu | native | 1327.1097 | 1336.0410 | 1.00× | 7122.0 | 1.00× | yes | + +### Numerics versus `native` + +| Tokens | Platform | Path | Mismatched elements (logp+LSE) | Relative L2 | +|---:|---|---|---:|---:| +| 2048 | mi300x | triton | 144 | 1.466e-08 | +| 2048 | mi300x | hip | 37 | 8.004e-09 | +| 2048 | h100 | triton | 72 | 1.136e-08 | +| 2048 | h100 | cuda | 57 | 1.100e-08 | + +### `validate=True` production entry point (cuda, hip, native, BF16) + +| Tokens | Platform | Path | validate=False (ms) | validate=True (ms) | Overhead | +|---:|---|---|---:|---:|---:| +| 2048 | mi300x | hip | 0.9116 | 1.0558 | 1.16× | +| 2048 | h100 | cuda | 0.7297 | 0.8614 | 1.18× | +| 2048 | cpu | native | 428.4042 | 432.0389 | 1.01× | + +`validate=True` adds host-side target-range checks and a non-finite LSE check that synchronizes the stream; the cost is a fixed per-call overhead. + +## Tile-stats kernel + +`deterministic_logp_tile_stats`, `hip_deterministic_logp_tile_stats` computes the per-row, per-tile FP32 `(max, sumexp)` partials that the operator all-gathers and merges; the PyTorch tile loop is what `native` uses for the same step. Tile maxima are bitwise equal; sums differ only by FP32 summation order. + +| Logits dtype | Tokens | Platform | Kernel | PyTorch tile loop (ms) | Kernel on FP32 (ms) | Kernel on stored dtype (ms) | Speedup | Loop peak MiB | Kernel peak MiB | Max bitwise | sumexp max rel | Repeat | +|---|---:|---|---|---:|---:|---:|---:|---:|---:|:---:|---:|:---:| +| bf16 | 2048 | mi300x | `hip_deterministic_logp_tile_stats` | 5.0990 | 0.5171 | 0.4669 | 9.86× | 56.6 | 1.0 | yes | 3.64e-07 | yes | +| bf16 | 2048 | h100 | `deterministic_logp_tile_stats` | 5.6172 | 0.5325 | 0.3713 | 10.55× | 56.6 | 1.0 | yes | 3.86e-07 | yes | +| fp32 | 2048 | mi300x | `hip_deterministic_logp_tile_stats` | 5.1025 | 0.5328 | 0.5331 | 9.58× | 56.6 | 1.0 | yes | 4.48e-07 | yes | +| fp32 | 2048 | h100 | `deterministic_logp_tile_stats` | 5.6570 | 0.5341 | 0.5346 | 10.59× | 56.6 | 1.0 | yes | 3.24e-07 | yes | + +## Distributed vocab-parallel logprob (BF16, NCCL/RCCL) + +### Forward + +| Topology | Tokens | Platform | Path | Median (ms) | p95 (ms) | Speedup vs native | Peak MiB/rank | logp max-abs vs FP64 | TP-replicated | Repeat | +|---|---:|---|---|---:|---:|---:|---:|---:|:---:|:---:| +| tp2 | 2048 | mi300x | native | 3.5877 | 3.6322 | 1.00× | 650.3 | 1.452e-06 | yes | yes | +| tp2 | 2048 | mi300x | triton | 0.8894 | 0.9138 | 4.03× | 3.0 | 1.452e-06 | yes | yes | +| tp2 | 2048 | mi300x | hip | 0.8487 | 0.8818 | 4.23× | 3.0 | 1.452e-06 | yes | yes | +| tp2 | 2048 | h100 | native | 3.9802 | 4.0415 | 1.00× | 650.3 | 1.178e-06 | yes | yes | +| tp2 | 2048 | h100 | triton | 0.9375 | 1.0272 | 4.25× | 3.0 | 1.178e-06 | yes | yes | +| tp2 | 2048 | h100 | cuda | 0.8378 | 0.8726 | 4.75× | 3.0 | 1.178e-06 | yes | yes | +| tp4 | 2048 | mi300x | native | 2.2691 | 2.3080 | 1.00× | 353.0 | 1.452e-06 | yes | yes | +| tp4 | 2048 | mi300x | triton | 0.9221 | 0.9652 | 2.46× | 2.5 | 1.452e-06 | yes | yes | +| tp4 | 2048 | mi300x | hip | 0.8714 | 0.9162 | 2.60× | 2.5 | 1.452e-06 | yes | yes | +| tp4 | 2048 | h100 | native | 2.5194 | 2.5661 | 1.00× | 353.0 | 1.178e-06 | yes | yes | +| tp4 | 2048 | h100 | triton | 0.9734 | 1.1620 | 2.59× | 2.5 | 1.178e-06 | yes | yes | +| tp4 | 2048 | h100 | cuda | 0.8695 | 0.8957 | 2.90× | 2.5 | 1.178e-06 | yes | yes | +| tp8 | 2048 | mi300x | native | 1.8066 | 1.8356 | 1.00× | 204.3 | 1.452e-06 | yes | yes | +| tp8 | 2048 | mi300x | triton | 1.1149 | 1.1961 | 1.62× | 2.3 | 1.452e-06 | yes | yes | +| tp8 | 2048 | mi300x | hip | 1.0631 | 1.1280 | 1.70× | 2.3 | 1.452e-06 | yes | yes | +| tp8 | 2048 | h100 | native | 1.8771 | 1.9184 | 1.00× | 204.3 | 1.178e-06 | yes | yes | +| tp8 | 2048 | h100 | triton | 1.0773 | 1.6099 | 1.74× | 2.3 | 1.178e-06 | yes | yes | +| tp8 | 2048 | h100 | cuda | 0.9952 | 1.0246 | 1.89× | 2.3 | 1.178e-06 | yes | yes | +| tp2_cp2 | 2048 | mi300x | native | 3.3882 | 3.5075 | 1.00× | 325.5 | 1.452e-06 | yes | yes | +| tp2_cp2 | 2048 | mi300x | triton | 0.8308 | 0.9605 | 4.08× | 1.5 | 1.452e-06 | yes | yes | +| tp2_cp2 | 2048 | mi300x | hip | 0.7809 | 0.8174 | 4.34× | 1.5 | 1.452e-06 | yes | yes | +| tp2_cp2 | 2048 | h100 | native | 3.5072 | 4.4380 | 1.00× | 325.5 | 1.178e-06 | yes | yes | +| tp2_cp2 | 2048 | h100 | triton | 0.8173 | 0.8876 | 4.29× | 1.5 | 1.178e-06 | yes | yes | +| tp2_cp2 | 2048 | h100 | cuda | 0.7557 | 0.7874 | 4.64× | 1.5 | 1.178e-06 | yes | yes | +| tp4_cp2 | 2048 | mi300x | native | 2.2382 | 2.2630 | 1.00× | 176.7 | 1.452e-06 | yes | yes | +| tp4_cp2 | 2048 | mi300x | triton | 0.8831 | 0.9622 | 2.53× | 1.3 | 1.452e-06 | yes | yes | +| tp4_cp2 | 2048 | mi300x | hip | 0.8386 | 0.8978 | 2.67× | 1.3 | 1.452e-06 | yes | yes | +| tp4_cp2 | 2048 | h100 | native | 2.3596 | 2.5023 | 1.00× | 176.7 | 1.178e-06 | yes | yes | +| tp4_cp2 | 2048 | h100 | triton | 0.8932 | 0.9222 | 2.64× | 1.3 | 1.178e-06 | yes | yes | +| tp4_cp2 | 2048 | h100 | cuda | 0.8499 | 0.9138 | 2.78× | 1.3 | 1.178e-06 | yes | yes | +| tp2_cp4 | 2048 | mi300x | native | 3.5261 | 3.6071 | 1.00× | 163.1 | 1.452e-06 | yes | yes | +| tp2_cp4 | 2048 | mi300x | triton | 0.8240 | 1.3511 | 4.28× | 0.8 | 1.452e-06 | yes | yes | +| tp2_cp4 | 2048 | mi300x | hip | 0.7547 | 0.7834 | 4.67× | 0.8 | 1.452e-06 | yes | yes | +| tp2_cp4 | 2048 | h100 | native | 3.5499 | 3.8794 | 1.00× | 163.1 | 1.178e-06 | yes | yes | +| tp2_cp4 | 2048 | h100 | triton | 0.8201 | 0.8657 | 4.33× | 0.8 | 1.178e-06 | yes | yes | +| tp2_cp4 | 2048 | h100 | cuda | 0.7561 | 1.5402 | 4.69× | 0.8 | 1.178e-06 | yes | yes | + +### Forward+backward + +| Topology | Tokens | Platform | Path | Median (ms) | p95 (ms) | Speedup vs native | Peak MiB/rank | Memory vs native | Grad finite | +|---|---:|---|---|---:|---:|---:|---:|---:|:---:| +| tp2 | 2048 | mi300x | native | 7.1482 | 7.5344 | 1.00× | 3857.9 | 1.00× | yes | +| tp2 | 2048 | mi300x | triton | 1.7667 | 1.8881 | 4.05× | 593.6 | 0.15× | yes | +| tp2 | 2048 | mi300x | hip | 1.5938 | 1.7476 | 4.49× | 593.6 | 0.15× | yes | +| tp2 | 2048 | h100 | native | 9.2129 | 9.3465 | 1.00× | 3860.9 | 1.00× | yes | +| tp2 | 2048 | h100 | triton | 1.8345 | 2.1763 | 5.02× | 594.1 | 0.15× | yes | +| tp2 | 2048 | h100 | cuda | 1.6801 | 2.0202 | 5.48× | 594.6 | 0.15× | yes | +| tp4 | 2048 | mi300x | native | 4.4688 | 4.5386 | 1.00× | 1929.0 | 1.00× | yes | +| tp4 | 2048 | mi300x | triton | 1.7425 | 2.1374 | 2.56× | 296.8 | 0.15× | yes | +| tp4 | 2048 | mi300x | hip | 1.5001 | 1.6030 | 2.98× | 296.8 | 0.15× | yes | +| tp4 | 2048 | h100 | native | 6.2461 | 6.5273 | 1.00× | 1929.0 | 1.00× | yes | +| tp4 | 2048 | h100 | triton | 2.2329 | 2.2542 | 2.80× | 296.8 | 0.15× | yes | +| tp4 | 2048 | h100 | cuda | 1.8626 | 1.9048 | 3.35× | 296.8 | 0.15× | yes | +| tp8 | 2048 | mi300x | native | 3.3579 | 3.9040 | 1.00× | 964.5 | 1.00× | yes | +| tp8 | 2048 | mi300x | triton | 1.8465 | 1.9464 | 1.82× | 148.5 | 0.15× | yes | +| tp8 | 2048 | mi300x | hip | 1.6417 | 1.8066 | 2.05× | 148.5 | 0.15× | yes | +| tp8 | 2048 | h100 | native | 3.9526 | 4.1235 | 1.00× | 964.5 | 1.00× | yes | +| tp8 | 2048 | h100 | triton | 1.8571 | 1.9126 | 2.13× | 148.5 | 0.15× | yes | +| tp8 | 2048 | h100 | cuda | 1.6378 | 2.7880 | 2.41× | 148.5 | 0.15× | yes | +| tp2_cp2 | 2048 | mi300x | native | 5.6092 | 5.7342 | 1.00× | 1929.0 | 1.00× | yes | +| tp2_cp2 | 2048 | mi300x | triton | 1.4580 | 1.7000 | 3.85× | 296.8 | 0.15× | yes | +| tp2_cp2 | 2048 | mi300x | hip | 1.3925 | 1.6368 | 4.03× | 296.8 | 0.15× | yes | +| tp2_cp2 | 2048 | h100 | native | 6.5482 | 6.6886 | 1.00× | 1929.0 | 1.00× | yes | +| tp2_cp2 | 2048 | h100 | triton | 1.5814 | 1.6376 | 4.14× | 296.8 | 0.15× | yes | +| tp2_cp2 | 2048 | h100 | cuda | 1.4018 | 1.4288 | 4.67× | 296.8 | 0.15× | yes | +| tp4_cp2 | 2048 | mi300x | native | 3.7864 | 3.8899 | 1.00× | 964.5 | 1.00× | yes | +| tp4_cp2 | 2048 | mi300x | triton | 1.6803 | 1.7273 | 2.25× | 148.4 | 0.15× | yes | +| tp4_cp2 | 2048 | mi300x | hip | 1.4029 | 1.4394 | 2.70× | 148.4 | 0.15× | yes | +| tp4_cp2 | 2048 | h100 | native | 4.4815 | 4.8207 | 1.00× | 964.5 | 1.00× | yes | +| tp4_cp2 | 2048 | h100 | triton | 1.6965 | 1.7702 | 2.64× | 148.4 | 0.15× | yes | +| tp4_cp2 | 2048 | h100 | cuda | 1.4813 | 1.4926 | 3.03× | 148.4 | 0.15× | yes | +| tp2_cp4 | 2048 | mi300x | native | 4.9700 | 5.0383 | 1.00× | 964.5 | 1.00× | yes | +| tp2_cp4 | 2048 | mi300x | triton | 1.5611 | 1.6497 | 3.18× | 148.4 | 0.15× | yes | +| tp2_cp4 | 2048 | mi300x | hip | 1.3959 | 1.5368 | 3.56× | 148.4 | 0.15× | yes | +| tp2_cp4 | 2048 | h100 | native | 5.6437 | 5.7185 | 1.00× | 964.5 | 1.00× | yes | +| tp2_cp4 | 2048 | h100 | triton | 1.6327 | 1.6834 | 3.46× | 148.4 | 0.15× | yes | +| tp2_cp4 | 2048 | h100 | cuda | 1.4039 | 3.7550 | 4.02× | 148.4 | 0.15× | yes | + +### Numerics versus `native` (distributed) + +| Topology | Tokens | Platform | Path | Mismatched elements (logp+LSE) | Relative L2 | +|---|---:|---|---|---:|---:| +| tp2 | 2048 | mi300x | triton | 137 | 1.454e-08 | +| tp2 | 2048 | mi300x | hip | 54 | 9.958e-09 | +| tp2 | 2048 | h100 | triton | 67 | 9.450e-09 | +| tp2 | 2048 | h100 | cuda | 67 | 1.050e-08 | +| tp4 | 2048 | mi300x | triton | 137 | 1.454e-08 | +| tp4 | 2048 | mi300x | hip | 54 | 9.958e-09 | +| tp4 | 2048 | h100 | triton | 67 | 9.450e-09 | +| tp4 | 2048 | h100 | cuda | 67 | 1.050e-08 | +| tp8 | 2048 | mi300x | triton | 137 | 1.454e-08 | +| tp8 | 2048 | mi300x | hip | 54 | 9.958e-09 | +| tp8 | 2048 | h100 | triton | 67 | 9.450e-09 | +| tp8 | 2048 | h100 | cuda | 67 | 1.050e-08 | +| tp2_cp2 | 2048 | mi300x | triton | 137 | 1.477e-08 | +| tp2_cp2 | 2048 | mi300x | hip | 54 | 1.069e-08 | +| tp2_cp2 | 2048 | h100 | triton | 67 | 9.804e-09 | +| tp2_cp2 | 2048 | h100 | cuda | 67 | 1.214e-08 | +| tp4_cp2 | 2048 | mi300x | triton | 137 | 1.477e-08 | +| tp4_cp2 | 2048 | mi300x | hip | 54 | 1.069e-08 | +| tp4_cp2 | 2048 | h100 | triton | 67 | 9.804e-09 | +| tp4_cp2 | 2048 | h100 | cuda | 67 | 1.214e-08 | +| tp2_cp4 | 2048 | mi300x | triton | 137 | 1.735e-08 | +| tp2_cp4 | 2048 | mi300x | hip | 54 | 1.162e-08 | +| tp2_cp4 | 2048 | h100 | triton | 67 | 1.075e-08 | +| tp2_cp4 | 2048 | h100 | cuda | 67 | 1.224e-08 | + +## Figures + +One line per backend and device across the full token sweep. The grid puts latency and peak memory, forward and forward+backward, on one page. + +![Single-device latency and memory grid](single_gpu_grid.png) + +The host and reference paths span three orders of magnitude, which flattens the kernel backends against each other. The second grid drops them and re-scales to the kernel backends alone, where the differences between Triton and the two vendor kernels are legible. + +![Kernel backends only](single_gpu_grid_kernels.png) + +![Single-device latency](single_gpu_latency.png) + +![Single-device peak memory](single_gpu_memory.png) + +![Distributed latency](distributed_logp_latency.png) diff --git a/benchmarks/results/pr328_rocm_mi300x/results.json b/benchmarks/results/pr328_rocm_mi300x/results.json new file mode 100644 index 00000000..89092ff4 --- /dev/null +++ b/benchmarks/results/pr328_rocm_mi300x/results.json @@ -0,0 +1,4134 @@ +{ + "config": { + "samples": 20, + "training_samples": 10, + "warmup": 5 + }, + "distributed": [ + { + "cp": 1, + "forward": { + "max_ms": 0.48768380656838417, + "median_ms": 0.4190076142549515, + "min_ms": 0.4007359966635704, + "p95_ms": 0.45819974038749933 + }, + "forward_peak_mib": 222.5634765625, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.2246337153243303e-06, + "logp_vs_fp64_rel_l2": 3.358224717653215e-08, + "lse_vs_fp64_max_abs": 1.2246337153243303e-06, + "mismatch_vs_reference": 14, + "path": "native", + "rel_l2_vs_reference": 1.4190225998557055e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 0.7625329308211803, + "median_ms": 0.6916869897395372, + "min_ms": 0.6808820180594921, + "p95_ms": 0.7598829921334982 + }, + "train_peak_mib": 296.75732421875 + }, + { + "cp": 1, + "forward": { + "max_ms": 3.3959737047553062, + "median_ms": 3.283221973106265, + "min_ms": 3.2631871290504932, + "p95_ms": 3.336320770904422 + }, + "forward_peak_mib": 81.8603515625, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.192615400213981e-06, + "logp_vs_fp64_rel_l2": 3.3594778807350763e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 4.6745226718485355, + "median_ms": 4.359946120530367, + "min_ms": 4.28913114592433, + "p95_ms": 4.57268045283854 + }, + "train_peak_mib": 482.30419921875 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.7916060276329517, + "median_ms": 0.7639192044734955, + "min_ms": 0.7495642639696598, + "p95_ms": 0.7835091790184379 + }, + "forward_peak_mib": 0.37548828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 9.623255152746424e-07, + "logp_vs_fp64_rel_l2": 3.1934274077238606e-08, + "lse_vs_fp64_max_abs": 9.623255152746424e-07, + "mismatch_vs_reference": 23, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.7002553523193623e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 2.098106313496828, + "median_ms": 1.546927960589528, + "min_ms": 1.3789492659270763, + "p95_ms": 1.8794237170368429 + }, + "train_peak_mib": 74.19873046875 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.7618209347128868, + "median_ms": 0.7162289693951607, + "min_ms": 0.707410741597414, + "p95_ms": 0.7614219095557928 + }, + "forward_peak_mib": 0.37548828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 9.63235819995134e-07, + "logp_vs_fp64_rel_l2": 3.3267799957748666e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "mismatch_vs_reference": 5, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.0160969201330157e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.3348530046641827, + "median_ms": 1.1684747878462076, + "min_ms": 1.128236297518015, + "p95_ms": 1.3319233199581504 + }, + "train_peak_mib": 74.19873046875 + }, + { + "cp": 1, + "forward": { + "max_ms": 1.3412716798484325, + "median_ms": 1.3117636553943157, + "min_ms": 1.2905574403703213, + "p95_ms": 1.3381803408265114 + }, + "forward_peak_mib": 1780.5078125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.3102978989626796e-08, + "lse_vs_fp64_max_abs": 1.3112566623618704e-06, + "mismatch_vs_reference": 88, + "path": "native", + "rel_l2_vs_reference": 1.261659473422803e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 3.034427296370268, + "median_ms": 2.72187776863575, + "min_ms": 2.6905969716608524, + "p95_ms": 2.901718509383499 + }, + "train_peak_mib": 2374.0498046875 + }, + { + "cp": 1, + "forward": { + "max_ms": 3.6768210120499134, + "median_ms": 3.5876790061593056, + "min_ms": 3.5461867228150368, + "p95_ms": 3.63216248806566 + }, + "forward_peak_mib": 650.3125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.314942529180228e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 7.70503468811512, + "median_ms": 7.148197619244456, + "min_ms": 7.095188833773136, + "p95_ms": 7.534388988278806 + }, + "train_peak_mib": 3857.9091796875 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.9566112421452999, + "median_ms": 0.8894458878785372, + "min_ms": 0.8791857399046421, + "p95_ms": 0.913844769820571 + }, + "forward_peak_mib": 3.001953125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.214240525949182e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 137, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.4538504826335258e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.9067618995904922, + "median_ms": 1.7666991334408522, + "min_ms": 1.579316332936287, + "p95_ms": 1.8881176132708788 + }, + "train_peak_mib": 593.5771484375 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.8942978456616402, + "median_ms": 0.8487105369567871, + "min_ms": 0.8369931019842625, + "p95_ms": 0.8818245492875576 + }, + "forward_peak_mib": 3.001953125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.3108948497165656e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 54, + "path": "ws2-rocm", + "rel_l2_vs_reference": 9.957929192632568e-09, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.7551067285239697, + "median_ms": 1.593762543052435, + "min_ms": 1.4162138104438782, + "p95_ms": 1.747589628212154 + }, + "train_peak_mib": 593.5771484375 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.4662228748202324, + "median_ms": 0.4294831305742264, + "min_ms": 0.41098007932305336, + "p95_ms": 0.46536598820239305 + }, + "forward_peak_mib": 111.2822265625, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.2246337153243303e-06, + "logp_vs_fp64_rel_l2": 3.358954781712721e-08, + "lse_vs_fp64_max_abs": 1.2246337153243303e-06, + "mismatch_vs_reference": 12, + "path": "native", + "rel_l2_vs_reference": 1.3632370710873585e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp4", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 0.7736789993941784, + "median_ms": 0.7229440379887819, + "min_ms": 0.6758440285921097, + "p95_ms": 0.7698032073676586 + }, + "train_peak_mib": 148.38232421875 + }, + { + "cp": 1, + "forward": { + "max_ms": 2.1146913059055805, + "median_ms": 2.0956681109964848, + "min_ms": 2.0710560493171215, + "p95_ms": 2.1139870397746563 + }, + "forward_peak_mib": 44.40966796875, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.192615400213981e-06, + "logp_vs_fp64_rel_l2": 3.3594778807350763e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp4", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 3.994573373347521, + "median_ms": 3.3612672705203295, + "min_ms": 3.3069918863475323, + "p95_ms": 3.754704259335994 + }, + "train_peak_mib": 241.158203125 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.9129159152507782, + "median_ms": 0.8004489354789257, + "min_ms": 0.7914560846984386, + "p95_ms": 0.8500181371346117 + }, + "forward_peak_mib": 0.31298828125, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 9.623255152746424e-07, + "logp_vs_fp64_rel_l2": 3.1934274077238606e-08, + "lse_vs_fp64_max_abs": 9.623255152746424e-07, + "mismatch_vs_reference": 23, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.7002553523193623e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp4", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.6431300900876522, + "median_ms": 1.5929813962429762, + "min_ms": 1.4386768452823162, + "p95_ms": 1.631475891917944 + }, + "train_peak_mib": 37.10498046875 + }, + { + "cp": 1, + "forward": { + "max_ms": 1.531684771180153, + "median_ms": 0.7509302813559771, + "min_ms": 0.737345777451992, + "p95_ms": 0.8267639204859739 + }, + "forward_peak_mib": 0.31298828125, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 9.63235819995134e-07, + "logp_vs_fp64_rel_l2": 3.3267799957748666e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "mismatch_vs_reference": 5, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.0160969201330157e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp4", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.3830838724970818, + "median_ms": 1.2959898449480534, + "min_ms": 1.1800932697951794, + "p95_ms": 1.3811147538945079 + }, + "train_peak_mib": 37.10498046875 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.769542995840311, + "median_ms": 0.7375199347734451, + "min_ms": 0.7233750075101852, + "p95_ms": 0.7685909979045391 + }, + "forward_peak_mib": 890.2578125, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.7274325472271812e-06, + "logp_vs_fp64_rel_l2": 3.3315658627415676e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 105, + "path": "native", + "rel_l2_vs_reference": 1.322253230879159e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp4", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 2.801039721816778, + "median_ms": 1.5266572590917349, + "min_ms": 1.4295442961156368, + "p95_ms": 2.2766862995922557 + }, + "train_peak_mib": 1187.0498046875 + }, + { + "cp": 1, + "forward": { + "max_ms": 2.3512630723416805, + "median_ms": 2.269135322421789, + "min_ms": 2.2396366111934185, + "p95_ms": 2.3080206010490656 + }, + "forward_peak_mib": 352.98681640625, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.314942529180228e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp4", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 4.546563141047955, + "median_ms": 4.468842409551144, + "min_ms": 4.372625146061182, + "p95_ms": 4.538554861210287 + }, + "train_peak_mib": 1928.99462890625 + }, + { + "cp": 1, + "forward": { + "max_ms": 1.1006738059222698, + "median_ms": 0.9221145883202553, + "min_ms": 0.9061959572136402, + "p95_ms": 0.9652029955759646 + }, + "forward_peak_mib": 2.501953125, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.214240525949182e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 137, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.4538504826335258e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp4", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 2.352474257349968, + "median_ms": 1.742522930726409, + "min_ms": 1.550281886011362, + "p95_ms": 2.1373696858063336 + }, + "train_peak_mib": 296.8271484375 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.9811879135668278, + "median_ms": 0.8714443538337946, + "min_ms": 0.854909885674715, + "p95_ms": 0.9162158239632845 + }, + "forward_peak_mib": 2.501953125, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.3108948497165656e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 54, + "path": "ws2-rocm", + "rel_l2_vs_reference": 9.957929192632568e-09, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp4", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.6562901437282562, + "median_ms": 1.5001380816102028, + "min_ms": 1.3153338804841042, + "p95_ms": 1.6030208440497518 + }, + "train_peak_mib": 296.8271484375 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.4549059085547924, + "median_ms": 0.40683988481760025, + "min_ms": 0.3881859593093395, + "p95_ms": 0.4503392381593585 + }, + "forward_peak_mib": 55.6416015625, + "grad_finite": true, + "local_vocab": 18992, + "logp_vs_fp64_max_abs": 1.192615400213981e-06, + "logp_vs_fp64_rel_l2": 3.3019515775307595e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "mismatch_vs_reference": 16, + "path": "native", + "rel_l2_vs_reference": 1.2835542362061829e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp8", + "tp": 8, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 0.8490509353578091, + "median_ms": 0.797174172475934, + "min_ms": 0.7719169370830059, + "p95_ms": 0.8395010139793158 + }, + "train_peak_mib": 74.19482421875 + }, + { + "cp": 1, + "forward": { + "max_ms": 1.7498787492513657, + "median_ms": 1.6366967465728521, + "min_ms": 1.585954800248146, + "p95_ms": 1.7218226799741387 + }, + "forward_peak_mib": 25.68408203125, + "grad_finite": true, + "local_vocab": 18992, + "logp_vs_fp64_max_abs": 1.192615400213981e-06, + "logp_vs_fp64_rel_l2": 3.3594778807350763e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp8", + "tp": 8, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 2.874980214983225, + "median_ms": 2.7750616427510977, + "min_ms": 2.6837559416890144, + "p95_ms": 2.871970064006746 + }, + "train_peak_mib": 120.58544921875 + }, + { + "cp": 1, + "forward": { + "max_ms": 2.6385802775621414, + "median_ms": 0.9444185998290777, + "min_ms": 0.8997973054647446, + "p95_ms": 2.2007678402587776 + }, + "forward_peak_mib": 0.28173828125, + "grad_finite": true, + "local_vocab": 18992, + "logp_vs_fp64_max_abs": 9.623255152746424e-07, + "logp_vs_fp64_rel_l2": 3.1934274077238606e-08, + "lse_vs_fp64_max_abs": 9.623255152746424e-07, + "mismatch_vs_reference": 23, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.7002553523193623e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp8", + "tp": 8, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.8622861243784428, + "median_ms": 1.6909562982618809, + "min_ms": 1.5980438329279423, + "p95_ms": 1.8021621042862535 + }, + "train_peak_mib": 18.55810546875 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.9400867857038975, + "median_ms": 0.8753803558647633, + "min_ms": 0.857312697917223, + "p95_ms": 0.9393630549311638 + }, + "forward_peak_mib": 0.28173828125, + "grad_finite": true, + "local_vocab": 18992, + "logp_vs_fp64_max_abs": 9.63235819995134e-07, + "logp_vs_fp64_rel_l2": 3.3267799957748666e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "mismatch_vs_reference": 5, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.0160969201330157e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 256, + "topology": "tp8", + "tp": 8, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.6698800027370453, + "median_ms": 1.4559680130332708, + "min_ms": 1.2819040566682816, + "p95_ms": 1.6370807774364948 + }, + "train_peak_mib": 18.55810546875 + }, + { + "cp": 1, + "forward": { + "max_ms": 0.5170977674424648, + "median_ms": 0.4875045269727707, + "min_ms": 0.4633176140487194, + "p95_ms": 0.5150433629751205 + }, + "forward_peak_mib": 445.1328125, + "grad_finite": true, + "local_vocab": 18992, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.363677481837659e-08, + "lse_vs_fp64_max_abs": 1.3207212923305178e-06, + "mismatch_vs_reference": 151, + "path": "native", + "rel_l2_vs_reference": 1.6072921840575048e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp8", + "tp": 8, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.050389837473631, + "median_ms": 0.9598308242857456, + "min_ms": 0.9048338979482651, + "p95_ms": 1.0330617194995284 + }, + "train_peak_mib": 593.5498046875 + }, + { + "cp": 1, + "forward": { + "max_ms": 1.8442091532051563, + "median_ms": 1.8065536860376596, + "min_ms": 1.7527742311358452, + "p95_ms": 1.8355898559093475 + }, + "forward_peak_mib": 204.32373046875, + "grad_finite": true, + "local_vocab": 18992, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.314942529180228e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp8", + "tp": 8, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 4.241500049829483, + "median_ms": 3.357867943122983, + "min_ms": 3.2875332981348038, + "p95_ms": 3.904002718627452 + }, + "train_peak_mib": 964.537109375 + }, + { + "cp": 1, + "forward": { + "max_ms": 1.1981800198554993, + "median_ms": 1.1148559860885143, + "min_ms": 1.0922616347670555, + "p95_ms": 1.1961441952735186 + }, + "forward_peak_mib": 2.251953125, + "grad_finite": true, + "local_vocab": 18992, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.214240525949182e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 137, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.4538504826335258e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp8", + "tp": 8, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.9532912410795689, + "median_ms": 1.8464724998921156, + "min_ms": 1.7132037319242954, + "p95_ms": 1.9464143086224794 + }, + "train_peak_mib": 148.4521484375 + }, + { + "cp": 1, + "forward": { + "max_ms": 1.148466020822525, + "median_ms": 1.0631137993186712, + "min_ms": 1.042358111590147, + "p95_ms": 1.128039206378162 + }, + "forward_peak_mib": 2.251953125, + "grad_finite": true, + "local_vocab": 18992, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.3108948497165656e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 54, + "path": "ws2-rocm", + "rel_l2_vs_reference": 9.957929192632568e-09, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 2048, + "topology": "tp8", + "tp": 8, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.8442701548337936, + "median_ms": 1.6416581347584724, + "min_ms": 1.5701819211244583, + "p95_ms": 1.806598319672048 + }, + "train_peak_mib": 148.4521484375 + }, + { + "cp": 2, + "forward": { + "max_ms": 0.9105419740080833, + "median_ms": 0.42372941970825195, + "min_ms": 0.4016370512545109, + "p95_ms": 0.5494242068380121 + }, + "forward_peak_mib": 111.28173828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.2246337153243303e-06, + "logp_vs_fp64_rel_l2": 3.5790710063228845e-08, + "lse_vs_fp64_max_abs": 1.2246337153243303e-06, + "mismatch_vs_reference": 14, + "path": "native", + "rel_l2_vs_reference": 1.8103588949760993e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 128, + "topology": "tp2_cp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 0.7477812469005585, + "median_ms": 0.6717976648360491, + "min_ms": 0.6506461650133133, + "p95_ms": 0.7439276669174433 + }, + "train_peak_mib": 148.37939453125 + }, + { + "cp": 2, + "forward": { + "max_ms": 3.271800000220537, + "median_ms": 3.2116554211825132, + "min_ms": 3.1774090602993965, + "p95_ms": 3.26847021933645 + }, + "forward_peak_mib": 41.2568359375, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.192615400213981e-06, + "logp_vs_fp64_rel_l2": 3.580900305363989e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 128, + "topology": "tp2_cp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 4.6104080975055695, + "median_ms": 4.5303895603865385, + "min_ms": 4.495627712458372, + "p95_ms": 4.583457973785698 + }, + "train_peak_mib": 241.18994140625 + }, + { + "cp": 2, + "forward": { + "max_ms": 0.7943799719214439, + "median_ms": 0.7632041815668344, + "min_ms": 0.7462790235877037, + "p95_ms": 0.7903264602646232 + }, + "forward_peak_mib": 0.18798828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 9.623255152746424e-07, + "logp_vs_fp64_rel_l2": 3.4143571581866415e-08, + "lse_vs_fp64_max_abs": 9.623255152746424e-07, + "mismatch_vs_reference": 23, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.920175576528477e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 128, + "topology": "tp2_cp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.5599783509969711, + "median_ms": 1.5244900714606047, + "min_ms": 1.417235005646944, + "p95_ms": 1.5520641580224037 + }, + "train_peak_mib": 37.1005859375 + }, + { + "cp": 2, + "forward": { + "max_ms": 0.7526082918047905, + "median_ms": 0.7347017526626587, + "min_ms": 0.6997189484536648, + "p95_ms": 0.748821534216404 + }, + "forward_peak_mib": 0.18798828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 9.63235819995134e-07, + "logp_vs_fp64_rel_l2": 3.519814568537689e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "mismatch_vs_reference": 5, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.4312143726175079e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 128, + "topology": "tp2_cp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 2.083924598991871, + "median_ms": 1.2248244602233171, + "min_ms": 1.148295123130083, + "p95_ms": 2.082325122319162 + }, + "train_peak_mib": 37.1005859375 + }, + { + "cp": 2, + "forward": { + "max_ms": 0.7630637846887112, + "median_ms": 0.7253275252878666, + "min_ms": 0.695972703397274, + "p95_ms": 0.7621122291311622 + }, + "forward_peak_mib": 890.25390625, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.322330169050237e-08, + "lse_vs_fp64_max_abs": 1.3112566623618704e-06, + "mismatch_vs_reference": 88, + "path": "native", + "rel_l2_vs_reference": 1.32910346411907e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 1024, + "topology": "tp2_cp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.48736871778965, + "median_ms": 1.447330228984356, + "min_ms": 1.4351606369018555, + "p95_ms": 1.4782787533476949 + }, + "train_peak_mib": 1187.025390625 + }, + { + "cp": 2, + "forward": { + "max_ms": 3.6023613065481186, + "median_ms": 3.3882129937410355, + "min_ms": 3.35833802819252, + "p95_ms": 3.5074760438874364 + }, + "forward_peak_mib": 325.482421875, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.318216084432305e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 1024, + "topology": "tp2_cp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 5.740798078477383, + "median_ms": 5.609216867014766, + "min_ms": 5.572086665779352, + "p95_ms": 5.734217865392566 + }, + "train_peak_mib": 1928.99169921875 + }, + { + "cp": 2, + "forward": { + "max_ms": 1.2611830607056618, + "median_ms": 0.8307832758873701, + "min_ms": 0.809862744063139, + "p95_ms": 0.960499467328191 + }, + "forward_peak_mib": 1.5009765625, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.238843099724084e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 137, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.4772159104804024e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 1024, + "topology": "tp2_cp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.7547458410263062, + "median_ms": 1.4580159913748503, + "min_ms": 1.441081054508686, + "p95_ms": 1.6999759711325169 + }, + "train_peak_mib": 296.7890625 + }, + { + "cp": 2, + "forward": { + "max_ms": 0.8257459849119186, + "median_ms": 0.7809000089764595, + "min_ms": 0.7658181712031364, + "p95_ms": 0.817440333776176 + }, + "forward_peak_mib": 1.5009765625, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.33456531944267e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 54, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.069130383334331e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 1024, + "topology": "tp2_cp2", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.6649230383336544, + "median_ms": 1.3925188686698675, + "min_ms": 1.26316724345088, + "p95_ms": 1.6368331853300333 + }, + "train_peak_mib": 296.7890625 + }, + { + "cp": 2, + "forward": { + "max_ms": 0.9470968507230282, + "median_ms": 0.4659618716686964, + "min_ms": 0.4259929992258549, + "p95_ms": 0.6475532660260799 + }, + "forward_peak_mib": 55.64111328125, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.2246337153243303e-06, + "logp_vs_fp64_rel_l2": 3.573556089056118e-08, + "lse_vs_fp64_max_abs": 1.2246337153243303e-06, + "mismatch_vs_reference": 11, + "path": "native", + "rel_l2_vs_reference": 1.693435683024808e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 128, + "topology": "tp4_cp2", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 0.8712639100849628, + "median_ms": 0.7841994520276785, + "min_ms": 0.7658470422029495, + "p95_ms": 0.8438089862465858 + }, + "train_peak_mib": 74.19189453125 + }, + { + "cp": 2, + "forward": { + "max_ms": 2.360045909881592, + "median_ms": 2.1834224462509155, + "min_ms": 2.1604690700769424, + "p95_ms": 2.347819460555911 + }, + "forward_peak_mib": 22.36865234375, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.192615400213981e-06, + "logp_vs_fp64_rel_l2": 3.580900305363989e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 128, + "topology": "tp4_cp2", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 3.557675052434206, + "median_ms": 3.459763713181019, + "min_ms": 3.397907130420208, + "p95_ms": 3.5327076679095626 + }, + "train_peak_mib": 120.59912109375 + }, + { + "cp": 2, + "forward": { + "max_ms": 0.9298510849475861, + "median_ms": 0.8165778126567602, + "min_ms": 0.7963827811181545, + "p95_ms": 0.8834703825414181 + }, + "forward_peak_mib": 0.15673828125, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 9.623255152746424e-07, + "logp_vs_fp64_rel_l2": 3.4143571581866415e-08, + "lse_vs_fp64_max_abs": 9.623255152746424e-07, + "mismatch_vs_reference": 23, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.920175576528477e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 128, + "topology": "tp4_cp2", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.890428364276886, + "median_ms": 1.6159303486347198, + "min_ms": 1.582070253789425, + "p95_ms": 1.8853715620934963 + }, + "train_peak_mib": 18.5537109375 + }, + { + "cp": 2, + "forward": { + "max_ms": 1.038712915033102, + "median_ms": 0.7646705489605665, + "min_ms": 0.7514357566833496, + "p95_ms": 0.8226943202316763 + }, + "forward_peak_mib": 0.15673828125, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 9.63235819995134e-07, + "logp_vs_fp64_rel_l2": 3.519814568537689e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "mismatch_vs_reference": 5, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.4312143726175079e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 128, + "topology": "tp4_cp2", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.5644128434360027, + "median_ms": 1.407951582223177, + "min_ms": 1.3164947740733624, + "p95_ms": 1.5445253113284707 + }, + "train_peak_mib": 18.5537109375 + }, + { + "cp": 2, + "forward": { + "max_ms": 1.428872812539339, + "median_ms": 0.49603707157075405, + "min_ms": 0.45394524931907654, + "p95_ms": 0.6215794477611786 + }, + "forward_peak_mib": 445.12890625, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.7274325472271812e-06, + "logp_vs_fp64_rel_l2": 3.337932457263568e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 101, + "path": "native", + "rel_l2_vs_reference": 1.386501020113035e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 1024, + "topology": "tp4_cp2", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.417775172740221, + "median_ms": 1.0220431722700596, + "min_ms": 0.8883103728294373, + "p95_ms": 1.257734047248959 + }, + "train_peak_mib": 593.525390625 + }, + { + "cp": 2, + "forward": { + "max_ms": 2.2686789743602276, + "median_ms": 2.2382440511137247, + "min_ms": 2.197904046624899, + "p95_ms": 2.263028477318585 + }, + "forward_peak_mib": 176.65673828125, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.318216084432305e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 1024, + "topology": "tp4_cp2", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 3.9137532003223896, + "median_ms": 3.7863540928810835, + "min_ms": 3.763199783861637, + "p95_ms": 3.889885521493852 + }, + "train_peak_mib": 964.51611328125 + }, + { + "cp": 2, + "forward": { + "max_ms": 1.006274949759245, + "median_ms": 0.8831163868308067, + "min_ms": 0.8671479299664497, + "p95_ms": 0.9622140787541866 + }, + "forward_peak_mib": 1.2509765625, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.238843099724084e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 137, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.4772159104804024e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 1024, + "topology": "tp4_cp2", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.7355279996991158, + "median_ms": 1.6802609898149967, + "min_ms": 1.5333471819758415, + "p95_ms": 1.7273348988965154 + }, + "train_peak_mib": 148.4140625 + }, + { + "cp": 2, + "forward": { + "max_ms": 0.9049340151250362, + "median_ms": 0.838590320199728, + "min_ms": 0.8162329904735088, + "p95_ms": 0.8977984543889761 + }, + "forward_peak_mib": 1.2509765625, + "grad_finite": true, + "local_vocab": 37984, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.33456531944267e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 54, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.069130383334331e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 1024, + "topology": "tp4_cp2", + "tp": 4, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.4411099255084991, + "median_ms": 1.402924070134759, + "min_ms": 1.2808619067072868, + "p95_ms": 1.4394335448741913 + }, + "train_peak_mib": 148.4140625 + }, + { + "cp": 4, + "forward": { + "max_ms": 0.9297211654484272, + "median_ms": 0.4601990804076195, + "min_ms": 0.4170001484453678, + "p95_ms": 0.5840539000928404 + }, + "forward_peak_mib": 55.64111328125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.2246337153243303e-06, + "logp_vs_fp64_rel_l2": 3.730282272543457e-08, + "lse_vs_fp64_max_abs": 1.2246337153243303e-06, + "mismatch_vs_reference": 14, + "path": "native", + "rel_l2_vs_reference": 2.3950332323559945e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 64, + "topology": "tp2_cp4", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.0157087817788124, + "median_ms": 0.8255115244537592, + "min_ms": 0.7808287627995014, + "p95_ms": 0.986437709070742 + }, + "train_peak_mib": 74.19140625 + }, + { + "cp": 4, + "forward": { + "max_ms": 3.438376821577549, + "median_ms": 3.382489550858736, + "min_ms": 3.308103885501623, + "p95_ms": 3.4154762281104922 + }, + "forward_peak_mib": 20.970703125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.192615400213981e-06, + "logp_vs_fp64_rel_l2": 3.646037610209301e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 64, + "topology": "tp2_cp4", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 5.002471152693033, + "median_ms": 4.629456205293536, + "min_ms": 4.59907203912735, + "p95_ms": 4.858391894958913 + }, + "train_peak_mib": 120.63427734375 + }, + { + "cp": 4, + "forward": { + "max_ms": 1.042709220200777, + "median_ms": 0.768306665122509, + "min_ms": 0.7598591037094593, + "p95_ms": 0.9768236195668579 + }, + "forward_peak_mib": 0.09423828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 9.623255152746424e-07, + "logp_vs_fp64_rel_l2": 3.534562677873104e-08, + "lse_vs_fp64_max_abs": 9.623255152746424e-07, + "mismatch_vs_reference": 23, + "path": "ws2-triton", + "rel_l2_vs_reference": 2.2173699061384028e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 64, + "topology": "tp2_cp4", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.9474220462143421, + "median_ms": 1.5305932611227036, + "min_ms": 1.3936497271060944, + "p95_ms": 1.927389926277101 + }, + "train_peak_mib": 18.552734375 + }, + { + "cp": 4, + "forward": { + "max_ms": 1.0852618142962456, + "median_ms": 0.7238299585878849, + "min_ms": 0.7138201035559177, + "p95_ms": 0.788106629624963 + }, + "forward_peak_mib": 0.09423828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 9.63235819995134e-07, + "logp_vs_fp64_rel_l2": 3.531619957105821e-08, + "lse_vs_fp64_max_abs": 9.63235819995134e-07, + "mismatch_vs_reference": 5, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.8104749470140388e-08, + "repeat_bitwise": true, + "tokens": 256, + "tokens_per_cp_rank": 64, + "topology": "tp2_cp4", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.569060143083334, + "median_ms": 1.2379284016788006, + "min_ms": 1.1319010518491268, + "p95_ms": 1.5026589157059789 + }, + "train_peak_mib": 18.552734375 + }, + { + "cp": 4, + "forward": { + "max_ms": 0.5426062270998955, + "median_ms": 0.4625171422958374, + "min_ms": 0.43542683124542236, + "p95_ms": 0.5310371518135071 + }, + "forward_peak_mib": 445.126953125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.339807416706557e-08, + "lse_vs_fp64_max_abs": 1.3112566623618704e-06, + "mismatch_vs_reference": 88, + "path": "native", + "rel_l2_vs_reference": 1.368420826566803e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 512, + "topology": "tp2_cp4", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.1636787094175816, + "median_ms": 0.9963898919522762, + "min_ms": 0.9112842381000519, + "p95_ms": 1.1164123192429543 + }, + "train_peak_mib": 593.51318359375 + }, + { + "cp": 4, + "forward": { + "max_ms": 3.6994353868067265, + "median_ms": 3.526117419824004, + "min_ms": 3.4705549478530884, + "p95_ms": 3.6070818547159433 + }, + "forward_peak_mib": 163.0673828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.323407831042024e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 512, + "topology": "tp2_cp4", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 5.049319937825203, + "median_ms": 4.9699717201292515, + "min_ms": 4.881730768829584, + "p95_ms": 5.038278875872493 + }, + "train_peak_mib": 964.53271484375 + }, + { + "cp": 4, + "forward": { + "max_ms": 1.389303244650364, + "median_ms": 0.8240439929068089, + "min_ms": 0.7940488867461681, + "p95_ms": 1.3511136174201965 + }, + "forward_peak_mib": 0.75048828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.277486808560087e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 137, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.7349402400487287e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 512, + "topology": "tp2_cp4", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.657932996749878, + "median_ms": 1.5610686969012022, + "min_ms": 1.4730580151081085, + "p95_ms": 1.649744505994022 + }, + "train_peak_mib": 148.39501953125 + }, + { + "cp": 4, + "forward": { + "max_ms": 0.8596866391599178, + "median_ms": 0.754661625251174, + "min_ms": 0.7468885742127895, + "p95_ms": 0.7833839161321521 + }, + "forward_peak_mib": 0.75048828125, + "grad_finite": true, + "local_vocab": 75968, + "logp_vs_fp64_max_abs": 1.4523163969215602e-06, + "logp_vs_fp64_rel_l2": 3.347644666288935e-08, + "lse_vs_fp64_max_abs": 1.2623544449752444e-06, + "mismatch_vs_reference": 54, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.161601588696792e-08, + "repeat_bitwise": true, + "tokens": 2048, + "tokens_per_cp_rank": 512, + "topology": "tp2_cp4", + "tp": 2, + "tp_replicated": true, + "train_fwd_bwd": { + "max_ms": 1.5739467926323414, + "median_ms": 1.3958783820271492, + "min_ms": 1.319149974733591, + "p95_ms": 1.5368028078228235 + }, + "train_peak_mib": 148.39501953125 + } + ], + "environment": { + "architecture": "gfx942:sramecc+:xnack-", + "extension_symbols": "hip_deterministic_logp_tile_stats, hip_deterministic_logp_backward", + "git_commit": "e9f1d2a5c67a283bd987978614b4444a9416c1f0", + "gpu": "AMD Instinct MI300X", + "gpu_count": 8, + "hip": "7.14.60850", + "native_collective": "torch.distributed ProcessGroupNCCL (RCCL on ROCm)", + "python": "3.12.3", + "torch": "2.12.0+rocm7.14.0a20260608" + }, + "single_gpu": { + "cases": [ + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.20935100317001343, + "median_ms": 0.12019799649715424, + "min_ms": 0.11268699914216995, + "p95_ms": 0.16653735563158992 + }, + "forward_peak_mib": 1.1611328125, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.365204219985256e-08 + }, + "lse_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.8718812743187373e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 0.4332039952278137, + "median_ms": 0.3826484978199005, + "min_ms": 0.3718720078468323, + "p95_ms": 0.42406404614448545 + }, + "train_peak_mib": 3.1904296875 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.28041699528694153, + "median_ms": 0.18904100358486176, + "min_ms": 0.18267099559307098, + "p95_ms": 0.2714727446436882 + }, + "forward_peak_mib": 1.744140625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.365204219985256e-08 + }, + "lse_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.8718812743187373e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 0.693149983882904, + "median_ms": 0.5962255001068115, + "min_ms": 0.5884339809417725, + "p95_ms": 0.6668128371238708 + }, + "train_peak_mib": 2.03466796875 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.3535250127315521, + "median_ms": 0.30254900455474854, + "min_ms": 0.29940399527549744, + "p95_ms": 0.3455716043710709 + }, + "forward_peak_mib": 0.00341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.365204219985256e-08 + }, + "lse_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.8718812743187373e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 0.3893379867076874, + "median_ms": 0.3491179943084717, + "min_ms": 0.34563401341438293, + "p95_ms": 0.3789542436599731 + }, + "train_peak_mib": 0.5830078125 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 5.050661087036133, + "median_ms": 4.9217095375061035, + "min_ms": 4.765758037567139, + "p95_ms": 5.02729377746582 + }, + "forward_peak_mib": 1.96630859375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.365204219985256e-08 + }, + "lse_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.8718812743187373e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 5.651031970977783, + "median_ms": 5.608389139175415, + "min_ms": 5.570312976837158, + "p95_ms": 5.641712045669555 + }, + "train_peak_mib": 3.91943359375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.453794002532959, + "median_ms": 0.39360450208187103, + "min_ms": 0.38184699416160583, + "p95_ms": 0.4045488551259041 + }, + "forward_peak_mib": 0.005859375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.365204219985256e-08 + }, + "lse_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.8718812743187373e-08 + }, + "mismatch_vs_reference": 0, + "path": "ws2-triton", + "rel_l2_vs_reference": 0.0, + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 0.8825510144233704, + "median_ms": 0.7837440073490143, + "min_ms": 0.7605699896812439, + "p95_ms": 0.8616399496793746 + }, + "train_peak_mib": 0.5859375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.3889780044555664, + "median_ms": 0.34681499004364014, + "min_ms": 0.33045101165771484, + "p95_ms": 0.3638219863176346 + }, + "forward_peak_mib": 0.005859375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.365204219985256e-08 + }, + "lse_vs_fp64": { + "max_abs": 5.384272654396227e-07, + "relative_l2": 3.8718812743187373e-08 + }, + "mismatch_vs_reference": 0, + "path": "ws2-rocm", + "rel_l2_vs_reference": 0.0, + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 0.6693540215492249, + "median_ms": 0.6335409879684448, + "min_ms": 0.627170979976654, + "p95_ms": 0.6575284570455551 + }, + "train_peak_mib": 0.5859375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.15094399452209473, + "median_ms": 0.12752900272607803, + "min_ms": 0.12198100239038467, + "p95_ms": 0.14850819632411003 + }, + "forward_peak_mib": 9.27685546875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 8.398243309670761e-07, + "relative_l2": 3.9344983122904206e-08 + }, + "lse_vs_fp64": { + "max_abs": 8.398243291907193e-07, + "relative_l2": 3.693892459584331e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 0.4181010127067566, + "median_ms": 0.38661450147628784, + "min_ms": 0.3778409957885742, + "p95_ms": 0.41184600740671157 + }, + "train_peak_mib": 25.50439453125 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.22136899828910828, + "median_ms": 0.186537005007267, + "min_ms": 0.1813099980354309, + "p95_ms": 0.20447135344147685 + }, + "forward_peak_mib": 13.91552734375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 8.398243309670761e-07, + "relative_l2": 3.9344983122904206e-08 + }, + "lse_vs_fp64": { + "max_abs": 8.398243291907193e-07, + "relative_l2": 3.693892459584331e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 0.7002000212669373, + "median_ms": 0.6067410111427307, + "min_ms": 0.5904369950294495, + "p95_ms": 0.6935661137104034 + }, + "train_peak_mib": 16.234375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.3513219952583313, + "median_ms": 0.3033300042152405, + "min_ms": 0.29960501194000244, + "p95_ms": 0.30980129539966583 + }, + "forward_peak_mib": 0.00341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 8.398243309670761e-07, + "relative_l2": 3.766025511891362e-08 + }, + "lse_vs_fp64": { + "max_abs": 8.398243291907193e-07, + "relative_l2": 3.535840565806921e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 0.37147098779678345, + "median_ms": 0.3569300025701523, + "min_ms": 0.3475959897041321, + "p95_ms": 0.3688393920660019 + }, + "train_peak_mib": 4.6396484375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 5.572154998779297, + "median_ms": 5.4453675746917725, + "min_ms": 5.37157678604126, + "p95_ms": 5.541255283355713 + }, + "forward_peak_mib": 6.22314453125, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 8.398243309670761e-07, + "relative_l2": 3.9344983122904206e-08 + }, + "lse_vs_fp64": { + "max_abs": 8.398243291907193e-07, + "relative_l2": 3.693892459584331e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 6.1761322021484375, + "median_ms": 6.16493558883667, + "min_ms": 6.118648052215576, + "p95_ms": 6.175754117965699 + }, + "train_peak_mib": 30.29052734375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.43841099739074707, + "median_ms": 0.3986324965953827, + "min_ms": 0.3909800052642822, + "p95_ms": 0.42318820357322695 + }, + "forward_peak_mib": 0.0107421875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.1808696349646652e-06, + "relative_l2": 4.670865489897585e-08 + }, + "lse_vs_fp64": { + "max_abs": 8.398243291907193e-07, + "relative_l2": 3.2524573103729204e-08 + }, + "mismatch_vs_reference": 2, + "path": "ws2-triton", + "rel_l2_vs_reference": 5.157416402269682e-08, + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 1.199422001838684, + "median_ms": 1.0126639902591705, + "min_ms": 0.978613018989563, + "p95_ms": 1.185198837518692 + }, + "train_peak_mib": 4.642578125 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.39446601271629333, + "median_ms": 0.35200299322605133, + "min_ms": 0.3450320065021515, + "p95_ms": 0.36558124870061875 + }, + "forward_peak_mib": 0.0107421875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 8.398243309670761e-07, + "relative_l2": 3.9344983122904206e-08 + }, + "lse_vs_fp64": { + "max_abs": 8.398243291907193e-07, + "relative_l2": 3.693892459584331e-08 + }, + "mismatch_vs_reference": 0, + "path": "ws2-rocm", + "rel_l2_vs_reference": 0.0, + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 0.9747679829597473, + "median_ms": 0.7770539820194244, + "min_ms": 0.7595679759979248, + "p95_ms": 0.9017950803041456 + }, + "train_peak_mib": 4.642578125 + }, + { + "batch_invariant": false, + "dtype": "bf16", + "forward": { + "max_ms": 0.1689710021018982, + "median_ms": 0.1626409962773323, + "min_ms": 0.13307799398899078, + "p95_ms": 0.16706815287470816 + }, + "forward_peak_mib": 37.10302734375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.8589741301866525e-06, + "relative_l2": 4.183495189450554e-08 + }, + "lse_vs_fp64": { + "max_abs": 9.292139733219074e-07, + "relative_l2": 3.427046566031066e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 0.5343539714813232, + "median_ms": 0.5017854869365692, + "min_ms": 0.4919700026512146, + "p95_ms": 0.5303696841001511 + }, + "train_peak_mib": 102.01025390625 + }, + { + "batch_invariant": false, + "dtype": "bf16", + "forward": { + "max_ms": 0.23274600505828857, + "median_ms": 0.19759349524974823, + "min_ms": 0.17562100291252136, + "p95_ms": 0.2175992012023926 + }, + "forward_peak_mib": 55.6513671875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.8589741301866525e-06, + "relative_l2": 4.183495189450554e-08 + }, + "lse_vs_fp64": { + "max_abs": 9.292139733219074e-07, + "relative_l2": 3.427046566031066e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 0.88919997215271, + "median_ms": 0.8541084825992584, + "min_ms": 0.8356819748878479, + "p95_ms": 0.8848016858100891 + }, + "train_peak_mib": 64.92529296875 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.35424599051475525, + "median_ms": 0.3021489977836609, + "min_ms": 0.2995249927043915, + "p95_ms": 0.3195387080311775 + }, + "forward_peak_mib": 0.00341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.165070511177646e-06, + "relative_l2": 4.043209774336302e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.0317191794229075e-06, + "relative_l2": 3.629703680141769e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 0.49056899547576904, + "median_ms": 0.4564579874277115, + "min_ms": 0.44145599007606506, + "p95_ms": 0.48285329937934873 + }, + "train_peak_mib": 18.5498046875 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 5.4874701499938965, + "median_ms": 5.397536516189575, + "min_ms": 5.278359889984131, + "p95_ms": 5.447852993011475 + }, + "forward_peak_mib": 20.78515625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 9.292139733219074e-07, + "relative_l2": 3.3261220288821645e-08 + }, + "lse_vs_fp64": { + "max_abs": 9.292139733219074e-07, + "relative_l2": 3.229389879068821e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 6.701952934265137, + "median_ms": 6.617066860198975, + "min_ms": 6.586542129516602, + "p95_ms": 6.6741920709609985 + }, + "train_peak_mib": 120.70654296875 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.44017401337623596, + "median_ms": 0.40299899876117706, + "min_ms": 0.3907009959220886, + "p95_ms": 0.4346934497356415 + }, + "forward_peak_mib": 0.0341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 9.292139733219074e-07, + "relative_l2": 3.208024712173244e-08 + }, + "lse_vs_fp64": { + "max_abs": 9.292139733219074e-07, + "relative_l2": 3.1245170134439244e-08 + }, + "mismatch_vs_reference": 4, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.8415464321385415e-08, + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 1.0324130058288574, + "median_ms": 0.984281986951828, + "min_ms": 0.9665150046348572, + "p95_ms": 1.0247517585754395 + }, + "train_peak_mib": 18.552734375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.4256319999694824, + "median_ms": 0.3501400053501129, + "min_ms": 0.34258899092674255, + "p95_ms": 0.37573989033699035 + }, + "forward_peak_mib": 0.0341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 9.292139733219074e-07, + "relative_l2": 3.3261220288821645e-08 + }, + "lse_vs_fp64": { + "max_abs": 9.292139733219074e-07, + "relative_l2": 3.229389879068821e-08 + }, + "mismatch_vs_reference": 0, + "path": "ws2-rocm", + "rel_l2_vs_reference": 0.0, + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 0.8069980144500732, + "median_ms": 0.7662184834480286, + "min_ms": 0.7476300001144409, + "p95_ms": 0.8005449950695038 + }, + "train_peak_mib": 18.552734375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.33245301246643066, + "median_ms": 0.29698099195957184, + "min_ms": 0.26515400409698486, + "p95_ms": 0.32309265732765197 + }, + "forward_peak_mib": 148.40771484375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.4593867163625873e-06, + "relative_l2": 3.7283481956277905e-08 + }, + "lse_vs_fp64": { + "max_abs": 9.801408431542313e-07, + "relative_l2": 3.050237461618506e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 0.626570999622345, + "median_ms": 0.5764960050582886, + "min_ms": 0.5228559970855713, + "p95_ms": 0.6142585605382919 + }, + "train_peak_mib": 408.03369140625 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.3854529857635498, + "median_ms": 0.303850993514061, + "min_ms": 0.2683979868888855, + "p95_ms": 0.33457099646329885 + }, + "forward_peak_mib": 222.59765625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.4593867163625873e-06, + "relative_l2": 3.7283481956277905e-08 + }, + "lse_vs_fp64": { + "max_abs": 9.801408431542313e-07, + "relative_l2": 3.050237461618506e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 1.1429779529571533, + "median_ms": 0.8844340145587921, + "min_ms": 0.7590069770812988, + "p95_ms": 1.0753051757812497 + }, + "train_peak_mib": 259.6923828125 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.36718499660491943, + "median_ms": 0.31983450055122375, + "min_ms": 0.3176719844341278, + "p95_ms": 0.3301939100027085 + }, + "forward_peak_mib": 0.00439453125, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.2187595608281754e-06, + "relative_l2": 3.6275552663920587e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2187595608281754e-06, + "relative_l2": 3.5171515827779054e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 0.517408013343811, + "median_ms": 0.46639250218868256, + "min_ms": 0.45699799060821533, + "p95_ms": 0.503797760605812 + }, + "train_peak_mib": 74.19091796875 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 5.568389892578125, + "median_ms": 5.465717315673828, + "min_ms": 5.350906848907471, + "p95_ms": 5.539391398429871 + }, + "forward_peak_mib": 79.03369140625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.4593867163625873e-06, + "relative_l2": 3.674698683930475e-08 + }, + "lse_vs_fp64": { + "max_abs": 9.801408431542313e-07, + "relative_l2": 2.9901949672921016e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 6.825016975402832, + "median_ms": 6.794530868530273, + "min_ms": 6.7457780838012695, + "p95_ms": 6.824007177352906 + }, + "train_peak_mib": 482.3720703125 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.46625199913978577, + "median_ms": 0.4062635004520416, + "min_ms": 0.3971889913082123, + "p95_ms": 0.42511413693428046 + }, + "forward_peak_mib": 0.12841796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.4593867163625873e-06, + "relative_l2": 3.478970586475828e-08 + }, + "lse_vs_fp64": { + "max_abs": 8.965857389853227e-07, + "relative_l2": 2.903246856398248e-08 + }, + "mismatch_vs_reference": 7, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.5511401243457844e-08, + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 1.1938129663467407, + "median_ms": 1.0162895321846008, + "min_ms": 0.9801759719848633, + "p95_ms": 1.1880444467067717 + }, + "train_peak_mib": 74.1943359375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.4069640040397644, + "median_ms": 0.35552799701690674, + "min_ms": 0.3464739918708801, + "p95_ms": 0.37290458977222446 + }, + "forward_peak_mib": 0.12841796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.4593867163625873e-06, + "relative_l2": 3.58337193479546e-08 + }, + "lse_vs_fp64": { + "max_abs": 9.801408431542313e-07, + "relative_l2": 3.0008002695414004e-08 + }, + "mismatch_vs_reference": 2, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.2665006080681408e-08, + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 0.9827389717102051, + "median_ms": 0.791395515203476, + "min_ms": 0.7667790055274963, + "p95_ms": 0.9700485289096832 + }, + "train_peak_mib": 74.1943359375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.8830310106277466, + "median_ms": 0.8542289733886719, + "min_ms": 0.8202589750289917, + "p95_ms": 0.8800632119178772 + }, + "forward_peak_mib": 593.630859375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.4498686553565676e-06, + "relative_l2": 3.3782278450435274e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.265310253018015e-06, + "relative_l2": 3.100727931388369e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 1.9153649806976318, + "median_ms": 1.8550350069999695, + "min_ms": 1.806203007698059, + "p95_ms": 1.909361982345581 + }, + "train_peak_mib": 1632.1318359375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.9521740078926086, + "median_ms": 0.9142979979515076, + "min_ms": 0.8730159997940063, + "p95_ms": 0.9433076441287995 + }, + "forward_peak_mib": 890.38916015625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.4498686553565676e-06, + "relative_l2": 3.3782278450435274e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.265310253018015e-06, + "relative_l2": 3.100727931388369e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 2.2674479484558105, + "median_ms": 2.255669951438904, + "min_ms": 2.2023909091949463, + "p95_ms": 2.267303538322449 + }, + "train_peak_mib": 1038.76806640625 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.3829689919948578, + "median_ms": 0.33712050318717957, + "min_ms": 0.33297398686408997, + "p95_ms": 0.3422111362218857 + }, + "forward_peak_mib": 0.01123046875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.767553932552346e-06, + "relative_l2": 3.649871630209298e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2392332031652131e-06, + "relative_l2": 3.379189024884559e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 0.5339930057525635, + "median_ms": 0.4645095020532608, + "min_ms": 0.4525119960308075, + "p95_ms": 0.5160924524068832 + }, + "train_peak_mib": 296.7607421875 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 5.67739200592041, + "median_ms": 5.543973922729492, + "min_ms": 5.352829933166504, + "p95_ms": 5.667003011703491 + }, + "forward_peak_mib": 312.21923828125, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.4498686553565676e-06, + "relative_l2": 3.3724263980110736e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.265310253018015e-06, + "relative_l2": 3.053555163043888e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 7.908987045288086, + "median_ms": 7.684312582015991, + "min_ms": 7.606616973876953, + "p95_ms": 7.843495869636535 + }, + "train_peak_mib": 1929.0439453125 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.4894070029258728, + "median_ms": 0.4432384967803955, + "min_ms": 0.43492600321769714, + "p95_ms": 0.4798547476530075 + }, + "forward_peak_mib": 0.5107421875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.39654445874271e-06, + "relative_l2": 3.263898958223032e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.265310253018015e-06, + "relative_l2": 2.91707358884744e-08 + }, + "mismatch_vs_reference": 43, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.757610490869294e-08, + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 1.2311480045318604, + "median_ms": 1.1473445296287537, + "min_ms": 1.0201549530029297, + "p95_ms": 1.223937636613846 + }, + "train_peak_mib": 296.77001953125 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.44249799847602844, + "median_ms": 0.40269799530506134, + "min_ms": 0.3962689936161041, + "p95_ms": 0.4225175932049751 + }, + "forward_peak_mib": 0.5107421875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.39654445874271e-06, + "relative_l2": 3.323884866087576e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.265310253018015e-06, + "relative_l2": 3.056490829487832e-08 + }, + "mismatch_vs_reference": 16, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.2006765443781159e-08, + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 1.4356119632720947, + "median_ms": 0.8617799878120422, + "min_ms": 0.8286309838294983, + "p95_ms": 1.252316182851791 + }, + "train_peak_mib": 296.77001953125 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 2.7969539165496826, + "median_ms": 2.302619457244873, + "min_ms": 2.263361930847168, + "p95_ms": 2.788733637332916 + }, + "forward_peak_mib": 2374.015625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.5566917461740104e-06, + "relative_l2": 3.33818552584321e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.3818783255459266e-06, + "relative_l2": 3.089300583262462e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 6.601564884185791, + "median_ms": 6.448056936264038, + "min_ms": 6.240708827972412, + "p95_ms": 6.571694302558899 + }, + "train_peak_mib": 6528.5244140625 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 2.7635040283203125, + "median_ms": 2.2864960432052612, + "min_ms": 2.2438929080963135, + "p95_ms": 2.6859070658683777 + }, + "forward_peak_mib": 3561.103515625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.5566917461740104e-06, + "relative_l2": 3.33818552584321e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.3818783255459266e-06, + "relative_l2": 3.089300583262462e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 7.106554985046387, + "median_ms": 7.031663656234741, + "min_ms": 6.984973907470703, + "p95_ms": 7.0932692527771 + }, + "train_peak_mib": 4154.619140625 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.6292549967765808, + "median_ms": 0.5970470011234283, + "min_ms": 0.594202995300293, + "p95_ms": 0.6152880758047103 + }, + "forward_peak_mib": 0.044921875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.5885667394854863e-06, + "relative_l2": 3.6268790641540615e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.5157620829597818e-06, + "relative_l2": 3.420512007493752e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 1.160202980041504, + "median_ms": 1.1344245076179504, + "min_ms": 1.1143759489059448, + "p95_ms": 1.1601490139961244 + }, + "train_peak_mib": 1187.0400390625 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 6.33885383605957, + "median_ms": 6.138717412948608, + "min_ms": 5.910459041595459, + "p95_ms": 6.307230234146118 + }, + "forward_peak_mib": 1244.96435546875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.5566917461740104e-06, + "relative_l2": 3.363042993957078e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.3818783255459266e-06, + "relative_l2": 3.0918720395375866e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 12.480816841125488, + "median_ms": 12.347699165344238, + "min_ms": 12.259528160095215, + "p95_ms": 12.468018198013306 + }, + "train_peak_mib": 7715.73779296875 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 1.002089023590088, + "median_ms": 0.9513734877109528, + "min_ms": 0.9448429942131042, + "p95_ms": 0.9800537467002869 + }, + "forward_peak_mib": 2.04296875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.5566917461740104e-06, + "relative_l2": 3.227948726455519e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.3182267863953712e-06, + "relative_l2": 2.972527831148482e-08 + }, + "mismatch_vs_reference": 169, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.7134613775374536e-08, + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 2.2043540477752686, + "median_ms": 2.1618911027908325, + "min_ms": 2.0272109508514404, + "p95_ms": 2.19452965259552 + }, + "train_peak_mib": 1187.0771484375 + }, + { + "batch_invariant": true, + "dtype": "bf16", + "forward": { + "max_ms": 0.9396349787712097, + "median_ms": 0.892345517873764, + "min_ms": 0.8854749798774719, + "p95_ms": 0.9161548167467117 + }, + "forward_peak_mib": 2.04296875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.6214356648447392e-06, + "relative_l2": 3.334502491966002e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.3182267863953712e-06, + "relative_l2": 3.0632826213818444e-08 + }, + "mismatch_vs_reference": 65, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.09810589785962e-08, + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 2.052248001098633, + "median_ms": 2.0358840227127075, + "min_ms": 2.0259690284729004, + "p95_ms": 2.051292598247528 + }, + "train_peak_mib": 1187.0771484375 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.1543090045452118, + "median_ms": 0.11416950076818466, + "min_ms": 0.10904199630022049, + "p95_ms": 0.12074360288679602 + }, + "forward_peak_mib": 0.58154296875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3678638666192455e-07, + "relative_l2": 8.545758893188667e-09 + }, + "lse_vs_fp64": { + "max_abs": 8.168879297443254e-07, + "relative_l2": 5.8742217369343824e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 0.5347949862480164, + "median_ms": 0.49237099289894104, + "min_ms": 0.47746899724006653, + "p95_ms": 0.5201029449701309 + }, + "train_peak_mib": 2.900390625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.2101919949054718, + "median_ms": 0.1837330013513565, + "min_ms": 0.17950700223445892, + "p95_ms": 0.2077571451663971 + }, + "forward_peak_mib": 1.16455078125, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3678638666192455e-07, + "relative_l2": 8.545758893188667e-09 + }, + "lse_vs_fp64": { + "max_abs": 8.168879297443254e-07, + "relative_l2": 5.8742217369343824e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 1.0200350284576416, + "median_ms": 0.8074195086956024, + "min_ms": 0.8034729957580566, + "p95_ms": 1.0148069202899932 + }, + "train_peak_mib": 2.322265625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.36177799105644226, + "median_ms": 0.2994849979877472, + "min_ms": 0.29600000381469727, + "p95_ms": 0.34925700277090077 + }, + "forward_peak_mib": 0.00341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3678638666192455e-07, + "relative_l2": 8.545758893188667e-09 + }, + "lse_vs_fp64": { + "max_abs": 8.168879297443254e-07, + "relative_l2": 5.8742217369343824e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 0.49481499195098877, + "median_ms": 0.45453500747680664, + "min_ms": 0.4421769976615906, + "p95_ms": 0.4892268911004066 + }, + "train_peak_mib": 1.162109375 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 4.817515850067139, + "median_ms": 4.744045972824097, + "min_ms": 4.700301170349121, + "p95_ms": 4.814090299606324 + }, + "forward_peak_mib": 1.38671875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3678638666192455e-07, + "relative_l2": 8.545758893188667e-09 + }, + "lse_vs_fp64": { + "max_abs": 8.168879297443254e-07, + "relative_l2": 5.8742217369343824e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 6.1407999992370605, + "median_ms": 6.0981974601745605, + "min_ms": 6.018377780914307, + "p95_ms": 6.135987257957458 + }, + "train_peak_mib": 3.62939453125 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.4266740083694458, + "median_ms": 0.38156700134277344, + "min_ms": 0.37199199199676514, + "p95_ms": 0.39862620979547503 + }, + "forward_peak_mib": 0.005859375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3678638666192455e-07, + "relative_l2": 8.545758893188667e-09 + }, + "lse_vs_fp64": { + "max_abs": 1.3678638666192455e-07, + "relative_l2": 9.836276637086277e-09 + }, + "mismatch_vs_reference": 1, + "path": "ws2-triton", + "rel_l2_vs_reference": 6.857849803488313e-08, + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 0.8591560125350952, + "median_ms": 0.7713660001754761, + "min_ms": 0.7473899722099304, + "p95_ms": 0.8526125490665436 + }, + "train_peak_mib": 1.1650390625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.3813669979572296, + "median_ms": 0.34324949979782104, + "min_ms": 0.3323740065097809, + "p95_ms": 0.3569719552993775 + }, + "forward_peak_mib": 0.005859375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3678638666192455e-07, + "relative_l2": 8.545758893188667e-09 + }, + "lse_vs_fp64": { + "max_abs": 8.168879297443254e-07, + "relative_l2": 5.8742217369343824e-08 + }, + "mismatch_vs_reference": 0, + "path": "ws2-rocm", + "rel_l2_vs_reference": 0.0, + "repeat_bitwise": true, + "tokens": 1, + "train_fwd_bwd": { + "max_ms": 0.7179059982299805, + "median_ms": 0.6274919807910919, + "min_ms": 0.615032970905304, + "p95_ms": 0.6916952908039092 + }, + "train_peak_mib": 1.1650390625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.14429500699043274, + "median_ms": 0.11721399798989296, + "min_ms": 0.11284700036048889, + "p95_ms": 0.13398180454969408 + }, + "forward_peak_mib": 4.64013671875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 7.819555989385663e-07, + "relative_l2": 3.607215816086566e-08 + }, + "lse_vs_fp64": { + "max_abs": 3.7962424670467954e-07, + "relative_l2": 1.7082649106274754e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 0.3843710124492645, + "median_ms": 0.3623580038547516, + "min_ms": 0.3551679849624634, + "p95_ms": 0.38348765373229976 + }, + "train_peak_mib": 23.18603515625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.22168999910354614, + "median_ms": 0.18493500351905823, + "min_ms": 0.17918600142002106, + "p95_ms": 0.2052492991089821 + }, + "forward_peak_mib": 9.27880859375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 7.819555989385663e-07, + "relative_l2": 3.607215816086566e-08 + }, + "lse_vs_fp64": { + "max_abs": 3.7962424670467954e-07, + "relative_l2": 1.7082649106274754e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 0.6065809726715088, + "median_ms": 0.572409987449646, + "min_ms": 0.5615940093994141, + "p95_ms": 0.5971350371837616 + }, + "train_peak_mib": 18.55224609375 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.3464750051498413, + "median_ms": 0.3008265048265457, + "min_ms": 0.2979219853878021, + "p95_ms": 0.3128326579928399 + }, + "forward_peak_mib": 0.00341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 7.819555989385663e-07, + "relative_l2": 3.418514829490082e-08 + }, + "lse_vs_fp64": { + "max_abs": 6.202060394144837e-07, + "relative_l2": 2.1631422949575374e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 0.37323498725891113, + "median_ms": 0.35190249979496, + "min_ms": 0.3442310094833374, + "p95_ms": 0.3687642410397529 + }, + "train_peak_mib": 9.2763671875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 5.581049919128418, + "median_ms": 5.453480005264282, + "min_ms": 5.361602783203125, + "p95_ms": 5.529634046554565 + }, + "forward_peak_mib": 1.58642578125, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 7.819555989385663e-07, + "relative_l2": 3.607215816086566e-08 + }, + "lse_vs_fp64": { + "max_abs": 3.7962424670467954e-07, + "relative_l2": 1.7082649106274754e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 6.8699631690979, + "median_ms": 6.778928995132446, + "min_ms": 6.7133307456970215, + "p95_ms": 6.835946750640869 + }, + "train_peak_mib": 27.97216796875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.42835599184036255, + "median_ms": 0.390780508518219, + "min_ms": 0.38400998711586, + "p95_ms": 0.40350588709115986 + }, + "forward_peak_mib": 0.0107421875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 7.819555989385663e-07, + "relative_l2": 3.607215816086566e-08 + }, + "lse_vs_fp64": { + "max_abs": 3.7962424670467954e-07, + "relative_l2": 1.7082649106274754e-08 + }, + "mismatch_vs_reference": 0, + "path": "ws2-triton", + "rel_l2_vs_reference": 0.0, + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 1.200382947921753, + "median_ms": 1.0127645134925842, + "min_ms": 0.9833800196647644, + "p95_ms": 1.1908830463886262 + }, + "train_peak_mib": 9.279296875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.39678898453712463, + "median_ms": 0.34459200501441956, + "min_ms": 0.3358590006828308, + "p95_ms": 0.3572490349411964 + }, + "forward_peak_mib": 0.0107421875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 7.819555989385663e-07, + "relative_l2": 3.607215816086566e-08 + }, + "lse_vs_fp64": { + "max_abs": 3.7962424670467954e-07, + "relative_l2": 1.7082649106274754e-08 + }, + "mismatch_vs_reference": 0, + "path": "ws2-rocm", + "rel_l2_vs_reference": 0.0, + "repeat_bitwise": true, + "tokens": 8, + "train_fwd_bwd": { + "max_ms": 0.9615880250930786, + "median_ms": 0.7646960020065308, + "min_ms": 0.7506340146064758, + "p95_ms": 0.8864344060420988 + }, + "train_peak_mib": 9.279296875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.16805000603199005, + "median_ms": 0.15777450054883957, + "min_ms": 0.13067400455474854, + "p95_ms": 0.16165650114417077 + }, + "forward_peak_mib": 18.55615234375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3984818885148798e-06, + "relative_l2": 3.32171929410509e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1600633094133173e-06, + "relative_l2": 3.572127579605161e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 0.542165994644165, + "median_ms": 0.47486498951911926, + "min_ms": 0.46753400564193726, + "p95_ms": 0.5188208922743797 + }, + "train_peak_mib": 92.73681640625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.23334699869155884, + "median_ms": 0.19080349802970886, + "min_ms": 0.16196100413799286, + "p95_ms": 0.21979905590415003 + }, + "forward_peak_mib": 37.1044921875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3984818885148798e-06, + "relative_l2": 3.32171929410509e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1600633094133173e-06, + "relative_l2": 3.572127579605161e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 0.8767020106315613, + "median_ms": 0.8396669924259186, + "min_ms": 0.823743999004364, + "p95_ms": 0.8703750014305114 + }, + "train_peak_mib": 74.19873046875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.35901200771331787, + "median_ms": 0.3029704988002777, + "min_ms": 0.29780200123786926, + "p95_ms": 0.32449566274881364 + }, + "forward_peak_mib": 0.00341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 8.62301916981778e-07, + "relative_l2": 3.133177230624908e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.0085128359094142e-06, + "relative_l2": 3.172992466391412e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 0.5286250114440918, + "median_ms": 0.485180988907814, + "min_ms": 0.4702180027961731, + "p95_ms": 0.5186201572418212 + }, + "train_peak_mib": 37.0966796875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 5.558335781097412, + "median_ms": 5.474370956420898, + "min_ms": 5.384798049926758, + "p95_ms": 5.557803058624267 + }, + "forward_peak_mib": 2.23828125, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3984818885148798e-06, + "relative_l2": 3.32171929410509e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1600633094133173e-06, + "relative_l2": 3.373302046539746e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 6.803424835205078, + "median_ms": 6.748803615570068, + "min_ms": 6.687372207641602, + "p95_ms": 6.789111709594726 + }, + "train_peak_mib": 111.43310546875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.46192601323127747, + "median_ms": 0.39490650594234467, + "min_ms": 0.38837599754333496, + "p95_ms": 0.4057923927903176 + }, + "forward_peak_mib": 0.0341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3984818885148798e-06, + "relative_l2": 3.32171929410509e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1600633094133173e-06, + "relative_l2": 3.373302046539746e-08 + }, + "mismatch_vs_reference": 0, + "path": "ws2-triton", + "rel_l2_vs_reference": 0.0, + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 1.2516590356826782, + "median_ms": 1.0681869983673096, + "min_ms": 1.0127040147781372, + "p95_ms": 1.2389685928821563 + }, + "train_peak_mib": 37.099609375 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.3988319933414459, + "median_ms": 0.3438510000705719, + "min_ms": 0.33401599526405334, + "p95_ms": 0.35868310481309895 + }, + "forward_peak_mib": 0.0341796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3984818885148798e-06, + "relative_l2": 3.32171929410509e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1600633094133173e-06, + "relative_l2": 3.373302046539746e-08 + }, + "mismatch_vs_reference": 0, + "path": "ws2-rocm", + "rel_l2_vs_reference": 0.0, + "repeat_bitwise": true, + "tokens": 32, + "train_fwd_bwd": { + "max_ms": 0.8316360116004944, + "median_ms": 0.7883504927158356, + "min_ms": 0.7761529684066772, + "p95_ms": 0.8204049050807952 + }, + "train_peak_mib": 37.099609375 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.3082979917526245, + "median_ms": 0.2700809985399246, + "min_ms": 0.21864500641822815, + "p95_ms": 0.30304638892412183 + }, + "forward_peak_mib": 74.22021484375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3963434604136182e-06, + "relative_l2": 3.895801114451707e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2196775962536321e-06, + "relative_l2": 3.098601801905746e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 0.6854979991912842, + "median_ms": 0.5381385087966919, + "min_ms": 0.5144850015640259, + "p95_ms": 0.6242975533008573 + }, + "train_peak_mib": 370.93994140625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.35492798686027527, + "median_ms": 0.27833350002765656, + "min_ms": 0.2521750032901764, + "p95_ms": 0.3129512906074524 + }, + "forward_peak_mib": 148.41015625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3963434604136182e-06, + "relative_l2": 3.895801114451707e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2196775962536321e-06, + "relative_l2": 3.098601801905746e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 1.0999130010604858, + "median_ms": 0.8765020072460175, + "min_ms": 0.794219970703125, + "p95_ms": 1.0784614980220795 + }, + "train_peak_mib": 296.78515625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.3689480125904083, + "median_ms": 0.3239609897136688, + "min_ms": 0.3197549879550934, + "p95_ms": 0.3447429448366165 + }, + "forward_peak_mib": 0.00439453125, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.2674927774014577e-06, + "relative_l2": 4.306072396111767e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.3520369215314076e-06, + "relative_l2": 3.688548182438791e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 0.5482540130615234, + "median_ms": 0.48772451281547546, + "min_ms": 0.479312002658844, + "p95_ms": 0.5361944526433945 + }, + "train_peak_mib": 148.37841796875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 5.440080165863037, + "median_ms": 5.354552984237671, + "min_ms": 5.318378925323486, + "p95_ms": 5.409367537498474 + }, + "forward_peak_mib": 4.84619140625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3963434604136182e-06, + "relative_l2": 3.8421598950283275e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2196775962536321e-06, + "relative_l2": 2.9435109902863417e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 6.947038173675537, + "median_ms": 6.898946523666382, + "min_ms": 6.871885776519775, + "p95_ms": 6.9410529851913445 + }, + "train_peak_mib": 445.2783203125 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.4412960112094879, + "median_ms": 0.40109600126743317, + "min_ms": 0.39162200689315796, + "p95_ms": 0.41389514654874804 + }, + "forward_peak_mib": 0.12841796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3963434604136182e-06, + "relative_l2": 3.8277114241343485e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.125072301988439e-06, + "relative_l2": 2.914173983910747e-08 + }, + "mismatch_vs_reference": 5, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.4822693151362585e-08, + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 1.2369179725646973, + "median_ms": 1.0605154633522034, + "min_ms": 1.0286879539489746, + "p95_ms": 1.230464094877243 + }, + "train_peak_mib": 148.3818359375 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.39907199144363403, + "median_ms": 0.3528845012187958, + "min_ms": 0.34455201029777527, + "p95_ms": 0.37250998914241795 + }, + "forward_peak_mib": 0.12841796875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3963434604136182e-06, + "relative_l2": 3.8421598950283275e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2196775962536321e-06, + "relative_l2": 2.9702644822176375e-08 + }, + "mismatch_vs_reference": 1, + "path": "ws2-rocm", + "rel_l2_vs_reference": 6.051339139114186e-09, + "repeat_bitwise": true, + "tokens": 128, + "train_fwd_bwd": { + "max_ms": 1.038902997970581, + "median_ms": 0.8048155009746552, + "min_ms": 0.787289023399353, + "p95_ms": 1.0180824041366576 + }, + "train_peak_mib": 148.3818359375 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.8246650099754333, + "median_ms": 0.7760325074195862, + "min_ms": 0.7508350014686584, + "p95_ms": 0.8231051206588745 + }, + "forward_peak_mib": 296.880859375, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3502658759989572e-06, + "relative_l2": 3.7873685403982327e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1845403324883819e-06, + "relative_l2": 3.04023501035505e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 1.8172589540481567, + "median_ms": 1.7723724842071533, + "min_ms": 1.719514012336731, + "p95_ms": 1.8108774960041045 + }, + "train_peak_mib": 1483.7568359375 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.9324650168418884, + "median_ms": 0.8163724839687347, + "min_ms": 0.788811981678009, + "p95_ms": 0.844249901175499 + }, + "forward_peak_mib": 593.63916015625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3502658759989572e-06, + "relative_l2": 3.7873685403982327e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1845403324883819e-06, + "relative_l2": 3.04023501035505e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 2.1884100437164307, + "median_ms": 2.1150814294815063, + "min_ms": 2.0706350803375244, + "p95_ms": 2.188031530380249 + }, + "train_peak_mib": 1187.1376953125 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.40275800228118896, + "median_ms": 0.36109650135040283, + "min_ms": 0.35961398482322693, + "p95_ms": 0.3949186071753502 + }, + "forward_peak_mib": 0.01123046875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.8262701786397884e-06, + "relative_l2": 3.9251382435907614e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2793943788835804e-06, + "relative_l2": 3.244174275029511e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 0.5963649749755859, + "median_ms": 0.554082989692688, + "min_ms": 0.5410839915275574, + "p95_ms": 0.5855132400989532 + }, + "train_peak_mib": 593.5107421875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 5.619105815887451, + "median_ms": 5.537344932556152, + "min_ms": 5.441641807556152, + "p95_ms": 5.605291271209717 + }, + "forward_peak_mib": 15.46923828125, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3502658759989572e-06, + "relative_l2": 3.7865335567684935e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1845403324883819e-06, + "relative_l2": 3.002057868447724e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 7.834836959838867, + "median_ms": 7.765453577041626, + "min_ms": 7.677241802215576, + "p95_ms": 7.826418471336365 + }, + "train_peak_mib": 1780.6689453125 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.5381199717521667, + "median_ms": 0.47991301119327545, + "min_ms": 0.46781501173973083, + "p95_ms": 0.5134969592094422 + }, + "forward_peak_mib": 0.5107421875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3502658759989572e-06, + "relative_l2": 3.740717508776414e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1845403324883819e-06, + "relative_l2": 2.8888100679186516e-08 + }, + "mismatch_vs_reference": 46, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.6010327429749704e-08, + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 1.4293630123138428, + "median_ms": 1.2251800298690796, + "min_ms": 1.1932519674301147, + "p95_ms": 1.355453187227249 + }, + "train_peak_mib": 593.52001953125 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.4631679952144623, + "median_ms": 0.43050000071525574, + "min_ms": 0.42539098858833313, + "p95_ms": 0.45030499547719954 + }, + "forward_peak_mib": 0.5107421875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.3502658759989572e-06, + "relative_l2": 3.786183308895877e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1845403324883819e-06, + "relative_l2": 3.0413377637587183e-08 + }, + "mismatch_vs_reference": 21, + "path": "ws2-rocm", + "rel_l2_vs_reference": 1.090919925911375e-08, + "repeat_bitwise": true, + "tokens": 512, + "train_fwd_bwd": { + "max_ms": 1.1435790061950684, + "median_ms": 1.0652225017547607, + "min_ms": 1.0553679466247559, + "p95_ms": 1.1227039635181426 + }, + "train_peak_mib": 593.52001953125 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 2.212686061859131, + "median_ms": 1.968884527683258, + "min_ms": 1.815858006477356, + "p95_ms": 2.167170453071594 + }, + "forward_peak_mib": 1187.015625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.7474648927873204e-06, + "relative_l2": 3.8132201080718383e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2706277345841954e-06, + "relative_l2": 3.087959234643079e-08 + }, + "path": "native", + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 6.234379768371582, + "median_ms": 6.0170769691467285, + "min_ms": 5.560098171234131, + "p95_ms": 6.218876385688782 + }, + "train_peak_mib": 5935.0244140625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 2.2421700954437256, + "median_ms": 2.091405987739563, + "min_ms": 1.8722610473632812, + "p95_ms": 2.2312479257583617 + }, + "forward_peak_mib": 2374.103515625, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.7474648927873204e-06, + "relative_l2": 3.8132201080718383e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2706277345841954e-06, + "relative_l2": 3.087959234643079e-08 + }, + "path": "ws1-pytorch", + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 7.097422122955322, + "median_ms": 6.734924077987671, + "min_ms": 6.4440507888793945, + "p95_ms": 7.07494285106659 + }, + "train_peak_mib": 4748.0400390625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 0.7211099863052368, + "median_ms": 0.6786679923534393, + "min_ms": 0.6683520078659058, + "p95_ms": 0.7013974994421005 + }, + "forward_peak_mib": 0.044921875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.9093045153795174e-06, + "relative_l2": 4.012909908798742e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.549798071209807e-06, + "relative_l2": 3.362522490350142e-08 + }, + "path": "ws1-triton", + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 2.07828688621521, + "median_ms": 1.84471994638443, + "min_ms": 1.7994730472564697, + "p95_ms": 2.0619906425476073 + }, + "train_peak_mib": 2374.0400390625 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 6.019101142883301, + "median_ms": 5.579066514968872, + "min_ms": 5.509542942047119, + "p95_ms": 5.652730774879456 + }, + "forward_peak_mib": 57.96435546875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.7474648927873204e-06, + "relative_l2": 3.797537729833619e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2706277345841954e-06, + "relative_l2": 3.092466091771656e-08 + }, + "path": "ws2-reference", + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 12.192548751831055, + "median_ms": 12.105359554290771, + "min_ms": 12.014043807983398, + "p95_ms": 12.180759477615357 + }, + "train_peak_mib": 7122.23779296875 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 1.120864987373352, + "median_ms": 1.0601550340652466, + "min_ms": 1.0489180088043213, + "p95_ms": 1.0850918352603913 + }, + "forward_peak_mib": 2.04296875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.7105002712014539e-06, + "relative_l2": 3.6994881489459246e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.1972776015056752e-06, + "relative_l2": 2.985702593475806e-08 + }, + "mismatch_vs_reference": 144, + "path": "ws2-triton", + "rel_l2_vs_reference": 1.4661270414586789e-08, + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 2.9037530422210693, + "median_ms": 2.82489550113678, + "min_ms": 2.676975965499878, + "p95_ms": 2.8798858880996705 + }, + "train_peak_mib": 2374.0771484375 + }, + { + "batch_invariant": true, + "dtype": "fp32", + "forward": { + "max_ms": 1.0352979898452759, + "median_ms": 0.99720099568367, + "min_ms": 0.990149974822998, + "p95_ms": 1.035107958316803 + }, + "forward_peak_mib": 2.04296875, + "grad_finite": true, + "logp_vs_fp64": { + "max_abs": 1.7474648927873204e-06, + "relative_l2": 3.790029424792315e-08 + }, + "lse_vs_fp64": { + "max_abs": 1.2706277345841954e-06, + "relative_l2": 3.082868776934362e-08 + }, + "mismatch_vs_reference": 37, + "path": "ws2-rocm", + "rel_l2_vs_reference": 8.004053417128302e-09, + "repeat_bitwise": true, + "tokens": 2048, + "train_fwd_bwd": { + "max_ms": 3.8982300758361816, + "median_ms": 2.7085630893707275, + "min_ms": 2.5650100708007812, + "p95_ms": 3.3876757740974415 + }, + "train_peak_mib": 2374.0771484375 + } + ], + "paths": [ + "native", + "ws1-pytorch", + "ws1-triton", + "ws2-reference", + "ws2-triton", + "ws2-rocm" + ], + "validate_overhead": [ + { + "tokens": 1, + "validate_false": { + "max_ms": 0.38557198643684387, + "median_ms": 0.3474554866552353, + "min_ms": 0.34082698822021484, + "p95_ms": 0.36266274005174637 + }, + "validate_true": { + "max_ms": 0.5197709798812866, + "median_ms": 0.502346009016037, + "min_ms": 0.4909690022468567, + "p95_ms": 0.5150913119316101 + } + }, + { + "tokens": 8, + "validate_false": { + "max_ms": 0.3962689936161041, + "median_ms": 0.35082100331783295, + "min_ms": 0.34082600474357605, + "p95_ms": 0.375033649802208 + }, + "validate_true": { + "max_ms": 0.5299869775772095, + "median_ms": 0.5066129863262177, + "min_ms": 0.4946550130844116, + "p95_ms": 0.524735403060913 + } + }, + { + "tokens": 32, + "validate_false": { + "max_ms": 0.3971090018749237, + "median_ms": 0.34933900833129883, + "min_ms": 0.34283000230789185, + "p95_ms": 0.36491445600986483 + }, + "validate_true": { + "max_ms": 0.5520200133323669, + "median_ms": 0.5064724981784821, + "min_ms": 0.4932529926300049, + "p95_ms": 0.5254570484161377 + } + }, + { + "tokens": 128, + "validate_false": { + "max_ms": 0.38757601380348206, + "median_ms": 0.35624949634075165, + "min_ms": 0.3489989936351776, + "p95_ms": 0.37349510788917545 + }, + "validate_true": { + "max_ms": 0.5195310115814209, + "median_ms": 0.5095964968204498, + "min_ms": 0.4995020031929016, + "p95_ms": 0.518656051158905 + } + }, + { + "tokens": 512, + "validate_false": { + "max_ms": 0.42298799753189087, + "median_ms": 0.4019169956445694, + "min_ms": 0.3943859934806824, + "p95_ms": 0.4149196416139603 + }, + "validate_true": { + "max_ms": 0.5864710211753845, + "median_ms": 0.5539029836654663, + "min_ms": 0.5408830046653748, + "p95_ms": 0.573798930644989 + } + }, + { + "tokens": 2048, + "validate_false": { + "max_ms": 1.5755000114440918, + "median_ms": 0.9115940034389496, + "min_ms": 0.8969720005989075, + "p95_ms": 1.4987020194530487 + }, + "validate_true": { + "max_ms": 1.080083966255188, + "median_ms": 1.0558279752731323, + "min_ms": 1.0382219552993774, + "p95_ms": 1.0768872916698455 + } + } + ] + }, + "tile_stats_component": [ + { + "dtype": "bf16", + "hip_fp32_input": { + "max_ms": 0.020750999450683594, + "median_ms": 0.014341499656438828, + "min_ms": 0.011858000420033932, + "p95_ms": 0.019952050223946572 + }, + "hip_native_dtype_input": { + "max_ms": 0.014059999957680702, + "median_ms": 0.012438500300049782, + "min_ms": 0.011857000179588795, + "p95_ms": 0.013946950202807784 + }, + "hip_peak_mib": 0.0009765625, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.015689849853516, + "median_ms": 4.467295408248901, + "min_ms": 4.213459014892578, + "p95_ms": 4.686615490913391 + }, + "pytorch_loop_peak_mib": 0.08154296875, + "repeat_bitwise": true, + "sumexp_max_rel": 1.457063660836866e-07, + "sumexp_rel_l2": 5.770375928877574e-08, + "tokens": 1 + }, + { + "dtype": "bf16", + "hip_fp32_input": { + "max_ms": 0.058646999299526215, + "median_ms": 0.01962899975478649, + "min_ms": 0.015502999536693096, + "p95_ms": 0.03893354982137681 + }, + "hip_native_dtype_input": { + "max_ms": 0.015863999724388123, + "median_ms": 0.012879000045359135, + "min_ms": 0.011978000402450562, + "p95_ms": 0.01392315020784736 + }, + "hip_peak_mib": 0.00390625, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.4044270515441895, + "median_ms": 5.02221941947937, + "min_ms": 4.736033916473389, + "p95_ms": 5.30616557598114 + }, + "pytorch_loop_peak_mib": 0.28125, + "repeat_bitwise": true, + "sumexp_max_rel": 1.9626060066002537e-07, + "sumexp_rel_l2": 5.3845673395049095e-08, + "tokens": 8 + }, + { + "dtype": "bf16", + "hip_fp32_input": { + "max_ms": 0.02275400049984455, + "median_ms": 0.01882850006222725, + "min_ms": 0.01510199997574091, + "p95_ms": 0.021879050694406033 + }, + "hip_native_dtype_input": { + "max_ms": 0.017464999109506607, + "median_ms": 0.015683500096201897, + "min_ms": 0.015583000145852566, + "p95_ms": 0.01731394995003939 + }, + "hip_peak_mib": 0.015625, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.307322978973389, + "median_ms": 5.033536434173584, + "min_ms": 4.758028030395508, + "p95_ms": 5.305229234695434 + }, + "pytorch_loop_peak_mib": 0.93310546875, + "repeat_bitwise": true, + "sumexp_max_rel": 2.358648032441124e-07, + "sumexp_rel_l2": 5.510006778248904e-08, + "tokens": 32 + }, + { + "dtype": "bf16", + "hip_fp32_input": { + "max_ms": 0.043744999915361404, + "median_ms": 0.04050050117075443, + "min_ms": 0.03941800072789192, + "p95_ms": 0.043555000238120554 + }, + "hip_native_dtype_input": { + "max_ms": 0.03417100012302399, + "median_ms": 0.03240850009024143, + "min_ms": 0.03132700175046921, + "p95_ms": 0.03356109857559204 + }, + "hip_peak_mib": 0.0625, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.240784168243408, + "median_ms": 4.972766637802124, + "min_ms": 4.74508810043335, + "p95_ms": 5.230013751983643 + }, + "pytorch_loop_peak_mib": 3.54052734375, + "repeat_bitwise": true, + "sumexp_max_rel": 3.4588379094202537e-07, + "sumexp_rel_l2": 5.557826589309245e-08, + "tokens": 128 + }, + { + "dtype": "bf16", + "hip_fp32_input": { + "max_ms": 0.14028899371623993, + "median_ms": 0.12786950170993805, + "min_ms": 0.12033899873495102, + "p95_ms": 0.13621635362505916 + }, + "hip_native_dtype_input": { + "max_ms": 0.12490599602460861, + "median_ms": 0.1133279986679554, + "min_ms": 0.10651800036430359, + "p95_ms": 0.12429705001413822 + }, + "hip_peak_mib": 0.25, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.42746114730835, + "median_ms": 5.089759588241577, + "min_ms": 4.867790222167969, + "p95_ms": 5.367902684211731 + }, + "pytorch_loop_peak_mib": 14.16064453125, + "repeat_bitwise": true, + "sumexp_max_rel": 2.514437369427469e-07, + "sumexp_rel_l2": 5.398326236955934e-08, + "tokens": 512 + }, + { + "dtype": "bf16", + "hip_fp32_input": { + "max_ms": 0.5528209805488586, + "median_ms": 0.5171479880809784, + "min_ms": 0.456277996301651, + "p95_ms": 0.5297968149185182 + }, + "hip_native_dtype_input": { + "max_ms": 0.5951640009880066, + "median_ms": 0.46687400341033936, + "min_ms": 0.45643800497055054, + "p95_ms": 0.4840007096529008 + }, + "hip_peak_mib": 1.0, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.4985270500183105, + "median_ms": 5.0989930629730225, + "min_ms": 4.875161170959473, + "p95_ms": 5.356805205345155 + }, + "pytorch_loop_peak_mib": 56.642578125, + "repeat_bitwise": true, + "sumexp_max_rel": 3.6352432175590366e-07, + "sumexp_rel_l2": 5.424340271221068e-08, + "tokens": 2048 + }, + { + "dtype": "fp32", + "hip_fp32_input": { + "max_ms": 0.024315999820828438, + "median_ms": 0.014480999670922756, + "min_ms": 0.012218000367283821, + "p95_ms": 0.0229460995644331 + }, + "hip_native_dtype_input": { + "max_ms": 0.013380000367760658, + "median_ms": 0.012277999892830849, + "min_ms": 0.011216999962925911, + "p95_ms": 0.01326600038446486 + }, + "hip_peak_mib": 0.0009765625, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.198602199554443, + "median_ms": 4.477289915084839, + "min_ms": 4.219387054443359, + "p95_ms": 4.824810409545899 + }, + "pytorch_loop_peak_mib": 0.08154296875, + "repeat_bitwise": true, + "sumexp_max_rel": 1.5778303463775956e-07, + "sumexp_rel_l2": 4.05095450038561e-08, + "tokens": 1 + }, + { + "dtype": "fp32", + "hip_fp32_input": { + "max_ms": 0.023475000634789467, + "median_ms": 0.018106999807059765, + "min_ms": 0.012178000062704086, + "p95_ms": 0.021115199476480488 + }, + "hip_native_dtype_input": { + "max_ms": 0.01734600029885769, + "median_ms": 0.012338500004261732, + "min_ms": 0.011617000214755535, + "p95_ms": 0.015100200008600953 + }, + "hip_peak_mib": 0.00390625, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.431427001953125, + "median_ms": 5.071452856063843, + "min_ms": 4.788833141326904, + "p95_ms": 5.370270824432373 + }, + "pytorch_loop_peak_mib": 0.28125, + "repeat_bitwise": true, + "sumexp_max_rel": 2.2257805198933056e-07, + "sumexp_rel_l2": 5.588434866821146e-08, + "tokens": 8 + }, + { + "dtype": "fp32", + "hip_fp32_input": { + "max_ms": 0.0389380007982254, + "median_ms": 0.019168000668287277, + "min_ms": 0.015022000297904015, + "p95_ms": 0.023030249774456037 + }, + "hip_native_dtype_input": { + "max_ms": 0.01642500050365925, + "median_ms": 0.01522199995815754, + "min_ms": 0.014941999688744545, + "p95_ms": 0.016310999635607004 + }, + "hip_peak_mib": 0.015625, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.403584957122803, + "median_ms": 5.046475410461426, + "min_ms": 4.80033016204834, + "p95_ms": 5.3525518655776985 + }, + "pytorch_loop_peak_mib": 0.93310546875, + "repeat_bitwise": true, + "sumexp_max_rel": 2.452264595831366e-07, + "sumexp_rel_l2": 5.466176954859379e-08, + "tokens": 32 + }, + { + "dtype": "fp32", + "hip_fp32_input": { + "max_ms": 0.047189999371767044, + "median_ms": 0.041881999000906944, + "min_ms": 0.04086099937558174, + "p95_ms": 0.0458960996940732 + }, + "hip_native_dtype_input": { + "max_ms": 0.07026399672031403, + "median_ms": 0.04184200055897236, + "min_ms": 0.04010000079870224, + "p95_ms": 0.04438599962741138 + }, + "hip_peak_mib": 0.0625, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.275114059448242, + "median_ms": 5.003592491149902, + "min_ms": 4.78242301940918, + "p95_ms": 5.232073163986207 + }, + "pytorch_loop_peak_mib": 3.54052734375, + "repeat_bitwise": true, + "sumexp_max_rel": 2.3305682361751678e-07, + "sumexp_rel_l2": 5.403736616575617e-08, + "tokens": 128 + }, + { + "dtype": "fp32", + "hip_fp32_input": { + "max_ms": 0.16576600074768066, + "median_ms": 0.12873099744319916, + "min_ms": 0.12374399602413177, + "p95_ms": 0.1486413061618805 + }, + "hip_native_dtype_input": { + "max_ms": 0.13616199791431427, + "median_ms": 0.1302734985947609, + "min_ms": 0.12626700103282928, + "p95_ms": 0.1339931555092335 + }, + "hip_peak_mib": 0.25, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.358519077301025, + "median_ms": 5.106103420257568, + "min_ms": 4.881009101867676, + "p95_ms": 5.346721720695496 + }, + "pytorch_loop_peak_mib": 14.16064453125, + "repeat_bitwise": true, + "sumexp_max_rel": 2.839864805537218e-07, + "sumexp_rel_l2": 5.3810339225841154e-08, + "tokens": 512 + }, + { + "dtype": "fp32", + "hip_fp32_input": { + "max_ms": 0.686178982257843, + "median_ms": 0.5328314900398254, + "min_ms": 0.4858019948005676, + "p95_ms": 0.5776804357767106 + }, + "hip_native_dtype_input": { + "max_ms": 0.5658810138702393, + "median_ms": 0.53313148021698, + "min_ms": 0.5255410075187683, + "p95_ms": 0.5386321574449539 + }, + "hip_peak_mib": 1.0, + "max_bitwise": true, + "pytorch_loop": { + "max_ms": 5.419569969177246, + "median_ms": 5.102478504180908, + "min_ms": 4.886377811431885, + "p95_ms": 5.320242595672608 + }, + "pytorch_loop_peak_mib": 56.642578125, + "repeat_bitwise": true, + "sumexp_max_rel": 4.4772318119612464e-07, + "sumexp_rel_l2": 5.43710522076467e-08, + "tokens": 2048 + } + ] +} diff --git a/benchmarks/results/pr328_rocm_mi300x/single_gpu_grid.png b/benchmarks/results/pr328_rocm_mi300x/single_gpu_grid.png new file mode 100644 index 00000000..dae029bb Binary files /dev/null and b/benchmarks/results/pr328_rocm_mi300x/single_gpu_grid.png differ diff --git a/benchmarks/results/pr328_rocm_mi300x/single_gpu_grid_kernels.png b/benchmarks/results/pr328_rocm_mi300x/single_gpu_grid_kernels.png new file mode 100644 index 00000000..f2fc6142 Binary files /dev/null and b/benchmarks/results/pr328_rocm_mi300x/single_gpu_grid_kernels.png differ diff --git a/benchmarks/results/pr328_rocm_mi300x/single_gpu_latency.png b/benchmarks/results/pr328_rocm_mi300x/single_gpu_latency.png new file mode 100644 index 00000000..d05dbbd2 Binary files /dev/null and b/benchmarks/results/pr328_rocm_mi300x/single_gpu_latency.png differ diff --git a/benchmarks/results/pr328_rocm_mi300x/single_gpu_memory.png b/benchmarks/results/pr328_rocm_mi300x/single_gpu_memory.png new file mode 100644 index 00000000..28c61d00 Binary files /dev/null and b/benchmarks/results/pr328_rocm_mi300x/single_gpu_memory.png differ diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu index 72d3f874..e7d836a3 100644 --- a/csrc/cuda/distributed/deterministic_collective.cu +++ b/csrc/cuda/distributed/deterministic_collective.cu @@ -949,7 +949,10 @@ class DeterministicCollectiveState { has_staged_input_ = true; } - void all_reduce(torch::Tensor& output, cudaStream_t stream) { + void all_reduce( + torch::Tensor& output, + cudaStream_t stream, + bool allow_owner_path = true) { check_tensor(output, "output"); TORCH_CHECK(has_staged_input_, "stage() must be called before all_reduce()"); TORCH_CHECK( @@ -960,7 +963,10 @@ class DeterministicCollectiveState { "all-reduce output size must match the staged input size"); const int64_t element_count = output.numel(); - if (staged_owner_path_) { + // Graph-replayed calls must avoid the owner-push branch because it writes + // to remote IPC frames. Callers that use the fused ABI pass false here; + // direct staged collectives retain the eager owner optimization. + if (allow_owner_path && staged_owner_path_) { if (rank_ == 0) { // Logical rank 0 is the topology-favorable reader in both traced TP // engines. It evaluates the original fixed tree exactly once, then @@ -1022,6 +1028,58 @@ class DeterministicCollectiveState { return; } + // Small collectives are launch-bound. The fast kernel folds the peer + // stage wait, fixed-tree reduction, and completion publication into one + // launch while preserving the exact reduction order. + if (staged_fast_path_ && staged_bytes_ <= kSingleBlockFastPathMaxBytes) { + if (element_count > 0) { + switch (output.scalar_type()) { + case at::ScalarType::Float: + launch_all_reduce_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + element_count, + world_size_, + stream); + break; + case at::ScalarType::Half: + launch_all_reduce_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + element_count, + world_size_, + stream); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: + launch_all_reduce_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + element_count, + world_size_, + stream); + break; +#endif + default: + TORCH_CHECK( + false, + "deterministic all-reduce supports float32, float16, and bfloat16; got ", + output.scalar_type()); + } + AT_CUDA_CHECK(cudaGetLastError()); + } else { + publish_done(stream); + } + has_staged_input_ = false; + return; + } + wait_for_staged_peers(stream); if (element_count == 0) { publish_done(stream); @@ -1090,10 +1148,11 @@ class DeterministicCollectiveState { output.numel() == input.numel(), "all-reduce output size must match the input size"); - // Route both ABI entry points through the same staged protocol so the - // topology-aware owner reduction is always applied. + // Use the graph-safe staged protocol for every message size. The fused + // two-slot protocol is intentionally kept available in the extension for + // experiments, but is not safe to replay across vLLM's many graph shapes. stage(input, stream); - all_reduce(output, stream); + all_reduce(output, stream, /*allow_owner_path=*/false); } void all_gather_fused( diff --git a/csrc/deterministic_logp_kernel.cu b/csrc/deterministic_logp_kernel.cu index c238e11d..aa3a23c8 100644 --- a/csrc/deterministic_logp_kernel.cu +++ b/csrc/deterministic_logp_kernel.cu @@ -2,11 +2,13 @@ #include #include #include +#include #include #include namespace { +constexpr int kDeterministicLogpTinyBlockSize = 64; constexpr int kDeterministicLogpSmallBlockSize = 128; constexpr int kDeterministicLogpMediumBlockSize = 256; constexpr int kDeterministicLogpLargeBlockSize = 512; @@ -27,7 +29,8 @@ __device__ __forceinline__ T deterministic_logp_shfl_down_32(T value, unsigned i template struct DeterministicLogpBlockTraits { static_assert( - BlockSize == kDeterministicLogpSmallBlockSize || + BlockSize == kDeterministicLogpTinyBlockSize || + BlockSize == kDeterministicLogpSmallBlockSize || BlockSize == kDeterministicLogpMediumBlockSize || BlockSize == kDeterministicLogpLargeBlockSize, "deterministic logp reduction topology requires a supported fixed block size"); @@ -358,3 +361,355 @@ torch::Tensor deterministic_logp_forward_indexed_fp32( auto output = torch::zeros({logits.size(0)}, logits.options().dtype(at::ScalarType::Float)); return deterministic_logp_forward_indexed_out(logits, token_ids, row_indices, output); } + +namespace { + +// Tuned on sm_90 (H100) for the Qwen3 shape (V=151936 split into 64 tiles of +// 2374 columns). Two things dominate: 16-byte vector loads, and a block small +// enough that a tile splits into several balanced chunks. The previous fixed +// block of 256 left only ~1.2 chunks per tile, so most of each pass ran with +// idle lanes and the kernel stalled well short of memory bandwidth. +#ifndef DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_NARROW +#define DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_NARROW 64 // 1- and 2-byte inputs +#endif +#ifndef DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_WIDE +#define DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_WIDE 128 // 4-byte inputs +#endif +#ifndef DETERMINISTIC_LOGP_TILE_VECTOR_BYTES +#define DETERMINISTIC_LOGP_TILE_VECTOR_BYTES 16 +#endif + +template +struct DeterministicLogpPacked; +template <> +struct DeterministicLogpPacked<16> { + using type = int4; +}; +template <> +struct DeterministicLogpPacked<8> { + using type = int2; +}; +template <> +struct DeterministicLogpPacked<4> { + using type = int; +}; +template <> +struct DeterministicLogpPacked<2> { + using type = short; +}; + +template +__device__ __forceinline__ void deterministicLogpLoadVector( + const scalar_t* __restrict__ pointer, + scalar_t (&out)[Vec]) { + constexpr int Bytes = Vec * static_cast(sizeof(scalar_t)); + constexpr int PieceBytes = Bytes < 16 ? Bytes : 16; + constexpr int Pieces = Bytes / PieceBytes; + using packed_t = typename DeterministicLogpPacked::type; + if ((reinterpret_cast(pointer) % alignof(packed_t)) == 0) { + packed_t packed[Pieces]; +#pragma unroll + for (int piece = 0; piece < Pieces; ++piece) { + packed[piece] = reinterpret_cast(pointer)[piece]; + } + __builtin_memcpy(out, packed, Bytes); + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + out[i] = pointer[i]; + } + } +} + +template +__device__ __forceinline__ void deterministicLogpStoreVector( + scalar_t* __restrict__ pointer, + const scalar_t (&in)[Vec]) { + constexpr int Bytes = Vec * static_cast(sizeof(scalar_t)); + constexpr int PieceBytes = Bytes < 16 ? Bytes : 16; + constexpr int Pieces = Bytes / PieceBytes; + using packed_t = typename DeterministicLogpPacked::type; + if ((reinterpret_cast(pointer) % alignof(packed_t)) == 0) { + packed_t packed[Pieces]; + __builtin_memcpy(packed, in, Bytes); +#pragma unroll + for (int piece = 0; piece < Pieces; ++piece) { + reinterpret_cast(pointer)[piece] = packed[piece]; + } + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + pointer[i] = in[i]; + } + } +} + +// Per-row, per-tile FP32 (max, sumexp) partials over the real-vocabulary part of +// the tile. The reduction tree is fixed by (BlockSize, Vec) and the position +// inside the tile, never by the row count or the tile's place in the global +// order, so the partials stay batch-invariant and TP-replicated. +template +__global__ void __launch_bounds__(BlockSize) deterministic_logp_tile_stats_kernel( + const scalar_t* __restrict__ logits, + float* __restrict__ tile_max, + float* __restrict__ tile_sum, + int64_t rows, + int64_t local_vocab, + int64_t vocab_start, + int64_t real_vocab, + int64_t tile_size, + int64_t local_tiles) { + constexpr int Chunk = BlockSize * Vec; + const int64_t tile_index = static_cast(blockIdx.y); + const int64_t row = static_cast(blockIdx.x); + if (row >= rows || tile_index >= local_tiles) { + return; + } + const int64_t col_begin = tile_index * tile_size; + const int64_t col_end = min(col_begin + tile_size, local_vocab); + // Columns at or beyond the real vocabulary are padding; hoisting the bound + // keeps the per-element predicate out of the inner loop. + const int64_t real_end = + min(col_end, max(real_vocab - vocab_start, static_cast(0))); + const scalar_t* __restrict__ row_pointer = logits + row * local_vocab; + + float local_max = -std::numeric_limits::infinity(); + for (int64_t base = col_begin + static_cast(threadIdx.x) * Vec; base < real_end; + base += Chunk) { + scalar_t values[Vec]; + if (base + Vec <= real_end) { + deterministicLogpLoadVector(row_pointer + base, values); +#pragma unroll + for (int i = 0; i < Vec; ++i) { + local_max = fmaxf(local_max, static_cast(values[i])); + } + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + if (base + i < real_end) { + local_max = fmaxf(local_max, static_cast(row_pointer[base + i])); + } + } + } + } + const float max_value = deterministicBlockReduceMax(local_max); + __shared__ float row_max; + if (threadIdx.x == 0) row_max = max_value; + __syncthreads(); + const float tile_max_value = row_max; + + float sum_value = 0.0f; + if (isfinite(tile_max_value)) { + float local_sum = 0.0f; + for (int64_t base = col_begin + static_cast(threadIdx.x) * Vec; base < real_end; + base += Chunk) { + scalar_t values[Vec]; + if (base + Vec <= real_end) { + deterministicLogpLoadVector(row_pointer + base, values); +#pragma unroll + for (int i = 0; i < Vec; ++i) { + local_sum += expf(static_cast(values[i]) - tile_max_value); + } + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + if (base + i < real_end) { + local_sum += + expf(static_cast(row_pointer[base + i]) - tile_max_value); + } + } + } + } + sum_value = deterministicBlockReduceSum(local_sum); + } + if (threadIdx.x == 0) { + const int64_t output_index = row * local_tiles + tile_index; + tile_max[output_index] = tile_max_value; + tile_sum[output_index] = sum_value; + } +} + + +#ifndef DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE +#define DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE 256 +#endif + +// grad = coef_logp * (onehot - p) + coef_lse * p, with p = exp(z - lse) on finite +// rows, 0 on non-finite rows and on padding columns. Purely elementwise, so the +// result does not depend on the launch geometry or on the batch. +template +__global__ void __launch_bounds__(BlockSize) deterministic_logp_backward_kernel( + const scalar_t* __restrict__ logits, + const float* __restrict__ lse, + const float* __restrict__ coef_logp, + const float* __restrict__ coef_lse, + const int64_t* __restrict__ target_local, + scalar_t* __restrict__ grad, + int64_t rows, + int64_t local_vocab, + int64_t vocab_start, + int64_t real_vocab, + bool has_lse_grad) { + constexpr int Chunk = BlockSize * Vec; + const int64_t row = static_cast(blockIdx.x); + if (row >= rows) { + return; + } + const float row_lse = lse[row]; + const bool finite_row = isfinite(row_lse); + const float lse_safe = finite_row ? row_lse : 0.0f; + const float g_logp = coef_logp[row]; + const float g_lse = has_lse_grad ? coef_lse[row] : 0.0f; + const int64_t hit = target_local[row]; + const int64_t real_end = + min(local_vocab, max(real_vocab - vocab_start, static_cast(0))); + const scalar_t* __restrict__ row_in = logits + row * local_vocab; + scalar_t* __restrict__ row_out = grad + row * local_vocab; + const int64_t stride = static_cast(gridDim.y) * Chunk; + + for (int64_t base = static_cast(blockIdx.y) * Chunk + + static_cast(threadIdx.x) * Vec; + base < local_vocab; + base += stride) { + scalar_t values[Vec]; + scalar_t outputs[Vec]; + const bool full = base + Vec <= local_vocab; + if (full) { + deterministicLogpLoadVector(row_in + base, values); + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + values[i] = + (base + i < local_vocab) ? row_in[base + i] : static_cast(0.0f); + } + } +#pragma unroll + for (int i = 0; i < Vec; ++i) { + const int64_t col = base + i; + float value = 0.0f; + if (col < real_end) { + const float p = + finite_row ? expf(static_cast(values[i]) - lse_safe) : 0.0f; + const float onehot = (col == hit) ? 1.0f : 0.0f; + value = g_logp * (onehot - p); + if (has_lse_grad) { + value = value + g_lse * p; + } + } + outputs[i] = static_cast(value); + } + if (full) { + deterministicLogpStoreVector(row_out + base, outputs); + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + if (base + i < local_vocab) { + row_out[base + i] = outputs[i]; + } + } + } + } +} + +} // namespace + +std::vector deterministic_logp_tile_stats( + torch::Tensor logits, + int64_t vocab_start, + int64_t real_vocab, + int64_t num_tiles) { + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA/ROCm tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be 2D [tokens, local_vocab]"); + TORCH_CHECK(logits.scalar_type() == at::ScalarType::Half || + logits.scalar_type() == at::ScalarType::BFloat16 || + logits.scalar_type() == at::ScalarType::Float, + "logits must be float16, bfloat16, or float32"); + TORCH_CHECK(vocab_start >= 0 && real_vocab > 0 && num_tiles > 0, + "invalid vocabulary metadata"); + auto input = logits.contiguous(); + const int64_t rows = input.size(0); + const int64_t local_vocab = input.size(1); + TORCH_CHECK(local_vocab > 0 && local_vocab % num_tiles == 0, + "local_vocab must be divisible by num_tiles"); + const int64_t tile_size = local_vocab / num_tiles; + auto options = input.options().dtype(at::ScalarType::Float); + auto tile_max = torch::empty({rows, num_tiles}, options); + auto tile_sum = torch::empty({rows, num_tiles}, options); + const dim3 grid(static_cast(rows), static_cast(num_tiles), 1); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), + "deterministic_logp_tile_stats", ([&] { + // 16-byte loads per thread; narrow inputs also want the smaller block + // so a tile splits into enough chunks to keep every lane busy. + constexpr int Vec = DETERMINISTIC_LOGP_TILE_VECTOR_BYTES / sizeof(scalar_t); + constexpr int BlockSize = sizeof(scalar_t) >= 4 + ? DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_WIDE + : DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_NARROW; + deterministic_logp_tile_stats_kernel + <<>>( + input.data_ptr(), tile_max.data_ptr(), + tile_sum.data_ptr(), rows, local_vocab, vocab_start, + real_vocab, tile_size, num_tiles); + })); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {tile_max, tile_sum}; +} + +torch::Tensor deterministic_logp_backward( + torch::Tensor logits, + torch::Tensor lse, + torch::Tensor coef_logp, + torch::Tensor coef_lse, + torch::Tensor target_local, + int64_t vocab_start, + int64_t real_vocab, + bool has_lse_grad) { + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA/ROCm tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be 2D [tokens, local_vocab]"); + TORCH_CHECK(logits.scalar_type() == at::ScalarType::Half || + logits.scalar_type() == at::ScalarType::BFloat16 || + logits.scalar_type() == at::ScalarType::Float, + "logits must be float16, bfloat16, or float32"); + TORCH_CHECK(vocab_start >= 0 && real_vocab > 0, "invalid vocabulary metadata"); + auto input = logits.contiguous(); + const int64_t rows = input.size(0); + const int64_t local_vocab = input.size(1); + auto check_row_vector = + [&](const torch::Tensor& tensor, at::ScalarType dtype, const char* name) { + TORCH_CHECK(tensor.is_cuda() && tensor.device() == input.device(), name, + " must live on the logits device"); + TORCH_CHECK(tensor.scalar_type() == dtype, name, " has the wrong dtype"); + TORCH_CHECK(tensor.dim() == 1 && tensor.size(0) == rows, name, + " must have one entry per token"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + }; + check_row_vector(lse, at::ScalarType::Float, "lse"); + check_row_vector(coef_logp, at::ScalarType::Float, "coef_logp"); + check_row_vector(coef_lse, at::ScalarType::Float, "coef_lse"); + check_row_vector(target_local, at::ScalarType::Long, "target_local"); + auto grad = torch::empty_like(input); + if (rows == 0 || local_vocab == 0) { + return grad; + } + TORCH_CHECK(rows <= std::numeric_limits::max(), "row count exceeds CUDA grid-x limit"); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), + "deterministic_logp_backward", ([&] { + constexpr int Vec = DETERMINISTIC_LOGP_TILE_VECTOR_BYTES / sizeof(scalar_t); + constexpr int BlockSize = DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE; + constexpr int Chunk = BlockSize * Vec; + const int64_t chunks = (local_vocab + Chunk - 1) / Chunk; + const dim3 grid(static_cast(rows), + static_cast(std::min(chunks, 65535)), 1); + deterministic_logp_backward_kernel + <<>>( + input.data_ptr(), lse.data_ptr(), + coef_logp.data_ptr(), coef_lse.data_ptr(), + target_local.data_ptr(), grad.data_ptr(), rows, + local_vocab, vocab_start, real_vocab, has_lse_grad); + })); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return grad; +} diff --git a/csrc/hip/hip_deterministic_logp_kernel.hip b/csrc/hip/hip_deterministic_logp_kernel.hip new file mode 100644 index 00000000..61753f59 --- /dev/null +++ b/csrc/hip/hip_deterministic_logp_kernel.hip @@ -0,0 +1,424 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// ROCm-tuned kernels for the WS2 vocab-parallel logprob backend +// (rocm-vocab-parallel-logp-ws2). This file is compiled only for ROCm builds; +// the shared csrc/deterministic_logp_kernel.cu keeps the SM90-tuned CUDA path. +// +// Every per-tile reduction below uses an element-to-thread assignment and an +// accumulation order fixed by (BlockSize, Vec) alone, measured from the start +// of the tile and independent of the storage dtype. A tile therefore produces +// identical bits whichever TP rank owns it, wherever it sits inside that +// rank's local shard, and whether the shard is BF16, FP16, or FP32. The fixed +// global tile-order merge in Python relies on exactly that. +// +// Tuned on MI300X (gfx942): 8-element vector loads straight from the stored +// shard (no FP32 copy), a 128-thread block per (row, tile), and a fused +// elementwise backward. Build-time knobs are injected by setup.py from the +// environment variables of the same name. + +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef DETERMINISTIC_LOGP_TILE_BLOCK_SIZE +#define DETERMINISTIC_LOGP_TILE_BLOCK_SIZE 128 +#endif +#ifndef DETERMINISTIC_LOGP_TILE_VECTOR_ELEMENTS +#define DETERMINISTIC_LOGP_TILE_VECTOR_ELEMENTS 8 +#endif +#ifndef DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE +#define DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE 256 +#endif + +namespace { + +constexpr int kTileBlockSize = DETERMINISTIC_LOGP_TILE_BLOCK_SIZE; +constexpr int kTileVectorElements = DETERMINISTIC_LOGP_TILE_VECTOR_ELEMENTS; +constexpr int kBackwardBlockSize = DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE; +// Logical 32-lane warps on a 64-wide wavefront keep the reduction tree identical +// to the shared CUDA kernel's topology. +constexpr int kLogicalWarpSize = 32; + +static_assert(kTileBlockSize % kLogicalWarpSize == 0, "tile block must be warp-aligned"); +static_assert(kTileBlockSize <= 1024, "tile block too large"); +static_assert(kTileVectorElements >= 1, "vector width must be at least one element"); +static_assert(kBackwardBlockSize % kLogicalWarpSize == 0, "backward block must be warp-aligned"); + +template +__device__ __forceinline__ T shfl_down_32(T value, unsigned int delta) { + return __shfl_down(value, delta, kLogicalWarpSize); +} + +// Fixed-tree block max. Threads without a warp value contribute -inf, so a tile +// with no real-vocabulary column reduces to the (-inf, 0) identity partial. +template +__device__ __forceinline__ float block_reduce_max(float val) { + constexpr int WarpCount = BlockSize / kLogicalWarpSize; + __shared__ float shared[WarpCount]; + const int lane = threadIdx.x & (kLogicalWarpSize - 1); + const int wid = threadIdx.x / kLogicalWarpSize; +#pragma unroll + for (int offset = kLogicalWarpSize / 2; offset > 0; offset >>= 1) { + val = fmaxf(val, shfl_down_32(val, offset)); + } + if (lane == 0) { + shared[wid] = val; + } + __syncthreads(); + const bool has_warp_value = threadIdx.x < WarpCount; + val = has_warp_value ? shared[threadIdx.x] : -std::numeric_limits::infinity(); + if (wid == 0) { +#pragma unroll + for (int offset = kLogicalWarpSize / 2; offset > 0; offset >>= 1) { + val = fmaxf(val, shfl_down_32(val, offset)); + } + } + return val; +} + +template +__device__ __forceinline__ float block_reduce_sum(float val) { + constexpr int WarpCount = BlockSize / kLogicalWarpSize; + __shared__ float shared[WarpCount]; + const int lane = threadIdx.x & (kLogicalWarpSize - 1); + const int wid = threadIdx.x / kLogicalWarpSize; +#pragma unroll + for (int offset = kLogicalWarpSize / 2; offset > 0; offset >>= 1) { + val += shfl_down_32(val, offset); + } + if (lane == 0) { + shared[wid] = val; + } + __syncthreads(); + const bool has_warp_value = threadIdx.x < WarpCount; + val = has_warp_value ? shared[threadIdx.x] : 0.0f; + if (wid == 0) { +#pragma unroll + for (int offset = kLogicalWarpSize / 2; offset > 0; offset >>= 1) { + val += shfl_down_32(val, offset); + } + } + return val; +} + +template +struct Packed; +template <> +struct Packed<16> { + using type = uint4; +}; +template <> +struct Packed<8> { + using type = uint2; +}; +template <> +struct Packed<4> { + using type = unsigned int; +}; +template <> +struct Packed<2> { + using type = unsigned short; +}; + +// Vector load/store of Vec consecutive elements in 16-byte (or smaller) pieces. +// Alignment only selects the instruction; the values and their order are the +// same either way. +template +__device__ __forceinline__ void load_vector(const scalar_t* __restrict__ ptr, scalar_t (&out)[Vec]) { + constexpr int Bytes = Vec * static_cast(sizeof(scalar_t)); + constexpr int PieceBytes = Bytes < 16 ? Bytes : 16; + constexpr int Pieces = Bytes / PieceBytes; + static_assert(Bytes % PieceBytes == 0, "vector width must split into whole pieces"); + using packed_t = typename Packed::type; + if ((reinterpret_cast(ptr) % alignof(packed_t)) == 0) { + packed_t packed[Pieces]; +#pragma unroll + for (int piece = 0; piece < Pieces; ++piece) { + packed[piece] = reinterpret_cast(ptr)[piece]; + } + __builtin_memcpy(out, packed, Bytes); + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + out[i] = ptr[i]; + } + } +} + +template +__device__ __forceinline__ void store_vector(scalar_t* __restrict__ ptr, const scalar_t (&in)[Vec]) { + constexpr int Bytes = Vec * static_cast(sizeof(scalar_t)); + constexpr int PieceBytes = Bytes < 16 ? Bytes : 16; + constexpr int Pieces = Bytes / PieceBytes; + using packed_t = typename Packed::type; + if ((reinterpret_cast(ptr) % alignof(packed_t)) == 0) { + packed_t packed[Pieces]; + __builtin_memcpy(packed, in, Bytes); +#pragma unroll + for (int piece = 0; piece < Pieces; ++piece) { + reinterpret_cast(ptr)[piece] = packed[piece]; + } + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + ptr[i] = in[i]; + } + } +} + +// Per-row, per-tile FP32 (max, sumexp) partials over the real-vocabulary part +// of the tile. Two passes over a tile that stays L2-resident: max, then sumexp. +template +__global__ void __launch_bounds__(BlockSize) hip_tile_stats_kernel( + const scalar_t* __restrict__ logits, + float* __restrict__ tile_max, + float* __restrict__ tile_sum, + int64_t rows, + int64_t local_vocab, + int64_t vocab_start, + int64_t real_vocab, + int64_t tile_size, + int64_t local_tiles) { + constexpr int Chunk = BlockSize * Vec; + const int64_t tile_index = static_cast(blockIdx.y); + const int64_t row = static_cast(blockIdx.x); + if (row >= rows || tile_index >= local_tiles) { + return; + } + const int64_t col_begin = tile_index * tile_size; + const int64_t col_end = min(col_begin + tile_size, local_vocab); + // Columns at or beyond the real vocabulary are padding and contribute nothing. + const int64_t real_end = min(col_end, max(real_vocab - vocab_start, static_cast(0))); + const scalar_t* __restrict__ row_ptr = logits + row * local_vocab; + + float local_max = -std::numeric_limits::infinity(); + for (int64_t base = col_begin + static_cast(threadIdx.x) * Vec; base < real_end; + base += Chunk) { + scalar_t values[Vec]; + if (base + Vec <= real_end) { + load_vector(row_ptr + base, values); +#pragma unroll + for (int i = 0; i < Vec; ++i) { + local_max = fmaxf(local_max, static_cast(values[i])); + } + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + if (base + i < real_end) { + local_max = fmaxf(local_max, static_cast(row_ptr[base + i])); + } + } + } + } + const float max_value = block_reduce_max(local_max); + __shared__ float row_max; + if (threadIdx.x == 0) { + row_max = max_value; + } + __syncthreads(); + const float tile_max_value = row_max; + + float sum_value = 0.0f; + if (isfinite(tile_max_value)) { + float local_sum = 0.0f; + for (int64_t base = col_begin + static_cast(threadIdx.x) * Vec; base < real_end; + base += Chunk) { + scalar_t values[Vec]; + if (base + Vec <= real_end) { + load_vector(row_ptr + base, values); +#pragma unroll + for (int i = 0; i < Vec; ++i) { + local_sum += expf(static_cast(values[i]) - tile_max_value); + } + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + if (base + i < real_end) { + local_sum += expf(static_cast(row_ptr[base + i]) - tile_max_value); + } + } + } + } + sum_value = block_reduce_sum(local_sum); + } + if (threadIdx.x == 0) { + const int64_t output_index = row * local_tiles + tile_index; + tile_max[output_index] = tile_max_value; + tile_sum[output_index] = sum_value; + } +} + +// grad = g_logp * (onehot - p) + g_lse * p with p = exp(z - lse) on finite rows, +// p = 0 otherwise, and 0 on padding columns. Elementwise, so the result does not +// depend on the launch geometry. +template +__global__ void __launch_bounds__(BlockSize) hip_backward_kernel( + const scalar_t* __restrict__ logits, + const float* __restrict__ lse, + const float* __restrict__ coef_logp, + const float* __restrict__ coef_lse, + const int64_t* __restrict__ target_local, + scalar_t* __restrict__ grad, + int64_t rows, + int64_t local_vocab, + int64_t vocab_start, + int64_t real_vocab, + bool has_lse_grad) { + constexpr int Chunk = BlockSize * Vec; + const int64_t row = static_cast(blockIdx.x); + if (row >= rows) { + return; + } + const float row_lse = lse[row]; + const bool finite_row = isfinite(row_lse); + const float lse_safe = finite_row ? row_lse : 0.0f; + const float g_logp = coef_logp[row]; + const float g_lse = has_lse_grad ? coef_lse[row] : 0.0f; + const int64_t hit = target_local[row]; + const int64_t real_end = min(local_vocab, max(real_vocab - vocab_start, static_cast(0))); + const scalar_t* __restrict__ row_in = logits + row * local_vocab; + scalar_t* __restrict__ row_out = grad + row * local_vocab; + const int64_t stride = static_cast(gridDim.y) * Chunk; + + for (int64_t base = static_cast(blockIdx.y) * Chunk + + static_cast(threadIdx.x) * Vec; + base < local_vocab; base += stride) { + scalar_t values[Vec]; + scalar_t outputs[Vec]; + const bool full = base + Vec <= local_vocab; + if (full) { + load_vector(row_in + base, values); + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + values[i] = (base + i < local_vocab) ? row_in[base + i] : static_cast(0.0f); + } + } +#pragma unroll + for (int i = 0; i < Vec; ++i) { + const int64_t col = base + i; + float value = 0.0f; + if (col < real_end) { + const float p = finite_row ? expf(static_cast(values[i]) - lse_safe) : 0.0f; + const float onehot = (col == hit) ? 1.0f : 0.0f; + value = g_logp * (onehot - p); + if (has_lse_grad) { + value = value + g_lse * p; + } + } + outputs[i] = static_cast(value); + } + if (full) { + store_vector(row_out + base, outputs); + } else { +#pragma unroll + for (int i = 0; i < Vec; ++i) { + if (base + i < local_vocab) { + row_out[base + i] = outputs[i]; + } + } + } + } +} + +void check_logits(const torch::Tensor& logits) { + TORCH_CHECK(logits.is_cuda(), "logits must be a ROCm tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be 2D [tokens, local_vocab]"); + TORCH_CHECK(logits.scalar_type() == at::ScalarType::Half || + logits.scalar_type() == at::ScalarType::BFloat16 || + logits.scalar_type() == at::ScalarType::Float, + "logits must be float16, bfloat16, or float32"); +} + +} // namespace + +std::vector hip_deterministic_logp_tile_stats( + torch::Tensor logits, + int64_t vocab_start, + int64_t real_vocab, + int64_t num_tiles) { + check_logits(logits); + TORCH_CHECK(vocab_start >= 0 && real_vocab > 0 && num_tiles > 0, + "invalid vocabulary metadata"); + auto input = logits.contiguous(); + const int64_t rows = input.size(0); + const int64_t local_vocab = input.size(1); + TORCH_CHECK(local_vocab > 0 && local_vocab % num_tiles == 0, + "local_vocab must be divisible by num_tiles"); + TORCH_CHECK(num_tiles <= 65535, "num_tiles must fit the launch grid"); + const int64_t tile_size = local_vocab / num_tiles; + auto options = input.options().dtype(at::ScalarType::Float); + auto tile_max = torch::empty({rows, num_tiles}, options); + auto tile_sum = torch::empty({rows, num_tiles}, options); + if (rows == 0) { + return {tile_max, tile_sum}; + } + const dim3 grid(static_cast(rows), static_cast(num_tiles), 1); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), + "hip_deterministic_logp_tile_stats", ([&] { + hip_tile_stats_kernel + <<>>( + input.data_ptr(), tile_max.data_ptr(), + tile_sum.data_ptr(), rows, local_vocab, vocab_start, + real_vocab, tile_size, num_tiles); + })); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {tile_max, tile_sum}; +} + +torch::Tensor hip_deterministic_logp_backward( + torch::Tensor logits, + torch::Tensor lse, + torch::Tensor coef_logp, + torch::Tensor coef_lse, + torch::Tensor target_local, + int64_t vocab_start, + int64_t real_vocab, + bool has_lse_grad) { + check_logits(logits); + TORCH_CHECK(vocab_start >= 0 && real_vocab > 0, "invalid vocabulary metadata"); + auto input = logits.contiguous(); + const int64_t rows = input.size(0); + const int64_t local_vocab = input.size(1); + auto check_row_vector = [&](const torch::Tensor& tensor, at::ScalarType dtype, const char* name) { + TORCH_CHECK(tensor.is_cuda() && tensor.device() == input.device(), name, + " must live on the logits device"); + TORCH_CHECK(tensor.scalar_type() == dtype, name, " has the wrong dtype"); + TORCH_CHECK(tensor.dim() == 1 && tensor.size(0) == rows, name, + " must have one entry per token"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + }; + check_row_vector(lse, at::ScalarType::Float, "lse"); + check_row_vector(coef_logp, at::ScalarType::Float, "coef_logp"); + check_row_vector(coef_lse, at::ScalarType::Float, "coef_lse"); + check_row_vector(target_local, at::ScalarType::Long, "target_local"); + auto grad = torch::empty_like(input); + if (rows == 0 || local_vocab == 0) { + return grad; + } + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), + "hip_deterministic_logp_backward", ([&] { + constexpr int Chunk = kBackwardBlockSize * kTileVectorElements; + const int64_t chunks = (local_vocab + Chunk - 1) / Chunk; + const dim3 grid(static_cast(rows), + static_cast(std::min(chunks, 65535)), 1); + hip_backward_kernel + <<>>( + input.data_ptr(), lse.data_ptr(), + coef_logp.data_ptr(), coef_lse.data_ptr(), + target_local.data_ptr(), grad.data_ptr(), rows, + local_vocab, vocab_start, real_vocab, has_lse_grad); + })); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return grad; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index aecc30ed..2f0d0878 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -79,6 +79,11 @@ torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) +std::vector deterministic_logp_tile_stats( + torch::Tensor logits, + int64_t vocab_start, + int64_t real_vocab, + int64_t num_tiles); torch::Tensor fused_logp_forward_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output); torch::Tensor fused_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids); torch::Tensor fused_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); @@ -92,7 +97,40 @@ torch::Tensor deterministic_logp_forward_out(torch::Tensor logits, torch::Tensor torch::Tensor deterministic_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids); torch::Tensor deterministic_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices); +std::vector deterministic_logp_tile_stats( + torch::Tensor logits, + int64_t vocab_start, + int64_t real_vocab, + int64_t num_tiles); +torch::Tensor deterministic_logp_backward( + torch::Tensor logits, + torch::Tensor lse, + torch::Tensor coef_logp, + torch::Tensor coef_lse, + torch::Tensor target_local, + int64_t vocab_start, + int64_t real_vocab, + bool has_lse_grad); + +// ROCm-tuned WS2 vocab-parallel logprob kernels (csrc/hip/hip_deterministic_logp_kernel.hip). +#if defined(__HIPCC__) || defined(__HIP_PLATFORM_AMD__) +std::vector hip_deterministic_logp_tile_stats( + torch::Tensor logits, + int64_t vocab_start, + int64_t real_vocab, + int64_t num_tiles); +torch::Tensor hip_deterministic_logp_backward( + torch::Tensor logits, + torch::Tensor lse, + torch::Tensor coef_logp, + torch::Tensor coef_lse, + torch::Tensor target_local, + int64_t vocab_start, + int64_t real_vocab, + bool has_lse_grad); +#endif +#if !defined(__HIPCC__) && !defined(__HIP_PLATFORM_AMD__) // Single-node TP=8 deterministic collectives. std::tuple, int64_t> deterministic_collective_ipc_meta( torch::Tensor& tensor); @@ -110,6 +148,7 @@ void deterministic_collective_reduce_scatter(int64_t handle, torch::Tensor& outp void deterministic_collective_all_gather(int64_t handle, torch::Tensor& output); void deterministic_collective_all_gather_fused( int64_t handle, torch::Tensor& input, torch::Tensor& output); +#endif // Batch-Invariant Deterministic GEMM Declarations bool det_gemm_sm90_compiled(); @@ -424,7 +463,18 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("deterministic_logp_forward_fp32", &deterministic_logp_forward_fp32, "Batch-invariant deterministic logp fp32"); 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"); + m.def("deterministic_logp_tile_stats", &deterministic_logp_tile_stats, + "Deterministic local vocab-tile FP32 max and sumexp partials"); + m.def("deterministic_logp_backward", &deterministic_logp_backward, + "Fused vocab-parallel selected-logprob/LSE backward on the local shard"); +#if defined(__HIPCC__) || defined(__HIP_PLATFORM_AMD__) + m.def("hip_deterministic_logp_tile_stats", &hip_deterministic_logp_tile_stats, + "ROCm-tuned vocab-tile FP32 max and sumexp partials read from the stored shard"); + m.def("hip_deterministic_logp_backward", &hip_deterministic_logp_backward, + "ROCm fused vocab-parallel selected-logprob/LSE backward on the local shard"); +#endif +#if !defined(__HIPCC__) && !defined(__HIP_PLATFORM_AMD__) // Single-node TP=8 fixed-tree collectives. m.def( "deterministic_collective_ipc_meta", @@ -462,6 +512,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_collective_all_gather_fused", &deterministic_collective_all_gather_fused, "Run a fused small-message deterministic rank-ordered all-gather"); +#endif // Prefix-shared attention uses NVIDIA PTX and falls back to PyTorch SDPA on ROCm. #if !defined(USE_ROCM) diff --git a/docs/operators/grpo-loss.md b/docs/operators/grpo-loss.md index 07fcac16..fbfe5751 100644 --- a/docs/operators/grpo-loss.md +++ b/docs/operators/grpo-loss.md @@ -15,6 +15,11 @@ fused ratio/KL kernel (logits → ratio/KL via online softmax), and the group-no logits --[ratio_kl op]--> (ratio, kl) --[group adv + clipped surrogate]--> loss ``` +The backends above consume dense `[B, T, V]` logits and reduce with a plain masked +mean, so they are single-shard only. For vocab-parallel TP, or when the reduction +must be bitwise reproducible across parallel degrees, see [Tensor and Data +Parallel](#tensor-and-data-parallel) below. + ## Entry Point ```python from rl_engine.kernels.registry import kernel_registry @@ -74,6 +79,117 @@ op mirrors this using `NativeRatioKLOp`. Gradients flow into `policy_logits` only (`ref_logits` is frozen; `old_logps` is cached). +## Tensor and Data Parallel + +`DistributedGRPOLossOp` +(`rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py`) +**Every TP × DP degree produces bit-identical loss, per-sequence totals, and +gradients.** It is a reference backend on top of the deterministic +[vocab-parallel logprob](batch-invariant-logp.md#tensor-parallel); the backends +above stay the default single-GPU path. + +The objective is elementwise — ratio, clipping, reference KL and the group-relative +advantage all act per token or per sequence — so the only place the parallel layout +can change the answer is the final sum over tokens. + +1. Selected logprobs come from `VocabParallelLogprobOp`, so ratio and KL inherit + cross-TP bitwise equality for free. TP needs no further handling here: by the + time the objective sees a logprob, the vocabulary has already been reduced. +2. A DP rank owns each of its sequences **whole**. `sequence_shard_bounds` is a + contiguous `[0, num_sequences)` partition in DP-rank order, exactly like + `ShardingSpec.vocab_shard_bounds` for the vocabulary. +3. Two nested reductions, each with a contract-fixed extent: `padded_seq_len` + token slots → one sequence total (entirely local, since the rank owns the + whole sequence); `num_sequences` totals → the scalar numerator. +4. Only per-sequence totals cross a rank boundary. They travel by `all_gather` + and are concatenated in DP-rank order into a `[num_sequences]` vector. The + collective moves bytes and placement is an exact copy, so every degree + performs identical arithmetic on identical inputs. `all_reduce` is excluded on + purpose: its combine order follows the collective's topology, not the declared + sequence order. +5. Advantages are replicated, not merged — rewards are one scalar per sequence, + so every rank normalizes every group over the identical `[num_sequences]` + tensor and keeps its own slice. This is what lets an advantage group straddle + DP ranks with no extra machinery. +6. The normalizer divides by a **global** active-token count, gathered as + integers so it is exact at every degree. + +Step 3 is why the determinism argument is short: because no sequence's token sum +is ever split across ranks, there is no partial-sequence state to merge and no +alignment rule to get wrong. + +```python +sequence_shard_bounds = ((0, 8),) # DP=1 +sequence_shard_bounds = ((0, 4), (4, 8)) # DP=2 +``` + +### Context parallelism is out of scope + +`cp_world_size` must be 1; anything else raises `LossContractError` rather than +silently reducing over a partial batch. CP is an attention-level concern — attention +is the only op with a cross-token dependency — and this operator consumes logits, by +which point CP has already been resolved upstream. Supporting it here would mean +splitting a sequence's token sum across ranks and reducing over the very axis the +logprob contract declares a *non-merge* axis. In a CP job that reduction belongs to +the caller. `cp_rank`/`cp_world_size` are carried for provenance only, mirroring +`ShardingSpec`. + +### Normalizer semantics + +`TokenNormalizer` makes the GRPO normalizer ambiguity explicit. The modes differ by +more than a scale factor once sequence lengths vary, so the choice is part of the +numerical identity and travels in the contract fingerprint. + +| Mode | Denominator | Notes | +| --- | --- | --- | +| `global_active_tokens` (default) | active tokens in the global batch | Matches `NativeGRPOLossOp`'s masked mean at DP=1. Long sequences weigh more. | +| `per_sequence_then_mean` | per-sequence count, then mean over live sequences | Original GRPO form; sequences weigh equally. | +| `fixed_constant` | declared constant | Dr.GRPO form; independent of the mask. | + +Usage goes through the contract-aware entry point: + +```python +from rl_engine.kernels.registry import kernel_registry + +dispatched = kernel_registry.get_loss_op(contract) # GRPOLossContract from +result = dispatched.op.apply( # rl_engine.kernels.loss_contract + policy_local_logits, # [n, local_vocab] differentiable + action_ids, # [n] + old_logps, # [n] + rewards, # [local_num_sequences] + contract=contract, + ref_local_logits=ref_local_logits, # required when beta > 0 + tp_group=tp_group, # vocab-parallel subgroup + dp_group=dp_group, # data-parallel subgroup +) +result.loss.backward() # gradients flow into policy_local_logits only + +loss, policy_loss, kl = result # unpacks like the single-GPU op +``` + +A preflight `all_gather_object` runs on **both** the DP and TP axes before any other +collective. Neither alone suffices: the loss is replicated across TP, so two TP +siblings disagreeing on `beta` would compute different losses for one sharded model, +and the logprob path's own preflight cannot see that because `beta` is not part of +the logprob contract. Other loud failures, with no silent fallback: sequence bounds +that are non-contiguous or leave a gap; a nested logprob contract whose token count +disagrees with the owned sequences; a determinism scope stronger or weaker than the +logprob path's; and population-std advantages over a singleton group. + +### Comparing configurations + +Compare `per_sequence_policy` / `per_sequence_kl`, not the scalar loss. Measured on +this operator's test inputs, regrouping the token sum — a real change of the +summation tree — moves the per-sequence vector in 12 of 12 seeds but the scalar loss +in only 5 of 12: averaging `num_sequences` totals into one fp32 number rounds most +reorderings away. A drift report that compares only the scalar will under-report +reduction differences. `GRPOLossResult` exposes both, plus `advantages`, +`per_sequence_active_tokens`, and a `provenance` dict. + +Bitwise here means *across parallel degrees on one PyTorch build and GPU model*. It +rests on PyTorch's reduction kernels being deterministic for a fixed shape on a fixed +device; cross-version and cross-architecture equality is neither tested nor claimed. + ## Accuracy Reference semantics (`NativeGRPOLossOp`): @@ -92,6 +208,11 @@ loss = masked_mean(policy, completion_mask) + beta * masked_mean(kl, completion_ The Triton op matches the native reference (forward and backward) to `atol=1e-4`. +For `DistributedGRPOLossOp`, the reference-equals-policy identity is exact rather +than approximate: with `ref_logits is policy_logits` and `old_logps == logp_policy` +the ratio is `exp(0) = 1` bitwise, so the result is invariant to the clip epsilon, +and the KL is exactly `0.0`. + ## Performance Notes The cost is dominated by the [`ratio_kl`](ratio-kl.md) stage (the vocab-dimension work); @@ -118,18 +239,31 @@ online — the forward peak is independent of `V`. ## Tests ```bash -python -m pytest tests/test_grpo_loss.py -v +python -m pytest tests/test_grpo_loss.py -v # single-GPU backends +python -m pytest tests/test_grpo_loss_contract.py -v # TP/DP/CP contract, CPU only +python -m pytest tests/test_distributed_grpo_loss.py -v ``` -Covers the native reference (group advantages + loss from logits), Triton forward/backward -vs native, masked-token invariance, an SGD loss step, and registry dispatch. Triton tests -skip without CUDA + Triton. +`test_grpo_loss.py` covers the native reference (group advantages + loss from logits), +Triton forward/backward vs native, masked-token invariance, an SGD loss step, and +registry dispatch. Triton tests skip without CUDA + Triton. + +`test_distributed_grpo_loss.py` covers every `(TP, DP)` combination reachable with +four ranks — `tp2`, `tp4`, `dp2`, `dp4`, `tp2xdp2` — each compared bitwise against a +single-rank GPU baseline, plus the KL=0 identity, run-to-run stability, negative +controls, and the DP- and TP-axis preflight guards. Larger degrees run unchanged on a +bigger node. Multi-rank tests need one GPU per rank and skip otherwise; they are +deliberately small (1000-token vocabulary, 8 sequences of 32 slots) and each worker +caps itself with `torch.cuda.set_per_process_memory_fraction`, so the suite can share +a node with a running training job. ## Implementation Files - `rl_engine/kernels/ops/pytorch/loss/grpo_loss.py` - `rl_engine/kernels/ops/triton/loss/grpo_loss.py` - `rl_engine/kernels/ops/triton/loss/ratio_kl.py`, `rl_engine/kernels/ops/pytorch/loss/ratio_kl.py` -- `rl_engine/kernels/registry.py` -- `tests/test_grpo_loss.py` +- `rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py` +- `rl_engine/kernels/loss_contract.py` +- `rl_engine/kernels/registry.py` (`register_loss_backend`, `get_loss_op`) +- `tests/test_grpo_loss.py`, `tests/test_grpo_loss_contract.py`, `tests/test_distributed_grpo_loss.py` - `benchmarks/benchmark_ratio_kl.py` diff --git a/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh b/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh index 4b3a601c..eb1cf61b 100755 --- a/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh +++ b/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh @@ -12,14 +12,37 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then # vLLM hooks; otherwise it silently falls back to native attention/FFN and # loses the R/R performance path. Preserve explicit ablation selections. strict_linear_logp=0 + rollout_batch_size="" + n_samples_per_prompt="1" + explicit_vllm_execution_config=0 previous_arg="" for current_arg in "$@"; do if [[ "${previous_arg}" == "--linear-logp-provider-mode" && "${current_arg}" == "strict" ]]; then strict_linear_logp=1 - break + elif [[ "${previous_arg}" == "--rollout-batch-size" ]]; then + rollout_batch_size="${current_arg}" + elif [[ "${previous_arg}" == "--n-samples-per-prompt" ]]; then + n_samples_per_prompt="${current_arg}" fi + + case "${current_arg}" in + --linear-logp-provider-mode=strict) + strict_linear_logp=1 + ;; + --rollout-batch-size=*) + rollout_batch_size="${current_arg#*=}" + ;; + --n-samples-per-prompt=*) + n_samples_per_prompt="${current_arg#*=}" + ;; + --vllm-enforce-eager|--vllm-optimization-level|--vllm-optimization-level=*|--vllm-compilation-config|--vllm-compilation-config=*) + explicit_vllm_execution_config=1 + ;; + esac previous_arg="${current_arg}" done + + strict_cudagraph_args=() if [[ "${strict_linear_logp}" == "1" ]]; then export RL_KERNEL_VLLM_INTEGRATION="${RL_KERNEL_VLLM_INTEGRATION:-1}" export RL_KERNEL_CUDA_ONLY="${RL_KERNEL_CUDA_ONLY:-1}" @@ -27,6 +50,42 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then export RL_KERNEL_ATTENTION_CASE="${RL_KERNEL_ATTENTION_CASE:-R/R}" export RL_KERNEL_FFN_CASE="${RL_KERNEL_FFN_CASE:-R/R}" export RL_KERNEL_LOGP_CASE="${RL_KERNEL_LOGP_CASE:-R/R}" + + # Strict rollout kernels preserve their arithmetic order under CUDA Graph. + # Capturing the complete decode graph removes the per-layer host-launch + # gaps that otherwise dominate small decode batches. Capture every exact + # batch size: padding a strict custom kernel to a larger sparse graph can + # access invalid slots and, more importantly, changes the tested contract. + # Explicit vLLM execution flags always win so callers can opt out. + if [[ "${explicit_vllm_execution_config}" == "0" ]]; then + if [[ "${rollout_batch_size}" =~ ^[1-9][0-9]*$ && "${n_samples_per_prompt}" =~ ^[1-9][0-9]*$ ]]; then + max_capture_size=$((rollout_batch_size * n_samples_per_prompt)) + if [[ -n "${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE:-}" ]]; then + max_capture_size="${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE}" + fi + if ! [[ "${max_capture_size}" =~ ^[1-9][0-9]*$ ]]; then + echo "RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE must be a positive integer" >&2 + exit 2 + fi + + capture_sizes="[" + for ((batch_size = 1; batch_size <= max_capture_size; batch_size++)); do + if ((batch_size > 1)); then + capture_sizes+="," + fi + capture_sizes+="${batch_size}" + done + capture_sizes+="]" + compilation_config="{\"cudagraph_mode\":\"FULL_DECODE_ONLY\",\"cudagraph_capture_sizes\":${capture_sizes},\"max_cudagraph_capture_size\":${max_capture_size}}" + strict_cudagraph_args=( + --vllm-optimization-level 0 + --vllm-compilation-config "${compilation_config}" + ) + echo "[RL-Kernel] strict vLLM full-decode CUDA Graph capture sizes: ${capture_sizes}" >&2 + else + echo "[RL-Kernel] strict CUDA Graph disabled: rollout batch size is unavailable" >&2 + fi + fi fi exec "${REAL_PYTHON}" "$@" \ --seed 1234 \ @@ -35,7 +94,8 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then --vllm-attention-backend flash_attn \ --vllm-disable-custom-all-reduce \ --deterministic-mode \ - --accumulate-allreduce-grads-in-fp32 + --accumulate-allreduce-grads-in-fp32 \ + "${strict_cudagraph_args[@]}" fi exec "${REAL_PYTHON}" "$@" diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8e0e865a..af89dc3f 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -167,6 +167,28 @@ def deterministic_logp_forward_indexed_fp32( token_ids: torch.Tensor, row_indices: torch.Tensor, ) -> torch.Tensor: ... +def deterministic_logp_tile_stats( + logits: torch.Tensor, + vocab_start: int, + real_vocab: int, + num_tiles: int, +) -> list[torch.Tensor]: ... +def hip_deterministic_logp_tile_stats( + logits: torch.Tensor, + vocab_start: int, + real_vocab: int, + num_tiles: int, +) -> list[torch.Tensor]: ... +def hip_deterministic_logp_backward( + logits: torch.Tensor, + lse: torch.Tensor, + coef_logp: torch.Tensor, + coef_lse: torch.Tensor, + target_local: torch.Tensor, + vocab_start: int, + real_vocab: int, + has_lse_grad: bool, +) -> torch.Tensor: ... def deterministic_attention_forward( q: torch.Tensor, k: torch.Tensor, diff --git a/rl_engine/integrations/megatron_runtime.py b/rl_engine/integrations/megatron_runtime.py index e296e824..72ac8919 100644 --- a/rl_engine/integrations/megatron_runtime.py +++ b/rl_engine/integrations/megatron_runtime.py @@ -6,6 +6,7 @@ from __future__ import annotations import importlib +import os from collections.abc import Callable, Iterable from types import MethodType from typing import Any @@ -290,7 +291,48 @@ def initialize_from_environment(_args: Any = None) -> MegatronIntegration: from rl_engine.integrations.ablation import integration_plan_from_environment - return install_megatron_integration(integration_plan_from_environment()) + plan = integration_plan_from_environment() + integration = install_megatron_integration(plan) + if plan.implementation_for("attention", "training") is Implementation.RL_KERNEL: + _precompile_strict_attention_training(_args) + return integration + + +def _precompile_strict_attention_training(args: Any) -> None: + """Warm FA4 CuTe fwd/bwd JIT outside Vime's actor_train timer.""" + + if args is None or torch.version.hip is not None or not torch.cuda.is_available(): + return + if os.getenv("RL_KERNEL_PRECOMPILE_FA4", "1") == "0": + return + + from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core + + attention_heads = int(getattr(args, "num_attention_heads", 0) or 0) + query_groups = int(getattr(args, "num_query_groups", 0) or attention_heads) + tp_size = int(getattr(args, "tensor_model_parallel_size", 1) or 1) + head_dim = int( + getattr(args, "kv_channels", 0) + or (int(getattr(args, "hidden_size", 0) or 0) // attention_heads) + ) + if ( + attention_heads <= 0 + or query_groups <= 0 + or tp_size <= 0 + or attention_heads % tp_size + or query_groups % tp_size + or head_dim <= 0 + ): + return + + params_dtype = getattr(args, "params_dtype", None) + dtype = params_dtype if params_dtype in (torch.float16, torch.bfloat16) else torch.bfloat16 + StrictFlashAttention4Core.precompile_training( + q_heads=attention_heads // tp_size, + kv_heads=query_groups // tp_size, + head_dim=head_dim, + dtype=dtype, + ) __all__ = ["initialize_from_environment", "install_megatron_integration"] diff --git a/rl_engine/integrations/vime/logp.py b/rl_engine/integrations/vime/logp.py new file mode 100644 index 00000000..68f5a6f1 --- /dev/null +++ b/rl_engine/integrations/vime/logp.py @@ -0,0 +1,264 @@ +# 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 " + f"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} != " + f"{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/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index c4da28be..6351f0ab 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -12,7 +12,11 @@ import torch -from rl_engine.distributed.collectives import DETERMINISTIC_ALL_REDUCE_OP +from rl_engine.distributed.collectives import ( + DETERMINISTIC_ALL_REDUCE_OP, + collective_for_group, + deterministic_all_reduce_inplace, +) from rl_engine.integrations.ablation import ( Implementation, IntegrationPlan, @@ -40,6 +44,8 @@ _STRICT_RMS_NORM_INIT_MARKER = "__rl_kernel_original_strict_rms_norm_init__" _STRICT_ROTARY_INIT_MARKER = "__rl_kernel_original_strict_rotary_init__" _STRICT_LM_HEAD_LINEAR_PATCH_MARKER = "__rl_kernel_original_lm_head_linear_apply__" +_STRICT_O_PROJ_COLLECTIVE_MARKER = "__rl_kernel_o_proj_collective__" +_STRICT_ROW_PARALLEL_PATCH_MARKER = "__rl_kernel_original_row_parallel_forward__" _RLK_ATTENTION_BACKEND: type[Any] | None = None _RLK_ATTENTION_IMPL: type[Any] | None = None _RLK_ATTENTION_BUILDER: type[Any] | None = None @@ -367,6 +373,7 @@ def _patch_qwen3_strict_model( rotary_cls: type[Any] | None = None, linear_method_cls: type[Any] | None = None, attention_cls: type[Any] | None = None, + row_parallel_cls: type[Any] | None = None, det_gemm: Any | None = None, ) -> None: """Align vLLM's RMSNorm and Attention projections with Megatron.""" @@ -374,7 +381,7 @@ def _patch_qwen3_strict_model( production_classes = rms_norm_cls is None or linear_method_cls is None or attention_cls is None if production_classes: from vllm.model_executor.layers.layernorm import RMSNorm - from vllm.model_executor.layers.linear import UnquantizedLinearMethod + from vllm.model_executor.layers.linear import RowParallelLinear, UnquantizedLinearMethod from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding from vllm.model_executor.models.qwen3 import Qwen3Attention @@ -382,6 +389,7 @@ def _patch_qwen3_strict_model( rotary_cls = RotaryEmbedding linear_method_cls = UnquantizedLinearMethod attention_cls = Qwen3Attention + row_parallel_cls = RowParallelLinear assert rms_norm_cls is not None assert linear_method_cls is not None assert attention_cls is not None @@ -435,6 +443,58 @@ def strict_rms_norm_forward_cuda( eps=instance.variance_epsilon, ) + def bind_o_proj_collective(module: Any) -> None: + if int(getattr(module, "tp_size", 1)) <= 1: + return + from vllm.distributed.parallel_state import get_tp_group + + coordinator = get_tp_group() + group = getattr(coordinator, "device_group", coordinator) + collective = collective_for_group(group) + if collective is None: + raise RuntimeError("strict rollout o_proj requires an initialized TP process group") + setattr(module, _STRICT_O_PROJ_COLLECTIVE_MARKER, collective) + + if row_parallel_cls is not None and not hasattr( + row_parallel_cls, _STRICT_ROW_PARALLEL_PATCH_MARKER + ): + row_parallel_forward = row_parallel_cls.forward + + def strict_row_parallel_forward(instance: Any, input_: torch.Tensor) -> Any: + collective = getattr(instance, _STRICT_O_PROJ_COLLECTIVE_MARKER, None) + if collective is None: + return row_parallel_forward(instance, input_) + + if instance.input_is_parallel: + input_parallel = input_ + else: + from vllm.distributed import split_tensor_along_last_dim + + input_parallel = split_tensor_along_last_dim( + input_, num_partitions=instance.tp_size + )[instance.tp_rank].contiguous() + + assert instance.quant_method is not None + bias_ = None if (instance.tp_rank > 0 or instance.skip_bias_add) else instance.bias + output_parallel = instance.quant_method.apply(instance, input_parallel, bias_) + + if instance.reduce_results and instance.tp_size > 1: + deterministic_all_reduce_inplace( + output_parallel, + collective_handle=int(collective._handle), + ) + output = output_parallel + else: + output = output_parallel + + if not instance.return_bias: + return output + output_bias = instance.bias if instance.skip_bias_add else None + return output, output_bias + + setattr(row_parallel_cls, _STRICT_ROW_PARALLEL_PATCH_MARKER, row_parallel_forward) + row_parallel_cls.forward = strict_row_parallel_forward + if not hasattr(rms_norm_cls, _STRICT_RMS_NORM_INIT_MARKER): rms_norm_init = rms_norm_cls.__init__ @@ -463,6 +523,7 @@ def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: attention_init(instance, *args, **kwargs) setattr(instance.qkv_proj, _STRICT_PROJECTION_MARKER, "qkv") setattr(instance.o_proj, _STRICT_PROJECTION_MARKER, "o_proj") + bind_o_proj_collective(instance.o_proj) setattr(attention_cls, _STRICT_MODEL_PATCH_MARKER, attention_init) linear_method_cls.apply = deterministic_linear_apply diff --git a/rl_engine/kernels/loss_contract.py b/rl_engine/kernels/loss_contract.py new file mode 100644 index 00000000..78f2e280 --- /dev/null +++ b/rl_engine/kernels/loss_contract.py @@ -0,0 +1,784 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for deterministic GRPO loss on the TP-aware logprob path. + +The GRPO objective consumes selected-token log-probabilities and reduces them +to a scalar:: + + ratio_t = exp(logp_policy_t - old_logp_t) + surrogate = -min(ratio_t * adv_t, clip(ratio_t) * adv_t) + loss = normalize(sum_t surrogate_t) + beta * normalize(sum_t kl_t) +""" + +from __future__ import annotations + +import hashlib +import json +import math +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, TypeVar + +from rl_engine.kernels.logprob_contract import ( + IMPLEMENTATION_KINDS, + RESERVED_DISPATCH_POLICIES, + DeterminismScope, + DowncastPoint, + LogprobContract, + LogprobDType, +) + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +class LossContractError(ValueError): + """Raised when loss metadata does not describe a valid GRPO invocation.""" + + +class TokenNormalizer(str, Enum): + """Denominator applied to the summed per-token loss terms. + + ``global_active_tokens``: divide by the number of active tokens in the + *global* batch, gathered across every DP rank. Long sequences therefore + contribute proportionally more. This matches the existing single-GPU + ``NativeGRPOLossOp`` masked mean at DP=1. + + ``per_sequence_then_mean``: divide each sequence's sum by that sequence's + own active-token count, then average over sequences that hold at least one + active token. Sequences are weighted equally regardless of length. + + ``fixed_constant``: divide by a declared constant, independent of the mask. + Requires ``LossReductionSpec.fixed_normalizer_constant``. + + The three differ by more than a scale factor once sequence lengths vary, so + the choice is part of the numerical identity and travels in the fingerprint. + """ + + GLOBAL_ACTIVE_TOKENS = "global_active_tokens" + PER_SEQUENCE_THEN_MEAN = "per_sequence_then_mean" + FIXED_CONSTANT = "fixed_constant" + + +class SummationOrder(str, Enum): + """Fixed combine order for per-token partials. + + ``sequence_major_fixed``: within one sequence, tokens combine over the full + ``padded_seq_len`` extent on the single rank that owns the sequence; + sequences then combine in ascending global sequence index. Both extents are + contract-fixed, so the floating-point grouping is identical at every DP + degree. Only *which rank* computes a sequence changes. + """ + + SEQUENCE_MAJOR_FIXED = "sequence_major_fixed" + + +class LossTransport(str, Enum): + """Collectives move partial sums only; they never reduce numerically. + + ``all_reduce`` is excluded on purpose: NCCL's reduction order depends on + world size and topology, so it would silently regroup the per-sequence + combines and break the cross-DP guarantee ``SummationOrder`` provides. + """ + + ALL_GATHER = "all_gather" + + +class KLEstimator(str, Enum): + """Per-token reference-KL estimator. + + ``k3_unbiased``: ``exp(logp_ref - logp_policy) - (logp_ref - logp_policy) - 1``, + the non-negative low-variance estimator used by the existing ratio/KL op. + + ``k1_log_ratio``: ``logp_policy - logp_ref``, the plain log-ratio. + """ + + K3_UNBIASED = "k3_unbiased" + K1_LOG_RATIO = "k1_log_ratio" + + +class ClipMode(str, Enum): + MIN_OF_UNCLIPPED_AND_CLIPPED = "min_of_unclipped_and_clipped" + + +class AdvantageNormalizer(str, Enum): + """Group-relative reward normalization. + + ``mean_std_population``: ``(r - mean) / std`` with the population (biased) + standard deviation, the original GRPO form. + + ``mean_only``: ``r - mean``, the Dr.GRPO form that drops the std divisor to + avoid its length/difficulty bias. + """ + + MEAN_STD_POPULATION = "mean_std_population" + MEAN_ONLY = "mean_only" + + +class VarianceFormula(str, Enum): + """Only the two-pass form is conformant. + + ``E[x^2] - E[x]^2`` cancels catastrophically once rewards share a large + offset, and its error depends on group size, so it cannot support a bitwise + claim. The two-pass form subtracts the mean before squaring. + """ + + TWO_PASS = "two_pass" + + +class GroupReplication(str, Enum): + """How advantage groups are evaluated when they span DP ranks. + + ``replicated_all_gather``: per-sequence rewards are all-gathered and *every* + rank normalizes *every* group identically, then keeps its own slice. + Rewards are one scalar per sequence, so replicating the whole computation is + cheaper than making a partial-statistic merge bitwise-reproducible. + """ + + REPLICATED_ALL_GATHER = "replicated_all_gather" + + +class LossPlacement(str, Enum): + REPLICATED = "replicated" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise LossContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LossContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise LossContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _non_negative_float(value: Any, field: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise LossContractError(f"{field} must be a real number; got {value!r}") + value = float(value) + if not math.isfinite(value) or value < 0.0: + raise LossContractError(f"{field} must be finite and non-negative; got {value!r}") + return value + + +@dataclass(frozen=True) +class ClipSpec: + """Asymmetric PPO-style ratio clipping bounds. + + Separate low/high epsilons cover the "clip-higher" variants; passing the + same value twice recovers the symmetric ``[1-eps, 1+eps]`` form. + """ + + clip_eps_low: float = 0.2 + clip_eps_high: float = 0.2 + mode: ClipMode = ClipMode.MIN_OF_UNCLIPPED_AND_CLIPPED + + def __post_init__(self) -> None: + object.__setattr__(self, "mode", _enum_value(ClipMode, self.mode, "clip.mode")) + low = _non_negative_float(self.clip_eps_low, "clip_eps_low") + high = _non_negative_float(self.clip_eps_high, "clip_eps_high") + if low >= 1.0: + raise LossContractError( + f"clip_eps_low={low} must be smaller than 1.0; the lower clip bound " + "1 - clip_eps_low must stay positive" + ) + object.__setattr__(self, "clip_eps_low", low) + object.__setattr__(self, "clip_eps_high", high) + + @property + def lower_bound(self) -> float: + return 1.0 - self.clip_eps_low + + @property + def upper_bound(self) -> float: + return 1.0 + self.clip_eps_high + + +@dataclass(frozen=True) +class AdvantageSpec: + """Group-relative advantage normalization semantics.""" + + normalizer: AdvantageNormalizer = AdvantageNormalizer.MEAN_STD_POPULATION + variance: VarianceFormula = VarianceFormula.TWO_PASS + std_eps: float = 1e-6 + replication: GroupReplication = GroupReplication.REPLICATED_ALL_GATHER + + def __post_init__(self) -> None: + object.__setattr__( + self, + "normalizer", + _enum_value(AdvantageNormalizer, self.normalizer, "advantage.normalizer"), + ) + object.__setattr__( + self, "variance", _enum_value(VarianceFormula, self.variance, "advantage.variance") + ) + object.__setattr__( + self, + "replication", + _enum_value(GroupReplication, self.replication, "advantage.replication"), + ) + std_eps = _non_negative_float(self.std_eps, "advantage.std_eps") + if std_eps <= 0.0: + raise LossContractError( + f"advantage.std_eps={std_eps} must be strictly positive; it is the floor " + "that keeps a zero-variance group from dividing by zero" + ) + object.__setattr__(self, "std_eps", std_eps) + + +@dataclass(frozen=True) +class ObjectiveSpec: + """The GRPO objective itself: clipping, reference KL, advantage shaping.""" + + clip: ClipSpec = field(default_factory=ClipSpec) + advantage: AdvantageSpec = field(default_factory=AdvantageSpec) + kl_estimator: KLEstimator = KLEstimator.K3_UNBIASED + beta: float = 0.0 + + def __post_init__(self) -> None: + if not isinstance(self.clip, ClipSpec): + raise LossContractError("objective.clip must be a ClipSpec") + if not isinstance(self.advantage, AdvantageSpec): + raise LossContractError("objective.advantage must be an AdvantageSpec") + object.__setattr__( + self, + "kl_estimator", + _enum_value(KLEstimator, self.kl_estimator, "objective.kl_estimator"), + ) + object.__setattr__(self, "beta", _non_negative_float(self.beta, "objective.beta")) + + @property + def uses_reference_model(self) -> bool: + """Whether reference logits are required at all. + + ``beta == 0`` drops the KL term from the loss, so a backend may skip the + reference forward entirely. The KL is still *reported*, so a caller + that wants the diagnostic must supply reference logits regardless. + """ + + return self.beta > 0.0 + + +@dataclass(frozen=True) +class LossReductionSpec: + """Deterministic token/sequence summation and normalizer semantics. + + Per-token loss terms are accumulated in fp32, combined in the order given by + ``summation_order``, moved between ranks by ``transport`` (never reduced by + it), and divided by the denominator selected by ``token_normalizer``. The + scalar is downcast, if at all, only at ``downcast_at``. + + ``determinism_scope`` reuses the logprob scale. ``cross_tp_bitwise`` here + means the scalar loss and its gradient are bitwise-identical across TP and + DP degrees, given a fixed vocab tile count. + """ + + token_normalizer: TokenNormalizer = TokenNormalizer.GLOBAL_ACTIVE_TOKENS + summation_order: SummationOrder = SummationOrder.SEQUENCE_MAJOR_FIXED + acc_dtype: LogprobDType = LogprobDType.FP32 + transport: LossTransport = LossTransport.ALL_GATHER + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + determinism_scope: DeterminismScope = DeterminismScope.CROSS_TP_BITWISE + fixed_normalizer_constant: int | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "token_normalizer", + _enum_value(TokenNormalizer, self.token_normalizer, "token_normalizer"), + ) + object.__setattr__( + self, + "summation_order", + _enum_value(SummationOrder, self.summation_order, "summation_order"), + ) + object.__setattr__( + self, "acc_dtype", _enum_value(LogprobDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__( + self, "transport", _enum_value(LossTransport, self.transport, "transport") + ) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__( + self, + "determinism_scope", + _enum_value(DeterminismScope, self.determinism_scope, "determinism_scope"), + ) + if self.acc_dtype is not LogprobDType.FP32: + raise LossContractError(f"loss accumulation must be fp32; got {self.acc_dtype.value}") + + needs_constant = self.token_normalizer is TokenNormalizer.FIXED_CONSTANT + if needs_constant: + object.__setattr__( + self, + "fixed_normalizer_constant", + _positive_int(self.fixed_normalizer_constant, "fixed_normalizer_constant"), + ) + elif self.fixed_normalizer_constant is not None: + raise LossContractError( + "fixed_normalizer_constant is only meaningful for " + f"token_normalizer={TokenNormalizer.FIXED_CONSTANT.value}; got " + f"{self.token_normalizer.value}" + ) + + +@dataclass(frozen=True) +class LossOutputSpec: + """Output surface: fp32 scalars replicated across every DP and TP rank.""" + + loss_dtype: LogprobDType = LogprobDType.FP32 + placement: LossPlacement = LossPlacement.REPLICATED + + def __post_init__(self) -> None: + object.__setattr__( + self, "loss_dtype", _enum_value(LogprobDType, self.loss_dtype, "loss_dtype") + ) + object.__setattr__( + self, "placement", _enum_value(LossPlacement, self.placement, "placement") + ) + if self.loss_dtype is not LogprobDType.FP32: + raise LossContractError(f"loss output must be fp32; got {self.loss_dtype.value}") + + +@dataclass(frozen=True) +class LossShardingSpec: + """Which sequences of the global batch this DP rank owns. + + The global batch is ``num_sequences`` sequences of ``padded_seq_len`` token + slots each. ``sequence_shard_bounds`` lists every DP rank's half-open + ``[start, end)`` sequence range, indexed by rank; the full table is required + on every rank and must form a contiguous ``[0, num_sequences)`` partition, + exactly as ``ShardingSpec.vocab_shard_bounds`` does for the vocabulary. + """ + + dp_rank: int + dp_world_size: int + num_sequences: int + padded_seq_len: int + sequence_shard_bounds: tuple[tuple[int, int], ...] + group_boundaries: tuple[int, ...] + cp_rank: int = 0 + cp_world_size: int = 1 + + def __post_init__(self) -> None: + dp_world_size = _positive_int(self.dp_world_size, "dp_world_size") + dp_rank = _non_negative_int(self.dp_rank, "dp_rank") + if dp_rank >= dp_world_size: + raise LossContractError( + f"dp_rank={dp_rank} must be smaller than dp_world_size={dp_world_size}" + ) + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + _non_negative_int(self.cp_rank, "cp_rank") + if cp_world_size != 1: + raise LossContractError( + f"cp_world_size={cp_world_size} is unsupported: context parallelism splits a " + "sequence's tokens across ranks, so the loss reduction would span an axis this " + "contract does not model. Reduce across CP outside this operator." + ) + if self.cp_rank != 0: + raise LossContractError(f"cp_rank must be 0 when cp_world_size=1; got {self.cp_rank}") + + _positive_int(self.num_sequences, "num_sequences") + _positive_int(self.padded_seq_len, "padded_seq_len") + object.__setattr__( + self, + "sequence_shard_bounds", + self._validated_bounds(self.sequence_shard_bounds, dp_world_size, self.num_sequences), + ) + object.__setattr__( + self, + "group_boundaries", + self._validated_groups(self.group_boundaries, self.num_sequences), + ) + + @staticmethod + def _validated_bounds( + raw: Any, dp_world_size: int, num_sequences: int + ) -> tuple[tuple[int, int], ...]: + try: + bounds = tuple((pair[0], pair[1]) for pair in raw) + except (TypeError, IndexError, KeyError) as exc: + raise LossContractError( + "sequence_shard_bounds must be an iterable of (start, end) integer pairs" + ) from exc + if len(bounds) != dp_world_size: + raise LossContractError( + "sequence_shard_bounds must declare exactly one (start, end) pair per DP rank; " + f"got {len(bounds)} pairs for dp_world_size={dp_world_size}" + ) + expected_start = 0 + for rank, (start, end) in enumerate(bounds): + for name, value in ( + (f"sequence_shard_bounds[{rank}][0]", start), + (f"sequence_shard_bounds[{rank}][1]", end), + ): + if isinstance(value, bool) or not isinstance(value, int): + raise LossContractError(f"{name} must be an integer; got {value!r}") + if end <= start: + raise LossContractError( + f"sequence_shard_bounds[{rank}] must satisfy end > start; got [{start}, {end})" + ) + if start != expected_start: + raise LossContractError( + "sequence_shard_bounds must form a contiguous [0, num_sequences) partition " + f"in DP-rank order; rank {rank} starts at {start}, expected {expected_start}" + ) + expected_start = end + if expected_start != num_sequences: + raise LossContractError( + "sequence_shard_bounds must cover num_sequences exactly; covered " + f"{expected_start}, declared {num_sequences}" + ) + return bounds + + @staticmethod + def _validated_groups(raw: Any, num_sequences: int) -> tuple[int, ...]: + try: + offsets = tuple(raw) + except TypeError as exc: + raise LossContractError("group_boundaries must be an iterable of integers") from exc + if len(offsets) < 2: + raise LossContractError( + "group_boundaries must hold num_groups + 1 offsets, so at least 2 entries" + ) + for index, value in enumerate(offsets): + if isinstance(value, bool) or not isinstance(value, int): + raise LossContractError(f"group_boundaries[{index}] must be an integer") + if offsets[0] != 0 or offsets[-1] != num_sequences: + raise LossContractError( + "group_boundaries must start at 0 and end at " + f"num_sequences={num_sequences}; got [{offsets[0]}, ..., {offsets[-1]}]" + ) + for index in range(1, len(offsets)): + if offsets[index] <= offsets[index - 1]: + raise LossContractError( + "group_boundaries must be strictly increasing; " + f"offset {index} is {offsets[index]} after {offsets[index - 1]}" + ) + return offsets + + @property + def num_groups(self) -> int: + return len(self.group_boundaries) - 1 + + @property + def group_sizes(self) -> tuple[int, ...]: + return tuple( + self.group_boundaries[i + 1] - self.group_boundaries[i] for i in range(self.num_groups) + ) + + @property + def local_sequence_start(self) -> int: + return self.sequence_shard_bounds[self.dp_rank][0] + + @property + def local_sequence_end(self) -> int: + return self.sequence_shard_bounds[self.dp_rank][1] + + @property + def local_num_sequences(self) -> int: + start, end = self.sequence_shard_bounds[self.dp_rank] + return end - start + + @property + def local_num_token_slots(self) -> int: + """Token slots this rank holds -- the row count of its logprob call.""" + + return self.local_num_sequences * self.padded_seq_len + + +@dataclass(frozen=True) +class GRPOLossContract: + """Complete semantic request for one deterministic GRPO loss invocation.""" + + logprob: LogprobContract + sharding: LossShardingSpec + objective: ObjectiveSpec = field(default_factory=ObjectiveSpec) + reduction: LossReductionSpec = field(default_factory=LossReductionSpec) + output: LossOutputSpec = field(default_factory=LossOutputSpec) + + def __post_init__(self) -> None: + if not isinstance(self.logprob, LogprobContract): + raise LossContractError("logprob must be a LogprobContract") + if not isinstance(self.sharding, LossShardingSpec): + raise LossContractError("sharding must be a LossShardingSpec") + if not isinstance(self.objective, ObjectiveSpec): + raise LossContractError("objective must be an ObjectiveSpec") + if not isinstance(self.reduction, LossReductionSpec): + raise LossContractError("reduction must be a LossReductionSpec") + if not isinstance(self.output, LossOutputSpec): + raise LossContractError("output must be a LossOutputSpec") + + # The nested logprob contract describes this rank's own rows, so its + # token count must match the cells this rank owns. Catching the + # mismatch here turns a silent shape error deep inside the reduction + # into a contract failure at construction. + expected_tokens = self.sharding.local_num_token_slots + if self.logprob.mask.num_tokens != expected_tokens: + raise LossContractError( + f"logprob.mask.num_tokens={self.logprob.mask.num_tokens} must equal the " + f"{expected_tokens} token slots this rank owns " + f"({self.sharding.local_num_sequences} sequences x " + f"{self.sharding.padded_seq_len} slots per sequence)" + ) + if self.logprob.reduction.determinism_scope is not self.reduction.determinism_scope: + raise LossContractError( + "the loss cannot claim a stronger or weaker determinism scope than the " + f"logprob path it consumes; loss={self.reduction.determinism_scope.value}, " + f"logprob={self.logprob.reduction.determinism_scope.value}" + ) + if self.logprob.sharding.cp_world_size != self.sharding.cp_world_size: + raise LossContractError( + f"cp_world_size disagrees between the logprob contract " + f"({self.logprob.sharding.cp_world_size}) and the loss sharding " + f"({self.sharding.cp_world_size})" + ) + if self.objective.advantage.normalizer is AdvantageNormalizer.MEAN_STD_POPULATION: + # A singleton group has zero population variance, so its advantage + # would collapse to 0 and the sequence would contribute nothing. + # Reject it rather than silently training on a dead group. + small = [index for index, size in enumerate(self.sharding.group_sizes) if size < 2] + if small: + raise LossContractError( + f"advantage normalizer {AdvantageNormalizer.MEAN_STD_POPULATION.value} " + f"needs at least 2 sequences per group; groups {small} are smaller" + ) + + @property + def global_token_slots(self) -> int: + return self.sharding.num_sequences * self.sharding.padded_seq_len + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "dp_rank": self.sharding.dp_rank, + "dp_world_size": self.sharding.dp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "num_sequences": self.sharding.num_sequences, + "padded_seq_len": self.sharding.padded_seq_len, + "sequence_shard_bounds": [list(pair) for pair in self.sharding.sequence_shard_bounds], + "group_boundaries": list(self.sharding.group_boundaries), + "local_sequence_start": self.sharding.local_sequence_start, + "local_sequence_end": self.sharding.local_sequence_end, + } + objective = { + "clip_eps_low": self.objective.clip.clip_eps_low, + "clip_eps_high": self.objective.clip.clip_eps_high, + "clip_mode": self.objective.clip.mode.value, + "kl_estimator": self.objective.kl_estimator.value, + "beta": self.objective.beta, + "advantage_normalizer": self.objective.advantage.normalizer.value, + "advantage_variance": self.objective.advantage.variance.value, + "advantage_std_eps": self.objective.advantage.std_eps, + "advantage_replication": self.objective.advantage.replication.value, + } + reduction = { + "token_normalizer": self.reduction.token_normalizer.value, + "summation_order": self.reduction.summation_order.value, + "acc_dtype": self.reduction.acc_dtype.value, + "transport": self.reduction.transport.value, + "downcast_at": self.reduction.downcast_at.value, + "determinism_scope": self.reduction.determinism_scope.value, + "fixed_normalizer_constant": self.reduction.fixed_normalizer_constant, + "cp_is_merge_axis": False, + "dp_is_merge_axis": True, + } + return { + "semantic_operator": "grpo_loss", + "logprob": self.logprob.to_dict(), + "sharding": sharding, + "objective": objective, + "reduction": reduction, + "output": { + "loss_dtype": self.output.loss_dtype.value, + "placement": self.output.placement.value, + }, + } + + def cross_rank_fingerprint(self) -> str: + """Rank-independent identity for preflight agreement across all ranks. + + Drops every rank-local field so all ``dp_world_size x tp_world_size`` + ranks of one logical invocation agree. That includes the nested + logprob contract's mask: each DP rank holds a different slice of + sequences, so its ``num_tokens`` and mask digest legitimately differ + even though the invocation is the same one. The global token geometry + is still pinned, by ``num_sequences``/``padded_seq_len`` and by the full + ``sequence_shard_bounds`` table, which every rank declares identically. + """ + + payload = self.to_dict() + logprob = payload["logprob"] + logprob.pop("mask", None) + logprob["sharding"] = { + key: value + for key, value in logprob["sharding"].items() + if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} + } + payload["sharding"] = { + key: value + for key, value in payload["sharding"].items() + if key not in {"dp_rank", "cp_rank", "local_sequence_start", "local_sequence_end"} + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class LossBackendCapability: + """Capabilities a concrete loss backend declares to contract-aware dispatch.""" + + backend_id: str + token_normalizers: frozenset[TokenNormalizer] + kl_estimators: frozenset[KLEstimator] + advantage_normalizers: frozenset[AdvantageNormalizer] + determinism_scopes: frozenset[DeterminismScope] + dp_world_sizes: tuple[int, ...] | None = None + supports_variable_group_sizes: bool = False + supports_asymmetric_clip: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise LossContractError("backend_id must be a non-empty string") + if self.backend_id.strip().lower() in RESERVED_DISPATCH_POLICIES: + raise LossContractError( + f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" + ) + object.__setattr__(self, "backend_id", self.backend_id.strip()) + for name, enum_type in ( + ("token_normalizers", TokenNormalizer), + ("kl_estimators", KLEstimator), + ("advantage_normalizers", AdvantageNormalizer), + ("determinism_scopes", DeterminismScope), + ): + try: + values = frozenset( + _enum_value(enum_type, value, name) for value in getattr(self, name) + ) + except TypeError as exc: + raise LossContractError(f"{name} must be an iterable of enum values") from exc + if not values: + raise LossContractError(f"{name} must not be empty") + object.__setattr__(self, name, values) + object.__setattr__( + self, + "dp_world_sizes", + self._validated_world_sizes(self.dp_world_sizes, "dp_world_sizes"), + ) + for flag_name in ("supports_variable_group_sizes", "supports_asymmetric_clip"): + if not isinstance(getattr(self, flag_name), bool): + raise LossContractError(f"{flag_name} must be a bool") + if self.implementation_kind not in IMPLEMENTATION_KINDS: + raise LossContractError( + f"implementation_kind must be one of: {', '.join(sorted(IMPLEMENTATION_KINDS))}" + ) + + @staticmethod + def _validated_world_sizes( + values: tuple[int, ...] | None, field: str + ) -> tuple[int, ...] | None: + if values is None: + return None + try: + sizes = tuple(values) + except TypeError as exc: + raise LossContractError(f"{field} must be an iterable of integers") from exc + if not sizes: + raise LossContractError(f"{field} must not be empty; use None for unrestricted") + for value in sizes: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LossContractError(f"{field} must contain positive values; got {value!r}") + if len(set(sizes)) != len(sizes): + raise LossContractError(f"{field} must not contain duplicates") + return sizes + + def incompatibilities(self, contract: GRPOLossContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + reduction = contract.reduction + objective = contract.objective + sharding = contract.sharding + if reduction.token_normalizer not in self.token_normalizers: + reasons.append(f"token_normalizer={reduction.token_normalizer.value} is unsupported") + if objective.kl_estimator not in self.kl_estimators: + reasons.append(f"kl_estimator={objective.kl_estimator.value} is unsupported") + if objective.advantage.normalizer not in self.advantage_normalizers: + reasons.append( + f"advantage normalizer={objective.advantage.normalizer.value} is unsupported" + ) + if reduction.determinism_scope not in self.determinism_scopes: + reasons.append(f"determinism_scope={reduction.determinism_scope.value} is unsupported") + if self.dp_world_sizes is not None and sharding.dp_world_size not in self.dp_world_sizes: + reasons.append(f"DP={sharding.dp_world_size} is unsupported") + if len(set(sharding.group_sizes)) > 1 and not self.supports_variable_group_sizes: + reasons.append("variable advantage group sizes are unsupported") + if ( + objective.clip.clip_eps_low != objective.clip.clip_eps_high + and not self.supports_asymmetric_clip + ): + reasons.append("asymmetric ratio clipping is unsupported") + return tuple(reasons) + + def supports(self, contract: GRPOLossContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "token_normalizers": sorted(item.value for item in self.token_normalizers), + "kl_estimators": sorted(item.value for item in self.kl_estimators), + "advantage_normalizers": sorted(item.value for item in self.advantage_normalizers), + "determinism_scopes": sorted(item.value for item in self.determinism_scopes), + "dp_world_sizes": list(self.dp_world_sizes) if self.dp_world_sizes else None, + "supports_variable_group_sizes": self.supports_variable_group_sizes, + "supports_asymmetric_clip": self.supports_asymmetric_clip, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class LossDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: LossBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "AdvantageNormalizer", + "AdvantageSpec", + "ClipMode", + "ClipSpec", + "GRPOLossContract", + "GroupReplication", + "KLEstimator", + "LossBackendCapability", + "LossContractError", + "LossDispatchResult", + "LossOutputSpec", + "LossPlacement", + "LossReductionSpec", + "LossShardingSpec", + "LossTransport", + "ObjectiveSpec", + "SummationOrder", + "TokenNormalizer", + "VarianceFormula", +] diff --git a/rl_engine/kernels/ops/cuda/attention/flash_attn.py b/rl_engine/kernels/ops/cuda/attention/flash_attn.py index cec8d2fe..e57cdeb5 100644 --- a/rl_engine/kernels/ops/cuda/attention/flash_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/flash_attn.py @@ -109,6 +109,76 @@ def __init__( self._op = op self._paged_op = paged_op + @classmethod + def precompile_training( + cls, + *, + q_heads: int, + kv_heads: int, + head_dim: int, + device: torch.device | None = None, + dtype: torch.dtype = torch.bfloat16, + sequence_length: int = 512, + ) -> None: + """Compile strict training FA4 forward/backward before timing. + + FA4 CuTe compiles its forward kernel on the first invocation and its + deterministic backward kernel on the first autograd backward. A + one-step RL workload would otherwise charge both compilations to + Vime's ``actor_train`` timer. These isolated tensors exercise the + same Qwen-style GQA and multi-block shape class without touching model + tensors, RNG state, or distributed collectives. + """ + if torch.version.hip is not None: + raise StrictFlashAttentionUnavailable("FA4 CUDA precompile is unavailable on ROCm") + if not torch.cuda.is_available(): + raise StrictFlashAttentionUnavailable( + "FA4 CUDA precompile requires an available CUDA device" + ) + if dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict FA4 training precompile requires FP16 or BF16") + if q_heads <= 0 or kv_heads <= 0 or q_heads % kv_heads != 0: + raise ValueError("Q/KV head counts must be positive and GQA-compatible") + if head_dim <= 0 or sequence_length <= 0: + raise ValueError("head_dim and sequence_length must be positive") + + target = torch.device("cuda", torch.cuda.current_device()) if device is None else device + if target.type != "cuda": + raise ValueError("strict FA4 training precompile requires a CUDA device") + + core = cls() + # Zeros deliberately avoid consuming RNG state. The tensors are + # independent from the model and only populate FA4's process-local + # JIT caches. + q = torch.zeros( + (1, sequence_length, q_heads, head_dim), + dtype=dtype, + device=target, + requires_grad=True, + ) + k = torch.zeros( + (1, sequence_length, kv_heads, head_dim), + dtype=dtype, + device=target, + requires_grad=True, + ) + v = torch.zeros_like(k, requires_grad=True) + positions = torch.arange(sequence_length, dtype=torch.int64, device=target).expand(1, -1) + with torch.enable_grad(): + result = core.forward_bshd_with_lse( + q, + k, + v, + causal=True, + scale=head_dim**-0.5, + query_position_ids=positions, + key_position_ids=positions, + output_dtype=dtype, + ) + result.out.sum().backward() + torch.cuda.synchronize(target) + del result, q, k, v, positions, core + @staticmethod def _validate_api(op: Callable[..., Any]) -> None: try: 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 a1e09e91..2b06f2ef 100644 --- a/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py @@ -148,3 +148,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/cuda/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/cuda/loss/vocab_parallel_logp.py new file mode 100644 index 00000000..ca1739ba --- /dev/null +++ b/rl_engine/kernels/ops/cuda/loss/vocab_parallel_logp.py @@ -0,0 +1,164 @@ +"""CUDA WS2 vocab-parallel logprob backend. + +The reference operator owns the contract: TP transport, fixed global tile-order +merge, target ownership, masking, entropy, and autograd semantics. This backend +keeps all of that and replaces the large per-shard tile scan with the SM90-tuned +CUDA kernel in ``csrc/deterministic_logp_kernel.cu``: + +* ``deterministic_logp_tile_stats`` computes the per-row, per-tile FP32 + ``(max, sumexp)`` partials straight from the FP16/BF16/FP32 shard. It filters + padding columns itself and reduces every tile with a fixed warp/block tree, so + the partials do not depend on the tile's position in the global order. + +* ``deterministic_logp_backward`` produces ``grad_logits`` for the selected + logprob and LSE outputs in one fused pass from the saved input shard. + +Both run through the shared :func:`apply_with_kernels` path, which reads the +stored BF16/FP16/FP32 shard directly instead of materializing an FP32 copy of +it. ``apply_with_entropy`` keeps the shared autograd path (with the CUDA tile +kernel) because the entropy gradient needs the full probability tensor anyway. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.kernels.logprob_contract import LogprobContract +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + DEFAULT_NUM_VOCAB_TILES, + VocabParallelLogprobOp, + apply_with_kernels, +) + +BACKEND_ID = "cuda-vocab-parallel-logp-ws2" + + +def native_tile_stats_available() -> bool: + """True when a CUDA-built ``rl_engine._C`` exposes both fused kernels.""" + + if torch.version.hip is not None: + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + except ImportError: + return False + return bool( + _EXT_AVAILABLE + and hasattr(_C, "deterministic_logp_tile_stats") + and hasattr(_C, "deterministic_logp_backward") + ) + + +def _native_cuda_tile_stats( + z_masked: torch.Tensor, + tile: int, + *, + vocab_start: int, + real_vocab_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """CUDA counterpart of ``_native_rocm_tile_stats``. + + TP transport and the global merge deliberately stay in Python so the issue + #241 reduction contract is identical across CUDA, ROCm, and the reference. + """ + + if torch.version.hip is not None or not z_masked.is_cuda: + raise RuntimeError("CUDA native tile stats require a CUDA tensor on a CUDA build") + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not hasattr(_C, "deterministic_logp_tile_stats"): + raise RuntimeError( + "CUDA vocab-parallel logprob native extension is unavailable; " + "build rl_engine._C with a CUDA toolchain" + ) + local_tiles = z_masked.shape[1] // tile + if local_tiles <= 0: + raise RuntimeError("CUDA native tile stats require at least one local vocab tile") + return tuple( + tensor.contiguous() + for tensor in _C.deterministic_logp_tile_stats( + z_masked, + int(vocab_start), + int(real_vocab_size), + int(local_tiles), + ) + ) + except (ImportError, AttributeError) as exc: + raise RuntimeError("CUDA vocab-parallel logprob native extension is unavailable") from exc + + +class _CudaKernels: + """``VocabParallelLogprobKernels`` over the CUDA extension symbols.""" + + @staticmethod + def tile_stats(shard, vocab_start, real_vocab, num_tiles): + from rl_engine.kernels.ops.base import _C + + tile_max, tile_sum = _C.deterministic_logp_tile_stats( + shard, vocab_start, real_vocab, num_tiles + ) + return tile_max, tile_sum + + @staticmethod + def backward( + shard, lse, coef_logp, coef_lse, target_local, vocab_start, real_vocab, has_lse_grad + ): + from rl_engine.kernels.ops.base import _C + + return _C.deterministic_logp_backward( + shard, lse, coef_logp, coef_lse, target_local, vocab_start, real_vocab, has_lse_grad + ) + + +class CudaVocabParallelLogprobOp(VocabParallelLogprobOp): + """Contract-preserving CUDA implementation with fused local reductions.""" + + op_class = "logprob" + is_batch_invariant = True + backend_id = BACKEND_ID + # Used by apply_with_entropy, which keeps the shared autograd path. + # staticmethod: the base reads ``self.use_native_tile_stats`` and calls it + # with the _native_rocm_tile_stats signature, so it must not bind ``self``. + use_native_tile_stats = staticmethod(_native_cuda_tile_stats) + + 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, + deterministic: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not deterministic: + return super().apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + deterministic=False, + ) + if not native_tile_stats_available(): + raise RuntimeError( + f"{BACKEND_ID} requires rl_engine._C built with a CUDA toolchain " + "(deterministic_logp_* symbols are missing); it does not fall back" + ) + return apply_with_kernels( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + kernels=_CudaKernels, + ) + + +__all__ = ["BACKEND_ID", "CudaVocabParallelLogprobOp", "native_tile_stats_available"] 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/distributed_grpo_loss.py b/rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py new file mode 100644 index 00000000..6add5721 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py @@ -0,0 +1,449 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic DP-aware GRPO loss.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Iterator + +import torch + +from rl_engine.kernels.loss_contract import ( + AdvantageNormalizer, + GRPOLossContract, + KLEstimator, + LossContractError, + TokenNormalizer, +) +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + DEFAULT_NUM_VOCAB_TILES, + VocabParallelLogprobOp, +) + +BACKEND_ID = "pytorch-distributed-grpo-loss-ws2" + +# Channel layout of the packed per-sequence tensor moved by the single fp32 +# all-gather; counts travel separately as integers so they stay exact. +_CH_POLICY = 0 +_CH_KL = 1 + + +@dataclass(frozen=True) +class GRPOLossResult: + """Scalar loss terms, per-sequence diagnostics, and bound provenance. + + Unpacks as ``loss, policy_loss, kl`` so it can stand in for the tuple the + single-GPU GRPO ops return. + + The ``per_sequence_*`` vectors are the reduction's last mesh-independent + intermediate, detached and replicated on every rank. They are the right + surface for comparing two configurations: the scalar loss averages + ``num_sequences`` totals into 24 mantissa bits and routinely rounds a real + reordering away, whereas the per-sequence vector preserves it. They also + give a drift report somewhere to point when one sequence is responsible. + """ + + loss: torch.Tensor + policy_loss: torch.Tensor + kl: torch.Tensor + advantages: torch.Tensor + per_sequence_policy: torch.Tensor + per_sequence_kl: torch.Tensor + per_sequence_active_tokens: torch.Tensor + provenance: dict[str, Any] = field(default_factory=dict) + + def __iter__(self) -> Iterator[torch.Tensor]: + yield from (self.loss, self.policy_loss, self.kl) + + +def _require_distributed_initialized(): + import torch.distributed as dist + + if not dist.is_available(): + raise LossContractError("distributed GRPO loss requires torch.distributed.") + if not dist.is_initialized(): + raise LossContractError( + "distributed GRPO loss requires an initialized process group when " + "the contract declares dp_world_size > 1." + ) + return dist + + +def _validate_invocation( + policy_local_logits: torch.Tensor, + ref_local_logits: torch.Tensor | None, + action_ids: torch.Tensor, + old_logps: torch.Tensor, + rewards: torch.Tensor, + contract: GRPOLossContract, + dp_group: Any, +) -> None: + sharding = contract.sharding + num_rows = sharding.local_num_token_slots + if policy_local_logits.dim() != 2: + raise LossContractError( + "policy_local_logits must be 2-D [num_tokens, local_vocab]; got " + f"{policy_local_logits.dim()}-D" + ) + if policy_local_logits.shape[0] != num_rows: + raise LossContractError( + f"policy_local_logits has {policy_local_logits.shape[0]} rows but this rank owns " + f"{num_rows} token slots" + ) + if ref_local_logits is not None and ref_local_logits.shape != policy_local_logits.shape: + raise LossContractError( + f"ref_local_logits shape {tuple(ref_local_logits.shape)} must match " + f"policy_local_logits shape {tuple(policy_local_logits.shape)}" + ) + if ref_local_logits is None and contract.objective.uses_reference_model: + raise LossContractError( + f"objective.beta={contract.objective.beta} puts the reference KL in the loss, " + "so ref_local_logits is required" + ) + for name, tensor in (("action_ids", action_ids), ("old_logps", old_logps)): + if tensor.dim() != 1 or tensor.shape[0] != num_rows: + raise LossContractError( + f"{name} must be 1-D with one entry per owned token slot; got shape " + f"{tuple(tensor.shape)} for {num_rows} slots" + ) + if rewards.dim() != 1 or rewards.shape[0] != sharding.local_num_sequences: + raise LossContractError( + "rewards must be 1-D with one entry per sequence this rank owns; got shape " + f"{tuple(rewards.shape)} for {sharding.local_num_sequences} sequences" + ) + + if sharding.dp_world_size > 1: + dist = _require_distributed_initialized() + group_world = dist.get_world_size(group=dp_group) + group_rank = dist.get_rank(group=dp_group) + if group_world != sharding.dp_world_size: + raise LossContractError( + f"dp_group world size {group_world} does not match the contract " + f"dp_world_size={sharding.dp_world_size}; pass the DP subgroup, " + "not the global group" + ) + if group_rank != sharding.dp_rank: + raise LossContractError( + f"dp_group rank {group_rank} does not match the contract dp_rank={sharding.dp_rank}" + ) + + +def _preflight_cross_rank_agreement( + contract: GRPOLossContract, dp_group: Any, tp_group: Any, num_vocab_tiles: int +) -> None: + """All-gather (fingerprint, backend id, vocab tile count) and abort on mismatch. + + Checked over the DP group *and* the TP group. Neither alone is sufficient: + the loss is replicated across TP, so two TP siblings that disagree on, say, + ``beta`` would compute different losses and produce inconsistent gradients + for one sharded model -- and the logprob path's own preflight cannot catch + that, because ``beta`` is not part of the logprob contract. Agreement + within both groups implies agreement across the whole DP x TP grid by + transitivity. + + Runs before any other collective, including the logprob path's own TP + preflight, so a rank that joined the wrong logical invocation fails here + rather than deadlocking a later reduction. + """ + + checks = [ + (axis, group, world_size) + for axis, group, world_size in ( + ("DP", dp_group, contract.sharding.dp_world_size), + ("TP", tp_group, contract.logprob.sharding.tp_world_size), + ) + if world_size > 1 + ] + if not checks: + return + + dist = _require_distributed_initialized() + payload = (contract.cross_rank_fingerprint(), BACKEND_ID, int(num_vocab_tiles)) + for axis, group, _ in checks: + gathered: list[Any] = [None] * dist.get_world_size(group=group) + dist.all_gather_object(gathered, payload, group=group) + mismatched = [(rank, other) for rank, other in enumerate(gathered) if other != payload] + if mismatched: + rank, other = mismatched[0] + raise LossContractError( + f"cross-rank preflight failed on the {axis} axis: this rank has {payload} " + f"but {axis} rank {rank} has {other}; every rank must agree on the contract " + "fingerprint, backend id and num_vocab_tiles before any collective" + ) + + +def _gather_global_rewards( + rewards: torch.Tensor, contract: GRPOLossContract, dp_group: Any +) -> torch.Tensor: + """Assemble the global ``[num_sequences]`` reward vector on every rank. + + ``sequence_shard_bounds`` is a contiguous partition in DP-rank order, so + concatenating the gathered slices in rank order reproduces the global vector + exactly -- no ownership arbitration is needed. + """ + + sharding = contract.sharding + local = rewards.float() + if sharding.dp_world_size == 1: + return local.contiguous() + + dist = _require_distributed_initialized() + max_local = max(end - start for start, end in sharding.sequence_shard_bounds) + padded = local.new_zeros(max_local) + padded[: local.shape[0]] = local + gathered = [torch.empty_like(padded) for _ in range(sharding.dp_world_size)] + dist.all_gather(gathered, padded.contiguous(), group=dp_group) + return torch.cat( + [ + gathered[rank][: end - start] + for rank, (start, end) in enumerate(sharding.sequence_shard_bounds) + ], + dim=0, + ) + + +def _group_advantages(global_rewards: torch.Tensor, contract: GRPOLossContract) -> torch.Tensor: + """Group-relative advantages over the global reward vector. + + Evaluated identically on every rank from an identically shaped input, so no + merge is involved and the result is bitwise-equal mesh-wide. The variance + is two-pass: centring before squaring keeps the result meaningful when the + rewards share a large offset, which ``E[x^2] - E[x]^2`` does not. + """ + + advantage = contract.objective.advantage + boundaries = contract.sharding.group_boundaries + parts: list[torch.Tensor] = [] + for index in range(len(boundaries) - 1): + start, end = boundaries[index], boundaries[index + 1] + group = global_rewards[start:end] + count = float(end - start) + centered = group - group.sum() / count + if advantage.normalizer is AdvantageNormalizer.MEAN_ONLY: + parts.append(centered) + continue + variance = (centered * centered).sum() / count + parts.append(centered / variance.clamp_min(advantage.std_eps**2).sqrt()) + return torch.cat(parts, dim=0) + + +def _sequence_totals(values: torch.Tensor, contract: GRPOLossContract) -> torch.Tensor: + """Reduce this rank's per-token values to one total per owned sequence. + + The reduced extent is ``padded_seq_len``, which the contract fixes, so this + sum is byte-for-byte the same work at every DP degree. + """ + + sharding = contract.sharding + view = values.reshape(sharding.local_num_sequences, sharding.padded_seq_len) + return view.sum(dim=1) + + +def _assemble_global_vector( + local_totals: torch.Tensor, contract: GRPOLossContract, dp_group: Any +) -> torch.Tensor: + """Place every rank's per-sequence totals into the fixed global vector. + + Returns ``[num_sequences, ...]``. The length is a property of the contract, + never of the DP degree, and filling it is pure placement -- no arithmetic + touches the gathered values -- so the downstream reduction sees identical + inputs at every degree. + """ + + sharding = contract.sharding + if sharding.dp_world_size == 1: + return local_totals + + dist = _require_distributed_initialized() + trailing = local_totals.shape[1:] + max_local = max(end - start for start, end in sharding.sequence_shard_bounds) + padded = local_totals.new_zeros((max_local, *trailing)) + padded[: local_totals.shape[0]] = local_totals.detach() + gathered = [torch.empty_like(padded) for _ in range(sharding.dp_world_size)] + dist.all_gather(gathered, padded.contiguous(), group=dp_group) + + # This rank's own slice comes from the live tensor: all_gather severs the + # graph, and the other ranks' slices are constants here anyway. + pieces = [] + for rank, (start, end) in enumerate(sharding.sequence_shard_bounds): + pieces.append(local_totals if rank == sharding.dp_rank else gathered[rank][: end - start]) + return torch.cat(pieces, dim=0) + + +def _normalized( + per_sequence_totals: torch.Tensor, + per_sequence_counts: torch.Tensor, + contract: GRPOLossContract, +) -> torch.Tensor: + """Apply the declared token normalizer to fixed-order sequence totals.""" + + reduction = contract.reduction + normalizer = reduction.token_normalizer + if normalizer is TokenNormalizer.FIXED_CONSTANT: + fixed_normalizer_constant = reduction.fixed_normalizer_constant + assert fixed_normalizer_constant is not None + return per_sequence_totals.sum() / float(fixed_normalizer_constant) + if normalizer is TokenNormalizer.GLOBAL_ACTIVE_TOKENS: + return per_sequence_totals.sum() / per_sequence_counts.sum().to(per_sequence_totals.dtype) + + # PER_SEQUENCE_THEN_MEAN: sequences with no active token contribute nothing + # and are excluded from the outer denominator rather than counted as zero. + live = per_sequence_counts > 0 + denominators = per_sequence_counts.to(per_sequence_totals.dtype).clamp_min(1.0) + per_sequence_means = torch.where( + live, per_sequence_totals / denominators, torch.zeros_like(per_sequence_totals) + ) + return per_sequence_means.sum() / live.sum().to(per_sequence_totals.dtype) + + +class DistributedGRPOLossOp: + """Deterministic GRPO loss over TP-sharded logits with a DP-invariant reduction. + + The WS2 reference (issue #241 PR5). ``policy_local_logits`` is this rank's + ``[n, local_vocab]`` vocabulary shard, not a dense ``[n, vocab]`` tensor: + tensor parallelism is delegated to the vocab-parallel logprob path, and this + operator owns the sum over tokens and sequences. + """ + + op_class = "grpo_loss" + is_batch_invariant = True + + def __init__(self) -> None: + self._logprob = VocabParallelLogprobOp() + + def __call__(self, *args: Any, **kwargs: Any) -> GRPOLossResult: + return self.apply(*args, **kwargs) + + def apply( + self, + policy_local_logits: torch.Tensor, + action_ids: torch.Tensor, + old_logps: torch.Tensor, + rewards: torch.Tensor, + *, + contract: GRPOLossContract, + ref_local_logits: torch.Tensor | None = None, + tp_group: Any = None, + dp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + ) -> GRPOLossResult: + if not isinstance(contract, GRPOLossContract): + raise LossContractError("contract must be a GRPOLossContract") + _validate_invocation( + policy_local_logits, + ref_local_logits, + action_ids, + old_logps, + rewards, + contract, + dp_group, + ) + sharding = contract.sharding + objective = contract.objective + if validate: + _preflight_cross_rank_agreement(contract, dp_group, tp_group, num_vocab_tiles) + + logp_policy, _ = self._logprob.apply( + policy_local_logits, + action_ids, + contract=contract.logprob, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + ) + active = torch.tensor( + contract.logprob.mask.active_mask, + dtype=torch.bool, + device=policy_local_logits.device, + ) + + delta = (logp_policy - old_logps.float()).masked_fill(~active, 0.0) + ratio = delta.exp() + + if ref_local_logits is None: + kl_terms = torch.zeros_like(logp_policy) + else: + with torch.no_grad(): + logp_ref, _ = self._logprob.apply( + ref_local_logits, + action_ids, + contract=contract.logprob, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=False, + ) + diff = (logp_ref - logp_policy).masked_fill(~active, 0.0) + if objective.kl_estimator is KLEstimator.K3_UNBIASED: + kl_terms = diff.exp() - diff - 1.0 + else: + kl_terms = -diff + kl_terms = kl_terms.masked_fill(~active, 0.0) + + global_rewards = _gather_global_rewards(rewards, contract, dp_group) + advantages = _group_advantages(global_rewards, contract) + adv_tokens = ( + advantages[sharding.local_sequence_start : sharding.local_sequence_end] + .reshape(-1, 1) + .expand(sharding.local_num_sequences, sharding.padded_seq_len) + .reshape(-1) + ) + + clip = objective.clip + unclipped = ratio * adv_tokens + clipped = ratio.clamp(clip.lower_bound, clip.upper_bound) * adv_tokens + policy_terms = (-torch.minimum(unclipped, clipped)).masked_fill(~active, 0.0) + + packed = torch.stack( + (_sequence_totals(policy_terms, contract), _sequence_totals(kl_terms, contract)), + dim=1, + ) + # Fixed-extent reductions: [local_seqs, padded_seq_len] -> [local_seqs], + # gathered into [num_sequences] -> scalar, at every DP degree. + totals = _assemble_global_vector(packed, contract, dp_group) + per_sequence_policy = totals[:, _CH_POLICY] + per_sequence_kl = totals[:, _CH_KL] + per_sequence_counts = _assemble_global_vector( + _sequence_totals(active.to(torch.long), contract), contract, dp_group + ) + + total_active = int(per_sequence_counts.sum().item()) + if validate and total_active == 0: + raise LossContractError( + "the global batch holds no active tokens; the loss normalizer would divide by zero" + ) + + policy_loss = _normalized(per_sequence_policy, per_sequence_counts, contract) + kl = _normalized(per_sequence_kl, per_sequence_counts, contract) + loss = policy_loss + objective.beta * kl + + provenance = { + "backend_id": BACKEND_ID, + "implementation_kind": "reference", + "num_vocab_tiles": int(num_vocab_tiles), + "padded_seq_len": int(sharding.padded_seq_len), + "num_sequences": int(sharding.num_sequences), + "global_active_tokens": total_active, + "reference_model_used": ref_local_logits is not None, + "requested_contract": contract.to_dict(), + "cross_rank_fingerprint": contract.cross_rank_fingerprint(), + } + return GRPOLossResult( + loss=loss, + policy_loss=policy_loss, + kl=kl, + advantages=advantages, + per_sequence_policy=per_sequence_policy.detach(), + per_sequence_kl=per_sequence_kl.detach(), + per_sequence_active_tokens=per_sequence_counts.detach(), + provenance=provenance, + ) + + +__all__ = [ + "BACKEND_ID", + "GRPOLossResult", + "DistributedGRPOLossOp", +] diff --git a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py index 3745382a..928e818a 100644 --- a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py @@ -33,19 +33,11 @@ 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. - -``deterministic=False`` trades the guarantee for speed: each rank reduces its -whole shard in one pass and the per-shard partials are merged in shard order, -so the floating-point grouping changes with the TP degree and nothing is -promised about reproducibility. ``num_vocab_tiles`` is ignored (no tile -alignment is required), and a contract declaring -``determinism_scope=cross_tp_bitwise`` is rejected loudly — the fast path -cannot honor it. Batch invariance is unaffected: rows never mix either way. """ from __future__ import annotations -from typing import Any +from typing import Any, Callable, Protocol import torch @@ -171,15 +163,12 @@ def _validate_active_targets( def _preflight_cross_rank_agreement( - contract: LogprobContract, tp_group: Any, num_vocab_tiles: int, deterministic: bool + contract: LogprobContract, + tp_group: Any, + num_vocab_tiles: int, + deterministic: bool, ) -> None: - """All-gather (fingerprint, backend id, tile count, mode) and abort on mismatch. - - The tile count travels as ``None`` when ``deterministic=False``: the fast - path never uses it, so ranks must not fail preflight over an irrelevant - value — but they must never disagree on the mode itself, or they would - issue different collectives. - """ + """All-gather the numerical mode and abort before mismatched collectives.""" dist = _require_distributed_initialized() payload = ( @@ -226,20 +215,64 @@ def _local_tile_stats(z_masked: torch.Tensor, tile: int) -> tuple[torch.Tensor, return torch.stack(m_parts, dim=1), torch.stack(s_parts, dim=1) +def _native_rocm_tile_stats( + z_masked: torch.Tensor, + tile: int, + *, + vocab_start: int, + real_vocab_size: int, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Use the HIP-compatible local partial kernel when the extension has it. + + TP transport and the global merge deliberately stay in Python. That keeps + the issue #241 reduction contract identical across CUDA, ROCm, and the + reference implementation while making the large local vocab scan native. + """ + + if torch.version.hip is None or not z_masked.is_cuda: + raise RuntimeError("ROCm native tile stats require a HIP CUDA tensor") + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not hasattr(_C, "deterministic_logp_tile_stats"): + raise RuntimeError( + "ROCm vocab-parallel logprob native extension is unavailable; " + "build rl_engine._C with a ROCm toolchain" + ) + local_tiles = z_masked.shape[1] // tile + if local_tiles <= 0: + raise RuntimeError("ROCm native tile stats require at least one local vocab tile") + # The ROCm build also ships a gfx942-tuned kernel with the same contract; + # one kernel serves every ROCm entry point so their partials stay bitwise equal. + tile_stats = getattr( + _C, "hip_deterministic_logp_tile_stats", _C.deterministic_logp_tile_stats + ) + return tuple( + tensor.contiguous() + for tensor in tile_stats( + z_masked, + int(vocab_start), + int(real_vocab_size), + int(local_tiles), + ) + ) + except (ImportError, AttributeError) as exc: + raise RuntimeError("ROCm vocab-parallel logprob native extension is unavailable") from exc + + def _gather_tile_stats( local_m: torch.Tensor, local_s: torch.Tensor, contract: LogprobContract, tp_group: Any, - tile_counts: list[int], + tile_counts: int | list[int], ) -> tuple[torch.Tensor, torch.Tensor]: - """Assemble every rank's partials in global shard order. - - ``tile_counts`` holds each rank's partial count: one per tile on the - deterministic path, exactly one per shard on the fast path. - """ + """Assemble partials in global shard and tile order.""" sharding = contract.sharding + if isinstance(tile_counts, int): + tile = tile_counts + 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() @@ -342,6 +375,7 @@ def forward( tile, with_entropy, with_entropy_grad, + use_native_tile_stats, ): z_masked = local_logits.float() sharding = contract.sharding @@ -354,14 +388,27 @@ def forward( safe_target = torch.where(active_mask, target_1d, torch.zeros_like(target_1d)) - if tile is not None: - local_m, local_s = _local_tile_stats(z_masked, tile) - tile_counts = [(end - start) // tile for start, end in sharding.vocab_shard_bounds] - else: - # Fast path: one (max, sumexp) partial over the whole shard. The - # grouping now depends on the shard layout, so no cross-TP claim. + if tile is None: local_m, local_s = _local_tile_stats(z_masked, z_masked.shape[1]) tile_counts = [1] * sharding.tp_world_size + elif use_native_tile_stats: + # A backend may pass its own tile-stats callable (same signature as + # _native_rocm_tile_stats); ``True`` selects the ROCm extension kernel. + tile_stats_fn = ( + use_native_tile_stats + if callable(use_native_tile_stats) + else _native_rocm_tile_stats + ) + local_m, local_s = tile_stats_fn( + z_masked, + tile, + vocab_start=sharding.local_vocab_start, + real_vocab_size=sharding.real_vocab_size, + ) + tile_counts = [(end - start) // tile for start, end in sharding.vocab_shard_bounds] + else: + local_m, local_s = _local_tile_stats(z_masked, tile) + tile_counts = [(end - start) // tile for start, end in sharding.vocab_shard_bounds] m_all, s_all = _gather_tile_stats(local_m, local_s, contract, tp_group, tile_counts) target_logit = _gather_target_logit(z_masked, safe_target, contract, tp_group) lse = _merge_tile_partials(m_all, s_all) @@ -401,7 +448,7 @@ 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 + return None, 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 @@ -432,20 +479,157 @@ def backward(ctx, grad_logp, grad_lse, grad_entropy): 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 + return grad.to(ctx.input_dtype), None, None, None, None, None, None, None, None -class VocabParallelLogprobOp: - """Vocab-parallel selected-token logprob (WS2 reference). +class VocabParallelLogprobKernels(Protocol): + """Native kernel pair a fused WS2 backend plugs into :func:`apply_with_kernels`. - Deterministic by default (cross-TP bitwise, tile-ordered merge); - ``deterministic=False`` selects the faster whole-shard reduction with no - reproducibility guarantee. + ``tile_stats`` returns FP32 ``(max, sumexp)`` partials of shape + ``[tokens, num_tiles]`` over the real-vocabulary part of each tile, with a + fixed per-tile reduction order that does not depend on the tile's position + in the local shard or on the storage dtype. ``backward`` returns + ``grad_logits`` in the shard dtype for + ``coef_logp * (onehot - p) + coef_lse * p`` with ``p = exp(z - lse)`` on + finite rows, ``0`` on non-finite rows and padding columns. """ + def tile_stats( + self, shard: torch.Tensor, vocab_start: int, real_vocab: int, num_tiles: int + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + def backward( + self, + shard: torch.Tensor, + lse: torch.Tensor, + coef_logp: torch.Tensor, + coef_lse: torch.Tensor, + target_local: torch.Tensor, + vocab_start: int, + real_vocab: int, + has_lse_grad: bool, + ) -> torch.Tensor: ... + + +class _KernelVocabParallelLogprobFunction(torch.autograd.Function): + """Selected logprob + LSE with native tile statistics and a fused native backward. + + The transport (``_gather_tile_stats``, ``_gather_target_logit``) and the + fixed tile-order merge are the reference helpers above, so every kernel + backend shares one contract and differs from the reference only by the + FP32 summation order inside a tile. The input shard is saved as-is (no + FP32 copy); the backward kernel recomputes ``p`` from it. + """ + + @staticmethod + def forward(ctx, local_logits, target_1d, active_mask, contract, tp_group, tile, kernels): + sharding = contract.sharding + shard = local_logits.contiguous() + local_tiles = sharding.local_vocab_size // tile + if local_tiles <= 0: + raise RuntimeError("native tile stats require at least one local vocab tile") + local_m, local_s = kernels.tile_stats( + shard, sharding.local_vocab_start, sharding.real_vocab_size, local_tiles + ) + m_all, s_all = _gather_tile_stats( + local_m.contiguous(), local_s.contiguous(), contract, tp_group, tile + ) + safe_target = torch.where(active_mask, target_1d, torch.zeros_like(target_1d)) + # Exact: the owner's logit is copied in its storage dtype and widened once. + target_logit = _gather_target_logit(shard, safe_target, contract, tp_group).float() + lse = _merge_tile_partials(m_all, s_all) + selected_logp = torch.where(active_mask, target_logit - lse, torch.zeros_like(lse)) + + ctx.save_for_backward(shard, lse, safe_target, active_mask) + ctx.kernels = kernels + ctx.local_vocab_start = sharding.local_vocab_start + ctx.real_vocab_size = sharding.real_vocab_size + ctx.set_materialize_grads(False) + return selected_logp, lse + + @staticmethod + def backward(ctx, grad_logp, grad_lse): + if not ctx.needs_input_grad[0] or (grad_logp is None and grad_lse is None): + return None, None, None, None, None, None, None + shard, lse, safe_target, active_mask = ctx.saved_tensors + n, local_vocab = shard.shape + start = ctx.local_vocab_start + if grad_logp is not None: + coef_logp = ( + torch.where(active_mask, grad_logp, torch.zeros_like(grad_logp)) + .float() + .contiguous() + ) + owns = (safe_target >= start) & (safe_target < start + local_vocab) + hit = owns & active_mask + target_local = torch.where( + hit, safe_target - start, torch.full_like(safe_target, -1) + ).contiguous() + else: + coef_logp = lse.new_zeros((n,)) + target_local = torch.full((n,), -1, dtype=torch.long, device=shard.device) + has_lse_grad = grad_lse is not None + coef_lse = grad_lse.float().contiguous() if has_lse_grad else lse.new_zeros((n,)) + grad = ctx.kernels.backward( + shard, + lse.contiguous(), + coef_logp, + coef_lse, + target_local, + start, + ctx.real_vocab_size, + has_lse_grad, + ) + return grad, None, None, None, None, None, None + + +def apply_with_kernels( + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any, + num_vocab_tiles: int, + validate: bool, + kernels: VocabParallelLogprobKernels, +) -> tuple[torch.Tensor, torch.Tensor]: + """Contract-checked selected-logprob/LSE forward through a native kernel pair.""" + + 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, True) + + selected_logp, lse = _KernelVocabParallelLogprobFunction.apply( + local_logits, target_1d, active_mask, contract, tp_group, tile, kernels + ) + + 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 + + +class VocabParallelLogprobOp: + """Deterministic vocab-parallel selected-token logprob (WS2 reference).""" + op_class = "logprob" is_batch_invariant = True backend_id = BACKEND_ID + # False: PyTorch tile loop. True: ROCm extension kernel. A callable with the + # _native_rocm_tile_stats signature: backend-specific tile statistics. + use_native_tile_stats: bool | Callable[..., tuple[torch.Tensor, torch.Tensor]] = False def __init__(self) -> None: pass @@ -492,6 +676,7 @@ def apply( deterministic=deterministic, with_entropy=False, with_entropy_grad=False, + use_native_tile_stats=self.use_native_tile_stats, ) return selected_logp, lse @@ -523,6 +708,7 @@ def apply_with_entropy( deterministic=deterministic, with_entropy=True, with_entropy_grad=with_entropy_grad, + use_native_tile_stats=self.use_native_tile_stats, ) def _apply( @@ -537,6 +723,7 @@ def _apply( deterministic: bool, with_entropy: bool, with_entropy_grad: bool, + use_native_tile_stats: bool | Callable[..., tuple[torch.Tensor, torch.Tensor]], ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if not isinstance(contract, LogprobContract): raise LogprobContractError("contract must be a LogprobContract") @@ -547,9 +734,8 @@ def _apply( elif contract.reduction.determinism_scope is DeterminismScope.CROSS_TP_BITWISE: raise LogprobContractError( "deterministic=False cannot honor determinism_scope=cross_tp_bitwise: " - "the fast path reduces each shard in one piece, so its floating-point " - "grouping changes with the TP degree; keep deterministic=True or relax " - "the contract to determinism_scope=fixed_topology" + "keep deterministic=True or relax the contract to " + "determinism_scope=fixed_topology" ) else: tile = None @@ -573,6 +759,7 @@ def _apply( tile, with_entropy, with_entropy_grad, + use_native_tile_stats if deterministic else False, ) if validate and bool((~torch.isfinite(lse) & active_mask).any().item()): @@ -586,5 +773,7 @@ def _apply( __all__ = [ "BACKEND_ID", "DEFAULT_NUM_VOCAB_TILES", + "VocabParallelLogprobKernels", "VocabParallelLogprobOp", + "apply_with_kernels", ] diff --git a/rl_engine/kernels/ops/rocm/loss/__init__.py b/rl_engine/kernels/ops/rocm/loss/__init__.py new file mode 100644 index 00000000..b7fbeaca --- /dev/null +++ b/rl_engine/kernels/ops/rocm/loss/__init__.py @@ -0,0 +1,5 @@ +"""ROCm logprob operators.""" + +from .vocab_parallel_logp import RocmVocabParallelLogprobOp + +__all__ = ["RocmVocabParallelLogprobOp"] diff --git a/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py new file mode 100644 index 00000000..bfabe8b1 --- /dev/null +++ b/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py @@ -0,0 +1,117 @@ +"""ROCm WS2 vocab-parallel logprob backend. + +The reference operator owns the contract: TP transport, fixed global tile-order +merge, target ownership, masking, entropy, and autograd semantics. This backend +keeps all of that and replaces the two large per-shard passes with HIP kernels +from ``csrc/hip/hip_deterministic_logp_kernel.hip`` (compiled only for ROCm; the +shared ``csrc/deterministic_logp_kernel.cu`` keeps the SM90-tuned CUDA path): + +* ``hip_deterministic_logp_tile_stats`` computes the per-row, per-tile FP32 + ``(max, sumexp)`` partials straight from the BF16/FP16/FP32 shard. The kernel + converts each element exactly, filters padding columns itself, and reduces + every tile with a fixed tree, so no FP32 copy of the logits is materialized. +* ``hip_deterministic_logp_backward`` produces ``grad_logits`` for the selected + logprob and LSE outputs in one fused pass from the saved input shard. + +``apply`` runs through the shared :func:`apply_with_kernels` path; +``apply_with_entropy`` keeps the shared autograd path (with the HIP tile +kernel) because the entropy gradient needs the full probability tensor anyway. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.kernels.logprob_contract import LogprobContract +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + DEFAULT_NUM_VOCAB_TILES, + VocabParallelLogprobOp, + apply_with_kernels, +) + +BACKEND_ID = "rocm-vocab-parallel-logp-ws2" + + +def _native_backward_available() -> bool: + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + except ImportError: + return False + return bool( + _EXT_AVAILABLE + and hasattr(_C, "hip_deterministic_logp_tile_stats") + and hasattr(_C, "hip_deterministic_logp_backward") + ) + + +class _HipKernels: + """``VocabParallelLogprobKernels`` over the ROCm extension symbols.""" + + @staticmethod + def tile_stats(shard, vocab_start, real_vocab, num_tiles): + from rl_engine.kernels.ops.base import _C + + tile_max, tile_sum = _C.hip_deterministic_logp_tile_stats( + shard, vocab_start, real_vocab, num_tiles + ) + return tile_max, tile_sum + + @staticmethod + def backward( + shard, lse, coef_logp, coef_lse, target_local, vocab_start, real_vocab, has_lse_grad + ): + from rl_engine.kernels.ops.base import _C + + return _C.hip_deterministic_logp_backward( + shard, lse, coef_logp, coef_lse, target_local, vocab_start, real_vocab, has_lse_grad + ) + + +class RocmVocabParallelLogprobOp(VocabParallelLogprobOp): + """Contract-preserving ROCm implementation with HIP local reductions.""" + + op_class = "logprob" + is_batch_invariant = True + backend_id = BACKEND_ID + use_native_tile_stats = True + + 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, + deterministic: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not deterministic: + return super().apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + deterministic=False, + ) + if not _native_backward_available(): + raise RuntimeError( + f"{BACKEND_ID} requires rl_engine._C built with a ROCm toolchain " + "(hip_deterministic_logp_* symbols are missing); it does not fall back" + ) + return apply_with_kernels( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + kernels=_HipKernels, + ) + + +__all__ = ["BACKEND_ID", "RocmVocabParallelLogprobOp"] 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/ops/triton/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/triton/loss/vocab_parallel_logp.py new file mode 100644 index 00000000..78dbf65b --- /dev/null +++ b/rl_engine/kernels/ops/triton/loss/vocab_parallel_logp.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Triton WS2 vocab-parallel logprob backend (``triton-vocab-parallel-logp-ws2``). + +Same contract, TP transport, and fixed tile-order merge as the PyTorch +reference; the two per-shard passes are Triton kernels, so the backend runs on +both CUDA and ROCm from one source: + +* ``_vocab_tile_stats_kernel``: one program per ``(row, tile)`` computes the + FP32 ``(max, sumexp)`` partial over the real-vocabulary part of the tile. The + tile is walked in ``BLOCK_V`` chunks measured from the tile start, and masked + lanes contribute the reduction identity, so the reduction order depends only + on ``BLOCK_V`` (not on the tile's position in the shard, the padding, or the + storage dtype) and TP=n stays bitwise equal to TP=1. +* ``_vocab_logp_backward_kernel``: fused elementwise + ``coef_logp * (onehot - p) + coef_lse * p`` with ``p = exp(z - lse)`` on + finite rows, ``0`` on non-finite rows and padding columns. +""" + +from __future__ import annotations + +from typing import Any + +import torch +import triton +import triton.language as tl + +from rl_engine.kernels.logprob_contract import LogprobContract +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + DEFAULT_NUM_VOCAB_TILES, + VocabParallelLogprobOp, + apply_with_kernels, +) + +BACKEND_ID = "triton-vocab-parallel-logp-ws2" +_BLOCK_V: int = 1024 +_NUM_WARPS: int = 4 +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +@triton.jit +def _vocab_tile_stats_kernel( + logits_ptr, # [rows, local_vocab] shard, any of fp16/bf16/fp32 + max_ptr, # [rows, local_tiles] fp32 + sum_ptr, # [rows, local_tiles] fp32 + local_vocab, + vocab_start, + real_vocab, + tile_size, + local_tiles, + stride_row, + BLOCK_V: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + col_begin = tile * tile_size + col_end = tl.minimum(col_begin + tile_size, local_vocab) + # Columns at or beyond the real vocabulary are padding and contribute nothing. + real_end = tl.minimum(col_end, tl.maximum(real_vocab - vocab_start, 0)) + row_base = row.to(tl.int64) * stride_row + + tile_max = tl.full((), float("-inf"), dtype=tl.float32) + for start in range(col_begin, col_end, BLOCK_V): + cols = start + tl.arange(0, BLOCK_V) + values = tl.load( + logits_ptr + row_base + cols, mask=cols < real_end, other=float("-inf") + ).to(tl.float32) + tile_max = tl.maximum(tile_max, tl.max(values)) + + empty = tile_max == float("-inf") + max_safe = tl.where(empty, 0.0, tile_max) + tile_sum = tl.zeros((), dtype=tl.float32) + for start in range(col_begin, col_end, BLOCK_V): + cols = start + tl.arange(0, BLOCK_V) + values = tl.load( + logits_ptr + row_base + cols, mask=cols < real_end, other=float("-inf") + ).to(tl.float32) + tile_sum += tl.sum(tl.exp(values - max_safe)) + tile_sum = tl.where(empty, 0.0, tile_sum) + + out = row.to(tl.int64) * local_tiles + tile + tl.store(max_ptr + out, tile_max) + tl.store(sum_ptr + out, tile_sum) + + +@triton.jit +def _vocab_logp_backward_kernel( + logits_ptr, # [rows, local_vocab] shard + lse_ptr, # [rows] fp32 merged vocabulary LSE + coef_logp_ptr, # [rows] fp32, upstream grad of selected logp (0 on inactive rows) + coef_lse_ptr, # [rows] fp32, upstream grad of LSE (ignored unless HAS_LSE_GRAD) + target_ptr, # [rows] int64 local column of the owned target, or -1 + grad_ptr, # [rows, local_vocab] output in the shard dtype + local_vocab, + vocab_start, + real_vocab, + stride_row, + HAS_LSE_GRAD: tl.constexpr, + BLOCK_V: tl.constexpr, +): + row = tl.program_id(0) + chunk = tl.program_id(1) + cols = chunk * BLOCK_V + tl.arange(0, BLOCK_V) + mask = cols < local_vocab + real_end = tl.minimum(local_vocab, tl.maximum(real_vocab - vocab_start, 0)) + row_base = row.to(tl.int64) * stride_row + + lse = tl.load(lse_ptr + row) + finite = (lse == lse) & (lse != float("inf")) & (lse != float("-inf")) + lse_safe = tl.where(finite, lse, 0.0) + g_logp = tl.load(coef_logp_ptr + row) + hit = tl.load(target_ptr + row) + + values = tl.load(logits_ptr + row_base + cols, mask=mask, other=0.0).to(tl.float32) + p = tl.where(finite, tl.exp(values - lse_safe), 0.0) + onehot = tl.where(cols == hit, 1.0, 0.0) + grad = g_logp * (onehot - p) + if HAS_LSE_GRAD: + g_lse = tl.load(coef_lse_ptr + row) + grad = grad + g_lse * p + grad = tl.where(cols < real_end, grad, 0.0) + tl.store(grad_ptr + row_base + cols, grad.to(grad_ptr.dtype.element_ty), mask=mask) + + +def _check_shard(shard: torch.Tensor) -> torch.Tensor: + if shard.device.type not in ("cuda", "hip", "xpu"): + raise RuntimeError(f"{BACKEND_ID} requires a GPU tensor, got device {shard.device}") + if shard.dim() != 2: + raise ValueError("logits must be 2D [tokens, local_vocab]") + if shard.dtype not in _SUPPORTED_DTYPES: + raise TypeError("logits must be float16, bfloat16, or float32") + return shard.contiguous() + + +def triton_vocab_tile_stats( + logits: torch.Tensor, vocab_start: int, real_vocab: int, num_tiles: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-row, per-tile FP32 ``(max, sumexp)`` partials over the real vocabulary.""" + + shard = _check_shard(logits) + if vocab_start < 0 or real_vocab <= 0 or num_tiles <= 0: + raise ValueError("invalid vocabulary metadata") + rows, local_vocab = shard.shape + if local_vocab <= 0 or local_vocab % num_tiles != 0: + raise ValueError("local_vocab must be divisible by num_tiles") + tile_max = torch.empty(rows, num_tiles, device=shard.device, dtype=torch.float32) + tile_sum = torch.empty(rows, num_tiles, device=shard.device, dtype=torch.float32) + if rows == 0: + return tile_max, tile_sum + _vocab_tile_stats_kernel[(rows, num_tiles)]( + shard, + tile_max, + tile_sum, + local_vocab, + int(vocab_start), + int(real_vocab), + local_vocab // num_tiles, + num_tiles, + shard.stride(0), + BLOCK_V=_BLOCK_V, + num_warps=_NUM_WARPS, + ) + return tile_max, tile_sum + + +def triton_vocab_logp_backward( + logits: torch.Tensor, + lse: torch.Tensor, + coef_logp: torch.Tensor, + coef_lse: torch.Tensor, + target_local: torch.Tensor, + vocab_start: int, + real_vocab: int, + has_lse_grad: bool, +) -> torch.Tensor: + """Fused ``grad_logits`` for the selected logprob and LSE outputs.""" + + shard = _check_shard(logits) + rows, local_vocab = shard.shape + for name, tensor, dtype in ( + ("lse", lse, torch.float32), + ("coef_logp", coef_logp, torch.float32), + ("coef_lse", coef_lse, torch.float32), + ("target_local", target_local, torch.long), + ): + if tensor.shape != (rows,) or tensor.dtype != dtype or tensor.device != shard.device: + raise ValueError(f"{name} must be a [{rows}] {dtype} tensor on the logits device") + grad = torch.empty_like(shard) + if rows == 0 or local_vocab == 0: + return grad + grid = (rows, triton.cdiv(local_vocab, _BLOCK_V)) + _vocab_logp_backward_kernel[grid]( + shard, + lse.contiguous(), + coef_logp.contiguous(), + coef_lse.contiguous(), + target_local.contiguous(), + grad, + local_vocab, + int(vocab_start), + int(real_vocab), + shard.stride(0), + HAS_LSE_GRAD=bool(has_lse_grad), + BLOCK_V=_BLOCK_V, + num_warps=_NUM_WARPS, + ) + return grad + + +class _TritonKernels: + """``VocabParallelLogprobKernels`` over the Triton launchers above.""" + + tile_stats = staticmethod(triton_vocab_tile_stats) + backward = staticmethod(triton_vocab_logp_backward) + + +def _entropy_tile_stats( + z_masked: torch.Tensor, tile: int, *, vocab_start: int, real_vocab_size: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Tile statistics for the shared entropy path (same kernel as ``apply``).""" + + return triton_vocab_tile_stats( + z_masked, vocab_start, real_vocab_size, z_masked.shape[1] // tile + ) + + +class TritonVocabParallelLogprobOp(VocabParallelLogprobOp): + """Contract-preserving Triton implementation of the WS2 vocab-parallel logprob.""" + + op_class = "logprob" + is_batch_invariant = True + backend_id = BACKEND_ID + use_native_tile_stats = staticmethod(_entropy_tile_stats) + + 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, + deterministic: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not deterministic: + return super().apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + deterministic=False, + ) + return apply_with_kernels( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + kernels=_TritonKernels, + ) + + +__all__ = [ + "BACKEND_ID", + "TritonVocabParallelLogprobOp", + "triton_vocab_logp_backward", + "triton_vocab_tile_stats", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 12ea9b21..8937a4da 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -40,6 +40,18 @@ from rl_engine.utils.logger import logger +def _rocm_vocab_logprob_native_available() -> bool: + """Return whether the strict ROCm tile kernel is actually loadable.""" + + if torch.version.hip is None: + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + except ImportError: + return False + return bool(_EXT_AVAILABLE and hasattr(_C, "deterministic_logp_tile_stats")) + + class _KernelEnumMeta(EnumMeta): """Metaclass to provide enhanced error messaging for backend lookups.""" @@ -111,6 +123,16 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_VOCAB_PARALLEL_LOGP = ( "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" ) + ROCM_VOCAB_PARALLEL_LOGP = ( + "rl_engine.kernels.ops.rocm.loss.vocab_parallel_logp.RocmVocabParallelLogprobOp" + ) + TRITON_VOCAB_PARALLEL_LOGP = ( + "rl_engine.kernels.ops.triton.loss.vocab_parallel_logp.TritonVocabParallelLogprobOp" + ) + # Deterministic GRPO loss on the TP-aware logprob path (WS2 #241 PR5) + PYTORCH_DISTRIBUTED_GRPO_LOSS = ( + "rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss.DistributedGRPOLossOp" + ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -299,6 +321,18 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ) +def _triton_vocab_logprob_available() -> bool: + """Return whether the Triton WS2 vocab-parallel kernels can run here.""" + + if not torch.cuda.is_available(): + return False + try: + import triton # noqa: F401 + except ImportError: + return False + return True + + def resolve_logp_op_type( logp_backend: Optional[str] = None, *, @@ -667,6 +701,54 @@ def __init__(self): prepend=True, ) + # Triton WS2 vocab-parallel production backend: one source for CUDA and + # ROCm, registered ahead of the reference wherever Triton can run. + triton_vocab_logprob_capability = LogprobBackendCapability( + backend_id="triton-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="production", + ) + if _triton_vocab_logprob_available(): + for triton_platform in ("cuda", "rocm"): + if triton_platform in self._priority_map: + self.register_logprob_backend( + OpBackend.TRITON_VOCAB_PARALLEL_LOGP, + triton_vocab_logprob_capability, + platform=triton_platform, + prepend=True, + ) + + rocm_vocab_logprob_capability = LogprobBackendCapability( + backend_id="rocm-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="production", + ) + if _rocm_vocab_logprob_native_available(): + self.register_logprob_backend( + OpBackend.ROCM_VOCAB_PARALLEL_LOGP, + rocm_vocab_logprob_capability, + platform="rocm", + prepend=True, + ) + def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() if rocm_attn_backend in {"flash_attn", "flash-attn", "flash_attention"}: diff --git a/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 79f882d9..a7f326d0 100644 --- a/setup.py +++ b/setup.py @@ -1,313 +1,334 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import warnings -from pathlib import Path - -from setuptools import find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - 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 - - -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]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - 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 = _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 = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - 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 = getattr(torch.version, "hip", None) is not None - - # 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'." - ) - - 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/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.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") - - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - 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", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - 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) - if os.name != "nt": - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - 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) +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os +import warnings +from pathlib import Path + +from setuptools import find_packages, setup + + +def _load_envs_module(): + envs_path = Path(__file__).with_name("envs.py") + spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load environment helpers from {envs_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +envs = _load_envs_module() + + +def _load_torch_extension_tools(): + try: + import torch + 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 + + +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]: + value = os.environ.get(name) + if value is None: + return [] + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {value!r}") + 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 = _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 = [] + torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") + 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 = getattr(torch.version, "hip", None) is not None + + # 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'." + ) + + 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/gemm/det_gemm_kernel.cu", + "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", + "csrc/cuda/attention/deterministic_attention.cu", + ] + if is_rocm: + # ROCm-tuned WS2 vocab-parallel logprob kernels; the shared + # deterministic_logp_kernel.cu keeps the SM90-tuned CUDA path. + cuda_sources.append("csrc/hip/hip_deterministic_logp_kernel.hip") + if not is_rocm: + cuda_sources.append("csrc/cuda/distributed/deterministic_collective.cu") + # 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") + + nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] + if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + nvcc_flags.append("--use_fast_math") + 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", + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + ) + ) + if is_rocm: + for tile_knob in ( + "DETERMINISTIC_LOGP_TILE_BLOCK_SIZE", + "DETERMINISTIC_LOGP_TILE_VECTOR_ELEMENTS", + "DETERMINISTIC_LOGP_BACKWARD_BLOCK_SIZE", + ): + nvcc_flags.extend(_cuda_define_from_env(tile_knob, tile_knob)) + else: + # Same idea for the CUDA tile-stats kernel: the defaults are tuned for + # sm_90, and a different architecture or vocabulary split may prefer + # another block size or vector width. + for tile_knob in ( + "DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_NARROW", + "DETERMINISTIC_LOGP_TILE_BLOCK_SIZE_WIDE", + "DETERMINISTIC_LOGP_TILE_VECTOR_BYTES", + ): + nvcc_flags.extend(_cuda_define_from_env(tile_knob, tile_knob)) + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + nvcc_flags.append("-lineinfo") + 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) + if os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") + + 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") - if "-lcuda" not in extra_link_args: - 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: - 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( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - 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() - if BuildExtension is None: - return {} - return {"build_ext": BuildExtension} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], - }, - entry_points={ - "console_scripts": [ - "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", - ], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") + if "-lcuda" not in extra_link_args: + 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: + 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( + name="rl_engine._C", + sources=cuda_sources, + include_dirs=[], + extra_compile_args={ + "cxx": cxx_flags, + "nvcc": nvcc_flags, + }, + 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() + if BuildExtension is None: + return {} + return {"build_ext": BuildExtension} + + +setup( + name="rl-engine", + version="0.1.0", + packages=find_packages(include=["rl_engine", "rl_engine.*"]), + install_requires=[ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", + ], + ext_modules=get_extensions(), + cmdclass=get_cmdclass(), + extras_require={ + "cuda": ["flashinfer"], + "rocm": ["aiter"], + "vllm": ["vllm>=0.6.0"], + "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], + }, + entry_points={ + "console_scripts": [ + "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", + ], + }, + python_requires=">=3.10", + include_package_data=True, + zip_safe=False, +) diff --git a/tests/test_distributed_grpo_loss.py b/tests/test_distributed_grpo_loss.py new file mode 100644 index 00000000..d4f8755f --- /dev/null +++ b/tests/test_distributed_grpo_loss.py @@ -0,0 +1,866 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic GRPO loss on the TP-aware logprob path. + +The headline claim under test is that the scalar loss and its gradient are +bitwise-identical across every TP x DP degree, given a fixed vocab tile count. +``TestMeshBitwise`` is the file's centre of gravity: it runs the reachable +degrees on real NCCL ranks and compares raw bit patterns against the single-rank +baseline. + +The remaining classes support that claim rather than duplicate it: the single- +rank tests pin the objective's algebraic identities (ratio exactly 1, KL exactly +0), the negative controls prove the bitwise comparisons are not vacuous, and the +guard tests check that a rank which disagrees about the invocation aborts +instead of corrupting the merge. + +Comparisons are made on ``per_sequence_policy``/``per_sequence_kl`` rather than +on the scalar loss alone. That is not a stylistic choice: measured over this +file's own inputs, regrouping the token sum moves the per-sequence vector in 12 +of 12 seeds but the scalar loss in only 5 of 12, because averaging +``NUM_SEQUENCES`` totals into one fp32 number rounds most reorderings away. +Asserting on the scalar alone would let a wrong reduction pass most of the time. + +Context parallelism is out of scope (see ``LossShardingSpec``); the contract +rejects ``cp_world_size > 1`` and ``TestGuards`` covers that. + +Multi-rank tests need one GPU per rank and skip otherwise. They stay small on +purpose -- a 1000-token vocabulary and 8 sequences of 32 slots -- so they can +share a node with a running training job; each worker additionally caps itself +with ``set_per_process_memory_fraction`` so a regression here cannot starve a +co-tenant. +""" + +from __future__ import annotations + +import math +import os +import queue +import tempfile +import traceback +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.loss_contract import ( + AdvantageNormalizer, + AdvantageSpec, + ClipSpec, + GRPOLossContract, + KLEstimator, + LossContractError, + LossReductionSpec, + LossShardingSpec, + ObjectiveSpec, + TokenNormalizer, +) +from rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss import ( + BACKEND_ID, + DistributedGRPOLossOp, +) + +# Global batch geometry, shared by every configuration so the comparisons are +# between meshes rather than between problems. The degrees below are chosen so +# that DP and CP each reach 4 while every shard stays tile-aligned. +NUM_SEQUENCES = 8 +PADDED_SEQ_LEN = 32 +NUM_TOKEN_SLOTS = NUM_SEQUENCES * PADDED_SEQ_LEN +# Deliberately straddles the DP=2 split at 4 and the DP=4 splits at 2/4/6, so the +# replicated-advantage path is exercised by groups that no single rank owns. +GROUP_BOUNDARIES = (0, 3, NUM_SEQUENCES) +REAL_VOCAB = 1000 +PADDED_VOCAB = 1024 +NUM_VOCAB_TILES = 32 +BETA = 0.04 +SEED = 20260811 + +_SPAWN_TIMEOUT_S = 600 +# Fraction of each card the workers may allocate. The test tensors need a few +# megabytes; the cap exists so a bug cannot balloon into a co-tenant job. +_MEMORY_FRACTION = 0.02 + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + """Raw bit pattern, so -0.0 vs 0.0 and NaN vs NaN compare honestly.""" + + view_dtype = { + torch.float32: torch.int32, + torch.bfloat16: torch.int16, + torch.float16: torch.int16, + }[tensor.dtype] + return tensor.contiguous().view(view_dtype) + + +def _scalar_bits(tensor: torch.Tensor) -> int: + return int(_bits(tensor.detach().float().cpu()).item()) + + +def _vector_bits(tensor: torch.Tensor) -> list[int]: + return _bits(tensor.detach().float().cpu()).tolist() + + +def _reduction_fingerprint(result) -> tuple: + """Everything a change of reduction order can move, at full resolution. + + The per-sequence vectors come first because they are the sensitive part; + the scalars are carried along so a normalizer bug is caught too. + """ + + return ( + _vector_bits(result.per_sequence_policy), + _vector_bits(result.per_sequence_kl), + _scalar_bits(result.loss), + _scalar_bits(result.policy_loss), + _scalar_bits(result.kl), + ) + + +def _cuda_device_count() -> int: + try: + return torch.cuda.device_count() + except Exception: # pragma: no cover - driver-level failures + return 0 + + +def _requires_gpus(count: int): + return pytest.mark.skipif( + _cuda_device_count() < count, + reason=f"needs {count} CUDA devices for a real {count}-rank mesh", + ) + + +# --------------------------------------------------------------------------- # +# Global problem definition +# --------------------------------------------------------------------------- # +def _global_active_mask() -> tuple[bool, ...]: + """Right-padded sequences of strictly decreasing length. + + The lengths are staggered so that no two DP shards hold the same number of + active tokens; equal counts would let a rank-local normalizer accidentally + agree with the global one and pass a test it should fail. + """ + + mask: list[bool] = [] + for seq in range(NUM_SEQUENCES): + real_len = PADDED_SEQ_LEN - seq * 3 + mask.extend(slot < real_len for slot in range(PADDED_SEQ_LEN)) + return tuple(mask) + + +GLOBAL_ACTIVE_MASK = _global_active_mask() + + +def _global_inputs(seed: int = SEED) -> dict[str, torch.Tensor]: + """Deterministic global tensors every rank slices its own view out of. + + ``old_logps`` is centred on ``-log(REAL_VOCAB)``, the scale of a selected + logprob under near-uniform logits, so the importance ratios land around 1 + with real spread. Leaving it centred on 0 would make every ratio ~1e-3 and + every per-token loss term nearly identical, and a sum of near-identical + values is almost invariant to how it is grouped -- which would quietly + drain the power out of every bitwise comparison in this file. + """ + + gen = torch.Generator().manual_seed(seed) + return { + "policy": torch.randn(NUM_TOKEN_SLOTS, PADDED_VOCAB, generator=gen), + "ref": torch.randn(NUM_TOKEN_SLOTS, PADDED_VOCAB, generator=gen), + "action_ids": torch.randint(0, REAL_VOCAB, (NUM_TOKEN_SLOTS,), generator=gen), + "old_logps": torch.randn(NUM_TOKEN_SLOTS, generator=gen) * 0.5 - math.log(REAL_VOCAB), + "rewards": torch.randn(NUM_SEQUENCES, generator=gen), + } + + +def _dp_bounds(dp: int) -> tuple[tuple[int, int], ...]: + """Contiguous sequence partition in DP-rank order.""" + + seqs = NUM_SEQUENCES // dp + return tuple((d * seqs, (d + 1) * seqs) for d in range(dp)) + + +def _owned_rows(bounds: tuple[int, int]) -> list[int]: + """Global token-slot indices this shard owns, in canonical local row order.""" + + start, end = bounds + return [ + seq * PADDED_SEQ_LEN + slot for seq in range(start, end) for slot in range(PADDED_SEQ_LEN) + ] + + +def _build_contract( + *, + tp_rank: int, + tp: int, + dp_rank: int, + dp: int, + objective: ObjectiveSpec | None = None, + reduction: LossReductionSpec | None = None, +) -> GRPOLossContract: + bounds = _dp_bounds(dp) + rows = _owned_rows(bounds[dp_rank]) + shard = PADDED_VOCAB // tp + logprob = LogprobContract( + role=LogprobRole.TRAIN, + dtype=LogprobDType.FP32, + mask=MaskSpec( + num_tokens=len(rows), + active_mask=tuple(GLOBAL_ACTIVE_MASK[row] for row in rows), + ), + sharding=ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp, + vocab_shard_bounds=tuple((r * shard, (r + 1) * shard) for r in range(tp)), + real_vocab_size=REAL_VOCAB, + padded_vocab_size=PADDED_VOCAB, + ), + reduction=ReductionSpec(), + ) + sharding = LossShardingSpec( + dp_rank=dp_rank, + dp_world_size=dp, + num_sequences=NUM_SEQUENCES, + padded_seq_len=PADDED_SEQ_LEN, + sequence_shard_bounds=bounds, + group_boundaries=GROUP_BOUNDARIES, + ) + return GRPOLossContract( + logprob=logprob, + sharding=sharding, + objective=objective if objective is not None else ObjectiveSpec(beta=BETA), + reduction=reduction if reduction is not None else LossReductionSpec(), + ) + + +def _rank_inputs( + globals_: dict[str, torch.Tensor], + *, + tp_rank: int, + tp: int, + bounds: tuple[int, int], + device: torch.device, +) -> dict[str, torch.Tensor]: + rows = _owned_rows(bounds) + shard = PADDED_VOCAB // tp + cols = slice(tp_rank * shard, (tp_rank + 1) * shard) + start, end = bounds + return { + "policy": globals_["policy"][rows, cols].clone().to(device).requires_grad_(True), + "ref": globals_["ref"][rows, cols].clone().to(device), + "action_ids": globals_["action_ids"][rows].clone().to(device), + "old_logps": globals_["old_logps"][rows].clone().to(device), + "rewards": globals_["rewards"][start:end].clone().to(device), + } + + +def _single_rank_setup( + *, + objective: ObjectiveSpec | None = None, + reduction: LossReductionSpec | None = None, + device: str = "cpu", + globals_: dict[str, torch.Tensor] | None = None, +) -> tuple[GRPOLossContract, dict[str, torch.Tensor]]: + """Contract and tensors for one rank owning the whole batch.""" + + contract = _build_contract( + tp_rank=0, tp=1, dp_rank=0, dp=1, objective=objective, reduction=reduction + ) + tensors = _rank_inputs( + globals_ if globals_ is not None else _global_inputs(), + tp_rank=0, + tp=1, + bounds=(0, NUM_SEQUENCES), + device=torch.device(device), + ) + return contract, tensors + + +def _run_single_rank( + *, + num_vocab_tiles: int = NUM_VOCAB_TILES, + with_reference: bool = True, + **setup, +): + """Baseline invocation: one rank, no collectives.""" + + contract, tensors = _single_rank_setup(**setup) + result = DistributedGRPOLossOp().apply( + tensors["policy"], + tensors["action_ids"], + tensors["old_logps"], + tensors["rewards"], + contract=contract, + ref_local_logits=tensors["ref"] if with_reference else None, + num_vocab_tiles=num_vocab_tiles, + ) + return result, tensors, contract + + +# --------------------------------------------------------------------------- # +# Multi-rank harness +# --------------------------------------------------------------------------- # +def _mesh_worker(rank, world_size, init_method, result_queue, tp, dp, scenario): + """One NCCL rank of a TP x DP mesh. + + Only plain Python values go back on the queue. Tensors sent over a + multiprocessing queue travel by shared memory and are lost when the sender + exits before the parent maps them, which surfaces as an empty queue rather + than an error. + """ + + payload = {"rank": rank} + try: + import torch.distributed as dist + + device_index = rank % torch.cuda.device_count() + torch.cuda.set_device(device_index) + torch.cuda.set_per_process_memory_fraction(_MEMORY_FRACTION, device_index) + device = torch.device("cuda", device_index) + dist.init_process_group( + backend="nccl", init_method=init_method, world_size=world_size, rank=rank + ) + + # Rank layout puts TP fastest: rank = dp_rank * tp + tp_rank. + tp_rank = rank % tp + dp_rank = rank // tp + # new_group is collective, so every rank builds every subgroup in the + # same order even though it only keeps one of each. + tp_group = None + dp_group = None + if tp > 1: + tp_groups = [ + dist.new_group(ranks=list(range(base * tp, (base + 1) * tp))) for base in range(dp) + ] + tp_group = tp_groups[dp_rank] + if dp > 1: + dp_groups = [ + dist.new_group(ranks=list(range(offset, world_size, tp))) for offset in range(tp) + ] + dp_group = dp_groups[tp_rank] + + objective = ObjectiveSpec(beta=BETA) + if scenario == "preflight_dp_mismatch" and dp_rank == 1: + # Perturb a pure fingerprint field. Changing the batch geometry + # instead would change this rank's local shapes, so it would fail + # while constructing its contract -- before the preflight -- and + # strand the other ranks inside the all-gather. + objective = ObjectiveSpec(beta=BETA * 2) + if scenario == "preflight_tp_mismatch" and tp_rank == 1: + # beta is invisible to the logprob contract, so only the loss + # preflight's TP-axis check can catch this one. + objective = ObjectiveSpec(beta=BETA * 2) + + contract = _build_contract( + tp_rank=tp_rank, tp=tp, dp_rank=dp_rank, dp=dp, objective=objective + ) + tensors = _rank_inputs( + _global_inputs(), + tp_rank=tp_rank, + tp=tp, + bounds=contract.sharding.sequence_shard_bounds[dp_rank], + device=device, + ) + op = DistributedGRPOLossOp() + result = op.apply( + tensors["policy"], + tensors["action_ids"], + tensors["old_logps"], + tensors["rewards"], + contract=contract, + ref_local_logits=tensors["ref"], + tp_group=tp_group, + dp_group=dp_group, + num_vocab_tiles=NUM_VOCAB_TILES, + ) + result.loss.backward() + + grad = tensors["policy"].grad + payload.update( + { + "ok": True, + "loss_bits": _scalar_bits(result.loss), + "policy_bits": _scalar_bits(result.policy_loss), + "kl_bits": _scalar_bits(result.kl), + "per_sequence_policy_bits": _vector_bits(result.per_sequence_policy), + "per_sequence_kl_bits": _vector_bits(result.per_sequence_kl), + "per_sequence_counts": result.per_sequence_active_tokens.cpu().tolist(), + "advantage_bits": _vector_bits(result.advantages), + "global_active": result.provenance["global_active_tokens"], + "backend_id": result.provenance["backend_id"], + "grad_nonzero": bool(grad.abs().sum().item() > 0.0), + # Gradient of one fixed global token slot, keyed by vocab shard + # so the parent can reassemble the full row across TP ranks. + "grad_row_bits": _vector_bits(grad[0]) if dp_rank == 0 else None, + "vocab_start": contract.logprob.sharding.local_vocab_start, + } + ) + except BaseException as exc: # the parent re-raises the text + payload.update( + { + "ok": False, + "error": f"{type(exc).__name__}: {exc}", + "tb": traceback.format_exc(), + } + ) + finally: + try: + import torch.distributed as dist + + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + except Exception: # pragma: no cover - teardown best effort + pass + try: + result_queue.put(payload) + except Exception: # pragma: no cover - queue already closed + pass + + +def _run_mesh(tp: int, dp: int, *, scenario: str = "correctness") -> list[dict]: + """Spawn a ``tp * dp`` NCCL mesh and collect one payload per rank.""" + + world_size = tp * dp + os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + ctx = mp.get_context("spawn") + result_queue = ctx.Queue() + with tempfile.TemporaryDirectory() as tmp: + init_method = f"file://{Path(tmp) / 'store'}" + spawned = mp.spawn( + _mesh_worker, + args=(world_size, init_method, result_queue, tp, dp, scenario), + nprocs=world_size, + join=False, + ) + payloads: list[dict] = [] + try: + for _ in range(world_size): + payloads.append(result_queue.get(timeout=_SPAWN_TIMEOUT_S)) + except queue.Empty: # pragma: no cover - only on a genuine hang + pytest.fail( + f"TP={tp} DP={dp}: only {len(payloads)}/{world_size} ranks reported " + f"within {_SPAWN_TIMEOUT_S}s" + ) + finally: + spawned.join(timeout=_SPAWN_TIMEOUT_S) + return sorted(payloads, key=lambda item: item["rank"]) + + +def _require_all_ok(payloads: list[dict]) -> None: + failures = [item for item in payloads if not item.get("ok")] + if failures: + head = failures[0] + pytest.fail(f"rank {head['rank']} failed: {head['error']}\n{head['tb']}") + + +def _assemble_grad_row(payloads: list[dict]) -> list[int]: + """Reassemble the fixed token slot's gradient row across TP vocab shards.""" + + contributions = [item for item in payloads if item["grad_row_bits"] is not None] + contributions.sort(key=lambda item: item["vocab_start"]) + row: list[int] = [] + for item in contributions: + row.extend(item["grad_row_bits"]) + return row + + +def _consensus(payloads: list[dict]) -> dict: + """Collapse the mesh's per-rank payloads into the one replicated answer.""" + + _require_all_ok(payloads) + replicated = ( + "loss_bits", + "policy_bits", + "kl_bits", + "per_sequence_policy_bits", + "per_sequence_kl_bits", + "per_sequence_counts", + "global_active", + "advantage_bits", + ) + for key in replicated: + values = {repr(item[key]) for item in payloads} + assert len(values) == 1, f"ranks disagree on {key}: {values}" + assert all(item["grad_nonzero"] for item in payloads), ( + "some rank produced a zero gradient; the all-gather likely severed the " + "graph for that rank's own block" + ) + head = payloads[0] + consensus = {key: head[key] for key in replicated} + consensus["grad_row"] = _assemble_grad_row(payloads) + return consensus + + +# --------------------------------------------------------------------------- # +# Single-rank behaviour +# --------------------------------------------------------------------------- # +class TestSingleRank: + def test_forward_and_backward_run(self): + result, tensors, contract = _run_single_rank() + result.loss.backward() + assert result.loss.dtype is torch.float32 + assert result.loss.shape == () + assert tensors["policy"].grad.abs().sum().item() > 0.0 + assert result.provenance["backend_id"] == BACKEND_ID + assert result.provenance["global_active_tokens"] == sum(GLOBAL_ACTIVE_MASK) + + def test_unpacks_as_the_legacy_triple(self): + result, _, _ = _run_single_rank() + loss, policy_loss, kl = result + assert _scalar_bits(loss) == _scalar_bits(result.loss) + assert _scalar_bits(policy_loss) == _scalar_bits(result.policy_loss) + assert _scalar_bits(kl) == _scalar_bits(result.kl) + + def test_run_to_run_bitwise_stability(self): + first, _, _ = _run_single_rank() + second, _, _ = _run_single_rank() + assert _scalar_bits(first.loss) == _scalar_bits(second.loss) + assert _scalar_bits(first.policy_loss) == _scalar_bits(second.policy_loss) + assert _scalar_bits(first.kl) == _scalar_bits(second.kl) + + def test_advantages_are_group_centred(self): + result, _, _ = _run_single_rank() + advantages = result.advantages + for start, end in zip(GROUP_BOUNDARIES[:-1], GROUP_BOUNDARIES[1:], strict=True): + assert advantages[start:end].sum().item() == pytest.approx(0.0, abs=1e-5) + + def test_inactive_tokens_do_not_affect_the_loss(self): + # Inactive slots carry real logits here, so a backend that forgot to + # mask them would shift the loss rather than merely change a padding. + baseline, _, _ = _run_single_rank() + perturbed_globals = _global_inputs() + inactive = [i for i, flag in enumerate(GLOBAL_ACTIVE_MASK) if not flag] + perturbed_globals["policy"][inactive] += 7.5 + perturbed_globals["old_logps"][inactive] -= 3.25 + perturbed, _, _ = _run_single_rank(globals_=perturbed_globals) + assert _scalar_bits(perturbed.loss) == _scalar_bits(baseline.loss) + + def test_dispatch_resolves_this_backend(self): + from rl_engine.kernels.registry import kernel_registry + + _, _, contract = _run_single_rank() + dispatched = kernel_registry.get_loss_op(contract) + assert dispatched.capability.backend_id == BACKEND_ID + assert isinstance(dispatched.op, DistributedGRPOLossOp) + assert dispatched.provenance["fallback"] is False + + +class TestKLZeroIdentity: + """Acceptance criterion: the reference-equals-policy identity is exact.""" + + def _identity_run(self, **overrides): + # old_logps set to the operator's own selected logprob, so the ratio is + # exp(0) = 1 exactly rather than approximately. + contract, tensors = _single_rank_setup(**overrides) + op = DistributedGRPOLossOp() + with torch.no_grad(): + logp, _ = op._logprob.apply( + tensors["policy"], + tensors["action_ids"], + contract=contract.logprob, + num_vocab_tiles=NUM_VOCAB_TILES, + ) + result = op.apply( + tensors["policy"], + tensors["action_ids"], + logp, + tensors["rewards"], + contract=contract, + ref_local_logits=tensors["policy"].detach(), + num_vocab_tiles=NUM_VOCAB_TILES, + ) + return result, tensors + + def test_kl_is_exactly_zero(self): + result, _ = self._identity_run() + assert _scalar_bits(result.kl) == _scalar_bits(torch.zeros(())) + + def test_loss_reduces_to_the_policy_term(self): + result, _ = self._identity_run() + assert _scalar_bits(result.loss) == _scalar_bits(result.policy_loss) + + def test_ratio_is_exactly_one_so_clipping_cannot_bind(self): + # A ratio that is only approximately 1 would land outside a sufficiently + # tight clip range and change the answer; an exact 1 cannot. + tight, _ = self._identity_run( + objective=ObjectiveSpec(beta=BETA, clip=ClipSpec(clip_eps_low=1e-7, clip_eps_high=1e-7)) + ) + loose, _ = self._identity_run( + objective=ObjectiveSpec(beta=BETA, clip=ClipSpec(clip_eps_low=0.9, clip_eps_high=0.9)) + ) + assert _scalar_bits(tight.loss) == _scalar_bits(loose.loss) + + def test_policy_term_matches_the_masked_advantage_mean(self): + result, _ = self._identity_run() + advantages = result.advantages + mask = torch.tensor(GLOBAL_ACTIVE_MASK) + per_token = advantages.reshape(-1, 1).expand(NUM_SEQUENCES, PADDED_SEQ_LEN).reshape(-1) + expected = -per_token.masked_fill(~mask, 0.0).sum() / mask.sum() + # Not a bitwise comparison: this reference sums the flat [N] vector, + # while the operator sums through its fixed tile structure. + assert result.policy_loss.item() == pytest.approx(expected.item(), abs=1e-6) + + +class TestNormalizers: + def test_global_active_tokens_matches_a_flat_masked_mean(self): + result, _, _ = _run_single_rank() + assert result.provenance["global_active_tokens"] == sum(GLOBAL_ACTIVE_MASK) + + def test_normalizers_disagree_on_unequal_sequence_lengths(self): + # If these ever agreed, the normalizer choice would be untestable and + # the sequence lengths in this file would have stopped being staggered. + token_mean, _, _ = _run_single_rank() + seq_mean, _, _ = _run_single_rank( + reduction=LossReductionSpec(token_normalizer=TokenNormalizer.PER_SEQUENCE_THEN_MEAN) + ) + fixed, _, _ = _run_single_rank( + reduction=LossReductionSpec( + token_normalizer=TokenNormalizer.FIXED_CONSTANT, + fixed_normalizer_constant=NUM_TOKEN_SLOTS, + ) + ) + values = { + _scalar_bits(token_mean.loss), + _scalar_bits(seq_mean.loss), + _scalar_bits(fixed.loss), + } + assert len(values) == 3 + + def test_fixed_constant_normalizer_scales_the_token_sum(self): + active = sum(GLOBAL_ACTIVE_MASK) + token_mean, _, _ = _run_single_rank() + fixed, _, _ = _run_single_rank( + reduction=LossReductionSpec( + token_normalizer=TokenNormalizer.FIXED_CONSTANT, + fixed_normalizer_constant=NUM_TOKEN_SLOTS, + ) + ) + assert fixed.policy_loss.item() == pytest.approx( + token_mean.policy_loss.item() * active / NUM_TOKEN_SLOTS, rel=1e-6 + ) + + def test_kl_estimators_differ(self): + k3, _, _ = _run_single_rank() + k1, _, _ = _run_single_rank( + objective=ObjectiveSpec(beta=BETA, kl_estimator=KLEstimator.K1_LOG_RATIO) + ) + assert _scalar_bits(k3.kl) != _scalar_bits(k1.kl) + # k3 is non-negative by construction; the plain log-ratio is not. + assert k3.kl.item() >= 0.0 + + def test_mean_only_advantage_skips_the_std_divisor(self): + std_normalized, _, _ = _run_single_rank() + mean_only, _, _ = _run_single_rank( + objective=ObjectiveSpec( + beta=BETA, + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY), + ) + ) + assert _scalar_bits(std_normalized.loss) != _scalar_bits(mean_only.loss) + + +class TestNegativeControls: + """Prove the bitwise assertions elsewhere are not comparing constants. + + Each control changes something that genuinely regroups or rescales the + reduction and asserts the compared surface notices. If one of these ever + starts passing trivially, the corresponding positive assertion has stopped + meaning anything. + """ + + def test_regrouping_the_token_sum_moves_the_per_sequence_vector(self): + # Split each sequence's token sum in half before adding, changing the + # summation tree without changing a single input value. The scalar loss + # absorbs that most of the time; the per-sequence vector does not, which + # is why it is the comparison surface everywhere else in this file. + import rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss as module + + baseline, _, _ = _run_single_rank() + original = module._sequence_totals + + def halved(values, contract): + view = values.reshape( + contract.sharding.local_num_sequences, 2, contract.sharding.padded_seq_len // 2 + ) + return view.sum(dim=2).sum(dim=1) + + module._sequence_totals = halved + try: + regrouped, _, _ = _run_single_rank() + finally: + module._sequence_totals = original + assert _vector_bits(baseline.per_sequence_policy) != _vector_bits( + regrouped.per_sequence_policy + ) + + def test_vocab_tile_count_perturbs_the_logprob_it_consumes(self): + baseline, _, _ = _run_single_rank() + retiled, _, _ = _run_single_rank(num_vocab_tiles=NUM_VOCAB_TILES * 2) + assert _reduction_fingerprint(baseline) != _reduction_fingerprint(retiled) + + def test_sequence_order_matters_to_the_scalar(self): + # An all_reduce would combine sequence totals in topology order rather + # than in global sequence order. Permuting the assembled grid is a + # stand-in for that mistake, and the loss must notice. + import rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss as module + + baseline, _, _ = _run_single_rank() + permutation = torch.tensor([5, 2, 7, 0, 3, 6, 1, 4]) + original = module._assemble_global_vector + module._assemble_global_vector = lambda totals, contract, group: original( + totals, contract, group + )[permutation] + try: + permuted, _, _ = _run_single_rank() + finally: + module._assemble_global_vector = original + assert _vector_bits(baseline.per_sequence_policy) != _vector_bits( + permuted.per_sequence_policy + ) + + def test_clip_epsilon_perturbs_the_loss(self): + baseline, _, _ = _run_single_rank() + clipped, _, _ = _run_single_rank( + objective=ObjectiveSpec(beta=BETA, clip=ClipSpec(clip_eps_low=0.01, clip_eps_high=0.01)) + ) + assert _reduction_fingerprint(baseline) != _reduction_fingerprint(clipped) + + +class TestGuards: + def test_reference_logits_required_when_beta_is_positive(self): + contract, tensors = _single_rank_setup() + with pytest.raises(LossContractError, match="ref_local_logits is required"): + DistributedGRPOLossOp().apply( + tensors["policy"], + tensors["action_ids"], + tensors["old_logps"], + tensors["rewards"], + contract=contract, + num_vocab_tiles=NUM_VOCAB_TILES, + ) + + def test_reference_optional_when_beta_is_zero(self): + result, _, _ = _run_single_rank(objective=ObjectiveSpec(beta=0.0), with_reference=False) + assert _scalar_bits(result.loss) == _scalar_bits(result.policy_loss) + assert _scalar_bits(result.kl) == _scalar_bits(torch.zeros(())) + assert result.provenance["reference_model_used"] is False + + def test_row_count_must_match_owned_slots(self): + contract, tensors = _single_rank_setup() + with pytest.raises(LossContractError, match="token slots"): + DistributedGRPOLossOp().apply( + tensors["policy"][:-1], + tensors["action_ids"][:-1], + tensors["old_logps"][:-1], + tensors["rewards"], + contract=contract, + ref_local_logits=tensors["ref"][:-1], + num_vocab_tiles=NUM_VOCAB_TILES, + ) + + def test_reward_count_must_match_owned_sequences(self): + contract, tensors = _single_rank_setup() + with pytest.raises(LossContractError, match="one entry per sequence"): + DistributedGRPOLossOp().apply( + tensors["policy"], + tensors["action_ids"], + tensors["old_logps"], + tensors["rewards"][:-1], + contract=contract, + ref_local_logits=tensors["ref"], + num_vocab_tiles=NUM_VOCAB_TILES, + ) + + +# --------------------------------------------------------------------------- # +# The claim: bitwise equality across the mesh +# --------------------------------------------------------------------------- # +# Every reachable (TP, DP) combination with at most 4 ranks. Larger degrees +# need a bigger node and run unchanged there via the same helper. +MESH_CONFIGS = [ + pytest.param(2, 1, id="tp2"), + pytest.param(4, 1, id="tp4"), + pytest.param(1, 2, id="dp2"), + pytest.param(1, 4, id="dp4"), + pytest.param(2, 2, id="tp2xdp2"), +] + + +@pytest.fixture(scope="module") +def gpu_baseline() -> dict: + """The TP=DP=1 answer, computed on one GPU so the mesh comparison is + device-for-device rather than CPU-versus-GPU. Built once for the module.""" + + if _cuda_device_count() < 1: + pytest.skip("needs a CUDA device") + return _consensus(_run_mesh(1, 1)) + + +@_requires_gpus(1) +class TestMeshBitwise: + @pytest.mark.parametrize(("tp", "dp"), MESH_CONFIGS) + def test_mesh_matches_the_single_rank_baseline(self, tp, dp, gpu_baseline): + world = tp * dp + if _cuda_device_count() < world: + pytest.skip(f"needs {world} CUDA devices for TP={tp} DP={dp}") + baseline = gpu_baseline + actual = _consensus(_run_mesh(tp, dp)) + label = f"TP={tp} DP={dp}" + + assert actual["global_active"] == baseline["global_active"] + assert actual["per_sequence_counts"] == baseline["per_sequence_counts"] + assert actual["advantage_bits"] == baseline["advantage_bits"] + # The sensitive comparison: per-sequence totals, before the scalar + # average rounds a reordering away. + assert ( + actual["per_sequence_policy_bits"] == baseline["per_sequence_policy_bits"] + ), f"{label} per-sequence policy totals differ from the single-rank baseline" + assert ( + actual["per_sequence_kl_bits"] == baseline["per_sequence_kl_bits"] + ), f"{label} per-sequence KL totals differ from the single-rank baseline" + assert actual["loss_bits"] == baseline["loss_bits"], f"{label} loss differs" + assert actual["policy_bits"] == baseline["policy_bits"] + assert actual["kl_bits"] == baseline["kl_bits"] + assert ( + actual["grad_row"] == baseline["grad_row"] + ), f"{label} gradient differs from the single-rank baseline" + + @_requires_gpus(4) + def test_pure_axes_agree_with_each_other(self): + # Transitively implied by both matching the baseline; kept because a + # direct TP-vs-DP comparison names the culprit when a shared drift moves + # both away from the baseline at once. + tp4 = _consensus(_run_mesh(4, 1)) + dp4 = _consensus(_run_mesh(1, 4)) + assert tp4["per_sequence_policy_bits"] == dp4["per_sequence_policy_bits"] + assert tp4["loss_bits"] == dp4["loss_bits"] + assert tp4["grad_row"] == dp4["grad_row"] + + +@_requires_gpus(2) +class TestMeshGuards: + def test_preflight_rejects_a_dp_rank_that_disagrees(self): + payloads = _run_mesh(1, 2, scenario="preflight_dp_mismatch") + assert all( + not item.get("ok") for item in payloads + ), "every rank must abort when one of them declares a different objective" + assert any("preflight" in item.get("error", "") for item in payloads) + + def test_preflight_rejects_a_tp_rank_that_disagrees(self): + # beta is not part of the logprob contract, so the logprob path's own TP + # preflight cannot see this; only the loss preflight's TP-axis check can. + # Without it two TP siblings would compute different losses for one + # sharded model and nothing would notice. + payloads = _run_mesh(2, 1, scenario="preflight_tp_mismatch") + assert all(not item.get("ok") for item in payloads) + assert any("TP axis" in item.get("error", "") for item in payloads) diff --git a/tests/test_distributed_logprob_comparison.py b/tests/test_distributed_logprob_comparison.py new file mode 100644 index 00000000..887b53bd --- /dev/null +++ b/tests/test_distributed_logprob_comparison.py @@ -0,0 +1,384 @@ +# 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 + + +@pytest.mark.skipif(torch.version.hip is None, reason="requires a ROCm PyTorch build") +@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires four ROCm GPUs") +def test_tp2_cp2_rocm_native_cli_emits_strict_report(tmp_path): + """Run the production ROCm backend across TP2 x CP2 and require provenance.""" + + from rl_engine.kernels.registry import _rocm_vocab_logprob_native_available + + if not _rocm_vocab_logprob_native_available(): + pytest.skip("requires the compiled ROCm logprob extension") + + script = ( + Path(__file__).resolve().parents[1] + / "rl_engine" + / "testing" + / "distributed_logprob_comparison.py" + ) + output = tmp_path / "tp2-cp2-rocm-native.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", + "cuda", + "--dist-backend", + "nccl", + "--backend", + "rocm-vocab-parallel-logp-ws2", + "--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=180, + env=environment, + ) + payload = json.loads(output.read_text(encoding="utf-8")) + assert payload["passed"] + assert {rank["actual_backend"] for rank in payload["ranks"]} == {"rocm-vocab-parallel-logp-ws2"} + assert {rank["fallback"] for rank in payload["ranks"]} == {False} + assert {rank["contract"]["sharding"]["cp_world_size"] for rank in payload["ranks"]} == {2} + assert json.loads(result.stdout)["case"]["cp_world_size"] == 2 diff --git a/tests/test_grpo_loss_contract.py b/tests/test_grpo_loss_contract.py new file mode 100644 index 00000000..2a17bce4 --- /dev/null +++ b/tests/test_grpo_loss_contract.py @@ -0,0 +1,392 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Unit tests for the WS2 deterministic GRPO loss contract. + +These are pure-Python contract checks: no tensors, no collectives, no GPU. +The distributed behaviour they describe is exercised in +``tests/test_distributed_grpo_loss.py``. +""" + +from __future__ import annotations + +import json + +import pytest + +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobContract, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.loss_contract import ( + AdvantageNormalizer, + AdvantageSpec, + ClipSpec, + GRPOLossContract, + KLEstimator, + LossBackendCapability, + LossContractError, + LossReductionSpec, + LossShardingSpec, + ObjectiveSpec, + TokenNormalizer, +) + +# Global batch geometry shared by every fixture below: 8 sequences of 8 token +# slots, in two equal advantage groups. Tests that care about variable or +# straddling groups override group_boundaries explicitly. +NUM_SEQUENCES = 8 +PADDED_SEQ_LEN = 8 +GROUP_BOUNDARIES = (0, 4, NUM_SEQUENCES) +REAL_VOCAB = 30 +PADDED_VOCAB = 32 + + +def _dp_bounds(dp: int) -> tuple[tuple[int, int], ...]: + """Contiguous sequence partition in DP-rank order.""" + + seqs = NUM_SEQUENCES // dp + return tuple((d * seqs, (d + 1) * seqs) for d in range(dp)) + + +def _logprob_contract(num_tokens: int, **overrides) -> LogprobContract: + kwargs = { + "role": LogprobRole.TRAIN, + "dtype": LogprobDType.FP32, + "mask": MaskSpec(num_tokens=num_tokens, active_mask=(True,) * num_tokens), + "sharding": ShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, PADDED_VOCAB),), + real_vocab_size=REAL_VOCAB, + padded_vocab_size=PADDED_VOCAB, + ), + "reduction": ReductionSpec(), + } + kwargs.update(overrides) + return LogprobContract(**kwargs) + + +def _sharding(*, dp_rank: int = 0, dp: int = 1, **overrides) -> LossShardingSpec: + kwargs = { + "dp_rank": dp_rank, + "dp_world_size": dp, + "num_sequences": NUM_SEQUENCES, + "padded_seq_len": PADDED_SEQ_LEN, + "sequence_shard_bounds": _dp_bounds(dp), + "group_boundaries": GROUP_BOUNDARIES, + } + kwargs.update(overrides) + return LossShardingSpec(**kwargs) + + +def _contract(*, dp_rank: int = 0, dp: int = 1, **overrides) -> GRPOLossContract: + sharding = overrides.pop("sharding", _sharding(dp_rank=dp_rank, dp=dp)) + kwargs = { + "logprob": _logprob_contract(sharding.local_num_token_slots), + "sharding": sharding, + "objective": ObjectiveSpec(), + "reduction": LossReductionSpec(), + } + kwargs.update(overrides) + return GRPOLossContract(**kwargs) + + +class TestSequenceOwnership: + def test_single_rank_owns_every_sequence(self): + sharding = _sharding() + assert sharding.local_sequence_start == 0 + assert sharding.local_sequence_end == NUM_SEQUENCES + assert sharding.local_num_token_slots == NUM_SEQUENCES * PADDED_SEQ_LEN + assert sharding.num_groups == 2 + assert sharding.group_sizes == (4, 4) + + @pytest.mark.parametrize("dp", [1, 2, 4, 8]) + def test_dp_shapes_partition_the_batch(self, dp): + total = 0 + for dp_rank in range(dp): + sharding = _sharding(dp_rank=dp_rank, dp=dp) + assert sharding.local_num_sequences == NUM_SEQUENCES // dp + total += sharding.local_num_token_slots + assert total == NUM_SEQUENCES * PADDED_SEQ_LEN + + def test_bounds_must_be_contiguous_in_rank_order(self): + with pytest.raises(LossContractError, match="contiguous"): + _sharding(dp=2, sequence_shard_bounds=((0, 4), (5, 8))) + + def test_bounds_must_cover_every_sequence(self): + with pytest.raises(LossContractError, match="cover num_sequences"): + _sharding(dp=2, sequence_shard_bounds=((0, 3), (3, 6))) + + def test_bounds_count_must_match_dp_world_size(self): + with pytest.raises( + LossContractError, match="exactly one \\(start, end\\) pair per DP rank" + ): + _sharding(dp=2, sequence_shard_bounds=((0, NUM_SEQUENCES),)) + + def test_empty_shard_rejected(self): + with pytest.raises(LossContractError, match="end > start"): + _sharding(dp=2, sequence_shard_bounds=((0, 0), (0, NUM_SEQUENCES))) + + def test_context_parallelism_is_rejected(self): + # CP splits a sequence's tokens across ranks, which this contract does + # not model; it must fail loudly rather than silently drop the rest. + with pytest.raises(LossContractError, match="cp_world_size=2 is unsupported"): + _sharding(cp_world_size=2) + + def test_cp_rank_must_be_zero(self): + with pytest.raises(LossContractError, match="cp_rank must be 0"): + _sharding(cp_rank=1) + + +class TestGroupBoundaries: + @pytest.mark.parametrize( + ("boundaries", "match"), + [ + ((1, NUM_SEQUENCES), "must start at 0"), + ((0, 3), "must start at 0"), + ((0, 2, 2, NUM_SEQUENCES), "strictly increasing"), + ((0,), "at least 2 entries"), + ], + ) + def test_malformed_boundaries_rejected(self, boundaries, match): + with pytest.raises(LossContractError, match=match): + _sharding(group_boundaries=boundaries) + + def test_variable_group_sizes_accepted(self): + sharding = _sharding(group_boundaries=(0, 7, NUM_SEQUENCES)) + assert sharding.group_sizes == (7, 1) + + def test_groups_may_straddle_dp_shards(self): + # A group split at 3 crosses the DP=2 shard boundary at 4, so no rank + # owns that group alone. The contract permits it; that is why the + # operator replicates advantages instead of merging partial statistics. + sharding = _sharding(dp=2, dp_rank=0, group_boundaries=(0, 3, NUM_SEQUENCES)) + assert sharding.sequence_shard_bounds == ((0, 4), (4, NUM_SEQUENCES)) + assert sharding.group_sizes == (3, 5) + + def test_population_std_rejects_singleton_groups(self): + # A one-sequence group has zero population variance, so its advantage + # would silently collapse to zero rather than fail. + with pytest.raises(LossContractError, match="at least 2 sequences per group"): + _contract( + sharding=_sharding(group_boundaries=(0, 7, NUM_SEQUENCES)), + objective=ObjectiveSpec(), + ) + + def test_mean_only_allows_singleton_groups(self): + contract = _contract( + sharding=_sharding(group_boundaries=(0, 7, NUM_SEQUENCES)), + objective=ObjectiveSpec( + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY) + ), + ) + assert contract.sharding.group_sizes == (7, 1) + + +class TestSpecValidation: + def test_accumulation_must_be_fp32(self): + with pytest.raises(LossContractError, match="must be fp32"): + LossReductionSpec(acc_dtype=LogprobDType.BF16) + + def test_fixed_constant_normalizer_requires_its_constant(self): + with pytest.raises(LossContractError, match="fixed_normalizer_constant"): + LossReductionSpec(token_normalizer=TokenNormalizer.FIXED_CONSTANT) + + def test_constant_rejected_for_other_normalizers(self): + with pytest.raises(LossContractError, match="only meaningful for"): + LossReductionSpec(fixed_normalizer_constant=32) + + def test_fixed_constant_normalizer_accepts_its_constant(self): + spec = LossReductionSpec( + token_normalizer=TokenNormalizer.FIXED_CONSTANT, fixed_normalizer_constant=32 + ) + assert spec.fixed_normalizer_constant == 32 + + def test_lower_clip_bound_must_stay_positive(self): + with pytest.raises(LossContractError, match="must be smaller than 1.0"): + ClipSpec(clip_eps_low=1.0) + + def test_asymmetric_clip_bounds(self): + clip = ClipSpec(clip_eps_low=0.2, clip_eps_high=0.28) + assert clip.lower_bound == pytest.approx(0.8) + assert clip.upper_bound == pytest.approx(1.28) + + def test_std_eps_must_be_positive(self): + with pytest.raises(LossContractError, match="strictly positive"): + AdvantageSpec(std_eps=0.0) + + def test_negative_beta_rejected(self): + with pytest.raises(LossContractError, match="non-negative"): + ObjectiveSpec(beta=-0.01) + + def test_uses_reference_model_tracks_beta(self): + assert not ObjectiveSpec(beta=0.0).uses_reference_model + assert ObjectiveSpec(beta=0.04).uses_reference_model + + +class TestContractCoherence: + def test_logprob_token_count_must_match_owned_slots(self): + sharding = _sharding(dp=2, dp_rank=0) + with pytest.raises(LossContractError, match="token slots this rank owns"): + GRPOLossContract( + logprob=_logprob_contract(sharding.local_num_token_slots + 1), + sharding=sharding, + ) + + def test_determinism_scope_must_agree_with_logprob_path(self): + tokens = _sharding().local_num_token_slots + with pytest.raises(LossContractError, match="stronger or weaker determinism scope"): + GRPOLossContract( + logprob=_logprob_contract( + tokens, + reduction=ReductionSpec(determinism_scope=DeterminismScope.FIXED_TOPOLOGY), + ), + sharding=_sharding(), + ) + + def test_global_token_slots(self): + assert _contract().global_token_slots == NUM_SEQUENCES * PADDED_SEQ_LEN + + +class TestFingerprint: + @pytest.mark.parametrize("dp", [2, 4, 8]) + def test_every_dp_rank_agrees(self, dp): + # This is the property the distributed preflight relies on. Each rank + # holds a different slice of sequences, so their nested logprob masks + # genuinely differ -- the fingerprint must still match. + fingerprints = { + _contract(dp_rank=rank, dp=dp).cross_rank_fingerprint() for rank in range(dp) + } + assert len(fingerprints) == 1 + + def test_dp_degree_changes_the_fingerprint(self): + # A different partition is a different logical invocation, so a rank + # that joined the wrong one must be caught rather than merged with. + assert _contract(dp=1).cross_rank_fingerprint() != _contract(dp=2).cross_rank_fingerprint() + + @pytest.mark.parametrize( + "overrides", + [ + { + "reduction": LossReductionSpec( + token_normalizer=TokenNormalizer.PER_SEQUENCE_THEN_MEAN + ) + }, + {"objective": ObjectiveSpec(beta=0.04)}, + {"objective": ObjectiveSpec(clip=ClipSpec(clip_eps_high=0.28))}, + {"objective": ObjectiveSpec(kl_estimator=KLEstimator.K1_LOG_RATIO)}, + { + "objective": ObjectiveSpec( + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY) + ) + }, + {"sharding": _sharding(group_boundaries=(0, 3, NUM_SEQUENCES))}, + ], + ids=["normalizer", "beta", "clip", "kl", "advantage", "groups"], + ) + def test_numerical_identity_changes_the_fingerprint(self, overrides): + assert ( + _contract().cross_rank_fingerprint() != _contract(**overrides).cross_rank_fingerprint() + ) + + def test_to_dict_is_json_serializable_and_stable(self): + contract = _contract() + first = json.dumps(contract.to_dict(), sort_keys=True) + second = json.dumps(contract.to_dict(), sort_keys=True) + assert first == second + payload = contract.to_dict() + assert payload["semantic_operator"] == "grpo_loss" + assert payload["reduction"]["cp_is_merge_axis"] is False + assert payload["reduction"]["dp_is_merge_axis"] is True + assert payload["logprob"]["reduction"]["cp_is_merge_axis"] is False + + +def _capability(**overrides) -> LossBackendCapability: + kwargs = { + "backend_id": "test-loss-backend", + "token_normalizers": frozenset({TokenNormalizer.GLOBAL_ACTIVE_TOKENS}), + "kl_estimators": frozenset({KLEstimator.K3_UNBIASED}), + "advantage_normalizers": frozenset({AdvantageNormalizer.MEAN_STD_POPULATION}), + "determinism_scopes": frozenset({DeterminismScope.CROSS_TP_BITWISE}), + "implementation_kind": "reference", + } + kwargs.update(overrides) + return LossBackendCapability(**kwargs) + + +class TestBackendCapability: + def test_matching_capability_supports_contract(self): + assert _capability().supports(_contract()) + + def test_backend_id_may_not_shadow_a_dispatch_policy(self): + with pytest.raises(LossContractError, match="reserved dispatch policy"): + _capability(backend_id="reference") + + def test_unsupported_normalizer_is_reported(self): + contract = _contract( + reduction=LossReductionSpec(token_normalizer=TokenNormalizer.PER_SEQUENCE_THEN_MEAN) + ) + reasons = _capability().incompatibilities(contract) + assert any("token_normalizer" in reason for reason in reasons) + + def test_unsupported_dp_degree_is_reported(self): + capability = _capability(dp_world_sizes=(1,)) + reasons = capability.incompatibilities(_contract(dp=2)) + assert any("DP=2" in reason for reason in reasons) + + def test_variable_group_sizes_gated(self): + contract = _contract( + sharding=_sharding(group_boundaries=(0, 7, NUM_SEQUENCES)), + objective=ObjectiveSpec( + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY) + ), + ) + capability = _capability(advantage_normalizers=frozenset({AdvantageNormalizer.MEAN_ONLY})) + assert any( + "variable advantage group sizes" in reason + for reason in capability.incompatibilities(contract) + ) + + def test_variable_group_sizes_allowed_when_declared(self): + contract = _contract( + sharding=_sharding(group_boundaries=(0, 7, NUM_SEQUENCES)), + objective=ObjectiveSpec( + advantage=AdvantageSpec(normalizer=AdvantageNormalizer.MEAN_ONLY) + ), + ) + capability = _capability( + advantage_normalizers=frozenset({AdvantageNormalizer.MEAN_ONLY}), + supports_variable_group_sizes=True, + ) + assert capability.supports(contract) + + def test_asymmetric_clip_gated(self): + contract = _contract(objective=ObjectiveSpec(clip=ClipSpec(clip_eps_high=0.28))) + assert any( + "asymmetric ratio clipping" in reason + for reason in _capability().incompatibilities(contract) + ) + assert _capability(supports_asymmetric_clip=True).supports(contract) + + def test_every_incompatibility_is_reported_at_once(self): + contract = _contract( + dp=2, + reduction=LossReductionSpec(token_normalizer=TokenNormalizer.PER_SEQUENCE_THEN_MEAN), + objective=ObjectiveSpec( + kl_estimator=KLEstimator.K1_LOG_RATIO, + clip=ClipSpec(clip_eps_high=0.28), + ), + ) + reasons = _capability(dp_world_sizes=(1,)).incompatibilities(contract) + assert len(reasons) >= 4 + + def test_to_dict_round_trips_declared_flags(self): + payload = _capability(dp_world_sizes=(1, 2)).to_dict() + assert payload["dp_world_sizes"] == [1, 2] + assert payload["implementation_kind"] == "reference" diff --git a/tests/test_logprob_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 index a961f312..060ddffd 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -248,13 +248,16 @@ def test_ignore_index_must_not_collide_with_the_real_vocabulary(): def _restrict_to_ws1_candidates(registry: KernelRegistry) -> None: - """Drop the #241 PR3 vocab-parallel reference so only WS1 backends remain.""" + """Drop all WS2 vocab-parallel backends so only WS1 backends remain.""" platform = registry._platform() + ws2_backends = { + OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP, + OpBackend.ROCM_VOCAB_PARALLEL_LOGP, + OpBackend.TRITON_VOCAB_PARALLEL_LOGP, + } registry._logprob_candidates[platform] = [ - backend - for backend in registry._logprob_candidates[platform] - if backend is not OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP + backend for backend in registry._logprob_candidates[platform] if backend not in ws2_backends ] @@ -297,7 +300,7 @@ def test_ws1_rejections_recorded_when_vocab_parallel_reference_resolves(): candidates.remove(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) candidates.append(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) - result = registry.get_logprob_op(_contract()) + result = registry.get_logprob_op(_contract(), requested_backend="reference") assert result.capability.backend_id == "pytorch-vocab-parallel-logp-ws2" assert result.provenance["fallback"] is True rejections = " | ".join(result.provenance["prior_rejections"]) diff --git a/tests/test_rocm_logprob_backend.py b/tests/test_rocm_logprob_backend.py new file mode 100644 index 00000000..934600b8 --- /dev/null +++ b/tests/test_rocm_logprob_backend.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest + +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import VocabParallelLogprobOp +from rl_engine.kernels.ops.rocm.loss.vocab_parallel_logp import RocmVocabParallelLogprobOp +from rl_engine.kernels.registry import KernelRegistry, OpBackend + + +def test_rocm_backend_preserves_ws2_operator_surface(): + assert issubclass(RocmVocabParallelLogprobOp, VocabParallelLogprobOp) + assert RocmVocabParallelLogprobOp.op_class == "logprob" + assert RocmVocabParallelLogprobOp.is_batch_invariant + + +def test_rocm_backend_is_gated_by_native_extension(monkeypatch): + registry = KernelRegistry() + registry._platform = lambda: "rocm" + candidates = registry._logprob_candidates["rocm"] + if OpBackend.ROCM_VOCAB_PARALLEL_LOGP in candidates: + assert candidates[0] is OpBackend.ROCM_VOCAB_PARALLEL_LOGP + capability = registry._logprob_capabilities["rocm"][OpBackend.ROCM_VOCAB_PARALLEL_LOGP] + assert capability.backend_id == "rocm-vocab-parallel-logp-ws2" + assert capability.implementation_kind == "production" + else: + assert candidates[0] in { + OpBackend.TRITON_VOCAB_PARALLEL_LOGP, + OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP, + } + + +def test_rocm_native_tile_kernel_is_hip_guarded_and_registered(): + source = (Path(__file__).resolve().parents[1] / "csrc" / "ops.cpp").read_text(encoding="utf-8") + kernel = ( + Path(__file__).resolve().parents[1] / "csrc" / "deterministic_logp_kernel.cu" + ).read_text(encoding="utf-8") + assert "deterministic_logp_tile_stats" in source + assert "deterministic_logp_tile_stats_kernel" in kernel + assert "atomic" not in kernel.lower() + hip_kernel = ( + Path(__file__).resolve().parents[1] / "csrc" / "hip" / "hip_deterministic_logp_kernel.hip" + ).read_text(encoding="utf-8") + assert "hip_deterministic_logp_tile_stats" in hip_kernel + assert "hip_deterministic_logp_backward" in hip_kernel + assert "atomic" not in hip_kernel.lower() + assert "hip_deterministic_logp_backward" in source + assert "__HIPCC__" in source + assert "deterministic_collective_all_gather" in source + + +def test_rocm_backend_import_does_not_require_native_extension(): + # Importing the wrapper must remain possible in CPU-only CI; capability + # loading, not module import, decides whether the native fast path exists. + op = RocmVocabParallelLogprobOp() + assert isinstance(op, VocabParallelLogprobOp) + + +def test_explicit_native_backend_fails_closed_when_extension_is_missing(monkeypatch): + import rl_engine.kernels.registry as registry_module + + monkeypatch.setattr(registry_module, "_rocm_vocab_logprob_native_available", lambda: False) + registry = KernelRegistry() + registry._platform = lambda: "rocm" + contract = LogprobContract( + role="train", + dtype="fp32", + mask=MaskSpec(num_tokens=1, active_mask=(True,)), + sharding=ShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, 8),), + real_vocab_size=8, + padded_vocab_size=8, + ), + reduction=ReductionSpec(), + ) + + with pytest.raises(RuntimeError, match="rocm-vocab-parallel-logp-ws2"): + registry.get_logprob_op( + contract, + requested_backend="rocm-vocab-parallel-logp-ws2", + ) + + +# --------------------------------------------------------------------------- GPU cases + + +def _native_rocm_available() -> bool: + import torch + + if torch.version.hip is None or not torch.cuda.is_available(): + return False + from rl_engine.kernels.registry import _rocm_vocab_logprob_native_available + + return _rocm_vocab_logprob_native_available() + + +def _triton_available() -> bool: + import torch + + if not torch.cuda.is_available(): + return False + try: + import triton # noqa: F401 + except ImportError: + return False + return True + + +def _kernel_backends() -> list: + """(id, op factory) for every fused kernel backend usable on this machine.""" + + backends = [] + if _native_rocm_available(): + backends.append(pytest.param(RocmVocabParallelLogprobOp, id="rocm-hip")) + if _triton_available(): + from rl_engine.kernels.ops.triton.loss.vocab_parallel_logp import ( + TritonVocabParallelLogprobOp, + ) + + backends.append(pytest.param(TritonVocabParallelLogprobOp, id="triton")) + return backends + + +@pytest.mark.skipif(not _kernel_backends(), reason="requires a fused WS2 kernel backend") +@pytest.mark.parametrize("op_class", _kernel_backends()) +class TestFusedKernelPath: + """Every fused kernel backend must agree with the reference op on the same contract.""" + + @staticmethod + def _case(real_vocab, padded_vocab, num_tokens, dtype, *, seed=3): + import torch + + device = torch.device("cuda", 0) + gen = torch.Generator(device="cpu").manual_seed(seed) + logits = (torch.randn(num_tokens, padded_vocab, generator=gen) * 3).to(device, dtype) + targets = torch.randint(0, real_vocab, (num_tokens,), generator=gen).to(device) + active = tuple(i % 4 != 2 for i in range(num_tokens)) + contract = LogprobContract( + role="train", + dtype={torch.bfloat16: "bf16", torch.float32: "fp32", torch.float16: "fp16"}[dtype], + mask=MaskSpec(num_tokens=num_tokens, active_mask=active), + sharding=ShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, padded_vocab),), + real_vocab_size=real_vocab, + padded_vocab_size=padded_vocab, + ), + reduction=ReductionSpec(), + ) + return logits, targets, torch.tensor(active, device=device), contract + + @pytest.mark.parametrize("dtype_name", ["fp32", "bf16", "fp16"]) + @pytest.mark.parametrize( + "real_vocab,padded_vocab,tiles", + [(1000, 1024, 32), (13, 16, 8), (151936, 151936, 64)], + ids=["partial-pad", "full-pad-tile", "qwen3"], + ) + def test_forward_backward_match_reference( + self, op_class, dtype_name, real_vocab, padded_vocab, tiles + ): + import torch + + dtype = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}[dtype_name] + logits, targets, active, contract = self._case(real_vocab, padded_vocab, 24, dtype) + ref_op, rocm_op = VocabParallelLogprobOp(), op_class() + outputs = {} + for name, op in (("ref", ref_op), ("rocm", rocm_op)): + leaf = logits.clone().requires_grad_(True) + logp, lse = op.apply(leaf, targets, contract=contract, num_vocab_tiles=tiles) + ((logp * active).sum() + 0.25 * lse.sum()).backward() + outputs[name] = (logp.detach(), lse.detach(), leaf.grad.detach()) + for ref, rocm in zip(outputs["ref"], outputs["rocm"]): + assert torch.isfinite(rocm).all() + torch.testing.assert_close(rocm.float(), ref.float(), rtol=2e-5, atol=2e-5) + # Padding columns never receive gradient. + grad = outputs["rocm"][2] + assert torch.equal(grad[:, real_vocab:], torch.zeros_like(grad[:, real_vocab:])) + # Inactive rows only carry the LSE gradient. + inactive = ~active + if inactive.any(): + row = inactive.nonzero()[0, 0] + p = torch.softmax(logits[row, :real_vocab].float(), dim=-1) + torch.testing.assert_close( + grad[row, :real_vocab].float(), 0.25 * p, rtol=2e-2, atol=2e-5 + ) + # Repeat is bitwise. + again = rocm_op.apply(logits, targets, contract=contract, num_vocab_tiles=tiles) + assert torch.equal(again[0], outputs["rocm"][0]) and torch.equal( + again[1], outputs["rocm"][1] + ) + + def test_logp_only_and_lse_only_gradients(self, op_class): + import torch + + logits, targets, active, contract = self._case(1000, 1024, 12, torch.bfloat16) + ref_op, rocm_op = VocabParallelLogprobOp(), op_class() + for which in ("logp", "lse"): + grads = [] + for op in (ref_op, rocm_op): + leaf = logits.clone().requires_grad_(True) + logp, lse = op.apply(leaf, targets, contract=contract, num_vocab_tiles=32) + ((logp * active).sum() if which == "logp" else lse.sum()).backward() + grads.append(leaf.grad.float()) + torch.testing.assert_close(grads[1], grads[0], rtol=2e-5, atol=2e-5) + + def test_tile_stats_read_input_dtype_exactly(self, op_class): + """BF16 input straight into the kernel equals the FP32 upcast path bitwise.""" + import torch + + from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import _local_tile_stats + + if op_class is RocmVocabParallelLogprobOp: + from rl_engine.kernels.ops.base import _C + + tile_stats = _C.hip_deterministic_logp_tile_stats + else: + from rl_engine.kernels.ops.triton.loss.vocab_parallel_logp import ( + triton_vocab_tile_stats, + ) + + tile_stats = triton_vocab_tile_stats + logits, _, _, _ = self._case(1000, 1024, 16, torch.bfloat16) + direct = tile_stats(logits, 0, 1000, 32) + upcast = tile_stats(logits.float(), 0, 1000, 32) + assert all(torch.equal(a, b) for a, b in zip(direct, upcast)) + # Same tile maxima as the PyTorch tile loop; sums differ only by order. + ids = torch.arange(1024, device=logits.device) + masked = logits.float().masked_fill((ids >= 1000).unsqueeze(0), float("-inf")) + ref_max, ref_sum = _local_tile_stats(masked, 32) + assert torch.equal(direct[0], ref_max) + torch.testing.assert_close(direct[1], ref_sum, rtol=1e-6, atol=0.0) + # A tile that is entirely padding is the (-inf, 0) identity partial. + tail = tile_stats(logits, 1024, 1000, 32) + assert torch.isneginf(tail[0]).all() and torch.equal(tail[1], torch.zeros_like(tail[1])) + + def test_entropy_path_still_available(self, op_class): + import torch + + logits, targets, active, contract = self._case(1000, 1024, 8, torch.float32) + ref_op, rocm_op = VocabParallelLogprobOp(), op_class() + ref = ref_op.apply_with_entropy(logits, targets, contract=contract, num_vocab_tiles=32) + rocm = rocm_op.apply_with_entropy(logits, targets, contract=contract, num_vocab_tiles=32) + for a, b in zip(ref, rocm): + torch.testing.assert_close(b, a, rtol=1e-5, atol=1e-5) + + +def test_triton_backend_registered_where_triton_runs(): + registry = KernelRegistry() + import rl_engine.kernels.registry as registry_module + + for platform in ("cuda", "rocm"): + candidates = registry._logprob_candidates[platform] + if registry_module._triton_vocab_logprob_available(): + assert OpBackend.TRITON_VOCAB_PARALLEL_LOGP in candidates + capability = registry._logprob_capabilities[platform][ + OpBackend.TRITON_VOCAB_PARALLEL_LOGP + ] + assert capability.backend_id == "triton-vocab-parallel-logp-ws2" + assert capability.implementation_kind == "production" + # The reference never outranks a production kernel backend. + assert candidates.index(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) > candidates.index( + OpBackend.TRITON_VOCAB_PARALLEL_LOGP + ) + else: + assert OpBackend.TRITON_VOCAB_PARALLEL_LOGP not in candidates + assert OpBackend.TRITON_VOCAB_PARALLEL_LOGP not in registry._logprob_candidates["cpu"] 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_vocab_parallel_logp.py b/tests/test_vocab_parallel_logp.py index eefc1d6a..a9033362 100644 --- a/tests/test_vocab_parallel_logp.py +++ b/tests/test_vocab_parallel_logp.py @@ -1,13 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Deterministic vocab-parallel TP logprob reference tests (issue #241 PR3). - -Bit-level determinism assertions compare raw bit patterns via -``tensor.view(torch.int32)`` rather than ``torch.equal``: value equality -treats ``-0.0 == 0.0`` as equal and ``NaN != NaN`` as different, neither of -which is what a bitwise claim means. -""" +"""Deterministic vocab-parallel TP logprob reference tests""" from __future__ import annotations @@ -34,6 +28,7 @@ BACKEND_ID, VocabParallelLogprobOp, ) +from rl_engine.kernels.ops.rocm.loss.vocab_parallel_logp import RocmVocabParallelLogprobOp from rl_engine.kernels.registry import KernelRegistry, OpBackend REAL_VOCAB = 27 @@ -60,7 +55,6 @@ def _contract( num_tokens: int = NUM_TOKENS, active: tuple[bool, ...] = ACTIVE, dtype: str = "fp32", - determinism_scope: str = "cross_tp_bitwise", ) -> LogprobContract: return LogprobContract( role="train", @@ -75,7 +69,7 @@ def _contract( real_vocab_size=real_vocab, padded_vocab_size=padded_vocab, ), - reduction=ReductionSpec(determinism_scope=determinism_scope), + reduction=ReductionSpec(), ) @@ -87,7 +81,11 @@ def _inputs(dtype=torch.float32, seed: int = 2026): def _bits(tensor: torch.Tensor) -> torch.Tensor: - view_dtype = {torch.float32: torch.int32, torch.bfloat16: torch.int16}[tensor.dtype] + view_dtype = { + torch.float32: torch.int32, + torch.bfloat16: torch.int16, + torch.float16: torch.int16, + }[tensor.dtype] return tensor.contiguous().view(view_dtype) @@ -216,63 +214,6 @@ def test_inactive_rows_zero_filled_lse_still_exported(self): assert torch.isfinite(lse[-1]) -class TestNonDeterministicPath: - """deterministic=False: the fast whole-shard reduction with no guarantee.""" - - def test_rejects_a_cross_tp_bitwise_contract(self): - logits, targets = _inputs() - with pytest.raises(LogprobContractError, match="cross_tp_bitwise"): - VocabParallelLogprobOp()(logits, targets, contract=_contract(), deterministic=False) - - def test_rejects_a_non_bool_flag(self): - logits, targets = _inputs() - with pytest.raises(LogprobContractError, match="deterministic must be a bool"): - VocabParallelLogprobOp()(logits, targets, contract=_contract(), deterministic=1) - - def test_matches_the_deterministic_path_within_tolerance(self): - tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] - logits, targets = _inputs() - relaxed = _contract(determinism_scope="fixed_topology") - logp_fast, lse_fast = VocabParallelLogprobOp()( - logits, targets, contract=relaxed, deterministic=False - ) - logp_det, lse_det = VocabParallelLogprobOp()( - logits, targets, contract=_contract(), num_vocab_tiles=NUM_TILES - ) - assert torch.allclose(logp_fast, logp_det, atol=tolerance["atol"], rtol=tolerance["rtol"]) - assert torch.allclose(lse_fast, lse_det, atol=tolerance["atol"], rtol=tolerance["rtol"]) - - def test_ignores_tile_constraints(self): - # 7 does not divide the padded vocab; the deterministic path rejects it, - # the fast path never looks at it. - logits, targets = _inputs() - relaxed = _contract(determinism_scope="fixed_topology") - logp, lse = VocabParallelLogprobOp()( - logits, targets, contract=relaxed, num_vocab_tiles=7, deterministic=False - ) - ref_lse = torch.logsumexp(logits[:, :REAL_VOCAB].float(), dim=-1) - assert torch.allclose(lse, ref_lse, atol=1e-5) - assert torch.isfinite(logp).all() - - def test_grads_match_autograd_oracle(self): - tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] - relaxed = _contract(determinism_scope="fixed_topology") - logits, targets = _inputs() - x = logits.clone().requires_grad_(True) - logp, lse = VocabParallelLogprobOp()(x, targets, contract=relaxed, deterministic=False) - (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()) - - class TestBackward: def test_grads_match_autograd_oracle(self): tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] @@ -345,7 +286,7 @@ def test_dispatch_resolves_reference_and_leaves_legacy_untouched(): registry = KernelRegistry() contract = _contract() - result = registry.get_logprob_op(contract) + result = registry.get_logprob_op(contract, requested_backend="reference") assert result.capability.backend_id == BACKEND_ID assert result.provenance["fallback"] is False assert isinstance(result.op, VocabParallelLogprobOp) @@ -364,125 +305,161 @@ def test_dispatch_resolves_reference_and_leaves_legacy_untouched(): assert OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP not in candidates -# --------------------------------------------------------------------------- -# Multi-rank gloo tests (spawn pattern from tests/test_linear_logp.py) -# --------------------------------------------------------------------------- +# Cross-TP bitwise determinism on real ranks (NCCL, one CUDA device per rank) +TP_REAL_VOCAB = 1000 +TP_PADDED_VOCAB = 1024 +TP_NUM_TILES = 32 # tile = 32 columns +TP_TILE = TP_PADDED_VOCAB // TP_NUM_TILES +TP_NUM_TOKENS = 48 +TP_ACTIVE = tuple(index % 7 != 5 for index in range(TP_NUM_TOKENS)) +TP_DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16} +_SPAWN_TIMEOUT_S = 300 -def _gloo_available() -> bool: - return torch.distributed.is_available() and torch.distributed.is_gloo_available() +def _cuda_device_count() -> int: + return torch.cuda.device_count() if torch.cuda.is_available() else 0 -requires_gloo = pytest.mark.skipif( - not _gloo_available(), reason="requires torch.distributed with the gloo backend" -) +def _requires_gpus(count: int): + return pytest.mark.skipif( + _cuda_device_count() < count, + reason=f"cross-TP determinism needs {count} CUDA devices to place one rank per device", + ) -_WORLD_SIZE = 4 -_UNEVEN_BOUNDS = ((0, 4), (4, 16), (16, 24), (24, 32)) # tile-aligned (tile=4) + +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_worker(rank, world_size, init_method, result_queue, scenario): +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, + backend_kind="pytorch", +): import torch.distributed as dist - torch.set_num_threads(1) try: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) dist.init_process_group( - backend="gloo", init_method=init_method, rank=rank, world_size=world_size - ) - dtype = torch.bfloat16 if scenario == "bf16" else torch.float32 - dtype_name = "bf16" if scenario == "bf16" else "fp32" - bounds = _UNEVEN_BOUNDS if scenario == "uneven" else _even_bounds(PADDED_VOCAB, world_size) - logits, targets = _inputs(dtype=dtype) - start, end = bounds[rank] - - op = VocabParallelLogprobOp() - tiles = 16 if scenario == "preflight" and rank == 0 else NUM_TILES - contract_tp = _contract( - tp_rank=rank, tp_world_size=world_size, bounds=bounds, dtype=dtype_name + backend="nccl", init_method=init_method, rank=rank, world_size=world_size ) + dtype = TP_DTYPES[dtype_name] + if backend_kind == "rocm": + op = RocmVocabParallelLogprobOp() + elif backend_kind == "triton": + from rl_engine.kernels.ops.triton.loss.vocab_parallel_logp import ( + TritonVocabParallelLogprobOp, + ) - if scenario == "preflight": + op = TritonVocabParallelLogprobOp() + else: + 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=contract_tp, + contract=_tp_contract(rank, world_size, bounds, dtype_name), tp_group=dist.group.WORLD, num_vocab_tiles=tiles, ) result_queue.put({"ok": False, "rank": rank, "traceback": "no error raised"}) - except LogprobContractError: - result_queue.put({"ok": True, "rank": rank}) - return - - if scenario == "mode_mismatch": - # Same relaxed contract everywhere; only rank 0 runs deterministic. - relaxed = _contract( - tp_rank=rank, - tp_world_size=world_size, - bounds=bounds, - determinism_scope="fixed_topology", - ) - try: - op( - logits[:, start:end].contiguous().clone(), - targets, - contract=relaxed, - tp_group=dist.group.WORLD, - num_vocab_tiles=NUM_TILES, - deterministic=(rank == 0), - ) - result_queue.put({"ok": False, "rank": rank, "traceback": "no error raised"}) - except LogprobContractError: - result_queue.put({"ok": True, "rank": rank}) - return - - if scenario == "nondeterministic": - relaxed = _contract( - tp_rank=rank, - tp_world_size=world_size, - bounds=bounds, - determinism_scope="fixed_topology", - ) - shard = logits[:, start:end].contiguous().clone().requires_grad_(True) - logp_tp, lse_tp = op( - shard, targets, contract=relaxed, tp_group=dist.group.WORLD, deterministic=False - ) - (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() - - full = logits.clone().requires_grad_(True) - relaxed_tp1 = _contract(determinism_scope="fixed_topology") - logp_one, lse_one = op(full, targets, contract=relaxed_tp1, deterministic=False) - (logp_one.sum() + 0.5 * lse_one.sum()).backward() - - result_queue.put( - { - "ok": True, - "rank": rank, - "logp_close": torch.allclose(logp_tp, logp_one, atol=1e-6), - "lse_close": torch.allclose(lse_tp, lse_one, atol=1e-6), - "grad_close": torch.allclose(shard.grad, full.grad[:, start:end], atol=1e-6), - "logp": logp_tp.detach().float(), - "lse": lse_tp.detach().float(), - } - ) + except LogprobContractError as exc: + result_queue.put({"ok": True, "rank": rank, "message": str(exc)}) return + start, end = bounds[rank] shard = logits[:, start:end].contiguous().clone().requires_grad_(True) + tp_contract = _tp_contract(rank, world_size, bounds, dtype_name) logp_tp, lse_tp = op( shard, targets, - contract=contract_tp, + contract=tp_contract, tp_group=dist.group.WORLD, - num_vocab_tiles=NUM_TILES, + num_vocab_tiles=TP_NUM_TILES, ) (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() - # In-process TP=1 run of the same op on the full logits: the cross-TP - # bitwise claim is TP=n output == TP=1 output, bit for bit. + # Same ranks, same inputs, run again: the collectives must not perturb bits. + rerun = logits[:, start:end].contiguous().clone() + logp_re, lse_re = op( + rerun, + targets, + contract=tp_contract, + tp_group=dist.group.WORLD, + num_vocab_tiles=TP_NUM_TILES, + ) + + # In-process TP=1 run on the full logits: the cross-TP claim is that a + # TP=n result equals the TP=1 result, bit for bit. full = logits.clone().requires_grad_(True) - contract_tp1 = _contract(dtype=dtype_name) - logp_one, lse_one = op(full, targets, contract=contract_tp1, num_vocab_tiles=NUM_TILES) + logp_one, lse_one = op( + full, + targets, + contract=_tp_contract(0, 1, ((0, TP_PADDED_VOCAB),), dtype_name), + num_vocab_tiles=TP_NUM_TILES, + ) (logp_one.sum() + 0.5 * lse_one.sum()).backward() result_queue.put( @@ -492,45 +469,61 @@ def _tp_worker(rank, world_size, init_method, result_queue, scenario): "logp_bits_match": _bitwise_equal(logp_tp, logp_one), "lse_bits_match": _bitwise_equal(lse_tp, lse_one), "grad_bits_match": _bitwise_equal(shard.grad, full.grad[:, start:end]), - "logp": logp_tp.detach().float(), - "lse": lse_tp.detach().float(), + "rerun_bits_match": ( + _bitwise_equal(logp_re, logp_tp) and _bitwise_equal(lse_re, lse_tp) + ), + "logp_bit_pattern": _bits(logp_tp.detach().float().cpu()).tolist(), + "lse_bit_pattern": _bits(lse_tp.detach().float().cpu()).tolist(), } ) - except Exception: + except Exception: # pragma: no cover - forwarded to the parent process result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) raise finally: - if torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() + import torch.distributed as dist + if dist.is_initialized(): + dist.destroy_process_group() -def _run_gloo_scenario(scenario): + +def _run_nccl_scenario( + world_size, scenario="correctness", uneven=False, dtype_name="fp32", backend_kind="pytorch" +): ctx = mp.get_context("spawn") with tempfile.TemporaryDirectory() as tmpdir: - init_method = (Path(tmpdir) / "gloo_init").as_uri() + init_method = (Path(tmpdir) / "nccl_init").as_uri() result_queue = ctx.Queue() processes = [ ctx.Process( - target=_tp_worker, - args=(rank, _WORLD_SIZE, init_method, result_queue, scenario), + target=_nccl_worker, + args=( + rank, + world_size, + init_method, + result_queue, + scenario, + uneven, + dtype_name, + backend_kind, + ), ) - for rank in range(_WORLD_SIZE) + for rank in range(world_size) ] results = [] try: for process in processes: process.start() - for _ in range(_WORLD_SIZE): + for _ in range(world_size): try: - results.append(result_queue.get(timeout=60)) + results.append(result_queue.get(timeout=_SPAWN_TIMEOUT_S)) except queue.Empty: for process in processes: if process.is_alive(): process.terminate() - pytest.fail("timed out waiting for vocab-parallel gloo workers") + pytest.fail(f"timed out waiting for NCCL workers (scenario={scenario})") finally: for process in processes: - process.join(timeout=10) + process.join(timeout=30) if process.is_alive(): process.terminate() results.sort(key=lambda item: item["rank"]) @@ -541,41 +534,110 @@ def _run_gloo_scenario(scenario): return results -@requires_gloo -@pytest.mark.parametrize("scenario", ["even", "uneven", "bf16"]) -def test_tp4_bitwise_identical_to_tp1(scenario): - results = _run_gloo_scenario(scenario) - for result in results: - assert result["logp_bits_match"], f"rank {result['rank']} logp bits differ from TP=1" - assert result["lse_bits_match"], f"rank {result['rank']} lse bits differ from TP=1" - assert result["grad_bits_match"], f"rank {result['rank']} grad bits differ from TP=1" - # Outputs are replicated: every rank must hold identical bits. - for other in results[1:]: - assert _bitwise_equal(results[0]["logp"], other["logp"]) - assert _bitwise_equal(results[0]["lse"], other["lse"]) - - -@requires_gloo -def test_preflight_rejects_mismatched_num_vocab_tiles(): - _run_gloo_scenario("preflight") - - -@requires_gloo -def test_preflight_rejects_mixed_deterministic_modes(): - _run_gloo_scenario("mode_mismatch") - - -@requires_gloo -def test_tp4_nondeterministic_matches_tp1_within_tolerance(): - """No bitwise claim: the fast path's grouping changes with the TP degree. - The values must still agree within fp32 tolerance, and the outputs stay - replicated bitwise across ranks — every rank merges the same partials.""" - - results = _run_gloo_scenario("nondeterministic") - for result in results: - assert result["logp_close"], f"rank {result['rank']} logp differs from TP=1 beyond atol" - assert result["lse_close"], f"rank {result['rank']} lse differs from TP=1 beyond atol" - assert result["grad_close"], f"rank {result['rank']} grads differ from TP=1 beyond atol" - for other in results[1:]: - assert _bitwise_equal(results[0]["logp"], other["logp"]) - assert _bitwise_equal(results[0]["lse"], other["lse"]) +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"] + + +@pytest.mark.skipif(torch.version.hip is None, reason="requires a ROCm PyTorch build") +class TestRocmNativeCrossTP: + """The ROCm production path must preserve the same TP contract as reference.""" + + @staticmethod + def _require_native() -> None: + from rl_engine.kernels.registry import _rocm_vocab_logprob_native_available + + if not _rocm_vocab_logprob_native_available(): + pytest.skip("requires the compiled ROCm logprob extension") + + @_requires_gpus(2) + @pytest.mark.parametrize("dtype_name", ["fp32", "bf16"]) + def test_tp2_native_matches_tp1_and_repeat(self, dtype_name): + self._require_native() + results = _run_nccl_scenario(2, dtype_name=dtype_name, backend_kind="rocm") + TestCrossTPBitwise._assert_matches_tp1(results) + + @_requires_gpus(4) + @pytest.mark.parametrize("dtype_name", ["fp32", "bf16"]) + def test_tp4_native_matches_tp1_and_repeat(self, dtype_name): + self._require_native() + results = _run_nccl_scenario(4, dtype_name=dtype_name, backend_kind="rocm") + TestCrossTPBitwise._assert_matches_tp1(results) + + +def _triton_available() -> bool: + if not torch.cuda.is_available(): + return False + try: + import triton # noqa: F401 + except ImportError: + return False + return True + + +@pytest.mark.skipif(not _triton_available(), reason="requires Triton on a CUDA/ROCm GPU") +class TestTritonNativeCrossTP: + """The Triton production path must preserve the same TP contract as reference.""" + + @_requires_gpus(2) + @pytest.mark.parametrize("dtype_name", ["fp32", "bf16"]) + def test_tp2_native_matches_tp1_and_repeat(self, dtype_name): + results = _run_nccl_scenario(2, dtype_name=dtype_name, backend_kind="triton") + TestCrossTPBitwise._assert_matches_tp1(results) + + @_requires_gpus(4) + @pytest.mark.parametrize("dtype_name", ["fp32", "bf16"]) + def test_tp4_native_matches_tp1_and_repeat(self, dtype_name): + results = _run_nccl_scenario(4, dtype_name=dtype_name, backend_kind="triton") + TestCrossTPBitwise._assert_matches_tp1(results)