diff --git a/docs/operators/lm_head.md b/docs/operators/lm_head.md index a7608385..79fd9047 100644 --- a/docs/operators/lm_head.md +++ b/docs/operators/lm_head.md @@ -1,18 +1,20 @@ # LM Head -The lm_head operator projects hidden states back to vocabulary logits — the final -layer of the Qwen3/Llama stack. It is a **WS1 ground-truth reference** (issue #108): -a pure-PyTorch definition of the "correct answer" that downstream fused CUDA/Triton +The lm_head operator projects hidden states back to vocabulary logits, the final +layer of the Qwen3/Llama stack. It is a WS1 ground-truth reference for issue #108: +a pure-PyTorch definition of the correct answer that downstream fused CUDA/Triton kernels are validated against. -- **LM Head** (`NativeLMHeadOp`): `out = hidden @ weight.t() (+ bias)`. +- **LM Head** (`NativeLMHeadOp`): mathematically `out = hidden @ weight.t() (+ bias)`. + The native reference implements this as row-wise fixed-K GEMV projections so the + reference path is batch-invariant. For Qwen3-8B the weight is the output projection `[vocab=151936, hidden=4096]` in the -HF `nn.Linear` `[out, in]` convention, so it is transposed internally. It is -**independent** from the embedding table (`tie_word_embeddings=false`) — the two -weights are not shared — and Qwen3 has **no bias** (`bias=None`). +HF `nn.Linear` `[out, in]` convention. It is independent from the embedding table +(`tie_word_embeddings=false`), and Qwen3 has no bias (`bias=None`). ## Entry Point + ```python from rl_engine.kernels.registry import kernel_registry @@ -24,80 +26,61 @@ logits = lm_head(hidden, weight, bias=b) # optional [vocab] bias The op exposes the WS1 dual-path contract: -- `forward(...)` — projects in the input dtype, returns the input dtype (Axis-B accuracy - candidate / dtype-behavior path). -- `forward_fp32(...)` — upcasts to fp32, accumulates in fp32, returns fp32 (the - ground-truth golden path). The matmul runs with autocast disabled and CUDA TF32 - turned off, so it stays a true fp32 reference regardless of the caller's ambient - precision context (the global `allow_tf32` flag is saved and restored around it). +- `forward(...)` projects in the input dtype and returns the input dtype. +- `forward_fp32(...)` upcasts to fp32, accumulates in fp32, and returns fp32. The + fixed-K projection runs with autocast disabled and CUDA TF32 turned off, so it + stays a true fp32 reference regardless of the caller's ambient precision context. ## Backends | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | PyTorch fallback | `NativeLMHeadOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA / ROCm / Triton | — | — | Planned: downstream fused kernels validate against this reference. | +| CUDA / ROCm / Triton | N/A | N/A | Planned: downstream fused kernels validate against this reference. | ## Tensor Contract | Argument | Shape | Dtype | Requirements | | --- | --- | --- | --- | -| `hidden` | `[B, S, hidden]` (any leading dims) | float (fp16/bf16/fp32) | Hidden states (Qwen3-8B `hidden=4096`). | -| `weight` | `[vocab, hidden]` | float | Output projection in HF `[out, in]` layout; transposed internally. Qwen3-8B `[151936, 4096]`. | -| `bias` | `[vocab]` or `None` | float | Optional; Qwen3 has none (`None`). | -| output | `hidden.shape[:-1] + (vocab,)` | `forward`: hidden dtype · `forward_fp32`: float32 | Logits. | - -Output dtype follows `hidden`. Pure function — no randomness, no in-place mutation, -device/dtype follow the inputs. - -> **Difference from the bare `matmul` op**: lm_head takes the weight in HF `[out, in]` -> layout and transposes it internally (`weight.t()`); the `matmul` op computes a bare -> `a @ b` with no transpose. Do not use them interchangeably. +| `hidden` | `[B, S, hidden]` or any leading dims | fp16/bf16/fp32 | Hidden states. | +| `weight` | `[vocab, hidden]` | fp16/bf16/fp32 | Output projection in HF `[out, in]` layout. | +| `bias` | `[vocab]` or `None` | fp16/bf16/fp32 | Optional; Qwen3 uses `None`. | +| output | `hidden.shape[:-1] + (vocab,)` | `forward`: hidden dtype; `forward_fp32`: fp32 | Logits. | -## Dispatch Behavior - -`kernel_registry.get_op("lm_head")` resolves through the `OpBackend` priority map. On -`cuda` / `rocm` / `cpu` the only registered backend today is the PyTorch native op -(`PYTORCH_NATIVE_LM_HEAD`), so every device dispatches to the fp32 reference. When fused -kernels land, they are prepended to the priority list and the native op becomes the fallback. +Output dtype follows `hidden`. The op is pure: no randomness and no in-place mutation. ## Accuracy Reference semantics (`forward_fp32`): ```python -out = hidden.float() @ weight.float().t() +flat_hidden = hidden.float().reshape(-1, hidden.size(-1)) +if flat_hidden.size(0) == 0: + flat_out = flat_hidden @ weight.float().t() +else: + rows = [torch.mv(weight.float(), row) for row in flat_hidden] + flat_out = torch.stack(rows) +out = flat_out.reshape(*hidden.shape[:-1], weight.size(0)) if bias is not None: out = out + bias.float() ``` - **Ground truth**: `forward_fp32` accumulates in and returns fp32, with autocast and - CUDA TF32 disabled so it is a true fp32 reference even if the caller has TF32 or - autocast enabled. + CUDA TF32 disabled. - **Dtype path**: `forward` runs the projection in the input dtype. Because this is a - reduction over `hidden`, low-precision accumulation **drifts** from the fp32 reference - (unlike the lossless embedding gather). Unlike `forward_fp32`, this path intentionally - follows the ambient precision context (it is the dtype-behavior path): with an fp32 - input it is bitwise-equal to the ground truth **when ambient TF32/autocast is off**, but - on a TF32-enabled GPU it tracks real hardware behavior and may drift. bf16/fp16 are - always checked with a tolerance. -- **Axis-B — accuracy tolerance**: measured as max absolute error relative to the output - peak magnitude. On a SMALL load point with the real `hidden=4096` reduction length, - bf16 drifts ~0.3–0.4% of peak and fp16 ~0.05%. Elementwise `rtol` is not used: many - logits are near zero while the accumulated error tracks the reduction length, not the - output value. -- **Axis-A — batch invariance**: a row's logits are independent of the rest of the batch, - so the output is bitwise-identical regardless of batch size or padding (`torch.equal`, - `atol=0`) — **provided the reduction order is fixed**. Multi-threaded CPU GEMM splits - the `hidden` reduction across threads by the `M = batch*seq` dimension, which silently - breaks bitwise batch invariance for large `hidden`; the tests pin a single thread to fix - the order. On GPU, cuBLAS likewise splits K by `M`, so a bitwise batch-invariant GEMM is - a downstream kernel concern, not a free property of `torch.matmul`. - -## Performance Notes - -Reference operator — no fused kernel or benchmark yet. Downstream fused kernels carry their -own benchmarks and are measured against this reference for correctness. + reduction over `hidden`, low-precision accumulation drifts from the fp32 reference and + is checked with tolerance. +- **Axis-A batch invariance**: a row's logits are bitwise-identical regardless of batch + size or padding. The native reference enforces this by flattening leading dimensions + and projecting each row through the same GEMV-shaped K reduction instead of relying on + batched GEMM, whose reduction tree can change with `M = batch * seq`. + +## Dispatch Behavior + +`kernel_registry.get_op("lm_head")` resolves through the `OpBackend` priority map. On +`cuda`, `rocm`, and `cpu`, the only registered backend today is the PyTorch native op +(`PYTORCH_NATIVE_LM_HEAD`), so every device dispatches to this reference. When fused +kernels land, they should preserve the same batch-invariant contract. ## Tests @@ -105,23 +88,13 @@ own benchmarks and are measured against this reference for correctness. python -m pytest tests/test_lm_head.py -v ``` -Covers: fp32 correctness vs naive matmul (bitwise, with ambient TF32 pinned off), -`forward_fp32` precision-context safety (true fp32 under ambient autocast + restores the -global TF32 flag on CPU; numerically beats a TF32 matmul on GPU), bf16/fp16 dtype-path -accuracy (relative-to-peak tolerance, with `bias`), output shape, bias semantics, Axis-A -batch invariance (slice + padding, single-thread reduction, all dtypes), input purity, -gradient flow to `hidden`/`weight` (closed-form check), registry dispatch, and a GPU-only -smoke test at the real Qwen3-8B dims (`vocab=151936, hidden=4096`) that skips when CUDA or -GPU memory is unavailable. +Covers fp32 correctness vs the fixed-K reference, precision-context safety, bf16/fp16 +accuracy, output shape, bias semantics, Axis-A batch invariance, input purity, gradient +flow to `hidden` and `weight`, registry dispatch, and a GPU-only smoke test at the real +Qwen3-8B dimensions. ## Implementation Files - `rl_engine/kernels/ops/pytorch/linear/lm_head.py` - `rl_engine/kernels/registry.py` - `tests/test_lm_head.py` - -## Known Limitations - -- PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). -- Axis-A bitwise batch invariance holds only with a fixed reduction order (single-thread on - CPU); a batch-invariant GEMM on GPU is a downstream concern. diff --git a/rl_engine/executors/deepspeed_trainer.py b/rl_engine/executors/deepspeed_trainer.py index 9eb276c8..7b59bebd 100644 --- a/rl_engine/executors/deepspeed_trainer.py +++ b/rl_engine/executors/deepspeed_trainer.py @@ -246,7 +246,9 @@ def _next_published_weight_version(self, consumed_weight_version: int) -> int: self._latest_published_weight_version = published return published - def _export_zero3_full_state_model(self) -> tuple[torch.nn.Module, Mapping[str, Any]]: + def _export_zero3_full_state_model( + self, + ) -> tuple[torch.nn.Module, Mapping[str, Any]]: model = getattr(self.engine, "module", self.model) rank = self._engine_rank() if rank != 0: @@ -439,7 +441,7 @@ def _linear_logp_op_for_device(device: torch.device | str) -> Any: resolved = torch.device(device) if resolved.type == "cpu": return NativeLinearLogpOp() - return kernel_registry.get_op("linear_logp") + return kernel_registry.get_op("linear_logp", device=resolved) def _unwrap_training_model(engine: Any, fallback_model: torch.nn.Module) -> torch.nn.Module: diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index f124cafb..4a0cee96 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -55,8 +55,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", "silu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", "swiglu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", - "embedding": f"{batch}x{seq}x{vocab}x{DEFAULT_HIDDEN}", - "lm_head": f"{batch}x{seq}x{vocab}", + "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", + "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", } try: @@ -169,9 +169,10 @@ def _make_embedding_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) + hidden_dim = _normalized_dim(args) return { "token_ids": _token_ids((batch, seq), vocab, args, device), - "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 0), + "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 0), } @@ -180,9 +181,10 @@ def _make_lm_head_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) + hidden_dim = _normalized_dim(args) return { - "hidden": _floating_tensor((batch, seq, DEFAULT_HIDDEN), args, dtype, device, 0), - "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 1), + "hidden": _floating_tensor((batch, seq, hidden_dim), args, dtype, device, 0), + "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 1), "bias": None, } diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 55a4a203..f1a4ee49 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -57,6 +57,26 @@ def _load_object(path: str) -> Any: }, grad_input_names=("hidden", "lm_head_weight"), ), + "embedding": OperatorSpec( + name="embedding", + op_class="elementwise", + gold_path="rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", + }, + grad_input_names=("weight",), + ), + "lm_head": OperatorSpec( + name="lm_head", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", + }, + grad_input_names=("hidden", "weight"), + ), } diff --git a/rl_engine/kernels/ops/pytorch/linear/lm_head.py b/rl_engine/kernels/ops/pytorch/linear/lm_head.py index 4604fee7..95bd62ea 100644 --- a/rl_engine/kernels/ops/pytorch/linear/lm_head.py +++ b/rl_engine/kernels/ops/pytorch/linear/lm_head.py @@ -28,7 +28,9 @@ class NativeLMHeadOp: input dtype and therefore drifts from the fp32 ``forward_fp32`` ground truth. Axis-B accuracy uses a tolerance (``torch.allclose``), not bitwise equality. Axis-A batch invariance still holds bitwise within a single dtype - (each output row reduces over ``hidden`` independently of the batch). + because each flattened output row is projected through the same GEMV-shaped + reduction over ``hidden``. The reduction path therefore depends on K and V, + not on how many other rows share the batch. """ def __init__(self) -> None: @@ -101,16 +103,25 @@ def _lm_head( """ h = hidden.to(compute_dtype) w = weight.to(compute_dtype) - # [..., hidden] @ [hidden, vocab] -> [..., vocab]; weight is [vocab, hidden] (HF [out, in]). if strict_fp32: with NativeLMHeadOp._strict_fp32_matmul(h.device.type): - out = h @ w.t() + out = NativeLMHeadOp._fixed_k_projection(h, w) else: - out = h @ w.t() + out = NativeLMHeadOp._fixed_k_projection(h, w) if bias is not None: out = out + bias.to(compute_dtype) return out.to(output_dtype) + @staticmethod + def _fixed_k_projection(hidden: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + flat_hidden = hidden.reshape(-1, hidden.size(-1)) + if flat_hidden.size(0) == 0: + flat_out = flat_hidden @ weight.t() + return flat_out.reshape(*hidden.shape[:-1], weight.size(0)) + rows = [torch.mv(weight, row) for row in flat_hidden] + flat_out = torch.stack(rows, dim=0) + return flat_out.reshape(*hidden.shape[:-1], weight.size(0)) + @staticmethod @contextmanager def _strict_fp32_matmul(device_type: str): diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e1..705b3b69 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -6,6 +6,8 @@ from enum import Enum, EnumMeta from typing import Any, Dict, Optional, Set, Type +import torch + from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -163,11 +165,18 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE, ], - "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], + "attn": [ + OpBackend.FLASH_ATTN, + OpBackend.TRITON_GENERIC, + OpBackend.PYTORCH_ATTN, + ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], - "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], + "linear_logp": [ + OpBackend.TRITON_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], @@ -179,7 +188,11 @@ def __init__(self): "rope": [OpBackend.PYTORCH_NATIVE_ROPE], }, "rocm": { - "logp": [OpBackend.ROCM_AITER, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], + "logp": [ + OpBackend.ROCM_AITER, + OpBackend.TRITON_GENERIC, + OpBackend.PYTORCH_NATIVE, + ], "logp_deterministic": [OpBackend.PYTORCH_NATIVE], "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [ @@ -191,7 +204,10 @@ def __init__(self): "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], - "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], + "linear_logp": [ + OpBackend.TRITON_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], @@ -237,7 +253,11 @@ def _adjust_priority_from_env(self): OpBackend.ROCM_FLASH_ATTN, OpBackend.TRITON_GENERIC, ] - elif rocm_attn_backend and rocm_attn_backend not in {"native", "pytorch", "sdpa"}: + elif rocm_attn_backend and rocm_attn_backend not in { + "native", + "pytorch", + "sdpa", + }: logger.warning( "Unknown RL_KERNEL_ROCM_ATTN_BACKEND=%s; using default ROCm attention priority.", rocm_attn_backend, @@ -281,16 +301,11 @@ def _adjust_priority_for_hardware(self): except Exception as e: logger.warning(f"Failed to probe device capability: {e}") - def get_op(self, op_type: str) -> Any: + def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: """Core distribution logic: Automatically select the best operator based on hardware and priority. """ - if device_ctx.is_rocm: - platform = "rocm" - elif device_ctx.device_type == "cuda": - platform = "cuda" - else: - platform = "cpu" + platform = self._platform_for_device(device) candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: @@ -314,6 +329,21 @@ def get_op(self, op_type: str) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def _platform_for_device(self, device: torch.device | str | None) -> str: + if device is None: + if device_ctx.is_rocm: + return "rocm" + if device_ctx.device_type == "cuda": + return "cuda" + return "cpu" + + resolved = torch.device(device) + if resolved.type == "cuda": + return "rocm" if torch.version.hip is not None else "cuda" + if resolved.type in self._priority_map: + return resolved.type + return "cpu" + def _load_backend(self, backend: OpBackend) -> Optional[Type]: """Dynamic loading technique: Import modules only when needed and check environment dependencies. diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index 7f2839e1..3d61a068 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -42,6 +42,28 @@ def test_rocm_attention_native_sdpa_opt_out(monkeypatch): assert registry._priority_map["rocm"]["attn"][1] == OpBackend.ROCM_FLASH_ATTN +def test_registry_explicit_device_selects_device_platform(monkeypatch): + registry = KernelRegistry() + loaded = [] + + class DummyOp: + pass + + def fake_load_backend(backend): + loaded.append(backend) + return DummyOp + + monkeypatch.setattr(registry, "_load_backend", fake_load_backend) + + op = registry.get_op("logp", device="cuda:0") + + assert isinstance(op, DummyOp) + expected = ( + OpBackend.ROCM_AITER if torch.version.hip is not None else OpBackend.CUDA_FUSED_LOGP_GENERIC + ) + assert loaded[0] == expected + + def test_executor_flow(): executor = RolloutExecutor() mock_input_ids = torch.ones((1, 16), dtype=torch.long) diff --git a/tests/test_deepspeed_training_worker.py b/tests/test_deepspeed_training_worker.py index 33b7d769..d0a61e81 100644 --- a/tests/test_deepspeed_training_worker.py +++ b/tests/test_deepspeed_training_worker.py @@ -393,7 +393,9 @@ def test_extract_hidden_states_rejects_ambiguous_multi_tensor_tuple(): _extract_hidden_states((torch.randn(2, 3, 11), torch.randn(2, 3, 5))) -def test_deepspeed_training_worker_routes_linear_logp_and_zeroes_masked_targets(monkeypatch): +def test_deepspeed_training_worker_routes_linear_logp_and_zeroes_masked_targets( + monkeypatch, +): _install_fake_deepspeed(monkeypatch) from rl_engine.executors import deepspeed_trainer @@ -463,6 +465,23 @@ def test_deepspeed_training_worker_routes_linear_logp_and_zeroes_masked_targets( assert math.isfinite(result.metrics["loss"]) +def test_deepspeed_linear_logp_device_lookup_uses_explicit_worker_device(monkeypatch): + from rl_engine.executors import deepspeed_trainer + + sentinel = object() + call = {} + + def fake_get_op(op_type, *, device=None): + call["op_type"] = op_type + call["device"] = device + return sentinel + + monkeypatch.setattr(deepspeed_trainer.kernel_registry, "get_op", fake_get_op) + + assert deepspeed_trainer._linear_logp_op_for_device("cuda:0") is sentinel + assert call == {"op_type": "linear_logp", "device": torch.device("cuda:0")} + + def test_deepspeed_training_worker_rejects_ignore_index_in_model_inputs(monkeypatch): _install_fake_deepspeed(monkeypatch) from rl_engine.executors import deepspeed_trainer @@ -503,7 +522,9 @@ def test_deepspeed_training_worker_rejects_ignore_index_in_model_inputs(monkeypa worker.train(_rollout()) -def test_deepspeed_zero3_training_gathers_lm_head_parameters_during_backward(monkeypatch): +def test_deepspeed_zero3_training_gathers_lm_head_parameters_during_backward( + monkeypatch, +): _install_fake_deepspeed_with_gather(monkeypatch) from rl_engine.executors import deepspeed_trainer diff --git a/tests/test_issue151_embedding_lm_head_invariance.py b/tests/test_issue151_embedding_lm_head_invariance.py new file mode 100644 index 00000000..58f558c7 --- /dev/null +++ b/tests/test_issue151_embedding_lm_head_invariance.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.linear.embedding import NativeEmbeddingOp +from rl_engine.kernels.ops.pytorch.loss.linear_logp import NativeLinearLogpOp +from rl_engine.testing import ( + SyntheticRLKernelBatch, + make_synthetic_rl_kernel_batch, + selected_logprobs_reference, +) + +VOCAB_SIZE = 4096 +HIDDEN_DIM = 128 +PROMPT_PROBE_POS = 1 +COMPLETION_PROBE_POS = 5 +PROMPT_PROBE_TOKEN = 1234 +COMPLETION_PROBE_TOKEN = 2345 + +BATCH_LAYOUTS = ( + dict( + num_prompts=1, + samples_per_prompt=2, + prompt_len=4, + completion_len=6, + vocab_size=VOCAB_SIZE, + valid_density=1.0, + seed=151, + ), + dict( + num_prompts=2, + samples_per_prompt=3, + prompt_len=4, + completion_len=8, + vocab_size=VOCAB_SIZE, + valid_density=0.5, + seed=152, + ), + dict( + num_prompts=3, + samples_per_prompt=4, + prompt_len=4, + completion_len=10, + vocab_size=VOCAB_SIZE, + valid_density=0.75, + seed=153, + ), +) + +CUDA_CASE = pytest.param( + "cuda", + torch.bfloat16, + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available"), +) + + +def _make_embedding_weight(*, device: str, dtype: torch.dtype, seed: int) -> torch.Tensor: + generator = torch.Generator(device=torch.device(device)) + generator.manual_seed(seed) + return torch.randn( + VOCAB_SIZE, + HIDDEN_DIM, + device=device, + dtype=torch.float32, + generator=generator, + ).to(dtype=dtype) + + +def _stamp_probe_tokens(batch: SyntheticRLKernelBatch) -> SyntheticRLKernelBatch: + completion_offset = COMPLETION_PROBE_POS - batch.prompt_len + if completion_offset < 0 or completion_offset >= batch.completion_len: + raise ValueError("completion probe position must fall inside completion tokens") + + input_ids = batch.input_ids.clone() + token_ids = batch.token_ids.clone() + completion_mask = batch.completion_mask.clone() + attention_mask = batch.attention_mask.clone() + + input_ids[:, PROMPT_PROBE_POS] = PROMPT_PROBE_TOKEN + input_ids[:, COMPLETION_PROBE_POS] = COMPLETION_PROBE_TOKEN + token_ids[:, completion_offset] = COMPLETION_PROBE_TOKEN + completion_mask[:, completion_offset] = True + attention_mask[:, COMPLETION_PROBE_POS] = True + valid_indices = completion_mask.reshape(-1).nonzero(as_tuple=False).squeeze(-1) + + return replace( + batch, + input_ids=input_ids, + token_ids=token_ids, + completion_mask=completion_mask, + attention_mask=attention_mask, + valid_indices=valid_indices, + metadata={ + **batch.metadata, + "valid_tokens": int(completion_mask.sum().item()), + "valid_density": float(completion_mask.float().mean().item()), + }, + ) + + +def _make_layout( + layout: dict[str, int | float], *, device: str, dtype: torch.dtype +) -> SyntheticRLKernelBatch: + batch = make_synthetic_rl_kernel_batch(device=device, dtype=dtype, **layout) + return _stamp_probe_tokens(batch) + + +def _permute_rows(batch: SyntheticRLKernelBatch, perm: torch.Tensor) -> SyntheticRLKernelBatch: + completion_mask = batch.completion_mask.index_select(0, perm) + return replace( + batch, + input_ids=batch.input_ids.index_select(0, perm), + attention_mask=batch.attention_mask.index_select(0, perm), + prompt_mask=batch.prompt_mask.index_select(0, perm), + completion_mask=completion_mask, + token_ids=batch.token_ids.index_select(0, perm), + rewards=batch.rewards.index_select(0, perm), + advantages=batch.advantages.index_select(0, perm), + old_logps=batch.old_logps.index_select(0, perm), + ref_logps=batch.ref_logps.index_select(0, perm), + valid_indices=completion_mask.reshape(-1).nonzero(as_tuple=False).squeeze(-1), + ) + + +def _assert_probe_vectors( + output: torch.Tensor, + *, + batch_size: int, + prompt_reference: torch.Tensor, + completion_reference: torch.Tensor, +) -> None: + assert torch.equal( + output[:, PROMPT_PROBE_POS, :], + prompt_reference.expand(batch_size, -1), + ) + assert torch.equal( + output[:, COMPLETION_PROBE_POS, :], + completion_reference.expand(batch_size, -1), + ) + + +@pytest.mark.parametrize("device,dtype", [("cpu", torch.float32), CUDA_CASE]) +def test_embedding_lookup_is_bitwise_identical_across_rl_batch_layouts( + device: str, dtype: torch.dtype +) -> None: + op = NativeEmbeddingOp() + weight = _make_embedding_weight(device=device, dtype=dtype, seed=2026) + prompt_reference = weight[PROMPT_PROBE_TOKEN].detach() + completion_reference = weight[COMPLETION_PROBE_TOKEN].detach() + + for layout in BATCH_LAYOUTS: + batch = _make_layout(layout, device=device, dtype=dtype) + output = op.forward(batch.input_ids, weight) + _assert_probe_vectors( + output, + batch_size=batch.batch_size, + prompt_reference=prompt_reference, + completion_reference=completion_reference, + ) + + +@pytest.mark.parametrize("device,dtype", [("cpu", torch.float32), CUDA_CASE]) +def test_embedding_lookup_is_row_order_invariant_under_rl_permutation( + device: str, dtype: torch.dtype +) -> None: + op = NativeEmbeddingOp() + weight = _make_embedding_weight(device=device, dtype=dtype, seed=2027) + batch = _make_layout(BATCH_LAYOUTS[2], device=device, dtype=dtype) + + perm = torch.arange(batch.batch_size - 1, -1, -1, device=torch.device(device)) + original = op.forward(batch.input_ids, weight) + permuted = op.forward(_permute_rows(batch, perm).input_ids, weight) + + assert torch.equal(permuted, original.index_select(0, perm)) + + +@pytest.mark.parametrize("device,dtype", [("cpu", torch.float32), CUDA_CASE]) +def test_embedding_lookup_active_tokens_ignore_padding_tail_mutations( + device: str, dtype: torch.dtype +) -> None: + op = NativeEmbeddingOp() + weight = _make_embedding_weight(device=device, dtype=dtype, seed=2028) + batch = _make_layout(BATCH_LAYOUTS[1], device=device, dtype=dtype) + inactive = ~batch.attention_mask + assert bool(inactive.any().item()) + + mutated_input_ids = batch.input_ids.clone() + generator = torch.Generator(device=torch.device(device)) + generator.manual_seed(404) + random_tokens = torch.randint( + 0, + VOCAB_SIZE, + mutated_input_ids.shape, + device=device, + generator=generator, + dtype=torch.long, + ) + mutated_input_ids[inactive] = random_tokens[inactive] + + baseline = op.forward(batch.input_ids, weight) + candidate = op.forward(mutated_input_ids, weight) + + assert torch.equal(candidate[batch.attention_mask], baseline[batch.attention_mask]) + + +def test_embedding_backward_weight_gradient_is_layout_invariant() -> None: + torch.manual_seed(2030) + op = NativeEmbeddingOp() + base_token_ids = torch.tensor([2, 5, 1, 5, 2, 3], dtype=torch.long) + base_upstream = torch.randn(6, HIDDEN_DIM) + layouts = [ + ((2, 3), [0, 1, 2, 3, 4, 5]), + ((3, 2), [5, 2, 1, 0, 4, 3]), + ((1, 6), [3, 0, 5, 2, 1, 4]), + ] + + canonical_grad = None + for lead_shape, order in layouts: + order_t = torch.tensor(order, dtype=torch.long) + token_ids = base_token_ids.index_select(0, order_t).reshape(lead_shape) + upstream = base_upstream.index_select(0, order_t).reshape(*lead_shape, HIDDEN_DIM) + weight = _make_embedding_weight(device="cpu", dtype=torch.float32, seed=2030) + weight.requires_grad_(True) + + (op.forward(token_ids, weight) * upstream).sum().backward() + grad = weight.grad.detach().clone() + + if canonical_grad is None: + canonical_grad = grad + else: + assert torch.allclose(grad, canonical_grad, atol=1e-6) + + +@pytest.mark.parametrize("device,dtype", [("cpu", torch.float32), CUDA_CASE]) +def test_lm_head_linear_logp_handoff_is_layout_invariant_for_rl_batches( + device: str, dtype: torch.dtype +) -> None: + target_device = torch.device(device) + generator = torch.Generator(device=target_device) + generator.manual_seed(2031) + op = NativeLinearLogpOp() + weight = torch.randn(29, 7, device=target_device, dtype=dtype, generator=generator) + bias = torch.randn(29, device=target_device, dtype=dtype, generator=generator) + base_hidden = torch.randn(6, 7, device=target_device, dtype=dtype, generator=generator) + base_target = torch.tensor([3, 7, 1, 9, 4, 6], device=target_device, dtype=torch.long) + base_mask = torch.tensor( + [True, False, True, True, False, True], device=target_device, dtype=torch.bool + ) + layouts = [ + ((2, 3), [0, 1, 2, 3, 4, 5]), + ((3, 2), [5, 1, 3, 0, 4, 2]), + ((1, 6), [2, 4, 1, 5, 0, 3]), + ] + + canonical = None + for lead_shape, order in layouts: + order_t = torch.tensor(order, device=target_device, dtype=torch.long) + hidden = base_hidden.index_select(0, order_t).reshape(*lead_shape, -1) + target = base_target.index_select(0, order_t).reshape(lead_shape) + mask = base_mask.index_select(0, order_t).reshape(lead_shape) + safe_target = target.masked_fill(~mask, 0) + padded_hidden = hidden.clone() + padding_values = torch.randn( + padded_hidden.shape, + device=target_device, + dtype=dtype, + generator=generator, + ) + padded_hidden[~mask] = padding_values[~mask] + + logps = op(padded_hidden, weight, safe_target, bias).masked_fill(~mask, 0.0) + logits = torch.nn.functional.linear(padded_hidden.float(), weight.float(), bias.float()) + expected = selected_logprobs_reference(logits, target, mask=mask) + + assert torch.allclose(logps, expected, atol=5e-2 if dtype is torch.bfloat16 else 1e-5) + recovered = logps.reshape(-1)[torch.argsort(order_t)].float() + if canonical is None: + canonical = recovered + else: + assert torch.allclose( + recovered, canonical, atol=5e-2 if dtype is torch.bfloat16 else 1e-6 + ) diff --git a/tests/test_lm_head.py b/tests/test_lm_head.py index 15571e2b..51f7c481 100644 --- a/tests/test_lm_head.py +++ b/tests/test_lm_head.py @@ -12,12 +12,10 @@ bitwise -- elementwise rtol is useless here because many logits are near zero while the accumulated error tracks the reduction length, not the output value. - * Axis-A (batch invariance): still bitwise within a single dtype, but only - once the CPU reduction order is pinned. Multi-threaded CPU GEMM splits the - K (=hidden) reduction across threads differently depending on the M - (=batch*seq) dimension, which silently breaks bitwise batch invariance for - large hidden. ``_single_thread`` fixes the reduction order; this is the - local stand-in for the planned testing/determinism.py::deterministic_context. + * Axis-A (batch invariance): bitwise within a single dtype because the native + reference projects each output row with the same GEMV-shaped fixed-K + reduction. The reduction path is independent of the flattened M + (=batch*seq) dimension. """ import contextlib @@ -93,26 +91,37 @@ def _rand_weight(vocab=_VOCAB, hidden=_HIDDEN, *, seed, dtype=torch.float32): return torch.randn(vocab, hidden, generator=gen, dtype=dtype) -# Correctness of the fp32 ground truth: forward_fp32 == naive fp32 matmul, -# bitwise. forward_fp32 disables autocast/TF32, so it is a true fp32 reference -# unconditionally. The fp32 dtype path (forward) follows the ambient precision -# context, so it is bitwise-equal to the ground truth only when TF32/autocast is -# off (the default here); only bf16/fp16 introduce drift. -def test_native_lm_head_fp32_matches_naive_matmul(): - """forward_fp32 (and fp32 forward, TF32 off) is bitwise-equal to a naive fp32 matmul.""" +def _fixed_k_reference(hidden: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor: + flat_hidden = hidden.reshape(-1, hidden.size(-1)) + if flat_hidden.size(0) == 0: + flat_out = flat_hidden @ weight.t() + else: + flat_out = torch.stack([torch.mv(weight, row) for row in flat_hidden], dim=0) + out = flat_out.reshape(*hidden.shape[:-1], weight.size(0)) + if bias is not None: + out = out + bias + return out + + +# Correctness of the fp32 ground truth: forward_fp32 == row-wise fixed-K fp32 +# projection, bitwise. This intentionally avoids batched GEMM as the golden +# source, because batched GEMM may choose a different K reduction tree when the +# flattened M (=batch*seq) dimension changes. +def test_native_lm_head_fp32_matches_fixed_k_reference(): + """forward_fp32 and fp32 forward are bitwise-equal to the fixed-K reference.""" hidden = _rand_hidden(2, 5, seed=1) weight = _rand_weight(seed=1) - # Pin TF32 off so the fp32 forward path and the naive reference below are both - # true fp32 regardless of the machine's global default, making this assertion - # deterministic. forward_fp32 disables TF32 internally and does not need this. + # Pin TF32 off so the fp32 forward path and fixed-K reference are both true + # fp32 regardless of the machine's global default. forward_fp32 disables TF32 + # internally and does not need this. prev_tf32 = torch.backends.cuda.matmul.allow_tf32 torch.backends.cuda.matmul.allow_tf32 = False try: - naive = hidden.float() @ weight.float().t() - assert torch.equal(NativeLMHeadOp().forward_fp32(hidden, weight), naive) + ref = _fixed_k_reference(hidden.float(), weight.float()) + assert torch.equal(NativeLMHeadOp().forward_fp32(hidden, weight), ref) # fp32 forward path computes in fp32 too -> bitwise equal to ground truth. - assert torch.equal(NativeLMHeadOp().forward(hidden, weight), naive) + assert torch.equal(NativeLMHeadOp().forward(hidden, weight), ref) finally: torch.backends.cuda.matmul.allow_tf32 = prev_tf32 @@ -124,7 +133,7 @@ def test_forward_fp32_ignores_ambient_autocast_and_restores_tf32(): weight = _rand_weight(vocab=7, hidden=32, seed=11) gen = torch.Generator().manual_seed(11) bias = torch.randn(7, generator=gen) - ref = hidden.float() @ weight.float().t() + bias.float() + ref = _fixed_k_reference(hidden.float(), weight.float(), bias.float()) prev_tf32 = torch.backends.cuda.matmul.allow_tf32 torch.backends.cuda.matmul.allow_tf32 = True @@ -203,6 +212,21 @@ def test_native_lm_head_output_shape(): assert out.shape == (3, 7, _VOCAB) +def test_native_lm_head_zero_rows_preserve_autograd_edges(): + op = NativeLMHeadOp() + hidden = torch.empty(0, _HIDDEN, requires_grad=True) + weight = _rand_weight(seed=31).requires_grad_(True) + + out = op.forward_fp32(hidden, weight) + out.sum().backward() + + assert out.shape == (0, _VOCAB) + assert hidden.grad is not None + assert hidden.grad.shape == hidden.shape + assert weight.grad is not None + assert torch.equal(weight.grad, torch.zeros_like(weight)) + + # Bias: None (Qwen3 default) is a plain matmul; a provided [vocab] bias is added. def test_native_lm_head_bias(): """bias=None is a plain matmul; a [vocab] bias is added elementwise.""" @@ -330,8 +354,6 @@ def test_native_lm_head_qwen3_8b_real_shape(): out = op.forward_fp32(hidden, weight) assert out.shape == (2, 16, _QWEN3_VOCAB) assert out.dtype == torch.float32 - # Bitwise equal to the naive fp32 matmul (same call, same inputs). - assert torch.equal(out, hidden @ weight.t()) - # NB: Axis-A bitwise is asserted on CPU (single-thread reduction) above. - # On GPU it is NOT free -- cuBLAS also splits K by the M dimension, so a - # batch-invariant GEMM is a downstream kernel concern, not validated here. + # Axis-A bitwise is already asserted above; this smoke test validates the + # fixed-K path at real Qwen3-8B vocab width and hidden reduction length. + assert torch.equal(out[:1], op.forward_fp32(hidden[:1], weight)) diff --git a/tests/test_op_checks.py b/tests/test_op_checks.py index cbecddea..e076e106 100644 --- a/tests/test_op_checks.py +++ b/tests/test_op_checks.py @@ -3,9 +3,18 @@ from __future__ import annotations +import argparse + import torch from rl_engine.kernels.gtest.op_checks import CandidateSpec, OperatorCase, run_operator_suite +from rl_engine.kernels.gtest.operator_specs import ( + make_candidate, + make_operator_case, + operator_names, +) +from rl_engine.kernels.ops.pytorch.linear.embedding import NativeEmbeddingOp +from rl_engine.kernels.ops.pytorch.linear.lm_head import NativeLMHeadOp from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp @@ -34,6 +43,54 @@ def _logp_backward_case(name: str, *, seed: int = 0) -> OperatorCase: ) +def _embedding_case(name: str, *, seed: int = 0) -> OperatorCase: + generator = torch.Generator().manual_seed(seed) + token_ids = torch.tensor([[1, 7, 3], [7, 0, 5]], dtype=torch.long) + weight = torch.randn(11, 5, generator=generator) + return OperatorCase( + name=name, + op_class="elementwise", + dtype=torch.float32, + inputs={"token_ids": token_ids, "weight": weight}, + gold_fn=NativeEmbeddingOp().forward_fp32, + grad_input_names=("weight",), + ) + + +def _lm_head_case(name: str, *, seed: int = 0) -> OperatorCase: + generator = torch.Generator().manual_seed(seed) + hidden = torch.randn(2, 3, 5, generator=generator) + weight = torch.randn(13, 5, generator=generator) + return OperatorCase( + name=name, + op_class="reduction", + dtype=torch.float32, + inputs={"hidden": hidden, "weight": weight, "bias": None}, + gold_fn=NativeLMHeadOp().forward_fp32, + grad_input_names=("hidden", "weight"), + ) + + +def _spec_args(op: str) -> argparse.Namespace: + return argparse.Namespace( + op=op, + candidate="native", + arch_key=None, + batch=1, + seq=2, + vocab=17, + seed=123, + input_mode="constant", + constant_value=0.5, + token_value=3, + normalized_dim=8, + k_dim=8, + n_dim=8, + theta=1.0e6, + eps=1.0e-6, + ) + + def test_logp_native_candidate_suite_passes(): report = run_operator_suite( "logp", @@ -51,6 +108,48 @@ def test_logp_native_candidate_suite_passes(): assert all(case.passed for case in report.candidates[0].cases) +def test_embedding_native_candidate_suite_passes_issue_108_helper(): + report = run_operator_suite( + "embedding", + candidates=[ + CandidateSpec(name="native-embedding", backend="pytorch", fn=NativeEmbeddingOp()) + ], + cases=[_embedding_case("fp32", seed=10)], + check_grad=True, + ) + + assert report.passed + assert report.candidates[0].cases[0].outputs[1].message == "gradient:weight" + + +def test_lm_head_native_candidate_suite_passes_issue_108_helper(): + report = run_operator_suite( + "lm_head", + candidates=[CandidateSpec(name="native-lm-head", backend="pytorch", fn=NativeLMHeadOp())], + cases=[_lm_head_case("fp32", seed=11)], + check_grad=True, + ) + + assert report.passed + messages = [output.message for output in report.candidates[0].cases[0].outputs] + assert "gradient:hidden" in messages + assert "gradient:weight" in messages + + +def test_issue151_ops_pass_shared_issue_108_spec_path(): + assert {"embedding", "lm_head"}.issubset(operator_names()) + + for op_name in ("embedding", "lm_head"): + args = _spec_args(op_name) + report = run_operator_suite( + op_name, + candidates=[make_candidate(args)], + cases=[make_operator_case(args, torch.float32, torch.device("cpu"))], + check_grad=True, + ) + assert report.passed + + def test_suite_reports_failure_for_bad_candidate(): def bad_logp(logits, token_ids): del token_ids diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index bb1a2220..db7aaed5 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -79,3 +79,22 @@ def test_constant_linear_logp_inputs_match_operator_contract(): assert torch.equal(inputs["lm_head_weight"], torch.full((17, 128), 0.51)) assert torch.equal(inputs["target_ids"], torch.full((1, 2), 3, dtype=torch.long)) assert inputs["bias"] is None + + +def test_constant_embedding_inputs_match_operator_contract(): + args = _args(input_mode="constant", constant_value=0.5, token_value=3) + inputs = make_operator_inputs("embedding", args, torch.float32, torch.device("cpu")) + + assert torch.equal(inputs["token_ids"], torch.full((1, 2), 3, dtype=torch.long)) + assert torch.equal(inputs["weight"], torch.full((17, 128), 0.5)) + assert operator_shape_name("embedding", args) == "1x2x17x128" + + +def test_constant_lm_head_inputs_match_operator_contract(): + args = _args(input_mode="constant", constant_value=0.5) + inputs = make_operator_inputs("lm_head", args, torch.float32, torch.device("cpu")) + + assert torch.equal(inputs["hidden"], torch.full((1, 2, 128), 0.5)) + assert torch.equal(inputs["weight"], torch.full((17, 128), 0.51)) + assert inputs["bias"] is None + assert operator_shape_name("lm_head", args) == "1x2x128x17"