From 3e6cd1d93466735864375254e393ef95fd88e052 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 03:30:37 +0000 Subject: [PATCH 01/14] Quarantine vendored trees under ops/_vendored (#78 phase 1) git-mv ops/quack -> ops/_vendored/quack and ops/linear_attention/flashqla -> ops/_vendored/flashqla. In-repo imports are rewritten to the new paths; the old paths remain as alias stubs (sys.modules aliasing) for one deprecation cycle. The lint/format/codespell excludes and the pyproject vendoring note collapse to the single _vendored/ prefix, and the ops/README.md map is updated. --- .codespellrc | 2 +- .pre-commit-config.yaml | 2 +- docs/src/content/docs/moe/kernels.mdx | 2 +- pyproject.toml | 5 +-- src/xorl/ops/README.md | 7 ++-- src/xorl/ops/_vendored/__init__.py | 7 ++++ .../flashqla/LICENSE | 0 .../flashqla/USAGE.md | 0 .../flashqla/VENDORED.md | 0 src/xorl/ops/_vendored/flashqla/__init__.py | 24 ++++++++++++ .../flashqla/ops/__init__.py | 0 .../flashqla/ops/gated_delta_rule/__init__.py | 0 .../ops/gated_delta_rule/chunk/__init__.py | 0 .../ops/gated_delta_rule/chunk/cp_context.py | 0 .../gated_delta_rule/chunk/hopper/__init__.py | 0 .../gated_delta_rule/chunk/hopper/cp_fwd.py | 0 .../chunk/hopper/fused_bwd.py | 0 .../chunk/hopper/fused_fwd.py | 0 .../chunk/hopper/kkt_solve.py | 0 .../chunk/hopper/prepare_h.py | 0 .../flashqla/ops/utils/__init__.py | 0 .../flashqla/ops/utils/cumsum.py | 0 .../flashqla/ops/utils/group_reduce.py | 0 .../flashqla/utils/__init__.py | 0 .../flashqla/utils/index.py | 0 .../flashqla/utils/math.py | 0 .../flashqla/utils/pack.py | 0 .../flashqla/utils/profiler.py | 0 .../ops/{ => _vendored}/quack/VENDORED.md | 0 src/xorl/ops/_vendored/quack/__init__.py | 34 +++++++++++++++++ .../{ => _vendored}/quack/_compile_worker.py | 0 .../{ => _vendored}/quack/_worker_protocol.py | 0 .../ops/{ => _vendored}/quack/activation.py | 0 .../ops/{ => _vendored}/quack/autotuner.py | 0 .../quack/blockscaled_gemm_utils.py | 0 .../{ => _vendored}/quack/broadcast_utils.py | 0 .../ops/{ => _vendored}/quack/cache_utils.py | 0 .../{ => _vendored}/quack/compile_utils.py | 0 .../ops/{ => _vendored}/quack/copy_utils.py | 0 .../{ => _vendored}/quack/cross_entropy.py | 0 .../{ => _vendored}/quack/cute_dsl_elf_fix.py | 0 .../quack/cute_dsl_mlir_threading.py | 0 .../{ => _vendored}/quack/cute_dsl_ptxas.py | 0 .../{ => _vendored}/quack/cute_dsl_utils.py | 0 .../{ => _vendored}/quack/epi_composable.py | 0 src/xorl/ops/{ => _vendored}/quack/epi_ops.py | 0 .../ops/{ => _vendored}/quack/epi_utils.py | 0 .../ops/{ => _vendored}/quack/fast_math.py | 0 src/xorl/ops/{ => _vendored}/quack/gemm.py | 0 .../ops/{ => _vendored}/quack/gemm_act.py | 0 .../ops/{ => _vendored}/quack/gemm_base.py | 0 .../quack/gemm_blockscaled_interface.py | 0 .../ops/{ => _vendored}/quack/gemm_config.py | 0 .../ops/{ => _vendored}/quack/gemm_dact.py | 0 .../{ => _vendored}/quack/gemm_default_epi.py | 0 .../{ => _vendored}/quack/gemm_interface.py | 0 .../{ => _vendored}/quack/gemm_norm_act.py | 0 .../ops/{ => _vendored}/quack/gemm_sm100.py | 0 .../ops/{ => _vendored}/quack/gemm_sm120.py | 0 .../ops/{ => _vendored}/quack/gemm_sm80.py | 0 .../ops/{ => _vendored}/quack/gemm_sm90.py | 0 .../{ => _vendored}/quack/gemm_sq_reduce.py | 0 .../{ => _vendored}/quack/gemm_symmetric.py | 0 .../quack/gemm_tvm_ffi_utils.py | 0 .../quack/gemm_wrapper_utils.py | 0 .../ops/{ => _vendored}/quack/layout_utils.py | 0 src/xorl/ops/{ => _vendored}/quack/linear.py | 0 .../quack/linear_cross_entropy.py | 0 src/xorl/ops/{ => _vendored}/quack/mlp.py | 0 .../ops/{ => _vendored}/quack/mx_utils.py | 0 .../{ => _vendored}/quack/nvmmh_heuristic.py | 0 .../ops/{ => _vendored}/quack/pipeline.py | 0 src/xorl/ops/{ => _vendored}/quack/reduce.py | 0 .../{ => _vendored}/quack/reduction_base.py | 0 .../{ => _vendored}/quack/rms_final_reduce.py | 0 src/xorl/ops/{ => _vendored}/quack/rmsnorm.py | 0 src/xorl/ops/{ => _vendored}/quack/rotary.py | 0 .../ops/{ => _vendored}/quack/rounding.py | 0 .../ops/{ => _vendored}/quack/sm100_utils.py | 0 .../ops/{ => _vendored}/quack/sm80_utils.py | 0 .../ops/{ => _vendored}/quack/sm90_utils.py | 0 src/xorl/ops/{ => _vendored}/quack/softmax.py | 0 .../{ => _vendored}/quack/sort/__init__.py | 0 .../quack/sort/bitonic_sort.py | 0 .../quack/sort/generate_sorting_networks.py | 0 .../quack/sort/sorting_networks.py | 0 .../ops/{ => _vendored}/quack/sort/utils.py | 0 .../quack/tensormap_manager.py | 0 .../{ => _vendored}/quack/tile_scheduler.py | 0 src/xorl/ops/{ => _vendored}/quack/topk.py | 0 src/xorl/ops/{ => _vendored}/quack/trace.py | 0 src/xorl/ops/{ => _vendored}/quack/utils.py | 0 .../ops/{ => _vendored}/quack/varlen_utils.py | 0 src/xorl/ops/group_gemm/kernel/quack.py | 2 +- src/xorl/ops/linear_attention/backend.py | 10 ++--- .../ops/linear_attention/flashqla/__init__.py | 28 ++++---------- src/xorl/ops/linear_attention/flashqla_cp.py | 6 +-- .../ops/linear_attention/gdn_decode_prep.py | 6 +-- src/xorl/ops/loss/causallm_loss.py | 2 +- src/xorl/ops/loss/fused_linear_logprob.py | 4 +- src/xorl/ops/quack/__init__.py | 38 ++++--------------- src/xorl/trainers/training_utils.py | 2 +- .../test_flashqla_cp_equivalence.py | 2 +- tests/ops/test_flashqla_contract_pin.py | 4 +- tests/ops/test_flashqla_gdn.py | 2 +- tests/ops/test_quack_process_safety.py | 4 +- tests/server/test_security.py | 6 +-- 107 files changed, 114 insertions(+), 85 deletions(-) create mode 100644 src/xorl/ops/_vendored/__init__.py rename src/xorl/ops/{linear_attention => _vendored}/flashqla/LICENSE (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/USAGE.md (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/VENDORED.md (100%) create mode 100644 src/xorl/ops/_vendored/flashqla/__init__.py rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/__init__.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/gated_delta_rule/__init__.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/gated_delta_rule/chunk/__init__.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/gated_delta_rule/chunk/cp_context.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/gated_delta_rule/chunk/hopper/__init__.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/gated_delta_rule/chunk/hopper/cp_fwd.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/gated_delta_rule/chunk/hopper/fused_bwd.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/gated_delta_rule/chunk/hopper/fused_fwd.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/gated_delta_rule/chunk/hopper/kkt_solve.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/gated_delta_rule/chunk/hopper/prepare_h.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/utils/__init__.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/utils/cumsum.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/ops/utils/group_reduce.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/utils/__init__.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/utils/index.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/utils/math.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/utils/pack.py (100%) rename src/xorl/ops/{linear_attention => _vendored}/flashqla/utils/profiler.py (100%) rename src/xorl/ops/{ => _vendored}/quack/VENDORED.md (100%) create mode 100644 src/xorl/ops/_vendored/quack/__init__.py rename src/xorl/ops/{ => _vendored}/quack/_compile_worker.py (100%) rename src/xorl/ops/{ => _vendored}/quack/_worker_protocol.py (100%) rename src/xorl/ops/{ => _vendored}/quack/activation.py (100%) rename src/xorl/ops/{ => _vendored}/quack/autotuner.py (100%) rename src/xorl/ops/{ => _vendored}/quack/blockscaled_gemm_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/broadcast_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/cache_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/compile_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/copy_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/cross_entropy.py (100%) rename src/xorl/ops/{ => _vendored}/quack/cute_dsl_elf_fix.py (100%) rename src/xorl/ops/{ => _vendored}/quack/cute_dsl_mlir_threading.py (100%) rename src/xorl/ops/{ => _vendored}/quack/cute_dsl_ptxas.py (100%) rename src/xorl/ops/{ => _vendored}/quack/cute_dsl_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/epi_composable.py (100%) rename src/xorl/ops/{ => _vendored}/quack/epi_ops.py (100%) rename src/xorl/ops/{ => _vendored}/quack/epi_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/fast_math.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_act.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_base.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_blockscaled_interface.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_config.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_dact.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_default_epi.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_interface.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_norm_act.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_sm100.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_sm120.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_sm80.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_sm90.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_sq_reduce.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_symmetric.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_tvm_ffi_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/gemm_wrapper_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/layout_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/linear.py (100%) rename src/xorl/ops/{ => _vendored}/quack/linear_cross_entropy.py (100%) rename src/xorl/ops/{ => _vendored}/quack/mlp.py (100%) rename src/xorl/ops/{ => _vendored}/quack/mx_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/nvmmh_heuristic.py (100%) rename src/xorl/ops/{ => _vendored}/quack/pipeline.py (100%) rename src/xorl/ops/{ => _vendored}/quack/reduce.py (100%) rename src/xorl/ops/{ => _vendored}/quack/reduction_base.py (100%) rename src/xorl/ops/{ => _vendored}/quack/rms_final_reduce.py (100%) rename src/xorl/ops/{ => _vendored}/quack/rmsnorm.py (100%) rename src/xorl/ops/{ => _vendored}/quack/rotary.py (100%) rename src/xorl/ops/{ => _vendored}/quack/rounding.py (100%) rename src/xorl/ops/{ => _vendored}/quack/sm100_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/sm80_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/sm90_utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/softmax.py (100%) rename src/xorl/ops/{ => _vendored}/quack/sort/__init__.py (100%) rename src/xorl/ops/{ => _vendored}/quack/sort/bitonic_sort.py (100%) rename src/xorl/ops/{ => _vendored}/quack/sort/generate_sorting_networks.py (100%) rename src/xorl/ops/{ => _vendored}/quack/sort/sorting_networks.py (100%) rename src/xorl/ops/{ => _vendored}/quack/sort/utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/tensormap_manager.py (100%) rename src/xorl/ops/{ => _vendored}/quack/tile_scheduler.py (100%) rename src/xorl/ops/{ => _vendored}/quack/topk.py (100%) rename src/xorl/ops/{ => _vendored}/quack/trace.py (100%) rename src/xorl/ops/{ => _vendored}/quack/utils.py (100%) rename src/xorl/ops/{ => _vendored}/quack/varlen_utils.py (100%) diff --git a/.codespellrc b/.codespellrc index ec6dbc3a..7496ad9a 100644 --- a/.codespellrc +++ b/.codespellrc @@ -1,3 +1,3 @@ [codespell] -skip = *.lock,*.json,submodules/*,.venv/*,.git,docs/node_modules/*,src/xorl/ops/quack/*,src/xorl/ops/linear_attention/flashqla/* +skip = *.lock,*.json,submodules/*,.venv/*,.git,docs/node_modules/*,src/xorl/ops/_vendored/* ignore-words-list = dout,te,subtile,parm,mot,numer,notin diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index aea412a2..250dd41c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ default_stages: [pre-commit, pre-push, manual] # and bi_families_v2.py is vendored byte-identical into the serving engine # with both copies sha256-gated, so it keeps that engine's formatting # (black, 88 columns). See src/xorl/ops/README.md. -exclude: ^(submodules/|\.venv/|src/xorl/ops/quack/|src/xorl/ops/linear_attention/flashqla/|src/xorl/ops/bi_families_v2\.py$) +exclude: ^(submodules/|\.venv/|src/xorl/ops/_vendored/|src/xorl/ops/bi_families_v2\.py$) repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 diff --git a/docs/src/content/docs/moe/kernels.mdx b/docs/src/content/docs/moe/kernels.mdx index eef1613f..5a363fbf 100644 --- a/docs/src/content/docs/moe/kernels.mdx +++ b/docs/src/content/docs/moe/kernels.mdx @@ -236,6 +236,6 @@ The retained tensors and avoided work differ by model and backend. Measure peak | [`src/xorl/ops/group_gemm/kernel/group_gemm.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/group_gemm/kernel/group_gemm.py) | `group_gemm_same_nk_kernel` — Triton autotuned grouped GEMM with optional STORE_ACTIVATIONS | | [`src/xorl/ops/__init__.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/__init__.py) | `triton_moe_forward`, `triton_moe_lora_forward`, backend registry | | [`src/xorl/ops/group_gemm/kernel/quack.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/group_gemm/kernel/quack.py) | `quack_group_gemm_same_nk` — CuTe GEMM wrapper; `XORL_QUACK_TUNED` env var | -| [`src/xorl/ops/quack/gemm_interface.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/quack/gemm_interface.py) | Quack GEMM interface — SM90/SM100 configs, autotuner, gated activation kernels | +| [`src/xorl/ops/_vendored/quack/gemm_interface.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/_vendored/quack/gemm_interface.py) | Quack GEMM interface — SM90/SM100 configs, autotuner, gated activation kernels | | [`src/xorl/models/layers/moe/backend/quack.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/models/layers/moe/backend/quack.py) | `quack_expert_forward` — non-EP quack MoE forward | | [`src/xorl/ops/moe/quack.py`](https://github.com/togethercomputer/xorl/blob/main/src/xorl/ops/moe/quack.py) | `QuackEPGroupGemm`, `QuackEPGroupGemmMoeAct` — EP + moe_act variants | diff --git a/pyproject.toml b/pyproject.toml index fefd626e..bcdcccdf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,7 +65,7 @@ dependencies = [ # PR #2303, merged post-v0.1.10 so absent from every PyPI wheel); we pin # the blessed prebuilt wheel of stock tile-ai/tilelang@a8d93798 (includes # #2303, cp38-abi3, CUDA 13.1) from togethercomputer/xorl-wheels — see - # src/xorl/ops/linear_attention/flashqla/VENDORED.md. Switch to PyPI + # src/xorl/ops/_vendored/flashqla/VENDORED.md. Switch to PyPI # tilelang>=0.1.11 once a release carrying #2303 ships; bump deliberately # + revalidate the GPU smokes (sparse-MLA bwd NaN history and the FlashQLA # contract-pin tests) when picking up a new build. @@ -148,8 +148,7 @@ line-length = 120 # is how pre-commit invokes ruff. force-exclude = true extend-exclude = [ - "src/xorl/ops/quack", - "src/xorl/ops/linear_attention/flashqla", + "src/xorl/ops/_vendored", "src/xorl/ops/bi_families_v2.py", ] diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index 94559c56..fe15a08f 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -10,7 +10,7 @@ Three kinds of code live here with three different rules: | kind | rule | | --- | --- | -| **Vendored** — `quack/`, `linear_attention/flashqla/` | Never hand-edit, lint, or reformat. Each tree carries a `VENDORED.md` with provenance and the local-patch ledger. First-party tooling skips them (`[tool.ruff]` excludes in `pyproject.toml`; top-level `exclude:` in `.pre-commit-config.yaml`). | +| **Vendored** — `_vendored/` (`quack/`, `flashqla/`) | Never hand-edit, lint, or reformat. Each tree carries a `VENDORED.md` with provenance and the local-patch ledger. First-party tooling skips them (`[tool.ruff]` excludes in `pyproject.toml`; top-level `exclude:` in `.pre-commit-config.yaml`). | | **Byte-contract-gated** — `bi_families_v2.py` | Vendored byte-identical into the serving engine; both copies are sha256-gated. Any edit here without the paired serving-side edit breaks the gate. It keeps the engine's formatting (black, 88 columns) and is excluded from all rewriting hooks. | | **First-party** — everything else | Normal rules. | @@ -29,14 +29,13 @@ Three kinds of code live here with three different rules: - `moe/` — MoE expert compute backends (triton/quack/native, LoRA variants). `ep_kernels/` (DeepEP sort/scatter) merges in here (#78 phase 5). - `linear_attention/` — GDN/linear-attention kernels; also currently hosts - the `GatedDeltaNet` layer class and the vendored `flashqla/` tree - (layer moves to `models/layers/`, vendored tree to `ops/_vendored/`, #78). + the `GatedDeltaNet` layer class (moves to `models/layers/`, #78 phase 4). - `ssm/` — Mamba-2 kernels; also currently hosts the `Mamba2Mixer` layer class (same plan as above). - `quantize/` — NF4/INT4/FP4/FP8 quantization codecs and fake-quant ops. - `glm5_kernels/`, `dsv4/` — model-family-specific kernels; planned home: `models/transformers/{glm5,deepseek_v4}/kernels/` (#78 phase 4). -- `quack/` — vendored (see above). +- `_vendored/` — vendored trees (see above); old paths (`ops/quack`, `ops/linear_attention/flashqla`) are alias stubs for one deprecation cycle. ## What does NOT belong here diff --git a/src/xorl/ops/_vendored/__init__.py b/src/xorl/ops/_vendored/__init__.py new file mode 100644 index 00000000..04ef9c29 --- /dev/null +++ b/src/xorl/ops/_vendored/__init__.py @@ -0,0 +1,7 @@ +"""Vendored third-party trees (issue #78 phase 1). + +Every subtree here is a snapshot of external code with its own +``VENDORED.md`` (provenance, upstream pin, local-patch ledger). Nothing in +this namespace is linted, formatted, or hand-edited; see +``src/xorl/ops/README.md`` for the edit policy. +""" diff --git a/src/xorl/ops/linear_attention/flashqla/LICENSE b/src/xorl/ops/_vendored/flashqla/LICENSE similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/LICENSE rename to src/xorl/ops/_vendored/flashqla/LICENSE diff --git a/src/xorl/ops/linear_attention/flashqla/USAGE.md b/src/xorl/ops/_vendored/flashqla/USAGE.md similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/USAGE.md rename to src/xorl/ops/_vendored/flashqla/USAGE.md diff --git a/src/xorl/ops/linear_attention/flashqla/VENDORED.md b/src/xorl/ops/_vendored/flashqla/VENDORED.md similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/VENDORED.md rename to src/xorl/ops/_vendored/flashqla/VENDORED.md diff --git a/src/xorl/ops/_vendored/flashqla/__init__.py b/src/xorl/ops/_vendored/flashqla/__init__.py new file mode 100644 index 00000000..a85bd53c --- /dev/null +++ b/src/xorl/ops/_vendored/flashqla/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) 2026 The Qwen team, Alibaba Group. +# Licensed under The MIT License [see LICENSE for details] + +__version__ = "0.1.0" + +# Local (xorl) modification: the kernels call `T.gemm_v1`, which stock tilelang >=0.1.9 no +# longer exposes. Re-add it (fast tl::gemm_ss template) before the kernels are traced. +from xorl.ops.linear_attention.tilelang_gemm_v1 import patch as _patch_gemm_v1 + + +_patch_gemm_v1() + +from xorl.ops.linear_attention.flashqla.ops.gated_delta_rule.chunk import ( # noqa: E402 + chunk_gated_delta_rule, + chunk_gated_delta_rule_bwd, + chunk_gated_delta_rule_fwd, +) + + +__all__ = [ + "chunk_gated_delta_rule_fwd", + "chunk_gated_delta_rule_bwd", + "chunk_gated_delta_rule", +] diff --git a/src/xorl/ops/linear_attention/flashqla/ops/__init__.py b/src/xorl/ops/_vendored/flashqla/ops/__init__.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/__init__.py rename to src/xorl/ops/_vendored/flashqla/ops/__init__.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/__init__.py b/src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/__init__.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/__init__.py rename to src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/__init__.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/__init__.py b/src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/__init__.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/__init__.py rename to src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/__init__.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/cp_context.py b/src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/cp_context.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/cp_context.py rename to src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/cp_context.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/__init__.py b/src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/__init__.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/__init__.py rename to src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/__init__.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/cp_fwd.py b/src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/cp_fwd.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/cp_fwd.py rename to src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/cp_fwd.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/fused_bwd.py b/src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/fused_bwd.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/fused_bwd.py rename to src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/fused_bwd.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/fused_fwd.py b/src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/fused_fwd.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/fused_fwd.py rename to src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/fused_fwd.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/kkt_solve.py b/src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/kkt_solve.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/kkt_solve.py rename to src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/kkt_solve.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/prepare_h.py b/src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/prepare_h.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/gated_delta_rule/chunk/hopper/prepare_h.py rename to src/xorl/ops/_vendored/flashqla/ops/gated_delta_rule/chunk/hopper/prepare_h.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/utils/__init__.py b/src/xorl/ops/_vendored/flashqla/ops/utils/__init__.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/utils/__init__.py rename to src/xorl/ops/_vendored/flashqla/ops/utils/__init__.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/utils/cumsum.py b/src/xorl/ops/_vendored/flashqla/ops/utils/cumsum.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/utils/cumsum.py rename to src/xorl/ops/_vendored/flashqla/ops/utils/cumsum.py diff --git a/src/xorl/ops/linear_attention/flashqla/ops/utils/group_reduce.py b/src/xorl/ops/_vendored/flashqla/ops/utils/group_reduce.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/ops/utils/group_reduce.py rename to src/xorl/ops/_vendored/flashqla/ops/utils/group_reduce.py diff --git a/src/xorl/ops/linear_attention/flashqla/utils/__init__.py b/src/xorl/ops/_vendored/flashqla/utils/__init__.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/utils/__init__.py rename to src/xorl/ops/_vendored/flashqla/utils/__init__.py diff --git a/src/xorl/ops/linear_attention/flashqla/utils/index.py b/src/xorl/ops/_vendored/flashqla/utils/index.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/utils/index.py rename to src/xorl/ops/_vendored/flashqla/utils/index.py diff --git a/src/xorl/ops/linear_attention/flashqla/utils/math.py b/src/xorl/ops/_vendored/flashqla/utils/math.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/utils/math.py rename to src/xorl/ops/_vendored/flashqla/utils/math.py diff --git a/src/xorl/ops/linear_attention/flashqla/utils/pack.py b/src/xorl/ops/_vendored/flashqla/utils/pack.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/utils/pack.py rename to src/xorl/ops/_vendored/flashqla/utils/pack.py diff --git a/src/xorl/ops/linear_attention/flashqla/utils/profiler.py b/src/xorl/ops/_vendored/flashqla/utils/profiler.py similarity index 100% rename from src/xorl/ops/linear_attention/flashqla/utils/profiler.py rename to src/xorl/ops/_vendored/flashqla/utils/profiler.py diff --git a/src/xorl/ops/quack/VENDORED.md b/src/xorl/ops/_vendored/quack/VENDORED.md similarity index 100% rename from src/xorl/ops/quack/VENDORED.md rename to src/xorl/ops/_vendored/quack/VENDORED.md diff --git a/src/xorl/ops/_vendored/quack/__init__.py b/src/xorl/ops/_vendored/quack/__init__.py new file mode 100644 index 00000000..939027b4 --- /dev/null +++ b/src/xorl/ops/_vendored/quack/__init__.py @@ -0,0 +1,34 @@ +__version__ = "0.4.1" + +import os + +# Two CuTeDSL workarounds, both must run before the first cute.compile call: +# - cutlass#3161: duplicate .text section flags break MCJIT in multi-process +# loads (see cute_dsl_elf_fix). +# - cutlass#3062: ir.Context spawns LLVM thread pools that leak across +# compiles, eventually exhausting pthreads (see cute_dsl_mlir_threading). +from . import cute_dsl_elf_fix +from . import cute_dsl_mlir_threading + +cute_dsl_elf_fix.patch() +cute_dsl_mlir_threading.patch() + +from .rmsnorm import rmsnorm # noqa: E402 +from .softmax import softmax # noqa: E402 +from .cross_entropy import cross_entropy # noqa: E402 +from .rounding import RoundingMode # noqa: E402 + + +if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None: + from . import cute_dsl_ptxas # noqa: F401 + + # Patch to dump ptx and then use system ptxas to compile to cubin + cute_dsl_ptxas.patch() + + +__all__ = [ + "rmsnorm", + "softmax", + "cross_entropy", + "RoundingMode", +] diff --git a/src/xorl/ops/quack/_compile_worker.py b/src/xorl/ops/_vendored/quack/_compile_worker.py similarity index 100% rename from src/xorl/ops/quack/_compile_worker.py rename to src/xorl/ops/_vendored/quack/_compile_worker.py diff --git a/src/xorl/ops/quack/_worker_protocol.py b/src/xorl/ops/_vendored/quack/_worker_protocol.py similarity index 100% rename from src/xorl/ops/quack/_worker_protocol.py rename to src/xorl/ops/_vendored/quack/_worker_protocol.py diff --git a/src/xorl/ops/quack/activation.py b/src/xorl/ops/_vendored/quack/activation.py similarity index 100% rename from src/xorl/ops/quack/activation.py rename to src/xorl/ops/_vendored/quack/activation.py diff --git a/src/xorl/ops/quack/autotuner.py b/src/xorl/ops/_vendored/quack/autotuner.py similarity index 100% rename from src/xorl/ops/quack/autotuner.py rename to src/xorl/ops/_vendored/quack/autotuner.py diff --git a/src/xorl/ops/quack/blockscaled_gemm_utils.py b/src/xorl/ops/_vendored/quack/blockscaled_gemm_utils.py similarity index 100% rename from src/xorl/ops/quack/blockscaled_gemm_utils.py rename to src/xorl/ops/_vendored/quack/blockscaled_gemm_utils.py diff --git a/src/xorl/ops/quack/broadcast_utils.py b/src/xorl/ops/_vendored/quack/broadcast_utils.py similarity index 100% rename from src/xorl/ops/quack/broadcast_utils.py rename to src/xorl/ops/_vendored/quack/broadcast_utils.py diff --git a/src/xorl/ops/quack/cache_utils.py b/src/xorl/ops/_vendored/quack/cache_utils.py similarity index 100% rename from src/xorl/ops/quack/cache_utils.py rename to src/xorl/ops/_vendored/quack/cache_utils.py diff --git a/src/xorl/ops/quack/compile_utils.py b/src/xorl/ops/_vendored/quack/compile_utils.py similarity index 100% rename from src/xorl/ops/quack/compile_utils.py rename to src/xorl/ops/_vendored/quack/compile_utils.py diff --git a/src/xorl/ops/quack/copy_utils.py b/src/xorl/ops/_vendored/quack/copy_utils.py similarity index 100% rename from src/xorl/ops/quack/copy_utils.py rename to src/xorl/ops/_vendored/quack/copy_utils.py diff --git a/src/xorl/ops/quack/cross_entropy.py b/src/xorl/ops/_vendored/quack/cross_entropy.py similarity index 100% rename from src/xorl/ops/quack/cross_entropy.py rename to src/xorl/ops/_vendored/quack/cross_entropy.py diff --git a/src/xorl/ops/quack/cute_dsl_elf_fix.py b/src/xorl/ops/_vendored/quack/cute_dsl_elf_fix.py similarity index 100% rename from src/xorl/ops/quack/cute_dsl_elf_fix.py rename to src/xorl/ops/_vendored/quack/cute_dsl_elf_fix.py diff --git a/src/xorl/ops/quack/cute_dsl_mlir_threading.py b/src/xorl/ops/_vendored/quack/cute_dsl_mlir_threading.py similarity index 100% rename from src/xorl/ops/quack/cute_dsl_mlir_threading.py rename to src/xorl/ops/_vendored/quack/cute_dsl_mlir_threading.py diff --git a/src/xorl/ops/quack/cute_dsl_ptxas.py b/src/xorl/ops/_vendored/quack/cute_dsl_ptxas.py similarity index 100% rename from src/xorl/ops/quack/cute_dsl_ptxas.py rename to src/xorl/ops/_vendored/quack/cute_dsl_ptxas.py diff --git a/src/xorl/ops/quack/cute_dsl_utils.py b/src/xorl/ops/_vendored/quack/cute_dsl_utils.py similarity index 100% rename from src/xorl/ops/quack/cute_dsl_utils.py rename to src/xorl/ops/_vendored/quack/cute_dsl_utils.py diff --git a/src/xorl/ops/quack/epi_composable.py b/src/xorl/ops/_vendored/quack/epi_composable.py similarity index 100% rename from src/xorl/ops/quack/epi_composable.py rename to src/xorl/ops/_vendored/quack/epi_composable.py diff --git a/src/xorl/ops/quack/epi_ops.py b/src/xorl/ops/_vendored/quack/epi_ops.py similarity index 100% rename from src/xorl/ops/quack/epi_ops.py rename to src/xorl/ops/_vendored/quack/epi_ops.py diff --git a/src/xorl/ops/quack/epi_utils.py b/src/xorl/ops/_vendored/quack/epi_utils.py similarity index 100% rename from src/xorl/ops/quack/epi_utils.py rename to src/xorl/ops/_vendored/quack/epi_utils.py diff --git a/src/xorl/ops/quack/fast_math.py b/src/xorl/ops/_vendored/quack/fast_math.py similarity index 100% rename from src/xorl/ops/quack/fast_math.py rename to src/xorl/ops/_vendored/quack/fast_math.py diff --git a/src/xorl/ops/quack/gemm.py b/src/xorl/ops/_vendored/quack/gemm.py similarity index 100% rename from src/xorl/ops/quack/gemm.py rename to src/xorl/ops/_vendored/quack/gemm.py diff --git a/src/xorl/ops/quack/gemm_act.py b/src/xorl/ops/_vendored/quack/gemm_act.py similarity index 100% rename from src/xorl/ops/quack/gemm_act.py rename to src/xorl/ops/_vendored/quack/gemm_act.py diff --git a/src/xorl/ops/quack/gemm_base.py b/src/xorl/ops/_vendored/quack/gemm_base.py similarity index 100% rename from src/xorl/ops/quack/gemm_base.py rename to src/xorl/ops/_vendored/quack/gemm_base.py diff --git a/src/xorl/ops/quack/gemm_blockscaled_interface.py b/src/xorl/ops/_vendored/quack/gemm_blockscaled_interface.py similarity index 100% rename from src/xorl/ops/quack/gemm_blockscaled_interface.py rename to src/xorl/ops/_vendored/quack/gemm_blockscaled_interface.py diff --git a/src/xorl/ops/quack/gemm_config.py b/src/xorl/ops/_vendored/quack/gemm_config.py similarity index 100% rename from src/xorl/ops/quack/gemm_config.py rename to src/xorl/ops/_vendored/quack/gemm_config.py diff --git a/src/xorl/ops/quack/gemm_dact.py b/src/xorl/ops/_vendored/quack/gemm_dact.py similarity index 100% rename from src/xorl/ops/quack/gemm_dact.py rename to src/xorl/ops/_vendored/quack/gemm_dact.py diff --git a/src/xorl/ops/quack/gemm_default_epi.py b/src/xorl/ops/_vendored/quack/gemm_default_epi.py similarity index 100% rename from src/xorl/ops/quack/gemm_default_epi.py rename to src/xorl/ops/_vendored/quack/gemm_default_epi.py diff --git a/src/xorl/ops/quack/gemm_interface.py b/src/xorl/ops/_vendored/quack/gemm_interface.py similarity index 100% rename from src/xorl/ops/quack/gemm_interface.py rename to src/xorl/ops/_vendored/quack/gemm_interface.py diff --git a/src/xorl/ops/quack/gemm_norm_act.py b/src/xorl/ops/_vendored/quack/gemm_norm_act.py similarity index 100% rename from src/xorl/ops/quack/gemm_norm_act.py rename to src/xorl/ops/_vendored/quack/gemm_norm_act.py diff --git a/src/xorl/ops/quack/gemm_sm100.py b/src/xorl/ops/_vendored/quack/gemm_sm100.py similarity index 100% rename from src/xorl/ops/quack/gemm_sm100.py rename to src/xorl/ops/_vendored/quack/gemm_sm100.py diff --git a/src/xorl/ops/quack/gemm_sm120.py b/src/xorl/ops/_vendored/quack/gemm_sm120.py similarity index 100% rename from src/xorl/ops/quack/gemm_sm120.py rename to src/xorl/ops/_vendored/quack/gemm_sm120.py diff --git a/src/xorl/ops/quack/gemm_sm80.py b/src/xorl/ops/_vendored/quack/gemm_sm80.py similarity index 100% rename from src/xorl/ops/quack/gemm_sm80.py rename to src/xorl/ops/_vendored/quack/gemm_sm80.py diff --git a/src/xorl/ops/quack/gemm_sm90.py b/src/xorl/ops/_vendored/quack/gemm_sm90.py similarity index 100% rename from src/xorl/ops/quack/gemm_sm90.py rename to src/xorl/ops/_vendored/quack/gemm_sm90.py diff --git a/src/xorl/ops/quack/gemm_sq_reduce.py b/src/xorl/ops/_vendored/quack/gemm_sq_reduce.py similarity index 100% rename from src/xorl/ops/quack/gemm_sq_reduce.py rename to src/xorl/ops/_vendored/quack/gemm_sq_reduce.py diff --git a/src/xorl/ops/quack/gemm_symmetric.py b/src/xorl/ops/_vendored/quack/gemm_symmetric.py similarity index 100% rename from src/xorl/ops/quack/gemm_symmetric.py rename to src/xorl/ops/_vendored/quack/gemm_symmetric.py diff --git a/src/xorl/ops/quack/gemm_tvm_ffi_utils.py b/src/xorl/ops/_vendored/quack/gemm_tvm_ffi_utils.py similarity index 100% rename from src/xorl/ops/quack/gemm_tvm_ffi_utils.py rename to src/xorl/ops/_vendored/quack/gemm_tvm_ffi_utils.py diff --git a/src/xorl/ops/quack/gemm_wrapper_utils.py b/src/xorl/ops/_vendored/quack/gemm_wrapper_utils.py similarity index 100% rename from src/xorl/ops/quack/gemm_wrapper_utils.py rename to src/xorl/ops/_vendored/quack/gemm_wrapper_utils.py diff --git a/src/xorl/ops/quack/layout_utils.py b/src/xorl/ops/_vendored/quack/layout_utils.py similarity index 100% rename from src/xorl/ops/quack/layout_utils.py rename to src/xorl/ops/_vendored/quack/layout_utils.py diff --git a/src/xorl/ops/quack/linear.py b/src/xorl/ops/_vendored/quack/linear.py similarity index 100% rename from src/xorl/ops/quack/linear.py rename to src/xorl/ops/_vendored/quack/linear.py diff --git a/src/xorl/ops/quack/linear_cross_entropy.py b/src/xorl/ops/_vendored/quack/linear_cross_entropy.py similarity index 100% rename from src/xorl/ops/quack/linear_cross_entropy.py rename to src/xorl/ops/_vendored/quack/linear_cross_entropy.py diff --git a/src/xorl/ops/quack/mlp.py b/src/xorl/ops/_vendored/quack/mlp.py similarity index 100% rename from src/xorl/ops/quack/mlp.py rename to src/xorl/ops/_vendored/quack/mlp.py diff --git a/src/xorl/ops/quack/mx_utils.py b/src/xorl/ops/_vendored/quack/mx_utils.py similarity index 100% rename from src/xorl/ops/quack/mx_utils.py rename to src/xorl/ops/_vendored/quack/mx_utils.py diff --git a/src/xorl/ops/quack/nvmmh_heuristic.py b/src/xorl/ops/_vendored/quack/nvmmh_heuristic.py similarity index 100% rename from src/xorl/ops/quack/nvmmh_heuristic.py rename to src/xorl/ops/_vendored/quack/nvmmh_heuristic.py diff --git a/src/xorl/ops/quack/pipeline.py b/src/xorl/ops/_vendored/quack/pipeline.py similarity index 100% rename from src/xorl/ops/quack/pipeline.py rename to src/xorl/ops/_vendored/quack/pipeline.py diff --git a/src/xorl/ops/quack/reduce.py b/src/xorl/ops/_vendored/quack/reduce.py similarity index 100% rename from src/xorl/ops/quack/reduce.py rename to src/xorl/ops/_vendored/quack/reduce.py diff --git a/src/xorl/ops/quack/reduction_base.py b/src/xorl/ops/_vendored/quack/reduction_base.py similarity index 100% rename from src/xorl/ops/quack/reduction_base.py rename to src/xorl/ops/_vendored/quack/reduction_base.py diff --git a/src/xorl/ops/quack/rms_final_reduce.py b/src/xorl/ops/_vendored/quack/rms_final_reduce.py similarity index 100% rename from src/xorl/ops/quack/rms_final_reduce.py rename to src/xorl/ops/_vendored/quack/rms_final_reduce.py diff --git a/src/xorl/ops/quack/rmsnorm.py b/src/xorl/ops/_vendored/quack/rmsnorm.py similarity index 100% rename from src/xorl/ops/quack/rmsnorm.py rename to src/xorl/ops/_vendored/quack/rmsnorm.py diff --git a/src/xorl/ops/quack/rotary.py b/src/xorl/ops/_vendored/quack/rotary.py similarity index 100% rename from src/xorl/ops/quack/rotary.py rename to src/xorl/ops/_vendored/quack/rotary.py diff --git a/src/xorl/ops/quack/rounding.py b/src/xorl/ops/_vendored/quack/rounding.py similarity index 100% rename from src/xorl/ops/quack/rounding.py rename to src/xorl/ops/_vendored/quack/rounding.py diff --git a/src/xorl/ops/quack/sm100_utils.py b/src/xorl/ops/_vendored/quack/sm100_utils.py similarity index 100% rename from src/xorl/ops/quack/sm100_utils.py rename to src/xorl/ops/_vendored/quack/sm100_utils.py diff --git a/src/xorl/ops/quack/sm80_utils.py b/src/xorl/ops/_vendored/quack/sm80_utils.py similarity index 100% rename from src/xorl/ops/quack/sm80_utils.py rename to src/xorl/ops/_vendored/quack/sm80_utils.py diff --git a/src/xorl/ops/quack/sm90_utils.py b/src/xorl/ops/_vendored/quack/sm90_utils.py similarity index 100% rename from src/xorl/ops/quack/sm90_utils.py rename to src/xorl/ops/_vendored/quack/sm90_utils.py diff --git a/src/xorl/ops/quack/softmax.py b/src/xorl/ops/_vendored/quack/softmax.py similarity index 100% rename from src/xorl/ops/quack/softmax.py rename to src/xorl/ops/_vendored/quack/softmax.py diff --git a/src/xorl/ops/quack/sort/__init__.py b/src/xorl/ops/_vendored/quack/sort/__init__.py similarity index 100% rename from src/xorl/ops/quack/sort/__init__.py rename to src/xorl/ops/_vendored/quack/sort/__init__.py diff --git a/src/xorl/ops/quack/sort/bitonic_sort.py b/src/xorl/ops/_vendored/quack/sort/bitonic_sort.py similarity index 100% rename from src/xorl/ops/quack/sort/bitonic_sort.py rename to src/xorl/ops/_vendored/quack/sort/bitonic_sort.py diff --git a/src/xorl/ops/quack/sort/generate_sorting_networks.py b/src/xorl/ops/_vendored/quack/sort/generate_sorting_networks.py similarity index 100% rename from src/xorl/ops/quack/sort/generate_sorting_networks.py rename to src/xorl/ops/_vendored/quack/sort/generate_sorting_networks.py diff --git a/src/xorl/ops/quack/sort/sorting_networks.py b/src/xorl/ops/_vendored/quack/sort/sorting_networks.py similarity index 100% rename from src/xorl/ops/quack/sort/sorting_networks.py rename to src/xorl/ops/_vendored/quack/sort/sorting_networks.py diff --git a/src/xorl/ops/quack/sort/utils.py b/src/xorl/ops/_vendored/quack/sort/utils.py similarity index 100% rename from src/xorl/ops/quack/sort/utils.py rename to src/xorl/ops/_vendored/quack/sort/utils.py diff --git a/src/xorl/ops/quack/tensormap_manager.py b/src/xorl/ops/_vendored/quack/tensormap_manager.py similarity index 100% rename from src/xorl/ops/quack/tensormap_manager.py rename to src/xorl/ops/_vendored/quack/tensormap_manager.py diff --git a/src/xorl/ops/quack/tile_scheduler.py b/src/xorl/ops/_vendored/quack/tile_scheduler.py similarity index 100% rename from src/xorl/ops/quack/tile_scheduler.py rename to src/xorl/ops/_vendored/quack/tile_scheduler.py diff --git a/src/xorl/ops/quack/topk.py b/src/xorl/ops/_vendored/quack/topk.py similarity index 100% rename from src/xorl/ops/quack/topk.py rename to src/xorl/ops/_vendored/quack/topk.py diff --git a/src/xorl/ops/quack/trace.py b/src/xorl/ops/_vendored/quack/trace.py similarity index 100% rename from src/xorl/ops/quack/trace.py rename to src/xorl/ops/_vendored/quack/trace.py diff --git a/src/xorl/ops/quack/utils.py b/src/xorl/ops/_vendored/quack/utils.py similarity index 100% rename from src/xorl/ops/quack/utils.py rename to src/xorl/ops/_vendored/quack/utils.py diff --git a/src/xorl/ops/quack/varlen_utils.py b/src/xorl/ops/_vendored/quack/varlen_utils.py similarity index 100% rename from src/xorl/ops/quack/varlen_utils.py rename to src/xorl/ops/_vendored/quack/varlen_utils.py diff --git a/src/xorl/ops/group_gemm/kernel/quack.py b/src/xorl/ops/group_gemm/kernel/quack.py index 5ac9c10f..28141d74 100644 --- a/src/xorl/ops/group_gemm/kernel/quack.py +++ b/src/xorl/ops/group_gemm/kernel/quack.py @@ -30,7 +30,7 @@ def _gemm_interface(): # quack backend (e.g. the torch-2.11 combined environment used by the # DSV4 exact lane, which carries SGLang's cutlass-dsl pin) must still be # able to import this package. - from xorl.ops.quack import gemm_interface # noqa: PLC0415 + from xorl.ops._vendored.quack import gemm_interface # noqa: PLC0415 return gemm_interface diff --git a/src/xorl/ops/linear_attention/backend.py b/src/xorl/ops/linear_attention/backend.py index b2789bca..a9013f53 100644 --- a/src/xorl/ops/linear_attention/backend.py +++ b/src/xorl/ops/linear_attention/backend.py @@ -6,10 +6,10 @@ (:mod:`xorl.ops.linear_attention.ops.gated_delta_rule`). Default. Supports Ulysses context parallelism via ``cp_context``. * ``flashqla`` — the vendored Qwen FlashQLA fused TileLang kernels - (:mod:`xorl.ops.linear_attention.flashqla`). Hopper (SM90) only, + (:mod:`xorl.ops._vendored.flashqla`). Hopper (SM90) only, requires ``tilelang``. Faster fwd/bwd. The single-GPU kernel is CP-unaware; under Ulysses CP it is driven by xorl's native CP - orchestration via :mod:`xorl.ops.linear_attention.flashqla_cp` + orchestration via :mod:`xorl.ops._vendored.flashqla_cp` (see :func:`flashqla_chunk_gated_delta_rule_cp`). Select the backend with the ``XORL_GDN_BACKEND`` environment variable @@ -109,7 +109,7 @@ def flashqla_chunk_gated_delta_rule(**kwargs: Any) -> Any: from xorl.ops.linear_attention import tilelang_gemm_v1 # noqa: PLC0415 tilelang_gemm_v1.patch() - from xorl.ops.linear_attention.flashqla import chunk_gated_delta_rule as _chunk # noqa: PLC0415 + from xorl.ops._vendored.flashqla import chunk_gated_delta_rule as _chunk # noqa: PLC0415 except Exception as exc: # pragma: no cover - environment dependent raise RuntimeError( f"Failed to import the FlashQLA GDN backend ({GDN_BACKEND_ENV}=flashqla). " @@ -127,7 +127,7 @@ def flashqla_chunk_gated_delta_rule_cp(**kwargs: Any) -> Any: Like :func:`flashqla_chunk_gated_delta_rule`, the import is deferred (FlashQLA pulls in ``tilelang`` and validates SM90 at import time). This path drives the FlashQLA interior with xorl's Ulysses/sequence-parallel orchestration; see - :mod:`xorl.ops.linear_attention.flashqla_cp`. + :mod:`xorl.ops._vendored.flashqla_cp`. """ global _flashqla_chunk_cp if _flashqla_chunk_cp is None: @@ -135,7 +135,7 @@ def flashqla_chunk_gated_delta_rule_cp(**kwargs: Any) -> Any: from xorl.ops.linear_attention import tilelang_gemm_v1 # noqa: PLC0415 tilelang_gemm_v1.patch() - from xorl.ops.linear_attention.flashqla_cp import ( # noqa: PLC0415 + from xorl.ops._vendored.flashqla_cp import ( # noqa: PLC0415 flashqla_chunk_gated_delta_rule_cp as _chunk_cp, ) except Exception as exc: # pragma: no cover - environment dependent diff --git a/src/xorl/ops/linear_attention/flashqla/__init__.py b/src/xorl/ops/linear_attention/flashqla/__init__.py index a85bd53c..1e80f8a4 100644 --- a/src/xorl/ops/linear_attention/flashqla/__init__.py +++ b/src/xorl/ops/linear_attention/flashqla/__init__.py @@ -1,24 +1,12 @@ -# Copyright (c) 2026 The Qwen team, Alibaba Group. -# Licensed under The MIT License [see LICENSE for details] +"""Moved to ``xorl.ops._vendored.flashqla`` (issue #78 phase 1). -__version__ = "0.1.0" +This stub aliases the package so old-path imports keep resolving for one +deprecation cycle. Deep submodule imports through the old path create a +duplicate module instance — switch to the new path. +""" -# Local (xorl) modification: the kernels call `T.gemm_v1`, which stock tilelang >=0.1.9 no -# longer exposes. Re-add it (fast tl::gemm_ss template) before the kernels are traced. -from xorl.ops.linear_attention.tilelang_gemm_v1 import patch as _patch_gemm_v1 +import importlib as _importlib +import sys as _sys -_patch_gemm_v1() - -from xorl.ops.linear_attention.flashqla.ops.gated_delta_rule.chunk import ( # noqa: E402 - chunk_gated_delta_rule, - chunk_gated_delta_rule_bwd, - chunk_gated_delta_rule_fwd, -) - - -__all__ = [ - "chunk_gated_delta_rule_fwd", - "chunk_gated_delta_rule_bwd", - "chunk_gated_delta_rule", -] +_sys.modules[__name__] = _importlib.import_module("xorl.ops._vendored.flashqla") diff --git a/src/xorl/ops/linear_attention/flashqla_cp.py b/src/xorl/ops/linear_attention/flashqla_cp.py index c560eda6..0d4817a6 100644 --- a/src/xorl/ops/linear_attention/flashqla_cp.py +++ b/src/xorl/ops/linear_attention/flashqla_cp.py @@ -30,13 +30,13 @@ import torch -from xorl.ops.linear_attention.flashqla.ops.gated_delta_rule.chunk import robust_kkt_solve -from xorl.ops.linear_attention.flashqla.ops.gated_delta_rule.chunk.hopper import ( +from xorl.ops._vendored.flashqla.ops.gated_delta_rule.chunk import robust_kkt_solve +from xorl.ops._vendored.flashqla.ops.gated_delta_rule.chunk.hopper import ( fused_gdr_bwd, fused_gdr_fwd, fused_gdr_h, ) -from xorl.ops.linear_attention.flashqla.ops.utils import chunk_local_cumsum as flashqla_chunk_local_cumsum +from xorl.ops._vendored.flashqla.ops.utils import chunk_local_cumsum as flashqla_chunk_local_cumsum from xorl.ops.linear_attention.modules.l2norm import l2norm_bwd, l2norm_fwd from xorl.ops.linear_attention.ops.common.chunk_o import chunk_bwd_dv_local from xorl.ops.linear_attention.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd diff --git a/src/xorl/ops/linear_attention/gdn_decode_prep.py b/src/xorl/ops/linear_attention/gdn_decode_prep.py index db5bd38e..a03852e3 100644 --- a/src/xorl/ops/linear_attention/gdn_decode_prep.py +++ b/src/xorl/ops/linear_attention/gdn_decode_prep.py @@ -42,9 +42,9 @@ def chunk_gated_delta_rule_fwd_decode( from xorl.ops.linear_attention import tilelang_gemm_v1 # noqa: PLC0415 tilelang_gemm_v1.patch() - from xorl.ops.linear_attention.flashqla.ops.gated_delta_rule.chunk.hopper import fused_gdr_fwd # noqa: PLC0415 - from xorl.ops.linear_attention.flashqla.ops.utils import chunk_local_cumsum # noqa: PLC0415 - from xorl.ops.linear_attention.flashqla.utils import l2norm # noqa: PLC0415 + from xorl.ops._vendored.flashqla.ops.gated_delta_rule.chunk.hopper import fused_gdr_fwd # noqa: PLC0415 + from xorl.ops._vendored.flashqla.ops.utils import chunk_local_cumsum # noqa: PLC0415 + from xorl.ops._vendored.flashqla.utils import l2norm # noqa: PLC0415 from xorl.ops.linear_attention.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd # noqa: PLC0415 if scale is None: diff --git a/src/xorl/ops/loss/causallm_loss.py b/src/xorl/ops/loss/causallm_loss.py index 1415a2a5..c00323d1 100644 --- a/src/xorl/ops/loss/causallm_loss.py +++ b/src/xorl/ops/loss/causallm_loss.py @@ -307,7 +307,7 @@ def _quack_linear_cross_entropy_loss( if hidden_states_flat.shape[-1] % 8 != 0 or weight.shape[0] % 8 != 0: raise ValueError("ce_mode='quack_linear' requires hidden and vocab dimensions to be divisible by 8") - from xorl.ops.quack.linear_cross_entropy import chunked_linear_cross_entropy # noqa: PLC0415 + from xorl.ops._vendored.quack.linear_cross_entropy import chunked_linear_cross_entropy # noqa: PLC0415 hidden_states_flat, labels_flat = _pad_quack_linear_rows( hidden_states_flat, diff --git a/src/xorl/ops/loss/fused_linear_logprob.py b/src/xorl/ops/loss/fused_linear_logprob.py index 78131549..63b794ed 100644 --- a/src/xorl/ops/loss/fused_linear_logprob.py +++ b/src/xorl/ops/loss/fused_linear_logprob.py @@ -56,8 +56,8 @@ try: - from xorl.ops.quack.cross_entropy import cross_entropy_bwd_out as _quack_ce_bwd_out - from xorl.ops.quack.cross_entropy import cross_entropy_fwd_out as _quack_ce_fwd_out + from xorl.ops._vendored.quack.cross_entropy import cross_entropy_bwd_out as _quack_ce_bwd_out + from xorl.ops._vendored.quack.cross_entropy import cross_entropy_fwd_out as _quack_ce_fwd_out _HAS_QUACK = True except Exception: # pragma: no cover - quack/CuTeDSL optional diff --git a/src/xorl/ops/quack/__init__.py b/src/xorl/ops/quack/__init__.py index 939027b4..d81e0324 100644 --- a/src/xorl/ops/quack/__init__.py +++ b/src/xorl/ops/quack/__init__.py @@ -1,34 +1,12 @@ -__version__ = "0.4.1" +"""Moved to ``xorl.ops._vendored.quack`` (issue #78 phase 1). -import os +This stub aliases the package so old-path imports keep resolving for one +deprecation cycle. Deep submodule imports through the old path create a +duplicate module instance — switch to the new path. +""" -# Two CuTeDSL workarounds, both must run before the first cute.compile call: -# - cutlass#3161: duplicate .text section flags break MCJIT in multi-process -# loads (see cute_dsl_elf_fix). -# - cutlass#3062: ir.Context spawns LLVM thread pools that leak across -# compiles, eventually exhausting pthreads (see cute_dsl_mlir_threading). -from . import cute_dsl_elf_fix -from . import cute_dsl_mlir_threading +import importlib as _importlib +import sys as _sys -cute_dsl_elf_fix.patch() -cute_dsl_mlir_threading.patch() -from .rmsnorm import rmsnorm # noqa: E402 -from .softmax import softmax # noqa: E402 -from .cross_entropy import cross_entropy # noqa: E402 -from .rounding import RoundingMode # noqa: E402 - - -if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None: - from . import cute_dsl_ptxas # noqa: F401 - - # Patch to dump ptx and then use system ptxas to compile to cubin - cute_dsl_ptxas.patch() - - -__all__ = [ - "rmsnorm", - "softmax", - "cross_entropy", - "RoundingMode", -] +_sys.modules[__name__] = _importlib.import_module("xorl.ops._vendored.quack") diff --git a/src/xorl/trainers/training_utils.py b/src/xorl/trainers/training_utils.py index 6c17f5dd..41966006 100644 --- a/src/xorl/trainers/training_utils.py +++ b/src/xorl/trainers/training_utils.py @@ -442,8 +442,8 @@ def _pp_quack_linear_ce_sum(hidden, labels, *, lm_head, num_chunks: int = 8): hidden (pipeline) and lm_head.weight (its FSDP unit reduce-scatters them). """ from xorl.models.module_utils import get_lm_head_weight # noqa: PLC0415 + from xorl.ops._vendored.quack.linear_cross_entropy import chunked_linear_cross_entropy # noqa: PLC0415 from xorl.ops.loss.causallm_loss import _chunk_size_from_num_chunks # noqa: PLC0415 - from xorl.ops.quack.linear_cross_entropy import chunked_linear_cross_entropy # noqa: PLC0415 weight = get_lm_head_weight(lm_head, fsdp_sharded_loss=False) h = hidden.reshape(-1, hidden.shape[-1]) diff --git a/tests/distributed/test_flashqla_cp_equivalence.py b/tests/distributed/test_flashqla_cp_equivalence.py index 4abf97f3..3eda4e57 100644 --- a/tests/distributed/test_flashqla_cp_equivalence.py +++ b/tests/distributed/test_flashqla_cp_equivalence.py @@ -19,7 +19,7 @@ pytest.skip("FlashQLA CP parity requires CUDA", allow_module_level=True) from xorl.distributed.parallel_state import init_parallel_state # noqa: E402 -from xorl.ops.linear_attention.flashqla_cp import flashqla_chunk_gated_delta_rule_cp # noqa: E402 +from xorl.ops._vendored.flashqla_cp import flashqla_chunk_gated_delta_rule_cp # noqa: E402 from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context # noqa: E402 from xorl.ops.linear_attention.ops.gated_delta_rule import chunk_gated_delta_rule # noqa: E402 from xorl.utils.device import get_nccl_backend # noqa: E402 diff --git a/tests/ops/test_flashqla_contract_pin.py b/tests/ops/test_flashqla_contract_pin.py index bc769369..b8aac332 100644 --- a/tests/ops/test_flashqla_contract_pin.py +++ b/tests/ops/test_flashqla_contract_pin.py @@ -40,7 +40,7 @@ def _flashqla_chunk_or_skip(): if "prefer_instruction" not in inspect.signature(_tl.copy).parameters: pytest.skip("tilelang lacks prefer_instruction (PR #2303); FlashQLA TMA path unavailable") - from xorl.ops.linear_attention.flashqla import chunk_gated_delta_rule as flashqla_chunk # noqa: PLC0415 + from xorl.ops._vendored.flashqla import chunk_gated_delta_rule as flashqla_chunk # noqa: PLC0415 return flashqla_chunk @@ -150,7 +150,7 @@ def fake_flashqla_chunk(**kwargs): def test_contract_lane_pins_autocp_off(monkeypatch): """The armed lane must skip intra_card_cp_preprocess; arg/env must restore it.""" fn = _flashqla_chunk_or_skip() - import xorl.ops.linear_attention.flashqla.ops.gated_delta_rule.chunk as chunk_mod # noqa: PLC0415 + import xorl.ops._vendored.flashqla.ops.gated_delta_rule.chunk as chunk_mod # noqa: PLC0415 calls = [] real = chunk_mod.intra_card_cp_preprocess diff --git a/tests/ops/test_flashqla_gdn.py b/tests/ops/test_flashqla_gdn.py index e59eb780..d4ef2909 100644 --- a/tests/ops/test_flashqla_gdn.py +++ b/tests/ops/test_flashqla_gdn.py @@ -22,7 +22,7 @@ def _flashqla_chunk_or_skip(): if "prefer_instruction" not in inspect.signature(tl.copy).parameters: pytest.skip("tilelang lacks the required prefer_instruction support") - from xorl.ops.linear_attention.flashqla import ( # noqa: PLC0415 + from xorl.ops._vendored.flashqla import ( # noqa: PLC0415 chunk_gated_delta_rule, ) diff --git a/tests/ops/test_quack_process_safety.py b/tests/ops/test_quack_process_safety.py index 95cde304..0175287c 100644 --- a/tests/ops/test_quack_process_safety.py +++ b/tests/ops/test_quack_process_safety.py @@ -7,8 +7,8 @@ import pytest -from xorl.ops.quack import _worker_protocol as worker_protocol -from xorl.ops.quack import cache_utils, cute_dsl_ptxas +from xorl.ops._vendored.quack import _worker_protocol as worker_protocol +from xorl.ops._vendored.quack import cache_utils, cute_dsl_ptxas def test_quack_process_and_cache_safety_policy(tmp_path, monkeypatch): diff --git a/tests/server/test_security.py b/tests/server/test_security.py index f306a174..52be96a3 100644 --- a/tests/server/test_security.py +++ b/tests/server/test_security.py @@ -6,7 +6,7 @@ import pytest import torch -from xorl.ops.quack._worker_protocol import recv_message, send_message +from xorl.ops._vendored.quack._worker_protocol import recv_message, send_message from xorl.server.security import ( build_http_endpoint_url, resolve_diagnostic_input, @@ -107,12 +107,12 @@ def _assert_diagnostic_input_requires_configured_root_and_regular_private_file(t def test_compile_worker_security_and_protocol_policy(): - from xorl.ops.quack._compile_worker import _resolve_compile_function + from xorl.ops._vendored.quack._compile_worker import _resolve_compile_function with pytest.raises(ValueError, match="Quack module"): _resolve_compile_function("os", "system") with pytest.raises(ValueError, match="safe qualified name"): - _resolve_compile_function("xorl.ops.quack.autotuner", "__builtins__.eval") + _resolve_compile_function("xorl.ops._vendored.quack.autotuner", "__builtins__.eval") _assert_compile_worker_protocol_roundtrips_safe_types_and_rejects_oversized_header() From ecc8cd78ab99f3502c6759b960154dcc6c016bb4 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 03:39:26 +0000 Subject: [PATCH 02/14] Move training objectives to xorl.objectives (#78 phase 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RL and supervised objective functions (policy_loss, grpo_loss, cispo_loss, importance_sampling_loss, opd_loss, opd_streaming_kl, causallm_loss, reducers, loss_output) and the loss registry move out of xorl.ops.loss into a new xorl.objectives package; ops/loss keeps only the cross-entropy/selected-logprob kernel stack. In-repo imports are rewritten; the old module paths remain as compat stubs, and xorl.ops.loss / xorl.ops re-export the objective API lazily (PEP 562) for one deprecation cycle — lazily because the objectives import the kernels, so an eager re-export is an import cycle. --- .../opd/vocab_parallel_kl_gathered.py | 6 +- src/xorl/models/module_utils.py | 2 +- src/xorl/objectives/__init__.py | 64 + src/xorl/objectives/causallm_loss.py | 864 +++++++++++ src/xorl/objectives/cispo_loss.py | 164 ++ src/xorl/objectives/grpo_loss.py | 262 ++++ .../objectives/importance_sampling_loss.py | 216 +++ src/xorl/objectives/loss_output.py | 21 + src/xorl/objectives/opd_loss.py | 1324 ++++++++++++++++ src/xorl/objectives/opd_streaming_kl.py | 820 ++++++++++ src/xorl/objectives/policy_loss.py | 394 +++++ src/xorl/objectives/reducers.py | 71 + src/xorl/ops/README.md | 8 +- src/xorl/ops/__init__.py | 15 +- src/xorl/ops/loss/__init__.py | 100 +- src/xorl/ops/loss/causallm_loss.py | 865 +---------- src/xorl/ops/loss/cispo_loss.py | 165 +- src/xorl/ops/loss/grpo_loss.py | 263 +--- src/xorl/ops/loss/importance_sampling_loss.py | 217 +-- src/xorl/ops/loss/loss_output.py | 22 +- src/xorl/ops/loss/opd_loss.py | 1325 +---------------- src/xorl/ops/loss/opd_streaming_kl.py | 821 +--------- src/xorl/ops/loss/policy_loss.py | 395 +---- src/xorl/ops/loss/reducers.py | 72 +- src/xorl/trainers/training_utils.py | 2 +- tests/distributed/test_bi_fused_lm_head_tp.py | 4 +- tests/distributed/test_lm_head_tp_fsdp_e2e.py | 6 +- tests/distributed/test_pp_byte_alignment.py | 2 +- .../test_ulysses_byte_alignment.py | 2 +- .../models/test_bi_fused_lm_head_tp_wiring.py | 6 +- ...st_glm52_exact_lm_head_loss_integration.py | 2 +- tests/ops/loss/test_causallm_z_loss.py | 2 +- tests/ops/loss/test_fp8_lm_head_ce.py | 4 +- tests/ops/loss/test_fused_linear_logprob.py | 8 +- tests/ops/loss/test_opd_loss.py | 2 +- tests/ops/loss/test_opd_verl_parity.py | 8 +- tests/ops/loss/test_streaming_forward_kl.py | 4 +- tests/ops/test_bi_fused_lm_head.py | 2 +- tests/ops/test_exact_sampling_transforms.py | 8 +- .../server/runner/test_model_runner_drgrpo.py | 2 +- tests/server/runner/test_opd_runner.py | 2 +- 41 files changed, 4309 insertions(+), 4233 deletions(-) create mode 100644 src/xorl/objectives/__init__.py create mode 100644 src/xorl/objectives/causallm_loss.py create mode 100644 src/xorl/objectives/cispo_loss.py create mode 100644 src/xorl/objectives/grpo_loss.py create mode 100644 src/xorl/objectives/importance_sampling_loss.py create mode 100644 src/xorl/objectives/loss_output.py create mode 100644 src/xorl/objectives/opd_loss.py create mode 100644 src/xorl/objectives/opd_streaming_kl.py create mode 100644 src/xorl/objectives/policy_loss.py create mode 100644 src/xorl/objectives/reducers.py diff --git a/certification/opd/vocab_parallel_kl_gathered.py b/certification/opd/vocab_parallel_kl_gathered.py index fb18d648..de1c3bf6 100755 --- a/certification/opd/vocab_parallel_kl_gathered.py +++ b/certification/opd/vocab_parallel_kl_gathered.py @@ -24,9 +24,9 @@ from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import Shard, distribute_tensor -from xorl.ops.loss.opd_loss import opd_vocab_parallel_loss_function -from xorl.ops.loss.opd_streaming_kl import streaming_reverse_kl_function -from xorl.ops.loss.reducers import TokenPartial +from xorl.objectives.opd_loss import opd_vocab_parallel_loss_function +from xorl.objectives.opd_streaming_kl import streaming_reverse_kl_function +from xorl.objectives.reducers import TokenPartial from xorl.ops.loss.vocab_parallel_reverse_kl import vocab_parallel_reverse_kl_gathered from xorl.server.runner.model_runner import ModelRunner diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index f10918cb..c8218f2b 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -36,8 +36,8 @@ parse_expert_full_key, parse_expert_key, ) +from xorl.objectives.reducers import TokenPartial from xorl.ops.loss import fsdp_sharded_causallm_loss_function, get_loss_function -from xorl.ops.loss.reducers import TokenPartial from xorl.utils import logging from xorl.utils.device import get_device_id, get_device_type, synchronize from xorl.utils.helper import empty_cache, get_dtype_size diff --git a/src/xorl/objectives/__init__.py b/src/xorl/objectives/__init__.py new file mode 100644 index 00000000..33f0ed3c --- /dev/null +++ b/src/xorl/objectives/__init__.py @@ -0,0 +1,64 @@ +"""Training objectives (issue #78 phase 2). + +The RL and supervised objective functions, their reducers, and the loss +registry. These consume the cross-entropy/selected-logprob KERNELS in +:mod:`xorl.ops.loss`; the kernels stay there — this package is the +trainer-facing API. +""" + +from typing import Callable, Dict + +from xorl.objectives.causallm_loss import causallm_loss_function, fsdp_sharded_causallm_loss_function +from xorl.objectives.cispo_loss import cispo_loss_function +from xorl.objectives.grpo_loss import drgrpo_loss_function +from xorl.objectives.importance_sampling_loss import importance_sampling_loss_function +from xorl.objectives.loss_output import LossOutput +from xorl.objectives.opd_loss import OPDLossMetrics, opd_loss_function, opd_vocab_parallel_loss_function +from xorl.objectives.policy_loss import policy_loss_function +from xorl.objectives.reducers import Reducer, SequencePartial, TokenPartial + + +# --------------------------------------------------------------------------- +# Loss function registry +# --------------------------------------------------------------------------- +LOSS_REGISTRY: Dict[str, Callable] = { + "causallm_loss": causallm_loss_function, + "cross_entropy": causallm_loss_function, # alias + "importance_sampling": importance_sampling_loss_function, + "cispo": cispo_loss_function, + "policy_loss": policy_loss_function, + "drgrpo": drgrpo_loss_function, + "opd_loss": opd_loss_function, +} + + +def get_loss_function(name: str) -> Callable: + """Look up a loss function by name.""" + if name not in LOSS_REGISTRY: + raise ValueError(f"Unknown loss function: {name}. Available: {list(LOSS_REGISTRY.keys())}") + return LOSS_REGISTRY[name] + + +def register_loss_function(name: str, fn: Callable) -> None: + """Register a custom loss function.""" + LOSS_REGISTRY[name] = fn + + +__all__ = [ + "LOSS_REGISTRY", + "LossOutput", + "OPDLossMetrics", + "Reducer", + "SequencePartial", + "TokenPartial", + "causallm_loss_function", + "cispo_loss_function", + "drgrpo_loss_function", + "fsdp_sharded_causallm_loss_function", + "get_loss_function", + "importance_sampling_loss_function", + "opd_loss_function", + "opd_vocab_parallel_loss_function", + "policy_loss_function", + "register_loss_function", +] diff --git a/src/xorl/objectives/causallm_loss.py b/src/xorl/objectives/causallm_loss.py new file mode 100644 index 00000000..ff74d1cc --- /dev/null +++ b/src/xorl/objectives/causallm_loss.py @@ -0,0 +1,864 @@ +from __future__ import annotations + +import math + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from xorl.objectives.loss_output import LossOutput +from xorl.objectives.reducers import Reducer, TokenPartial +from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.loss.compiled_cross_entropy import ( + compiled_ce_and_lse_sq_function, + compiled_cross_entropy_function, +) +from xorl.ops.loss.per_token_ce import ( + LogprobProbability, + LogprobTemperature, + LogprobTopK, + compute_per_token_ce, + normalize_logprob_temperature, + resolve_bi_fused_lm_head_tp_groups, +) +from xorl.ops.loss.vocab_parallel_cross_entropy import ( + _backward_kernel as _vocab_parallel_ce_backward_kernel, +) +from xorl.ops.loss.vocab_parallel_cross_entropy import ( + _forward_kernel as _vocab_parallel_ce_forward_kernel, +) +from xorl.ops.loss.vocab_parallel_cross_entropy import ( + vocab_parallel_cross_entropy, + vocab_parallel_cross_entropy_with_lm_head, +) + + +_MODULE_LM_HEAD_MIN_CHUNK_ROWS = 128 + + +def _all_gather_cat_same_shape(x: torch.Tensor, *, dim: int, group: dist.ProcessGroup) -> torch.Tensor: + world_size = dist.get_world_size(group) + gathered = [torch.empty_like(x) for _ in range(world_size)] + dist.all_gather(gathered, x.contiguous(), group=group) + return torch.cat(gathered, dim=dim) + + +class _FSDPShardedCausalLMLoss(torch.autograd.Function): + """Sequence-streaming vocab-parallel CE for FSDP-sharded lm_head. + + The lm_head shard group is also the sequence-parallel group for the GLM + 128K configuration, so each rank owns different sequence tokens and a + different vocab shard. The usual vocab-parallel CE needs all ranks to see + the same token batch; this Function gathers one small sequence chunk at a + time and avoids saving those gathered chunks for backward. + """ + + @staticmethod + def forward( + ctx, + hidden_states: torch.Tensor, + local_weight: torch.Tensor, + labels: torch.Tensor, + global_valid_tokens: torch.Tensor, + sequence_group: dist.ProcessGroup, + vocab_group: dist.ProcessGroup, + num_chunks: int, + ignore_index: int, + loss_reduce_group: "dist.ProcessGroup | None" = None, + loss_reduce_divisor: float = 1.0, + ) -> torch.Tensor: + if hidden_states.dim() != 3: + raise ValueError(f"Expected hidden_states to have shape [B, S, H], got {tuple(hidden_states.shape)}") + if labels.shape != hidden_states.shape[:2]: + raise ValueError(f"Expected labels shape {tuple(hidden_states.shape[:2])}, got {tuple(labels.shape)}") + + ctx.save_for_backward(hidden_states, local_weight, labels, global_valid_tokens) + ctx.sequence_group = sequence_group + ctx.vocab_group = vocab_group + ctx.num_chunks = num_chunks + ctx.ignore_index = ignore_index + + vocab_rank = dist.get_rank(vocab_group) + local_vocab_size = local_weight.shape[0] + vocab_offset = vocab_rank * local_vocab_size + local_seq_len = hidden_states.shape[1] + chunk_size = max(1, math.ceil(local_seq_len / num_chunks)) + denom = global_valid_tokens.clamp(min=1.0) + loss = hidden_states.new_zeros((), dtype=torch.float32) + + for start in range(0, local_seq_len, chunk_size): + end = min(start + chunk_size, local_seq_len) + gathered_hidden = _all_gather_cat_same_shape(hidden_states[:, start:end, :], dim=1, group=sequence_group) + gathered_labels = _all_gather_cat_same_shape(labels[:, start:end], dim=1, group=sequence_group) + hidden_flat = gathered_hidden.reshape(-1, gathered_hidden.shape[-1]) + labels_flat = gathered_labels.reshape(-1) + per_token_ce, _, _, _, _, valid_mask = _vocab_parallel_ce_forward_kernel( + hidden_flat, + local_weight, + labels_flat, + vocab_group, + vocab_offset, + local_vocab_size, + ignore_index, + ) + loss = loss + (per_token_ce * valid_mask.float()).sum() / denom + del gathered_hidden, gathered_labels, hidden_flat, labels_flat, per_token_ce, valid_mask + + # lm-head-TP: each replica group computes the CE over its own sequence + # shard and over the full (TP-split) vocab, so the per-replica losses must + # be summed across replicas. The divisor removes the within-TP-group + # duplication (every TP rank computed the same per-replica loss). Backward + # intentionally stays unscaled: each rank returns its local weight-shard + # gradient, which the caller combines with an all-reduce over the replica + # group to reconstruct the full gradient. + if loss_reduce_group is not None: + dist.all_reduce(loss, op=dist.ReduceOp.SUM, group=loss_reduce_group) + if loss_reduce_divisor != 1.0: + loss = loss / loss_reduce_divisor + + return loss + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + hidden_states, local_weight, labels, global_valid_tokens = ctx.saved_tensors + sequence_group = ctx.sequence_group + vocab_group = ctx.vocab_group + vocab_rank = dist.get_rank(vocab_group) + local_vocab_size = local_weight.shape[0] + vocab_offset = vocab_rank * local_vocab_size + local_seq_len = hidden_states.shape[1] + chunk_size = max(1, math.ceil(local_seq_len / ctx.num_chunks)) + denom = global_valid_tokens.clamp(min=1.0) + + grad_hidden = torch.zeros_like(hidden_states) if ctx.needs_input_grad[0] else None + grad_weight = torch.zeros_like(local_weight) if ctx.needs_input_grad[1] else None + sp_world = dist.get_world_size(sequence_group) + sp_rank = dist.get_rank(sequence_group) + + for start in range(0, local_seq_len, chunk_size): + end = min(start + chunk_size, local_seq_len) + gathered_hidden = _all_gather_cat_same_shape(hidden_states[:, start:end, :], dim=1, group=sequence_group) + gathered_labels = _all_gather_cat_same_shape(labels[:, start:end], dim=1, group=sequence_group) + hidden_flat = gathered_hidden.reshape(-1, gathered_hidden.shape[-1]) + labels_flat = gathered_labels.reshape(-1) + _, global_max, global_sumexp, target_in_range, safe_local_target, valid_mask = ( + _vocab_parallel_ce_forward_kernel( + hidden_flat, + local_weight, + labels_flat, + vocab_group, + vocab_offset, + local_vocab_size, + ctx.ignore_index, + ) + ) + per_token_grad = grad_output.to(hidden_flat.dtype) * valid_mask.to(hidden_flat.dtype) / denom + chunk_grad_hidden, chunk_grad_weight = _vocab_parallel_ce_backward_kernel( + per_token_grad, + hidden_flat, + local_weight, + global_max, + global_sumexp, + target_in_range, + safe_local_target, + valid_mask, + vocab_group, + ctx.needs_input_grad[1], + ) + if grad_hidden is not None: + chunk_grad_hidden = chunk_grad_hidden.view_as(gathered_hidden) + local_grad = chunk_grad_hidden.chunk(sp_world, dim=1)[sp_rank].contiguous() + grad_hidden[:, start:end, :] = local_grad.to(grad_hidden.dtype) + if grad_weight is not None and chunk_grad_weight is not None: + grad_weight.add_(chunk_grad_weight.to(grad_weight.dtype)) + del ( + gathered_hidden, + gathered_labels, + hidden_flat, + labels_flat, + global_max, + global_sumexp, + target_in_range, + safe_local_target, + valid_mask, + per_token_grad, + chunk_grad_hidden, + chunk_grad_weight, + ) + + # Grads for: hidden_states, local_weight, then None for labels, + # global_valid_tokens, sequence_group, vocab_group, num_chunks, + # ignore_index, loss_reduce_group, loss_reduce_divisor. + return grad_hidden, grad_weight, None, None, None, None, None, None, None, None + + +def fsdp_sharded_causallm_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + sp_group: dist.ProcessGroup, + fsdp_group: dist.ProcessGroup, + num_chunks: int, + ignore_index: int = -100, + lm_head_fp32: bool = False, + global_valid_tokens: torch.Tensor | None = None, + sequence_group: "dist.ProcessGroup | None" = None, + vocab_group: "dist.ProcessGroup | None" = None, + loss_reduce_group: "dist.ProcessGroup | None" = None, + loss_reduce_divisor: float = 1.0, +) -> "LossOutput": + # sequence_group/vocab_group default to sp_group/fsdp_group (the FSDP-sharded + # lm_head case where the sequence-parallel group is also the vocab shard + # group). lm-head-only TP passes a dedicated lm_head_tp_group for both and a + # loss_reduce_group (+ divisor) to sum the per-replica losses. + if sequence_group is None: + sequence_group = sp_group + if vocab_group is None: + vocab_group = fsdp_group + if lm_head_fp32: + hidden_states = hidden_states.float() + + local_weight = weight.to_local() if hasattr(weight, "to_local") else weight + if local_weight.dtype != hidden_states.dtype: + if lm_head_fp32: + local_weight = local_weight.float() + else: + local_weight = local_weight.to(hidden_states.dtype) + + if global_valid_tokens is None: + global_valid_tokens = (labels != ignore_index).sum().to(hidden_states.device, dtype=torch.float32) + dist.all_reduce(global_valid_tokens, op=dist.ReduceOp.SUM, group=fsdp_group) + else: + global_valid_tokens = global_valid_tokens.detach().to(hidden_states.device, dtype=torch.float32) + loss = _FSDPShardedCausalLMLoss.apply( + hidden_states, + local_weight, + labels, + global_valid_tokens, + sequence_group, + vocab_group, + int(num_chunks), + int(ignore_index), + loss_reduce_group, + float(loss_reduce_divisor), + ) + return LossOutput(loss=loss) + + +def _chunked_lm_head_cross_entropy( + hidden_states_flat: torch.Tensor, + labels_flat: torch.Tensor, + *, + lm_head: torch.nn.Module, + ignore_index: int, + num_chunks: int, + z_loss_enabled: bool, + valid_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Compute CE by calling the lm_head module in chunks. + + FP8 training wraps ``lm_head`` with ``FP8Linear``. The compiled CE helpers + operate on ``lm_head.weight`` directly, which bypasses that module, so FP8 + training uses this path to keep the output head matmul on FP8 compute. + """ + + if hidden_states_flat.shape[0] == 0: + empty = hidden_states_flat.new_empty((0,), dtype=torch.float32) + return empty, empty if z_loss_enabled else None + + chunk_count = max(1, int(num_chunks)) + chunk_size = max(_MODULE_LM_HEAD_MIN_CHUNK_ROWS, math.ceil(hidden_states_flat.shape[0] / chunk_count)) + ce_chunks: list[torch.Tensor] = [] + lse_sq_chunks: list[torch.Tensor] = [] + for start in range(0, hidden_states_flat.shape[0], chunk_size): + end = min(start + chunk_size, hidden_states_flat.shape[0]) + logits = lm_head(hidden_states_flat[start:end]).float() + labels = labels_flat[start:end] + ce_chunks.append(F.cross_entropy(logits, labels, reduction="none", ignore_index=ignore_index)) + if z_loss_enabled: + lse = torch.logsumexp(logits, dim=-1) + lse_sq_chunks.append((lse * lse) * valid_mask[start:end].to(lse.dtype)) + + per_token_ce = torch.cat(ce_chunks, dim=0) + per_token_lse_sq = torch.cat(lse_sq_chunks, dim=0) if z_loss_enabled else None + return per_token_ce, per_token_lse_sq + + +def _ceil_to_multiple(value: int, multiple: int) -> int: + return ((value + multiple - 1) // multiple) * multiple + + +def _chunk_size_from_num_chunks(num_tokens: int, num_chunks: int) -> int: + if num_chunks <= 0: + return _ceil_to_multiple(num_tokens, 8) + return _ceil_to_multiple((num_tokens + num_chunks - 1) // num_chunks, 8) + + +def _quack_linear_cross_entropy_loss( + hidden_states_flat: torch.Tensor, + weight: torch.Tensor, + labels_flat: torch.Tensor, + ignore_index: int, + num_chunks: int, + loss_reducer: TokenPartial, +) -> torch.Tensor: + if not hidden_states_flat.is_cuda: + raise ValueError("ce_mode='quack_linear' requires CUDA tensors") + if hidden_states_flat.shape[-1] % 8 != 0 or weight.shape[0] % 8 != 0: + raise ValueError("ce_mode='quack_linear' requires hidden and vocab dimensions to be divisible by 8") + + from xorl.ops._vendored.quack.linear_cross_entropy import chunked_linear_cross_entropy # noqa: PLC0415 + + hidden_states_flat, labels_flat = _pad_quack_linear_rows( + hidden_states_flat, + labels_flat, + ignore_index=ignore_index, + ) + valid_count = (labels_flat != ignore_index).sum() + if valid_count.item() == 0: + return (hidden_states_flat.sum() + weight.sum()) * 0.0 + + chunk_size = _chunk_size_from_num_chunks(hidden_states_flat.shape[0], num_chunks) + loss_sum = chunked_linear_cross_entropy( + hidden_states_flat, + weight, + labels_flat, + chunk_size=chunk_size, + ignore_index=ignore_index, + reduction="sum", + ) + scale = loss_reducer.scale.to(device=loss_sum.device, dtype=loss_sum.dtype) + return loss_sum / scale.clamp(min=1.0) + + +def _pad_quack_linear_rows( + hidden_states_flat: torch.Tensor, + labels_flat: torch.Tensor, + *, + ignore_index: int, + multiple: int = 8, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pad rows for Quack CE kernels without changing scalar CE semantics.""" + + remainder = hidden_states_flat.shape[0] % multiple + if remainder == 0: + return hidden_states_flat, labels_flat + pad_rows = multiple - remainder + hidden_pad = hidden_states_flat.new_zeros((pad_rows, hidden_states_flat.shape[-1])) + label_pad = labels_flat.new_full((pad_rows,), ignore_index) + return torch.cat((hidden_states_flat, hidden_pad), dim=0), torch.cat((labels_flat, label_pad), dim=0) + + +def _fused_quack_per_token_ce( + hidden_states_flat: torch.Tensor, + weight: torch.Tensor, + labels_flat: torch.Tensor, + ignore_index: int, + num_chunks: int, + tp_group, + lm_head_fp32: bool, +) -> torch.Tensor: + """Per-token CE ``[N]`` via the fused chunked cuBLAS + CuTeDSL path. + + ``fused_selected_logprob_ce`` keeps the logits tile bounded to + ``[chunk, V_local]`` and never materializes the full ``[N, V]`` logits. + Mirrors the dispatch in ``ops.loss.per_token_ce.compute_per_token_ce`` so + ``ce_mode='fused_quack'`` works from this entry point too — without it, + fused_quack fell through to the eager full-logits path and OOM'd at large + vocab / long context. + """ + if not hidden_states_flat.is_cuda: + raise ValueError("ce_mode='fused_quack' requires CUDA tensors") + + from xorl.ops.loss.fused_linear_logprob import fused_selected_logprob_ce # noqa: PLC0415 + + local_weight = weight.to_local() if hasattr(weight, "to_local") else weight + hidden = hidden_states_flat + if lm_head_fp32: + hidden = hidden.float() + local_weight = local_weight.float() + chunk_size = _chunk_size_from_num_chunks(hidden.shape[0], num_chunks) + return fused_selected_logprob_ce( + hidden, + local_weight, + labels_flat, + tp_group=tp_group, + ignore_index=ignore_index, + chunk_size=chunk_size, + ) + + +def _bi_fused_per_token_ce_checked( + hidden_states_flat: torch.Tensor, + weight: torch.Tensor, + labels_flat: torch.Tensor, + ignore_index: int, + lm_head_fp32: bool, + z_loss_enabled: bool, +) -> torch.Tensor: + """Guarded entry for ``ce_mode='bi_fused'`` (the batch-invariant lm-head + contract). The contract IS the fp32-class lm-head computation, so it + requires ``lm_head_fp32`` semantics without materializing the fp32 weight.""" + from xorl.ops.loss.bi_fused_lm_head import bi_fused_per_token_ce + + if z_loss_enabled: + raise NotImplementedError("ce_mode='bi_fused' does not support softmax_auxiliary_loss") + if not lm_head_fp32: + raise NotImplementedError( + "ce_mode='bi_fused' implements the fp32-class lm-head contract; set lm_head_fp32: true" + ) + local_weight = weight.to_local() if hasattr(weight, "to_local") else weight + return bi_fused_per_token_ce(hidden_states_flat, local_weight, labels_flat, ignore_index) + + +def _quack_linear_per_token_cross_entropy( + hidden_states_flat: torch.Tensor, + weight: torch.Tensor, + labels_flat: torch.Tensor, + ignore_index: int, + num_chunks: int, + lm_head_fp32: bool, +) -> torch.Tensor: + """Per-token return path for ``ce_mode='quack_linear'``. + + The scalar training path keeps using Quack's chunked linear CE reduction. + ``return_per_token=True`` callers need one CE value per input row, so route + those through the existing fused selected-logprob kernel and still avoid + full-logit materialization. + """ + + if not hidden_states_flat.is_cuda: + raise ValueError("ce_mode='quack_linear' requires CUDA tensors") + if hidden_states_flat.shape[-1] % 8 != 0 or weight.shape[0] % 8 != 0: + raise ValueError("ce_mode='quack_linear' requires hidden and vocab dimensions to be divisible by 8") + return _fused_quack_per_token_ce( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + num_chunks, + tp_group=None, + lm_head_fp32=lm_head_fp32, + ) + + +def causallm_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + return_per_token: bool = False, + ce_mode: str = "compiled", + num_chunks: int = 8, + tp_group=None, + use_compile: bool = False, + lm_head_fp32: bool = False, + loss_reducer: Reducer | None = None, + z_loss_coef: float = 0.0, + lm_head: torch.nn.Module | None = None, + logprob_temperature: LogprobTemperature = 1.0, + logprob_top_k: LogprobTopK = TOP_K_ALL, + logprob_top_p: LogprobProbability = 1.0, + logprob_min_p: LogprobProbability = 0.0, +) -> "LossOutput": + """ + Compute causal language modeling loss. + + Supports multiple computation modes: + - "compiled": RECOMMENDED. torch.compile (1.6x speed, 16% memory) + - "eager": Simple F.cross_entropy baseline (may OOM at 32K) + + Args: + hidden_states: Model hidden states, shape (batch, seq_len, hidden_dim) + weight: LM head weight matrix, shape (vocab_size, hidden_dim). + With TP, this is the local shard [vocab_size/tp, hidden_dim]. + labels: Target labels, shape (batch, seq_len). Labels are assumed to be + already next-token aligned (labels[i] is the target for hidden_states[i]). + ignore_index: Index to ignore in loss computation (default: -100) + return_per_token: If True, return per-token logprobs and losses (default: False) + ce_mode: Cross-entropy mode - "compiled" (default) or "eager" + num_chunks: Number of chunks for compiled mode (default: 8). + tp_group: TP process group for vocab-parallel cross-entropy (default: None). + loss_reducer: Optional ``(values, mask) -> scalar``. When supplied, the + returned loss is a partial share under the reducer's denominator + (sum across micro-batches + all-reduce across ranks recovers the + globally-correct loss). When None, falls back to a local token mean. + Z-loss (when enabled) is reduced through the same reducer so the + two terms compose consistently. + z_loss_coef: If > 0, add the Z-loss auxiliary term used in OLMo / + PaLM-style training: + z_loss = coef * sum(logsumexp(logits)^2 * mask) / num_valid_tokens + where ``mask = labels != ignore_index``. Equivalent to OLMo's + ``cross_entropy_loss(..., reduction="sum")`` path divided by + ``batch_size_in_tokens``. Encourages log(Z) to stay near zero, + stabilizing training at large vocab / high LR. Not supported + in the TP path. + logprob_temperature: Temperature for selected-token logprobs. ``1.0`` + returns raw model logprobs; a rollout temperature such as + ``0.7`` returns behavior-policy logprobs using + ``log_softmax(logits / temperature)``. Exact LM heads also + accept contiguous FP32 temperatures aligned with labels. + + Returns: + LossOutput with loss, and optionally per_token_logprobs/per_token_loss. + When ``z_loss_coef > 0``, ``LossOutput.metrics`` contains + ``{"ce_loss": , "z_loss": }``. + """ + # Store original shape before flattening for per-token outputs + original_shape = labels.shape + + # Flatten the labels and hidden_states + labels_flat = labels.view(-1) + hidden_states_flat = hidden_states.view(-1, hidden_states.size(-1)) + valid_mask = labels_flat != ignore_index + bi_fused_tp_groups = resolve_bi_fused_lm_head_tp_groups(ce_mode, tp_group, lm_head) + has_explicit_loss_reducer = loss_reducer is not None + + if loss_reducer is None: + scale = valid_mask.sum().float() + if bi_fused_tp_groups is not None: + dedicated_group, replica_group = bi_fused_tp_groups + dist.all_reduce(scale, op=dist.ReduceOp.SUM, group=dedicated_group) + if replica_group is not None: + dist.all_reduce(scale, op=dist.ReduceOp.SUM, group=replica_group) + loss_reducer = TokenPartial(scale=scale) + + mask_flat = valid_mask.float() + if isinstance(logprob_temperature, torch.Tensor): + if not logprob_temperature.is_contiguous(): + raise ValueError("per-row logprob_temperature must be contiguous") + if tuple(logprob_temperature.shape) not in (tuple(labels.shape), (labels_flat.shape[0],)): + raise ValueError( + "per-row logprob_temperature must match labels or flattened labels, got " + f"{tuple(logprob_temperature.shape)} for labels {tuple(labels.shape)}" + ) + logprob_temperature = logprob_temperature.reshape(-1) + logprob_temperature = normalize_logprob_temperature( + logprob_temperature, + rows=labels_flat.shape[0], + device=hidden_states.device, + ) + + def _flatten_sampling_metadata(value, name: str): + if not isinstance(value, torch.Tensor): + return value + if not value.is_contiguous(): + raise ValueError(f"per-row {name} must be contiguous") + if tuple(value.shape) not in (tuple(labels.shape), (labels_flat.shape[0],)): + raise ValueError(f"per-row {name} must match labels or flattened labels") + return value.reshape(-1) + + logprob_top_k = _flatten_sampling_metadata(logprob_top_k, "logprob_top_ks") + logprob_top_p = _flatten_sampling_metadata(logprob_top_p, "logprob_top_ps") + logprob_min_p = _flatten_sampling_metadata(logprob_min_p, "logprob_min_ps") + has_temperature_transform = isinstance(logprob_temperature, torch.Tensor) or logprob_temperature != 1.0 + has_sampling_filter = ( + isinstance(logprob_top_k, torch.Tensor) + or isinstance(logprob_top_p, torch.Tensor) + or isinstance(logprob_min_p, torch.Tensor) + or int(logprob_top_k) < TOP_K_ALL + or float(logprob_top_p) != 1.0 + or float(logprob_min_p) != 0.0 + ) + exact_lm_head = bool( + lm_head is not None + and (getattr(lm_head, "_glm52_exact_tp16_lm_head", False) or getattr(lm_head, "_dsv4_exact_tp8_lm_head", False)) + ) + if ce_mode == "bi_fused": + if tp_group is not None and not exact_lm_head and bi_fused_tp_groups is None: + raise NotImplementedError( + "ce_mode='bi_fused' supports TP only through the dedicated vocabulary-sharded LM-head TP path" + ) + if lm_head is not None and not lm_head_fp32 and not exact_lm_head: + raise NotImplementedError("ce_mode='bi_fused' does not support FP8 lm_head modules") + if exact_lm_head: + if z_loss_coef > 0.0: + raise NotImplementedError("The exact GLM-5.2 active-LoRA lm head does not support Z-loss") + per_token_ce = compute_per_token_ce( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + use_compile=use_compile, + lm_head_fp32=lm_head_fp32, + lm_head=lm_head, + logprob_temperature=logprob_temperature, + logprob_top_k=logprob_top_k, + logprob_top_p=logprob_top_p, + logprob_min_p=logprob_min_p, + ) + loss = loss_reducer(per_token_ce, mask_flat) + if return_per_token: + return LossOutput( + loss=loss, + per_token_logprobs=-per_token_ce.detach().view(original_shape), + per_token_loss=per_token_ce.view(original_shape), + ) + return LossOutput(loss=loss) + if bi_fused_tp_groups is not None: + if z_loss_coef > 0.0: + raise NotImplementedError("ce_mode='bi_fused' does not support softmax_auxiliary_loss") + per_token_ce = compute_per_token_ce( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + use_compile=use_compile, + lm_head_fp32=lm_head_fp32, + lm_head=lm_head, + logprob_temperature=logprob_temperature, + logprob_top_k=logprob_top_k, + logprob_top_p=logprob_top_p, + logprob_min_p=logprob_min_p, + ) + local_loss = loss_reducer(per_token_ce, mask_flat) + if has_explicit_loss_reducer: + if return_per_token: + return LossOutput( + loss=local_loss, + per_token_logprobs=-per_token_ce.detach().view(original_shape), + per_token_loss=per_token_ce.view(original_shape), + ) + return LossOutput(loss=local_loss) + + # Standalone/default-reducer calls historically return the full scalar + # on every rank. Explicit reducers instead promise a local partial, and + # their caller owns detached reporting aggregation. + global_loss = local_loss.detach().clone() + dedicated_group, replica_group = bi_fused_tp_groups + dist.all_reduce(global_loss, op=dist.ReduceOp.SUM, group=dedicated_group) + if replica_group is not None: + dist.all_reduce(global_loss, op=dist.ReduceOp.SUM, group=replica_group) + loss = local_loss + (global_loss - local_loss.detach()) + if return_per_token: + return LossOutput( + loss=loss, + per_token_logprobs=-per_token_ce.detach().view(original_shape), + per_token_loss=per_token_ce.view(original_shape), + ) + return LossOutput(loss=loss) + if has_temperature_transform or has_sampling_filter: + if z_loss_coef > 0.0: + raise NotImplementedError("logprob_temperature is not supported with softmax_auxiliary_loss") + per_token_ce = compute_per_token_ce( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + use_compile=use_compile, + lm_head_fp32=lm_head_fp32, + lm_head=lm_head, + logprob_temperature=logprob_temperature, + logprob_top_k=logprob_top_k, + logprob_top_p=logprob_top_p, + logprob_min_p=logprob_min_p, + ) + loss = loss_reducer(per_token_ce, mask_flat) + if return_per_token: + return LossOutput( + loss=loss, + per_token_logprobs=-per_token_ce.detach().view(original_shape), + per_token_loss=per_token_ce.view(original_shape), + ) + return LossOutput(loss=loss) + + # Vocab-parallel cross-entropy for tensor parallelism + if tp_group is not None: + if z_loss_coef > 0.0: + raise NotImplementedError( + "softmax_auxiliary_loss (Z-loss) is not yet supported with tensor parallelism. " + "Disable softmax_auxiliary_loss or run without TP." + ) + if lm_head is not None and not lm_head_fp32: + per_token_ce = vocab_parallel_cross_entropy_with_lm_head( + hidden_states_flat, + lm_head, + labels_flat, + tp_group, + ignore_index=ignore_index, + num_chunks=num_chunks, + use_compile=use_compile, + ) + else: + # lm_head_fp32 takes precedence over the FP8 lm_head module: compute + # the vocab-parallel CE in fp32 from the master weight (FP8 module + # bypassed). Extract local weight from DTensor if needed. + local_weight = weight.to_local() if hasattr(weight, "to_local") else weight + if lm_head_fp32: + hidden_states_flat = hidden_states_flat.float() + local_weight = local_weight.float() + elif local_weight.dtype != hidden_states_flat.dtype: + local_weight = local_weight.to(hidden_states_flat.dtype) + + per_token_ce = vocab_parallel_cross_entropy( + hidden_states_flat, + local_weight, + labels_flat, + tp_group, + ignore_index=ignore_index, + num_chunks=num_chunks, + use_compile=use_compile, + ) + + loss = loss_reducer(per_token_ce, mask_flat) + if return_per_token: + return LossOutput( + loss=loss, + per_token_logprobs=-per_token_ce.detach().view(original_shape), + per_token_loss=per_token_ce.view(original_shape), + ) + return LossOutput(loss=loss) + + z_loss_enabled = z_loss_coef > 0.0 + # lm_head_fp32 takes precedence over the FP8 lm_head module: an FP32 lm_head + # must not be FP8-quantized, so route to the fp32 weight-CE path below + # (compiled/eager honor lm_head_fp32) instead of _chunked_lm_head_cross_entropy + # (which calls FP8Linear.forward). The FP8 lm_head otherwise catastrophically + # mis-scores rare near-certain tokens (R1). + use_lm_head_module = lm_head is not None and not lm_head_fp32 + + if ce_mode == "quack_linear" and not return_per_token: + if z_loss_enabled: + raise NotImplementedError("ce_mode='quack_linear' does not support softmax_auxiliary_loss") + if lm_head_fp32: + raise NotImplementedError("ce_mode='quack_linear' does not support lm_head_fp32=True") + if not isinstance(loss_reducer, TokenPartial): + raise NotImplementedError("ce_mode='quack_linear' currently supports only TokenPartial loss reduction") + return LossOutput( + loss=_quack_linear_cross_entropy_loss( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + num_chunks, + loss_reducer, + ) + ) + + if return_per_token: + # Compute cross-entropy based on mode (and Z-loss when enabled). + per_token_lse_sq = None + if use_lm_head_module: + per_token_ce, per_token_lse_sq = _chunked_lm_head_cross_entropy( + hidden_states_flat, + labels_flat, + lm_head=lm_head, + ignore_index=ignore_index, + num_chunks=num_chunks, + z_loss_enabled=z_loss_enabled, + valid_mask=valid_mask, + ) + elif ce_mode == "compiled": + if z_loss_enabled: + per_token_ce, per_token_lse_sq = compiled_ce_and_lse_sq_function( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 + ) + else: + per_token_ce = compiled_cross_entropy_function( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 + ) + elif ce_mode == "fused_quack": + if z_loss_enabled: + raise NotImplementedError("ce_mode='fused_quack' does not support softmax_auxiliary_loss") + per_token_ce = _fused_quack_per_token_ce( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, tp_group, lm_head_fp32 + ) + elif ce_mode == "bi_fused": + per_token_ce = _bi_fused_per_token_ce_checked( + hidden_states_flat, weight, labels_flat, ignore_index, lm_head_fp32, z_loss_enabled + ) + elif ce_mode == "quack_linear": + if z_loss_enabled: + raise NotImplementedError("ce_mode='quack_linear' does not support softmax_auxiliary_loss") + per_token_ce = _quack_linear_per_token_cross_entropy( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + num_chunks, + lm_head_fp32, + ) + else: # eager mode + if lm_head_fp32: + logits_flat = (hidden_states_flat.float() @ weight.float().t()).float() + else: + logits_flat = (hidden_states_flat @ weight.t()).float() + per_token_ce = F.cross_entropy(logits_flat, labels_flat, reduction="none", ignore_index=ignore_index) + if z_loss_enabled: + lse = torch.logsumexp(logits_flat, dim=-1) + per_token_lse_sq = (lse * lse) * valid_mask.to(lse.dtype) + + ce_loss = loss_reducer(per_token_ce, mask_flat) + if z_loss_enabled: + z_loss = loss_reducer(per_token_lse_sq, mask_flat) + loss = ce_loss + z_loss_coef * z_loss + metrics = {"ce_loss": ce_loss.detach(), "z_loss": z_loss.detach()} + else: + loss = ce_loss + metrics = None + return LossOutput( + loss=loss, + per_token_logprobs=-per_token_ce.detach().view(original_shape), + per_token_loss=per_token_ce.view(original_shape), + metrics=metrics, + ) + else: + # Always use reduction="none" + manual mean to avoid NaN when all labels + # are ignore_index (reduction="mean" returns NaN for 0 valid elements). + # Keeping the autograd graph intact is critical for FSDP2: all ranks must + # trigger reduce-scatter for every parameter, including lm_head weight. + per_token_lse_sq = None + if use_lm_head_module: + per_token_ce, per_token_lse_sq = _chunked_lm_head_cross_entropy( + hidden_states_flat, + labels_flat, + lm_head=lm_head, + ignore_index=ignore_index, + num_chunks=num_chunks, + z_loss_enabled=z_loss_enabled, + valid_mask=valid_mask, + ) + elif ce_mode == "compiled": + if z_loss_enabled: + per_token_ce, per_token_lse_sq = compiled_ce_and_lse_sq_function( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 + ) + else: + per_token_ce = compiled_cross_entropy_function( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 + ) + elif ce_mode == "fused_quack": + if z_loss_enabled: + raise NotImplementedError("ce_mode='fused_quack' does not support softmax_auxiliary_loss") + per_token_ce = _fused_quack_per_token_ce( + hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, tp_group, lm_head_fp32 + ) + elif ce_mode == "bi_fused": + per_token_ce = _bi_fused_per_token_ce_checked( + hidden_states_flat, weight, labels_flat, ignore_index, lm_head_fp32, z_loss_enabled + ) + else: # eager mode + if lm_head_fp32: + logits_flat = (hidden_states_flat.float() @ weight.float().t()).float() + else: + logits_flat = (hidden_states_flat @ weight.t()).float() + per_token_ce = F.cross_entropy(logits_flat, labels_flat, reduction="none", ignore_index=ignore_index) + if z_loss_enabled: + lse = torch.logsumexp(logits_flat, dim=-1) + per_token_lse_sq = (lse * lse) * valid_mask.to(lse.dtype) + + ce_loss = loss_reducer(per_token_ce, mask_flat) + if z_loss_enabled: + z_loss = loss_reducer(per_token_lse_sq, mask_flat) + loss = ce_loss + z_loss_coef * z_loss + return LossOutput(loss=loss, metrics={"ce_loss": ce_loss.detach(), "z_loss": z_loss.detach()}) + return LossOutput(loss=ce_loss) diff --git a/src/xorl/objectives/cispo_loss.py b/src/xorl/objectives/cispo_loss.py new file mode 100644 index 00000000..acf49dfc --- /dev/null +++ b/src/xorl/objectives/cispo_loss.py @@ -0,0 +1,164 @@ +from __future__ import annotations + +from typing import Any, Dict, Optional + +import torch +import torch.distributed as dist + +from xorl.objectives.importance_sampling_loss import K3_DEBUG_THRESHOLDS +from xorl.objectives.loss_output import LossOutput +from xorl.objectives.reducers import Reducer, TokenPartial +from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.loss.per_token_ce import compute_per_token_ce + + +def cispo_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + old_logprobs: torch.Tensor, + advantages: torch.Tensor, + ignore_index: int = -100, + clip_low_threshold: float = 0.0, + clip_high_threshold: float = 4.0, + num_chunks: int = 8, + ce_mode: str = "compiled", + tp_group: Optional[dist.ProcessGroup] = None, + compute_kl_stats: bool = False, + lm_head_fp32: bool = False, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, + lm_head: Optional[torch.nn.Module] = None, + logprob_temperature: float = 1.0, + logprob_top_k: int | torch.Tensor = TOP_K_ALL, + logprob_top_p: float | torch.Tensor = 1.0, + logprob_min_p: float | torch.Tensor = 0.0, +) -> LossOutput: + """Compute Tinker-compatible CISPO. + + CISPO clips the importance ratio and uses it as a detached coefficient on + ``log p_theta``. Unlike PPO objective clipping, this retains a gradient for + every valid token:: + + ratio = exp(target_logprobs - sampling_logprobs) + clipped = clamp(ratio, clip_low_threshold, clip_high_threshold) + loss = -(clipped.detach() * target_logprobs * advantages).mean() + + The default absolute ratio bounds, ``[0, 4]``, match Tinker's one-sided + CISPO default and the MiniMax-M1 prescription of disabling the lower bound. + """ + if clip_low_threshold < 0.0: + raise ValueError("clip_low_threshold must be non-negative") + if clip_high_threshold < clip_low_threshold: + raise ValueError("clip_high_threshold must be >= clip_low_threshold") + + original_shape = labels.shape + hidden_size = hidden_states.size(-1) + + labels_flat = labels.reshape(-1) + hidden_states_flat = hidden_states.reshape(-1, hidden_size) + old_logprobs_flat = old_logprobs.reshape(-1) + advantages_flat = advantages.reshape(-1) + + valid_mask = labels_flat != ignore_index + valid_mask_f = valid_mask.float() + valid_count = valid_mask.sum() + + if loss_reducer is None: + loss_reducer = TokenPartial(scale=valid_count.float()) + if metric_reducer is None: + metric_reducer = TokenPartial(scale=valid_count.float()) + + per_token_ce = compute_per_token_ce( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + lm_head_fp32=lm_head_fp32, + lm_head=lm_head, + logprob_temperature=logprob_temperature, + logprob_top_k=logprob_top_k, + logprob_top_p=logprob_top_p, + logprob_min_p=logprob_min_p, + ) + + current_support = torch.isfinite(per_token_ce) + new_logprobs_flat = -per_token_ce.detach() + log_ratio = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) + ratio = torch.exp(log_ratio) + clipped_ratio = torch.clamp(ratio, clip_low_threshold, clip_high_threshold) + + coefficient = (clipped_ratio.detach() * advantages_flat).masked_fill(~valid_mask | ~current_support, 0.0) + safe_per_token_ce = torch.where(current_support, per_token_ce, torch.zeros_like(per_token_ce)) + per_token_loss_flat = coefficient * safe_per_token_ce + loss = loss_reducer(per_token_loss_flat, valid_mask_f) + + is_clipped = (clipped_ratio != ratio) & valid_mask + if valid_mask.any(): + ratio_min = ratio.masked_fill(~valid_mask, float("inf")).min() + ratio_max = ratio.masked_fill(~valid_mask, float("-inf")).max() + else: + ratio_min = ratio.new_tensor(float("inf")) + ratio_max = ratio.new_tensor(float("-inf")) + + metrics: Dict[str, Any] = { + "ratio_mean": metric_reducer(ratio, valid_mask_f).detach(), + "ratio_min": ratio_min.detach(), + "ratio_max": ratio_max.detach(), + "clip_fraction": metric_reducer(is_clipped.float(), valid_mask_f).detach(), + } + + metric_ops = {"ratio_min": "min", "ratio_max": "max"} + if compute_kl_stats: + with torch.no_grad(): + metric_log_ratio = torch.where( + current_support, + log_ratio, + torch.full_like(log_ratio, -20.0), + ) + per_token_k3 = ratio - metric_log_ratio - 1.0 + if valid_mask.any(): + k3_max = per_token_k3.masked_fill(~valid_mask, float("-inf")).max() + logratio_min = metric_log_ratio.masked_fill(~valid_mask, float("inf")).min() + logratio_max = metric_log_ratio.masked_fill(~valid_mask, float("-inf")).max() + abs_logratio_max = metric_log_ratio.abs().masked_fill(~valid_mask, float("-inf")).max() + else: + k3_max = per_token_k3.new_tensor(float("-inf")) + logratio_min = log_ratio.new_tensor(float("inf")) + logratio_max = log_ratio.new_tensor(float("-inf")) + abs_logratio_max = log_ratio.new_tensor(float("-inf")) + metrics["kl_sample_train_k3"] = metric_reducer(per_token_k3, valid_mask_f) + metrics["kl_k3_debug_mean"] = metric_reducer(per_token_k3, valid_mask_f) + metrics["kl_k3_debug_max"] = k3_max + metrics["kl_k3_debug_abs_logratio_mean"] = metric_reducer(metric_log_ratio.abs(), valid_mask_f) + metrics["kl_k3_debug_abs_logratio_max"] = abs_logratio_max + metrics["kl_k3_debug_logratio_mean"] = metric_reducer(metric_log_ratio, valid_mask_f) + metrics["kl_k3_debug_logratio_min"] = logratio_min + metrics["kl_k3_debug_logratio_max"] = logratio_max + metrics["kl_k3_debug_frac_logratio_positive"] = metric_reducer((metric_log_ratio > 0).float(), valid_mask_f) + for suffix, threshold in K3_DEBUG_THRESHOLDS: + metrics[f"kl_k3_debug_frac_gt_{suffix}"] = metric_reducer( + (per_token_k3 > threshold).float(), valid_mask_f + ) + metrics["entropy_sample"] = metric_reducer(-old_logprobs_flat, valid_mask_f) + metrics["valid_tokens"] = valid_count.item() + metrics["current_support_fraction"] = metric_reducer(current_support.float(), valid_mask_f) + metric_ops.update( + { + "kl_k3_debug_max": "max", + "kl_k3_debug_abs_logratio_max": "max", + "kl_k3_debug_logratio_min": "min", + "kl_k3_debug_logratio_max": "max", + } + ) + + return LossOutput( + loss=loss, + per_token_logprobs=new_logprobs_flat.view(original_shape), + per_token_loss=per_token_loss_flat.detach().view(original_shape), + metrics=metrics, + metric_ops=metric_ops, + ) diff --git a/src/xorl/objectives/grpo_loss.py b/src/xorl/objectives/grpo_loss.py new file mode 100644 index 00000000..9692cf4b --- /dev/null +++ b/src/xorl/objectives/grpo_loss.py @@ -0,0 +1,262 @@ +""" +DR-GRPO: "Done Right" GRPO Loss for RL Training. + +Reference: Liu et al., "Understanding R1-Zero-Like Training" (2025). +https://arxiv.org/abs/2503.20783 +""" + +from typing import List, Literal, Tuple + +import torch +import torch.distributed as dist + +from xorl.objectives.loss_output import LossOutput +from xorl.objectives.reducers import Reducer, TokenPartial +from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.loss.per_token_ce import compute_per_token_ce + + +KLType = Literal["k1", "k2", "k3"] +RatioType = Literal["token", "sequence"] + + +def compute_ratio( + logprobs: torch.Tensor, + generator_logprobs: torch.Tensor, + mask: torch.Tensor, + metric_reducer: Reducer, + ratio_type: RatioType = "token", +) -> Tuple[torch.Tensor, torch.Tensor, List[Tuple[str, torch.Tensor]]]: + """Importance sampling ratio r = π_θ/π_old. + + token: r_t = exp(logprobs_t - generator_logprobs_t) + sequence: r_seq = exp(mean_t[logprobs - generator_logprobs]), uses reparameterization. + """ + current_support = torch.isfinite(logprobs) + raw_token_log_ratio = torch.where( + current_support, + logprobs - generator_logprobs.detach(), + torch.full_like(logprobs, -torch.inf), + ) + if ratio_type == "token": + log_ratio = raw_token_log_ratio + ratio = torch.exp(log_ratio) + elif ratio_type == "sequence": + seq_lengths = mask.sum(dim=-1).clamp(min=1) + finite_token_log_ratio = torch.where( + current_support, + raw_token_log_ratio, + torch.zeros_like(raw_token_log_ratio), + ) + seq_log_ratio = (finite_token_log_ratio * mask).sum(dim=-1) / seq_lengths + seq_supported = (current_support | ~mask.bool()).all(dim=-1) + seq_log_ratio = torch.where( + seq_supported, + seq_log_ratio, + torch.full_like(seq_log_ratio, -torch.inf), + ) + + # Reparameterization: forward uses seq ratio, backward uses token grads + safe_logprobs = torch.where(current_support, logprobs, torch.zeros_like(logprobs)) + log_ratio = safe_logprobs - safe_logprobs.detach() + seq_log_ratio.detach().unsqueeze(-1) + ratio = torch.exp(log_ratio) + else: + raise ValueError(f"Unknown ratio_type: {ratio_type}") + + with torch.no_grad(): + metric_log_ratio = torch.where( + torch.isfinite(log_ratio), + log_ratio, + torch.full_like(log_ratio, -20.0), + ) + metrics = [ + ("loss/ratio/mean", metric_reducer(ratio, mask)), + ("loss/kl_policy/mean", metric_reducer(-metric_log_ratio, mask)), + ("loss/current_support_fraction", metric_reducer(current_support.float(), mask)), + ] + + return ratio, log_ratio, metrics + + +def compute_kl( + policy_logprobs: torch.Tensor, + ref_logprobs: torch.Tensor, + mask: torch.Tensor, + metric_reducer: Reducer, + kl_type: KLType = "k3", +) -> Tuple[torch.Tensor, List[Tuple[str, torch.Tensor]]]: + """KL divergence using Schulman's estimators (k1, k2, k3).""" + policy_support = torch.isfinite(policy_logprobs) + ref_support = torch.isfinite(ref_logprobs) + safe_policy_logprobs = torch.where(policy_support, policy_logprobs, torch.zeros_like(policy_logprobs)) + safe_ref_logprobs = torch.where(ref_support, ref_logprobs, torch.zeros_like(ref_logprobs)) + finite_log_ratio = safe_policy_logprobs - safe_ref_logprobs.detach() + log_ratio = torch.where( + policy_support & ref_support, + finite_log_ratio, + torch.where( + policy_support, + torch.full_like(finite_log_ratio, 20.0), + torch.where(ref_support, torch.full_like(finite_log_ratio, -20.0), torch.zeros_like(finite_log_ratio)), + ), + ) + + if kl_type == "k1": + kl = log_ratio + elif kl_type == "k2": + kl = 0.5 * log_ratio.square() + elif kl_type == "k3": + neg_log_ratio = torch.clamp(-log_ratio, min=-10.0, max=10.0) + ratio = torch.exp(neg_log_ratio) + kl = ratio - neg_log_ratio - 1 + else: + raise ValueError(f"Unknown kl_type: {kl_type}") + + with torch.no_grad(): + metrics = [("loss/kl_ref/mean", metric_reducer(kl, mask))] + + return kl, metrics + + +def pg_ppo_clip( + ratio: torch.Tensor, + advantages: torch.Tensor, + mask: torch.Tensor, + metric_reducer: Reducer, + clip_low: float = 0.2, + clip_high: float = 0.2, +) -> Tuple[torch.Tensor, List[Tuple[str, torch.Tensor]]]: + """PPO clipped surrogate: L = max(-r*A, -clip(r, 1-ε_low, 1+ε_high)*A).""" + clipped_ratio = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) + unclipped_loss = -ratio * advantages + clipped_loss = -clipped_ratio * advantages + pg_loss = torch.maximum(unclipped_loss, clipped_loss) + + with torch.no_grad(): + mask_bool = mask.bool() + clipped_high = (ratio > 1 + clip_high) & mask_bool + clipped_low = (ratio < 1 - clip_low) & mask_bool + pos_adv = advantages > 0 + neg_adv = advantages < 0 + + metrics = [ + ("loss/clip/clipped_ratio/mean", metric_reducer(clipped_ratio, mask)), + ("loss/clip/high_fraction", metric_reducer((clipped_high & pos_adv).float(), mask)), + ("loss/clip/low_fraction", metric_reducer((clipped_low & neg_adv).float(), mask)), + ] + + return pg_loss, metrics + + +def drgrpo_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + old_logprobs: torch.Tensor, + advantages: torch.Tensor, + ref_logprobs: torch.Tensor | None = None, + ignore_index: int = -100, + clip_low: float = 0.2, + clip_high: float = 0.28, + beta: float = 0.1, + ratio_type: RatioType = "token", + kl_type: KLType = "k3", + ce_mode: str = "compiled", + num_chunks: int = 8, + tp_group: dist.ProcessGroup | None = None, + lm_head_fp32: bool = False, + loss_reducer: Reducer | None = None, + metric_reducer: Reducer | None = None, + lm_head: torch.nn.Module | None = None, + logprob_temperature: float = 1.0, + logprob_top_k: int | torch.Tensor = TOP_K_ALL, + logprob_top_p: float | torch.Tensor = 1.0, + logprob_min_p: float | torch.Tensor = 0.0, +) -> LossOutput: + """DR-GRPO loss for RL training. + + Per-token: L_t = max(-r*A, -clip(r, 1-ε, 1+ε)*A) + β*KL + Aggregated: ``loss_reducer(per_token_loss, mask)``. Defaults to + ``TokenPartial(scale=loss_mask.sum())`` — the local active-token mean. + + Args: + hidden_states: (B, S, H) model hidden states. + weight: (V, H) or (V/tp, H) LM head weight. + labels: (B, S) target token IDs, already next-token aligned. + old_logprobs: (B, S) log probs from generation policy. + advantages: (B, S) per-token advantages. + ref_logprobs: (B, S) reference model log probs for KL (required if beta > 0). + ignore_index: Token ID to ignore (default: -100). + clip_low: Lower clip bound (default: 0.2). + clip_high: Upper clip bound (default: 0.28). + beta: KL penalty coefficient (default: 0.1). + ratio_type: Ratio type: "token" or "sequence" (default: "token"). + kl_type: KL estimator: "k1", "k2", "k3" (default: "k3"). + ce_mode: Cross-entropy mode: "compiled" or "eager". + num_chunks: Chunks for compiled mode. + tp_group: TP process group for vocab-parallel CE. + lm_head_fp32: Compute LM head in FP32. + loss_reducer / metric_reducer: Both default to + ``TokenPartial(scale=loss_mask.sum())`` (legacy local active-token + mean; does not compose across mbs/ranks). Pass shared global-scale + reducers to make summed partial shares recover the global value. + logprob_temperature: Temperature applied to trainer logits before + selected-token logprob calculation. ``1.0`` is raw policy logprobs; + setting this to the rollout temperature yields behavior-policy + semantics for DR-GRPO ratios and K3 stats. + + Returns: + LossOutput with loss, per_token_logprobs, per_token_loss, and metrics. + """ + if beta > 0 and ref_logprobs is None: + raise ValueError("ref_logprobs required when beta > 0") + + B, S = labels.shape + H = hidden_states.size(-1) + + labels_flat = labels.reshape(-1) + hidden_states_flat = hidden_states.reshape(-1, H) + + per_token_ce = compute_per_token_ce( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + lm_head_fp32=lm_head_fp32, + lm_head=lm_head, + logprob_temperature=logprob_temperature, + logprob_top_k=logprob_top_k, + logprob_top_p=logprob_top_p, + logprob_min_p=logprob_min_p, + ) + logprobs = -per_token_ce.view(B, S) + + loss_mask = (labels != ignore_index).float() + + if metric_reducer is None: + metric_reducer = TokenPartial(scale=loss_mask.sum()) + if loss_reducer is None: + loss_reducer = TokenPartial(scale=loss_mask.sum()) + + ratio, _, ratio_m = compute_ratio(logprobs, old_logprobs, loss_mask, metric_reducer, ratio_type) + + pg_loss, clip_m = pg_ppo_clip(ratio, advantages, loss_mask, metric_reducer, clip_low, clip_high) + + kl_m: List[Tuple[str, torch.Tensor]] = [] + if beta > 0: + kl, kl_m = compute_kl(logprobs, ref_logprobs, loss_mask, metric_reducer, kl_type) + pg_loss = pg_loss + beta * kl + + loss = loss_reducer(pg_loss, loss_mask) + + metrics = dict(ratio_m + clip_m + kl_m) + + return LossOutput( + loss=loss, + per_token_logprobs=logprobs.detach(), + per_token_loss=pg_loss.detach(), + metrics=metrics, + ) diff --git a/src/xorl/objectives/importance_sampling_loss.py b/src/xorl/objectives/importance_sampling_loss.py new file mode 100644 index 00000000..7af4d275 --- /dev/null +++ b/src/xorl/objectives/importance_sampling_loss.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from typing import Any, Dict, Optional + +import torch +import torch.distributed as dist + +from xorl.objectives.loss_output import LossOutput +from xorl.objectives.reducers import Reducer, TokenPartial +from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.loss.per_token_ce import compute_per_token_ce + + +K3_DEBUG_THRESHOLDS = ( + ("1e_minus_6", 1e-6), + ("1e_minus_4", 1e-4), + ("1e_minus_3", 1e-3), + ("1e_minus_2", 1e-2), + ("1e_minus_1", 1e-1), + ("1", 1.0), +) + + +def importance_sampling_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + old_logprobs: torch.Tensor, + advantages: torch.Tensor, + ignore_index: int = -100, + num_chunks: int = 8, + ce_mode: str = "compiled", + return_per_token: bool = False, + tp_group: Optional[dist.ProcessGroup] = None, + compute_kl_stats: bool = False, + lm_head_fp32: bool = False, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, + lm_head: Optional[torch.nn.Module] = None, + logprob_temperature: float = 1.0, + logprob_top_k: int | torch.Tensor = TOP_K_ALL, + logprob_top_p: float | torch.Tensor = 1.0, + logprob_min_p: float | torch.Tensor = 0.0, +) -> "LossOutput": + """ + Compute importance sampling loss for GRPO/RL training. + + This implements the Tinker-style importance sampling loss: + prob_ratio = exp(new_logprobs - old_logprobs) + loss = -(prob_ratio * advantages).mean() + + Supports multiple computation modes: + - "compiled": RECOMMENDED. torch.compile (1.6x speed, 16% memory) + - "eager": Simple F.cross_entropy baseline (may OOM at 32K) + + Args: + hidden_states: Model hidden states, shape (batch, seq_len, hidden_dim) + weight: LM head weight matrix, shape (vocab_size, hidden_dim) + labels: Target token IDs, shape (batch, seq_len). Already next-token aligned. + old_logprobs: Old policy log probabilities from sampling, shape (batch, seq_len) + advantages: Per-token advantages, shape (batch, seq_len) + ignore_index: Index to ignore in loss computation (default: -100) + num_chunks: Number of chunks for compiled mode (default: 8). + ce_mode: Cross-entropy mode - "compiled" (default) or "eager" + return_per_token: If True, returns per-token logprobs and per-token CE loss. + Useful for custom loss computations. + compute_kl_stats: If True, compute and return KL statistics in metrics dict: + - kl_sample_train_k3: Schulman's K3 estimator: mean(exp(log_ratio) - log_ratio - 1) + where log_ratio = new_logprobs - old_logprobs. Non-negative, unbiased, lower variance. + - entropy_sample: -mean(old_logprobs) over valid tokens + - valid_tokens: Count of valid tokens + loss_reducer: Reduces per-token loss to a scalar partial share. None => + ``TokenPartial(scale=valid_mask.sum())`` (legacy local token-mean; does + not compose across micro-batches/ranks). Pass a shared global-scale + reducer to make summed partial shares recover the global loss. + metric_reducer: Reduces per-token /mean metrics (ratio_mean, + kl_sample_train_k3, entropy_sample). ratio_min/ratio_max stay local + scalars and bypass it. + logprob_temperature: Temperature applied to trainer logits before + selected-token logprob calculation. ``1.0`` is raw policy logprobs; + setting this to the rollout temperature yields behavior-policy + semantics for the sampled-token ratio. + + Returns: + LossOutput with loss, per_token_logprobs, per_token_loss, and metrics. + """ + original_shape = labels.shape + H = hidden_states.size(-1) + + # Flatten tensors + labels_flat = labels.reshape(-1) + hidden_states_flat = hidden_states.reshape(-1, H) + old_logprobs_flat = old_logprobs.reshape(-1) + advantages_flat = advantages.reshape(-1) + + # Valid/action mask + valid_mask = labels_flat != ignore_index + valid_mask_f = valid_mask.float() + valid_count = valid_mask.sum() + + if loss_reducer is None: + loss_reducer = TokenPartial(scale=valid_count.float()) + if metric_reducer is None: + metric_reducer = TokenPartial(scale=valid_count.float()) + + # ---- Cross-entropy computation ---- + per_token_ce = compute_per_token_ce( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + lm_head_fp32=lm_head_fp32, + lm_head=lm_head, + logprob_temperature=logprob_temperature, + logprob_top_k=logprob_top_k, + logprob_top_p=logprob_top_p, + logprob_min_p=logprob_min_p, + ) + + current_support = torch.isfinite(per_token_ce) + # new logprobs = log p(target) = -CE + new_logprobs_flat = -per_token_ce.detach() + + # ---- ratio computation (no sanitization) ---- + delta = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) + delta = torch.where(current_support, delta.clamp(min=-20.0, max=20.0), torch.full_like(delta, -torch.inf)) + ratio = torch.exp(delta) + + # ---- Per-token policy gradient loss: -(ratio * advantages) ---- + per_token_pg = -(ratio * advantages_flat) + per_token_pg = per_token_pg.masked_fill(~valid_mask, 0.0) + + # ---- Option B: value from true PG, grad from weighted CE surrogate ---- + true_pg = loss_reducer(per_token_pg, valid_mask_f) + + w = (ratio.detach() * advantages_flat).masked_fill(~valid_mask | ~current_support, 0.0) + safe_per_token_ce = torch.where(current_support, per_token_ce, torch.zeros_like(per_token_ce)) + surrogate = loss_reducer(w * safe_per_token_ce, valid_mask_f) + + loss = true_pg.detach() + surrogate - surrogate.detach() + + # ±inf identity on empty ranks lets cross-rank MIN/MAX-allreduce ignore empty contributors. + if valid_mask.any(): + ratio_min = ratio.masked_fill(~valid_mask, float("inf")).min() + ratio_max = ratio.masked_fill(~valid_mask, float("-inf")).max() + else: + ratio_min = ratio.new_tensor(float("inf")) + ratio_max = ratio.new_tensor(float("-inf")) + metrics: Dict[str, Any] = { + "ratio_mean": metric_reducer(ratio, valid_mask_f).detach(), + "ratio_min": ratio_min.detach(), + "ratio_max": ratio_max.detach(), + } + + if compute_kl_stats: + with torch.no_grad(): + raw_log_ratio = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) + log_ratio_full = torch.where( + current_support, + raw_log_ratio, + torch.full_like(raw_log_ratio, -20.0), + ) + ratio_full = torch.where(current_support, torch.exp(raw_log_ratio), torch.zeros_like(raw_log_ratio)) + per_token_k3 = ratio_full - log_ratio_full - 1.0 + if valid_mask.any(): + k3_max = per_token_k3.masked_fill(~valid_mask, float("-inf")).max() + logratio_min = log_ratio_full.masked_fill(~valid_mask, float("inf")).min() + logratio_max = log_ratio_full.masked_fill(~valid_mask, float("-inf")).max() + abs_logratio_max = log_ratio_full.abs().masked_fill(~valid_mask, float("-inf")).max() + else: + k3_max = per_token_k3.new_tensor(float("-inf")) + logratio_min = log_ratio_full.new_tensor(float("inf")) + logratio_max = log_ratio_full.new_tensor(float("-inf")) + abs_logratio_max = log_ratio_full.new_tensor(float("-inf")) + metrics["kl_sample_train_k3"] = metric_reducer(per_token_k3, valid_mask_f) + metrics["kl_k3_debug_mean"] = metric_reducer(per_token_k3, valid_mask_f) + metrics["kl_k3_debug_max"] = k3_max + metrics["kl_k3_debug_abs_logratio_mean"] = metric_reducer(log_ratio_full.abs(), valid_mask_f) + metrics["kl_k3_debug_abs_logratio_max"] = abs_logratio_max + metrics["kl_k3_debug_logratio_mean"] = metric_reducer(log_ratio_full, valid_mask_f) + metrics["kl_k3_debug_logratio_min"] = logratio_min + metrics["kl_k3_debug_logratio_max"] = logratio_max + metrics["kl_k3_debug_frac_logratio_positive"] = metric_reducer((log_ratio_full > 0).float(), valid_mask_f) + for suffix, threshold in K3_DEBUG_THRESHOLDS: + metrics[f"kl_k3_debug_frac_gt_{suffix}"] = metric_reducer( + (per_token_k3 > threshold).float(), valid_mask_f + ) + metrics["entropy_sample"] = metric_reducer(-old_logprobs_flat, valid_mask_f) + metrics["valid_tokens"] = valid_count.item() + metrics["current_support_fraction"] = metric_reducer(current_support.float(), valid_mask_f) + + # Reshape per-token outputs + per_token_logprobs = new_logprobs_flat.view(original_shape) + per_token_loss = per_token_pg.view(original_shape) + + metric_ops = {"ratio_min": "min", "ratio_max": "max"} + if compute_kl_stats: + metric_ops.update( + { + "kl_k3_debug_max": "max", + "kl_k3_debug_abs_logratio_max": "max", + "kl_k3_debug_logratio_min": "min", + "kl_k3_debug_logratio_max": "max", + } + ) + + return LossOutput( + loss=loss, + per_token_logprobs=per_token_logprobs, + per_token_loss=per_token_loss, + metrics=metrics, + metric_ops=metric_ops, + ) diff --git a/src/xorl/objectives/loss_output.py b/src/xorl/objectives/loss_output.py new file mode 100644 index 00000000..41a79eee --- /dev/null +++ b/src/xorl/objectives/loss_output.py @@ -0,0 +1,21 @@ +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch + + +@dataclass +class LossOutput: + """Standardized return type for all loss functions. + + ``metric_ops`` tags ``metrics`` keys whose cross-mb / cross-rank composition + isn't the default mean (``"min"``/``"max"``). The sidecar (rather than a + tagged-value type in ``metrics``) keeps the metrics dict directly + JSON-serializable for untagged consumers. + """ + + loss: torch.Tensor + per_token_logprobs: Optional[torch.Tensor] = None + per_token_loss: Optional[torch.Tensor] = None + metrics: Optional[Dict[str, Any]] = None + metric_ops: Optional[Dict[str, str]] = None diff --git a/src/xorl/objectives/opd_loss.py b/src/xorl/objectives/opd_loss.py new file mode 100644 index 00000000..570ad800 --- /dev/null +++ b/src/xorl/objectives/opd_loss.py @@ -0,0 +1,1324 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, Optional + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from xorl.objectives.loss_output import LossOutput +from xorl.objectives.opd_streaming_kl import ( + streaming_forward_kl_function, + streaming_forward_kl_lowmem_function, + streaming_full_vocab_diagnostics, + streaming_reverse_kl_function, + streaming_reverse_kl_lowmem_function, +) +from xorl.objectives.reducers import Reducer, TokenPartial +from xorl.ops.loss.compiled_cross_entropy import ( + compiled_forward_kl_full_function, + compiled_forward_kl_full_with_diag_function, + compiled_reverse_kl_function, + compiled_reverse_kl_with_diag_function, + compiled_sampled_token_logprobs_function, +) +from xorl.ops.loss.vocab_parallel_reverse_kl import vocab_parallel_reverse_kl_gathered + + +# OPD loss modes. +# +# `reverse_kl_full` (default) and `forward_kl_full` are full-vocabulary KL +# variants computed from the *full* teacher distribution — xorl ships teacher +# hidden states + the cached teacher LM head separately, so the full p_T is +# materialized at loss time. VERL's `forward_kl_topk` is intentionally NOT +# ported: it's a bandwidth workaround for teacher-emits-truncated-logprobs +# pipelines, strictly worse than `forward_kl_full` for the same student compute. +# +# Single-sample KL estimators are ported from verl/trainer/ppo/core_algos.py:: +# kl_penalty as a different (cheap, sampled-token-only) point in the +# memory/quality trade-off. A trailing "+" applies the k2 straight-through- +# gradient trick (same forward, k2 gradient). +LOSS_MODE_REVERSE_KL_FULL = "reverse_kl_full" +LOSS_MODE_FORWARD_KL_FULL = "forward_kl_full" +_ESTIMATOR_MODES = {"kl", "k1", "abs", "mse", "k2", "low_var_kl", "k3"} + + +def _strip_estimator_plus(loss_mode: str) -> tuple[str, bool]: + if loss_mode.endswith("+"): + return loss_mode[:-1], True + return loss_mode, False + + +def is_estimator_loss_mode(loss_mode: str) -> bool: + base, _ = _strip_estimator_plus(loss_mode) + return base in _ESTIMATOR_MODES + + +def _kl_penalty_estimator( + logprob: torch.Tensor, + ref_logprob: torch.Tensor, + loss_mode: str, +) -> torch.Tensor: + """Single-sample KL estimators, byte-for-byte port of VERL's kl_penalty. + + See verl/trainer/ppo/core_algos.py::kl_penalty and ::kl_penalty_forward. + A trailing "+" (e.g. "k3+") applies the k2 straight-through trick: forward + value from the chosen estimator, but backward as if it were 0.5*(logp-ref)^2. + """ + base, straight_through = _strip_estimator_plus(loss_mode) + forward = _kl_penalty_forward(logprob, ref_logprob, base) + if not straight_through or base in ("mse", "k2"): + return forward + backward_score = 0.5 * (logprob - ref_logprob).square() + return backward_score - backward_score.detach() + forward.detach() + + +def _kl_penalty_forward( + logprob: torch.Tensor, + ref_logprob: torch.Tensor, + loss_mode: str, +) -> torch.Tensor: + if loss_mode in ("kl", "k1"): + return logprob - ref_logprob + if loss_mode == "abs": + return (logprob - ref_logprob).abs() + if loss_mode in ("mse", "k2"): + return 0.5 * (logprob - ref_logprob).square() + if loss_mode in ("low_var_kl", "k3"): + kl = ref_logprob - logprob + kl = torch.clamp(kl, min=-20, max=20) + ratio = torch.exp(kl) + kld = (ratio - kl - 1).contiguous() + return torch.clamp(kld, min=-10, max=10) + raise ValueError(f"Unknown KL estimator '{loss_mode}'") + + +@dataclass(frozen=True) +class OPDLossMetrics: + """OPD loss metrics emitted per micro-batch. + + All fields are always present in `to_dict()` (defaulting to 0.0 / 0 when + not applicable to the active loss_mode). This is intentional — dict-keyed + distributed all-reduces deadlock when ranks have different key sets, so + every loss path on every rank must contribute the same dict shape. + """ + + valid_tokens: int + opd_kl: float = 0.0 + opd_weighted_kl: float = 0.0 + opd_hidden_match_loss: float = 0.0 + opd_hidden_match_raw_loss: float = 0.0 + opd_hidden_match_weight_mean: float = 0.0 + opd_hidden_match_pos_loss: float = 0.0 + opd_hidden_match_neg_loss: float = 0.0 + opd_hidden_match_pos_raw_loss: float = 0.0 + opd_hidden_match_neg_raw_loss: float = 0.0 + opd_hidden_match_neg_minus_pos_raw: float = 0.0 + opd_hidden_match_pos_weight_mean: float = 0.0 + opd_hidden_match_neg_weight_mean: float = 0.0 + opd_teacher_weight_mean: float = 0.0 + opd_num_teachers: int = 0 + # ---- Full-vocab diagnostics (reverse_kl_full / forward_kl_full) ---- + opd_teacher_entropy: float = 0.0 + opd_student_entropy: float = 0.0 + opd_top1_agreement: float = 0.0 + # ---- KL-estimator diagnostic (k1/abs/...) ---- + opd_abs_loss: float = 0.0 + # ---- Distillation-loss range ---- + opd_loss_min: float = 0.0 + opd_loss_max: float = 0.0 + opd_loss_abs_mean: float = 0.0 + # Fraction of valid tokens whose pre-clamp |KL| >= loss_max_clamp. Clamped + # tokens pass ZERO gradient (gradient-dead mass) — watch this when the loss + # plateaus while opd_loss_max sits pinned at the clamp. + opd_loss_clamp_frac: float = 0.0 + # ---- Region / sample-correctness KL splits (diagnostic) ---- + # All `*_per_valid` fields are masked SUMS divided by the micro-batch's TOTAL + # valid-token count (not per-region means), and every `opd_frac_*` is a + # region-token count over the same denominator. Both are linear functionals + # of per-token values, so the valid-token-weighted mean aggregation across + # micro-batches / ranks recomposes them EXACTLY; derive the human-readable + # region mean as `per_valid / frac` after aggregation (the client does this). + # Regions come from `diag_region_ids` (0=prompt, 1=buffer, 2=answer, -1=n/a); + # correctness from `diag_sample_ok` (1=sampled answer correct, 0=wrong, + # -1=unknown), both client-provided per-token tensors. + opd_kl_prompt_per_valid: float = 0.0 + opd_kl_buffer_per_valid: float = 0.0 + opd_kl_answer_per_valid: float = 0.0 + opd_frac_prompt: float = 0.0 + opd_frac_buffer: float = 0.0 + opd_frac_answer: float = 0.0 + opd_kl_answer_correct_per_valid: float = 0.0 + opd_kl_answer_wrong_per_valid: float = 0.0 + opd_frac_answer_correct: float = 0.0 + opd_frac_answer_wrong: float = 0.0 + opd_student_entropy_answer_correct_per_valid: float = 0.0 + opd_student_entropy_answer_wrong_per_valid: float = 0.0 + opd_teacher_entropy_answer_correct_per_valid: float = 0.0 + opd_teacher_entropy_answer_wrong_per_valid: float = 0.0 + # ---- Multi-layer OPRD (all-layer hidden matching) ---- + opd_oprd_loss: float = 0.0 + opd_oprd_raw_loss: float = 0.0 + opd_oprd_num_layers: int = 0 + # ---- Vocab-parallel OPD debugging ---- + # These are raw numerator/count diagnostics from the vocab-parallel group, + # not values used by normal metric aggregation. They stay in the fixed + # metric key set so optional debug consumers do not introduce rank-skewed + # metric dictionaries. + opd_vocab_parallel_group_tokens: int = 0 + opd_vocab_parallel_kl_sum: float = 0.0 + opd_vocab_parallel_weighted_kl_sum: float = 0.0 + # ---- PG-mode (use_policy_gradient=True) ---- + opd_pg_clipfrac: float = 0.0 + opd_pg_clipfrac_lower: float = 0.0 + opd_ppo_kl: float = 0.0 + + def to_dict(self) -> dict[str, int | float]: + return { + "valid_tokens": self.valid_tokens, + "opd_kl": self.opd_kl, + "opd_weighted_kl": self.opd_weighted_kl, + "opd_hidden_match_loss": self.opd_hidden_match_loss, + "opd_hidden_match_raw_loss": self.opd_hidden_match_raw_loss, + "opd_hidden_match_weight_mean": self.opd_hidden_match_weight_mean, + "opd_hidden_match_pos_loss": self.opd_hidden_match_pos_loss, + "opd_hidden_match_neg_loss": self.opd_hidden_match_neg_loss, + "opd_hidden_match_pos_raw_loss": self.opd_hidden_match_pos_raw_loss, + "opd_hidden_match_neg_raw_loss": self.opd_hidden_match_neg_raw_loss, + "opd_hidden_match_neg_minus_pos_raw": self.opd_hidden_match_neg_minus_pos_raw, + "opd_hidden_match_pos_weight_mean": self.opd_hidden_match_pos_weight_mean, + "opd_hidden_match_neg_weight_mean": self.opd_hidden_match_neg_weight_mean, + "opd_teacher_weight_mean": self.opd_teacher_weight_mean, + "opd_num_teachers": self.opd_num_teachers, + "opd_teacher_entropy": self.opd_teacher_entropy, + "opd_student_entropy": self.opd_student_entropy, + "opd_top1_agreement": self.opd_top1_agreement, + "opd_abs_loss": self.opd_abs_loss, + "opd_loss_min": self.opd_loss_min, + "opd_loss_max": self.opd_loss_max, + "opd_loss_abs_mean": self.opd_loss_abs_mean, + "opd_loss_clamp_frac": self.opd_loss_clamp_frac, + "opd_kl_prompt_per_valid": self.opd_kl_prompt_per_valid, + "opd_kl_buffer_per_valid": self.opd_kl_buffer_per_valid, + "opd_kl_answer_per_valid": self.opd_kl_answer_per_valid, + "opd_frac_prompt": self.opd_frac_prompt, + "opd_frac_buffer": self.opd_frac_buffer, + "opd_frac_answer": self.opd_frac_answer, + "opd_kl_answer_correct_per_valid": self.opd_kl_answer_correct_per_valid, + "opd_kl_answer_wrong_per_valid": self.opd_kl_answer_wrong_per_valid, + "opd_frac_answer_correct": self.opd_frac_answer_correct, + "opd_frac_answer_wrong": self.opd_frac_answer_wrong, + "opd_student_entropy_answer_correct_per_valid": self.opd_student_entropy_answer_correct_per_valid, + "opd_student_entropy_answer_wrong_per_valid": self.opd_student_entropy_answer_wrong_per_valid, + "opd_teacher_entropy_answer_correct_per_valid": self.opd_teacher_entropy_answer_correct_per_valid, + "opd_teacher_entropy_answer_wrong_per_valid": self.opd_teacher_entropy_answer_wrong_per_valid, + "opd_oprd_loss": self.opd_oprd_loss, + "opd_oprd_raw_loss": self.opd_oprd_raw_loss, + "opd_oprd_num_layers": self.opd_oprd_num_layers, + "opd_vocab_parallel_group_tokens": self.opd_vocab_parallel_group_tokens, + "opd_vocab_parallel_kl_sum": self.opd_vocab_parallel_kl_sum, + "opd_vocab_parallel_weighted_kl_sum": self.opd_vocab_parallel_weighted_kl_sum, + "opd_pg_clipfrac": self.opd_pg_clipfrac, + "opd_pg_clipfrac_lower": self.opd_pg_clipfrac_lower, + "opd_ppo_kl": self.opd_ppo_kl, + } + + +def _as_flat_optional_weights( + teacher_weights: Optional[torch.Tensor], + valid_mask: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + if teacher_weights is None: + return torch.ones(valid_mask.sum(), dtype=dtype, device=valid_mask.device) + weights_flat = teacher_weights.reshape(-1).to(device=valid_mask.device, dtype=dtype) + return weights_flat[valid_mask] + + +def _zero_loss_with_graph(hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Build a 0-valued loss that still flows gradients through hidden_states + weight. + + Always returns fp32 so the dtype matches the normal-return path + (`total_weighted_kl / denom`, fp32). A dtype mismatch between the early-return + branch on no-valid-token ranks and the fp32 normal branch corrupts NCCL + all_reduce in the trainer's loss-reporting path. + """ + + def anchor(tensor: torch.Tensor) -> torch.Tensor: + if tensor.numel() == 0: + return torch.zeros((), dtype=torch.float32, device=tensor.device) + return tensor.reshape(-1)[:1].float().sum() * 0.0 + + return anchor(hidden_states) + anchor(weight) + + +def _denominator_tensor( + denominator: torch.Tensor | int | float | None, + *, + fallback: torch.Tensor, + device: torch.device | str, +) -> torch.Tensor: + if denominator is None: + return fallback.to(device=device, dtype=torch.float32) + if torch.is_tensor(denominator): + return denominator.to(device=device, dtype=torch.float32) + return torch.tensor(float(denominator), device=device, dtype=torch.float32) + + +def _gather_1d_nograd(local_values: torch.Tensor, *, group) -> torch.Tensor: + if local_values.ndim != 1: + raise ValueError(f"expected rank-1 tensor, got shape {tuple(local_values.shape)}") + if not dist.is_available() or not dist.is_initialized(): + return local_values.detach() + world = dist.get_world_size(group) + if world == 1: + return local_values.detach() + + local_count = torch.tensor([local_values.shape[0]], dtype=torch.long, device=local_values.device) + counts = torch.empty(world, dtype=torch.long, device=local_values.device) + dist.all_gather_into_tensor(counts, local_count, group=group) + max_count = int(counts.max().item()) if counts.numel() else 0 + if max_count == 0: + return local_values.detach().new_empty((0,)) + + padded = local_values.detach().new_zeros((max_count,)) + if local_values.numel(): + padded[: local_values.shape[0]].copy_(local_values.detach().contiguous()) + gathered = local_values.detach().new_empty((world * max_count,)) + dist.all_gather_into_tensor(gathered, padded.contiguous(), group=group) + pieces = [] + for rank, n_rows in enumerate(int(x) for x in counts.detach().cpu().tolist()): + if n_rows: + start = rank * max_count + pieces.append(gathered[start : start + n_rows]) + if not pieces: + return local_values.detach().new_empty((0,)) + return torch.cat(pieces, dim=0) + + +def _gathered_local_bounds(n_local: int, *, device: torch.device, group) -> tuple[int, int]: + """Return this rank's [start, end) slice in the gathered 1-D token order.""" + if not dist.is_available() or not dist.is_initialized(): + return 0, int(n_local) + world = dist.get_world_size(group) + if world == 1: + return 0, int(n_local) + + rank = dist.get_rank(group) + local_count = torch.tensor([int(n_local)], dtype=torch.long, device=device) + counts = torch.empty(world, dtype=torch.long, device=device) + dist.all_gather_into_tensor(counts, local_count, group=group) + start = int(counts[:rank].sum().item()) + return start, start + int(n_local) + + +def _value_divided_without_grad_scale(loss: torch.Tensor, divisor: int) -> torch.Tensor: + if divisor <= 1: + return loss + reported = loss / float(divisor) + return loss + (reported - loss).detach() + + +def _oprd_hidden_distance( + student_layer_hidden_states: torch.Tensor, + teacher_layer_hidden_states: torch.Tensor, + expected_rows: int, + *, + layer_chunk_size: int = 4, +) -> tuple[torch.Tensor, int]: + """Return per-token all-layer OPRD MSE without materializing all layers as fp32.""" + if student_layer_hidden_states.shape != teacher_layer_hidden_states.shape: + raise ValueError( + "OPRD requires matching student/teacher layer shapes, got " + f"student={tuple(student_layer_hidden_states.shape)} teacher={tuple(teacher_layer_hidden_states.shape)}" + ) + if student_layer_hidden_states.ndim != 3: + raise ValueError( + "OPRD layer tensors must be rank-3 [valid_tokens, layers, hidden], got " + f"{tuple(student_layer_hidden_states.shape)}" + ) + if student_layer_hidden_states.shape[0] != expected_rows: + raise ValueError( + f"OPRD layer tensors must have {expected_rows} rows, got {student_layer_hidden_states.shape[0]}" + ) + + return _oprd_hidden_distance_from_fetcher( + student_layer_hidden_states, + teacher_layer_fetcher=lambda start, end: teacher_layer_hidden_states.detach()[:, start:end, :], + expected_rows=expected_rows, + num_layers=int(student_layer_hidden_states.shape[1]), + layer_chunk_size=layer_chunk_size, + ) + + +def _oprd_hidden_distance_from_fetcher( + student_layer_hidden_states: torch.Tensor, + *, + teacher_layer_fetcher: Callable[[int, int], torch.Tensor], + expected_rows: int, + num_layers: int, + layer_chunk_size: int = 4, +) -> tuple[torch.Tensor, int]: + """Return per-token all-layer OPRD MSE while fetching teacher layers by chunk.""" + if student_layer_hidden_states.ndim != 3: + raise ValueError( + "OPRD layer tensors must be rank-3 [valid_tokens, layers, hidden], got " + f"{tuple(student_layer_hidden_states.shape)}" + ) + if student_layer_hidden_states.shape[0] != expected_rows: + raise ValueError( + f"OPRD layer tensors must have {expected_rows} rows, got {student_layer_hidden_states.shape[0]}" + ) + if int(student_layer_hidden_states.shape[1]) != int(num_layers): + raise ValueError( + f"OPRD student layer count {student_layer_hidden_states.shape[1]} " + f"does not match teacher layer count {num_layers}" + ) + if expected_rows == 0: + return torch.empty(0, dtype=torch.float32, device=student_layer_hidden_states.device), num_layers + if num_layers <= 0: + raise ValueError("OPRD layer tensors must include at least one layer") + + chunk_size = max(1, int(layer_chunk_size)) + hidden_distance_sum = None + for start in range(0, num_layers, chunk_size): + end = min(start + chunk_size, num_layers) + student_chunk = student_layer_hidden_states[:, start:end, :].float() + teacher_chunk = teacher_layer_fetcher(start, end).detach() + expected_shape = student_layer_hidden_states.shape[0], end - start, student_layer_hidden_states.shape[2] + if tuple(teacher_chunk.shape) != tuple(expected_shape): + raise ValueError( + "OPRD teacher layer fetcher returned wrong shape: " + f"got {tuple(teacher_chunk.shape)}, expected {tuple(expected_shape)}" + ) + teacher_chunk = teacher_chunk.float() + chunk_distance = (student_chunk - teacher_chunk).square().mean(dim=-1).sum(dim=-1) + hidden_distance_sum = chunk_distance if hidden_distance_sum is None else hidden_distance_sum + chunk_distance + return hidden_distance_sum / float(num_layers), num_layers + + +def opd_vocab_parallel_loss_function( + student_hidden_flat: torch.Tensor, + student_weight_local: torch.Tensor, + labels: torch.Tensor, + teacher_hidden_flat: torch.Tensor, + teacher_weight_local: torch.Tensor, + teacher_weights: Optional[torch.Tensor] = None, + hidden_match_weights: Optional[torch.Tensor] = None, + ignore_index: int = -100, + lm_head_fp32: bool = False, + teacher_lm_head_fp32: bool = True, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, + loss_mode: str = LOSS_MODE_REVERSE_KL_FULL, + loss_max_clamp: Optional[float] = None, + use_task_rewards: bool = False, + distillation_loss_coef: float = 1.0, + hidden_match_coef: float = 0.0, + kl_loss_weight: float = 1.0, + hidden_match_mode: str = "cosine", + teacher_layer_hidden_states: Optional[torch.Tensor] = None, + teacher_layer_fetcher: Optional[Callable[[int, int], torch.Tensor]] = None, + teacher_layer_num_layers: Optional[int] = None, + oprd_layer_chunk_size: int = 4, + student_layer_hidden_states: Optional[torch.Tensor] = None, + group=None, + debug_token_outputs: bool = False, +) -> LossOutput: + """OPD reverse-KL over vocab-sharded student/teacher LM heads. + + Inputs are already restricted to this rank's valid tokens for one teacher. + The vocab-parallel KL gathers token activations across ``group`` and returns + a replicated full-token KL vector. Its loss value is divided by the group + size for detached reporting only; gradients remain the full global-token sum + and are normalized later by the trainer's global-valid-token scale. + """ + if loss_mode != LOSS_MODE_REVERSE_KL_FULL: + raise ValueError("vocab-parallel OPD currently supports only loss_mode='reverse_kl_full'") + if labels.ndim != 1: + raise ValueError(f"vocab-parallel OPD labels must be rank-1, got {tuple(labels.shape)}") + if student_hidden_flat.ndim != 2 or teacher_hidden_flat.ndim != 2: + raise ValueError( + "vocab-parallel OPD hidden tensors must be rank-2, got " + f"student={tuple(student_hidden_flat.shape)} teacher={tuple(teacher_hidden_flat.shape)}" + ) + if student_hidden_flat.shape[0] != labels.shape[0] or teacher_hidden_flat.shape[0] != labels.shape[0]: + raise ValueError( + "vocab-parallel OPD token counts must match, got " + f"student={student_hidden_flat.shape[0]} teacher={teacher_hidden_flat.shape[0]} labels={labels.shape[0]}" + ) + if student_hidden_flat.shape[-1] != student_weight_local.shape[-1]: + raise ValueError( + "student hidden size " + f"({student_hidden_flat.shape[-1]}) must match local student head width " + f"({student_weight_local.shape[-1]})" + ) + if teacher_hidden_flat.shape[-1] != teacher_weight_local.shape[-1]: + raise ValueError( + "teacher hidden size " + f"({teacher_hidden_flat.shape[-1]}) must match local teacher head width " + f"({teacher_weight_local.shape[-1]})" + ) + if teacher_weight_local.shape[0] != student_weight_local.shape[0]: + raise ValueError( + "local student and teacher vocab shards must have the same rows, got " + f"{student_weight_local.shape[0]} and {teacher_weight_local.shape[0]}" + ) + + local_tokens = int(labels.shape[0]) + device = student_hidden_flat.device + if teacher_weights is None: + local_token_weights = torch.ones(local_tokens, dtype=torch.float32, device=device) + else: + if teacher_weights.shape[0] != local_tokens: + raise ValueError(f"teacher_weights must have {local_tokens} rows, got {tuple(teacher_weights.shape)}") + local_token_weights = teacher_weights.to(device=device, dtype=torch.float32) + if hidden_match_weights is None: + local_hidden_weights = local_token_weights + else: + if hidden_match_weights.shape[0] != local_tokens: + raise ValueError( + f"hidden_match_weights must have {local_tokens} rows, got {tuple(hidden_match_weights.shape)}" + ) + local_hidden_weights = hidden_match_weights.to(device=device, dtype=torch.float32) + + if lm_head_fp32: + student_hidden_for_kl = student_hidden_flat.float() + student_weight_for_kl = student_weight_local.float() + else: + student_hidden_for_kl = student_hidden_flat + student_weight_for_kl = student_weight_local + if teacher_lm_head_fp32: + teacher_hidden_for_kl = teacher_hidden_flat.float() + teacher_weight_for_kl = teacher_weight_local.float() + else: + teacher_hidden_for_kl = teacher_hidden_flat + teacher_weight_for_kl = teacher_weight_local + + token_kl = vocab_parallel_reverse_kl_gathered( + local_student_hidden=student_hidden_for_kl, + student_weight_local=student_weight_for_kl, + local_teacher_hidden=teacher_hidden_for_kl.detach(), + teacher_weight_local=teacher_weight_for_kl.detach(), + local_labels=labels.to(device=device), + ignore_index=ignore_index, + group=group, + ) + full_token_weights = _gather_1d_nograd(local_token_weights.to(token_kl.device), group=group) + if full_token_weights.shape[0] != token_kl.shape[0]: + raise ValueError( + f"gathered teacher_weights rows {full_token_weights.shape[0]} do not match KL rows {token_kl.shape[0]}" + ) + + clamp_frac = 0.0 + if loss_max_clamp is not None and token_kl.numel(): + clamp_frac = (token_kl.detach().abs() >= float(loss_max_clamp)).float().mean().item() + token_kl = token_kl.clamp(min=-loss_max_clamp, max=loss_max_clamp) + + if loss_reducer is None: + scale = torch.tensor(max(int(token_kl.numel()), 1), dtype=torch.float32, device=token_kl.device) + loss_reducer = TokenPartial(scale=scale) + if metric_reducer is None: + scale = torch.tensor(max(local_tokens, 1), dtype=torch.float32, device=device) + metric_reducer = TokenPartial(scale=scale) + + weighted_token_kl = token_kl * full_token_weights.to(token_kl.device) + valid_full = torch.ones_like(weighted_token_kl, dtype=torch.float32) + raw_kl_loss = loss_reducer(weighted_token_kl, valid_full) + if float(kl_loss_weight) == 0.0: + loss = raw_kl_loss.detach() * 0.0 + else: + scaled_kl_loss = float(kl_loss_weight) * raw_kl_loss + world = dist.get_world_size(group) if dist.is_available() and dist.is_initialized() else 1 + loss = _value_divided_without_grad_scale(scaled_kl_loss, world) + + hidden_match_metric = 0.0 + hidden_match_raw_metric = 0.0 + hidden_match_weight_mean = local_hidden_weights.mean().item() if local_tokens else 0.0 + hidden_match_pos_metric = 0.0 + hidden_match_neg_metric = 0.0 + hidden_match_pos_raw_metric = 0.0 + hidden_match_neg_raw_metric = 0.0 + hidden_match_neg_minus_pos_raw = 0.0 + hidden_match_pos_weight_mean = 0.0 + hidden_match_neg_weight_mean = 0.0 + oprd_metric = 0.0 + oprd_raw_metric = 0.0 + oprd_num_layers = 0 + + hidden_match_coef = float(hidden_match_coef or 0.0) + local_valid = torch.ones(local_tokens, dtype=torch.float32, device=device) + use_oprd = ( + hidden_match_coef + and student_layer_hidden_states is not None + and (teacher_layer_hidden_states is not None or teacher_layer_fetcher is not None) + ) + if use_oprd and local_tokens: + if teacher_layer_fetcher is not None: + hidden_distance, oprd_num_layers = _oprd_hidden_distance_from_fetcher( + student_layer_hidden_states, + teacher_layer_fetcher=teacher_layer_fetcher, + expected_rows=local_tokens, + num_layers=teacher_layer_num_layers + if teacher_layer_num_layers is not None + else int(student_layer_hidden_states.shape[1]), + layer_chunk_size=oprd_layer_chunk_size, + ) + else: + hidden_distance, oprd_num_layers = _oprd_hidden_distance( + student_layer_hidden_states, + teacher_layer_hidden_states, + local_tokens, + layer_chunk_size=oprd_layer_chunk_size, + ) + weighted_hidden_match = hidden_distance * local_hidden_weights.to(hidden_distance.device) + hidden_match_loss = loss_reducer(weighted_hidden_match, local_valid.to(hidden_distance.device)) + loss = loss + hidden_match_coef * hidden_match_loss + oprd_metric = metric_reducer(weighted_hidden_match.detach(), local_valid.to(hidden_distance.device)).item() + oprd_raw_metric = metric_reducer(hidden_distance.detach(), local_valid.to(hidden_distance.device)).item() + hidden_match_metric = oprd_metric + hidden_match_raw_metric = oprd_raw_metric + elif hidden_match_coef and local_tokens: + if student_hidden_flat.shape[-1] != teacher_hidden_flat.shape[-1]: + raise ValueError( + "hidden_match requires matching hidden sizes, got " + f"student={student_hidden_flat.shape[-1]} teacher={teacher_hidden_flat.shape[-1]}" + ) + if str(hidden_match_mode).lower() == "mse": + hidden_distance = ((student_hidden_flat.float() - teacher_hidden_flat.float()) ** 2).mean(dim=-1) + else: + hidden_distance = 1.0 - F.cosine_similarity( + student_hidden_flat.float(), + teacher_hidden_flat.float(), + dim=-1, + eps=1e-6, + ) + hidden_weights_on_device = local_hidden_weights.to(hidden_distance.device) + weighted_hidden_match = hidden_distance * hidden_weights_on_device + hidden_match_loss = loss_reducer(weighted_hidden_match, local_valid.to(hidden_distance.device)) + loss = loss + hidden_match_coef * hidden_match_loss + hidden_match_metric = metric_reducer( + weighted_hidden_match.detach(), local_valid.to(hidden_distance.device) + ).item() + hidden_match_raw_metric = metric_reducer( + hidden_distance.detach(), local_valid.to(hidden_distance.device) + ).item() + pos_weights = torch.clamp(hidden_weights_on_device, min=0.0) + neg_weights = torch.clamp(-hidden_weights_on_device, min=0.0) + hidden_match_pos_metric = metric_reducer((hidden_distance * pos_weights).detach(), local_valid).item() + hidden_match_neg_metric = metric_reducer((hidden_distance * neg_weights).detach(), local_valid).item() + pos_weight_sum = pos_weights.sum() + neg_weight_sum = neg_weights.sum() + if pos_weight_sum.item() > 0: + hidden_match_pos_raw_metric = ((hidden_distance.detach() * pos_weights).sum() / pos_weight_sum).item() + if neg_weight_sum.item() > 0: + hidden_match_neg_raw_metric = ((hidden_distance.detach() * neg_weights).sum() / neg_weight_sum).item() + hidden_match_neg_minus_pos_raw = hidden_match_neg_raw_metric - hidden_match_pos_raw_metric + hidden_match_pos_weight_mean = pos_weights.mean().item() + hidden_match_neg_weight_mean = neg_weights.mean().item() + + if use_task_rewards: + loss = loss * float(distillation_loss_coef) + + detached_token_kl = token_kl.detach() + full_valid_count = max(float(detached_token_kl.numel()), 1.0) + metrics_kwargs: dict = { + "valid_tokens": local_tokens, + "opd_kl": detached_token_kl.sum().item() / full_valid_count, + "opd_weighted_kl": weighted_token_kl.detach().sum().item() / full_valid_count, + "opd_vocab_parallel_group_tokens": int(detached_token_kl.numel()), + "opd_vocab_parallel_kl_sum": detached_token_kl.sum().item(), + "opd_vocab_parallel_weighted_kl_sum": weighted_token_kl.detach().sum().item(), + "opd_hidden_match_loss": hidden_match_metric, + "opd_hidden_match_raw_loss": hidden_match_raw_metric, + "opd_hidden_match_weight_mean": hidden_match_weight_mean, + "opd_hidden_match_pos_loss": hidden_match_pos_metric, + "opd_hidden_match_neg_loss": hidden_match_neg_metric, + "opd_hidden_match_pos_raw_loss": hidden_match_pos_raw_metric, + "opd_hidden_match_neg_raw_loss": hidden_match_neg_raw_metric, + "opd_hidden_match_neg_minus_pos_raw": hidden_match_neg_minus_pos_raw, + "opd_hidden_match_pos_weight_mean": hidden_match_pos_weight_mean, + "opd_hidden_match_neg_weight_mean": hidden_match_neg_weight_mean, + "opd_teacher_weight_mean": full_token_weights.mean().item() if full_token_weights.numel() else 0.0, + "opd_loss_clamp_frac": clamp_frac, + "opd_oprd_loss": oprd_metric, + "opd_oprd_raw_loss": oprd_raw_metric, + "opd_oprd_num_layers": oprd_num_layers, + } + if detached_token_kl.numel(): + metrics_kwargs["opd_loss_min"] = detached_token_kl.min().item() + metrics_kwargs["opd_loss_max"] = detached_token_kl.max().item() + metrics_kwargs["opd_loss_abs_mean"] = detached_token_kl.abs().mean().item() + + metrics = OPDLossMetrics(**metrics_kwargs).to_dict() + if debug_token_outputs: + local_start, local_end = _gathered_local_bounds(local_tokens, device=token_kl.device, group=group) + metrics["_opd_debug_local_token_kl"] = detached_token_kl[local_start:local_end].detach() + metrics["_opd_debug_local_weighted_token_kl"] = weighted_token_kl.detach()[local_start:local_end] + metrics["_opd_debug_local_token_weight"] = full_token_weights.detach()[local_start:local_end] + + return LossOutput(loss=loss, metrics=metrics) + + +def opd_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_lm_head_weight: torch.Tensor, + teacher_weights: Optional[torch.Tensor] = None, + hidden_match_weights: Optional[torch.Tensor] = None, + ignore_index: int = -100, + num_chunks: int = 8, + lm_head_fp32: bool = False, + teacher_lm_head_fp32: bool = True, + kl_backend: str = "torch_compile", + vocab_chunk_size: int | None = 32768, + streaming_lowmem: bool = False, + return_per_token: bool = False, + normalization_denominator: Optional[torch.Tensor | int | float] = None, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, + loss_mode: str = LOSS_MODE_REVERSE_KL_FULL, + log_prob_min_clamp: Optional[float] = None, + loss_max_clamp: Optional[float] = None, + emit_full_vocab_diagnostics: bool = False, + use_policy_gradient: bool = False, + old_logprobs: Optional[torch.Tensor] = None, + clip_ratio_low: float = 0.2, + clip_ratio_high: float = 0.2, + use_task_rewards: bool = False, + distillation_loss_coef: float = 1.0, + hidden_match_coef: float = 0.0, + kl_loss_weight: float = 1.0, + hidden_match_mode: str = "cosine", + teacher_layer_hidden_states: Optional[torch.Tensor] = None, + teacher_layer_fetcher: Optional[Callable[[int, int], torch.Tensor]] = None, + teacher_layer_num_layers: Optional[int] = None, + oprd_layer_chunk_size: int = 4, + student_layer_hidden_states: Optional[torch.Tensor] = None, + diag_region_ids: Optional[torch.Tensor] = None, + diag_sample_ok: Optional[torch.Tensor] = None, +) -> LossOutput: + """Compute the OPD distillation loss for one micro-batch. + + Supported `loss_mode`s: + - "reverse_kl_full" (default): full-vocabulary KL(student||teacher), + i.e. Σ_v p_S(v)·(log p_S(v) - log p_T(v)). + - "forward_kl_full": full-vocabulary KL(teacher||student), + i.e. Σ_v p_T(v)·(log p_T(v) - log p_S(v)). The full-distribution + counterpart of VERL's truncated `forward_kl_topk` — possible here + because xorl materializes the full teacher distribution locally via the + cached teacher LM head. + - Single-sample estimators "kl"/"k1"/"abs"/"mse"/"k2"/"low_var_kl"/"k3" + (with optional "+" suffix for k2 straight-through gradient). These use + only the per-token (student_logprob, teacher_logprob) at the sampled + token; a cheaper memory point for RL-style updates. + + Set `emit_full_vocab_diagnostics=True` to also emit per-step teacher_entropy, + student_entropy, and top1_agreement. Adds a single extra full-vocab pass over + the logits, so opt-in. Only applies to the two full-vocab loss modes. + + Expected shapes: + hidden_states: [batch, seq, student_hidden_dim] + weight: [vocab_size, student_hidden_dim] + labels: [batch, seq], with ignore_index masking tokens out of the loss + teacher_hidden_states: [batch, seq, teacher_hidden_dim] + teacher_lm_head_weight: [vocab_size, teacher_hidden_dim] + teacher_weights: optional [batch, seq] per-token multipliers applied + to the KL term after KL computation and before final normalization. + hidden_match_weights: optional [batch, seq] per-token multipliers for + hidden-state matching. Defaults to teacher_weights when omitted. + teacher_layer_hidden_states / student_layer_hidden_states: optional + multi-layer OPRD tensors, each shaped [valid_tokens, L, d] and already + restricted to the valid (response-masked) positions IN THE SAME ORDER + as the internal valid_mask flatten (i.e. row i corresponds to the i-th + True entry of labels.reshape(-1) != ignore_index). When BOTH are given + the OPRD term — hidden_match_coef * mean_L( (1/d)||student_l - + teacher_l.detach()||^2 ), weighted by hidden_match_weights and reduced + by the existing reducer — REPLACES the single-layer hidden-match term to + avoid double-counting. When None, single-layer behavior is unchanged. + + `diag_region_ids` / `diag_sample_ok` are optional [batch, seq] int tensors + (labels-aligned, like teacher_weights) that drive metrics-only KL splits: + region 0=prompt / 1=buffer / 2=answer (-1 = unattributed) and per-sample + sampled-answer correctness 1/0 (-1 = unknown), broadcast over the sample's + positions by the client. They never touch the loss value. + + `log_prob_min_clamp` is a generic stability clamp on student log-probs. + Currently used by `forward_kl_full` (where p_T(v) · log p_S(v) terms can + blow up if the student puts ~0 mass on a teacher-preferred token); reverse + KL doesn't need it (its weights are p_S(v)). + `loss_max_clamp` symmetrically clamps the *unweighted* per-token loss before + teacher-weighting. + + Teacher tensors are detached by construction. Only the student hidden states + and student LM head receive gradients. + """ + if hidden_states.shape[:-1] != labels.shape: + raise ValueError(f"hidden_states shape {hidden_states.shape} is incompatible with labels {labels.shape}") + if teacher_hidden_states.shape[:-1] != labels.shape: + raise ValueError( + f"teacher_hidden_states shape {teacher_hidden_states.shape} is incompatible with labels {labels.shape}" + ) + if weight.shape[0] != teacher_lm_head_weight.shape[0]: + raise ValueError( + f"student vocab size ({weight.shape[0]}) must match teacher vocab size ({teacher_lm_head_weight.shape[0]})" + ) + if hidden_states.shape[-1] != weight.shape[-1]: + raise ValueError( + f"student hidden size ({hidden_states.shape[-1]}) must match student head width ({weight.shape[-1]})" + ) + if teacher_hidden_states.shape[-1] != teacher_lm_head_weight.shape[-1]: + raise ValueError( + "teacher hidden size " + f"({teacher_hidden_states.shape[-1]}) must match teacher head width ({teacher_lm_head_weight.shape[-1]})" + ) + + original_shape = labels.shape + labels_flat = labels.reshape(-1) + valid_mask = labels_flat != ignore_index + valid_count = valid_mask.sum() + + if valid_count.item() == 0: + loss = _zero_loss_with_graph(hidden_states, weight) + per_token_loss = ( + torch.zeros(original_shape, dtype=torch.float32, device=labels.device) if return_per_token else None + ) + return LossOutput(loss=loss, per_token_loss=per_token_loss, metrics=OPDLossMetrics(valid_tokens=0).to_dict()) + + student_hidden_flat = hidden_states.reshape(-1, hidden_states.size(-1))[valid_mask] + teacher_hidden_flat = teacher_hidden_states.reshape(-1, teacher_hidden_states.size(-1))[valid_mask].detach() + labels_valid = labels_flat[valid_mask] + token_weights = _as_flat_optional_weights(teacher_weights, valid_mask, torch.float32) + hidden_weights = ( + _as_flat_optional_weights(hidden_match_weights, valid_mask, torch.float32) + if hidden_match_weights is not None + else token_weights + ) + + default_scale = _denominator_tensor( + normalization_denominator, + fallback=valid_count, + device=hidden_states.device, + ) + if loss_reducer is None: + loss_reducer = TokenPartial(scale=default_scale) + if metric_reducer is None: + metric_reducer = TokenPartial(scale=default_scale) + + region_flat = None + if diag_region_ids is not None: + region_flat = diag_region_ids.reshape(-1).to(device=labels.device)[valid_mask] + sample_ok_flat = None + if diag_sample_ok is not None: + sample_ok_flat = diag_sample_ok.reshape(-1).to(device=labels.device)[valid_mask] + + backend = kl_backend.lower() + is_compile_backend = backend in {"torch_compile", "compile", "auto_chunker"} + is_streaming_backend = backend in {"streaming", "tilelang"} + + # Full-vocab diagnostics: the compile backend has fused *_with_diag kernels; + # the streaming backend covers the full-vocab KL modes via a separate no-grad + # streaming pass (same chunked-vocab shape, never materializes full logits). + streaming_diag_supported = is_streaming_backend and loss_mode in ( + LOSS_MODE_REVERSE_KL_FULL, + LOSS_MODE_FORWARD_KL_FULL, + ) + diag_enabled = bool(emit_full_vocab_diagnostics) and (is_compile_backend or streaming_diag_supported) + if loss_mode == LOSS_MODE_FORWARD_KL_FULL and not (is_compile_backend or is_streaming_backend): + raise ValueError( + "loss_mode='forward_kl_full' requires kl_backend in {'torch_compile', 'streaming', 'tilelang'}" + ) + + teacher_entropy_per_tok = None + student_entropy_per_tok = None + top1_agreement_per_tok = None + student_lp_at_label: Optional[torch.Tensor] = None + pg_clipfrac = 0.0 + pg_clipfrac_lower = 0.0 + ppo_kl = 0.0 + + if loss_mode == LOSS_MODE_REVERSE_KL_FULL: + if is_compile_backend: + if not torch.is_tensor(teacher_lm_head_weight): + raise ValueError("torch_compile OPD KL backend requires a materialized teacher LM head tensor") + if diag_enabled: + ( + token_kl, + teacher_entropy_per_tok, + student_entropy_per_tok, + top1_agreement_per_tok, + ) = compiled_reverse_kl_with_diag_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + num_chunks=num_chunks, + lm_head_fp32=lm_head_fp32, + teacher_lm_head_fp32=teacher_lm_head_fp32, + ) + else: + token_kl = compiled_reverse_kl_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + num_chunks=num_chunks, + lm_head_fp32=lm_head_fp32, + teacher_lm_head_fp32=teacher_lm_head_fp32, + ) + elif is_streaming_backend and streaming_lowmem: + # Memory-lean streaming KL for the fp32 lm-head regime: keep the + # student/teacher lm-head weights in their native (bf16) dtype and + # upcast each vocab chunk to fp32 inside the kernel, instead of + # holding two full fp32 weight copies + a full fp32 grad buffer. Only + # the cheap [N,H] hidden states are upcast here. Gradient-identical to + # the streaming path below (slicing commutes with the elementwise + # upcast; vocab chunks partition grad rows disjointly), but saves + # several GB of resident lm-head memory -- the AMDAHL-029..033 1-node + # blocker. The fp32 chunked matmul matches lm_head_fp32=True / + # teacher_lm_head_fp32=True (the OPD default); the diagnostics pass + # below is metrics-only and stays on the native-dtype weight. + if lm_head_fp32: + student_hidden_flat = student_hidden_flat.float() + if teacher_lm_head_fp32: + teacher_hidden_flat = teacher_hidden_flat.float() + compute_dtype = torch.float32 if (lm_head_fp32 or teacher_lm_head_fp32) else weight.dtype + token_kl = streaming_reverse_kl_lowmem_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + vocab_chunk_size=vocab_chunk_size, + compute_dtype=compute_dtype, + inplace_weight_grad=False, + ) + if diag_enabled: + ( + teacher_entropy_per_tok, + student_entropy_per_tok, + top1_agreement_per_tok, + ) = streaming_full_vocab_diagnostics( + student_hidden_states=student_hidden_flat.detach(), + student_weight=weight.detach(), + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + vocab_chunk_size=vocab_chunk_size, + ) + elif is_streaming_backend: + if lm_head_fp32: + student_hidden_flat = student_hidden_flat.float() + weight = weight.float() + if teacher_lm_head_fp32: + teacher_hidden_flat = teacher_hidden_flat.float() + if torch.is_tensor(teacher_lm_head_weight): + teacher_lm_head_weight = teacher_lm_head_weight.float() + token_kl = streaming_reverse_kl_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + vocab_chunk_size=vocab_chunk_size, + ) + if diag_enabled: + ( + teacher_entropy_per_tok, + student_entropy_per_tok, + top1_agreement_per_tok, + ) = streaming_full_vocab_diagnostics( + student_hidden_states=student_hidden_flat.detach(), + student_weight=weight.detach(), + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + vocab_chunk_size=vocab_chunk_size, + ) + else: + raise ValueError( + f"Unsupported OPD KL backend '{kl_backend}'. Expected 'torch_compile', 'streaming', or 'tilelang'." + ) + elif loss_mode == LOSS_MODE_FORWARD_KL_FULL: + if is_compile_backend: + if not torch.is_tensor(teacher_lm_head_weight): + raise ValueError("forward_kl_full requires a materialized teacher LM head tensor") + if diag_enabled: + ( + token_kl, + teacher_entropy_per_tok, + student_entropy_per_tok, + top1_agreement_per_tok, + ) = compiled_forward_kl_full_with_diag_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + log_prob_min_clamp=log_prob_min_clamp, + num_chunks=num_chunks, + lm_head_fp32=lm_head_fp32, + teacher_lm_head_fp32=teacher_lm_head_fp32, + ) + else: + token_kl = compiled_forward_kl_full_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + log_prob_min_clamp=log_prob_min_clamp, + num_chunks=num_chunks, + lm_head_fp32=lm_head_fp32, + teacher_lm_head_fp32=teacher_lm_head_fp32, + ) + elif is_streaming_backend: + # Streaming forward-KL: never materializes the full-vocab fp32 logits + # (~9.27 GiB at 151936 vocab), unblocking the large-microbatch OOM that + # forces forward_kl_full onto the compile backend. The streaming kernel + # accumulates the teacher-weighted log-ratio online over vocab chunks; + # it has no student-log-prob materialization, so `log_prob_min_clamp` + # (a clamp on log p_S) is unsupported here -- fail loud rather than + # silently ignore it. + if log_prob_min_clamp is not None: + raise ValueError( + "log_prob_min_clamp is not supported by the streaming forward_kl_full backend " + "(it never materializes student log-probs); use kl_backend='torch_compile' to clamp." + ) + if streaming_lowmem: + if lm_head_fp32: + student_hidden_flat = student_hidden_flat.float() + if teacher_lm_head_fp32: + teacher_hidden_flat = teacher_hidden_flat.float() + compute_dtype = torch.float32 if (lm_head_fp32 or teacher_lm_head_fp32) else weight.dtype + token_kl = streaming_forward_kl_lowmem_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + vocab_chunk_size=vocab_chunk_size, + compute_dtype=compute_dtype, + inplace_weight_grad=False, + ) + else: + if lm_head_fp32: + student_hidden_flat = student_hidden_flat.float() + weight = weight.float() + if teacher_lm_head_fp32: + teacher_hidden_flat = teacher_hidden_flat.float() + if torch.is_tensor(teacher_lm_head_weight): + teacher_lm_head_weight = teacher_lm_head_weight.float() + token_kl = streaming_forward_kl_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + vocab_chunk_size=vocab_chunk_size, + ) + if diag_enabled: + ( + teacher_entropy_per_tok, + student_entropy_per_tok, + top1_agreement_per_tok, + ) = streaming_full_vocab_diagnostics( + student_hidden_states=student_hidden_flat.detach(), + student_weight=weight.detach(), + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + vocab_chunk_size=vocab_chunk_size, + ) + else: + raise ValueError( + f"Unsupported OPD KL backend '{kl_backend}'. Expected 'torch_compile', 'streaming', or 'tilelang'." + ) + elif is_estimator_loss_mode(loss_mode): + if not torch.is_tensor(teacher_lm_head_weight): + raise ValueError(f"loss_mode='{loss_mode}' requires a materialized teacher LM head tensor") + student_lp, teacher_lp = compiled_sampled_token_logprobs_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + num_chunks=num_chunks, + lm_head_fp32=lm_head_fp32, + teacher_lm_head_fp32=teacher_lm_head_fp32, + ) + token_kl = _kl_penalty_estimator(student_lp, teacher_lp, loss_mode) + # `student_lp` is reused by PG mode below, no need for a second forward. + student_lp_at_label = student_lp + else: + raise ValueError( + f"Unsupported OPD loss_mode '{loss_mode}'. Expected 'reverse_kl_full', " + f"'forward_kl_full', or one of {_ESTIMATOR_MODES} (with optional '+' suffix)." + ) + + # Symmetric loss clamp before teacher-weighting. Matches + # verl/trainer/distillation/losses.py:228 (`distillation_losses.clamp(...)`). + clamp_frac = 0.0 + if loss_max_clamp is not None: + clamp_frac = (token_kl.detach().abs() >= float(loss_max_clamp)).float().mean().item() + token_kl = token_kl.clamp(min=-loss_max_clamp, max=loss_max_clamp) + # Post-clamp, pre-PG-surrogate per-token KL for the metrics-only region / + # correctness splits below (PG mode replaces token_kl with the surrogate). + kl_for_diag = token_kl.detach() + + # Policy-gradient OPD: treat -token_kl as advantage, apply PPO-clip on the + # student's sampled-token log-prob ratio. Mirrors VERL's `use_policy_gradient` + # path in verl/trainer/distillation/losses.py:259-281. + if use_policy_gradient: + if old_logprobs is None: + raise ValueError("use_policy_gradient=True requires old_logprobs to be provided") + if student_lp_at_label is None: + # Need student logprob at sampled-token position. Reuse the same + # compiled-sampled-logp helper used by estimator modes. + student_lp_at_label, _ = compiled_sampled_token_logprobs_function( + student_hidden_states=student_hidden_flat, + student_weight=weight, + teacher_hidden_states=teacher_hidden_flat, + teacher_weight=teacher_lm_head_weight, + labels=labels_valid, + ignore_index=ignore_index, + num_chunks=num_chunks, + lm_head_fp32=lm_head_fp32, + teacher_lm_head_fp32=teacher_lm_head_fp32, + ) + old_logprobs_flat = old_logprobs.reshape(-1).to(device=labels.device) + old_logprobs_valid = old_logprobs_flat[valid_mask].to(student_lp_at_label.dtype) + advantages = (-token_kl).detach() + log_ratio = student_lp_at_label - old_logprobs_valid + ratio = log_ratio.exp() + # min over the two surrogate losses, per PPO. NOTE: token_kl is replaced + # by the PG surrogate; further teacher-weighting / coef scaling below + # still applies element-wise. + pg_losses1 = -ratio * advantages + pg_losses2 = -ratio.clamp(1.0 - clip_ratio_low, 1.0 + clip_ratio_high) * advantages + pg_losses = torch.maximum(pg_losses1, pg_losses2) + is_clipped = (pg_losses2 > pg_losses1).to(token_kl.dtype) + pg_clipfrac = is_clipped.mean().item() if is_clipped.numel() > 0 else 0.0 + # Parity with VERL's `actor/ppo_kl` and `actor/pg_clipfrac_lower` + # (`core_algos.py:1365-1367`, surfaced under `distillation/` prefix in + # `distillation_loss`). `ppo_kl` = mean(old_logp - new_logp); `clipfrac_lower` + # = fraction of tokens where the LOWER clamp bound binds (1-clip_ratio_low) + # AND the advantage is positive (the unclipped surrogate would be smaller + # than the clipped one, so the clip is "hurting" gradient flow into a + # promising direction). + ppo_kl = (-log_ratio).mean().item() if log_ratio.numel() > 0 else 0.0 + is_clipped_lower = ((ratio < (1.0 - clip_ratio_low)) & (advantages < 0)).to(token_kl.dtype) + pg_clipfrac_lower = is_clipped_lower.mean().item() if is_clipped_lower.numel() > 0 else 0.0 + token_kl = pg_losses # downstream weighting / reduction unchanged + + token_weights_on_device = token_weights.to(token_kl.device) + weighted_token_kl = token_kl * token_weights_on_device + valid_ones = torch.ones_like(weighted_token_kl, dtype=torch.float32) + valid_count_float = max(float(valid_count.item()), 1.0) + + # kl_loss_weight scales the logit-KL term. Set 0.0 to supervise ONLY on hidden + # states (loss = hidden_match_coef * hidden_match): in self-distillation the + # student/teacher share the LM head, so matching answer-position hiddens matches + # the logits via that head while forcing the buffer to set up the answer hidden. + if float(kl_loss_weight) == 0.0 and not use_policy_gradient: + # Hidden-only supervision: the KL term is scaled by 0, so it contributes no + # gradient -- yet `0.0 * reduce(kl)` still backprops zeros through the full + # vocab-parallel KL and its LM-head matmul (the single most expensive term; + # the only gradient path to the shared LM head here). Detach it so the + # backward skips that subgraph entirely. token_kl/weighted_token_kl remain + # computed above, so opd_kl and the KL diagnostics are unaffected -- only the + # (zero) gradient is elided. + loss = loss_reducer(weighted_token_kl, valid_ones).detach() * 0.0 + else: + loss = float(kl_loss_weight) * loss_reducer(weighted_token_kl, valid_ones) + hidden_match_metric = 0.0 + hidden_match_raw_metric = 0.0 + hidden_match_weight_mean = 0.0 + hidden_match_pos_metric = 0.0 + hidden_match_neg_metric = 0.0 + hidden_match_pos_raw_metric = 0.0 + hidden_match_neg_raw_metric = 0.0 + hidden_match_neg_minus_pos_raw = 0.0 + hidden_match_pos_weight_mean = 0.0 + hidden_match_neg_weight_mean = 0.0 + weighted_hidden_match = None + oprd_metric = 0.0 + oprd_raw_metric = 0.0 + oprd_num_layers = 0 + hidden_match_coef = float(hidden_match_coef or 0.0) + # Multi-layer OPRD: when BOTH per-layer tensors are present, this REPLACES the + # single-layer hidden term below (the `else` branch) so the same coefficient + # isn't applied twice. Layer tensors arrive [valid_tokens, L, d] already + # restricted to the valid positions in valid_mask order. + use_oprd = ( + hidden_match_coef + and student_layer_hidden_states is not None + and (teacher_layer_hidden_states is not None or teacher_layer_fetcher is not None) + ) + if use_oprd: + if teacher_layer_fetcher is not None: + hidden_distance, oprd_num_layers = _oprd_hidden_distance_from_fetcher( + student_layer_hidden_states, + teacher_layer_fetcher=teacher_layer_fetcher, + expected_rows=int(student_hidden_flat.shape[0]), + num_layers=teacher_layer_num_layers + if teacher_layer_num_layers is not None + else int(student_layer_hidden_states.shape[1]), + layer_chunk_size=oprd_layer_chunk_size, + ) + else: + hidden_distance, oprd_num_layers = _oprd_hidden_distance( + student_layer_hidden_states, + teacher_layer_hidden_states, + int(student_hidden_flat.shape[0]), + layer_chunk_size=oprd_layer_chunk_size, + ) + hidden_weights_on_device = hidden_weights.to(hidden_distance.device) + weighted_hidden_match = hidden_distance * hidden_weights_on_device + hidden_match_loss = loss_reducer(weighted_hidden_match, valid_ones) + loss = loss + hidden_match_coef * hidden_match_loss + oprd_metric = metric_reducer(weighted_hidden_match.detach(), valid_ones).item() + oprd_raw_metric = metric_reducer(hidden_distance.detach(), valid_ones).item() + hidden_match_metric = oprd_metric + hidden_match_raw_metric = oprd_raw_metric + hidden_match_weight_mean = hidden_weights.mean().item() + elif hidden_match_coef: + if student_hidden_flat.shape[-1] != teacher_hidden_flat.shape[-1]: + raise ValueError( + "hidden_match requires matching hidden sizes, got " + f"student={student_hidden_flat.shape[-1]} teacher={teacher_hidden_flat.shape[-1]}" + ) + if str(hidden_match_mode).lower() == "mse": + # Magnitude-aware: MSE->0 implies student hidden == teacher hidden, so in + # self-distillation (shared LM head) the logits match too. Cosine alone + # matches only direction (ignores magnitude) and decouples from generation + # (eval accuracy collapses while cosine-distance keeps falling). + hidden_distance = ((student_hidden_flat.float() - teacher_hidden_flat.float()) ** 2).mean(dim=-1) + else: + hidden_distance = 1.0 - F.cosine_similarity( + student_hidden_flat.float(), + teacher_hidden_flat.float(), + dim=-1, + eps=1e-6, + ) + hidden_weights_on_device = hidden_weights.to(hidden_distance.device) + weighted_hidden_match = hidden_distance * hidden_weights_on_device + hidden_match_loss = loss_reducer(weighted_hidden_match, valid_ones) + loss = loss + hidden_match_coef * hidden_match_loss + hidden_match_metric = metric_reducer(weighted_hidden_match.detach(), valid_ones).item() + hidden_match_raw_metric = metric_reducer(hidden_distance.detach(), valid_ones).item() + hidden_match_weight_mean = hidden_weights.mean().item() + pos_weights = torch.clamp(hidden_weights_on_device, min=0.0) + neg_weights = torch.clamp(-hidden_weights_on_device, min=0.0) + hidden_match_pos_metric = metric_reducer((hidden_distance * pos_weights).detach(), valid_ones).item() + hidden_match_neg_metric = metric_reducer((hidden_distance * neg_weights).detach(), valid_ones).item() + pos_weight_sum = pos_weights.sum() + neg_weight_sum = neg_weights.sum() + if pos_weight_sum.item() > 0: + hidden_match_pos_raw_metric = ((hidden_distance.detach() * pos_weights).sum() / pos_weight_sum).item() + if neg_weight_sum.item() > 0: + hidden_match_neg_raw_metric = ((hidden_distance.detach() * neg_weights).sum() / neg_weight_sum).item() + hidden_match_neg_minus_pos_raw = hidden_match_neg_raw_metric - hidden_match_pos_raw_metric + hidden_match_pos_weight_mean = pos_weights.mean().item() + hidden_match_neg_weight_mean = neg_weights.mean().item() + + # Task-reward mixing: when use_task_rewards=True the caller is mixing this + # distillation loss with a separate task-RL loss outside this function and + # provides the coefficient here. When False, coef is ignored (== VERL semantics: + # `distillation_loss_coef if use_task_rewards else 1.0`). + if use_task_rewards: + loss = loss * float(distillation_loss_coef) + + per_token_loss = None + if return_per_token: + per_token_flat = torch.zeros(labels_flat.shape, dtype=torch.float32, device=labels.device) + per_token_contrib = float(kl_loss_weight) * weighted_token_kl + if weighted_hidden_match is not None: + per_token_contrib = per_token_contrib + hidden_match_coef * weighted_hidden_match + per_token_flat[valid_mask] = per_token_contrib.detach().to(per_token_flat.device) + per_token_loss = per_token_flat.view(original_shape) + + detached_token_kl = token_kl.detach() + metrics_kwargs: dict = { + "valid_tokens": int(valid_count.item()), + "opd_kl": detached_token_kl.sum().item() / valid_count_float, + "opd_weighted_kl": metric_reducer(weighted_token_kl.detach(), valid_ones).item(), + "opd_hidden_match_loss": hidden_match_metric, + "opd_hidden_match_raw_loss": hidden_match_raw_metric, + "opd_hidden_match_weight_mean": hidden_match_weight_mean, + "opd_hidden_match_pos_loss": hidden_match_pos_metric, + "opd_hidden_match_neg_loss": hidden_match_neg_metric, + "opd_hidden_match_pos_raw_loss": hidden_match_pos_raw_metric, + "opd_hidden_match_neg_raw_loss": hidden_match_neg_raw_metric, + "opd_hidden_match_neg_minus_pos_raw": hidden_match_neg_minus_pos_raw, + "opd_hidden_match_pos_weight_mean": hidden_match_pos_weight_mean, + "opd_hidden_match_neg_weight_mean": hidden_match_neg_weight_mean, + "opd_teacher_weight_mean": token_weights.mean().item(), + } + # Distillation-loss range metrics (parity with VERL compute_distillation_loss_range). + metrics_kwargs["opd_loss_min"] = detached_token_kl.min().item() + metrics_kwargs["opd_loss_max"] = detached_token_kl.max().item() + metrics_kwargs["opd_loss_abs_mean"] = detached_token_kl.abs().mean().item() + metrics_kwargs["opd_loss_clamp_frac"] = clamp_frac + + # Region / sample-correctness KL splits (metrics only; see OPDLossMetrics for + # the per-valid normalization + exact-aggregation contract). + def _masked_split(mask: torch.Tensor) -> tuple[float, float]: + m = mask.to(kl_for_diag.dtype) + return ( + (kl_for_diag * m).sum().item() / valid_count_float, + m.sum().item() / valid_count_float, + ) + + if region_flat is not None: + for region_value, region_name in ((0, "prompt"), (1, "buffer"), (2, "answer")): + kl_per_valid, frac = _masked_split(region_flat == region_value) + metrics_kwargs[f"opd_kl_{region_name}_per_valid"] = kl_per_valid + metrics_kwargs[f"opd_frac_{region_name}"] = frac + if sample_ok_flat is not None: + answer_mask = region_flat == 2 + for ok_value, ok_name in ((1, "correct"), (0, "wrong")): + split_mask = answer_mask & (sample_ok_flat == ok_value) + kl_per_valid, frac = _masked_split(split_mask) + metrics_kwargs[f"opd_kl_answer_{ok_name}_per_valid"] = kl_per_valid + metrics_kwargs[f"opd_frac_answer_{ok_name}"] = frac + if student_entropy_per_tok is not None and teacher_entropy_per_tok is not None: + m = split_mask.to(kl_for_diag.dtype) + metrics_kwargs[f"opd_student_entropy_answer_{ok_name}_per_valid"] = ( + student_entropy_per_tok.detach() * m + ).sum().item() / valid_count_float + metrics_kwargs[f"opd_teacher_entropy_answer_{ok_name}_per_valid"] = ( + teacher_entropy_per_tok.detach() * m + ).sum().item() / valid_count_float + metrics_kwargs["opd_oprd_loss"] = oprd_metric + metrics_kwargs["opd_oprd_raw_loss"] = oprd_raw_metric + metrics_kwargs["opd_oprd_num_layers"] = oprd_num_layers + if is_estimator_loss_mode(loss_mode): + # k1 can be negative; mirror VERL's distillation/abs_loss metric. + metrics_kwargs["opd_abs_loss"] = detached_token_kl.abs().mean().item() + if use_policy_gradient: + metrics_kwargs["opd_pg_clipfrac"] = float(pg_clipfrac) + metrics_kwargs["opd_pg_clipfrac_lower"] = float(pg_clipfrac_lower) + metrics_kwargs["opd_ppo_kl"] = float(ppo_kl) + # All flattened tensors here are over already-filtered valid tokens, so + # every position counts as "response-masked valid" and mean = sum / valid_count. + if teacher_entropy_per_tok is not None and student_entropy_per_tok is not None: + metrics_kwargs["opd_teacher_entropy"] = teacher_entropy_per_tok.detach().mean().item() + metrics_kwargs["opd_student_entropy"] = student_entropy_per_tok.detach().mean().item() + if top1_agreement_per_tok is not None: + metrics_kwargs["opd_top1_agreement"] = top1_agreement_per_tok.detach().mean().item() + + metrics = OPDLossMetrics(**metrics_kwargs).to_dict() + + return LossOutput(loss=loss, per_token_loss=per_token_loss, metrics=metrics) diff --git a/src/xorl/objectives/opd_streaming_kl.py b/src/xorl/objectives/opd_streaming_kl.py new file mode 100644 index 00000000..44bafaf4 --- /dev/null +++ b/src/xorl/objectives/opd_streaming_kl.py @@ -0,0 +1,820 @@ +from __future__ import annotations + +import math + +import torch + + +def _iter_weight_chunks(teacher_weight, vocab_size: int, chunk_rows: int): + if hasattr(teacher_weight, "iter_device_chunks"): + yield from teacher_weight.iter_device_chunks(chunk_rows) + return + for start, end in _iter_ranges(vocab_size, chunk_rows): + yield start, end, teacher_weight[start:end] + + +def _chunk_size(vocab_size: int, requested: int) -> int: + if requested <= 0 or requested >= vocab_size: + return vocab_size + return requested + + +def _iter_ranges(vocab_size: int, requested_chunk_size: int): + chunk_size = _chunk_size(vocab_size, requested_chunk_size) + for start in range(0, vocab_size, chunk_size): + yield start, min(start + chunk_size, vocab_size) + + +def _accum_dtype(reference: torch.Tensor) -> torch.dtype: + """fp32 accumulation by default, but preserve fp64 inputs (e.g. gradcheck). + + The streaming kernels accumulate in fp32 -- in production the lm-head tensors + are already fp32, so the per-chunk `.float()` upcast is a no-op. When a caller + passes fp64 inputs (the strongest gradcheck regime), truncating to fp32 would + swamp the numerical Jacobian with rounding noise; honoring fp64 there is + strictly more accurate and changes nothing for the fp32 production path. + """ + return torch.float64 if reference.dtype == torch.float64 else torch.float32 + + +def _update_online_logsumexp( + running_max: torch.Tensor, + running_sumexp: torch.Tensor, + logits: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + chunk_max = logits.max(dim=-1, keepdim=True).values + new_max = torch.maximum(running_max, chunk_max) + prev_scale = torch.where( + torch.isfinite(running_max), + (running_max - new_max).exp(), + torch.zeros_like(running_sumexp), + ) + chunk_sumexp = (logits - new_max).exp().sum(dim=-1, keepdim=True) + return new_max, running_sumexp * prev_scale + chunk_sumexp + + +class _StreamingReverseKL(torch.autograd.Function): + """Exact KL(student || teacher) over vocab chunks. + + This is the TileLang-facing OPD path: it exposes the same execution shape a + native kernel will use (stream vocab blocks, save only per-token statistics, + recompute logits in backward) while keeping a pure PyTorch implementation as + the portable fallback. + """ + + @staticmethod + def forward( + ctx, + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + labels: torch.Tensor, + teacher_weight, + ignore_index: int, + vocab_chunk_size: int, + ) -> torch.Tensor: + teacher_shape = tuple(int(x) for x in teacher_weight.shape) + if student_weight.shape[0] != teacher_shape[0]: + raise ValueError( + f"student vocab size ({student_weight.shape[0]}) must match teacher vocab size ({teacher_shape[0]})" + ) + + vocab_size = int(student_weight.shape[0]) + token_count = int(student_hidden_states.shape[0]) + valid = labels != ignore_index + neg_inf = -float("inf") + + # Single fused pass over vocab chunks. Reverse-KL decomposes as + # KL = sum_v p_s(v)*(s_v - t_v) - logZ_s + logZ_t + # and sum_v p_s(v)*(s_v - t_v) = A / Z_s with + # A = sum_v exp(s_v - s_max) * (s_v - t_v), + # which is online-accumulable (flash-attention style) alongside the + # student/teacher log-sum-exp normalizers. This halves the forward GEMM + # vs the prior two-pass form (a logsumexp pass + a probability pass) + # while still never materializing the full logits. The backward + # (recomputes logits from the saved normalizers) is unchanged, so the + # gradients are identical up to fp32 summation order. + s_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=torch.float32) + t_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=torch.float32) + s_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=torch.float32) + t_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=torch.float32) + weighted_diff = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=torch.float32) + + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): + s_logits = (student_hidden_states @ student_weight[start:end].t()).float() + t_logits = (teacher_hidden_states @ t_weight.t()).float() + # Student: online log-sum-exp + online sum_v exp(s_v - s_max)*(s_v - t_v). + chunk_s_max = s_logits.max(dim=-1, keepdim=True).values + new_s_max = torch.maximum(s_max, chunk_s_max) + s_scale = torch.where( + torch.isfinite(s_max), + (s_max - new_s_max).exp(), + torch.zeros_like(s_sumexp), + ) + exp_s = (s_logits - new_s_max).exp() + s_sumexp = s_sumexp * s_scale + exp_s.sum(dim=-1, keepdim=True) + weighted_diff = weighted_diff * s_scale + (exp_s * (s_logits - t_logits)).sum(dim=-1, keepdim=True) + s_max = new_s_max + # Teacher: online log-sum-exp normalizer only. + t_max, t_sumexp = _update_online_logsumexp(t_max, t_sumexp, t_logits) + + s_logz = s_sumexp.log() + s_max + t_logz = t_sumexp.log() + t_max + kl = (weighted_diff / s_sumexp - s_logz + t_logz).squeeze(-1) + kl = kl * valid.to(kl.dtype) + ctx.save_for_backward( + student_hidden_states, + student_weight, + teacher_hidden_states, + labels, + s_logz, + t_logz, + kl, + ) + ctx.teacher_weight = teacher_weight + ctx.ignore_index = ignore_index + ctx.vocab_chunk_size = vocab_chunk_size + return kl + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + ( + student_hidden_states, + student_weight, + teacher_hidden_states, + labels, + s_logz, + t_logz, + kl, + ) = ctx.saved_tensors + teacher_weight = ctx.teacher_weight + + valid = (labels != ctx.ignore_index).to(dtype=torch.float32, device=grad_output.device) + scale = grad_output.to(dtype=torch.float32) * valid + vocab_size = int(student_weight.shape[0]) + + grad_hidden = None + if ctx.needs_input_grad[0]: + grad_hidden = torch.zeros_like(student_hidden_states, dtype=torch.float32) + + grad_weight = None + if ctx.needs_input_grad[1]: + grad_weight = torch.zeros_like(student_weight, dtype=torch.float32) + + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, ctx.vocab_chunk_size): + s_weight = student_weight[start:end] + s_logits = (student_hidden_states @ s_weight.t()).float() + t_logits = (teacher_hidden_states @ t_weight.t()).float() + s_log_probs = s_logits - s_logz + t_log_probs = t_logits - t_logz + s_probs = s_log_probs.exp() + + # d KL(p_s || p_t) / d student_logits_i = + # p_s_i * (log p_s_i - log p_t_i - KL) + grad_logits = s_probs * (s_log_probs - t_log_probs - kl.unsqueeze(1)) + grad_logits = grad_logits * scale.unsqueeze(1) + + if grad_hidden is not None: + grad_hidden = grad_hidden + grad_logits @ s_weight.float() + if grad_weight is not None: + grad_weight[start:end] = grad_logits.t() @ student_hidden_states.float() + + if grad_hidden is not None: + grad_hidden = grad_hidden.to(student_hidden_states.dtype) + if grad_weight is not None: + grad_weight = grad_weight.to(student_weight.dtype) + + if hasattr(teacher_weight, "clear_device_cache"): + teacher_weight.clear_device_cache() + + return grad_hidden, grad_weight, None, None, None, None, None + + +class _StreamingForwardKL(torch.autograd.Function): + """Exact forward KL(teacher || student) over vocab chunks. + + The forward-KL counterpart of `_StreamingReverseKL`: same streaming shape + (stream vocab blocks, save only per-token statistics, recompute logits in + backward), never materializing the full [tokens, vocab] logits. This unblocks + `loss_mode='forward_kl_full'` on the streaming backend, where the compile + backend OOMs on the materialized full-vocab fp32 logits. + + Forward KL = KL(p_T || p_S) = sum_v p_T(v)*(log p_T(v) - log p_S(v)). Writing + log p_S(v) = s_v - s_logz and log p_T(v) = t_v - t_logz it telescopes to a + streaming-friendly form that never takes the log of a near-zero student prob: + + KL = s_logz - t_logz + (sum_v exp(t_v - t_max)*(t_v - s_v)) / t_sumexp + + The weighted term sum_v exp(t_v - t_max)*(t_v - s_v) is accumulated online + (flash-attention style) against the *teacher* running max/sumexp -- mirroring + `_StreamingReverseKL` but weighting by the teacher exp and using (t_v - s_v) + rather than the student exp and (s_v - t_v). The student/teacher log-sum-exp + normalizers are accumulated in the same pass. + + Backward w.r.t. the student logits is exact and far simpler than reverse KL: + the p_T*log p_T term is constant in z_S, so + d/d z_S[k] (-sum_v p_T(v) log p_S(v)) = p_S(k) - p_T(k), + i.e. grad_logits = (p_S - p_T)*grad_output*valid. No per-token KL term and no + p_S weighting of log-ratios. The teacher is detached, so no teacher grads. + """ + + @staticmethod + def forward( + ctx, + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + labels: torch.Tensor, + teacher_weight, + ignore_index: int, + vocab_chunk_size: int, + ) -> torch.Tensor: + teacher_shape = tuple(int(x) for x in teacher_weight.shape) + if student_weight.shape[0] != teacher_shape[0]: + raise ValueError( + f"student vocab size ({student_weight.shape[0]}) must match teacher vocab size ({teacher_shape[0]})" + ) + + vocab_size = int(student_weight.shape[0]) + token_count = int(student_hidden_states.shape[0]) + valid = labels != ignore_index + neg_inf = -float("inf") + acc = _accum_dtype(student_hidden_states) + + # Single fused pass over vocab chunks. Forward-KL decomposes as + # KL = s_logz - t_logz + sum_v p_T(v)*(t_v - s_v) + # and sum_v p_T(v)*(t_v - s_v) = B / Z_t with + # B = sum_v exp(t_v - t_max) * (t_v - s_v), + # which is online-accumulable (flash-attention style) alongside the + # student/teacher log-sum-exp normalizers. The teacher running max/sumexp + # weight the difference (mirror image of reverse KL's student-weighted A). + s_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=acc) + t_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=acc) + s_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=acc) + t_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=acc) + weighted_diff = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=acc) + + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): + s_logits = (student_hidden_states @ student_weight[start:end].t()).to(acc) + t_logits = (teacher_hidden_states @ t_weight.t()).to(acc) + # Teacher: online log-sum-exp + online sum_v exp(t_v - t_max)*(t_v - s_v). + chunk_t_max = t_logits.max(dim=-1, keepdim=True).values + new_t_max = torch.maximum(t_max, chunk_t_max) + t_scale = torch.where( + torch.isfinite(t_max), + (t_max - new_t_max).exp(), + torch.zeros_like(t_sumexp), + ) + exp_t = (t_logits - new_t_max).exp() + t_sumexp = t_sumexp * t_scale + exp_t.sum(dim=-1, keepdim=True) + weighted_diff = weighted_diff * t_scale + (exp_t * (t_logits - s_logits)).sum(dim=-1, keepdim=True) + t_max = new_t_max + # Student: online log-sum-exp normalizer only. + s_max, s_sumexp = _update_online_logsumexp(s_max, s_sumexp, s_logits) + + s_logz = s_sumexp.log() + s_max + t_logz = t_sumexp.log() + t_max + kl = (s_logz - t_logz + weighted_diff / t_sumexp).squeeze(-1) + kl = kl * valid.to(kl.dtype) + # Backward only needs s_logz/t_logz to recompute p_S and p_T per chunk; the + # forward-KL gradient (p_S - p_T) does not depend on the per-token KL value. + ctx.save_for_backward( + student_hidden_states, + student_weight, + teacher_hidden_states, + labels, + s_logz, + t_logz, + ) + ctx.teacher_weight = teacher_weight + ctx.ignore_index = ignore_index + ctx.vocab_chunk_size = vocab_chunk_size + return kl + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + ( + student_hidden_states, + student_weight, + teacher_hidden_states, + labels, + s_logz, + t_logz, + ) = ctx.saved_tensors + teacher_weight = ctx.teacher_weight + acc = _accum_dtype(student_hidden_states) + + valid = (labels != ctx.ignore_index).to(dtype=acc, device=grad_output.device) + scale = grad_output.to(dtype=acc) * valid + vocab_size = int(student_weight.shape[0]) + + grad_hidden = None + if ctx.needs_input_grad[0]: + grad_hidden = torch.zeros_like(student_hidden_states, dtype=acc) + + grad_weight = None + if ctx.needs_input_grad[1]: + grad_weight = torch.zeros_like(student_weight, dtype=acc) + + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, ctx.vocab_chunk_size): + s_weight = student_weight[start:end] + s_logits = (student_hidden_states @ s_weight.t()).to(acc) + t_logits = (teacher_hidden_states @ t_weight.t()).to(acc) + s_probs = (s_logits - s_logz).exp() + t_probs = (t_logits - t_logz).exp() + + # d KL(p_t || p_s) / d student_logits_k = p_s_k - p_t_k. + grad_logits = (s_probs - t_probs) * scale.unsqueeze(1) + + if grad_hidden is not None: + grad_hidden = grad_hidden + grad_logits @ s_weight.to(acc) + if grad_weight is not None: + grad_weight[start:end] = grad_logits.t() @ student_hidden_states.to(acc) + + if grad_hidden is not None: + grad_hidden = grad_hidden.to(student_hidden_states.dtype) + if grad_weight is not None: + grad_weight = grad_weight.to(student_weight.dtype) + + if hasattr(teacher_weight, "clear_device_cache"): + teacher_weight.clear_device_cache() + + return grad_hidden, grad_weight, None, None, None, None, None + + +def streaming_forward_kl_function( + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + vocab_chunk_size: int | None = 32768, +) -> torch.Tensor: + """Compute per-token forward KL(teacher||student) without materializing full-vocab logits.""" + if vocab_chunk_size is None: + vocab_chunk_size = 32768 + elif vocab_chunk_size <= 0: + vocab_chunk_size = int(student_weight.shape[0]) + if not math.isfinite(float(vocab_chunk_size)): + raise ValueError(f"Invalid vocab_chunk_size={vocab_chunk_size}") + return _StreamingForwardKL.apply( + student_hidden_states, + student_weight, + teacher_hidden_states.detach(), + labels, + teacher_weight.detach() if torch.is_tensor(teacher_weight) else teacher_weight, + int(ignore_index), + int(vocab_chunk_size), + ) + + +@torch.no_grad() +def streaming_full_vocab_diagnostics( + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_weight, + vocab_chunk_size: int | None = 32768, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Per-token (teacher_entropy, student_entropy, top1_agreement) diagnostics. + + Streams vocab chunks exactly like `_StreamingReverseKL` (one logsumexp pass, + one probability pass) so the full logits are never materialized. No-grad and + opt-in: this is the streaming-backend counterpart of the compile backend's + `*_with_diag` kernels, costing one extra full-vocab pass per micro-batch. + """ + if vocab_chunk_size is None: + vocab_chunk_size = 32768 + elif vocab_chunk_size <= 0: + vocab_chunk_size = int(student_weight.shape[0]) + vocab_size = int(student_weight.shape[0]) + token_count = int(student_hidden_states.shape[0]) + device = student_hidden_states.device + neg_inf = -float("inf") + + s_max = torch.full((token_count, 1), neg_inf, device=device, dtype=torch.float32) + t_max = torch.full((token_count, 1), neg_inf, device=device, dtype=torch.float32) + s_sumexp = torch.zeros((token_count, 1), device=device, dtype=torch.float32) + t_sumexp = torch.zeros((token_count, 1), device=device, dtype=torch.float32) + s_best_idx = torch.zeros(token_count, device=device, dtype=torch.long) + t_best_idx = torch.zeros(token_count, device=device, dtype=torch.long) + s_best_val = torch.full((token_count,), neg_inf, device=device, dtype=torch.float32) + t_best_val = torch.full((token_count,), neg_inf, device=device, dtype=torch.float32) + + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): + s_weight = student_weight[start:end].to(student_hidden_states.dtype) + t_weight = t_weight.to(teacher_hidden_states.dtype) + s_logits = (student_hidden_states @ s_weight.t()).float() + t_logits = (teacher_hidden_states @ t_weight.t()).float() + s_max, s_sumexp = _update_online_logsumexp(s_max, s_sumexp, s_logits) + t_max, t_sumexp = _update_online_logsumexp(t_max, t_sumexp, t_logits) + s_chunk_val, s_chunk_idx = s_logits.max(dim=-1) + t_chunk_val, t_chunk_idx = t_logits.max(dim=-1) + s_better = s_chunk_val > s_best_val + t_better = t_chunk_val > t_best_val + s_best_idx = torch.where(s_better, s_chunk_idx + start, s_best_idx) + s_best_val = torch.maximum(s_best_val, s_chunk_val) + t_best_idx = torch.where(t_better, t_chunk_idx + start, t_best_idx) + t_best_val = torch.maximum(t_best_val, t_chunk_val) + + s_logz = s_sumexp.log() + s_max + t_logz = t_sumexp.log() + t_max + s_plogp = torch.zeros(token_count, device=device, dtype=torch.float32) + t_plogp = torch.zeros(token_count, device=device, dtype=torch.float32) + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): + s_weight = student_weight[start:end].to(student_hidden_states.dtype) + t_weight = t_weight.to(teacher_hidden_states.dtype) + s_logits = (student_hidden_states @ s_weight.t()).float() + t_logits = (teacher_hidden_states @ t_weight.t()).float() + s_log_probs = s_logits - s_logz + t_log_probs = t_logits - t_logz + s_plogp = s_plogp + (s_log_probs.exp() * s_log_probs).sum(dim=-1) + t_plogp = t_plogp + (t_log_probs.exp() * t_log_probs).sum(dim=-1) + + if hasattr(teacher_weight, "clear_device_cache"): + teacher_weight.clear_device_cache() + + teacher_entropy = -t_plogp + student_entropy = -s_plogp + top1_agreement = (s_best_idx == t_best_idx).float() + return teacher_entropy, student_entropy, top1_agreement + + +def streaming_reverse_kl_function( + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + vocab_chunk_size: int | None = 32768, +) -> torch.Tensor: + """Compute per-token reverse KL without materializing full-vocab logits.""" + if vocab_chunk_size is None: + vocab_chunk_size = 32768 + elif vocab_chunk_size <= 0: + vocab_chunk_size = int(student_weight.shape[0]) + if not math.isfinite(float(vocab_chunk_size)): + raise ValueError(f"Invalid vocab_chunk_size={vocab_chunk_size}") + return _StreamingReverseKL.apply( + student_hidden_states, + student_weight, + teacher_hidden_states.detach(), + labels, + teacher_weight.detach() if torch.is_tensor(teacher_weight) else teacher_weight, + int(ignore_index), + int(vocab_chunk_size), + ) + + +class _StreamingReverseKLLowMem(torch.autograd.Function): + """Memory-lean reverse KL(student||teacher), bit-exact with _StreamingReverseKL + under an fp32 lm-head, but without holding full fp32 copies of the lm-head + weights or a second full [V,H] grad buffer. + + Two memory levers, both gradient-identical to the current OPD fp32 path: + + 1. Per-chunk fp32 upcast. The student/teacher lm-head weights stay in their + native (bf16) dtype; each [chunk, H] vocab slice is upcast to + `compute_dtype` (fp32) inside the loop, right before the matmul. Slicing + commutes with the elementwise upcast, so the per-chunk fp32 matmul is + identical to multiplying by a whole pre-upcast fp32 weight -- but the + two ~2 GB fp32 weight copies (`weight.float()`) are never resident. + + 2. In-place weight grad. The weight gradient is streamed straight into the + leaf `student_weight.grad` (created lazily, in the weight's native + dtype), one vocab chunk at a time, instead of allocating a full [V,H] + buffer that autograd then adds into `.grad`. Vocab chunks partition the + grad rows disjointly, so this is exact; it removes the full extra buffer + (and its grad-accumulation doubling). `student_weight` must be the leaf + parameter when `inplace_weight_grad=True`. + + This is the AMDAHL-029..033 1-node unblock: it removes ~2-4 GB of resident + lm-head memory that tipped the streaming-KL backward over on one node. + """ + + @staticmethod + def forward( + ctx, + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + labels: torch.Tensor, + teacher_weight, + ignore_index: int, + vocab_chunk_size: int, + compute_dtype: torch.dtype, + inplace_weight_grad: bool, + ) -> torch.Tensor: + teacher_shape = tuple(int(x) for x in teacher_weight.shape) + if student_weight.shape[0] != teacher_shape[0]: + raise ValueError( + f"student vocab size ({student_weight.shape[0]}) must match teacher vocab size ({teacher_shape[0]})" + ) + vocab_size = int(student_weight.shape[0]) + token_count = int(student_hidden_states.shape[0]) + valid = labels != ignore_index + neg_inf = -float("inf") + dev = student_hidden_states.device + sh = student_hidden_states.to(compute_dtype) + th = teacher_hidden_states.to(compute_dtype) + + s_max = torch.full((token_count, 1), neg_inf, device=dev, dtype=torch.float32) + t_max = torch.full((token_count, 1), neg_inf, device=dev, dtype=torch.float32) + s_sumexp = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) + t_sumexp = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): + s_logits = (sh @ student_weight[start:end].to(compute_dtype).t()).float() + t_logits = (th @ t_weight.to(compute_dtype).t()).float() + s_max, s_sumexp = _update_online_logsumexp(s_max, s_sumexp, s_logits) + t_max, t_sumexp = _update_online_logsumexp(t_max, t_sumexp, t_logits) + + s_logz = s_sumexp.log() + s_max + t_logz = t_sumexp.log() + t_max + kl = torch.zeros(token_count, device=dev, dtype=torch.float32) + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): + s_logits = (sh @ student_weight[start:end].to(compute_dtype).t()).float() + t_logits = (th @ t_weight.to(compute_dtype).t()).float() + s_log_probs = s_logits - s_logz + t_log_probs = t_logits - t_logz + s_probs = s_log_probs.exp() + kl = kl + (s_probs * (s_log_probs - t_log_probs)).sum(dim=-1) + + kl = kl * valid.to(kl.dtype) + ctx.save_for_backward(student_hidden_states, student_weight, teacher_hidden_states, labels, s_logz, t_logz, kl) + ctx.teacher_weight = teacher_weight + ctx.ignore_index = ignore_index + ctx.vocab_chunk_size = vocab_chunk_size + ctx.compute_dtype = compute_dtype + ctx.inplace_weight_grad = inplace_weight_grad + ctx.student_weight_ref = student_weight + return kl + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + (student_hidden_states, student_weight, teacher_hidden_states, labels, s_logz, t_logz, kl) = ctx.saved_tensors + teacher_weight = ctx.teacher_weight + cdt = ctx.compute_dtype + valid = (labels != ctx.ignore_index).to(dtype=torch.float32, device=grad_output.device) + scale = grad_output.to(dtype=torch.float32) * valid + vocab_size = int(student_weight.shape[0]) + sh = student_hidden_states.to(cdt) + th = teacher_hidden_states.to(cdt) + + grad_hidden = None + if ctx.needs_input_grad[0]: + grad_hidden = torch.zeros_like(student_hidden_states, dtype=torch.float32) + + # Weight grad: either stream in place into the leaf .grad (no second full + # buffer), or build one buffer in the weight's native dtype and return it. + accumulate_weight = ctx.needs_input_grad[1] or ctx.inplace_weight_grad + inplace = ctx.inplace_weight_grad + grad_weight = None + if accumulate_weight: + if inplace: + wparam = ctx.student_weight_ref + if wparam.grad is None: + wparam.grad = torch.zeros_like(wparam) + else: + grad_weight = torch.zeros_like(student_weight) + + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, ctx.vocab_chunk_size): + s_weight = student_weight[start:end].to(cdt) + s_logits = (sh @ s_weight.t()).float() + t_logits = (th @ t_weight.to(cdt).t()).float() + s_log_probs = s_logits - s_logz + t_log_probs = t_logits - t_logz + s_probs = s_log_probs.exp() + grad_logits = s_probs * (s_log_probs - t_log_probs - kl.unsqueeze(1)) + grad_logits = grad_logits * scale.unsqueeze(1) + if grad_hidden is not None: + grad_hidden = grad_hidden + grad_logits @ s_weight + if accumulate_weight: + chunk = (grad_logits.t() @ sh).to(student_weight.dtype) + if inplace: + ctx.student_weight_ref.grad[start:end].add_(chunk) + else: + grad_weight[start:end] = chunk + + if grad_hidden is not None: + grad_hidden = grad_hidden.to(student_hidden_states.dtype) + if hasattr(teacher_weight, "clear_device_cache"): + teacher_weight.clear_device_cache() + + return grad_hidden, grad_weight, None, None, None, None, None, None, None + + +def streaming_reverse_kl_lowmem_function( + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + vocab_chunk_size: int | None = 32768, + compute_dtype: torch.dtype = torch.float32, + inplace_weight_grad: bool = False, +) -> torch.Tensor: + """Memory-lean per-token reverse KL: native-dtype lm-head weights upcast per + vocab chunk to `compute_dtype`, gradient-identical to the fp32 path. + + Pass the lm-head weights in their stored (bf16) dtype -- do NOT pre-cast them + to fp32. `student_hidden_states` may be fp32 (cheap; [N,H]). Set + `inplace_weight_grad=True` only when `student_weight` is the leaf parameter + whose `.grad` the optimizer reads (saves the second full [V,H] buffer). + """ + if vocab_chunk_size is None: + vocab_chunk_size = 32768 + elif vocab_chunk_size <= 0: + vocab_chunk_size = int(student_weight.shape[0]) + if not math.isfinite(float(vocab_chunk_size)): + raise ValueError(f"Invalid vocab_chunk_size={vocab_chunk_size}") + return _StreamingReverseKLLowMem.apply( + student_hidden_states, + student_weight, + teacher_hidden_states.detach(), + labels, + teacher_weight.detach() if torch.is_tensor(teacher_weight) else teacher_weight, + int(ignore_index), + int(vocab_chunk_size), + compute_dtype, + bool(inplace_weight_grad), + ) + + +class _StreamingForwardKLLowMem(torch.autograd.Function): + """Memory-lean forward KL(teacher||student), bit-exact with _StreamingForwardKL + under an fp32 lm-head, but without holding full fp32 copies of the lm-head + weights or a second full [V,H] grad buffer. + + The forward-KL counterpart of `_StreamingReverseKLLowMem`, with the identical + two memory levers (both gradient-identical to the fp32 forward-KL path): + + 1. Per-chunk fp32 upcast. The student/teacher lm-head weights stay in their + native (bf16) dtype; each [chunk, H] vocab slice is upcast to + `compute_dtype` (fp32) inside the loop before the matmul. Slicing + commutes with the elementwise upcast, so this matches multiplying by a + whole pre-upcast fp32 weight without ever holding the fp32 copies. + + 2. In-place weight grad. The weight gradient is streamed straight into the + leaf `student_weight.grad` one vocab chunk at a time (chunks partition + the grad rows disjointly, so it is exact), avoiding a full [V,H] buffer. + `student_weight` must be the leaf parameter when `inplace_weight_grad`. + + Forward uses the same telescoped, never-log-a-near-zero-prob form as + `_StreamingForwardKL`: + KL = s_logz - t_logz + (sum_v exp(t_v - t_max)*(t_v - s_v)) / t_sumexp. + Backward is exact and simple: grad_logits = (p_S - p_T)*grad_output*valid. + """ + + @staticmethod + def forward( + ctx, + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + labels: torch.Tensor, + teacher_weight, + ignore_index: int, + vocab_chunk_size: int, + compute_dtype: torch.dtype, + inplace_weight_grad: bool, + ) -> torch.Tensor: + teacher_shape = tuple(int(x) for x in teacher_weight.shape) + if student_weight.shape[0] != teacher_shape[0]: + raise ValueError( + f"student vocab size ({student_weight.shape[0]}) must match teacher vocab size ({teacher_shape[0]})" + ) + vocab_size = int(student_weight.shape[0]) + token_count = int(student_hidden_states.shape[0]) + valid = labels != ignore_index + neg_inf = -float("inf") + dev = student_hidden_states.device + sh = student_hidden_states.to(compute_dtype) + th = teacher_hidden_states.to(compute_dtype) + + s_max = torch.full((token_count, 1), neg_inf, device=dev, dtype=torch.float32) + t_max = torch.full((token_count, 1), neg_inf, device=dev, dtype=torch.float32) + s_sumexp = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) + t_sumexp = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) + weighted_diff = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) + # Single fused pass (same telescoped form as `_StreamingForwardKL`): online + # student/teacher log-sum-exp plus the online teacher-weighted (t_v - s_v) + # accumulation, so the KL is computed in one sweep over vocab chunks rather + # than a second full-vocab GEMM pass. KL = s_logz - t_logz + weighted_diff / t_sumexp. + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): + s_logits = (sh @ student_weight[start:end].to(compute_dtype).t()).float() + t_logits = (th @ t_weight.to(compute_dtype).t()).float() + chunk_t_max = t_logits.max(dim=-1, keepdim=True).values + new_t_max = torch.maximum(t_max, chunk_t_max) + t_scale = torch.where( + torch.isfinite(t_max), + (t_max - new_t_max).exp(), + torch.zeros_like(t_sumexp), + ) + exp_t = (t_logits - new_t_max).exp() + t_sumexp = t_sumexp * t_scale + exp_t.sum(dim=-1, keepdim=True) + weighted_diff = weighted_diff * t_scale + (exp_t * (t_logits - s_logits)).sum(dim=-1, keepdim=True) + t_max = new_t_max + s_max, s_sumexp = _update_online_logsumexp(s_max, s_sumexp, s_logits) + + s_logz = s_sumexp.log() + s_max + t_logz = t_sumexp.log() + t_max + kl = (s_logz - t_logz + weighted_diff / t_sumexp).squeeze(-1) + kl = kl * valid.to(kl.dtype) + ctx.save_for_backward(student_hidden_states, student_weight, teacher_hidden_states, labels, s_logz, t_logz) + ctx.teacher_weight = teacher_weight + ctx.ignore_index = ignore_index + ctx.vocab_chunk_size = vocab_chunk_size + ctx.compute_dtype = compute_dtype + ctx.inplace_weight_grad = inplace_weight_grad + ctx.student_weight_ref = student_weight + return kl + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + (student_hidden_states, student_weight, teacher_hidden_states, labels, s_logz, t_logz) = ctx.saved_tensors + teacher_weight = ctx.teacher_weight + cdt = ctx.compute_dtype + valid = (labels != ctx.ignore_index).to(dtype=torch.float32, device=grad_output.device) + scale = grad_output.to(dtype=torch.float32) * valid + vocab_size = int(student_weight.shape[0]) + sh = student_hidden_states.to(cdt) + th = teacher_hidden_states.to(cdt) + + grad_hidden = None + if ctx.needs_input_grad[0]: + grad_hidden = torch.zeros_like(student_hidden_states, dtype=torch.float32) + + accumulate_weight = ctx.needs_input_grad[1] or ctx.inplace_weight_grad + inplace = ctx.inplace_weight_grad + grad_weight = None + if accumulate_weight: + if inplace: + wparam = ctx.student_weight_ref + if wparam.grad is None: + wparam.grad = torch.zeros_like(wparam) + else: + grad_weight = torch.zeros_like(student_weight) + + for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, ctx.vocab_chunk_size): + s_weight = student_weight[start:end].to(cdt) + s_logits = (sh @ s_weight.t()).float() + t_logits = (th @ t_weight.to(cdt).t()).float() + s_probs = (s_logits - s_logz).exp() + t_probs = (t_logits - t_logz).exp() + grad_logits = (s_probs - t_probs) * scale.unsqueeze(1) + if grad_hidden is not None: + grad_hidden = grad_hidden + grad_logits @ s_weight + if accumulate_weight: + chunk = (grad_logits.t() @ sh).to(student_weight.dtype) + if inplace: + ctx.student_weight_ref.grad[start:end].add_(chunk) + else: + grad_weight[start:end] = chunk + + if grad_hidden is not None: + grad_hidden = grad_hidden.to(student_hidden_states.dtype) + if hasattr(teacher_weight, "clear_device_cache"): + teacher_weight.clear_device_cache() + + return grad_hidden, grad_weight, None, None, None, None, None, None, None + + +def streaming_forward_kl_lowmem_function( + student_hidden_states: torch.Tensor, + student_weight: torch.Tensor, + teacher_hidden_states: torch.Tensor, + teacher_weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + vocab_chunk_size: int | None = 32768, + compute_dtype: torch.dtype = torch.float32, + inplace_weight_grad: bool = False, +) -> torch.Tensor: + """Memory-lean per-token forward KL: native-dtype lm-head weights upcast per + vocab chunk to `compute_dtype`, gradient-identical to the fp32 path. + + Pass the lm-head weights in their stored (bf16) dtype -- do NOT pre-cast them + to fp32. `student_hidden_states` may be fp32 (cheap; [N,H]). Set + `inplace_weight_grad=True` only when `student_weight` is the leaf parameter + whose `.grad` the optimizer reads (saves the second full [V,H] buffer). + """ + if vocab_chunk_size is None: + vocab_chunk_size = 32768 + elif vocab_chunk_size <= 0: + vocab_chunk_size = int(student_weight.shape[0]) + if not math.isfinite(float(vocab_chunk_size)): + raise ValueError(f"Invalid vocab_chunk_size={vocab_chunk_size}") + return _StreamingForwardKLLowMem.apply( + student_hidden_states, + student_weight, + teacher_hidden_states.detach(), + labels, + teacher_weight.detach() if torch.is_tensor(teacher_weight) else teacher_weight, + int(ignore_index), + int(vocab_chunk_size), + compute_dtype, + bool(inplace_weight_grad), + ) diff --git a/src/xorl/objectives/policy_loss.py b/src/xorl/objectives/policy_loss.py new file mode 100644 index 00000000..59ff3dc6 --- /dev/null +++ b/src/xorl/objectives/policy_loss.py @@ -0,0 +1,394 @@ +""" +Policy Loss with PPO Clipping and TIS Correction. + +This module provides the policy loss functions including: +- PPO-style clipped policy gradient loss +- Temporal Importance Sampling (TIS) correction +- Combined policy_loss_function +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, Optional, Tuple + +import torch +import torch.distributed as dist + +from xorl.objectives.loss_output import LossOutput +from xorl.objectives.reducers import Reducer, TokenPartial +from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.loss.per_token_ce import compute_per_token_ce + + +logger = logging.getLogger(__name__) + + +K3_DEBUG_THRESHOLDS = ( + ("1e_minus_6", 1e-6), + ("1e_minus_4", 1e-4), + ("1e_minus_3", 1e-3), + ("1e_minus_2", 1e-2), + ("1e_minus_1", 1e-1), + ("1", 1.0), +) + + +@torch.compile(dynamic=True) +def compute_ppo_loss( + ppo_kl: torch.Tensor, + advantages: torch.Tensor, + eps_clip: float = 0.2, + eps_clip_high: float = 0.2, + eps_clip_c: Optional[float] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + PPO-style clipped policy loss. + + Args: + ppo_kl: KL divergence tensor (old_log_probs - new_log_probs) + advantages: Per-token advantages + eps_clip: Lower clip ratio (default: 0.2) + eps_clip_high: Upper clip ratio (default: 0.2) + eps_clip_c: Dual-clip ratio for negative advantages (optional) + + Returns: + pg_losses: Clipped policy gradient losses + is_clipped: Per-token boolean mask of clipped tokens + ratio: Importance sampling ratio exp(-ppo_kl) + """ + ratio = (-ppo_kl).exp() + pg_losses1 = -ratio * advantages + pg_losses2 = -ratio.clamp(1 - eps_clip, 1 + eps_clip_high) * advantages + clip_pg_losses1 = torch.maximum(pg_losses1, pg_losses2) + is_clipped = torch.gt(pg_losses2, pg_losses1) + + # Optional dual-clip for negative advantages + if eps_clip_c is not None: + assert eps_clip_c > 1.0, f"eps_clip_c must be > 1.0, got {eps_clip_c}" + pg_losses3 = -eps_clip_c * advantages + clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1) + pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1) + # Also mark dual-clipped tokens + is_dual_clipped = (advantages < 0) & torch.lt(pg_losses3, clip_pg_losses1) + is_clipped = is_clipped | is_dual_clipped + else: + pg_losses = clip_pg_losses1 + + return pg_losses, is_clipped, ratio + + +def apply_tis_correction( + pg_loss: torch.Tensor, + train_log_probs: torch.Tensor, + rollout_log_probs: torch.Tensor, + valid_mask: torch.Tensor, + metric_reducer: Reducer, + tis_clip_low: float = 0.1, + tis_clip_high: float = 2.0, +) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + """ + Apply Temporal Importance Sampling (TIS) correction. + + TIS corrects for the distribution shift between rollout time and training time. + The TIS weight is: exp(train_log_probs - rollout_log_probs) + + Args: + pg_loss: Policy gradient loss tensor + train_log_probs: Log probabilities from current training step + rollout_log_probs: Log probabilities from rollout/inference + valid_mask: Mask for valid tokens + metric_reducer: Reducer applied to per-token mean metrics (tis_mean, + tis_clipfrac). min/max are local reductions and bypass it. + tis_clip_low: Lower bound for TIS clipping (default: 0.1) + tis_clip_high: Upper bound for TIS clipping (default: 2.0) + + Returns: + Tuple of (corrected_loss, metrics_dict) + """ + # Compute TIS weights: ratio of train vs rollout distributions + tis = torch.exp(train_log_probs - rollout_log_probs) + + # Clip TIS weights to prevent extreme values + tis_clipped = torch.clamp(tis, min=tis_clip_low, max=tis_clip_high) + + # Apply TIS correction to loss + corrected_loss = pg_loss * tis_clipped + + valid_mask_f = valid_mask.float() + tis_clipfrac_per_token = (tis_clipped != tis).float() + # ±inf identity on empty ranks lets cross-rank MIN/MAX-allreduce ignore empty contributors. + if valid_mask.any(): + tis_min = tis.masked_fill(~valid_mask, float("inf")).min() + tis_max = tis.masked_fill(~valid_mask, float("-inf")).max() + else: + tis_min = tis.new_tensor(float("inf")) + tis_max = tis.new_tensor(float("-inf")) + tis_metrics = { + "tis_mean": metric_reducer(tis, valid_mask_f), + "tis_min": tis_min, + "tis_max": tis_max, + "tis_clipfrac": metric_reducer(tis_clipfrac_per_token, valid_mask_f), + } + + return corrected_loss, tis_metrics + + +def policy_loss_function( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + old_logprobs: torch.Tensor, + advantages: torch.Tensor, + rollout_logprobs: Optional[torch.Tensor] = None, + ignore_index: int = -100, + eps_clip: float = 0.2, + eps_clip_high: float = 0.2, + eps_clip_c: Optional[float] = None, + tis_clip_low: float = 0.1, + tis_clip_high: float = 2.0, + use_tis: bool = False, + use_liger: bool = True, + num_chunks: int = 8, + ce_mode: str = "compiled", + compute_kl_stats: bool = False, + tp_group: Optional[dist.ProcessGroup] = None, + lm_head_fp32: bool = False, + icepop_beta: Optional[float] = None, + loss_reducer: Optional[Reducer] = None, + metric_reducer: Optional[Reducer] = None, + lm_head: Optional[torch.nn.Module] = None, + logprob_temperature: float = 1.0, + logprob_top_k: int | torch.Tensor = TOP_K_ALL, + logprob_top_p: float | torch.Tensor = 1.0, + logprob_min_p: float | torch.Tensor = 0.0, +) -> "LossOutput": + """ + Policy loss with PPO clipping, optional IcePop masking, and optional TIS correction. + + This implements the loss function which includes: + 1. PPO-style clipping on the importance sampling ratio + 2. Optional IcePop hard masking (GLM-5): zeros gradient for tokens where ratio is outside [1/β, β] + 3. Optional Temporal Importance Sampling (TIS) correction for off-policy data + + Supports multiple computation modes: + - "compiled": RECOMMENDED. torch.compile (1.6x speed, 16% memory) + - "eager": Simple F.cross_entropy baseline (may OOM at 32K) + + Args: + hidden_states: Model hidden states, shape (batch, seq_len, hidden_dim) + weight: LM head weight matrix, shape (vocab_size, hidden_dim) + labels: Target token IDs, shape (batch, seq_len). Already next-token aligned. + old_logprobs: Old policy log probabilities from sampling, shape (batch, seq_len) + advantages: Per-token advantages, shape (batch, seq_len) + rollout_logprobs: Optional rollout log probabilities for TIS correction, shape (batch, seq_len) + ignore_index: Index to ignore in loss computation (default: -100) + eps_clip: Lower clip ratio for PPO (default: 0.2) + eps_clip_high: Upper clip ratio for PPO (default: 0.2) + eps_clip_c: Dual-clip ratio for negative advantages (optional) + tis_clip_low: Lower bound for TIS clipping (default: 0.1) + tis_clip_high: Upper bound for TIS clipping (default: 2.0) + use_tis: Whether to apply TIS correction (default: False) + use_liger: Kept for API compatibility (ignored) + num_chunks: Number of chunks for auto_chunker (default: 8). Only used when ce_mode="compiled". + ce_mode: Cross-entropy mode - "compiled" (recommended) or "eager" + tp_group: TP process group for vocab-parallel cross-entropy (default: None) + compute_kl_stats: If True, compute and return full KL statistics in metrics dict + (kl_sample_train_k3, entropy_sample, ratio stats). + If False (default), only return valid_tokens and pg_clipfrac. + loss_reducer: Reduces per-token loss to a scalar partial share. None => + ``TokenPartial(scale=valid_mask.sum())`` (legacy local token-mean; does + not compose across micro-batches/ranks). Pass a shared global-scale + reducer to make summed partial shares recover the global loss. + metric_reducer: Reduces per-token /mean metrics (pg_clipfrac, icepop_maskfrac, + tis_mean, tis_clipfrac, kl_sample_train_k3, entropy_sample, ratio_mean) + the same way. ratio_min/ratio_max/tis_min/tis_max stay local scalars. + logprob_temperature: Temperature applied to trainer logits before + selected-token logprob calculation. ``1.0`` is raw policy logprobs; + setting this to the rollout temperature yields Slime-style behavior + logprob semantics for PPO ratios and K3 stats. + + Returns: + LossOutput with loss, per_token_logprobs (new logprobs), and metrics. + """ + + # Store original shape + original_shape = labels.shape + + # Flatten tensors + labels_flat = labels.view(-1) + hidden_states_flat = hidden_states.view(-1, hidden_states.size(-1)) + old_logprobs_flat = old_logprobs.view(-1) + advantages_flat = advantages.view(-1) + + # Create mask for valid tokens (use labels != ignore_index) + valid_mask = labels_flat != ignore_index + valid_mask_f = valid_mask.float() + valid_count = valid_mask.sum() + + if loss_reducer is None: + loss_reducer = TokenPartial(scale=valid_count.float()) + if metric_reducer is None: + metric_reducer = TokenPartial(scale=valid_count.float()) + + # Compute cross-entropy (supports vocab-parallel TP via tp_group) + per_token_ce = compute_per_token_ce( + hidden_states_flat, + weight, + labels_flat, + ignore_index, + ce_mode, + num_chunks, + tp_group=tp_group, + lm_head_fp32=lm_head_fp32, + lm_head=lm_head, + logprob_temperature=logprob_temperature, + logprob_top_k=logprob_top_k, + logprob_top_p=logprob_top_p, + logprob_min_p=logprob_min_p, + ) + + current_support = torch.isfinite(per_token_ce) + new_logprobs_flat = -per_token_ce.detach() + + # Compute PPO KL: old_log_probs - new_log_probs + ppo_kl = old_logprobs_flat - new_logprobs_flat + + # Mask invalid positions + ppo_kl = ppo_kl.masked_fill(~valid_mask, 0.0) + advantages_masked = advantages_flat.masked_fill(~valid_mask, 0.0) + + # Computed BEFORE compute_ppo_loss to avoid torch.compile interference. + _kl_stats = None + if compute_kl_stats: + with torch.no_grad(): + _raw_log_ratio = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) + _log_ratio_full = torch.where( + current_support, + _raw_log_ratio, + torch.full_like(_raw_log_ratio, -20.0), + ) + _ratio_full = torch.where(current_support, torch.exp(_raw_log_ratio), torch.zeros_like(_raw_log_ratio)) + _per_token_k3 = _ratio_full - _log_ratio_full - 1.0 + # ±inf identity on empty ranks lets cross-rank MIN/MAX-allreduce ignore empty contributors. + if valid_mask.any(): + _ratio_min = _ratio_full.masked_fill(~valid_mask, float("inf")).min() + _ratio_max = _ratio_full.masked_fill(~valid_mask, float("-inf")).max() + _k3_max = _per_token_k3.masked_fill(~valid_mask, float("-inf")).max() + _logratio_min = _log_ratio_full.masked_fill(~valid_mask, float("inf")).min() + _logratio_max = _log_ratio_full.masked_fill(~valid_mask, float("-inf")).max() + _abs_logratio_max = _log_ratio_full.abs().masked_fill(~valid_mask, float("-inf")).max() + else: + _ratio_min = _ratio_full.new_tensor(float("inf")) + _ratio_max = _ratio_full.new_tensor(float("-inf")) + _k3_max = _per_token_k3.new_tensor(float("-inf")) + _logratio_min = _log_ratio_full.new_tensor(float("inf")) + _logratio_max = _log_ratio_full.new_tensor(float("-inf")) + _abs_logratio_max = _log_ratio_full.new_tensor(float("-inf")) + _kl_stats = { + "kl_sample_train_k3": metric_reducer(_per_token_k3, valid_mask_f), + "kl_k3_debug_mean": metric_reducer(_per_token_k3, valid_mask_f), + "kl_k3_debug_max": _k3_max, + "kl_k3_debug_abs_logratio_mean": metric_reducer(_log_ratio_full.abs(), valid_mask_f), + "kl_k3_debug_abs_logratio_max": _abs_logratio_max, + "kl_k3_debug_logratio_mean": metric_reducer(_log_ratio_full, valid_mask_f), + "kl_k3_debug_logratio_min": _logratio_min, + "kl_k3_debug_logratio_max": _logratio_max, + "kl_k3_debug_frac_logratio_positive": metric_reducer((_log_ratio_full > 0).float(), valid_mask_f), + "entropy_sample": metric_reducer(-old_logprobs_flat, valid_mask_f), + "ratio_mean": metric_reducer(_ratio_full, valid_mask_f), + "ratio_min": _ratio_min, + "ratio_max": _ratio_max, + "current_support_fraction": metric_reducer(current_support.float(), valid_mask_f), + } + for suffix, threshold in K3_DEBUG_THRESHOLDS: + _kl_stats[f"kl_k3_debug_frac_gt_{suffix}"] = metric_reducer( + (_per_token_k3 > threshold).float(), valid_mask_f + ) + + # Compute PPO-style clipped loss (returns per-token losses, clip mask, and ratio) + pg_losses, is_clipped, ratio = compute_ppo_loss( + ppo_kl=ppo_kl, + advantages=advantages_masked, + eps_clip=eps_clip, + eps_clip_high=eps_clip_high, + eps_clip_c=eps_clip_c, + ) + + # IcePop hard masking (GLM-5, arXiv:2602.15763): + # Zero gradient for tokens where ratio is outside [1/β, β] + icepop_mask = None + if icepop_beta is not None: + if use_tis: + logger.warning( + "IcePop and TIS are both enabled. IcePop makes TIS redundant " + "when using inference logprobs as old_logprobs." + ) + ratio_d = ratio.detach() + icepop_mask = (ratio_d >= 1.0 / icepop_beta) & (ratio_d <= icepop_beta) + + # Apply TIS correction if enabled and rollout_logprobs provided + tis_metrics = {} + if use_tis and rollout_logprobs is not None: + rollout_logprobs_flat = rollout_logprobs.view(-1) + pg_losses, tis_metrics = apply_tis_correction( + pg_loss=pg_losses, + train_log_probs=new_logprobs_flat, + rollout_log_probs=rollout_logprobs_flat, + valid_mask=valid_mask, + metric_reducer=metric_reducer, + tis_clip_low=tis_clip_low, + tis_clip_high=tis_clip_high, + ) + + # True loss value (for logging): partial share under loss_reducer. + true_loss = loss_reducer(pg_losses, valid_mask_f) + + # Gradient-active mask: tokens that are not clipped, not IcePop-masked, and valid + gradient_active = ~is_clipped & valid_mask & current_support + if icepop_mask is not None: + gradient_active = gradient_active & icepop_mask + + # Surrogate: gradient weight = ratio * advantages, zeroed for inactive tokens + gradient_weight = (ratio.detach() * advantages_flat).masked_fill(~gradient_active, 0.0) + safe_per_token_ce = torch.where(current_support, per_token_ce, torch.zeros_like(per_token_ce)) + surrogate = loss_reducer(gradient_weight * safe_per_token_ce, valid_mask_f) + + # Combine: forward value from true_loss, gradient from surrogate + loss_with_grad = true_loss.detach() + surrogate - surrogate.detach() + + # Return training logprobs reshaped + new_logprobs = new_logprobs_flat.view(original_shape) + + with torch.no_grad(): + metrics: Dict[str, Any] = { + "valid_tokens": valid_count.item(), + "pg_clipfrac": metric_reducer(is_clipped.float(), valid_mask_f), + } + + if icepop_mask is not None: + metrics["icepop_maskfrac"] = metric_reducer((~icepop_mask).float(), valid_mask_f) + + if _kl_stats is not None: + metrics.update(_kl_stats) + + metrics.update(tis_metrics) + + metric_ops: Dict[str, str] = {} + if _kl_stats is not None: + metric_ops["ratio_min"] = "min" + metric_ops["ratio_max"] = "max" + metric_ops["kl_k3_debug_max"] = "max" + metric_ops["kl_k3_debug_abs_logratio_max"] = "max" + metric_ops["kl_k3_debug_logratio_min"] = "min" + metric_ops["kl_k3_debug_logratio_max"] = "max" + if tis_metrics: + metric_ops["tis_min"] = "min" + metric_ops["tis_max"] = "max" + + return LossOutput( + loss=loss_with_grad, + per_token_logprobs=new_logprobs, + metrics=metrics, + metric_ops=metric_ops or None, + ) diff --git a/src/xorl/objectives/reducers.py b/src/xorl/objectives/reducers.py new file mode 100644 index 00000000..48b414c8 --- /dev/null +++ b/src/xorl/objectives/reducers.py @@ -0,0 +1,71 @@ +"""Reducer protocol and canonical denominator policies for loss aggregation. + +A ``Reducer`` collapses a ``(B, S)`` tensor to a scalar over a +caller-supplied denominator policy. Partial shares sum across micro-batches +and ``all_reduce(SUM)`` across ranks to the globally-correct value. + +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + +import torch + + +@runtime_checkable +class Reducer(Protocol): + """``(values, mask) -> scalar`` partial share over a pre-computed + denominator. Partial shares sum across micro-batches and ``all_reduce(SUM)`` + across ranks. + """ + + def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: ... + + +@dataclass(frozen=True) +class TokenPartial: + """Flat masked sum divided by a caller-supplied ``scale``.""" + + scale: torch.Tensor + + def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + return (values * mask).sum() / self.scale.clamp(min=1.0) + + +@dataclass(frozen=True) +class SequencePartial: + """Sum of per-segment token-means, divided by a caller-supplied ``scale``. + + Segment boundaries are flat across ``(values * mask).reshape(-1)``: + + - ``cu_seqlens_local: (N+1,)`` — shard-local segment extents. Under CP each + rank's slice sums to its segment's local contribution. + - ``seq_lengths_global: (N,)`` — pre-CP-shard token count per segment, used + as the per-segment denominator so partial shares from each CP rank sum + to the correct per-segment mean. + """ + + scale: torch.Tensor + cu_seqlens_local: torch.Tensor + seq_lengths_global: torch.Tensor + + def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + flat = (values * mask).reshape(-1) + seg_lengths_local = self.cu_seqlens_local.diff() + n_segments = seg_lengths_local.numel() + seg_ids = torch.repeat_interleave( + torch.arange(n_segments, device=flat.device), + seg_lengths_local, + ) + seg_sums = torch.zeros(n_segments, dtype=flat.dtype, device=flat.device).index_add(0, seg_ids, flat) + seg_means = seg_sums / self.seq_lengths_global.clamp(min=1.0) + return seg_means.sum() / self.scale.clamp(min=1.0) + + +__all__ = [ + "Reducer", + "SequencePartial", + "TokenPartial", +] diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index fe15a08f..ee37606e 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -23,9 +23,9 @@ Three kinds of code live here with three different rules: byte-pinned programs shared with the serving engine. The `bi_` / `exact_` / `canonical_` / `class_b` prefixes are historical names for the same concept. Planned home: `ops/exact/` (#78 phase 3). -- `loss/` — the CE/selected-logprob kernel stack **plus** (for now) the RL - objective functions; the objectives move to `xorl/objectives/` - (#78 phase 2). +- `loss/` — the CE/selected-logprob kernel stack. The RL/supervised + objective functions live in `xorl/objectives/` (#78 phase 2); old module + paths here are compat stubs for one deprecation cycle. - `moe/` — MoE expert compute backends (triton/quack/native, LoRA variants). `ep_kernels/` (DeepEP sort/scatter) merges in here (#78 phase 5). - `linear_attention/` — GDN/linear-attention kernels; also currently hosts @@ -40,5 +40,5 @@ Three kinds of code live here with three different rules: ## What does NOT belong here New `nn.Module` layer classes (→ `models/layers/`), RL objectives -(→ `xorl/objectives/` once it exists), and orchestration logic. `ops/` is +(→ `xorl/objectives/`), and orchestration logic. `ops/` is for kernels and the autograd boundaries directly over them. diff --git a/src/xorl/ops/__init__.py b/src/xorl/ops/__init__.py index 65be904a..36d8f446 100644 --- a/src/xorl/ops/__init__.py +++ b/src/xorl/ops/__init__.py @@ -3,10 +3,6 @@ chunk_gated_delta_rule, fused_recurrent_gated_delta_rule, ) -from .loss import ( - causallm_loss_function, - importance_sampling_loss_function, -) from .moe.quack import quack_moe_forward from .moe.triton import TritonMoeExpertsFunction, triton_moe_forward from .moe.triton_lora import ( @@ -16,6 +12,17 @@ from .ssm import Mamba2Mixer, ssd_chunked +def __getattr__(name: str): + # Objective functions moved to xorl.objectives (#78 phase 2); resolve the + # historical xorl.ops re-exports lazily to avoid an import cycle (the + # objectives import the loss kernels under this package). + if name in ("causallm_loss_function", "importance_sampling_loss_function"): + import xorl.objectives as _objectives + + return getattr(_objectives, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = [ "TritonMoeExpertsFunction", "triton_moe_forward", diff --git a/src/xorl/ops/loss/__init__.py b/src/xorl/ops/loss/__init__.py index e96383ee..5094f03a 100644 --- a/src/xorl/ops/loss/__init__.py +++ b/src/xorl/ops/loss/__init__.py @@ -1,24 +1,14 @@ -""" -Loss functions for training. +"""Cross-entropy / selected-logprob KERNELS. -This module provides various loss functions for language model training: -- causallm_loss_function: Standard causal language modeling loss -- importance_sampling_loss_function: Importance sampling loss for GRPO/RL -- cispo_loss_function: Clipped IS-weight policy optimization -- policy_loss_function: PPO-style policy loss with TIS correction -- drgrpo_loss_function: DR-GRPO loss with PPO clipping and KL penalty +The RL and supervised objective functions moved to :mod:`xorl.objectives` +(issue #78 phase 2). Their public API is re-exported here lazily (PEP 562) +for one deprecation cycle so ``from xorl.ops.loss import ...`` keeps working +without creating an import cycle (the objectives import the kernels in this +package). """ -from typing import Callable, Dict, Literal +from typing import Literal -from xorl.ops.loss.causallm_loss import causallm_loss_function, fsdp_sharded_causallm_loss_function -from xorl.ops.loss.cispo_loss import cispo_loss_function -from xorl.ops.loss.grpo_loss import drgrpo_loss_function -from xorl.ops.loss.importance_sampling_loss import importance_sampling_loss_function -from xorl.ops.loss.loss_output import LossOutput -from xorl.ops.loss.opd_loss import OPDLossMetrics, opd_loss_function, opd_vocab_parallel_loss_function -from xorl.ops.loss.policy_loss import policy_loss_function -from xorl.ops.loss.reducers import Reducer, SequencePartial, TokenPartial from xorl.ops.loss.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy @@ -27,50 +17,38 @@ # ``bi_fused`` runs the shared batch-invariant projection and fixed-order LSE. CrossEntropyMode = Literal["eager", "compiled", "bi_fused", "quack_linear", "fused_quack"] - -# --------------------------------------------------------------------------- -# Loss function registry -# --------------------------------------------------------------------------- -LOSS_REGISTRY: Dict[str, Callable] = { - "causallm_loss": causallm_loss_function, - "cross_entropy": causallm_loss_function, # alias - "importance_sampling": importance_sampling_loss_function, - "cispo": cispo_loss_function, - "policy_loss": policy_loss_function, - "drgrpo": drgrpo_loss_function, - "opd_loss": opd_loss_function, -} - - -def get_loss_function(name: str) -> Callable: - """Look up a loss function by name.""" - if name not in LOSS_REGISTRY: - raise ValueError(f"Unknown loss function: {name}. Available: {list(LOSS_REGISTRY.keys())}") - return LOSS_REGISTRY[name] - - -def register_loss_function(name: str, fn: Callable) -> None: - """Register a custom loss function.""" - LOSS_REGISTRY[name] = fn - - -__all__ = [ +_OBJECTIVE_EXPORTS = frozenset( + { + "LOSS_REGISTRY", + "LossOutput", + "OPDLossMetrics", + "Reducer", + "SequencePartial", + "TokenPartial", + "causallm_loss_function", + "cispo_loss_function", + "drgrpo_loss_function", + "fsdp_sharded_causallm_loss_function", + "get_loss_function", + "importance_sampling_loss_function", + "opd_loss_function", + "opd_vocab_parallel_loss_function", + "policy_loss_function", + "register_loss_function", + } +) + + +def __getattr__(name: str): + if name in _OBJECTIVE_EXPORTS: + import xorl.objectives as _objectives + + return getattr(_objectives, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ # noqa: F822 (objective names resolve via __getattr__) "CrossEntropyMode", - "LossOutput", - "OPDLossMetrics", - "LOSS_REGISTRY", - "Reducer", - "SequencePartial", - "TokenPartial", - "get_loss_function", - "register_loss_function", - "causallm_loss_function", - "fsdp_sharded_causallm_loss_function", - "drgrpo_loss_function", - "importance_sampling_loss_function", - "cispo_loss_function", - "opd_loss_function", - "opd_vocab_parallel_loss_function", - "policy_loss_function", "vocab_parallel_cross_entropy", + *sorted(_OBJECTIVE_EXPORTS), ] diff --git a/src/xorl/ops/loss/causallm_loss.py b/src/xorl/ops/loss/causallm_loss.py index c00323d1..e0c1fb50 100644 --- a/src/xorl/ops/loss/causallm_loss.py +++ b/src/xorl/ops/loss/causallm_loss.py @@ -1,864 +1,3 @@ -from __future__ import annotations +"""Moved to ``xorl.objectives.causallm_loss`` (issue #78 phase 2); compat stub.""" -import math - -import torch -import torch.distributed as dist -import torch.nn.functional as F - -from xorl.ops.exact_sampling_transforms import TOP_K_ALL -from xorl.ops.loss.compiled_cross_entropy import ( - compiled_ce_and_lse_sq_function, - compiled_cross_entropy_function, -) -from xorl.ops.loss.loss_output import LossOutput -from xorl.ops.loss.per_token_ce import ( - LogprobProbability, - LogprobTemperature, - LogprobTopK, - compute_per_token_ce, - normalize_logprob_temperature, - resolve_bi_fused_lm_head_tp_groups, -) -from xorl.ops.loss.reducers import Reducer, TokenPartial -from xorl.ops.loss.vocab_parallel_cross_entropy import ( - _backward_kernel as _vocab_parallel_ce_backward_kernel, -) -from xorl.ops.loss.vocab_parallel_cross_entropy import ( - _forward_kernel as _vocab_parallel_ce_forward_kernel, -) -from xorl.ops.loss.vocab_parallel_cross_entropy import ( - vocab_parallel_cross_entropy, - vocab_parallel_cross_entropy_with_lm_head, -) - - -_MODULE_LM_HEAD_MIN_CHUNK_ROWS = 128 - - -def _all_gather_cat_same_shape(x: torch.Tensor, *, dim: int, group: dist.ProcessGroup) -> torch.Tensor: - world_size = dist.get_world_size(group) - gathered = [torch.empty_like(x) for _ in range(world_size)] - dist.all_gather(gathered, x.contiguous(), group=group) - return torch.cat(gathered, dim=dim) - - -class _FSDPShardedCausalLMLoss(torch.autograd.Function): - """Sequence-streaming vocab-parallel CE for FSDP-sharded lm_head. - - The lm_head shard group is also the sequence-parallel group for the GLM - 128K configuration, so each rank owns different sequence tokens and a - different vocab shard. The usual vocab-parallel CE needs all ranks to see - the same token batch; this Function gathers one small sequence chunk at a - time and avoids saving those gathered chunks for backward. - """ - - @staticmethod - def forward( - ctx, - hidden_states: torch.Tensor, - local_weight: torch.Tensor, - labels: torch.Tensor, - global_valid_tokens: torch.Tensor, - sequence_group: dist.ProcessGroup, - vocab_group: dist.ProcessGroup, - num_chunks: int, - ignore_index: int, - loss_reduce_group: "dist.ProcessGroup | None" = None, - loss_reduce_divisor: float = 1.0, - ) -> torch.Tensor: - if hidden_states.dim() != 3: - raise ValueError(f"Expected hidden_states to have shape [B, S, H], got {tuple(hidden_states.shape)}") - if labels.shape != hidden_states.shape[:2]: - raise ValueError(f"Expected labels shape {tuple(hidden_states.shape[:2])}, got {tuple(labels.shape)}") - - ctx.save_for_backward(hidden_states, local_weight, labels, global_valid_tokens) - ctx.sequence_group = sequence_group - ctx.vocab_group = vocab_group - ctx.num_chunks = num_chunks - ctx.ignore_index = ignore_index - - vocab_rank = dist.get_rank(vocab_group) - local_vocab_size = local_weight.shape[0] - vocab_offset = vocab_rank * local_vocab_size - local_seq_len = hidden_states.shape[1] - chunk_size = max(1, math.ceil(local_seq_len / num_chunks)) - denom = global_valid_tokens.clamp(min=1.0) - loss = hidden_states.new_zeros((), dtype=torch.float32) - - for start in range(0, local_seq_len, chunk_size): - end = min(start + chunk_size, local_seq_len) - gathered_hidden = _all_gather_cat_same_shape(hidden_states[:, start:end, :], dim=1, group=sequence_group) - gathered_labels = _all_gather_cat_same_shape(labels[:, start:end], dim=1, group=sequence_group) - hidden_flat = gathered_hidden.reshape(-1, gathered_hidden.shape[-1]) - labels_flat = gathered_labels.reshape(-1) - per_token_ce, _, _, _, _, valid_mask = _vocab_parallel_ce_forward_kernel( - hidden_flat, - local_weight, - labels_flat, - vocab_group, - vocab_offset, - local_vocab_size, - ignore_index, - ) - loss = loss + (per_token_ce * valid_mask.float()).sum() / denom - del gathered_hidden, gathered_labels, hidden_flat, labels_flat, per_token_ce, valid_mask - - # lm-head-TP: each replica group computes the CE over its own sequence - # shard and over the full (TP-split) vocab, so the per-replica losses must - # be summed across replicas. The divisor removes the within-TP-group - # duplication (every TP rank computed the same per-replica loss). Backward - # intentionally stays unscaled: each rank returns its local weight-shard - # gradient, which the caller combines with an all-reduce over the replica - # group to reconstruct the full gradient. - if loss_reduce_group is not None: - dist.all_reduce(loss, op=dist.ReduceOp.SUM, group=loss_reduce_group) - if loss_reduce_divisor != 1.0: - loss = loss / loss_reduce_divisor - - return loss - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - hidden_states, local_weight, labels, global_valid_tokens = ctx.saved_tensors - sequence_group = ctx.sequence_group - vocab_group = ctx.vocab_group - vocab_rank = dist.get_rank(vocab_group) - local_vocab_size = local_weight.shape[0] - vocab_offset = vocab_rank * local_vocab_size - local_seq_len = hidden_states.shape[1] - chunk_size = max(1, math.ceil(local_seq_len / ctx.num_chunks)) - denom = global_valid_tokens.clamp(min=1.0) - - grad_hidden = torch.zeros_like(hidden_states) if ctx.needs_input_grad[0] else None - grad_weight = torch.zeros_like(local_weight) if ctx.needs_input_grad[1] else None - sp_world = dist.get_world_size(sequence_group) - sp_rank = dist.get_rank(sequence_group) - - for start in range(0, local_seq_len, chunk_size): - end = min(start + chunk_size, local_seq_len) - gathered_hidden = _all_gather_cat_same_shape(hidden_states[:, start:end, :], dim=1, group=sequence_group) - gathered_labels = _all_gather_cat_same_shape(labels[:, start:end], dim=1, group=sequence_group) - hidden_flat = gathered_hidden.reshape(-1, gathered_hidden.shape[-1]) - labels_flat = gathered_labels.reshape(-1) - _, global_max, global_sumexp, target_in_range, safe_local_target, valid_mask = ( - _vocab_parallel_ce_forward_kernel( - hidden_flat, - local_weight, - labels_flat, - vocab_group, - vocab_offset, - local_vocab_size, - ctx.ignore_index, - ) - ) - per_token_grad = grad_output.to(hidden_flat.dtype) * valid_mask.to(hidden_flat.dtype) / denom - chunk_grad_hidden, chunk_grad_weight = _vocab_parallel_ce_backward_kernel( - per_token_grad, - hidden_flat, - local_weight, - global_max, - global_sumexp, - target_in_range, - safe_local_target, - valid_mask, - vocab_group, - ctx.needs_input_grad[1], - ) - if grad_hidden is not None: - chunk_grad_hidden = chunk_grad_hidden.view_as(gathered_hidden) - local_grad = chunk_grad_hidden.chunk(sp_world, dim=1)[sp_rank].contiguous() - grad_hidden[:, start:end, :] = local_grad.to(grad_hidden.dtype) - if grad_weight is not None and chunk_grad_weight is not None: - grad_weight.add_(chunk_grad_weight.to(grad_weight.dtype)) - del ( - gathered_hidden, - gathered_labels, - hidden_flat, - labels_flat, - global_max, - global_sumexp, - target_in_range, - safe_local_target, - valid_mask, - per_token_grad, - chunk_grad_hidden, - chunk_grad_weight, - ) - - # Grads for: hidden_states, local_weight, then None for labels, - # global_valid_tokens, sequence_group, vocab_group, num_chunks, - # ignore_index, loss_reduce_group, loss_reduce_divisor. - return grad_hidden, grad_weight, None, None, None, None, None, None, None, None - - -def fsdp_sharded_causallm_loss_function( - hidden_states: torch.Tensor, - weight: torch.Tensor, - labels: torch.Tensor, - sp_group: dist.ProcessGroup, - fsdp_group: dist.ProcessGroup, - num_chunks: int, - ignore_index: int = -100, - lm_head_fp32: bool = False, - global_valid_tokens: torch.Tensor | None = None, - sequence_group: "dist.ProcessGroup | None" = None, - vocab_group: "dist.ProcessGroup | None" = None, - loss_reduce_group: "dist.ProcessGroup | None" = None, - loss_reduce_divisor: float = 1.0, -) -> "LossOutput": - # sequence_group/vocab_group default to sp_group/fsdp_group (the FSDP-sharded - # lm_head case where the sequence-parallel group is also the vocab shard - # group). lm-head-only TP passes a dedicated lm_head_tp_group for both and a - # loss_reduce_group (+ divisor) to sum the per-replica losses. - if sequence_group is None: - sequence_group = sp_group - if vocab_group is None: - vocab_group = fsdp_group - if lm_head_fp32: - hidden_states = hidden_states.float() - - local_weight = weight.to_local() if hasattr(weight, "to_local") else weight - if local_weight.dtype != hidden_states.dtype: - if lm_head_fp32: - local_weight = local_weight.float() - else: - local_weight = local_weight.to(hidden_states.dtype) - - if global_valid_tokens is None: - global_valid_tokens = (labels != ignore_index).sum().to(hidden_states.device, dtype=torch.float32) - dist.all_reduce(global_valid_tokens, op=dist.ReduceOp.SUM, group=fsdp_group) - else: - global_valid_tokens = global_valid_tokens.detach().to(hidden_states.device, dtype=torch.float32) - loss = _FSDPShardedCausalLMLoss.apply( - hidden_states, - local_weight, - labels, - global_valid_tokens, - sequence_group, - vocab_group, - int(num_chunks), - int(ignore_index), - loss_reduce_group, - float(loss_reduce_divisor), - ) - return LossOutput(loss=loss) - - -def _chunked_lm_head_cross_entropy( - hidden_states_flat: torch.Tensor, - labels_flat: torch.Tensor, - *, - lm_head: torch.nn.Module, - ignore_index: int, - num_chunks: int, - z_loss_enabled: bool, - valid_mask: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Compute CE by calling the lm_head module in chunks. - - FP8 training wraps ``lm_head`` with ``FP8Linear``. The compiled CE helpers - operate on ``lm_head.weight`` directly, which bypasses that module, so FP8 - training uses this path to keep the output head matmul on FP8 compute. - """ - - if hidden_states_flat.shape[0] == 0: - empty = hidden_states_flat.new_empty((0,), dtype=torch.float32) - return empty, empty if z_loss_enabled else None - - chunk_count = max(1, int(num_chunks)) - chunk_size = max(_MODULE_LM_HEAD_MIN_CHUNK_ROWS, math.ceil(hidden_states_flat.shape[0] / chunk_count)) - ce_chunks: list[torch.Tensor] = [] - lse_sq_chunks: list[torch.Tensor] = [] - for start in range(0, hidden_states_flat.shape[0], chunk_size): - end = min(start + chunk_size, hidden_states_flat.shape[0]) - logits = lm_head(hidden_states_flat[start:end]).float() - labels = labels_flat[start:end] - ce_chunks.append(F.cross_entropy(logits, labels, reduction="none", ignore_index=ignore_index)) - if z_loss_enabled: - lse = torch.logsumexp(logits, dim=-1) - lse_sq_chunks.append((lse * lse) * valid_mask[start:end].to(lse.dtype)) - - per_token_ce = torch.cat(ce_chunks, dim=0) - per_token_lse_sq = torch.cat(lse_sq_chunks, dim=0) if z_loss_enabled else None - return per_token_ce, per_token_lse_sq - - -def _ceil_to_multiple(value: int, multiple: int) -> int: - return ((value + multiple - 1) // multiple) * multiple - - -def _chunk_size_from_num_chunks(num_tokens: int, num_chunks: int) -> int: - if num_chunks <= 0: - return _ceil_to_multiple(num_tokens, 8) - return _ceil_to_multiple((num_tokens + num_chunks - 1) // num_chunks, 8) - - -def _quack_linear_cross_entropy_loss( - hidden_states_flat: torch.Tensor, - weight: torch.Tensor, - labels_flat: torch.Tensor, - ignore_index: int, - num_chunks: int, - loss_reducer: TokenPartial, -) -> torch.Tensor: - if not hidden_states_flat.is_cuda: - raise ValueError("ce_mode='quack_linear' requires CUDA tensors") - if hidden_states_flat.shape[-1] % 8 != 0 or weight.shape[0] % 8 != 0: - raise ValueError("ce_mode='quack_linear' requires hidden and vocab dimensions to be divisible by 8") - - from xorl.ops._vendored.quack.linear_cross_entropy import chunked_linear_cross_entropy # noqa: PLC0415 - - hidden_states_flat, labels_flat = _pad_quack_linear_rows( - hidden_states_flat, - labels_flat, - ignore_index=ignore_index, - ) - valid_count = (labels_flat != ignore_index).sum() - if valid_count.item() == 0: - return (hidden_states_flat.sum() + weight.sum()) * 0.0 - - chunk_size = _chunk_size_from_num_chunks(hidden_states_flat.shape[0], num_chunks) - loss_sum = chunked_linear_cross_entropy( - hidden_states_flat, - weight, - labels_flat, - chunk_size=chunk_size, - ignore_index=ignore_index, - reduction="sum", - ) - scale = loss_reducer.scale.to(device=loss_sum.device, dtype=loss_sum.dtype) - return loss_sum / scale.clamp(min=1.0) - - -def _pad_quack_linear_rows( - hidden_states_flat: torch.Tensor, - labels_flat: torch.Tensor, - *, - ignore_index: int, - multiple: int = 8, -) -> tuple[torch.Tensor, torch.Tensor]: - """Pad rows for Quack CE kernels without changing scalar CE semantics.""" - - remainder = hidden_states_flat.shape[0] % multiple - if remainder == 0: - return hidden_states_flat, labels_flat - pad_rows = multiple - remainder - hidden_pad = hidden_states_flat.new_zeros((pad_rows, hidden_states_flat.shape[-1])) - label_pad = labels_flat.new_full((pad_rows,), ignore_index) - return torch.cat((hidden_states_flat, hidden_pad), dim=0), torch.cat((labels_flat, label_pad), dim=0) - - -def _fused_quack_per_token_ce( - hidden_states_flat: torch.Tensor, - weight: torch.Tensor, - labels_flat: torch.Tensor, - ignore_index: int, - num_chunks: int, - tp_group, - lm_head_fp32: bool, -) -> torch.Tensor: - """Per-token CE ``[N]`` via the fused chunked cuBLAS + CuTeDSL path. - - ``fused_selected_logprob_ce`` keeps the logits tile bounded to - ``[chunk, V_local]`` and never materializes the full ``[N, V]`` logits. - Mirrors the dispatch in ``ops.loss.per_token_ce.compute_per_token_ce`` so - ``ce_mode='fused_quack'`` works from this entry point too — without it, - fused_quack fell through to the eager full-logits path and OOM'd at large - vocab / long context. - """ - if not hidden_states_flat.is_cuda: - raise ValueError("ce_mode='fused_quack' requires CUDA tensors") - - from xorl.ops.loss.fused_linear_logprob import fused_selected_logprob_ce # noqa: PLC0415 - - local_weight = weight.to_local() if hasattr(weight, "to_local") else weight - hidden = hidden_states_flat - if lm_head_fp32: - hidden = hidden.float() - local_weight = local_weight.float() - chunk_size = _chunk_size_from_num_chunks(hidden.shape[0], num_chunks) - return fused_selected_logprob_ce( - hidden, - local_weight, - labels_flat, - tp_group=tp_group, - ignore_index=ignore_index, - chunk_size=chunk_size, - ) - - -def _bi_fused_per_token_ce_checked( - hidden_states_flat: torch.Tensor, - weight: torch.Tensor, - labels_flat: torch.Tensor, - ignore_index: int, - lm_head_fp32: bool, - z_loss_enabled: bool, -) -> torch.Tensor: - """Guarded entry for ``ce_mode='bi_fused'`` (the batch-invariant lm-head - contract). The contract IS the fp32-class lm-head computation, so it - requires ``lm_head_fp32`` semantics without materializing the fp32 weight.""" - from xorl.ops.loss.bi_fused_lm_head import bi_fused_per_token_ce - - if z_loss_enabled: - raise NotImplementedError("ce_mode='bi_fused' does not support softmax_auxiliary_loss") - if not lm_head_fp32: - raise NotImplementedError( - "ce_mode='bi_fused' implements the fp32-class lm-head contract; set lm_head_fp32: true" - ) - local_weight = weight.to_local() if hasattr(weight, "to_local") else weight - return bi_fused_per_token_ce(hidden_states_flat, local_weight, labels_flat, ignore_index) - - -def _quack_linear_per_token_cross_entropy( - hidden_states_flat: torch.Tensor, - weight: torch.Tensor, - labels_flat: torch.Tensor, - ignore_index: int, - num_chunks: int, - lm_head_fp32: bool, -) -> torch.Tensor: - """Per-token return path for ``ce_mode='quack_linear'``. - - The scalar training path keeps using Quack's chunked linear CE reduction. - ``return_per_token=True`` callers need one CE value per input row, so route - those through the existing fused selected-logprob kernel and still avoid - full-logit materialization. - """ - - if not hidden_states_flat.is_cuda: - raise ValueError("ce_mode='quack_linear' requires CUDA tensors") - if hidden_states_flat.shape[-1] % 8 != 0 or weight.shape[0] % 8 != 0: - raise ValueError("ce_mode='quack_linear' requires hidden and vocab dimensions to be divisible by 8") - return _fused_quack_per_token_ce( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - num_chunks, - tp_group=None, - lm_head_fp32=lm_head_fp32, - ) - - -def causallm_loss_function( - hidden_states: torch.Tensor, - weight: torch.Tensor, - labels: torch.Tensor, - ignore_index: int = -100, - return_per_token: bool = False, - ce_mode: str = "compiled", - num_chunks: int = 8, - tp_group=None, - use_compile: bool = False, - lm_head_fp32: bool = False, - loss_reducer: Reducer | None = None, - z_loss_coef: float = 0.0, - lm_head: torch.nn.Module | None = None, - logprob_temperature: LogprobTemperature = 1.0, - logprob_top_k: LogprobTopK = TOP_K_ALL, - logprob_top_p: LogprobProbability = 1.0, - logprob_min_p: LogprobProbability = 0.0, -) -> "LossOutput": - """ - Compute causal language modeling loss. - - Supports multiple computation modes: - - "compiled": RECOMMENDED. torch.compile (1.6x speed, 16% memory) - - "eager": Simple F.cross_entropy baseline (may OOM at 32K) - - Args: - hidden_states: Model hidden states, shape (batch, seq_len, hidden_dim) - weight: LM head weight matrix, shape (vocab_size, hidden_dim). - With TP, this is the local shard [vocab_size/tp, hidden_dim]. - labels: Target labels, shape (batch, seq_len). Labels are assumed to be - already next-token aligned (labels[i] is the target for hidden_states[i]). - ignore_index: Index to ignore in loss computation (default: -100) - return_per_token: If True, return per-token logprobs and losses (default: False) - ce_mode: Cross-entropy mode - "compiled" (default) or "eager" - num_chunks: Number of chunks for compiled mode (default: 8). - tp_group: TP process group for vocab-parallel cross-entropy (default: None). - loss_reducer: Optional ``(values, mask) -> scalar``. When supplied, the - returned loss is a partial share under the reducer's denominator - (sum across micro-batches + all-reduce across ranks recovers the - globally-correct loss). When None, falls back to a local token mean. - Z-loss (when enabled) is reduced through the same reducer so the - two terms compose consistently. - z_loss_coef: If > 0, add the Z-loss auxiliary term used in OLMo / - PaLM-style training: - z_loss = coef * sum(logsumexp(logits)^2 * mask) / num_valid_tokens - where ``mask = labels != ignore_index``. Equivalent to OLMo's - ``cross_entropy_loss(..., reduction="sum")`` path divided by - ``batch_size_in_tokens``. Encourages log(Z) to stay near zero, - stabilizing training at large vocab / high LR. Not supported - in the TP path. - logprob_temperature: Temperature for selected-token logprobs. ``1.0`` - returns raw model logprobs; a rollout temperature such as - ``0.7`` returns behavior-policy logprobs using - ``log_softmax(logits / temperature)``. Exact LM heads also - accept contiguous FP32 temperatures aligned with labels. - - Returns: - LossOutput with loss, and optionally per_token_logprobs/per_token_loss. - When ``z_loss_coef > 0``, ``LossOutput.metrics`` contains - ``{"ce_loss": , "z_loss": }``. - """ - # Store original shape before flattening for per-token outputs - original_shape = labels.shape - - # Flatten the labels and hidden_states - labels_flat = labels.view(-1) - hidden_states_flat = hidden_states.view(-1, hidden_states.size(-1)) - valid_mask = labels_flat != ignore_index - bi_fused_tp_groups = resolve_bi_fused_lm_head_tp_groups(ce_mode, tp_group, lm_head) - has_explicit_loss_reducer = loss_reducer is not None - - if loss_reducer is None: - scale = valid_mask.sum().float() - if bi_fused_tp_groups is not None: - dedicated_group, replica_group = bi_fused_tp_groups - dist.all_reduce(scale, op=dist.ReduceOp.SUM, group=dedicated_group) - if replica_group is not None: - dist.all_reduce(scale, op=dist.ReduceOp.SUM, group=replica_group) - loss_reducer = TokenPartial(scale=scale) - - mask_flat = valid_mask.float() - if isinstance(logprob_temperature, torch.Tensor): - if not logprob_temperature.is_contiguous(): - raise ValueError("per-row logprob_temperature must be contiguous") - if tuple(logprob_temperature.shape) not in (tuple(labels.shape), (labels_flat.shape[0],)): - raise ValueError( - "per-row logprob_temperature must match labels or flattened labels, got " - f"{tuple(logprob_temperature.shape)} for labels {tuple(labels.shape)}" - ) - logprob_temperature = logprob_temperature.reshape(-1) - logprob_temperature = normalize_logprob_temperature( - logprob_temperature, - rows=labels_flat.shape[0], - device=hidden_states.device, - ) - - def _flatten_sampling_metadata(value, name: str): - if not isinstance(value, torch.Tensor): - return value - if not value.is_contiguous(): - raise ValueError(f"per-row {name} must be contiguous") - if tuple(value.shape) not in (tuple(labels.shape), (labels_flat.shape[0],)): - raise ValueError(f"per-row {name} must match labels or flattened labels") - return value.reshape(-1) - - logprob_top_k = _flatten_sampling_metadata(logprob_top_k, "logprob_top_ks") - logprob_top_p = _flatten_sampling_metadata(logprob_top_p, "logprob_top_ps") - logprob_min_p = _flatten_sampling_metadata(logprob_min_p, "logprob_min_ps") - has_temperature_transform = isinstance(logprob_temperature, torch.Tensor) or logprob_temperature != 1.0 - has_sampling_filter = ( - isinstance(logprob_top_k, torch.Tensor) - or isinstance(logprob_top_p, torch.Tensor) - or isinstance(logprob_min_p, torch.Tensor) - or int(logprob_top_k) < TOP_K_ALL - or float(logprob_top_p) != 1.0 - or float(logprob_min_p) != 0.0 - ) - exact_lm_head = bool( - lm_head is not None - and (getattr(lm_head, "_glm52_exact_tp16_lm_head", False) or getattr(lm_head, "_dsv4_exact_tp8_lm_head", False)) - ) - if ce_mode == "bi_fused": - if tp_group is not None and not exact_lm_head and bi_fused_tp_groups is None: - raise NotImplementedError( - "ce_mode='bi_fused' supports TP only through the dedicated vocabulary-sharded LM-head TP path" - ) - if lm_head is not None and not lm_head_fp32 and not exact_lm_head: - raise NotImplementedError("ce_mode='bi_fused' does not support FP8 lm_head modules") - if exact_lm_head: - if z_loss_coef > 0.0: - raise NotImplementedError("The exact GLM-5.2 active-LoRA lm head does not support Z-loss") - per_token_ce = compute_per_token_ce( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - ce_mode, - num_chunks, - tp_group=tp_group, - use_compile=use_compile, - lm_head_fp32=lm_head_fp32, - lm_head=lm_head, - logprob_temperature=logprob_temperature, - logprob_top_k=logprob_top_k, - logprob_top_p=logprob_top_p, - logprob_min_p=logprob_min_p, - ) - loss = loss_reducer(per_token_ce, mask_flat) - if return_per_token: - return LossOutput( - loss=loss, - per_token_logprobs=-per_token_ce.detach().view(original_shape), - per_token_loss=per_token_ce.view(original_shape), - ) - return LossOutput(loss=loss) - if bi_fused_tp_groups is not None: - if z_loss_coef > 0.0: - raise NotImplementedError("ce_mode='bi_fused' does not support softmax_auxiliary_loss") - per_token_ce = compute_per_token_ce( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - ce_mode, - num_chunks, - tp_group=tp_group, - use_compile=use_compile, - lm_head_fp32=lm_head_fp32, - lm_head=lm_head, - logprob_temperature=logprob_temperature, - logprob_top_k=logprob_top_k, - logprob_top_p=logprob_top_p, - logprob_min_p=logprob_min_p, - ) - local_loss = loss_reducer(per_token_ce, mask_flat) - if has_explicit_loss_reducer: - if return_per_token: - return LossOutput( - loss=local_loss, - per_token_logprobs=-per_token_ce.detach().view(original_shape), - per_token_loss=per_token_ce.view(original_shape), - ) - return LossOutput(loss=local_loss) - - # Standalone/default-reducer calls historically return the full scalar - # on every rank. Explicit reducers instead promise a local partial, and - # their caller owns detached reporting aggregation. - global_loss = local_loss.detach().clone() - dedicated_group, replica_group = bi_fused_tp_groups - dist.all_reduce(global_loss, op=dist.ReduceOp.SUM, group=dedicated_group) - if replica_group is not None: - dist.all_reduce(global_loss, op=dist.ReduceOp.SUM, group=replica_group) - loss = local_loss + (global_loss - local_loss.detach()) - if return_per_token: - return LossOutput( - loss=loss, - per_token_logprobs=-per_token_ce.detach().view(original_shape), - per_token_loss=per_token_ce.view(original_shape), - ) - return LossOutput(loss=loss) - if has_temperature_transform or has_sampling_filter: - if z_loss_coef > 0.0: - raise NotImplementedError("logprob_temperature is not supported with softmax_auxiliary_loss") - per_token_ce = compute_per_token_ce( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - ce_mode, - num_chunks, - tp_group=tp_group, - use_compile=use_compile, - lm_head_fp32=lm_head_fp32, - lm_head=lm_head, - logprob_temperature=logprob_temperature, - logprob_top_k=logprob_top_k, - logprob_top_p=logprob_top_p, - logprob_min_p=logprob_min_p, - ) - loss = loss_reducer(per_token_ce, mask_flat) - if return_per_token: - return LossOutput( - loss=loss, - per_token_logprobs=-per_token_ce.detach().view(original_shape), - per_token_loss=per_token_ce.view(original_shape), - ) - return LossOutput(loss=loss) - - # Vocab-parallel cross-entropy for tensor parallelism - if tp_group is not None: - if z_loss_coef > 0.0: - raise NotImplementedError( - "softmax_auxiliary_loss (Z-loss) is not yet supported with tensor parallelism. " - "Disable softmax_auxiliary_loss or run without TP." - ) - if lm_head is not None and not lm_head_fp32: - per_token_ce = vocab_parallel_cross_entropy_with_lm_head( - hidden_states_flat, - lm_head, - labels_flat, - tp_group, - ignore_index=ignore_index, - num_chunks=num_chunks, - use_compile=use_compile, - ) - else: - # lm_head_fp32 takes precedence over the FP8 lm_head module: compute - # the vocab-parallel CE in fp32 from the master weight (FP8 module - # bypassed). Extract local weight from DTensor if needed. - local_weight = weight.to_local() if hasattr(weight, "to_local") else weight - if lm_head_fp32: - hidden_states_flat = hidden_states_flat.float() - local_weight = local_weight.float() - elif local_weight.dtype != hidden_states_flat.dtype: - local_weight = local_weight.to(hidden_states_flat.dtype) - - per_token_ce = vocab_parallel_cross_entropy( - hidden_states_flat, - local_weight, - labels_flat, - tp_group, - ignore_index=ignore_index, - num_chunks=num_chunks, - use_compile=use_compile, - ) - - loss = loss_reducer(per_token_ce, mask_flat) - if return_per_token: - return LossOutput( - loss=loss, - per_token_logprobs=-per_token_ce.detach().view(original_shape), - per_token_loss=per_token_ce.view(original_shape), - ) - return LossOutput(loss=loss) - - z_loss_enabled = z_loss_coef > 0.0 - # lm_head_fp32 takes precedence over the FP8 lm_head module: an FP32 lm_head - # must not be FP8-quantized, so route to the fp32 weight-CE path below - # (compiled/eager honor lm_head_fp32) instead of _chunked_lm_head_cross_entropy - # (which calls FP8Linear.forward). The FP8 lm_head otherwise catastrophically - # mis-scores rare near-certain tokens (R1). - use_lm_head_module = lm_head is not None and not lm_head_fp32 - - if ce_mode == "quack_linear" and not return_per_token: - if z_loss_enabled: - raise NotImplementedError("ce_mode='quack_linear' does not support softmax_auxiliary_loss") - if lm_head_fp32: - raise NotImplementedError("ce_mode='quack_linear' does not support lm_head_fp32=True") - if not isinstance(loss_reducer, TokenPartial): - raise NotImplementedError("ce_mode='quack_linear' currently supports only TokenPartial loss reduction") - return LossOutput( - loss=_quack_linear_cross_entropy_loss( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - num_chunks, - loss_reducer, - ) - ) - - if return_per_token: - # Compute cross-entropy based on mode (and Z-loss when enabled). - per_token_lse_sq = None - if use_lm_head_module: - per_token_ce, per_token_lse_sq = _chunked_lm_head_cross_entropy( - hidden_states_flat, - labels_flat, - lm_head=lm_head, - ignore_index=ignore_index, - num_chunks=num_chunks, - z_loss_enabled=z_loss_enabled, - valid_mask=valid_mask, - ) - elif ce_mode == "compiled": - if z_loss_enabled: - per_token_ce, per_token_lse_sq = compiled_ce_and_lse_sq_function( - hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 - ) - else: - per_token_ce = compiled_cross_entropy_function( - hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 - ) - elif ce_mode == "fused_quack": - if z_loss_enabled: - raise NotImplementedError("ce_mode='fused_quack' does not support softmax_auxiliary_loss") - per_token_ce = _fused_quack_per_token_ce( - hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, tp_group, lm_head_fp32 - ) - elif ce_mode == "bi_fused": - per_token_ce = _bi_fused_per_token_ce_checked( - hidden_states_flat, weight, labels_flat, ignore_index, lm_head_fp32, z_loss_enabled - ) - elif ce_mode == "quack_linear": - if z_loss_enabled: - raise NotImplementedError("ce_mode='quack_linear' does not support softmax_auxiliary_loss") - per_token_ce = _quack_linear_per_token_cross_entropy( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - num_chunks, - lm_head_fp32, - ) - else: # eager mode - if lm_head_fp32: - logits_flat = (hidden_states_flat.float() @ weight.float().t()).float() - else: - logits_flat = (hidden_states_flat @ weight.t()).float() - per_token_ce = F.cross_entropy(logits_flat, labels_flat, reduction="none", ignore_index=ignore_index) - if z_loss_enabled: - lse = torch.logsumexp(logits_flat, dim=-1) - per_token_lse_sq = (lse * lse) * valid_mask.to(lse.dtype) - - ce_loss = loss_reducer(per_token_ce, mask_flat) - if z_loss_enabled: - z_loss = loss_reducer(per_token_lse_sq, mask_flat) - loss = ce_loss + z_loss_coef * z_loss - metrics = {"ce_loss": ce_loss.detach(), "z_loss": z_loss.detach()} - else: - loss = ce_loss - metrics = None - return LossOutput( - loss=loss, - per_token_logprobs=-per_token_ce.detach().view(original_shape), - per_token_loss=per_token_ce.view(original_shape), - metrics=metrics, - ) - else: - # Always use reduction="none" + manual mean to avoid NaN when all labels - # are ignore_index (reduction="mean" returns NaN for 0 valid elements). - # Keeping the autograd graph intact is critical for FSDP2: all ranks must - # trigger reduce-scatter for every parameter, including lm_head weight. - per_token_lse_sq = None - if use_lm_head_module: - per_token_ce, per_token_lse_sq = _chunked_lm_head_cross_entropy( - hidden_states_flat, - labels_flat, - lm_head=lm_head, - ignore_index=ignore_index, - num_chunks=num_chunks, - z_loss_enabled=z_loss_enabled, - valid_mask=valid_mask, - ) - elif ce_mode == "compiled": - if z_loss_enabled: - per_token_ce, per_token_lse_sq = compiled_ce_and_lse_sq_function( - hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 - ) - else: - per_token_ce = compiled_cross_entropy_function( - hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, lm_head_fp32=lm_head_fp32 - ) - elif ce_mode == "fused_quack": - if z_loss_enabled: - raise NotImplementedError("ce_mode='fused_quack' does not support softmax_auxiliary_loss") - per_token_ce = _fused_quack_per_token_ce( - hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, tp_group, lm_head_fp32 - ) - elif ce_mode == "bi_fused": - per_token_ce = _bi_fused_per_token_ce_checked( - hidden_states_flat, weight, labels_flat, ignore_index, lm_head_fp32, z_loss_enabled - ) - else: # eager mode - if lm_head_fp32: - logits_flat = (hidden_states_flat.float() @ weight.float().t()).float() - else: - logits_flat = (hidden_states_flat @ weight.t()).float() - per_token_ce = F.cross_entropy(logits_flat, labels_flat, reduction="none", ignore_index=ignore_index) - if z_loss_enabled: - lse = torch.logsumexp(logits_flat, dim=-1) - per_token_lse_sq = (lse * lse) * valid_mask.to(lse.dtype) - - ce_loss = loss_reducer(per_token_ce, mask_flat) - if z_loss_enabled: - z_loss = loss_reducer(per_token_lse_sq, mask_flat) - loss = ce_loss + z_loss_coef * z_loss - return LossOutput(loss=loss, metrics={"ce_loss": ce_loss.detach(), "z_loss": z_loss.detach()}) - return LossOutput(loss=ce_loss) +from xorl.objectives.causallm_loss import * # noqa: F401,F403 diff --git a/src/xorl/ops/loss/cispo_loss.py b/src/xorl/ops/loss/cispo_loss.py index eaacfa31..20a52535 100644 --- a/src/xorl/ops/loss/cispo_loss.py +++ b/src/xorl/ops/loss/cispo_loss.py @@ -1,164 +1,3 @@ -from __future__ import annotations +"""Moved to ``xorl.objectives.cispo_loss`` (issue #78 phase 2); compat stub.""" -from typing import Any, Dict, Optional - -import torch -import torch.distributed as dist - -from xorl.ops.exact_sampling_transforms import TOP_K_ALL -from xorl.ops.loss.importance_sampling_loss import K3_DEBUG_THRESHOLDS -from xorl.ops.loss.loss_output import LossOutput -from xorl.ops.loss.per_token_ce import compute_per_token_ce -from xorl.ops.loss.reducers import Reducer, TokenPartial - - -def cispo_loss_function( - hidden_states: torch.Tensor, - weight: torch.Tensor, - labels: torch.Tensor, - old_logprobs: torch.Tensor, - advantages: torch.Tensor, - ignore_index: int = -100, - clip_low_threshold: float = 0.0, - clip_high_threshold: float = 4.0, - num_chunks: int = 8, - ce_mode: str = "compiled", - tp_group: Optional[dist.ProcessGroup] = None, - compute_kl_stats: bool = False, - lm_head_fp32: bool = False, - loss_reducer: Optional[Reducer] = None, - metric_reducer: Optional[Reducer] = None, - lm_head: Optional[torch.nn.Module] = None, - logprob_temperature: float = 1.0, - logprob_top_k: int | torch.Tensor = TOP_K_ALL, - logprob_top_p: float | torch.Tensor = 1.0, - logprob_min_p: float | torch.Tensor = 0.0, -) -> LossOutput: - """Compute Tinker-compatible CISPO. - - CISPO clips the importance ratio and uses it as a detached coefficient on - ``log p_theta``. Unlike PPO objective clipping, this retains a gradient for - every valid token:: - - ratio = exp(target_logprobs - sampling_logprobs) - clipped = clamp(ratio, clip_low_threshold, clip_high_threshold) - loss = -(clipped.detach() * target_logprobs * advantages).mean() - - The default absolute ratio bounds, ``[0, 4]``, match Tinker's one-sided - CISPO default and the MiniMax-M1 prescription of disabling the lower bound. - """ - if clip_low_threshold < 0.0: - raise ValueError("clip_low_threshold must be non-negative") - if clip_high_threshold < clip_low_threshold: - raise ValueError("clip_high_threshold must be >= clip_low_threshold") - - original_shape = labels.shape - hidden_size = hidden_states.size(-1) - - labels_flat = labels.reshape(-1) - hidden_states_flat = hidden_states.reshape(-1, hidden_size) - old_logprobs_flat = old_logprobs.reshape(-1) - advantages_flat = advantages.reshape(-1) - - valid_mask = labels_flat != ignore_index - valid_mask_f = valid_mask.float() - valid_count = valid_mask.sum() - - if loss_reducer is None: - loss_reducer = TokenPartial(scale=valid_count.float()) - if metric_reducer is None: - metric_reducer = TokenPartial(scale=valid_count.float()) - - per_token_ce = compute_per_token_ce( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - ce_mode, - num_chunks, - tp_group=tp_group, - lm_head_fp32=lm_head_fp32, - lm_head=lm_head, - logprob_temperature=logprob_temperature, - logprob_top_k=logprob_top_k, - logprob_top_p=logprob_top_p, - logprob_min_p=logprob_min_p, - ) - - current_support = torch.isfinite(per_token_ce) - new_logprobs_flat = -per_token_ce.detach() - log_ratio = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) - ratio = torch.exp(log_ratio) - clipped_ratio = torch.clamp(ratio, clip_low_threshold, clip_high_threshold) - - coefficient = (clipped_ratio.detach() * advantages_flat).masked_fill(~valid_mask | ~current_support, 0.0) - safe_per_token_ce = torch.where(current_support, per_token_ce, torch.zeros_like(per_token_ce)) - per_token_loss_flat = coefficient * safe_per_token_ce - loss = loss_reducer(per_token_loss_flat, valid_mask_f) - - is_clipped = (clipped_ratio != ratio) & valid_mask - if valid_mask.any(): - ratio_min = ratio.masked_fill(~valid_mask, float("inf")).min() - ratio_max = ratio.masked_fill(~valid_mask, float("-inf")).max() - else: - ratio_min = ratio.new_tensor(float("inf")) - ratio_max = ratio.new_tensor(float("-inf")) - - metrics: Dict[str, Any] = { - "ratio_mean": metric_reducer(ratio, valid_mask_f).detach(), - "ratio_min": ratio_min.detach(), - "ratio_max": ratio_max.detach(), - "clip_fraction": metric_reducer(is_clipped.float(), valid_mask_f).detach(), - } - - metric_ops = {"ratio_min": "min", "ratio_max": "max"} - if compute_kl_stats: - with torch.no_grad(): - metric_log_ratio = torch.where( - current_support, - log_ratio, - torch.full_like(log_ratio, -20.0), - ) - per_token_k3 = ratio - metric_log_ratio - 1.0 - if valid_mask.any(): - k3_max = per_token_k3.masked_fill(~valid_mask, float("-inf")).max() - logratio_min = metric_log_ratio.masked_fill(~valid_mask, float("inf")).min() - logratio_max = metric_log_ratio.masked_fill(~valid_mask, float("-inf")).max() - abs_logratio_max = metric_log_ratio.abs().masked_fill(~valid_mask, float("-inf")).max() - else: - k3_max = per_token_k3.new_tensor(float("-inf")) - logratio_min = log_ratio.new_tensor(float("inf")) - logratio_max = log_ratio.new_tensor(float("-inf")) - abs_logratio_max = log_ratio.new_tensor(float("-inf")) - metrics["kl_sample_train_k3"] = metric_reducer(per_token_k3, valid_mask_f) - metrics["kl_k3_debug_mean"] = metric_reducer(per_token_k3, valid_mask_f) - metrics["kl_k3_debug_max"] = k3_max - metrics["kl_k3_debug_abs_logratio_mean"] = metric_reducer(metric_log_ratio.abs(), valid_mask_f) - metrics["kl_k3_debug_abs_logratio_max"] = abs_logratio_max - metrics["kl_k3_debug_logratio_mean"] = metric_reducer(metric_log_ratio, valid_mask_f) - metrics["kl_k3_debug_logratio_min"] = logratio_min - metrics["kl_k3_debug_logratio_max"] = logratio_max - metrics["kl_k3_debug_frac_logratio_positive"] = metric_reducer((metric_log_ratio > 0).float(), valid_mask_f) - for suffix, threshold in K3_DEBUG_THRESHOLDS: - metrics[f"kl_k3_debug_frac_gt_{suffix}"] = metric_reducer( - (per_token_k3 > threshold).float(), valid_mask_f - ) - metrics["entropy_sample"] = metric_reducer(-old_logprobs_flat, valid_mask_f) - metrics["valid_tokens"] = valid_count.item() - metrics["current_support_fraction"] = metric_reducer(current_support.float(), valid_mask_f) - metric_ops.update( - { - "kl_k3_debug_max": "max", - "kl_k3_debug_abs_logratio_max": "max", - "kl_k3_debug_logratio_min": "min", - "kl_k3_debug_logratio_max": "max", - } - ) - - return LossOutput( - loss=loss, - per_token_logprobs=new_logprobs_flat.view(original_shape), - per_token_loss=per_token_loss_flat.detach().view(original_shape), - metrics=metrics, - metric_ops=metric_ops, - ) +from xorl.objectives.cispo_loss import * # noqa: F401,F403 diff --git a/src/xorl/ops/loss/grpo_loss.py b/src/xorl/ops/loss/grpo_loss.py index faba6488..c5311b4f 100644 --- a/src/xorl/ops/loss/grpo_loss.py +++ b/src/xorl/ops/loss/grpo_loss.py @@ -1,262 +1,3 @@ -""" -DR-GRPO: "Done Right" GRPO Loss for RL Training. +"""Moved to ``xorl.objectives.grpo_loss`` (issue #78 phase 2); compat stub.""" -Reference: Liu et al., "Understanding R1-Zero-Like Training" (2025). -https://arxiv.org/abs/2503.20783 -""" - -from typing import List, Literal, Tuple - -import torch -import torch.distributed as dist - -from xorl.ops.exact_sampling_transforms import TOP_K_ALL -from xorl.ops.loss.loss_output import LossOutput -from xorl.ops.loss.per_token_ce import compute_per_token_ce -from xorl.ops.loss.reducers import Reducer, TokenPartial - - -KLType = Literal["k1", "k2", "k3"] -RatioType = Literal["token", "sequence"] - - -def compute_ratio( - logprobs: torch.Tensor, - generator_logprobs: torch.Tensor, - mask: torch.Tensor, - metric_reducer: Reducer, - ratio_type: RatioType = "token", -) -> Tuple[torch.Tensor, torch.Tensor, List[Tuple[str, torch.Tensor]]]: - """Importance sampling ratio r = π_θ/π_old. - - token: r_t = exp(logprobs_t - generator_logprobs_t) - sequence: r_seq = exp(mean_t[logprobs - generator_logprobs]), uses reparameterization. - """ - current_support = torch.isfinite(logprobs) - raw_token_log_ratio = torch.where( - current_support, - logprobs - generator_logprobs.detach(), - torch.full_like(logprobs, -torch.inf), - ) - if ratio_type == "token": - log_ratio = raw_token_log_ratio - ratio = torch.exp(log_ratio) - elif ratio_type == "sequence": - seq_lengths = mask.sum(dim=-1).clamp(min=1) - finite_token_log_ratio = torch.where( - current_support, - raw_token_log_ratio, - torch.zeros_like(raw_token_log_ratio), - ) - seq_log_ratio = (finite_token_log_ratio * mask).sum(dim=-1) / seq_lengths - seq_supported = (current_support | ~mask.bool()).all(dim=-1) - seq_log_ratio = torch.where( - seq_supported, - seq_log_ratio, - torch.full_like(seq_log_ratio, -torch.inf), - ) - - # Reparameterization: forward uses seq ratio, backward uses token grads - safe_logprobs = torch.where(current_support, logprobs, torch.zeros_like(logprobs)) - log_ratio = safe_logprobs - safe_logprobs.detach() + seq_log_ratio.detach().unsqueeze(-1) - ratio = torch.exp(log_ratio) - else: - raise ValueError(f"Unknown ratio_type: {ratio_type}") - - with torch.no_grad(): - metric_log_ratio = torch.where( - torch.isfinite(log_ratio), - log_ratio, - torch.full_like(log_ratio, -20.0), - ) - metrics = [ - ("loss/ratio/mean", metric_reducer(ratio, mask)), - ("loss/kl_policy/mean", metric_reducer(-metric_log_ratio, mask)), - ("loss/current_support_fraction", metric_reducer(current_support.float(), mask)), - ] - - return ratio, log_ratio, metrics - - -def compute_kl( - policy_logprobs: torch.Tensor, - ref_logprobs: torch.Tensor, - mask: torch.Tensor, - metric_reducer: Reducer, - kl_type: KLType = "k3", -) -> Tuple[torch.Tensor, List[Tuple[str, torch.Tensor]]]: - """KL divergence using Schulman's estimators (k1, k2, k3).""" - policy_support = torch.isfinite(policy_logprobs) - ref_support = torch.isfinite(ref_logprobs) - safe_policy_logprobs = torch.where(policy_support, policy_logprobs, torch.zeros_like(policy_logprobs)) - safe_ref_logprobs = torch.where(ref_support, ref_logprobs, torch.zeros_like(ref_logprobs)) - finite_log_ratio = safe_policy_logprobs - safe_ref_logprobs.detach() - log_ratio = torch.where( - policy_support & ref_support, - finite_log_ratio, - torch.where( - policy_support, - torch.full_like(finite_log_ratio, 20.0), - torch.where(ref_support, torch.full_like(finite_log_ratio, -20.0), torch.zeros_like(finite_log_ratio)), - ), - ) - - if kl_type == "k1": - kl = log_ratio - elif kl_type == "k2": - kl = 0.5 * log_ratio.square() - elif kl_type == "k3": - neg_log_ratio = torch.clamp(-log_ratio, min=-10.0, max=10.0) - ratio = torch.exp(neg_log_ratio) - kl = ratio - neg_log_ratio - 1 - else: - raise ValueError(f"Unknown kl_type: {kl_type}") - - with torch.no_grad(): - metrics = [("loss/kl_ref/mean", metric_reducer(kl, mask))] - - return kl, metrics - - -def pg_ppo_clip( - ratio: torch.Tensor, - advantages: torch.Tensor, - mask: torch.Tensor, - metric_reducer: Reducer, - clip_low: float = 0.2, - clip_high: float = 0.2, -) -> Tuple[torch.Tensor, List[Tuple[str, torch.Tensor]]]: - """PPO clipped surrogate: L = max(-r*A, -clip(r, 1-ε_low, 1+ε_high)*A).""" - clipped_ratio = torch.clamp(ratio, 1 - clip_low, 1 + clip_high) - unclipped_loss = -ratio * advantages - clipped_loss = -clipped_ratio * advantages - pg_loss = torch.maximum(unclipped_loss, clipped_loss) - - with torch.no_grad(): - mask_bool = mask.bool() - clipped_high = (ratio > 1 + clip_high) & mask_bool - clipped_low = (ratio < 1 - clip_low) & mask_bool - pos_adv = advantages > 0 - neg_adv = advantages < 0 - - metrics = [ - ("loss/clip/clipped_ratio/mean", metric_reducer(clipped_ratio, mask)), - ("loss/clip/high_fraction", metric_reducer((clipped_high & pos_adv).float(), mask)), - ("loss/clip/low_fraction", metric_reducer((clipped_low & neg_adv).float(), mask)), - ] - - return pg_loss, metrics - - -def drgrpo_loss_function( - hidden_states: torch.Tensor, - weight: torch.Tensor, - labels: torch.Tensor, - old_logprobs: torch.Tensor, - advantages: torch.Tensor, - ref_logprobs: torch.Tensor | None = None, - ignore_index: int = -100, - clip_low: float = 0.2, - clip_high: float = 0.28, - beta: float = 0.1, - ratio_type: RatioType = "token", - kl_type: KLType = "k3", - ce_mode: str = "compiled", - num_chunks: int = 8, - tp_group: dist.ProcessGroup | None = None, - lm_head_fp32: bool = False, - loss_reducer: Reducer | None = None, - metric_reducer: Reducer | None = None, - lm_head: torch.nn.Module | None = None, - logprob_temperature: float = 1.0, - logprob_top_k: int | torch.Tensor = TOP_K_ALL, - logprob_top_p: float | torch.Tensor = 1.0, - logprob_min_p: float | torch.Tensor = 0.0, -) -> LossOutput: - """DR-GRPO loss for RL training. - - Per-token: L_t = max(-r*A, -clip(r, 1-ε, 1+ε)*A) + β*KL - Aggregated: ``loss_reducer(per_token_loss, mask)``. Defaults to - ``TokenPartial(scale=loss_mask.sum())`` — the local active-token mean. - - Args: - hidden_states: (B, S, H) model hidden states. - weight: (V, H) or (V/tp, H) LM head weight. - labels: (B, S) target token IDs, already next-token aligned. - old_logprobs: (B, S) log probs from generation policy. - advantages: (B, S) per-token advantages. - ref_logprobs: (B, S) reference model log probs for KL (required if beta > 0). - ignore_index: Token ID to ignore (default: -100). - clip_low: Lower clip bound (default: 0.2). - clip_high: Upper clip bound (default: 0.28). - beta: KL penalty coefficient (default: 0.1). - ratio_type: Ratio type: "token" or "sequence" (default: "token"). - kl_type: KL estimator: "k1", "k2", "k3" (default: "k3"). - ce_mode: Cross-entropy mode: "compiled" or "eager". - num_chunks: Chunks for compiled mode. - tp_group: TP process group for vocab-parallel CE. - lm_head_fp32: Compute LM head in FP32. - loss_reducer / metric_reducer: Both default to - ``TokenPartial(scale=loss_mask.sum())`` (legacy local active-token - mean; does not compose across mbs/ranks). Pass shared global-scale - reducers to make summed partial shares recover the global value. - logprob_temperature: Temperature applied to trainer logits before - selected-token logprob calculation. ``1.0`` is raw policy logprobs; - setting this to the rollout temperature yields behavior-policy - semantics for DR-GRPO ratios and K3 stats. - - Returns: - LossOutput with loss, per_token_logprobs, per_token_loss, and metrics. - """ - if beta > 0 and ref_logprobs is None: - raise ValueError("ref_logprobs required when beta > 0") - - B, S = labels.shape - H = hidden_states.size(-1) - - labels_flat = labels.reshape(-1) - hidden_states_flat = hidden_states.reshape(-1, H) - - per_token_ce = compute_per_token_ce( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - ce_mode, - num_chunks, - tp_group=tp_group, - lm_head_fp32=lm_head_fp32, - lm_head=lm_head, - logprob_temperature=logprob_temperature, - logprob_top_k=logprob_top_k, - logprob_top_p=logprob_top_p, - logprob_min_p=logprob_min_p, - ) - logprobs = -per_token_ce.view(B, S) - - loss_mask = (labels != ignore_index).float() - - if metric_reducer is None: - metric_reducer = TokenPartial(scale=loss_mask.sum()) - if loss_reducer is None: - loss_reducer = TokenPartial(scale=loss_mask.sum()) - - ratio, _, ratio_m = compute_ratio(logprobs, old_logprobs, loss_mask, metric_reducer, ratio_type) - - pg_loss, clip_m = pg_ppo_clip(ratio, advantages, loss_mask, metric_reducer, clip_low, clip_high) - - kl_m: List[Tuple[str, torch.Tensor]] = [] - if beta > 0: - kl, kl_m = compute_kl(logprobs, ref_logprobs, loss_mask, metric_reducer, kl_type) - pg_loss = pg_loss + beta * kl - - loss = loss_reducer(pg_loss, loss_mask) - - metrics = dict(ratio_m + clip_m + kl_m) - - return LossOutput( - loss=loss, - per_token_logprobs=logprobs.detach(), - per_token_loss=pg_loss.detach(), - metrics=metrics, - ) +from xorl.objectives.grpo_loss import * # noqa: F401,F403 diff --git a/src/xorl/ops/loss/importance_sampling_loss.py b/src/xorl/ops/loss/importance_sampling_loss.py index e61e616b..bb153c26 100644 --- a/src/xorl/ops/loss/importance_sampling_loss.py +++ b/src/xorl/ops/loss/importance_sampling_loss.py @@ -1,216 +1,3 @@ -from __future__ import annotations +"""Moved to ``xorl.objectives.importance_sampling_loss`` (issue #78 phase 2); compat stub.""" -from typing import Any, Dict, Optional - -import torch -import torch.distributed as dist - -from xorl.ops.exact_sampling_transforms import TOP_K_ALL -from xorl.ops.loss.loss_output import LossOutput -from xorl.ops.loss.per_token_ce import compute_per_token_ce -from xorl.ops.loss.reducers import Reducer, TokenPartial - - -K3_DEBUG_THRESHOLDS = ( - ("1e_minus_6", 1e-6), - ("1e_minus_4", 1e-4), - ("1e_minus_3", 1e-3), - ("1e_minus_2", 1e-2), - ("1e_minus_1", 1e-1), - ("1", 1.0), -) - - -def importance_sampling_loss_function( - hidden_states: torch.Tensor, - weight: torch.Tensor, - labels: torch.Tensor, - old_logprobs: torch.Tensor, - advantages: torch.Tensor, - ignore_index: int = -100, - num_chunks: int = 8, - ce_mode: str = "compiled", - return_per_token: bool = False, - tp_group: Optional[dist.ProcessGroup] = None, - compute_kl_stats: bool = False, - lm_head_fp32: bool = False, - loss_reducer: Optional[Reducer] = None, - metric_reducer: Optional[Reducer] = None, - lm_head: Optional[torch.nn.Module] = None, - logprob_temperature: float = 1.0, - logprob_top_k: int | torch.Tensor = TOP_K_ALL, - logprob_top_p: float | torch.Tensor = 1.0, - logprob_min_p: float | torch.Tensor = 0.0, -) -> "LossOutput": - """ - Compute importance sampling loss for GRPO/RL training. - - This implements the Tinker-style importance sampling loss: - prob_ratio = exp(new_logprobs - old_logprobs) - loss = -(prob_ratio * advantages).mean() - - Supports multiple computation modes: - - "compiled": RECOMMENDED. torch.compile (1.6x speed, 16% memory) - - "eager": Simple F.cross_entropy baseline (may OOM at 32K) - - Args: - hidden_states: Model hidden states, shape (batch, seq_len, hidden_dim) - weight: LM head weight matrix, shape (vocab_size, hidden_dim) - labels: Target token IDs, shape (batch, seq_len). Already next-token aligned. - old_logprobs: Old policy log probabilities from sampling, shape (batch, seq_len) - advantages: Per-token advantages, shape (batch, seq_len) - ignore_index: Index to ignore in loss computation (default: -100) - num_chunks: Number of chunks for compiled mode (default: 8). - ce_mode: Cross-entropy mode - "compiled" (default) or "eager" - return_per_token: If True, returns per-token logprobs and per-token CE loss. - Useful for custom loss computations. - compute_kl_stats: If True, compute and return KL statistics in metrics dict: - - kl_sample_train_k3: Schulman's K3 estimator: mean(exp(log_ratio) - log_ratio - 1) - where log_ratio = new_logprobs - old_logprobs. Non-negative, unbiased, lower variance. - - entropy_sample: -mean(old_logprobs) over valid tokens - - valid_tokens: Count of valid tokens - loss_reducer: Reduces per-token loss to a scalar partial share. None => - ``TokenPartial(scale=valid_mask.sum())`` (legacy local token-mean; does - not compose across micro-batches/ranks). Pass a shared global-scale - reducer to make summed partial shares recover the global loss. - metric_reducer: Reduces per-token /mean metrics (ratio_mean, - kl_sample_train_k3, entropy_sample). ratio_min/ratio_max stay local - scalars and bypass it. - logprob_temperature: Temperature applied to trainer logits before - selected-token logprob calculation. ``1.0`` is raw policy logprobs; - setting this to the rollout temperature yields behavior-policy - semantics for the sampled-token ratio. - - Returns: - LossOutput with loss, per_token_logprobs, per_token_loss, and metrics. - """ - original_shape = labels.shape - H = hidden_states.size(-1) - - # Flatten tensors - labels_flat = labels.reshape(-1) - hidden_states_flat = hidden_states.reshape(-1, H) - old_logprobs_flat = old_logprobs.reshape(-1) - advantages_flat = advantages.reshape(-1) - - # Valid/action mask - valid_mask = labels_flat != ignore_index - valid_mask_f = valid_mask.float() - valid_count = valid_mask.sum() - - if loss_reducer is None: - loss_reducer = TokenPartial(scale=valid_count.float()) - if metric_reducer is None: - metric_reducer = TokenPartial(scale=valid_count.float()) - - # ---- Cross-entropy computation ---- - per_token_ce = compute_per_token_ce( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - ce_mode, - num_chunks, - tp_group=tp_group, - lm_head_fp32=lm_head_fp32, - lm_head=lm_head, - logprob_temperature=logprob_temperature, - logprob_top_k=logprob_top_k, - logprob_top_p=logprob_top_p, - logprob_min_p=logprob_min_p, - ) - - current_support = torch.isfinite(per_token_ce) - # new logprobs = log p(target) = -CE - new_logprobs_flat = -per_token_ce.detach() - - # ---- ratio computation (no sanitization) ---- - delta = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) - delta = torch.where(current_support, delta.clamp(min=-20.0, max=20.0), torch.full_like(delta, -torch.inf)) - ratio = torch.exp(delta) - - # ---- Per-token policy gradient loss: -(ratio * advantages) ---- - per_token_pg = -(ratio * advantages_flat) - per_token_pg = per_token_pg.masked_fill(~valid_mask, 0.0) - - # ---- Option B: value from true PG, grad from weighted CE surrogate ---- - true_pg = loss_reducer(per_token_pg, valid_mask_f) - - w = (ratio.detach() * advantages_flat).masked_fill(~valid_mask | ~current_support, 0.0) - safe_per_token_ce = torch.where(current_support, per_token_ce, torch.zeros_like(per_token_ce)) - surrogate = loss_reducer(w * safe_per_token_ce, valid_mask_f) - - loss = true_pg.detach() + surrogate - surrogate.detach() - - # ±inf identity on empty ranks lets cross-rank MIN/MAX-allreduce ignore empty contributors. - if valid_mask.any(): - ratio_min = ratio.masked_fill(~valid_mask, float("inf")).min() - ratio_max = ratio.masked_fill(~valid_mask, float("-inf")).max() - else: - ratio_min = ratio.new_tensor(float("inf")) - ratio_max = ratio.new_tensor(float("-inf")) - metrics: Dict[str, Any] = { - "ratio_mean": metric_reducer(ratio, valid_mask_f).detach(), - "ratio_min": ratio_min.detach(), - "ratio_max": ratio_max.detach(), - } - - if compute_kl_stats: - with torch.no_grad(): - raw_log_ratio = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) - log_ratio_full = torch.where( - current_support, - raw_log_ratio, - torch.full_like(raw_log_ratio, -20.0), - ) - ratio_full = torch.where(current_support, torch.exp(raw_log_ratio), torch.zeros_like(raw_log_ratio)) - per_token_k3 = ratio_full - log_ratio_full - 1.0 - if valid_mask.any(): - k3_max = per_token_k3.masked_fill(~valid_mask, float("-inf")).max() - logratio_min = log_ratio_full.masked_fill(~valid_mask, float("inf")).min() - logratio_max = log_ratio_full.masked_fill(~valid_mask, float("-inf")).max() - abs_logratio_max = log_ratio_full.abs().masked_fill(~valid_mask, float("-inf")).max() - else: - k3_max = per_token_k3.new_tensor(float("-inf")) - logratio_min = log_ratio_full.new_tensor(float("inf")) - logratio_max = log_ratio_full.new_tensor(float("-inf")) - abs_logratio_max = log_ratio_full.new_tensor(float("-inf")) - metrics["kl_sample_train_k3"] = metric_reducer(per_token_k3, valid_mask_f) - metrics["kl_k3_debug_mean"] = metric_reducer(per_token_k3, valid_mask_f) - metrics["kl_k3_debug_max"] = k3_max - metrics["kl_k3_debug_abs_logratio_mean"] = metric_reducer(log_ratio_full.abs(), valid_mask_f) - metrics["kl_k3_debug_abs_logratio_max"] = abs_logratio_max - metrics["kl_k3_debug_logratio_mean"] = metric_reducer(log_ratio_full, valid_mask_f) - metrics["kl_k3_debug_logratio_min"] = logratio_min - metrics["kl_k3_debug_logratio_max"] = logratio_max - metrics["kl_k3_debug_frac_logratio_positive"] = metric_reducer((log_ratio_full > 0).float(), valid_mask_f) - for suffix, threshold in K3_DEBUG_THRESHOLDS: - metrics[f"kl_k3_debug_frac_gt_{suffix}"] = metric_reducer( - (per_token_k3 > threshold).float(), valid_mask_f - ) - metrics["entropy_sample"] = metric_reducer(-old_logprobs_flat, valid_mask_f) - metrics["valid_tokens"] = valid_count.item() - metrics["current_support_fraction"] = metric_reducer(current_support.float(), valid_mask_f) - - # Reshape per-token outputs - per_token_logprobs = new_logprobs_flat.view(original_shape) - per_token_loss = per_token_pg.view(original_shape) - - metric_ops = {"ratio_min": "min", "ratio_max": "max"} - if compute_kl_stats: - metric_ops.update( - { - "kl_k3_debug_max": "max", - "kl_k3_debug_abs_logratio_max": "max", - "kl_k3_debug_logratio_min": "min", - "kl_k3_debug_logratio_max": "max", - } - ) - - return LossOutput( - loss=loss, - per_token_logprobs=per_token_logprobs, - per_token_loss=per_token_loss, - metrics=metrics, - metric_ops=metric_ops, - ) +from xorl.objectives.importance_sampling_loss import * # noqa: F401,F403 diff --git a/src/xorl/ops/loss/loss_output.py b/src/xorl/ops/loss/loss_output.py index 41a79eee..77b8f5da 100644 --- a/src/xorl/ops/loss/loss_output.py +++ b/src/xorl/ops/loss/loss_output.py @@ -1,21 +1,3 @@ -from dataclasses import dataclass -from typing import Any, Dict, Optional +"""Moved to ``xorl.objectives.loss_output`` (issue #78 phase 2); compat stub.""" -import torch - - -@dataclass -class LossOutput: - """Standardized return type for all loss functions. - - ``metric_ops`` tags ``metrics`` keys whose cross-mb / cross-rank composition - isn't the default mean (``"min"``/``"max"``). The sidecar (rather than a - tagged-value type in ``metrics``) keeps the metrics dict directly - JSON-serializable for untagged consumers. - """ - - loss: torch.Tensor - per_token_logprobs: Optional[torch.Tensor] = None - per_token_loss: Optional[torch.Tensor] = None - metrics: Optional[Dict[str, Any]] = None - metric_ops: Optional[Dict[str, str]] = None +from xorl.objectives.loss_output import * # noqa: F401,F403 diff --git a/src/xorl/ops/loss/opd_loss.py b/src/xorl/ops/loss/opd_loss.py index 01ba9851..9b16a563 100644 --- a/src/xorl/ops/loss/opd_loss.py +++ b/src/xorl/ops/loss/opd_loss.py @@ -1,1324 +1,3 @@ -from __future__ import annotations +"""Moved to ``xorl.objectives.opd_loss`` (issue #78 phase 2); compat stub.""" -from dataclasses import dataclass -from typing import Callable, Optional - -import torch -import torch.distributed as dist -import torch.nn.functional as F - -from xorl.ops.loss.compiled_cross_entropy import ( - compiled_forward_kl_full_function, - compiled_forward_kl_full_with_diag_function, - compiled_reverse_kl_function, - compiled_reverse_kl_with_diag_function, - compiled_sampled_token_logprobs_function, -) -from xorl.ops.loss.loss_output import LossOutput -from xorl.ops.loss.opd_streaming_kl import ( - streaming_forward_kl_function, - streaming_forward_kl_lowmem_function, - streaming_full_vocab_diagnostics, - streaming_reverse_kl_function, - streaming_reverse_kl_lowmem_function, -) -from xorl.ops.loss.reducers import Reducer, TokenPartial -from xorl.ops.loss.vocab_parallel_reverse_kl import vocab_parallel_reverse_kl_gathered - - -# OPD loss modes. -# -# `reverse_kl_full` (default) and `forward_kl_full` are full-vocabulary KL -# variants computed from the *full* teacher distribution — xorl ships teacher -# hidden states + the cached teacher LM head separately, so the full p_T is -# materialized at loss time. VERL's `forward_kl_topk` is intentionally NOT -# ported: it's a bandwidth workaround for teacher-emits-truncated-logprobs -# pipelines, strictly worse than `forward_kl_full` for the same student compute. -# -# Single-sample KL estimators are ported from verl/trainer/ppo/core_algos.py:: -# kl_penalty as a different (cheap, sampled-token-only) point in the -# memory/quality trade-off. A trailing "+" applies the k2 straight-through- -# gradient trick (same forward, k2 gradient). -LOSS_MODE_REVERSE_KL_FULL = "reverse_kl_full" -LOSS_MODE_FORWARD_KL_FULL = "forward_kl_full" -_ESTIMATOR_MODES = {"kl", "k1", "abs", "mse", "k2", "low_var_kl", "k3"} - - -def _strip_estimator_plus(loss_mode: str) -> tuple[str, bool]: - if loss_mode.endswith("+"): - return loss_mode[:-1], True - return loss_mode, False - - -def is_estimator_loss_mode(loss_mode: str) -> bool: - base, _ = _strip_estimator_plus(loss_mode) - return base in _ESTIMATOR_MODES - - -def _kl_penalty_estimator( - logprob: torch.Tensor, - ref_logprob: torch.Tensor, - loss_mode: str, -) -> torch.Tensor: - """Single-sample KL estimators, byte-for-byte port of VERL's kl_penalty. - - See verl/trainer/ppo/core_algos.py::kl_penalty and ::kl_penalty_forward. - A trailing "+" (e.g. "k3+") applies the k2 straight-through trick: forward - value from the chosen estimator, but backward as if it were 0.5*(logp-ref)^2. - """ - base, straight_through = _strip_estimator_plus(loss_mode) - forward = _kl_penalty_forward(logprob, ref_logprob, base) - if not straight_through or base in ("mse", "k2"): - return forward - backward_score = 0.5 * (logprob - ref_logprob).square() - return backward_score - backward_score.detach() + forward.detach() - - -def _kl_penalty_forward( - logprob: torch.Tensor, - ref_logprob: torch.Tensor, - loss_mode: str, -) -> torch.Tensor: - if loss_mode in ("kl", "k1"): - return logprob - ref_logprob - if loss_mode == "abs": - return (logprob - ref_logprob).abs() - if loss_mode in ("mse", "k2"): - return 0.5 * (logprob - ref_logprob).square() - if loss_mode in ("low_var_kl", "k3"): - kl = ref_logprob - logprob - kl = torch.clamp(kl, min=-20, max=20) - ratio = torch.exp(kl) - kld = (ratio - kl - 1).contiguous() - return torch.clamp(kld, min=-10, max=10) - raise ValueError(f"Unknown KL estimator '{loss_mode}'") - - -@dataclass(frozen=True) -class OPDLossMetrics: - """OPD loss metrics emitted per micro-batch. - - All fields are always present in `to_dict()` (defaulting to 0.0 / 0 when - not applicable to the active loss_mode). This is intentional — dict-keyed - distributed all-reduces deadlock when ranks have different key sets, so - every loss path on every rank must contribute the same dict shape. - """ - - valid_tokens: int - opd_kl: float = 0.0 - opd_weighted_kl: float = 0.0 - opd_hidden_match_loss: float = 0.0 - opd_hidden_match_raw_loss: float = 0.0 - opd_hidden_match_weight_mean: float = 0.0 - opd_hidden_match_pos_loss: float = 0.0 - opd_hidden_match_neg_loss: float = 0.0 - opd_hidden_match_pos_raw_loss: float = 0.0 - opd_hidden_match_neg_raw_loss: float = 0.0 - opd_hidden_match_neg_minus_pos_raw: float = 0.0 - opd_hidden_match_pos_weight_mean: float = 0.0 - opd_hidden_match_neg_weight_mean: float = 0.0 - opd_teacher_weight_mean: float = 0.0 - opd_num_teachers: int = 0 - # ---- Full-vocab diagnostics (reverse_kl_full / forward_kl_full) ---- - opd_teacher_entropy: float = 0.0 - opd_student_entropy: float = 0.0 - opd_top1_agreement: float = 0.0 - # ---- KL-estimator diagnostic (k1/abs/...) ---- - opd_abs_loss: float = 0.0 - # ---- Distillation-loss range ---- - opd_loss_min: float = 0.0 - opd_loss_max: float = 0.0 - opd_loss_abs_mean: float = 0.0 - # Fraction of valid tokens whose pre-clamp |KL| >= loss_max_clamp. Clamped - # tokens pass ZERO gradient (gradient-dead mass) — watch this when the loss - # plateaus while opd_loss_max sits pinned at the clamp. - opd_loss_clamp_frac: float = 0.0 - # ---- Region / sample-correctness KL splits (diagnostic) ---- - # All `*_per_valid` fields are masked SUMS divided by the micro-batch's TOTAL - # valid-token count (not per-region means), and every `opd_frac_*` is a - # region-token count over the same denominator. Both are linear functionals - # of per-token values, so the valid-token-weighted mean aggregation across - # micro-batches / ranks recomposes them EXACTLY; derive the human-readable - # region mean as `per_valid / frac` after aggregation (the client does this). - # Regions come from `diag_region_ids` (0=prompt, 1=buffer, 2=answer, -1=n/a); - # correctness from `diag_sample_ok` (1=sampled answer correct, 0=wrong, - # -1=unknown), both client-provided per-token tensors. - opd_kl_prompt_per_valid: float = 0.0 - opd_kl_buffer_per_valid: float = 0.0 - opd_kl_answer_per_valid: float = 0.0 - opd_frac_prompt: float = 0.0 - opd_frac_buffer: float = 0.0 - opd_frac_answer: float = 0.0 - opd_kl_answer_correct_per_valid: float = 0.0 - opd_kl_answer_wrong_per_valid: float = 0.0 - opd_frac_answer_correct: float = 0.0 - opd_frac_answer_wrong: float = 0.0 - opd_student_entropy_answer_correct_per_valid: float = 0.0 - opd_student_entropy_answer_wrong_per_valid: float = 0.0 - opd_teacher_entropy_answer_correct_per_valid: float = 0.0 - opd_teacher_entropy_answer_wrong_per_valid: float = 0.0 - # ---- Multi-layer OPRD (all-layer hidden matching) ---- - opd_oprd_loss: float = 0.0 - opd_oprd_raw_loss: float = 0.0 - opd_oprd_num_layers: int = 0 - # ---- Vocab-parallel OPD debugging ---- - # These are raw numerator/count diagnostics from the vocab-parallel group, - # not values used by normal metric aggregation. They stay in the fixed - # metric key set so optional debug consumers do not introduce rank-skewed - # metric dictionaries. - opd_vocab_parallel_group_tokens: int = 0 - opd_vocab_parallel_kl_sum: float = 0.0 - opd_vocab_parallel_weighted_kl_sum: float = 0.0 - # ---- PG-mode (use_policy_gradient=True) ---- - opd_pg_clipfrac: float = 0.0 - opd_pg_clipfrac_lower: float = 0.0 - opd_ppo_kl: float = 0.0 - - def to_dict(self) -> dict[str, int | float]: - return { - "valid_tokens": self.valid_tokens, - "opd_kl": self.opd_kl, - "opd_weighted_kl": self.opd_weighted_kl, - "opd_hidden_match_loss": self.opd_hidden_match_loss, - "opd_hidden_match_raw_loss": self.opd_hidden_match_raw_loss, - "opd_hidden_match_weight_mean": self.opd_hidden_match_weight_mean, - "opd_hidden_match_pos_loss": self.opd_hidden_match_pos_loss, - "opd_hidden_match_neg_loss": self.opd_hidden_match_neg_loss, - "opd_hidden_match_pos_raw_loss": self.opd_hidden_match_pos_raw_loss, - "opd_hidden_match_neg_raw_loss": self.opd_hidden_match_neg_raw_loss, - "opd_hidden_match_neg_minus_pos_raw": self.opd_hidden_match_neg_minus_pos_raw, - "opd_hidden_match_pos_weight_mean": self.opd_hidden_match_pos_weight_mean, - "opd_hidden_match_neg_weight_mean": self.opd_hidden_match_neg_weight_mean, - "opd_teacher_weight_mean": self.opd_teacher_weight_mean, - "opd_num_teachers": self.opd_num_teachers, - "opd_teacher_entropy": self.opd_teacher_entropy, - "opd_student_entropy": self.opd_student_entropy, - "opd_top1_agreement": self.opd_top1_agreement, - "opd_abs_loss": self.opd_abs_loss, - "opd_loss_min": self.opd_loss_min, - "opd_loss_max": self.opd_loss_max, - "opd_loss_abs_mean": self.opd_loss_abs_mean, - "opd_loss_clamp_frac": self.opd_loss_clamp_frac, - "opd_kl_prompt_per_valid": self.opd_kl_prompt_per_valid, - "opd_kl_buffer_per_valid": self.opd_kl_buffer_per_valid, - "opd_kl_answer_per_valid": self.opd_kl_answer_per_valid, - "opd_frac_prompt": self.opd_frac_prompt, - "opd_frac_buffer": self.opd_frac_buffer, - "opd_frac_answer": self.opd_frac_answer, - "opd_kl_answer_correct_per_valid": self.opd_kl_answer_correct_per_valid, - "opd_kl_answer_wrong_per_valid": self.opd_kl_answer_wrong_per_valid, - "opd_frac_answer_correct": self.opd_frac_answer_correct, - "opd_frac_answer_wrong": self.opd_frac_answer_wrong, - "opd_student_entropy_answer_correct_per_valid": self.opd_student_entropy_answer_correct_per_valid, - "opd_student_entropy_answer_wrong_per_valid": self.opd_student_entropy_answer_wrong_per_valid, - "opd_teacher_entropy_answer_correct_per_valid": self.opd_teacher_entropy_answer_correct_per_valid, - "opd_teacher_entropy_answer_wrong_per_valid": self.opd_teacher_entropy_answer_wrong_per_valid, - "opd_oprd_loss": self.opd_oprd_loss, - "opd_oprd_raw_loss": self.opd_oprd_raw_loss, - "opd_oprd_num_layers": self.opd_oprd_num_layers, - "opd_vocab_parallel_group_tokens": self.opd_vocab_parallel_group_tokens, - "opd_vocab_parallel_kl_sum": self.opd_vocab_parallel_kl_sum, - "opd_vocab_parallel_weighted_kl_sum": self.opd_vocab_parallel_weighted_kl_sum, - "opd_pg_clipfrac": self.opd_pg_clipfrac, - "opd_pg_clipfrac_lower": self.opd_pg_clipfrac_lower, - "opd_ppo_kl": self.opd_ppo_kl, - } - - -def _as_flat_optional_weights( - teacher_weights: Optional[torch.Tensor], - valid_mask: torch.Tensor, - dtype: torch.dtype, -) -> torch.Tensor: - if teacher_weights is None: - return torch.ones(valid_mask.sum(), dtype=dtype, device=valid_mask.device) - weights_flat = teacher_weights.reshape(-1).to(device=valid_mask.device, dtype=dtype) - return weights_flat[valid_mask] - - -def _zero_loss_with_graph(hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - """Build a 0-valued loss that still flows gradients through hidden_states + weight. - - Always returns fp32 so the dtype matches the normal-return path - (`total_weighted_kl / denom`, fp32). A dtype mismatch between the early-return - branch on no-valid-token ranks and the fp32 normal branch corrupts NCCL - all_reduce in the trainer's loss-reporting path. - """ - - def anchor(tensor: torch.Tensor) -> torch.Tensor: - if tensor.numel() == 0: - return torch.zeros((), dtype=torch.float32, device=tensor.device) - return tensor.reshape(-1)[:1].float().sum() * 0.0 - - return anchor(hidden_states) + anchor(weight) - - -def _denominator_tensor( - denominator: torch.Tensor | int | float | None, - *, - fallback: torch.Tensor, - device: torch.device | str, -) -> torch.Tensor: - if denominator is None: - return fallback.to(device=device, dtype=torch.float32) - if torch.is_tensor(denominator): - return denominator.to(device=device, dtype=torch.float32) - return torch.tensor(float(denominator), device=device, dtype=torch.float32) - - -def _gather_1d_nograd(local_values: torch.Tensor, *, group) -> torch.Tensor: - if local_values.ndim != 1: - raise ValueError(f"expected rank-1 tensor, got shape {tuple(local_values.shape)}") - if not dist.is_available() or not dist.is_initialized(): - return local_values.detach() - world = dist.get_world_size(group) - if world == 1: - return local_values.detach() - - local_count = torch.tensor([local_values.shape[0]], dtype=torch.long, device=local_values.device) - counts = torch.empty(world, dtype=torch.long, device=local_values.device) - dist.all_gather_into_tensor(counts, local_count, group=group) - max_count = int(counts.max().item()) if counts.numel() else 0 - if max_count == 0: - return local_values.detach().new_empty((0,)) - - padded = local_values.detach().new_zeros((max_count,)) - if local_values.numel(): - padded[: local_values.shape[0]].copy_(local_values.detach().contiguous()) - gathered = local_values.detach().new_empty((world * max_count,)) - dist.all_gather_into_tensor(gathered, padded.contiguous(), group=group) - pieces = [] - for rank, n_rows in enumerate(int(x) for x in counts.detach().cpu().tolist()): - if n_rows: - start = rank * max_count - pieces.append(gathered[start : start + n_rows]) - if not pieces: - return local_values.detach().new_empty((0,)) - return torch.cat(pieces, dim=0) - - -def _gathered_local_bounds(n_local: int, *, device: torch.device, group) -> tuple[int, int]: - """Return this rank's [start, end) slice in the gathered 1-D token order.""" - if not dist.is_available() or not dist.is_initialized(): - return 0, int(n_local) - world = dist.get_world_size(group) - if world == 1: - return 0, int(n_local) - - rank = dist.get_rank(group) - local_count = torch.tensor([int(n_local)], dtype=torch.long, device=device) - counts = torch.empty(world, dtype=torch.long, device=device) - dist.all_gather_into_tensor(counts, local_count, group=group) - start = int(counts[:rank].sum().item()) - return start, start + int(n_local) - - -def _value_divided_without_grad_scale(loss: torch.Tensor, divisor: int) -> torch.Tensor: - if divisor <= 1: - return loss - reported = loss / float(divisor) - return loss + (reported - loss).detach() - - -def _oprd_hidden_distance( - student_layer_hidden_states: torch.Tensor, - teacher_layer_hidden_states: torch.Tensor, - expected_rows: int, - *, - layer_chunk_size: int = 4, -) -> tuple[torch.Tensor, int]: - """Return per-token all-layer OPRD MSE without materializing all layers as fp32.""" - if student_layer_hidden_states.shape != teacher_layer_hidden_states.shape: - raise ValueError( - "OPRD requires matching student/teacher layer shapes, got " - f"student={tuple(student_layer_hidden_states.shape)} teacher={tuple(teacher_layer_hidden_states.shape)}" - ) - if student_layer_hidden_states.ndim != 3: - raise ValueError( - "OPRD layer tensors must be rank-3 [valid_tokens, layers, hidden], got " - f"{tuple(student_layer_hidden_states.shape)}" - ) - if student_layer_hidden_states.shape[0] != expected_rows: - raise ValueError( - f"OPRD layer tensors must have {expected_rows} rows, got {student_layer_hidden_states.shape[0]}" - ) - - return _oprd_hidden_distance_from_fetcher( - student_layer_hidden_states, - teacher_layer_fetcher=lambda start, end: teacher_layer_hidden_states.detach()[:, start:end, :], - expected_rows=expected_rows, - num_layers=int(student_layer_hidden_states.shape[1]), - layer_chunk_size=layer_chunk_size, - ) - - -def _oprd_hidden_distance_from_fetcher( - student_layer_hidden_states: torch.Tensor, - *, - teacher_layer_fetcher: Callable[[int, int], torch.Tensor], - expected_rows: int, - num_layers: int, - layer_chunk_size: int = 4, -) -> tuple[torch.Tensor, int]: - """Return per-token all-layer OPRD MSE while fetching teacher layers by chunk.""" - if student_layer_hidden_states.ndim != 3: - raise ValueError( - "OPRD layer tensors must be rank-3 [valid_tokens, layers, hidden], got " - f"{tuple(student_layer_hidden_states.shape)}" - ) - if student_layer_hidden_states.shape[0] != expected_rows: - raise ValueError( - f"OPRD layer tensors must have {expected_rows} rows, got {student_layer_hidden_states.shape[0]}" - ) - if int(student_layer_hidden_states.shape[1]) != int(num_layers): - raise ValueError( - f"OPRD student layer count {student_layer_hidden_states.shape[1]} " - f"does not match teacher layer count {num_layers}" - ) - if expected_rows == 0: - return torch.empty(0, dtype=torch.float32, device=student_layer_hidden_states.device), num_layers - if num_layers <= 0: - raise ValueError("OPRD layer tensors must include at least one layer") - - chunk_size = max(1, int(layer_chunk_size)) - hidden_distance_sum = None - for start in range(0, num_layers, chunk_size): - end = min(start + chunk_size, num_layers) - student_chunk = student_layer_hidden_states[:, start:end, :].float() - teacher_chunk = teacher_layer_fetcher(start, end).detach() - expected_shape = student_layer_hidden_states.shape[0], end - start, student_layer_hidden_states.shape[2] - if tuple(teacher_chunk.shape) != tuple(expected_shape): - raise ValueError( - "OPRD teacher layer fetcher returned wrong shape: " - f"got {tuple(teacher_chunk.shape)}, expected {tuple(expected_shape)}" - ) - teacher_chunk = teacher_chunk.float() - chunk_distance = (student_chunk - teacher_chunk).square().mean(dim=-1).sum(dim=-1) - hidden_distance_sum = chunk_distance if hidden_distance_sum is None else hidden_distance_sum + chunk_distance - return hidden_distance_sum / float(num_layers), num_layers - - -def opd_vocab_parallel_loss_function( - student_hidden_flat: torch.Tensor, - student_weight_local: torch.Tensor, - labels: torch.Tensor, - teacher_hidden_flat: torch.Tensor, - teacher_weight_local: torch.Tensor, - teacher_weights: Optional[torch.Tensor] = None, - hidden_match_weights: Optional[torch.Tensor] = None, - ignore_index: int = -100, - lm_head_fp32: bool = False, - teacher_lm_head_fp32: bool = True, - loss_reducer: Optional[Reducer] = None, - metric_reducer: Optional[Reducer] = None, - loss_mode: str = LOSS_MODE_REVERSE_KL_FULL, - loss_max_clamp: Optional[float] = None, - use_task_rewards: bool = False, - distillation_loss_coef: float = 1.0, - hidden_match_coef: float = 0.0, - kl_loss_weight: float = 1.0, - hidden_match_mode: str = "cosine", - teacher_layer_hidden_states: Optional[torch.Tensor] = None, - teacher_layer_fetcher: Optional[Callable[[int, int], torch.Tensor]] = None, - teacher_layer_num_layers: Optional[int] = None, - oprd_layer_chunk_size: int = 4, - student_layer_hidden_states: Optional[torch.Tensor] = None, - group=None, - debug_token_outputs: bool = False, -) -> LossOutput: - """OPD reverse-KL over vocab-sharded student/teacher LM heads. - - Inputs are already restricted to this rank's valid tokens for one teacher. - The vocab-parallel KL gathers token activations across ``group`` and returns - a replicated full-token KL vector. Its loss value is divided by the group - size for detached reporting only; gradients remain the full global-token sum - and are normalized later by the trainer's global-valid-token scale. - """ - if loss_mode != LOSS_MODE_REVERSE_KL_FULL: - raise ValueError("vocab-parallel OPD currently supports only loss_mode='reverse_kl_full'") - if labels.ndim != 1: - raise ValueError(f"vocab-parallel OPD labels must be rank-1, got {tuple(labels.shape)}") - if student_hidden_flat.ndim != 2 or teacher_hidden_flat.ndim != 2: - raise ValueError( - "vocab-parallel OPD hidden tensors must be rank-2, got " - f"student={tuple(student_hidden_flat.shape)} teacher={tuple(teacher_hidden_flat.shape)}" - ) - if student_hidden_flat.shape[0] != labels.shape[0] or teacher_hidden_flat.shape[0] != labels.shape[0]: - raise ValueError( - "vocab-parallel OPD token counts must match, got " - f"student={student_hidden_flat.shape[0]} teacher={teacher_hidden_flat.shape[0]} labels={labels.shape[0]}" - ) - if student_hidden_flat.shape[-1] != student_weight_local.shape[-1]: - raise ValueError( - "student hidden size " - f"({student_hidden_flat.shape[-1]}) must match local student head width " - f"({student_weight_local.shape[-1]})" - ) - if teacher_hidden_flat.shape[-1] != teacher_weight_local.shape[-1]: - raise ValueError( - "teacher hidden size " - f"({teacher_hidden_flat.shape[-1]}) must match local teacher head width " - f"({teacher_weight_local.shape[-1]})" - ) - if teacher_weight_local.shape[0] != student_weight_local.shape[0]: - raise ValueError( - "local student and teacher vocab shards must have the same rows, got " - f"{student_weight_local.shape[0]} and {teacher_weight_local.shape[0]}" - ) - - local_tokens = int(labels.shape[0]) - device = student_hidden_flat.device - if teacher_weights is None: - local_token_weights = torch.ones(local_tokens, dtype=torch.float32, device=device) - else: - if teacher_weights.shape[0] != local_tokens: - raise ValueError(f"teacher_weights must have {local_tokens} rows, got {tuple(teacher_weights.shape)}") - local_token_weights = teacher_weights.to(device=device, dtype=torch.float32) - if hidden_match_weights is None: - local_hidden_weights = local_token_weights - else: - if hidden_match_weights.shape[0] != local_tokens: - raise ValueError( - f"hidden_match_weights must have {local_tokens} rows, got {tuple(hidden_match_weights.shape)}" - ) - local_hidden_weights = hidden_match_weights.to(device=device, dtype=torch.float32) - - if lm_head_fp32: - student_hidden_for_kl = student_hidden_flat.float() - student_weight_for_kl = student_weight_local.float() - else: - student_hidden_for_kl = student_hidden_flat - student_weight_for_kl = student_weight_local - if teacher_lm_head_fp32: - teacher_hidden_for_kl = teacher_hidden_flat.float() - teacher_weight_for_kl = teacher_weight_local.float() - else: - teacher_hidden_for_kl = teacher_hidden_flat - teacher_weight_for_kl = teacher_weight_local - - token_kl = vocab_parallel_reverse_kl_gathered( - local_student_hidden=student_hidden_for_kl, - student_weight_local=student_weight_for_kl, - local_teacher_hidden=teacher_hidden_for_kl.detach(), - teacher_weight_local=teacher_weight_for_kl.detach(), - local_labels=labels.to(device=device), - ignore_index=ignore_index, - group=group, - ) - full_token_weights = _gather_1d_nograd(local_token_weights.to(token_kl.device), group=group) - if full_token_weights.shape[0] != token_kl.shape[0]: - raise ValueError( - f"gathered teacher_weights rows {full_token_weights.shape[0]} do not match KL rows {token_kl.shape[0]}" - ) - - clamp_frac = 0.0 - if loss_max_clamp is not None and token_kl.numel(): - clamp_frac = (token_kl.detach().abs() >= float(loss_max_clamp)).float().mean().item() - token_kl = token_kl.clamp(min=-loss_max_clamp, max=loss_max_clamp) - - if loss_reducer is None: - scale = torch.tensor(max(int(token_kl.numel()), 1), dtype=torch.float32, device=token_kl.device) - loss_reducer = TokenPartial(scale=scale) - if metric_reducer is None: - scale = torch.tensor(max(local_tokens, 1), dtype=torch.float32, device=device) - metric_reducer = TokenPartial(scale=scale) - - weighted_token_kl = token_kl * full_token_weights.to(token_kl.device) - valid_full = torch.ones_like(weighted_token_kl, dtype=torch.float32) - raw_kl_loss = loss_reducer(weighted_token_kl, valid_full) - if float(kl_loss_weight) == 0.0: - loss = raw_kl_loss.detach() * 0.0 - else: - scaled_kl_loss = float(kl_loss_weight) * raw_kl_loss - world = dist.get_world_size(group) if dist.is_available() and dist.is_initialized() else 1 - loss = _value_divided_without_grad_scale(scaled_kl_loss, world) - - hidden_match_metric = 0.0 - hidden_match_raw_metric = 0.0 - hidden_match_weight_mean = local_hidden_weights.mean().item() if local_tokens else 0.0 - hidden_match_pos_metric = 0.0 - hidden_match_neg_metric = 0.0 - hidden_match_pos_raw_metric = 0.0 - hidden_match_neg_raw_metric = 0.0 - hidden_match_neg_minus_pos_raw = 0.0 - hidden_match_pos_weight_mean = 0.0 - hidden_match_neg_weight_mean = 0.0 - oprd_metric = 0.0 - oprd_raw_metric = 0.0 - oprd_num_layers = 0 - - hidden_match_coef = float(hidden_match_coef or 0.0) - local_valid = torch.ones(local_tokens, dtype=torch.float32, device=device) - use_oprd = ( - hidden_match_coef - and student_layer_hidden_states is not None - and (teacher_layer_hidden_states is not None or teacher_layer_fetcher is not None) - ) - if use_oprd and local_tokens: - if teacher_layer_fetcher is not None: - hidden_distance, oprd_num_layers = _oprd_hidden_distance_from_fetcher( - student_layer_hidden_states, - teacher_layer_fetcher=teacher_layer_fetcher, - expected_rows=local_tokens, - num_layers=teacher_layer_num_layers - if teacher_layer_num_layers is not None - else int(student_layer_hidden_states.shape[1]), - layer_chunk_size=oprd_layer_chunk_size, - ) - else: - hidden_distance, oprd_num_layers = _oprd_hidden_distance( - student_layer_hidden_states, - teacher_layer_hidden_states, - local_tokens, - layer_chunk_size=oprd_layer_chunk_size, - ) - weighted_hidden_match = hidden_distance * local_hidden_weights.to(hidden_distance.device) - hidden_match_loss = loss_reducer(weighted_hidden_match, local_valid.to(hidden_distance.device)) - loss = loss + hidden_match_coef * hidden_match_loss - oprd_metric = metric_reducer(weighted_hidden_match.detach(), local_valid.to(hidden_distance.device)).item() - oprd_raw_metric = metric_reducer(hidden_distance.detach(), local_valid.to(hidden_distance.device)).item() - hidden_match_metric = oprd_metric - hidden_match_raw_metric = oprd_raw_metric - elif hidden_match_coef and local_tokens: - if student_hidden_flat.shape[-1] != teacher_hidden_flat.shape[-1]: - raise ValueError( - "hidden_match requires matching hidden sizes, got " - f"student={student_hidden_flat.shape[-1]} teacher={teacher_hidden_flat.shape[-1]}" - ) - if str(hidden_match_mode).lower() == "mse": - hidden_distance = ((student_hidden_flat.float() - teacher_hidden_flat.float()) ** 2).mean(dim=-1) - else: - hidden_distance = 1.0 - F.cosine_similarity( - student_hidden_flat.float(), - teacher_hidden_flat.float(), - dim=-1, - eps=1e-6, - ) - hidden_weights_on_device = local_hidden_weights.to(hidden_distance.device) - weighted_hidden_match = hidden_distance * hidden_weights_on_device - hidden_match_loss = loss_reducer(weighted_hidden_match, local_valid.to(hidden_distance.device)) - loss = loss + hidden_match_coef * hidden_match_loss - hidden_match_metric = metric_reducer( - weighted_hidden_match.detach(), local_valid.to(hidden_distance.device) - ).item() - hidden_match_raw_metric = metric_reducer( - hidden_distance.detach(), local_valid.to(hidden_distance.device) - ).item() - pos_weights = torch.clamp(hidden_weights_on_device, min=0.0) - neg_weights = torch.clamp(-hidden_weights_on_device, min=0.0) - hidden_match_pos_metric = metric_reducer((hidden_distance * pos_weights).detach(), local_valid).item() - hidden_match_neg_metric = metric_reducer((hidden_distance * neg_weights).detach(), local_valid).item() - pos_weight_sum = pos_weights.sum() - neg_weight_sum = neg_weights.sum() - if pos_weight_sum.item() > 0: - hidden_match_pos_raw_metric = ((hidden_distance.detach() * pos_weights).sum() / pos_weight_sum).item() - if neg_weight_sum.item() > 0: - hidden_match_neg_raw_metric = ((hidden_distance.detach() * neg_weights).sum() / neg_weight_sum).item() - hidden_match_neg_minus_pos_raw = hidden_match_neg_raw_metric - hidden_match_pos_raw_metric - hidden_match_pos_weight_mean = pos_weights.mean().item() - hidden_match_neg_weight_mean = neg_weights.mean().item() - - if use_task_rewards: - loss = loss * float(distillation_loss_coef) - - detached_token_kl = token_kl.detach() - full_valid_count = max(float(detached_token_kl.numel()), 1.0) - metrics_kwargs: dict = { - "valid_tokens": local_tokens, - "opd_kl": detached_token_kl.sum().item() / full_valid_count, - "opd_weighted_kl": weighted_token_kl.detach().sum().item() / full_valid_count, - "opd_vocab_parallel_group_tokens": int(detached_token_kl.numel()), - "opd_vocab_parallel_kl_sum": detached_token_kl.sum().item(), - "opd_vocab_parallel_weighted_kl_sum": weighted_token_kl.detach().sum().item(), - "opd_hidden_match_loss": hidden_match_metric, - "opd_hidden_match_raw_loss": hidden_match_raw_metric, - "opd_hidden_match_weight_mean": hidden_match_weight_mean, - "opd_hidden_match_pos_loss": hidden_match_pos_metric, - "opd_hidden_match_neg_loss": hidden_match_neg_metric, - "opd_hidden_match_pos_raw_loss": hidden_match_pos_raw_metric, - "opd_hidden_match_neg_raw_loss": hidden_match_neg_raw_metric, - "opd_hidden_match_neg_minus_pos_raw": hidden_match_neg_minus_pos_raw, - "opd_hidden_match_pos_weight_mean": hidden_match_pos_weight_mean, - "opd_hidden_match_neg_weight_mean": hidden_match_neg_weight_mean, - "opd_teacher_weight_mean": full_token_weights.mean().item() if full_token_weights.numel() else 0.0, - "opd_loss_clamp_frac": clamp_frac, - "opd_oprd_loss": oprd_metric, - "opd_oprd_raw_loss": oprd_raw_metric, - "opd_oprd_num_layers": oprd_num_layers, - } - if detached_token_kl.numel(): - metrics_kwargs["opd_loss_min"] = detached_token_kl.min().item() - metrics_kwargs["opd_loss_max"] = detached_token_kl.max().item() - metrics_kwargs["opd_loss_abs_mean"] = detached_token_kl.abs().mean().item() - - metrics = OPDLossMetrics(**metrics_kwargs).to_dict() - if debug_token_outputs: - local_start, local_end = _gathered_local_bounds(local_tokens, device=token_kl.device, group=group) - metrics["_opd_debug_local_token_kl"] = detached_token_kl[local_start:local_end].detach() - metrics["_opd_debug_local_weighted_token_kl"] = weighted_token_kl.detach()[local_start:local_end] - metrics["_opd_debug_local_token_weight"] = full_token_weights.detach()[local_start:local_end] - - return LossOutput(loss=loss, metrics=metrics) - - -def opd_loss_function( - hidden_states: torch.Tensor, - weight: torch.Tensor, - labels: torch.Tensor, - teacher_hidden_states: torch.Tensor, - teacher_lm_head_weight: torch.Tensor, - teacher_weights: Optional[torch.Tensor] = None, - hidden_match_weights: Optional[torch.Tensor] = None, - ignore_index: int = -100, - num_chunks: int = 8, - lm_head_fp32: bool = False, - teacher_lm_head_fp32: bool = True, - kl_backend: str = "torch_compile", - vocab_chunk_size: int | None = 32768, - streaming_lowmem: bool = False, - return_per_token: bool = False, - normalization_denominator: Optional[torch.Tensor | int | float] = None, - loss_reducer: Optional[Reducer] = None, - metric_reducer: Optional[Reducer] = None, - loss_mode: str = LOSS_MODE_REVERSE_KL_FULL, - log_prob_min_clamp: Optional[float] = None, - loss_max_clamp: Optional[float] = None, - emit_full_vocab_diagnostics: bool = False, - use_policy_gradient: bool = False, - old_logprobs: Optional[torch.Tensor] = None, - clip_ratio_low: float = 0.2, - clip_ratio_high: float = 0.2, - use_task_rewards: bool = False, - distillation_loss_coef: float = 1.0, - hidden_match_coef: float = 0.0, - kl_loss_weight: float = 1.0, - hidden_match_mode: str = "cosine", - teacher_layer_hidden_states: Optional[torch.Tensor] = None, - teacher_layer_fetcher: Optional[Callable[[int, int], torch.Tensor]] = None, - teacher_layer_num_layers: Optional[int] = None, - oprd_layer_chunk_size: int = 4, - student_layer_hidden_states: Optional[torch.Tensor] = None, - diag_region_ids: Optional[torch.Tensor] = None, - diag_sample_ok: Optional[torch.Tensor] = None, -) -> LossOutput: - """Compute the OPD distillation loss for one micro-batch. - - Supported `loss_mode`s: - - "reverse_kl_full" (default): full-vocabulary KL(student||teacher), - i.e. Σ_v p_S(v)·(log p_S(v) - log p_T(v)). - - "forward_kl_full": full-vocabulary KL(teacher||student), - i.e. Σ_v p_T(v)·(log p_T(v) - log p_S(v)). The full-distribution - counterpart of VERL's truncated `forward_kl_topk` — possible here - because xorl materializes the full teacher distribution locally via the - cached teacher LM head. - - Single-sample estimators "kl"/"k1"/"abs"/"mse"/"k2"/"low_var_kl"/"k3" - (with optional "+" suffix for k2 straight-through gradient). These use - only the per-token (student_logprob, teacher_logprob) at the sampled - token; a cheaper memory point for RL-style updates. - - Set `emit_full_vocab_diagnostics=True` to also emit per-step teacher_entropy, - student_entropy, and top1_agreement. Adds a single extra full-vocab pass over - the logits, so opt-in. Only applies to the two full-vocab loss modes. - - Expected shapes: - hidden_states: [batch, seq, student_hidden_dim] - weight: [vocab_size, student_hidden_dim] - labels: [batch, seq], with ignore_index masking tokens out of the loss - teacher_hidden_states: [batch, seq, teacher_hidden_dim] - teacher_lm_head_weight: [vocab_size, teacher_hidden_dim] - teacher_weights: optional [batch, seq] per-token multipliers applied - to the KL term after KL computation and before final normalization. - hidden_match_weights: optional [batch, seq] per-token multipliers for - hidden-state matching. Defaults to teacher_weights when omitted. - teacher_layer_hidden_states / student_layer_hidden_states: optional - multi-layer OPRD tensors, each shaped [valid_tokens, L, d] and already - restricted to the valid (response-masked) positions IN THE SAME ORDER - as the internal valid_mask flatten (i.e. row i corresponds to the i-th - True entry of labels.reshape(-1) != ignore_index). When BOTH are given - the OPRD term — hidden_match_coef * mean_L( (1/d)||student_l - - teacher_l.detach()||^2 ), weighted by hidden_match_weights and reduced - by the existing reducer — REPLACES the single-layer hidden-match term to - avoid double-counting. When None, single-layer behavior is unchanged. - - `diag_region_ids` / `diag_sample_ok` are optional [batch, seq] int tensors - (labels-aligned, like teacher_weights) that drive metrics-only KL splits: - region 0=prompt / 1=buffer / 2=answer (-1 = unattributed) and per-sample - sampled-answer correctness 1/0 (-1 = unknown), broadcast over the sample's - positions by the client. They never touch the loss value. - - `log_prob_min_clamp` is a generic stability clamp on student log-probs. - Currently used by `forward_kl_full` (where p_T(v) · log p_S(v) terms can - blow up if the student puts ~0 mass on a teacher-preferred token); reverse - KL doesn't need it (its weights are p_S(v)). - `loss_max_clamp` symmetrically clamps the *unweighted* per-token loss before - teacher-weighting. - - Teacher tensors are detached by construction. Only the student hidden states - and student LM head receive gradients. - """ - if hidden_states.shape[:-1] != labels.shape: - raise ValueError(f"hidden_states shape {hidden_states.shape} is incompatible with labels {labels.shape}") - if teacher_hidden_states.shape[:-1] != labels.shape: - raise ValueError( - f"teacher_hidden_states shape {teacher_hidden_states.shape} is incompatible with labels {labels.shape}" - ) - if weight.shape[0] != teacher_lm_head_weight.shape[0]: - raise ValueError( - f"student vocab size ({weight.shape[0]}) must match teacher vocab size ({teacher_lm_head_weight.shape[0]})" - ) - if hidden_states.shape[-1] != weight.shape[-1]: - raise ValueError( - f"student hidden size ({hidden_states.shape[-1]}) must match student head width ({weight.shape[-1]})" - ) - if teacher_hidden_states.shape[-1] != teacher_lm_head_weight.shape[-1]: - raise ValueError( - "teacher hidden size " - f"({teacher_hidden_states.shape[-1]}) must match teacher head width ({teacher_lm_head_weight.shape[-1]})" - ) - - original_shape = labels.shape - labels_flat = labels.reshape(-1) - valid_mask = labels_flat != ignore_index - valid_count = valid_mask.sum() - - if valid_count.item() == 0: - loss = _zero_loss_with_graph(hidden_states, weight) - per_token_loss = ( - torch.zeros(original_shape, dtype=torch.float32, device=labels.device) if return_per_token else None - ) - return LossOutput(loss=loss, per_token_loss=per_token_loss, metrics=OPDLossMetrics(valid_tokens=0).to_dict()) - - student_hidden_flat = hidden_states.reshape(-1, hidden_states.size(-1))[valid_mask] - teacher_hidden_flat = teacher_hidden_states.reshape(-1, teacher_hidden_states.size(-1))[valid_mask].detach() - labels_valid = labels_flat[valid_mask] - token_weights = _as_flat_optional_weights(teacher_weights, valid_mask, torch.float32) - hidden_weights = ( - _as_flat_optional_weights(hidden_match_weights, valid_mask, torch.float32) - if hidden_match_weights is not None - else token_weights - ) - - default_scale = _denominator_tensor( - normalization_denominator, - fallback=valid_count, - device=hidden_states.device, - ) - if loss_reducer is None: - loss_reducer = TokenPartial(scale=default_scale) - if metric_reducer is None: - metric_reducer = TokenPartial(scale=default_scale) - - region_flat = None - if diag_region_ids is not None: - region_flat = diag_region_ids.reshape(-1).to(device=labels.device)[valid_mask] - sample_ok_flat = None - if diag_sample_ok is not None: - sample_ok_flat = diag_sample_ok.reshape(-1).to(device=labels.device)[valid_mask] - - backend = kl_backend.lower() - is_compile_backend = backend in {"torch_compile", "compile", "auto_chunker"} - is_streaming_backend = backend in {"streaming", "tilelang"} - - # Full-vocab diagnostics: the compile backend has fused *_with_diag kernels; - # the streaming backend covers the full-vocab KL modes via a separate no-grad - # streaming pass (same chunked-vocab shape, never materializes full logits). - streaming_diag_supported = is_streaming_backend and loss_mode in ( - LOSS_MODE_REVERSE_KL_FULL, - LOSS_MODE_FORWARD_KL_FULL, - ) - diag_enabled = bool(emit_full_vocab_diagnostics) and (is_compile_backend or streaming_diag_supported) - if loss_mode == LOSS_MODE_FORWARD_KL_FULL and not (is_compile_backend or is_streaming_backend): - raise ValueError( - "loss_mode='forward_kl_full' requires kl_backend in {'torch_compile', 'streaming', 'tilelang'}" - ) - - teacher_entropy_per_tok = None - student_entropy_per_tok = None - top1_agreement_per_tok = None - student_lp_at_label: Optional[torch.Tensor] = None - pg_clipfrac = 0.0 - pg_clipfrac_lower = 0.0 - ppo_kl = 0.0 - - if loss_mode == LOSS_MODE_REVERSE_KL_FULL: - if is_compile_backend: - if not torch.is_tensor(teacher_lm_head_weight): - raise ValueError("torch_compile OPD KL backend requires a materialized teacher LM head tensor") - if diag_enabled: - ( - token_kl, - teacher_entropy_per_tok, - student_entropy_per_tok, - top1_agreement_per_tok, - ) = compiled_reverse_kl_with_diag_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - num_chunks=num_chunks, - lm_head_fp32=lm_head_fp32, - teacher_lm_head_fp32=teacher_lm_head_fp32, - ) - else: - token_kl = compiled_reverse_kl_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - num_chunks=num_chunks, - lm_head_fp32=lm_head_fp32, - teacher_lm_head_fp32=teacher_lm_head_fp32, - ) - elif is_streaming_backend and streaming_lowmem: - # Memory-lean streaming KL for the fp32 lm-head regime: keep the - # student/teacher lm-head weights in their native (bf16) dtype and - # upcast each vocab chunk to fp32 inside the kernel, instead of - # holding two full fp32 weight copies + a full fp32 grad buffer. Only - # the cheap [N,H] hidden states are upcast here. Gradient-identical to - # the streaming path below (slicing commutes with the elementwise - # upcast; vocab chunks partition grad rows disjointly), but saves - # several GB of resident lm-head memory -- the AMDAHL-029..033 1-node - # blocker. The fp32 chunked matmul matches lm_head_fp32=True / - # teacher_lm_head_fp32=True (the OPD default); the diagnostics pass - # below is metrics-only and stays on the native-dtype weight. - if lm_head_fp32: - student_hidden_flat = student_hidden_flat.float() - if teacher_lm_head_fp32: - teacher_hidden_flat = teacher_hidden_flat.float() - compute_dtype = torch.float32 if (lm_head_fp32 or teacher_lm_head_fp32) else weight.dtype - token_kl = streaming_reverse_kl_lowmem_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - vocab_chunk_size=vocab_chunk_size, - compute_dtype=compute_dtype, - inplace_weight_grad=False, - ) - if diag_enabled: - ( - teacher_entropy_per_tok, - student_entropy_per_tok, - top1_agreement_per_tok, - ) = streaming_full_vocab_diagnostics( - student_hidden_states=student_hidden_flat.detach(), - student_weight=weight.detach(), - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - vocab_chunk_size=vocab_chunk_size, - ) - elif is_streaming_backend: - if lm_head_fp32: - student_hidden_flat = student_hidden_flat.float() - weight = weight.float() - if teacher_lm_head_fp32: - teacher_hidden_flat = teacher_hidden_flat.float() - if torch.is_tensor(teacher_lm_head_weight): - teacher_lm_head_weight = teacher_lm_head_weight.float() - token_kl = streaming_reverse_kl_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - vocab_chunk_size=vocab_chunk_size, - ) - if diag_enabled: - ( - teacher_entropy_per_tok, - student_entropy_per_tok, - top1_agreement_per_tok, - ) = streaming_full_vocab_diagnostics( - student_hidden_states=student_hidden_flat.detach(), - student_weight=weight.detach(), - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - vocab_chunk_size=vocab_chunk_size, - ) - else: - raise ValueError( - f"Unsupported OPD KL backend '{kl_backend}'. Expected 'torch_compile', 'streaming', or 'tilelang'." - ) - elif loss_mode == LOSS_MODE_FORWARD_KL_FULL: - if is_compile_backend: - if not torch.is_tensor(teacher_lm_head_weight): - raise ValueError("forward_kl_full requires a materialized teacher LM head tensor") - if diag_enabled: - ( - token_kl, - teacher_entropy_per_tok, - student_entropy_per_tok, - top1_agreement_per_tok, - ) = compiled_forward_kl_full_with_diag_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - log_prob_min_clamp=log_prob_min_clamp, - num_chunks=num_chunks, - lm_head_fp32=lm_head_fp32, - teacher_lm_head_fp32=teacher_lm_head_fp32, - ) - else: - token_kl = compiled_forward_kl_full_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - log_prob_min_clamp=log_prob_min_clamp, - num_chunks=num_chunks, - lm_head_fp32=lm_head_fp32, - teacher_lm_head_fp32=teacher_lm_head_fp32, - ) - elif is_streaming_backend: - # Streaming forward-KL: never materializes the full-vocab fp32 logits - # (~9.27 GiB at 151936 vocab), unblocking the large-microbatch OOM that - # forces forward_kl_full onto the compile backend. The streaming kernel - # accumulates the teacher-weighted log-ratio online over vocab chunks; - # it has no student-log-prob materialization, so `log_prob_min_clamp` - # (a clamp on log p_S) is unsupported here -- fail loud rather than - # silently ignore it. - if log_prob_min_clamp is not None: - raise ValueError( - "log_prob_min_clamp is not supported by the streaming forward_kl_full backend " - "(it never materializes student log-probs); use kl_backend='torch_compile' to clamp." - ) - if streaming_lowmem: - if lm_head_fp32: - student_hidden_flat = student_hidden_flat.float() - if teacher_lm_head_fp32: - teacher_hidden_flat = teacher_hidden_flat.float() - compute_dtype = torch.float32 if (lm_head_fp32 or teacher_lm_head_fp32) else weight.dtype - token_kl = streaming_forward_kl_lowmem_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - vocab_chunk_size=vocab_chunk_size, - compute_dtype=compute_dtype, - inplace_weight_grad=False, - ) - else: - if lm_head_fp32: - student_hidden_flat = student_hidden_flat.float() - weight = weight.float() - if teacher_lm_head_fp32: - teacher_hidden_flat = teacher_hidden_flat.float() - if torch.is_tensor(teacher_lm_head_weight): - teacher_lm_head_weight = teacher_lm_head_weight.float() - token_kl = streaming_forward_kl_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - vocab_chunk_size=vocab_chunk_size, - ) - if diag_enabled: - ( - teacher_entropy_per_tok, - student_entropy_per_tok, - top1_agreement_per_tok, - ) = streaming_full_vocab_diagnostics( - student_hidden_states=student_hidden_flat.detach(), - student_weight=weight.detach(), - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - vocab_chunk_size=vocab_chunk_size, - ) - else: - raise ValueError( - f"Unsupported OPD KL backend '{kl_backend}'. Expected 'torch_compile', 'streaming', or 'tilelang'." - ) - elif is_estimator_loss_mode(loss_mode): - if not torch.is_tensor(teacher_lm_head_weight): - raise ValueError(f"loss_mode='{loss_mode}' requires a materialized teacher LM head tensor") - student_lp, teacher_lp = compiled_sampled_token_logprobs_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - num_chunks=num_chunks, - lm_head_fp32=lm_head_fp32, - teacher_lm_head_fp32=teacher_lm_head_fp32, - ) - token_kl = _kl_penalty_estimator(student_lp, teacher_lp, loss_mode) - # `student_lp` is reused by PG mode below, no need for a second forward. - student_lp_at_label = student_lp - else: - raise ValueError( - f"Unsupported OPD loss_mode '{loss_mode}'. Expected 'reverse_kl_full', " - f"'forward_kl_full', or one of {_ESTIMATOR_MODES} (with optional '+' suffix)." - ) - - # Symmetric loss clamp before teacher-weighting. Matches - # verl/trainer/distillation/losses.py:228 (`distillation_losses.clamp(...)`). - clamp_frac = 0.0 - if loss_max_clamp is not None: - clamp_frac = (token_kl.detach().abs() >= float(loss_max_clamp)).float().mean().item() - token_kl = token_kl.clamp(min=-loss_max_clamp, max=loss_max_clamp) - # Post-clamp, pre-PG-surrogate per-token KL for the metrics-only region / - # correctness splits below (PG mode replaces token_kl with the surrogate). - kl_for_diag = token_kl.detach() - - # Policy-gradient OPD: treat -token_kl as advantage, apply PPO-clip on the - # student's sampled-token log-prob ratio. Mirrors VERL's `use_policy_gradient` - # path in verl/trainer/distillation/losses.py:259-281. - if use_policy_gradient: - if old_logprobs is None: - raise ValueError("use_policy_gradient=True requires old_logprobs to be provided") - if student_lp_at_label is None: - # Need student logprob at sampled-token position. Reuse the same - # compiled-sampled-logp helper used by estimator modes. - student_lp_at_label, _ = compiled_sampled_token_logprobs_function( - student_hidden_states=student_hidden_flat, - student_weight=weight, - teacher_hidden_states=teacher_hidden_flat, - teacher_weight=teacher_lm_head_weight, - labels=labels_valid, - ignore_index=ignore_index, - num_chunks=num_chunks, - lm_head_fp32=lm_head_fp32, - teacher_lm_head_fp32=teacher_lm_head_fp32, - ) - old_logprobs_flat = old_logprobs.reshape(-1).to(device=labels.device) - old_logprobs_valid = old_logprobs_flat[valid_mask].to(student_lp_at_label.dtype) - advantages = (-token_kl).detach() - log_ratio = student_lp_at_label - old_logprobs_valid - ratio = log_ratio.exp() - # min over the two surrogate losses, per PPO. NOTE: token_kl is replaced - # by the PG surrogate; further teacher-weighting / coef scaling below - # still applies element-wise. - pg_losses1 = -ratio * advantages - pg_losses2 = -ratio.clamp(1.0 - clip_ratio_low, 1.0 + clip_ratio_high) * advantages - pg_losses = torch.maximum(pg_losses1, pg_losses2) - is_clipped = (pg_losses2 > pg_losses1).to(token_kl.dtype) - pg_clipfrac = is_clipped.mean().item() if is_clipped.numel() > 0 else 0.0 - # Parity with VERL's `actor/ppo_kl` and `actor/pg_clipfrac_lower` - # (`core_algos.py:1365-1367`, surfaced under `distillation/` prefix in - # `distillation_loss`). `ppo_kl` = mean(old_logp - new_logp); `clipfrac_lower` - # = fraction of tokens where the LOWER clamp bound binds (1-clip_ratio_low) - # AND the advantage is positive (the unclipped surrogate would be smaller - # than the clipped one, so the clip is "hurting" gradient flow into a - # promising direction). - ppo_kl = (-log_ratio).mean().item() if log_ratio.numel() > 0 else 0.0 - is_clipped_lower = ((ratio < (1.0 - clip_ratio_low)) & (advantages < 0)).to(token_kl.dtype) - pg_clipfrac_lower = is_clipped_lower.mean().item() if is_clipped_lower.numel() > 0 else 0.0 - token_kl = pg_losses # downstream weighting / reduction unchanged - - token_weights_on_device = token_weights.to(token_kl.device) - weighted_token_kl = token_kl * token_weights_on_device - valid_ones = torch.ones_like(weighted_token_kl, dtype=torch.float32) - valid_count_float = max(float(valid_count.item()), 1.0) - - # kl_loss_weight scales the logit-KL term. Set 0.0 to supervise ONLY on hidden - # states (loss = hidden_match_coef * hidden_match): in self-distillation the - # student/teacher share the LM head, so matching answer-position hiddens matches - # the logits via that head while forcing the buffer to set up the answer hidden. - if float(kl_loss_weight) == 0.0 and not use_policy_gradient: - # Hidden-only supervision: the KL term is scaled by 0, so it contributes no - # gradient -- yet `0.0 * reduce(kl)` still backprops zeros through the full - # vocab-parallel KL and its LM-head matmul (the single most expensive term; - # the only gradient path to the shared LM head here). Detach it so the - # backward skips that subgraph entirely. token_kl/weighted_token_kl remain - # computed above, so opd_kl and the KL diagnostics are unaffected -- only the - # (zero) gradient is elided. - loss = loss_reducer(weighted_token_kl, valid_ones).detach() * 0.0 - else: - loss = float(kl_loss_weight) * loss_reducer(weighted_token_kl, valid_ones) - hidden_match_metric = 0.0 - hidden_match_raw_metric = 0.0 - hidden_match_weight_mean = 0.0 - hidden_match_pos_metric = 0.0 - hidden_match_neg_metric = 0.0 - hidden_match_pos_raw_metric = 0.0 - hidden_match_neg_raw_metric = 0.0 - hidden_match_neg_minus_pos_raw = 0.0 - hidden_match_pos_weight_mean = 0.0 - hidden_match_neg_weight_mean = 0.0 - weighted_hidden_match = None - oprd_metric = 0.0 - oprd_raw_metric = 0.0 - oprd_num_layers = 0 - hidden_match_coef = float(hidden_match_coef or 0.0) - # Multi-layer OPRD: when BOTH per-layer tensors are present, this REPLACES the - # single-layer hidden term below (the `else` branch) so the same coefficient - # isn't applied twice. Layer tensors arrive [valid_tokens, L, d] already - # restricted to the valid positions in valid_mask order. - use_oprd = ( - hidden_match_coef - and student_layer_hidden_states is not None - and (teacher_layer_hidden_states is not None or teacher_layer_fetcher is not None) - ) - if use_oprd: - if teacher_layer_fetcher is not None: - hidden_distance, oprd_num_layers = _oprd_hidden_distance_from_fetcher( - student_layer_hidden_states, - teacher_layer_fetcher=teacher_layer_fetcher, - expected_rows=int(student_hidden_flat.shape[0]), - num_layers=teacher_layer_num_layers - if teacher_layer_num_layers is not None - else int(student_layer_hidden_states.shape[1]), - layer_chunk_size=oprd_layer_chunk_size, - ) - else: - hidden_distance, oprd_num_layers = _oprd_hidden_distance( - student_layer_hidden_states, - teacher_layer_hidden_states, - int(student_hidden_flat.shape[0]), - layer_chunk_size=oprd_layer_chunk_size, - ) - hidden_weights_on_device = hidden_weights.to(hidden_distance.device) - weighted_hidden_match = hidden_distance * hidden_weights_on_device - hidden_match_loss = loss_reducer(weighted_hidden_match, valid_ones) - loss = loss + hidden_match_coef * hidden_match_loss - oprd_metric = metric_reducer(weighted_hidden_match.detach(), valid_ones).item() - oprd_raw_metric = metric_reducer(hidden_distance.detach(), valid_ones).item() - hidden_match_metric = oprd_metric - hidden_match_raw_metric = oprd_raw_metric - hidden_match_weight_mean = hidden_weights.mean().item() - elif hidden_match_coef: - if student_hidden_flat.shape[-1] != teacher_hidden_flat.shape[-1]: - raise ValueError( - "hidden_match requires matching hidden sizes, got " - f"student={student_hidden_flat.shape[-1]} teacher={teacher_hidden_flat.shape[-1]}" - ) - if str(hidden_match_mode).lower() == "mse": - # Magnitude-aware: MSE->0 implies student hidden == teacher hidden, so in - # self-distillation (shared LM head) the logits match too. Cosine alone - # matches only direction (ignores magnitude) and decouples from generation - # (eval accuracy collapses while cosine-distance keeps falling). - hidden_distance = ((student_hidden_flat.float() - teacher_hidden_flat.float()) ** 2).mean(dim=-1) - else: - hidden_distance = 1.0 - F.cosine_similarity( - student_hidden_flat.float(), - teacher_hidden_flat.float(), - dim=-1, - eps=1e-6, - ) - hidden_weights_on_device = hidden_weights.to(hidden_distance.device) - weighted_hidden_match = hidden_distance * hidden_weights_on_device - hidden_match_loss = loss_reducer(weighted_hidden_match, valid_ones) - loss = loss + hidden_match_coef * hidden_match_loss - hidden_match_metric = metric_reducer(weighted_hidden_match.detach(), valid_ones).item() - hidden_match_raw_metric = metric_reducer(hidden_distance.detach(), valid_ones).item() - hidden_match_weight_mean = hidden_weights.mean().item() - pos_weights = torch.clamp(hidden_weights_on_device, min=0.0) - neg_weights = torch.clamp(-hidden_weights_on_device, min=0.0) - hidden_match_pos_metric = metric_reducer((hidden_distance * pos_weights).detach(), valid_ones).item() - hidden_match_neg_metric = metric_reducer((hidden_distance * neg_weights).detach(), valid_ones).item() - pos_weight_sum = pos_weights.sum() - neg_weight_sum = neg_weights.sum() - if pos_weight_sum.item() > 0: - hidden_match_pos_raw_metric = ((hidden_distance.detach() * pos_weights).sum() / pos_weight_sum).item() - if neg_weight_sum.item() > 0: - hidden_match_neg_raw_metric = ((hidden_distance.detach() * neg_weights).sum() / neg_weight_sum).item() - hidden_match_neg_minus_pos_raw = hidden_match_neg_raw_metric - hidden_match_pos_raw_metric - hidden_match_pos_weight_mean = pos_weights.mean().item() - hidden_match_neg_weight_mean = neg_weights.mean().item() - - # Task-reward mixing: when use_task_rewards=True the caller is mixing this - # distillation loss with a separate task-RL loss outside this function and - # provides the coefficient here. When False, coef is ignored (== VERL semantics: - # `distillation_loss_coef if use_task_rewards else 1.0`). - if use_task_rewards: - loss = loss * float(distillation_loss_coef) - - per_token_loss = None - if return_per_token: - per_token_flat = torch.zeros(labels_flat.shape, dtype=torch.float32, device=labels.device) - per_token_contrib = float(kl_loss_weight) * weighted_token_kl - if weighted_hidden_match is not None: - per_token_contrib = per_token_contrib + hidden_match_coef * weighted_hidden_match - per_token_flat[valid_mask] = per_token_contrib.detach().to(per_token_flat.device) - per_token_loss = per_token_flat.view(original_shape) - - detached_token_kl = token_kl.detach() - metrics_kwargs: dict = { - "valid_tokens": int(valid_count.item()), - "opd_kl": detached_token_kl.sum().item() / valid_count_float, - "opd_weighted_kl": metric_reducer(weighted_token_kl.detach(), valid_ones).item(), - "opd_hidden_match_loss": hidden_match_metric, - "opd_hidden_match_raw_loss": hidden_match_raw_metric, - "opd_hidden_match_weight_mean": hidden_match_weight_mean, - "opd_hidden_match_pos_loss": hidden_match_pos_metric, - "opd_hidden_match_neg_loss": hidden_match_neg_metric, - "opd_hidden_match_pos_raw_loss": hidden_match_pos_raw_metric, - "opd_hidden_match_neg_raw_loss": hidden_match_neg_raw_metric, - "opd_hidden_match_neg_minus_pos_raw": hidden_match_neg_minus_pos_raw, - "opd_hidden_match_pos_weight_mean": hidden_match_pos_weight_mean, - "opd_hidden_match_neg_weight_mean": hidden_match_neg_weight_mean, - "opd_teacher_weight_mean": token_weights.mean().item(), - } - # Distillation-loss range metrics (parity with VERL compute_distillation_loss_range). - metrics_kwargs["opd_loss_min"] = detached_token_kl.min().item() - metrics_kwargs["opd_loss_max"] = detached_token_kl.max().item() - metrics_kwargs["opd_loss_abs_mean"] = detached_token_kl.abs().mean().item() - metrics_kwargs["opd_loss_clamp_frac"] = clamp_frac - - # Region / sample-correctness KL splits (metrics only; see OPDLossMetrics for - # the per-valid normalization + exact-aggregation contract). - def _masked_split(mask: torch.Tensor) -> tuple[float, float]: - m = mask.to(kl_for_diag.dtype) - return ( - (kl_for_diag * m).sum().item() / valid_count_float, - m.sum().item() / valid_count_float, - ) - - if region_flat is not None: - for region_value, region_name in ((0, "prompt"), (1, "buffer"), (2, "answer")): - kl_per_valid, frac = _masked_split(region_flat == region_value) - metrics_kwargs[f"opd_kl_{region_name}_per_valid"] = kl_per_valid - metrics_kwargs[f"opd_frac_{region_name}"] = frac - if sample_ok_flat is not None: - answer_mask = region_flat == 2 - for ok_value, ok_name in ((1, "correct"), (0, "wrong")): - split_mask = answer_mask & (sample_ok_flat == ok_value) - kl_per_valid, frac = _masked_split(split_mask) - metrics_kwargs[f"opd_kl_answer_{ok_name}_per_valid"] = kl_per_valid - metrics_kwargs[f"opd_frac_answer_{ok_name}"] = frac - if student_entropy_per_tok is not None and teacher_entropy_per_tok is not None: - m = split_mask.to(kl_for_diag.dtype) - metrics_kwargs[f"opd_student_entropy_answer_{ok_name}_per_valid"] = ( - student_entropy_per_tok.detach() * m - ).sum().item() / valid_count_float - metrics_kwargs[f"opd_teacher_entropy_answer_{ok_name}_per_valid"] = ( - teacher_entropy_per_tok.detach() * m - ).sum().item() / valid_count_float - metrics_kwargs["opd_oprd_loss"] = oprd_metric - metrics_kwargs["opd_oprd_raw_loss"] = oprd_raw_metric - metrics_kwargs["opd_oprd_num_layers"] = oprd_num_layers - if is_estimator_loss_mode(loss_mode): - # k1 can be negative; mirror VERL's distillation/abs_loss metric. - metrics_kwargs["opd_abs_loss"] = detached_token_kl.abs().mean().item() - if use_policy_gradient: - metrics_kwargs["opd_pg_clipfrac"] = float(pg_clipfrac) - metrics_kwargs["opd_pg_clipfrac_lower"] = float(pg_clipfrac_lower) - metrics_kwargs["opd_ppo_kl"] = float(ppo_kl) - # All flattened tensors here are over already-filtered valid tokens, so - # every position counts as "response-masked valid" and mean = sum / valid_count. - if teacher_entropy_per_tok is not None and student_entropy_per_tok is not None: - metrics_kwargs["opd_teacher_entropy"] = teacher_entropy_per_tok.detach().mean().item() - metrics_kwargs["opd_student_entropy"] = student_entropy_per_tok.detach().mean().item() - if top1_agreement_per_tok is not None: - metrics_kwargs["opd_top1_agreement"] = top1_agreement_per_tok.detach().mean().item() - - metrics = OPDLossMetrics(**metrics_kwargs).to_dict() - - return LossOutput(loss=loss, per_token_loss=per_token_loss, metrics=metrics) +from xorl.objectives.opd_loss import * # noqa: F401,F403 diff --git a/src/xorl/ops/loss/opd_streaming_kl.py b/src/xorl/ops/loss/opd_streaming_kl.py index 44bafaf4..c3bb570d 100644 --- a/src/xorl/ops/loss/opd_streaming_kl.py +++ b/src/xorl/ops/loss/opd_streaming_kl.py @@ -1,820 +1,3 @@ -from __future__ import annotations +"""Moved to ``xorl.objectives.opd_streaming_kl`` (issue #78 phase 2); compat stub.""" -import math - -import torch - - -def _iter_weight_chunks(teacher_weight, vocab_size: int, chunk_rows: int): - if hasattr(teacher_weight, "iter_device_chunks"): - yield from teacher_weight.iter_device_chunks(chunk_rows) - return - for start, end in _iter_ranges(vocab_size, chunk_rows): - yield start, end, teacher_weight[start:end] - - -def _chunk_size(vocab_size: int, requested: int) -> int: - if requested <= 0 or requested >= vocab_size: - return vocab_size - return requested - - -def _iter_ranges(vocab_size: int, requested_chunk_size: int): - chunk_size = _chunk_size(vocab_size, requested_chunk_size) - for start in range(0, vocab_size, chunk_size): - yield start, min(start + chunk_size, vocab_size) - - -def _accum_dtype(reference: torch.Tensor) -> torch.dtype: - """fp32 accumulation by default, but preserve fp64 inputs (e.g. gradcheck). - - The streaming kernels accumulate in fp32 -- in production the lm-head tensors - are already fp32, so the per-chunk `.float()` upcast is a no-op. When a caller - passes fp64 inputs (the strongest gradcheck regime), truncating to fp32 would - swamp the numerical Jacobian with rounding noise; honoring fp64 there is - strictly more accurate and changes nothing for the fp32 production path. - """ - return torch.float64 if reference.dtype == torch.float64 else torch.float32 - - -def _update_online_logsumexp( - running_max: torch.Tensor, - running_sumexp: torch.Tensor, - logits: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - chunk_max = logits.max(dim=-1, keepdim=True).values - new_max = torch.maximum(running_max, chunk_max) - prev_scale = torch.where( - torch.isfinite(running_max), - (running_max - new_max).exp(), - torch.zeros_like(running_sumexp), - ) - chunk_sumexp = (logits - new_max).exp().sum(dim=-1, keepdim=True) - return new_max, running_sumexp * prev_scale + chunk_sumexp - - -class _StreamingReverseKL(torch.autograd.Function): - """Exact KL(student || teacher) over vocab chunks. - - This is the TileLang-facing OPD path: it exposes the same execution shape a - native kernel will use (stream vocab blocks, save only per-token statistics, - recompute logits in backward) while keeping a pure PyTorch implementation as - the portable fallback. - """ - - @staticmethod - def forward( - ctx, - student_hidden_states: torch.Tensor, - student_weight: torch.Tensor, - teacher_hidden_states: torch.Tensor, - labels: torch.Tensor, - teacher_weight, - ignore_index: int, - vocab_chunk_size: int, - ) -> torch.Tensor: - teacher_shape = tuple(int(x) for x in teacher_weight.shape) - if student_weight.shape[0] != teacher_shape[0]: - raise ValueError( - f"student vocab size ({student_weight.shape[0]}) must match teacher vocab size ({teacher_shape[0]})" - ) - - vocab_size = int(student_weight.shape[0]) - token_count = int(student_hidden_states.shape[0]) - valid = labels != ignore_index - neg_inf = -float("inf") - - # Single fused pass over vocab chunks. Reverse-KL decomposes as - # KL = sum_v p_s(v)*(s_v - t_v) - logZ_s + logZ_t - # and sum_v p_s(v)*(s_v - t_v) = A / Z_s with - # A = sum_v exp(s_v - s_max) * (s_v - t_v), - # which is online-accumulable (flash-attention style) alongside the - # student/teacher log-sum-exp normalizers. This halves the forward GEMM - # vs the prior two-pass form (a logsumexp pass + a probability pass) - # while still never materializing the full logits. The backward - # (recomputes logits from the saved normalizers) is unchanged, so the - # gradients are identical up to fp32 summation order. - s_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=torch.float32) - t_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=torch.float32) - s_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=torch.float32) - t_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=torch.float32) - weighted_diff = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=torch.float32) - - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): - s_logits = (student_hidden_states @ student_weight[start:end].t()).float() - t_logits = (teacher_hidden_states @ t_weight.t()).float() - # Student: online log-sum-exp + online sum_v exp(s_v - s_max)*(s_v - t_v). - chunk_s_max = s_logits.max(dim=-1, keepdim=True).values - new_s_max = torch.maximum(s_max, chunk_s_max) - s_scale = torch.where( - torch.isfinite(s_max), - (s_max - new_s_max).exp(), - torch.zeros_like(s_sumexp), - ) - exp_s = (s_logits - new_s_max).exp() - s_sumexp = s_sumexp * s_scale + exp_s.sum(dim=-1, keepdim=True) - weighted_diff = weighted_diff * s_scale + (exp_s * (s_logits - t_logits)).sum(dim=-1, keepdim=True) - s_max = new_s_max - # Teacher: online log-sum-exp normalizer only. - t_max, t_sumexp = _update_online_logsumexp(t_max, t_sumexp, t_logits) - - s_logz = s_sumexp.log() + s_max - t_logz = t_sumexp.log() + t_max - kl = (weighted_diff / s_sumexp - s_logz + t_logz).squeeze(-1) - kl = kl * valid.to(kl.dtype) - ctx.save_for_backward( - student_hidden_states, - student_weight, - teacher_hidden_states, - labels, - s_logz, - t_logz, - kl, - ) - ctx.teacher_weight = teacher_weight - ctx.ignore_index = ignore_index - ctx.vocab_chunk_size = vocab_chunk_size - return kl - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - ( - student_hidden_states, - student_weight, - teacher_hidden_states, - labels, - s_logz, - t_logz, - kl, - ) = ctx.saved_tensors - teacher_weight = ctx.teacher_weight - - valid = (labels != ctx.ignore_index).to(dtype=torch.float32, device=grad_output.device) - scale = grad_output.to(dtype=torch.float32) * valid - vocab_size = int(student_weight.shape[0]) - - grad_hidden = None - if ctx.needs_input_grad[0]: - grad_hidden = torch.zeros_like(student_hidden_states, dtype=torch.float32) - - grad_weight = None - if ctx.needs_input_grad[1]: - grad_weight = torch.zeros_like(student_weight, dtype=torch.float32) - - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, ctx.vocab_chunk_size): - s_weight = student_weight[start:end] - s_logits = (student_hidden_states @ s_weight.t()).float() - t_logits = (teacher_hidden_states @ t_weight.t()).float() - s_log_probs = s_logits - s_logz - t_log_probs = t_logits - t_logz - s_probs = s_log_probs.exp() - - # d KL(p_s || p_t) / d student_logits_i = - # p_s_i * (log p_s_i - log p_t_i - KL) - grad_logits = s_probs * (s_log_probs - t_log_probs - kl.unsqueeze(1)) - grad_logits = grad_logits * scale.unsqueeze(1) - - if grad_hidden is not None: - grad_hidden = grad_hidden + grad_logits @ s_weight.float() - if grad_weight is not None: - grad_weight[start:end] = grad_logits.t() @ student_hidden_states.float() - - if grad_hidden is not None: - grad_hidden = grad_hidden.to(student_hidden_states.dtype) - if grad_weight is not None: - grad_weight = grad_weight.to(student_weight.dtype) - - if hasattr(teacher_weight, "clear_device_cache"): - teacher_weight.clear_device_cache() - - return grad_hidden, grad_weight, None, None, None, None, None - - -class _StreamingForwardKL(torch.autograd.Function): - """Exact forward KL(teacher || student) over vocab chunks. - - The forward-KL counterpart of `_StreamingReverseKL`: same streaming shape - (stream vocab blocks, save only per-token statistics, recompute logits in - backward), never materializing the full [tokens, vocab] logits. This unblocks - `loss_mode='forward_kl_full'` on the streaming backend, where the compile - backend OOMs on the materialized full-vocab fp32 logits. - - Forward KL = KL(p_T || p_S) = sum_v p_T(v)*(log p_T(v) - log p_S(v)). Writing - log p_S(v) = s_v - s_logz and log p_T(v) = t_v - t_logz it telescopes to a - streaming-friendly form that never takes the log of a near-zero student prob: - - KL = s_logz - t_logz + (sum_v exp(t_v - t_max)*(t_v - s_v)) / t_sumexp - - The weighted term sum_v exp(t_v - t_max)*(t_v - s_v) is accumulated online - (flash-attention style) against the *teacher* running max/sumexp -- mirroring - `_StreamingReverseKL` but weighting by the teacher exp and using (t_v - s_v) - rather than the student exp and (s_v - t_v). The student/teacher log-sum-exp - normalizers are accumulated in the same pass. - - Backward w.r.t. the student logits is exact and far simpler than reverse KL: - the p_T*log p_T term is constant in z_S, so - d/d z_S[k] (-sum_v p_T(v) log p_S(v)) = p_S(k) - p_T(k), - i.e. grad_logits = (p_S - p_T)*grad_output*valid. No per-token KL term and no - p_S weighting of log-ratios. The teacher is detached, so no teacher grads. - """ - - @staticmethod - def forward( - ctx, - student_hidden_states: torch.Tensor, - student_weight: torch.Tensor, - teacher_hidden_states: torch.Tensor, - labels: torch.Tensor, - teacher_weight, - ignore_index: int, - vocab_chunk_size: int, - ) -> torch.Tensor: - teacher_shape = tuple(int(x) for x in teacher_weight.shape) - if student_weight.shape[0] != teacher_shape[0]: - raise ValueError( - f"student vocab size ({student_weight.shape[0]}) must match teacher vocab size ({teacher_shape[0]})" - ) - - vocab_size = int(student_weight.shape[0]) - token_count = int(student_hidden_states.shape[0]) - valid = labels != ignore_index - neg_inf = -float("inf") - acc = _accum_dtype(student_hidden_states) - - # Single fused pass over vocab chunks. Forward-KL decomposes as - # KL = s_logz - t_logz + sum_v p_T(v)*(t_v - s_v) - # and sum_v p_T(v)*(t_v - s_v) = B / Z_t with - # B = sum_v exp(t_v - t_max) * (t_v - s_v), - # which is online-accumulable (flash-attention style) alongside the - # student/teacher log-sum-exp normalizers. The teacher running max/sumexp - # weight the difference (mirror image of reverse KL's student-weighted A). - s_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=acc) - t_max = torch.full((token_count, 1), neg_inf, device=student_hidden_states.device, dtype=acc) - s_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=acc) - t_sumexp = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=acc) - weighted_diff = torch.zeros((token_count, 1), device=student_hidden_states.device, dtype=acc) - - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): - s_logits = (student_hidden_states @ student_weight[start:end].t()).to(acc) - t_logits = (teacher_hidden_states @ t_weight.t()).to(acc) - # Teacher: online log-sum-exp + online sum_v exp(t_v - t_max)*(t_v - s_v). - chunk_t_max = t_logits.max(dim=-1, keepdim=True).values - new_t_max = torch.maximum(t_max, chunk_t_max) - t_scale = torch.where( - torch.isfinite(t_max), - (t_max - new_t_max).exp(), - torch.zeros_like(t_sumexp), - ) - exp_t = (t_logits - new_t_max).exp() - t_sumexp = t_sumexp * t_scale + exp_t.sum(dim=-1, keepdim=True) - weighted_diff = weighted_diff * t_scale + (exp_t * (t_logits - s_logits)).sum(dim=-1, keepdim=True) - t_max = new_t_max - # Student: online log-sum-exp normalizer only. - s_max, s_sumexp = _update_online_logsumexp(s_max, s_sumexp, s_logits) - - s_logz = s_sumexp.log() + s_max - t_logz = t_sumexp.log() + t_max - kl = (s_logz - t_logz + weighted_diff / t_sumexp).squeeze(-1) - kl = kl * valid.to(kl.dtype) - # Backward only needs s_logz/t_logz to recompute p_S and p_T per chunk; the - # forward-KL gradient (p_S - p_T) does not depend on the per-token KL value. - ctx.save_for_backward( - student_hidden_states, - student_weight, - teacher_hidden_states, - labels, - s_logz, - t_logz, - ) - ctx.teacher_weight = teacher_weight - ctx.ignore_index = ignore_index - ctx.vocab_chunk_size = vocab_chunk_size - return kl - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - ( - student_hidden_states, - student_weight, - teacher_hidden_states, - labels, - s_logz, - t_logz, - ) = ctx.saved_tensors - teacher_weight = ctx.teacher_weight - acc = _accum_dtype(student_hidden_states) - - valid = (labels != ctx.ignore_index).to(dtype=acc, device=grad_output.device) - scale = grad_output.to(dtype=acc) * valid - vocab_size = int(student_weight.shape[0]) - - grad_hidden = None - if ctx.needs_input_grad[0]: - grad_hidden = torch.zeros_like(student_hidden_states, dtype=acc) - - grad_weight = None - if ctx.needs_input_grad[1]: - grad_weight = torch.zeros_like(student_weight, dtype=acc) - - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, ctx.vocab_chunk_size): - s_weight = student_weight[start:end] - s_logits = (student_hidden_states @ s_weight.t()).to(acc) - t_logits = (teacher_hidden_states @ t_weight.t()).to(acc) - s_probs = (s_logits - s_logz).exp() - t_probs = (t_logits - t_logz).exp() - - # d KL(p_t || p_s) / d student_logits_k = p_s_k - p_t_k. - grad_logits = (s_probs - t_probs) * scale.unsqueeze(1) - - if grad_hidden is not None: - grad_hidden = grad_hidden + grad_logits @ s_weight.to(acc) - if grad_weight is not None: - grad_weight[start:end] = grad_logits.t() @ student_hidden_states.to(acc) - - if grad_hidden is not None: - grad_hidden = grad_hidden.to(student_hidden_states.dtype) - if grad_weight is not None: - grad_weight = grad_weight.to(student_weight.dtype) - - if hasattr(teacher_weight, "clear_device_cache"): - teacher_weight.clear_device_cache() - - return grad_hidden, grad_weight, None, None, None, None, None - - -def streaming_forward_kl_function( - student_hidden_states: torch.Tensor, - student_weight: torch.Tensor, - teacher_hidden_states: torch.Tensor, - teacher_weight: torch.Tensor, - labels: torch.Tensor, - ignore_index: int = -100, - vocab_chunk_size: int | None = 32768, -) -> torch.Tensor: - """Compute per-token forward KL(teacher||student) without materializing full-vocab logits.""" - if vocab_chunk_size is None: - vocab_chunk_size = 32768 - elif vocab_chunk_size <= 0: - vocab_chunk_size = int(student_weight.shape[0]) - if not math.isfinite(float(vocab_chunk_size)): - raise ValueError(f"Invalid vocab_chunk_size={vocab_chunk_size}") - return _StreamingForwardKL.apply( - student_hidden_states, - student_weight, - teacher_hidden_states.detach(), - labels, - teacher_weight.detach() if torch.is_tensor(teacher_weight) else teacher_weight, - int(ignore_index), - int(vocab_chunk_size), - ) - - -@torch.no_grad() -def streaming_full_vocab_diagnostics( - student_hidden_states: torch.Tensor, - student_weight: torch.Tensor, - teacher_hidden_states: torch.Tensor, - teacher_weight, - vocab_chunk_size: int | None = 32768, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Per-token (teacher_entropy, student_entropy, top1_agreement) diagnostics. - - Streams vocab chunks exactly like `_StreamingReverseKL` (one logsumexp pass, - one probability pass) so the full logits are never materialized. No-grad and - opt-in: this is the streaming-backend counterpart of the compile backend's - `*_with_diag` kernels, costing one extra full-vocab pass per micro-batch. - """ - if vocab_chunk_size is None: - vocab_chunk_size = 32768 - elif vocab_chunk_size <= 0: - vocab_chunk_size = int(student_weight.shape[0]) - vocab_size = int(student_weight.shape[0]) - token_count = int(student_hidden_states.shape[0]) - device = student_hidden_states.device - neg_inf = -float("inf") - - s_max = torch.full((token_count, 1), neg_inf, device=device, dtype=torch.float32) - t_max = torch.full((token_count, 1), neg_inf, device=device, dtype=torch.float32) - s_sumexp = torch.zeros((token_count, 1), device=device, dtype=torch.float32) - t_sumexp = torch.zeros((token_count, 1), device=device, dtype=torch.float32) - s_best_idx = torch.zeros(token_count, device=device, dtype=torch.long) - t_best_idx = torch.zeros(token_count, device=device, dtype=torch.long) - s_best_val = torch.full((token_count,), neg_inf, device=device, dtype=torch.float32) - t_best_val = torch.full((token_count,), neg_inf, device=device, dtype=torch.float32) - - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): - s_weight = student_weight[start:end].to(student_hidden_states.dtype) - t_weight = t_weight.to(teacher_hidden_states.dtype) - s_logits = (student_hidden_states @ s_weight.t()).float() - t_logits = (teacher_hidden_states @ t_weight.t()).float() - s_max, s_sumexp = _update_online_logsumexp(s_max, s_sumexp, s_logits) - t_max, t_sumexp = _update_online_logsumexp(t_max, t_sumexp, t_logits) - s_chunk_val, s_chunk_idx = s_logits.max(dim=-1) - t_chunk_val, t_chunk_idx = t_logits.max(dim=-1) - s_better = s_chunk_val > s_best_val - t_better = t_chunk_val > t_best_val - s_best_idx = torch.where(s_better, s_chunk_idx + start, s_best_idx) - s_best_val = torch.maximum(s_best_val, s_chunk_val) - t_best_idx = torch.where(t_better, t_chunk_idx + start, t_best_idx) - t_best_val = torch.maximum(t_best_val, t_chunk_val) - - s_logz = s_sumexp.log() + s_max - t_logz = t_sumexp.log() + t_max - s_plogp = torch.zeros(token_count, device=device, dtype=torch.float32) - t_plogp = torch.zeros(token_count, device=device, dtype=torch.float32) - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): - s_weight = student_weight[start:end].to(student_hidden_states.dtype) - t_weight = t_weight.to(teacher_hidden_states.dtype) - s_logits = (student_hidden_states @ s_weight.t()).float() - t_logits = (teacher_hidden_states @ t_weight.t()).float() - s_log_probs = s_logits - s_logz - t_log_probs = t_logits - t_logz - s_plogp = s_plogp + (s_log_probs.exp() * s_log_probs).sum(dim=-1) - t_plogp = t_plogp + (t_log_probs.exp() * t_log_probs).sum(dim=-1) - - if hasattr(teacher_weight, "clear_device_cache"): - teacher_weight.clear_device_cache() - - teacher_entropy = -t_plogp - student_entropy = -s_plogp - top1_agreement = (s_best_idx == t_best_idx).float() - return teacher_entropy, student_entropy, top1_agreement - - -def streaming_reverse_kl_function( - student_hidden_states: torch.Tensor, - student_weight: torch.Tensor, - teacher_hidden_states: torch.Tensor, - teacher_weight: torch.Tensor, - labels: torch.Tensor, - ignore_index: int = -100, - vocab_chunk_size: int | None = 32768, -) -> torch.Tensor: - """Compute per-token reverse KL without materializing full-vocab logits.""" - if vocab_chunk_size is None: - vocab_chunk_size = 32768 - elif vocab_chunk_size <= 0: - vocab_chunk_size = int(student_weight.shape[0]) - if not math.isfinite(float(vocab_chunk_size)): - raise ValueError(f"Invalid vocab_chunk_size={vocab_chunk_size}") - return _StreamingReverseKL.apply( - student_hidden_states, - student_weight, - teacher_hidden_states.detach(), - labels, - teacher_weight.detach() if torch.is_tensor(teacher_weight) else teacher_weight, - int(ignore_index), - int(vocab_chunk_size), - ) - - -class _StreamingReverseKLLowMem(torch.autograd.Function): - """Memory-lean reverse KL(student||teacher), bit-exact with _StreamingReverseKL - under an fp32 lm-head, but without holding full fp32 copies of the lm-head - weights or a second full [V,H] grad buffer. - - Two memory levers, both gradient-identical to the current OPD fp32 path: - - 1. Per-chunk fp32 upcast. The student/teacher lm-head weights stay in their - native (bf16) dtype; each [chunk, H] vocab slice is upcast to - `compute_dtype` (fp32) inside the loop, right before the matmul. Slicing - commutes with the elementwise upcast, so the per-chunk fp32 matmul is - identical to multiplying by a whole pre-upcast fp32 weight -- but the - two ~2 GB fp32 weight copies (`weight.float()`) are never resident. - - 2. In-place weight grad. The weight gradient is streamed straight into the - leaf `student_weight.grad` (created lazily, in the weight's native - dtype), one vocab chunk at a time, instead of allocating a full [V,H] - buffer that autograd then adds into `.grad`. Vocab chunks partition the - grad rows disjointly, so this is exact; it removes the full extra buffer - (and its grad-accumulation doubling). `student_weight` must be the leaf - parameter when `inplace_weight_grad=True`. - - This is the AMDAHL-029..033 1-node unblock: it removes ~2-4 GB of resident - lm-head memory that tipped the streaming-KL backward over on one node. - """ - - @staticmethod - def forward( - ctx, - student_hidden_states: torch.Tensor, - student_weight: torch.Tensor, - teacher_hidden_states: torch.Tensor, - labels: torch.Tensor, - teacher_weight, - ignore_index: int, - vocab_chunk_size: int, - compute_dtype: torch.dtype, - inplace_weight_grad: bool, - ) -> torch.Tensor: - teacher_shape = tuple(int(x) for x in teacher_weight.shape) - if student_weight.shape[0] != teacher_shape[0]: - raise ValueError( - f"student vocab size ({student_weight.shape[0]}) must match teacher vocab size ({teacher_shape[0]})" - ) - vocab_size = int(student_weight.shape[0]) - token_count = int(student_hidden_states.shape[0]) - valid = labels != ignore_index - neg_inf = -float("inf") - dev = student_hidden_states.device - sh = student_hidden_states.to(compute_dtype) - th = teacher_hidden_states.to(compute_dtype) - - s_max = torch.full((token_count, 1), neg_inf, device=dev, dtype=torch.float32) - t_max = torch.full((token_count, 1), neg_inf, device=dev, dtype=torch.float32) - s_sumexp = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) - t_sumexp = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): - s_logits = (sh @ student_weight[start:end].to(compute_dtype).t()).float() - t_logits = (th @ t_weight.to(compute_dtype).t()).float() - s_max, s_sumexp = _update_online_logsumexp(s_max, s_sumexp, s_logits) - t_max, t_sumexp = _update_online_logsumexp(t_max, t_sumexp, t_logits) - - s_logz = s_sumexp.log() + s_max - t_logz = t_sumexp.log() + t_max - kl = torch.zeros(token_count, device=dev, dtype=torch.float32) - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): - s_logits = (sh @ student_weight[start:end].to(compute_dtype).t()).float() - t_logits = (th @ t_weight.to(compute_dtype).t()).float() - s_log_probs = s_logits - s_logz - t_log_probs = t_logits - t_logz - s_probs = s_log_probs.exp() - kl = kl + (s_probs * (s_log_probs - t_log_probs)).sum(dim=-1) - - kl = kl * valid.to(kl.dtype) - ctx.save_for_backward(student_hidden_states, student_weight, teacher_hidden_states, labels, s_logz, t_logz, kl) - ctx.teacher_weight = teacher_weight - ctx.ignore_index = ignore_index - ctx.vocab_chunk_size = vocab_chunk_size - ctx.compute_dtype = compute_dtype - ctx.inplace_weight_grad = inplace_weight_grad - ctx.student_weight_ref = student_weight - return kl - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - (student_hidden_states, student_weight, teacher_hidden_states, labels, s_logz, t_logz, kl) = ctx.saved_tensors - teacher_weight = ctx.teacher_weight - cdt = ctx.compute_dtype - valid = (labels != ctx.ignore_index).to(dtype=torch.float32, device=grad_output.device) - scale = grad_output.to(dtype=torch.float32) * valid - vocab_size = int(student_weight.shape[0]) - sh = student_hidden_states.to(cdt) - th = teacher_hidden_states.to(cdt) - - grad_hidden = None - if ctx.needs_input_grad[0]: - grad_hidden = torch.zeros_like(student_hidden_states, dtype=torch.float32) - - # Weight grad: either stream in place into the leaf .grad (no second full - # buffer), or build one buffer in the weight's native dtype and return it. - accumulate_weight = ctx.needs_input_grad[1] or ctx.inplace_weight_grad - inplace = ctx.inplace_weight_grad - grad_weight = None - if accumulate_weight: - if inplace: - wparam = ctx.student_weight_ref - if wparam.grad is None: - wparam.grad = torch.zeros_like(wparam) - else: - grad_weight = torch.zeros_like(student_weight) - - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, ctx.vocab_chunk_size): - s_weight = student_weight[start:end].to(cdt) - s_logits = (sh @ s_weight.t()).float() - t_logits = (th @ t_weight.to(cdt).t()).float() - s_log_probs = s_logits - s_logz - t_log_probs = t_logits - t_logz - s_probs = s_log_probs.exp() - grad_logits = s_probs * (s_log_probs - t_log_probs - kl.unsqueeze(1)) - grad_logits = grad_logits * scale.unsqueeze(1) - if grad_hidden is not None: - grad_hidden = grad_hidden + grad_logits @ s_weight - if accumulate_weight: - chunk = (grad_logits.t() @ sh).to(student_weight.dtype) - if inplace: - ctx.student_weight_ref.grad[start:end].add_(chunk) - else: - grad_weight[start:end] = chunk - - if grad_hidden is not None: - grad_hidden = grad_hidden.to(student_hidden_states.dtype) - if hasattr(teacher_weight, "clear_device_cache"): - teacher_weight.clear_device_cache() - - return grad_hidden, grad_weight, None, None, None, None, None, None, None - - -def streaming_reverse_kl_lowmem_function( - student_hidden_states: torch.Tensor, - student_weight: torch.Tensor, - teacher_hidden_states: torch.Tensor, - teacher_weight: torch.Tensor, - labels: torch.Tensor, - ignore_index: int = -100, - vocab_chunk_size: int | None = 32768, - compute_dtype: torch.dtype = torch.float32, - inplace_weight_grad: bool = False, -) -> torch.Tensor: - """Memory-lean per-token reverse KL: native-dtype lm-head weights upcast per - vocab chunk to `compute_dtype`, gradient-identical to the fp32 path. - - Pass the lm-head weights in their stored (bf16) dtype -- do NOT pre-cast them - to fp32. `student_hidden_states` may be fp32 (cheap; [N,H]). Set - `inplace_weight_grad=True` only when `student_weight` is the leaf parameter - whose `.grad` the optimizer reads (saves the second full [V,H] buffer). - """ - if vocab_chunk_size is None: - vocab_chunk_size = 32768 - elif vocab_chunk_size <= 0: - vocab_chunk_size = int(student_weight.shape[0]) - if not math.isfinite(float(vocab_chunk_size)): - raise ValueError(f"Invalid vocab_chunk_size={vocab_chunk_size}") - return _StreamingReverseKLLowMem.apply( - student_hidden_states, - student_weight, - teacher_hidden_states.detach(), - labels, - teacher_weight.detach() if torch.is_tensor(teacher_weight) else teacher_weight, - int(ignore_index), - int(vocab_chunk_size), - compute_dtype, - bool(inplace_weight_grad), - ) - - -class _StreamingForwardKLLowMem(torch.autograd.Function): - """Memory-lean forward KL(teacher||student), bit-exact with _StreamingForwardKL - under an fp32 lm-head, but without holding full fp32 copies of the lm-head - weights or a second full [V,H] grad buffer. - - The forward-KL counterpart of `_StreamingReverseKLLowMem`, with the identical - two memory levers (both gradient-identical to the fp32 forward-KL path): - - 1. Per-chunk fp32 upcast. The student/teacher lm-head weights stay in their - native (bf16) dtype; each [chunk, H] vocab slice is upcast to - `compute_dtype` (fp32) inside the loop before the matmul. Slicing - commutes with the elementwise upcast, so this matches multiplying by a - whole pre-upcast fp32 weight without ever holding the fp32 copies. - - 2. In-place weight grad. The weight gradient is streamed straight into the - leaf `student_weight.grad` one vocab chunk at a time (chunks partition - the grad rows disjointly, so it is exact), avoiding a full [V,H] buffer. - `student_weight` must be the leaf parameter when `inplace_weight_grad`. - - Forward uses the same telescoped, never-log-a-near-zero-prob form as - `_StreamingForwardKL`: - KL = s_logz - t_logz + (sum_v exp(t_v - t_max)*(t_v - s_v)) / t_sumexp. - Backward is exact and simple: grad_logits = (p_S - p_T)*grad_output*valid. - """ - - @staticmethod - def forward( - ctx, - student_hidden_states: torch.Tensor, - student_weight: torch.Tensor, - teacher_hidden_states: torch.Tensor, - labels: torch.Tensor, - teacher_weight, - ignore_index: int, - vocab_chunk_size: int, - compute_dtype: torch.dtype, - inplace_weight_grad: bool, - ) -> torch.Tensor: - teacher_shape = tuple(int(x) for x in teacher_weight.shape) - if student_weight.shape[0] != teacher_shape[0]: - raise ValueError( - f"student vocab size ({student_weight.shape[0]}) must match teacher vocab size ({teacher_shape[0]})" - ) - vocab_size = int(student_weight.shape[0]) - token_count = int(student_hidden_states.shape[0]) - valid = labels != ignore_index - neg_inf = -float("inf") - dev = student_hidden_states.device - sh = student_hidden_states.to(compute_dtype) - th = teacher_hidden_states.to(compute_dtype) - - s_max = torch.full((token_count, 1), neg_inf, device=dev, dtype=torch.float32) - t_max = torch.full((token_count, 1), neg_inf, device=dev, dtype=torch.float32) - s_sumexp = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) - t_sumexp = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) - weighted_diff = torch.zeros((token_count, 1), device=dev, dtype=torch.float32) - # Single fused pass (same telescoped form as `_StreamingForwardKL`): online - # student/teacher log-sum-exp plus the online teacher-weighted (t_v - s_v) - # accumulation, so the KL is computed in one sweep over vocab chunks rather - # than a second full-vocab GEMM pass. KL = s_logz - t_logz + weighted_diff / t_sumexp. - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, vocab_chunk_size): - s_logits = (sh @ student_weight[start:end].to(compute_dtype).t()).float() - t_logits = (th @ t_weight.to(compute_dtype).t()).float() - chunk_t_max = t_logits.max(dim=-1, keepdim=True).values - new_t_max = torch.maximum(t_max, chunk_t_max) - t_scale = torch.where( - torch.isfinite(t_max), - (t_max - new_t_max).exp(), - torch.zeros_like(t_sumexp), - ) - exp_t = (t_logits - new_t_max).exp() - t_sumexp = t_sumexp * t_scale + exp_t.sum(dim=-1, keepdim=True) - weighted_diff = weighted_diff * t_scale + (exp_t * (t_logits - s_logits)).sum(dim=-1, keepdim=True) - t_max = new_t_max - s_max, s_sumexp = _update_online_logsumexp(s_max, s_sumexp, s_logits) - - s_logz = s_sumexp.log() + s_max - t_logz = t_sumexp.log() + t_max - kl = (s_logz - t_logz + weighted_diff / t_sumexp).squeeze(-1) - kl = kl * valid.to(kl.dtype) - ctx.save_for_backward(student_hidden_states, student_weight, teacher_hidden_states, labels, s_logz, t_logz) - ctx.teacher_weight = teacher_weight - ctx.ignore_index = ignore_index - ctx.vocab_chunk_size = vocab_chunk_size - ctx.compute_dtype = compute_dtype - ctx.inplace_weight_grad = inplace_weight_grad - ctx.student_weight_ref = student_weight - return kl - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - (student_hidden_states, student_weight, teacher_hidden_states, labels, s_logz, t_logz) = ctx.saved_tensors - teacher_weight = ctx.teacher_weight - cdt = ctx.compute_dtype - valid = (labels != ctx.ignore_index).to(dtype=torch.float32, device=grad_output.device) - scale = grad_output.to(dtype=torch.float32) * valid - vocab_size = int(student_weight.shape[0]) - sh = student_hidden_states.to(cdt) - th = teacher_hidden_states.to(cdt) - - grad_hidden = None - if ctx.needs_input_grad[0]: - grad_hidden = torch.zeros_like(student_hidden_states, dtype=torch.float32) - - accumulate_weight = ctx.needs_input_grad[1] or ctx.inplace_weight_grad - inplace = ctx.inplace_weight_grad - grad_weight = None - if accumulate_weight: - if inplace: - wparam = ctx.student_weight_ref - if wparam.grad is None: - wparam.grad = torch.zeros_like(wparam) - else: - grad_weight = torch.zeros_like(student_weight) - - for start, end, t_weight in _iter_weight_chunks(teacher_weight, vocab_size, ctx.vocab_chunk_size): - s_weight = student_weight[start:end].to(cdt) - s_logits = (sh @ s_weight.t()).float() - t_logits = (th @ t_weight.to(cdt).t()).float() - s_probs = (s_logits - s_logz).exp() - t_probs = (t_logits - t_logz).exp() - grad_logits = (s_probs - t_probs) * scale.unsqueeze(1) - if grad_hidden is not None: - grad_hidden = grad_hidden + grad_logits @ s_weight - if accumulate_weight: - chunk = (grad_logits.t() @ sh).to(student_weight.dtype) - if inplace: - ctx.student_weight_ref.grad[start:end].add_(chunk) - else: - grad_weight[start:end] = chunk - - if grad_hidden is not None: - grad_hidden = grad_hidden.to(student_hidden_states.dtype) - if hasattr(teacher_weight, "clear_device_cache"): - teacher_weight.clear_device_cache() - - return grad_hidden, grad_weight, None, None, None, None, None, None, None - - -def streaming_forward_kl_lowmem_function( - student_hidden_states: torch.Tensor, - student_weight: torch.Tensor, - teacher_hidden_states: torch.Tensor, - teacher_weight: torch.Tensor, - labels: torch.Tensor, - ignore_index: int = -100, - vocab_chunk_size: int | None = 32768, - compute_dtype: torch.dtype = torch.float32, - inplace_weight_grad: bool = False, -) -> torch.Tensor: - """Memory-lean per-token forward KL: native-dtype lm-head weights upcast per - vocab chunk to `compute_dtype`, gradient-identical to the fp32 path. - - Pass the lm-head weights in their stored (bf16) dtype -- do NOT pre-cast them - to fp32. `student_hidden_states` may be fp32 (cheap; [N,H]). Set - `inplace_weight_grad=True` only when `student_weight` is the leaf parameter - whose `.grad` the optimizer reads (saves the second full [V,H] buffer). - """ - if vocab_chunk_size is None: - vocab_chunk_size = 32768 - elif vocab_chunk_size <= 0: - vocab_chunk_size = int(student_weight.shape[0]) - if not math.isfinite(float(vocab_chunk_size)): - raise ValueError(f"Invalid vocab_chunk_size={vocab_chunk_size}") - return _StreamingForwardKLLowMem.apply( - student_hidden_states, - student_weight, - teacher_hidden_states.detach(), - labels, - teacher_weight.detach() if torch.is_tensor(teacher_weight) else teacher_weight, - int(ignore_index), - int(vocab_chunk_size), - compute_dtype, - bool(inplace_weight_grad), - ) +from xorl.objectives.opd_streaming_kl import * # noqa: F401,F403 diff --git a/src/xorl/ops/loss/policy_loss.py b/src/xorl/ops/loss/policy_loss.py index 107bd817..9485f5e7 100644 --- a/src/xorl/ops/loss/policy_loss.py +++ b/src/xorl/ops/loss/policy_loss.py @@ -1,394 +1,3 @@ -""" -Policy Loss with PPO Clipping and TIS Correction. +"""Moved to ``xorl.objectives.policy_loss`` (issue #78 phase 2); compat stub.""" -This module provides the policy loss functions including: -- PPO-style clipped policy gradient loss -- Temporal Importance Sampling (TIS) correction -- Combined policy_loss_function -""" - -from __future__ import annotations - -import logging -from typing import Any, Dict, Optional, Tuple - -import torch -import torch.distributed as dist - -from xorl.ops.exact_sampling_transforms import TOP_K_ALL -from xorl.ops.loss.loss_output import LossOutput -from xorl.ops.loss.per_token_ce import compute_per_token_ce -from xorl.ops.loss.reducers import Reducer, TokenPartial - - -logger = logging.getLogger(__name__) - - -K3_DEBUG_THRESHOLDS = ( - ("1e_minus_6", 1e-6), - ("1e_minus_4", 1e-4), - ("1e_minus_3", 1e-3), - ("1e_minus_2", 1e-2), - ("1e_minus_1", 1e-1), - ("1", 1.0), -) - - -@torch.compile(dynamic=True) -def compute_ppo_loss( - ppo_kl: torch.Tensor, - advantages: torch.Tensor, - eps_clip: float = 0.2, - eps_clip_high: float = 0.2, - eps_clip_c: Optional[float] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - PPO-style clipped policy loss. - - Args: - ppo_kl: KL divergence tensor (old_log_probs - new_log_probs) - advantages: Per-token advantages - eps_clip: Lower clip ratio (default: 0.2) - eps_clip_high: Upper clip ratio (default: 0.2) - eps_clip_c: Dual-clip ratio for negative advantages (optional) - - Returns: - pg_losses: Clipped policy gradient losses - is_clipped: Per-token boolean mask of clipped tokens - ratio: Importance sampling ratio exp(-ppo_kl) - """ - ratio = (-ppo_kl).exp() - pg_losses1 = -ratio * advantages - pg_losses2 = -ratio.clamp(1 - eps_clip, 1 + eps_clip_high) * advantages - clip_pg_losses1 = torch.maximum(pg_losses1, pg_losses2) - is_clipped = torch.gt(pg_losses2, pg_losses1) - - # Optional dual-clip for negative advantages - if eps_clip_c is not None: - assert eps_clip_c > 1.0, f"eps_clip_c must be > 1.0, got {eps_clip_c}" - pg_losses3 = -eps_clip_c * advantages - clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1) - pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1) - # Also mark dual-clipped tokens - is_dual_clipped = (advantages < 0) & torch.lt(pg_losses3, clip_pg_losses1) - is_clipped = is_clipped | is_dual_clipped - else: - pg_losses = clip_pg_losses1 - - return pg_losses, is_clipped, ratio - - -def apply_tis_correction( - pg_loss: torch.Tensor, - train_log_probs: torch.Tensor, - rollout_log_probs: torch.Tensor, - valid_mask: torch.Tensor, - metric_reducer: Reducer, - tis_clip_low: float = 0.1, - tis_clip_high: float = 2.0, -) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - """ - Apply Temporal Importance Sampling (TIS) correction. - - TIS corrects for the distribution shift between rollout time and training time. - The TIS weight is: exp(train_log_probs - rollout_log_probs) - - Args: - pg_loss: Policy gradient loss tensor - train_log_probs: Log probabilities from current training step - rollout_log_probs: Log probabilities from rollout/inference - valid_mask: Mask for valid tokens - metric_reducer: Reducer applied to per-token mean metrics (tis_mean, - tis_clipfrac). min/max are local reductions and bypass it. - tis_clip_low: Lower bound for TIS clipping (default: 0.1) - tis_clip_high: Upper bound for TIS clipping (default: 2.0) - - Returns: - Tuple of (corrected_loss, metrics_dict) - """ - # Compute TIS weights: ratio of train vs rollout distributions - tis = torch.exp(train_log_probs - rollout_log_probs) - - # Clip TIS weights to prevent extreme values - tis_clipped = torch.clamp(tis, min=tis_clip_low, max=tis_clip_high) - - # Apply TIS correction to loss - corrected_loss = pg_loss * tis_clipped - - valid_mask_f = valid_mask.float() - tis_clipfrac_per_token = (tis_clipped != tis).float() - # ±inf identity on empty ranks lets cross-rank MIN/MAX-allreduce ignore empty contributors. - if valid_mask.any(): - tis_min = tis.masked_fill(~valid_mask, float("inf")).min() - tis_max = tis.masked_fill(~valid_mask, float("-inf")).max() - else: - tis_min = tis.new_tensor(float("inf")) - tis_max = tis.new_tensor(float("-inf")) - tis_metrics = { - "tis_mean": metric_reducer(tis, valid_mask_f), - "tis_min": tis_min, - "tis_max": tis_max, - "tis_clipfrac": metric_reducer(tis_clipfrac_per_token, valid_mask_f), - } - - return corrected_loss, tis_metrics - - -def policy_loss_function( - hidden_states: torch.Tensor, - weight: torch.Tensor, - labels: torch.Tensor, - old_logprobs: torch.Tensor, - advantages: torch.Tensor, - rollout_logprobs: Optional[torch.Tensor] = None, - ignore_index: int = -100, - eps_clip: float = 0.2, - eps_clip_high: float = 0.2, - eps_clip_c: Optional[float] = None, - tis_clip_low: float = 0.1, - tis_clip_high: float = 2.0, - use_tis: bool = False, - use_liger: bool = True, - num_chunks: int = 8, - ce_mode: str = "compiled", - compute_kl_stats: bool = False, - tp_group: Optional[dist.ProcessGroup] = None, - lm_head_fp32: bool = False, - icepop_beta: Optional[float] = None, - loss_reducer: Optional[Reducer] = None, - metric_reducer: Optional[Reducer] = None, - lm_head: Optional[torch.nn.Module] = None, - logprob_temperature: float = 1.0, - logprob_top_k: int | torch.Tensor = TOP_K_ALL, - logprob_top_p: float | torch.Tensor = 1.0, - logprob_min_p: float | torch.Tensor = 0.0, -) -> "LossOutput": - """ - Policy loss with PPO clipping, optional IcePop masking, and optional TIS correction. - - This implements the loss function which includes: - 1. PPO-style clipping on the importance sampling ratio - 2. Optional IcePop hard masking (GLM-5): zeros gradient for tokens where ratio is outside [1/β, β] - 3. Optional Temporal Importance Sampling (TIS) correction for off-policy data - - Supports multiple computation modes: - - "compiled": RECOMMENDED. torch.compile (1.6x speed, 16% memory) - - "eager": Simple F.cross_entropy baseline (may OOM at 32K) - - Args: - hidden_states: Model hidden states, shape (batch, seq_len, hidden_dim) - weight: LM head weight matrix, shape (vocab_size, hidden_dim) - labels: Target token IDs, shape (batch, seq_len). Already next-token aligned. - old_logprobs: Old policy log probabilities from sampling, shape (batch, seq_len) - advantages: Per-token advantages, shape (batch, seq_len) - rollout_logprobs: Optional rollout log probabilities for TIS correction, shape (batch, seq_len) - ignore_index: Index to ignore in loss computation (default: -100) - eps_clip: Lower clip ratio for PPO (default: 0.2) - eps_clip_high: Upper clip ratio for PPO (default: 0.2) - eps_clip_c: Dual-clip ratio for negative advantages (optional) - tis_clip_low: Lower bound for TIS clipping (default: 0.1) - tis_clip_high: Upper bound for TIS clipping (default: 2.0) - use_tis: Whether to apply TIS correction (default: False) - use_liger: Kept for API compatibility (ignored) - num_chunks: Number of chunks for auto_chunker (default: 8). Only used when ce_mode="compiled". - ce_mode: Cross-entropy mode - "compiled" (recommended) or "eager" - tp_group: TP process group for vocab-parallel cross-entropy (default: None) - compute_kl_stats: If True, compute and return full KL statistics in metrics dict - (kl_sample_train_k3, entropy_sample, ratio stats). - If False (default), only return valid_tokens and pg_clipfrac. - loss_reducer: Reduces per-token loss to a scalar partial share. None => - ``TokenPartial(scale=valid_mask.sum())`` (legacy local token-mean; does - not compose across micro-batches/ranks). Pass a shared global-scale - reducer to make summed partial shares recover the global loss. - metric_reducer: Reduces per-token /mean metrics (pg_clipfrac, icepop_maskfrac, - tis_mean, tis_clipfrac, kl_sample_train_k3, entropy_sample, ratio_mean) - the same way. ratio_min/ratio_max/tis_min/tis_max stay local scalars. - logprob_temperature: Temperature applied to trainer logits before - selected-token logprob calculation. ``1.0`` is raw policy logprobs; - setting this to the rollout temperature yields Slime-style behavior - logprob semantics for PPO ratios and K3 stats. - - Returns: - LossOutput with loss, per_token_logprobs (new logprobs), and metrics. - """ - - # Store original shape - original_shape = labels.shape - - # Flatten tensors - labels_flat = labels.view(-1) - hidden_states_flat = hidden_states.view(-1, hidden_states.size(-1)) - old_logprobs_flat = old_logprobs.view(-1) - advantages_flat = advantages.view(-1) - - # Create mask for valid tokens (use labels != ignore_index) - valid_mask = labels_flat != ignore_index - valid_mask_f = valid_mask.float() - valid_count = valid_mask.sum() - - if loss_reducer is None: - loss_reducer = TokenPartial(scale=valid_count.float()) - if metric_reducer is None: - metric_reducer = TokenPartial(scale=valid_count.float()) - - # Compute cross-entropy (supports vocab-parallel TP via tp_group) - per_token_ce = compute_per_token_ce( - hidden_states_flat, - weight, - labels_flat, - ignore_index, - ce_mode, - num_chunks, - tp_group=tp_group, - lm_head_fp32=lm_head_fp32, - lm_head=lm_head, - logprob_temperature=logprob_temperature, - logprob_top_k=logprob_top_k, - logprob_top_p=logprob_top_p, - logprob_min_p=logprob_min_p, - ) - - current_support = torch.isfinite(per_token_ce) - new_logprobs_flat = -per_token_ce.detach() - - # Compute PPO KL: old_log_probs - new_log_probs - ppo_kl = old_logprobs_flat - new_logprobs_flat - - # Mask invalid positions - ppo_kl = ppo_kl.masked_fill(~valid_mask, 0.0) - advantages_masked = advantages_flat.masked_fill(~valid_mask, 0.0) - - # Computed BEFORE compute_ppo_loss to avoid torch.compile interference. - _kl_stats = None - if compute_kl_stats: - with torch.no_grad(): - _raw_log_ratio = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0) - _log_ratio_full = torch.where( - current_support, - _raw_log_ratio, - torch.full_like(_raw_log_ratio, -20.0), - ) - _ratio_full = torch.where(current_support, torch.exp(_raw_log_ratio), torch.zeros_like(_raw_log_ratio)) - _per_token_k3 = _ratio_full - _log_ratio_full - 1.0 - # ±inf identity on empty ranks lets cross-rank MIN/MAX-allreduce ignore empty contributors. - if valid_mask.any(): - _ratio_min = _ratio_full.masked_fill(~valid_mask, float("inf")).min() - _ratio_max = _ratio_full.masked_fill(~valid_mask, float("-inf")).max() - _k3_max = _per_token_k3.masked_fill(~valid_mask, float("-inf")).max() - _logratio_min = _log_ratio_full.masked_fill(~valid_mask, float("inf")).min() - _logratio_max = _log_ratio_full.masked_fill(~valid_mask, float("-inf")).max() - _abs_logratio_max = _log_ratio_full.abs().masked_fill(~valid_mask, float("-inf")).max() - else: - _ratio_min = _ratio_full.new_tensor(float("inf")) - _ratio_max = _ratio_full.new_tensor(float("-inf")) - _k3_max = _per_token_k3.new_tensor(float("-inf")) - _logratio_min = _log_ratio_full.new_tensor(float("inf")) - _logratio_max = _log_ratio_full.new_tensor(float("-inf")) - _abs_logratio_max = _log_ratio_full.new_tensor(float("-inf")) - _kl_stats = { - "kl_sample_train_k3": metric_reducer(_per_token_k3, valid_mask_f), - "kl_k3_debug_mean": metric_reducer(_per_token_k3, valid_mask_f), - "kl_k3_debug_max": _k3_max, - "kl_k3_debug_abs_logratio_mean": metric_reducer(_log_ratio_full.abs(), valid_mask_f), - "kl_k3_debug_abs_logratio_max": _abs_logratio_max, - "kl_k3_debug_logratio_mean": metric_reducer(_log_ratio_full, valid_mask_f), - "kl_k3_debug_logratio_min": _logratio_min, - "kl_k3_debug_logratio_max": _logratio_max, - "kl_k3_debug_frac_logratio_positive": metric_reducer((_log_ratio_full > 0).float(), valid_mask_f), - "entropy_sample": metric_reducer(-old_logprobs_flat, valid_mask_f), - "ratio_mean": metric_reducer(_ratio_full, valid_mask_f), - "ratio_min": _ratio_min, - "ratio_max": _ratio_max, - "current_support_fraction": metric_reducer(current_support.float(), valid_mask_f), - } - for suffix, threshold in K3_DEBUG_THRESHOLDS: - _kl_stats[f"kl_k3_debug_frac_gt_{suffix}"] = metric_reducer( - (_per_token_k3 > threshold).float(), valid_mask_f - ) - - # Compute PPO-style clipped loss (returns per-token losses, clip mask, and ratio) - pg_losses, is_clipped, ratio = compute_ppo_loss( - ppo_kl=ppo_kl, - advantages=advantages_masked, - eps_clip=eps_clip, - eps_clip_high=eps_clip_high, - eps_clip_c=eps_clip_c, - ) - - # IcePop hard masking (GLM-5, arXiv:2602.15763): - # Zero gradient for tokens where ratio is outside [1/β, β] - icepop_mask = None - if icepop_beta is not None: - if use_tis: - logger.warning( - "IcePop and TIS are both enabled. IcePop makes TIS redundant " - "when using inference logprobs as old_logprobs." - ) - ratio_d = ratio.detach() - icepop_mask = (ratio_d >= 1.0 / icepop_beta) & (ratio_d <= icepop_beta) - - # Apply TIS correction if enabled and rollout_logprobs provided - tis_metrics = {} - if use_tis and rollout_logprobs is not None: - rollout_logprobs_flat = rollout_logprobs.view(-1) - pg_losses, tis_metrics = apply_tis_correction( - pg_loss=pg_losses, - train_log_probs=new_logprobs_flat, - rollout_log_probs=rollout_logprobs_flat, - valid_mask=valid_mask, - metric_reducer=metric_reducer, - tis_clip_low=tis_clip_low, - tis_clip_high=tis_clip_high, - ) - - # True loss value (for logging): partial share under loss_reducer. - true_loss = loss_reducer(pg_losses, valid_mask_f) - - # Gradient-active mask: tokens that are not clipped, not IcePop-masked, and valid - gradient_active = ~is_clipped & valid_mask & current_support - if icepop_mask is not None: - gradient_active = gradient_active & icepop_mask - - # Surrogate: gradient weight = ratio * advantages, zeroed for inactive tokens - gradient_weight = (ratio.detach() * advantages_flat).masked_fill(~gradient_active, 0.0) - safe_per_token_ce = torch.where(current_support, per_token_ce, torch.zeros_like(per_token_ce)) - surrogate = loss_reducer(gradient_weight * safe_per_token_ce, valid_mask_f) - - # Combine: forward value from true_loss, gradient from surrogate - loss_with_grad = true_loss.detach() + surrogate - surrogate.detach() - - # Return training logprobs reshaped - new_logprobs = new_logprobs_flat.view(original_shape) - - with torch.no_grad(): - metrics: Dict[str, Any] = { - "valid_tokens": valid_count.item(), - "pg_clipfrac": metric_reducer(is_clipped.float(), valid_mask_f), - } - - if icepop_mask is not None: - metrics["icepop_maskfrac"] = metric_reducer((~icepop_mask).float(), valid_mask_f) - - if _kl_stats is not None: - metrics.update(_kl_stats) - - metrics.update(tis_metrics) - - metric_ops: Dict[str, str] = {} - if _kl_stats is not None: - metric_ops["ratio_min"] = "min" - metric_ops["ratio_max"] = "max" - metric_ops["kl_k3_debug_max"] = "max" - metric_ops["kl_k3_debug_abs_logratio_max"] = "max" - metric_ops["kl_k3_debug_logratio_min"] = "min" - metric_ops["kl_k3_debug_logratio_max"] = "max" - if tis_metrics: - metric_ops["tis_min"] = "min" - metric_ops["tis_max"] = "max" - - return LossOutput( - loss=loss_with_grad, - per_token_logprobs=new_logprobs, - metrics=metrics, - metric_ops=metric_ops or None, - ) +from xorl.objectives.policy_loss import * # noqa: F401,F403 diff --git a/src/xorl/ops/loss/reducers.py b/src/xorl/ops/loss/reducers.py index 48b414c8..1f924d39 100644 --- a/src/xorl/ops/loss/reducers.py +++ b/src/xorl/ops/loss/reducers.py @@ -1,71 +1,3 @@ -"""Reducer protocol and canonical denominator policies for loss aggregation. +"""Moved to ``xorl.objectives.reducers`` (issue #78 phase 2); compat stub.""" -A ``Reducer`` collapses a ``(B, S)`` tensor to a scalar over a -caller-supplied denominator policy. Partial shares sum across micro-batches -and ``all_reduce(SUM)`` across ranks to the globally-correct value. - -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Protocol, runtime_checkable - -import torch - - -@runtime_checkable -class Reducer(Protocol): - """``(values, mask) -> scalar`` partial share over a pre-computed - denominator. Partial shares sum across micro-batches and ``all_reduce(SUM)`` - across ranks. - """ - - def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: ... - - -@dataclass(frozen=True) -class TokenPartial: - """Flat masked sum divided by a caller-supplied ``scale``.""" - - scale: torch.Tensor - - def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: - return (values * mask).sum() / self.scale.clamp(min=1.0) - - -@dataclass(frozen=True) -class SequencePartial: - """Sum of per-segment token-means, divided by a caller-supplied ``scale``. - - Segment boundaries are flat across ``(values * mask).reshape(-1)``: - - - ``cu_seqlens_local: (N+1,)`` — shard-local segment extents. Under CP each - rank's slice sums to its segment's local contribution. - - ``seq_lengths_global: (N,)`` — pre-CP-shard token count per segment, used - as the per-segment denominator so partial shares from each CP rank sum - to the correct per-segment mean. - """ - - scale: torch.Tensor - cu_seqlens_local: torch.Tensor - seq_lengths_global: torch.Tensor - - def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: - flat = (values * mask).reshape(-1) - seg_lengths_local = self.cu_seqlens_local.diff() - n_segments = seg_lengths_local.numel() - seg_ids = torch.repeat_interleave( - torch.arange(n_segments, device=flat.device), - seg_lengths_local, - ) - seg_sums = torch.zeros(n_segments, dtype=flat.dtype, device=flat.device).index_add(0, seg_ids, flat) - seg_means = seg_sums / self.seq_lengths_global.clamp(min=1.0) - return seg_means.sum() / self.scale.clamp(min=1.0) - - -__all__ = [ - "Reducer", - "SequencePartial", - "TokenPartial", -] +from xorl.objectives.reducers import * # noqa: F401,F403 diff --git a/src/xorl/trainers/training_utils.py b/src/xorl/trainers/training_utils.py index 41966006..1eb9adea 100644 --- a/src/xorl/trainers/training_utils.py +++ b/src/xorl/trainers/training_utils.py @@ -442,8 +442,8 @@ def _pp_quack_linear_ce_sum(hidden, labels, *, lm_head, num_chunks: int = 8): hidden (pipeline) and lm_head.weight (its FSDP unit reduce-scatters them). """ from xorl.models.module_utils import get_lm_head_weight # noqa: PLC0415 + from xorl.objectives.causallm_loss import _chunk_size_from_num_chunks # noqa: PLC0415 from xorl.ops._vendored.quack.linear_cross_entropy import chunked_linear_cross_entropy # noqa: PLC0415 - from xorl.ops.loss.causallm_loss import _chunk_size_from_num_chunks # noqa: PLC0415 weight = get_lm_head_weight(lm_head, fsdp_sharded_loss=False) h = hidden.reshape(-1, hidden.shape[-1]) diff --git a/tests/distributed/test_bi_fused_lm_head_tp.py b/tests/distributed/test_bi_fused_lm_head_tp.py index 203f522a..7ae0e239 100644 --- a/tests/distributed/test_bi_fused_lm_head_tp.py +++ b/tests/distributed/test_bi_fused_lm_head_tp.py @@ -12,13 +12,13 @@ import torch.distributed as dist import xorl.distributed.parallel_state as parallel_state_impl +from xorl.objectives.causallm_loss import causallm_loss_function +from xorl.objectives.reducers import TokenPartial from xorl.ops import bi_families_v2 from xorl.ops.loss.bi_fused_lm_head import ( bi_fused_per_token_ce, bi_fused_vocab_parallel_per_token_ce, ) -from xorl.ops.loss.causallm_loss import causallm_loss_function -from xorl.ops.loss.reducers import TokenPartial _HIDDEN = 128 diff --git a/tests/distributed/test_lm_head_tp_fsdp_e2e.py b/tests/distributed/test_lm_head_tp_fsdp_e2e.py index b2026988..e115a88c 100644 --- a/tests/distributed/test_lm_head_tp_fsdp_e2e.py +++ b/tests/distributed/test_lm_head_tp_fsdp_e2e.py @@ -23,9 +23,9 @@ import torch.distributed as dist # noqa: E402 from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state # noqa: E402 -from xorl.ops.loss.causallm_loss import fsdp_sharded_causallm_loss_function # noqa: E402 -from xorl.ops.loss.opd_loss import opd_vocab_parallel_loss_function # noqa: E402 -from xorl.ops.loss.reducers import TokenPartial # noqa: E402 +from xorl.objectives.causallm_loss import fsdp_sharded_causallm_loss_function # noqa: E402 +from xorl.objectives.opd_loss import opd_vocab_parallel_loss_function # noqa: E402 +from xorl.objectives.reducers import TokenPartial # noqa: E402 from xorl.server.runner.model_runner import ModelRunner # noqa: E402 from xorl.trainers.training_utils import sync_lm_head_tp_gradient, sync_lm_head_tp_parameters # noqa: E402 diff --git a/tests/distributed/test_pp_byte_alignment.py b/tests/distributed/test_pp_byte_alignment.py index 04bc1853..e448d769 100644 --- a/tests/distributed/test_pp_byte_alignment.py +++ b/tests/distributed/test_pp_byte_alignment.py @@ -47,7 +47,7 @@ from xorl.models.layers.normalization import set_rmsnorm_mode from xorl.models.transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config from xorl.models.transformers.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM -from xorl.ops.loss.causallm_loss import causallm_loss_function +from xorl.objectives.causallm_loss import causallm_loss_function from xorl.trainers.training_utils import forward_only_pp, pad_micro_batches_for_pp from xorl.utils.device import get_nccl_backend diff --git a/tests/distributed/test_ulysses_byte_alignment.py b/tests/distributed/test_ulysses_byte_alignment.py index b29efb2a..7848bddd 100644 --- a/tests/distributed/test_ulysses_byte_alignment.py +++ b/tests/distributed/test_ulysses_byte_alignment.py @@ -43,7 +43,7 @@ from xorl.models.layers.normalization import set_rmsnorm_mode from xorl.models.transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config from xorl.models.transformers.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM -from xorl.ops.loss.causallm_loss import causallm_loss_function +from xorl.objectives.causallm_loss import causallm_loss_function THIS_DIR = Path(__file__).resolve().parent diff --git a/tests/models/test_bi_fused_lm_head_tp_wiring.py b/tests/models/test_bi_fused_lm_head_tp_wiring.py index f700d3c8..4e618422 100644 --- a/tests/models/test_bi_fused_lm_head_tp_wiring.py +++ b/tests/models/test_bi_fused_lm_head_tp_wiring.py @@ -5,11 +5,11 @@ from torch import nn import xorl.models.module_utils as module_utils -import xorl.ops.loss.causallm_loss as causallm_loss_impl +import xorl.objectives.causallm_loss as causallm_loss_impl import xorl.ops.loss.per_token_ce as per_token_ce_impl import xorl.trainers.training_utils as training_utils -from xorl.ops.loss.loss_output import LossOutput -from xorl.ops.loss.reducers import TokenPartial +from xorl.objectives.loss_output import LossOutput +from xorl.objectives.reducers import TokenPartial def test_causallm_routes_bi_fused_tp_before_ordinary_vocab_ce(monkeypatch): diff --git a/tests/models/test_glm52_exact_lm_head_loss_integration.py b/tests/models/test_glm52_exact_lm_head_loss_integration.py index 3df4f334..58be2420 100644 --- a/tests/models/test_glm52_exact_lm_head_loss_integration.py +++ b/tests/models/test_glm52_exact_lm_head_loss_integration.py @@ -125,7 +125,7 @@ def _fake_exact(_hidden, _weight, _labels, **kwargs): def test_causallm_exact_head_admits_its_tp_group_and_rejects_z_loss(monkeypatch: pytest.MonkeyPatch) -> None: - causallm_impl = importlib.import_module("xorl.ops.loss.causallm_loss") + causallm_impl = importlib.import_module("xorl.objectives.causallm_loss") lm_head = nn.Module() lm_head._glm52_exact_tp16_lm_head = True hidden = torch.arange(8, dtype=torch.float32).reshape(1, 2, 4).to(torch.bfloat16).requires_grad_(True) diff --git a/tests/ops/loss/test_causallm_z_loss.py b/tests/ops/loss/test_causallm_z_loss.py index 29dbd388..360e7ca8 100644 --- a/tests/ops/loss/test_causallm_z_loss.py +++ b/tests/ops/loss/test_causallm_z_loss.py @@ -4,7 +4,7 @@ import torch from tests.ops.loss.conftest import assert_close -from xorl.ops.loss.causallm_loss import causallm_loss_function +from xorl.objectives.causallm_loss import causallm_loss_function def _reference_z_loss(hidden_states, weight, labels, ignore_index=-100): diff --git a/tests/ops/loss/test_fp8_lm_head_ce.py b/tests/ops/loss/test_fp8_lm_head_ce.py index 0433d4b2..d246a760 100644 --- a/tests/ops/loss/test_fp8_lm_head_ce.py +++ b/tests/ops/loss/test_fp8_lm_head_ce.py @@ -4,8 +4,8 @@ import torch.nn as nn import torch.nn.functional as F -from xorl.ops.loss.causallm_loss import causallm_loss_function -from xorl.ops.loss.importance_sampling_loss import importance_sampling_loss_function +from xorl.objectives.causallm_loss import causallm_loss_function +from xorl.objectives.importance_sampling_loss import importance_sampling_loss_function from xorl.ops.loss.per_token_ce import compute_per_token_ce diff --git a/tests/ops/loss/test_fused_linear_logprob.py b/tests/ops/loss/test_fused_linear_logprob.py index 6de7e8d9..1132da13 100644 --- a/tests/ops/loss/test_fused_linear_logprob.py +++ b/tests/ops/loss/test_fused_linear_logprob.py @@ -14,8 +14,8 @@ import torch import torch.nn.functional as F +from xorl.objectives.importance_sampling_loss import importance_sampling_loss_function from xorl.ops.loss.fused_linear_logprob import fused_selected_logprob_ce -from xorl.ops.loss.importance_sampling_loss import importance_sampling_loss_function from xorl.ops.loss.per_token_ce import compute_per_token_ce @@ -115,7 +115,7 @@ def _assert_irregular_tail_shape_matches_eager(): def _assert_loss_dispatchers_match_eager(): - from xorl.ops.loss.causallm_loss import causallm_loss_function # noqa: PLC0415 + from xorl.objectives.causallm_loss import causallm_loss_function # noqa: PLC0415 h, w, _, labels = _make_inputs(96, 192, 800, torch.bfloat16, has_bias=False) fused = compute_per_token_ce(h, w, labels, ignore_index=-100, ce_mode="fused_quack") @@ -147,7 +147,7 @@ def _assert_production_vocab_paths_are_finite_and_match_eager(): covered V <= 65536-class shapes where every cluster block owns columns.""" # The Qwen integration case and the largest direct GPT-OSS case cover the # two production boundaries that the small-vocabulary dispatcher case misses. - from xorl.ops.loss.causallm_loss import causallm_loss_function # noqa: PLC0415 + from xorl.objectives.causallm_loss import causallm_loss_function # noqa: PLC0415 N, H, V = 512, 1024, 151936 h, w, _, labels = _make_inputs(N, H, V, torch.bfloat16, has_bias=False) @@ -180,7 +180,7 @@ def test_causallm_fused_quack_does_not_materialize_full_logits(): the full logits tile. This is the assertion that actually distinguishes the fixed code from the broken fall-through (a loss-match test passes either way, since the old fall-through == eager).""" - from xorl.ops.loss.causallm_loss import causallm_loss_function # noqa: PLC0415 + from xorl.objectives.causallm_loss import causallm_loss_function # noqa: PLC0415 N, H, V = 16384, 2048, 50000 h = torch.randn(1, N, H, device="cuda", dtype=torch.bfloat16, requires_grad=True) diff --git a/tests/ops/loss/test_opd_loss.py b/tests/ops/loss/test_opd_loss.py index 5beb2611..2d26a366 100644 --- a/tests/ops/loss/test_opd_loss.py +++ b/tests/ops/loss/test_opd_loss.py @@ -5,8 +5,8 @@ from tests._helpers.opd import reference_opd_loss from tests.ops.loss.conftest import assert_close from xorl.distillation.teacher_store import TeacherHeadShardView, TeacherHeadStore, prepare_lm_head_teacher_store +from xorl.objectives.opd_loss import _oprd_hidden_distance, _oprd_hidden_distance_from_fetcher from xorl.ops.loss import TokenPartial, opd_loss_function -from xorl.ops.loss.opd_loss import _oprd_hidden_distance, _oprd_hidden_distance_from_fetcher pytestmark = pytest.mark.cpu diff --git a/tests/ops/loss/test_opd_verl_parity.py b/tests/ops/loss/test_opd_verl_parity.py index 5c1136bc..9569875d 100644 --- a/tests/ops/loss/test_opd_verl_parity.py +++ b/tests/ops/loss/test_opd_verl_parity.py @@ -15,14 +15,14 @@ import torch import torch.nn.functional as F -from xorl.ops.loss.compiled_cross_entropy import ( - compiled_sampled_token_logprobs_function, -) -from xorl.ops.loss.opd_loss import ( +from xorl.objectives.opd_loss import ( LOSS_MODE_FORWARD_KL_FULL, _kl_penalty_estimator, opd_loss_function, ) +from xorl.ops.loss.compiled_cross_entropy import ( + compiled_sampled_token_logprobs_function, +) pytestmark = [pytest.mark.cpu] diff --git a/tests/ops/loss/test_streaming_forward_kl.py b/tests/ops/loss/test_streaming_forward_kl.py index ad89a627..5feb5ed1 100644 --- a/tests/ops/loss/test_streaming_forward_kl.py +++ b/tests/ops/loss/test_streaming_forward_kl.py @@ -13,11 +13,11 @@ import torch.nn.functional as F from tests.ops.loss.conftest import assert_close -from xorl.ops.loss import opd_loss_function -from xorl.ops.loss.opd_streaming_kl import ( +from xorl.objectives.opd_streaming_kl import ( streaming_forward_kl_function, streaming_forward_kl_lowmem_function, ) +from xorl.ops.loss import opd_loss_function pytestmark = pytest.mark.cpu diff --git a/tests/ops/test_bi_fused_lm_head.py b/tests/ops/test_bi_fused_lm_head.py index c0d27de5..a4af5844 100644 --- a/tests/ops/test_bi_fused_lm_head.py +++ b/tests/ops/test_bi_fused_lm_head.py @@ -1,8 +1,8 @@ import pytest import torch +from xorl.objectives.causallm_loss import causallm_loss_function from xorl.ops.loss.bi_fused_lm_head import bi_fused_per_token_ce -from xorl.ops.loss.causallm_loss import causallm_loss_function requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/ops/test_exact_sampling_transforms.py b/tests/ops/test_exact_sampling_transforms.py index 05dfe4bf..ae4fd10d 100644 --- a/tests/ops/test_exact_sampling_transforms.py +++ b/tests/ops/test_exact_sampling_transforms.py @@ -224,13 +224,13 @@ def test_identity_row_metadata_collapses_to_no_filter_switch(): @pytest.mark.parametrize( "module_name,function_name,extra", [ - ("xorl.ops.loss.policy_loss", "policy_loss_function", {"compute_kl_stats": True}), + ("xorl.objectives.policy_loss", "policy_loss_function", {"compute_kl_stats": True}), ( - "xorl.ops.loss.importance_sampling_loss", + "xorl.objectives.importance_sampling_loss", "importance_sampling_loss_function", {"compute_kl_stats": True}, ), - ("xorl.ops.loss.cispo_loss", "cispo_loss_function", {"compute_kl_stats": True}), + ("xorl.objectives.cispo_loss", "cispo_loss_function", {"compute_kl_stats": True}), ], ) def test_rl_surrogates_are_finite_and_zero_gradient_outside_current_support( @@ -261,7 +261,7 @@ def test_rl_surrogates_are_finite_and_zero_gradient_outside_current_support( @pytest.mark.parametrize("ratio_type", ["token", "sequence"]) def test_drgrpo_is_finite_and_zero_gradient_outside_current_support(monkeypatch, ratio_type): - module = importlib.import_module("xorl.ops.loss.grpo_loss") + module = importlib.import_module("xorl.objectives.grpo_loss") ce = torch.tensor([math.inf, 0.4], dtype=torch.float32, requires_grad=True) monkeypatch.setattr(module, "compute_per_token_ce", lambda *args, **kwargs: ce) output = module.drgrpo_loss_function( diff --git a/tests/server/runner/test_model_runner_drgrpo.py b/tests/server/runner/test_model_runner_drgrpo.py index 63bbd044..389b8f54 100644 --- a/tests/server/runner/test_model_runner_drgrpo.py +++ b/tests/server/runner/test_model_runner_drgrpo.py @@ -4,8 +4,8 @@ import torch import torch.nn.functional as F +import xorl.objectives.causallm_loss as causallm_loss_impl import xorl.ops.loss.bi_fused_lm_head as bi_fused_lm_head_impl -import xorl.ops.loss.causallm_loss as causallm_loss_impl import xorl.ops.loss.per_token_ce as per_token_ce_impl import xorl.server.runner.model_runner as model_runner_module from xorl.server.runner.model_runner import ModelRunner diff --git a/tests/server/runner/test_opd_runner.py b/tests/server/runner/test_opd_runner.py index 035fb824..cf59daea 100644 --- a/tests/server/runner/test_opd_runner.py +++ b/tests/server/runner/test_opd_runner.py @@ -9,7 +9,7 @@ from tests._helpers.opd import make_teacher_files from xorl.data.constants import IGNORE_INDEX from xorl.distillation import MooncakeHiddenStore, TeacherActivationCache -from xorl.ops.loss.opd_loss import OPDLossMetrics +from xorl.objectives.opd_loss import OPDLossMetrics from xorl.server.runner.model_runner import ModelRunner From 1bfa317e4a7cde77ca7e04b64c4e80bfd31349f9 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 03:43:36 +0000 Subject: [PATCH 03/14] Gather the serving-parity contract family under ops/exact (#78 phase 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Physically moves the seven first-party contract modules (rope_class_b, canonical_moe_leaf, canonical_moe_cast, kernel_config_pin, bi_gemm_configs, block_fp8_native, fused_silu_and_mul) into xorl.ops.exact, rewriting in-repo imports and leaving compat stubs at the old paths. Three members are aliased in place rather than moved, deviating deliberately from the original #78 sketch: - bi_families_v2.py is sha256-gated byte-identical against the serving copy; - batch_invariant_ops.py turned out to be vendored-adapted from SGLang's srt/batch_invariant_ops — splitting or moving it would destroy side-by-side diffability with the serving twin, so the planned 4-way split is dropped; - exact_sampling_transforms.py is under active rewrite in #74; the canonical direction flips to ops/exact/sampling_transforms.py after #74 lands. --- src/xorl/checkpoint/checkpointer.py | 2 +- src/xorl/distributed/canonical_moe.py | 4 +- src/xorl/models/layers/rope.py | 4 +- .../deepseek_v3/modeling_deepseek_v3.py | 2 +- .../deepseek_v4/native_payload.py | 6 +- .../glm4_moe/modeling_glm4_moe.py | 2 +- .../glm5/exact_absorbed_kv_b_qlora.py | 2 +- .../transformers/glm5/exact_dense_mlp.py | 2 +- .../glm5/exact_fullparam_admission.py | 2 +- .../glm5/exact_fullparam_experts.py | 2 +- .../transformers/glm5/exact_fullparam_fp8.py | 4 +- .../transformers/glm5/exact_gate_up_qlora.py | 2 +- .../models/transformers/glm5/exact_qlora.py | 2 +- .../glm5/exact_routed_experts_qlora.py | 2 +- .../glm5/exact_shared_expert_qlora.py | 4 +- src/xorl/models/transformers/glm5/indexer.py | 2 +- .../models/transformers/glm5/modeling_glm5.py | 4 +- .../models/transformers/glm5/native_fp8.py | 4 +- src/xorl/models/transformers/glm5/qlora.py | 2 +- .../transformers/llama3/modeling_llama3.py | 2 +- .../transformers/olmo2/modeling_olmo2.py | 2 +- .../transformers/qwen2/modeling_qwen2.py | 2 +- .../transformers/qwen3/modeling_qwen3.py | 2 +- .../transformers/qwen3_5/modeling_qwen3_5.py | 2 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 2 +- .../qwen3_moe/modeling_qwen3_moe.py | 2 +- src/xorl/ops/README.md | 15 +- src/xorl/ops/batch_invariant_ops.py | 2 +- src/xorl/ops/bi_gemm_configs.py | 148 +----- src/xorl/ops/block_fp8_native.py | 489 +----------------- src/xorl/ops/canonical_moe_cast.py | 124 +---- src/xorl/ops/canonical_moe_leaf.py | 164 +----- src/xorl/ops/exact/__init__.py | 21 + src/xorl/ops/exact/batch_invariant.py | 3 + src/xorl/ops/exact/bi_gemm_configs.py | 147 ++++++ src/xorl/ops/exact/block_fp8_native.py | 488 +++++++++++++++++ src/xorl/ops/exact/canonical_moe_cast.py | 123 +++++ src/xorl/ops/exact/canonical_moe_leaf.py | 163 ++++++ src/xorl/ops/exact/families_v2.py | 3 + src/xorl/ops/exact/fused_silu_and_mul.py | 359 +++++++++++++ src/xorl/ops/exact/kernel_config_pin.py | 182 +++++++ src/xorl/ops/exact/rope_class_b.py | 161 ++++++ src/xorl/ops/exact/sampling_transforms.py | 3 + src/xorl/ops/fused_silu_and_mul.py | 360 +------------ src/xorl/ops/kernel_config_pin.py | 183 +------ src/xorl/ops/rope_class_b.py | 162 +----- .../test_glm52_fullparam_construction_gate.py | 4 +- .../test_glm52_fullparam_ep16_combine.py | 2 +- .../test_native_fp8_fsdp2_materialization.py | 2 +- .../test_torch_parallelize_policies.py | 2 +- .../test_glm52_exact_absorbed_kv_b_qlora.py | 2 +- .../test_glm52_exact_dense_mlp_composition.py | 2 +- .../models/test_glm52_exact_fullparam_fp8.py | 4 +- .../models/test_glm52_exact_gate_up_qlora.py | 4 +- tests/models/test_glm52_exact_qlora.py | 2 +- .../test_glm52_exact_shared_expert_qlora.py | 2 +- .../models/test_glm52_fullparam_admission.py | 2 +- tests/models/test_glm52_native_fp8.py | 2 +- tests/models/test_glm52_qlora.py | 2 +- tests/models/test_glm5_one_round_swiglu.py | 2 +- tests/models/test_module_utils_broadcast.py | 2 +- tests/models/test_op_parity_dense.py | 2 +- .../test_qwen35_hybrid_ulysses_admission.py | 4 +- tests/ops/test_bi_gemm_config_table.py | 2 +- tests/ops/test_block_fp8_frozen_dgrad.py | 4 +- tests/ops/test_block_fp8_native.py | 2 +- tests/ops/test_canonical_moe_cast.py | 2 +- tests/ops/test_fused_silu_and_mul_scoping.py | 2 +- tests/ops/test_rope_class_b_contract.py | 2 +- .../test_glm52_fullparam_step_publish.py | 2 +- 70 files changed, 1738 insertions(+), 1688 deletions(-) create mode 100644 src/xorl/ops/exact/__init__.py create mode 100644 src/xorl/ops/exact/batch_invariant.py create mode 100644 src/xorl/ops/exact/bi_gemm_configs.py create mode 100644 src/xorl/ops/exact/block_fp8_native.py create mode 100644 src/xorl/ops/exact/canonical_moe_cast.py create mode 100644 src/xorl/ops/exact/canonical_moe_leaf.py create mode 100644 src/xorl/ops/exact/families_v2.py create mode 100644 src/xorl/ops/exact/fused_silu_and_mul.py create mode 100644 src/xorl/ops/exact/kernel_config_pin.py create mode 100644 src/xorl/ops/exact/rope_class_b.py create mode 100644 src/xorl/ops/exact/sampling_transforms.py diff --git a/src/xorl/checkpoint/checkpointer.py b/src/xorl/checkpoint/checkpointer.py index 048aa6ff..7a20ab69 100644 --- a/src/xorl/checkpoint/checkpointer.py +++ b/src/xorl/checkpoint/checkpointer.py @@ -1040,7 +1040,7 @@ def load( # the global expert dimension and therefore matches DCP metadata. expected_model_state = load_state["model"].state_dict() if any("packed_weight_f32" in name or name.endswith("weight_scale_inv") for name in expected_model_state): - from xorl.ops.block_fp8_native import validate_native_fp8_dcp_checkpoint # noqa: PLC0415 + from xorl.ops.exact.block_fp8_native import validate_native_fp8_dcp_checkpoint # noqa: PLC0415 validate_native_fp8_dcp_checkpoint( checkpoint_dir, diff --git a/src/xorl/distributed/canonical_moe.py b/src/xorl/distributed/canonical_moe.py index d730ba76..63eaa91c 100644 --- a/src/xorl/distributed/canonical_moe.py +++ b/src/xorl/distributed/canonical_moe.py @@ -26,7 +26,7 @@ import torch import torch.distributed as dist -from xorl.ops.canonical_moe_leaf import canonical_moe_leaf_fp32_v1_op +from xorl.ops.exact.canonical_moe_leaf import canonical_moe_leaf_fp32_v1_op CANONICAL_MOE_FOLD_VERSION = "canonical_moe_fold_fp64_v3" @@ -677,7 +677,7 @@ def _canonical_moe_cast_fp64_to_transport( if folded_fp64.dtype is not torch.float64: raise TypeError(f"Canonical MoE final cast requires FP64 input, got {folded_fp64.dtype}") if folded_fp64.is_cuda and transport_dtype in (torch.bfloat16, torch.float16): - from xorl.ops.canonical_moe_cast import canonical_moe_fp64_to_lowp_rne # noqa: PLC0415 + from xorl.ops.exact.canonical_moe_cast import canonical_moe_fp64_to_lowp_rne # noqa: PLC0415 return canonical_moe_fp64_to_lowp_rne(folded_fp64.contiguous(), transport_dtype) return folded_fp64.to(transport_dtype) diff --git a/src/xorl/models/layers/rope.py b/src/xorl/models/layers/rope.py index fbb79a36..b8e83be8 100644 --- a/src/xorl/models/layers/rope.py +++ b/src/xorl/models/layers/rope.py @@ -637,8 +637,8 @@ def rope_class_b_enabled() -> bool: def stock_fused_apply_rotary_pos_emb(q, k, cos, sin, *, interleaved: bool = False, doubled: bool = True): - """Class-B RoPE application backed by the compiled expression in ``xorl.ops.rope_class_b``.""" - from xorl.ops.rope_class_b import class_b_apply_rotary_pos_emb # noqa: PLC0415 + """Class-B RoPE application backed by the compiled expression in ``xorl.ops.exact.rope_class_b``.""" + from xorl.ops.exact.rope_class_b import class_b_apply_rotary_pos_emb # noqa: PLC0415 return class_b_apply_rotary_pos_emb(q, k, cos, sin, interleaved=interleaved, doubled=doubled) diff --git a/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py b/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py index 3c7f0394..0233373d 100644 --- a/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py +++ b/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py @@ -24,7 +24,7 @@ has_packed_expert_weights, ) from xorl.models.transformers.qwen3_5_shared import qwen3_5_apply_rotary_pos_emb -from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul from xorl.utils import logging diff --git a/src/xorl/models/transformers/deepseek_v4/native_payload.py b/src/xorl/models/transformers/deepseek_v4/native_payload.py index 0c93edac..c74bdeb7 100644 --- a/src/xorl/models/transformers/deepseek_v4/native_payload.py +++ b/src/xorl/models/transformers/deepseek_v4/native_payload.py @@ -214,7 +214,7 @@ def _dequantize_native_block_fp8(payload: Dsv4NativeBlockFp8Payload) -> torch.Te def _native_block_fp8_value(input: torch.Tensor, payload: Dsv4NativeBlockFp8Payload) -> torch.Tensor: - from xorl.ops.block_fp8_native import _sglang_native_block_fp8_linear_value # noqa: PLC0415 + from xorl.ops.exact.block_fp8_native import _sglang_native_block_fp8_linear_value # noqa: PLC0415 rows = input.numel() // payload.in_features weight, scales = payload() @@ -1142,7 +1142,7 @@ def _native_fp8_slice_value( input_start: int = 0, input_end: int | None = None, ) -> torch.Tensor: - from xorl.ops.block_fp8_native import ( # noqa: PLC0415 + from xorl.ops.exact.block_fp8_native import ( # noqa: PLC0415 _sglang_native_block_fp8_linear_value, ) @@ -1245,7 +1245,7 @@ def _dsv4_native_shared_tp_forward( raise ValueError("DSV4 shared intermediate size must divide the TP width") width = intermediate // tp_size start, end = tp_rank * width, (tp_rank + 1) * width - from xorl.ops.block_fp8_native import ( # noqa: PLC0415 + from xorl.ops.exact.block_fp8_native import ( # noqa: PLC0415 _sglang_native_block_fp8_linear_value, ) diff --git a/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py b/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py index b3354c20..4bbf0f72 100644 --- a/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py +++ b/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py @@ -44,7 +44,7 @@ from xorl.models.transformers.glm4_moe import parallelize from xorl.models.transformers.glm4_moe.checkpoint_handler import Glm4MoeCheckpointHandler from xorl.models.transformers.glm4_moe.configuration_glm4_moe import Glm4MoeConfig -from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul from xorl.utils import logging diff --git a/src/xorl/models/transformers/glm5/exact_absorbed_kv_b_qlora.py b/src/xorl/models/transformers/glm5/exact_absorbed_kv_b_qlora.py index 268a686a..c7c54d41 100644 --- a/src/xorl/models/transformers/glm5/exact_absorbed_kv_b_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_absorbed_kv_b_qlora.py @@ -20,7 +20,7 @@ from torch import Tensor, nn from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling -from xorl.ops.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear GLM52_EXACT_TP1_ABSORBED_KV_B_QLORA_CONTRACT_VERSION = "glm52_exact_tp1_absorbed_kv_b_qlora_v2" diff --git a/src/xorl/models/transformers/glm5/exact_dense_mlp.py b/src/xorl/models/transformers/glm5/exact_dense_mlp.py index b3336d29..bfa33685 100644 --- a/src/xorl/models/transformers/glm5/exact_dense_mlp.py +++ b/src/xorl/models/transformers/glm5/exact_dense_mlp.py @@ -16,7 +16,7 @@ ) from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul GLM52_EXACT_TP1_DENSE_MLP_CONTRACT_VERSION = "glm52_exact_tp1_dense_mlp_qlora_v2" diff --git a/src/xorl/models/transformers/glm5/exact_fullparam_admission.py b/src/xorl/models/transformers/glm5/exact_fullparam_admission.py index 41f3fc00..7f6cbc5a 100644 --- a/src/xorl/models/transformers/glm5/exact_fullparam_admission.py +++ b/src/xorl/models/transformers/glm5/exact_fullparam_admission.py @@ -42,7 +42,7 @@ Glm52NativeBlockFP8Experts, validate_glm52_native_fp8_config, ) -from xorl.ops.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear logger = logging.getLogger(__name__) diff --git a/src/xorl/models/transformers/glm5/exact_fullparam_experts.py b/src/xorl/models/transformers/glm5/exact_fullparam_experts.py index 7817e553..4854169a 100644 --- a/src/xorl/models/transformers/glm5/exact_fullparam_experts.py +++ b/src/xorl/models/transformers/glm5/exact_fullparam_experts.py @@ -30,7 +30,7 @@ quantize_expert_masters_to_serving_bytes, ) from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8Experts -from xorl.ops.block_fp8_native import pack_fp8_as_float32 +from xorl.ops.exact.block_fp8_native import pack_fp8_as_float32 logger = logging.getLogger(__name__) diff --git a/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py b/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py index fcd6e405..6cf1f694 100644 --- a/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py +++ b/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py @@ -25,12 +25,12 @@ import torch.nn.functional as F from torch import Tensor, nn -from xorl.ops.block_fp8_native import ( +from xorl.ops.exact.block_fp8_native import ( _sglang_native_block_fp8_linear_value, pack_fp8_as_float32, unpack_float32_as_fp8, ) -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul logger = logging.getLogger(__name__) diff --git a/src/xorl/models/transformers/glm5/exact_gate_up_qlora.py b/src/xorl/models/transformers/glm5/exact_gate_up_qlora.py index c24d19d1..323ded28 100644 --- a/src/xorl/models/transformers/glm5/exact_gate_up_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_gate_up_qlora.py @@ -17,7 +17,7 @@ from torch import Tensor, nn from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling -from xorl.ops.block_fp8_native import NativeBlockFP8Linear, _sglang_native_block_fp8_linear_value +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear, _sglang_native_block_fp8_linear_value GLM52_EXACT_TP1_GATE_UP_QLORA_CONTRACT_VERSION = "glm52_exact_tp1_fused_gate_up_qlora_v2" diff --git a/src/xorl/models/transformers/glm5/exact_qlora.py b/src/xorl/models/transformers/glm5/exact_qlora.py index f55c6bfd..47c61d4a 100644 --- a/src/xorl/models/transformers/glm5/exact_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_qlora.py @@ -17,7 +17,7 @@ from torch import Tensor, nn from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling -from xorl.ops.block_fp8_native import _sglang_native_block_fp8_linear_value +from xorl.ops.exact.block_fp8_native import _sglang_native_block_fp8_linear_value from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear diff --git a/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py b/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py index eebeb2de..91af0267 100644 --- a/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py @@ -33,7 +33,7 @@ from xorl.models.layers.moe.experts import MoEExperts from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8Experts -from xorl.ops.block_fp8_native import pack_fp8_as_float32 +from xorl.ops.exact.block_fp8_native import pack_fp8_as_float32 GLM52_EXACT_EP16_ROUTED_QLORA_CONTRACT_VERSION = "glm52_exact_ep16_routed_qlora_v2" diff --git a/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py b/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py index 0925ba83..882de8c2 100644 --- a/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py @@ -25,8 +25,8 @@ from torch import Tensor, nn from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling -from xorl.ops.block_fp8_native import NativeBlockFP8Linear, _sglang_native_block_fp8_linear_value -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear, _sglang_native_block_fp8_linear_value +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul GLM52_EXACT_TP16_SHARED_EXPERT_QLORA_CONTRACT_VERSION = "glm52_exact_tp16_shared_expert_qlora_v2" diff --git a/src/xorl/models/transformers/glm5/indexer.py b/src/xorl/models/transformers/glm5/indexer.py index 19effa47..7bb11464 100644 --- a/src/xorl/models/transformers/glm5/indexer.py +++ b/src/xorl/models/transformers/glm5/indexer.py @@ -26,7 +26,7 @@ select_glm52_logical_indices, ) from xorl.ops.batch_invariant_ops import bi_bf16_fp32_linear, matmul_persistent -from xorl.ops.rope_class_b import build_class_b_cos_sin +from xorl.ops.exact.rope_class_b import build_class_b_cos_sin GLM52_FUSED_INDEX_K_IMPORT = "sglang.kernels.ops.quantization.dsv32.elementwise.fused_k_indexer_norm_rope" diff --git a/src/xorl/models/transformers/glm5/modeling_glm5.py b/src/xorl/models/transformers/glm5/modeling_glm5.py index e1720171..7bb2e375 100644 --- a/src/xorl/models/transformers/glm5/modeling_glm5.py +++ b/src/xorl/models/transformers/glm5/modeling_glm5.py @@ -69,8 +69,8 @@ from xorl.models.transformers.glm5.rotary import glm5_apply_rotary_pos_emb from xorl.models.transformers.glm5.sparse_mla import sparse_mla_dispatch from xorl.models.transformers.glm5.support import validate_glm5_sequence_parallel -from xorl.ops.block_fp8_native import NativeBlockFP8Linear -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul from xorl.utils import logging diff --git a/src/xorl/models/transformers/glm5/native_fp8.py b/src/xorl/models/transformers/glm5/native_fp8.py index 81cc9842..128a7dbc 100644 --- a/src/xorl/models/transformers/glm5/native_fp8.py +++ b/src/xorl/models/transformers/glm5/native_fp8.py @@ -13,7 +13,7 @@ from torch import nn from xorl.models.layers.moe.experts import MoEExperts -from xorl.ops.block_fp8_native import ( +from xorl.ops.exact.block_fp8_native import ( NATIVE_BLOCK_FP8_CONTRACT_VERSION, NativeBlockFP8Linear, pack_fp8_as_float32, @@ -1003,7 +1003,7 @@ def checkpoint_split_bytes(self) -> tuple[torch.Tensor, ...]: ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul # noqa: PLC0415 + from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul # noqa: PLC0415 gate_up = self.gate_up_proj(hidden_states) activated = exact_fp32_silu_and_mul(gate_up) diff --git a/src/xorl/models/transformers/glm5/qlora.py b/src/xorl/models/transformers/glm5/qlora.py index 0786392a..bc44bf6f 100644 --- a/src/xorl/models/transformers/glm5/qlora.py +++ b/src/xorl/models/transformers/glm5/qlora.py @@ -33,7 +33,7 @@ replace_glm52_native_fp8_modules, validate_glm52_native_fp8_config, ) -from xorl.ops.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear from xorl.qlora.modules.moe_experts import BlockFP8QLoRAMoeExperts diff --git a/src/xorl/models/transformers/llama3/modeling_llama3.py b/src/xorl/models/transformers/llama3/modeling_llama3.py index d114325e..9951aaa8 100644 --- a/src/xorl/models/transformers/llama3/modeling_llama3.py +++ b/src/xorl/models/transformers/llama3/modeling_llama3.py @@ -24,7 +24,7 @@ from xorl.models.transformers.llama3 import parallelize from xorl.models.transformers.llama3.checkpoint_handler import Llama3CheckpointHandler from xorl.models.transformers.llama3.configuration_llama3 import Llama3Config -from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul from xorl.utils import logging diff --git a/src/xorl/models/transformers/olmo2/modeling_olmo2.py b/src/xorl/models/transformers/olmo2/modeling_olmo2.py index 74879ac7..89945be6 100644 --- a/src/xorl/models/transformers/olmo2/modeling_olmo2.py +++ b/src/xorl/models/transformers/olmo2/modeling_olmo2.py @@ -26,7 +26,7 @@ from xorl.models.transformers.olmo2 import parallelize from xorl.models.transformers.olmo2.checkpoint_handler import Olmo2CheckpointHandler from xorl.models.transformers.olmo2.configuration_olmo2 import Olmo2Config -from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul from xorl.utils import logging diff --git a/src/xorl/models/transformers/qwen2/modeling_qwen2.py b/src/xorl/models/transformers/qwen2/modeling_qwen2.py index 65b358e7..3562b3d4 100644 --- a/src/xorl/models/transformers/qwen2/modeling_qwen2.py +++ b/src/xorl/models/transformers/qwen2/modeling_qwen2.py @@ -31,7 +31,7 @@ from xorl.models.transformers.qwen2 import parallelize from xorl.models.transformers.qwen2.checkpoint_handler import Qwen2CheckpointHandler from xorl.models.transformers.qwen2.configuration_qwen2 import Qwen2Config -from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul _RUNTIME_CONFIG_ATTRS = { diff --git a/src/xorl/models/transformers/qwen3/modeling_qwen3.py b/src/xorl/models/transformers/qwen3/modeling_qwen3.py index 1841273e..87eee810 100644 --- a/src/xorl/models/transformers/qwen3/modeling_qwen3.py +++ b/src/xorl/models/transformers/qwen3/modeling_qwen3.py @@ -30,7 +30,7 @@ from xorl.models.transformers.qwen3 import parallelize from xorl.models.transformers.qwen3.checkpoint_handler import Qwen3CheckpointHandler from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config -from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul from xorl.utils import logging diff --git a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py index 2cf90883..b5aa20a6 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -42,7 +42,7 @@ has_linear_attention_layers, qwen3_5_apply_rotary_pos_emb, ) -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul from xorl.ops.linear_attention import GatedDeltaNet from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils import logging diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index 3645d3ab..bc4380bc 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -43,7 +43,7 @@ has_linear_attention_layers, qwen3_5_apply_rotary_pos_emb, ) -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul from xorl.ops.linear_attention import GatedDeltaNet from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils import logging diff --git a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index e228a2c7..5160fd42 100644 --- a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py @@ -47,7 +47,7 @@ from xorl.models.transformers.qwen3_moe import parallelize from xorl.models.transformers.qwen3_moe.checkpoint_handler import Qwen3MoeCheckpointHandler from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig -from xorl.ops.fused_silu_and_mul import fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul from xorl.utils import logging diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index ee37606e..79ace5a9 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -11,18 +11,17 @@ Three kinds of code live here with three different rules: | kind | rule | | --- | --- | | **Vendored** — `_vendored/` (`quack/`, `flashqla/`) | Never hand-edit, lint, or reformat. Each tree carries a `VENDORED.md` with provenance and the local-patch ledger. First-party tooling skips them (`[tool.ruff]` excludes in `pyproject.toml`; top-level `exclude:` in `.pre-commit-config.yaml`). | -| **Byte-contract-gated** — `bi_families_v2.py` | Vendored byte-identical into the serving engine; both copies are sha256-gated. Any edit here without the paired serving-side edit breaks the gate. It keeps the engine's formatting (black, 88 columns) and is excluded from all rewriting hooks. | +| **Byte-contract-gated** — `bi_families_v2.py` (sha256-gated), `batch_invariant_ops.py` (parity-diffable twin of SGLang's copy; edits must consider the serving side) | Vendored byte-identical into the serving engine; both copies are sha256-gated. Any edit here without the paired serving-side edit breaks the gate. It keeps the engine's formatting (black, 88 columns) and is excluded from all rewriting hooks. | | **First-party** — everything else | Normal rules. | ## Map (current) -- `bi_families_v2.py`, `batch_invariant_ops.py`, `bi_gemm_configs.py`, - `exact_sampling_transforms.py`, `rope_class_b.py`, `canonical_moe_leaf.py`, - `canonical_moe_cast.py`, `fused_silu_and_mul.py`, `kernel_config_pin.py`, - `block_fp8_native.py` — the **serving-parity (exact) contract family**: - byte-pinned programs shared with the serving engine. The `bi_` / `exact_` / - `canonical_` / `class_b` prefixes are historical names for the same - concept. Planned home: `ops/exact/` (#78 phase 3). +- `exact/` — the **serving-parity (exact) contract family**: byte-pinned + programs shared with the serving engine (#78 phase 3). Three members are + aliased rather than moved: `bi_families_v2.py` (sha256-gated), + `batch_invariant_ops.py` (diffable parity twin of SGLang's copy), and + `exact_sampling_transforms.py` (in-flight in #74). Old root-level module + paths are compat stubs for one deprecation cycle. - `loss/` — the CE/selected-logprob kernel stack. The RL/supervised objective functions live in `xorl/objectives/` (#78 phase 2); old module paths here are compat stubs for one deprecation cycle. diff --git a/src/xorl/ops/batch_invariant_ops.py b/src/xorl/ops/batch_invariant_ops.py index 6a2c6286..79f4d56b 100644 --- a/src/xorl/ops/batch_invariant_ops.py +++ b/src/xorl/ops/batch_invariant_ops.py @@ -23,7 +23,7 @@ import triton.language as tl from triton.runtime.errors import OutOfResources -from xorl.ops.bi_gemm_configs import baseline_mm_config, lookup_mm_config +from xorl.ops.exact.bi_gemm_configs import baseline_mm_config, lookup_mm_config # --- Stubs for SGLang-internal imports --------------------------------------- diff --git a/src/xorl/ops/bi_gemm_configs.py b/src/xorl/ops/bi_gemm_configs.py index f366058f..bc72e1aa 100644 --- a/src/xorl/ops/bi_gemm_configs.py +++ b/src/xorl/ops/bi_gemm_configs.py @@ -1,147 +1,3 @@ -"""Shape-keyed tile configs for the batch-invariant persistent matmul (gap-map R1). +"""Moved to ``xorl.ops.exact.bi_gemm_configs`` (issue #78 phase 3); compat stub.""" -Only BLOCK_SIZE_K is bit-relevant in ``matmul_kernel_persistent``: it sets the -per-element K-reduction order. BLOCK_SIZE_M/N, GROUP_SIZE_M, num_stages and -num_warps are bit-neutral tuning axes (warps split M/N, never K; no split-K), -so they may vary per shape without moving the contract bits. Every entry below -was admitted only after torch.equal against the pinned baseline on multiple -seeds plus a cross-M row-invariance check (xorl -experiments/k3_tests/tools/tune_bi_gemm.py, H100, 2026-07-07). - -The shape-keyed table is the production configuration. The pinned baseline is -retained only as an internal launch-failure fallback; it is not a user-selectable -mode. -""" - -# The bit-relevant axis: pinned per dtype, NEVER shape-keyed. -PINNED_BLOCK_K = { - "torch.bfloat16": 64, - "torch.float16": 64, - "torch.float32": 32, -} - -# The pre-R1 one-config-per-dtype baselines (also the launch-failure fallback). -BASELINE_CONFIG = { - "torch.bfloat16": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}, - "torch.float16": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}, - "torch.float32": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}, -} - -# Class defaults on (dtype, M-bucket) for shapes without an exact (K, N) entry. -# Buckets: decode/graph shapes (M<=256) want tiny BLOCK_M; batch shapes want -# the wide 128x256 tile. Measured on delphi-class shapes; bit-neutral either way. -CLASS_DEFAULTS = { - "torch.bfloat16": ( - (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), - (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), - ), - "torch.float32": ( - (256, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 4, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), - ), -} - -# Exact-shape entries: {(dtype, K, N): ((max_M, cfg), ..., (None, cfg))}. -# Generated from offline tuner results. Do not hand-edit these values; rerun -# the tuning and generation pipeline instead. -# BEGIN GENERATED TABLE -TABLE = { - ("torch.bfloat16", 3840, 3840): ( - (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 8, "num_stages": 12, "num_warps": 4}), - (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "GROUP_SIZE_M": 1, "num_stages": 12, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), - ), - ("torch.bfloat16", 3840, 5376): ( - (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 1, "num_stages": 12, "num_warps": 4}), - (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 1, "num_stages": 8, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), - ), - ("torch.bfloat16", 3840, 8192): ( - (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 8, "num_stages": 8, "num_warps": 4}), - (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 1, "num_stages": 8, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), - ), - ("torch.bfloat16", 3840, 11520): ( - (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 5, "num_warps": 4}), - (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), - ), - ("torch.bfloat16", 3840, 30720): ( - (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), - (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), - ), - ("torch.bfloat16", 3840, 128256): ( - (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 5, "num_warps": 4}), - (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), - ), - ("torch.bfloat16", 15360, 3840): ( - (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 1, "num_stages": 20, "num_warps": 4}), - (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "GROUP_SIZE_M": 1, "num_stages": 12, "num_warps": 4}), - (8192, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 4, "num_warps": 8}), - (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 1, "num_stages": 4, "num_warps": 8}), - ), - ("torch.float32", 3840, 8192): ( - (8192, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 5, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), - ), -} -# END GENERATED TABLE - - -# Conservative H100 dynamic-smem budget. The persistent kernel needs -# num_stages * (BM*BK + BK*BN) * in_elt for the pipeline plus BM*BN*out_elt for -# the epilogue store staging when the output is wider than the inputs (the BI -# lm-head's fp32-out form) — a config that fits bf16-out can OOM fp32-out. -_SMEM_BUDGET = 220 * 1024 - -_ELT = {"torch.bfloat16": 2, "torch.float16": 2, "torch.float32": 4} - - -def _fits(cfg, block_k: int, in_elt: int, out_elt: int) -> bool: - pipe = cfg["num_stages"] * (cfg["BLOCK_SIZE_M"] * block_k + block_k * cfg["BLOCK_SIZE_N"]) * in_elt - # same-width outputs reuse pipeline smem; wider outputs stage a full C tile - epilogue = cfg["BLOCK_SIZE_M"] * cfg["BLOCK_SIZE_N"] * out_elt if out_elt > in_elt else 0 - return pipe + epilogue <= _SMEM_BUDGET - - -def lookup_mm_config(dtype, M: int, N: int, K: int, out_itemsize: int | None = None): - """Tile config for the persistent matmul at (dtype, M, N, K). - - Returns the full launch config including the pinned BLOCK_SIZE_K. Exact - (K, N) entries win, then (dtype, M-bucket) class defaults, then the pinned - baseline. All candidates are bit-identical; this only picks speed. Pass - ``out_itemsize`` when the output dtype is wider than the inputs so the - epilogue smem staging is budgeted (oversized configs fall back to baseline). - """ - key = str(dtype) - block_k = PINNED_BLOCK_K.get(key) - if block_k is None: - raise ValueError(f"Unsupported dtype {dtype} for batch-invariant matmul") - in_elt = _ELT[key] - out_elt = out_itemsize if out_itemsize is not None else in_elt - cfg = None - buckets = TABLE.get((key, K, N)) - exact = buckets is not None - if buckets is None: - buckets = CLASS_DEFAULTS.get(key) - if buckets is not None: - for max_m, bucket_cfg in buckets: - if max_m is None or M <= max_m: - # exact entries were measured/launch-validated at their out - # dtype; class defaults get the conservative smem model (the - # launch sites also fall back to baseline on OutOfResources) - if exact or _fits(bucket_cfg, block_k, in_elt, out_elt): - cfg = bucket_cfg - break - if cfg is None: - cfg = BASELINE_CONFIG[key] - return dict(cfg, BLOCK_SIZE_K=block_k) - - -def baseline_mm_config(dtype): - """The pinned pre-R1 config for dtype — the launch-failure fallback.""" - key = str(dtype) - return dict(BASELINE_CONFIG[key], BLOCK_SIZE_K=PINNED_BLOCK_K[key]) +from xorl.ops.exact.bi_gemm_configs import * # noqa: F401,F403 diff --git a/src/xorl/ops/block_fp8_native.py b/src/xorl/ops/block_fp8_native.py index f0c78158..d89b77b5 100644 --- a/src/xorl/ops/block_fp8_native.py +++ b/src/xorl/ops/block_fp8_native.py @@ -1,488 +1,3 @@ -"""Frozen native block-FP8 state and exact SGLang dense dispatch. +"""Moved to ``xorl.ops.exact.block_fp8_native`` (issue #78 phase 3); compat stub.""" -The packed parameter uses the same public storage convention as XoRL QLoRA: -four float8 bytes are viewed as one float32 element. This makes the payload a -normal reshardable parameter while retaining every checkpoint byte. It must -be wrapped without an FSDP mixed-precision policy; casting the packed float32 -values numerically would corrupt the embedded bytes. -""" - -from __future__ import annotations - -import logging - -import torch -from torch import nn - - -logger = logging.getLogger(__name__) - -NATIVE_BLOCK_FP8_CONTRACT_VERSION = "xorl_native_block_fp8_sglang_v1" -NATIVE_BLOCK_FP8_FROZEN_DGRAD_CONTRACT_VERSION = "xorl_native_block_fp8_frozen_dgrad_v1" -_FP8_DTYPE = torch.float8_e4m3fn - - -def _sglang_native_block_fp8_linear_value( - input: torch.Tensor, - weight: torch.Tensor, - weight_scale_inv: torch.Tensor, - *, - block_size: tuple[int, int] = (128, 128), -) -> torch.Tensor: - """Run the shared exact-value SGLang W8A8 dispatch. - - Gradient policy belongs to the caller: :class:`NativeBlockFP8Linear` - remains scoring-only, while the exact QLoRA wrapper supplies its validated - surrogate through a custom autograd boundary. - """ - - if input.device.type != "cuda": - raise RuntimeError("Native block-FP8 forward requires CUDA and the pinned SGLang Triton kernel") - if input.dtype is not torch.bfloat16: - raise TypeError(f"Native block-FP8 forward requires BF16 activations, got {input.dtype}") - if weight.dtype is not _FP8_DTYPE: - raise TypeError(f"Native block-FP8 weight must remain float8_e4m3fn, got {weight.dtype}") - if weight_scale_inv.dtype is not torch.float32: - raise TypeError(f"Native block-FP8 scales must remain FP32, got {weight_scale_inv.dtype}") - if input.device != weight.device or input.device != weight_scale_inv.device: - raise RuntimeError("Native block-FP8 weight, scales, and input must be on the same CUDA device") - - try: - from sglang.srt.layers.quantization.fp8_utils import ( # noqa: PLC0415 - triton_w8a8_block_fp8_linear, - ) - except Exception as exc: - raise RuntimeError("Pinned public SGLang block-FP8 kernel is required for native FP8") from exc - - return triton_w8a8_block_fp8_linear( - input, - weight, - list(block_size), - weight_scale_inv, - ) - - -def pack_fp8_as_float32(weight: torch.Tensor) -> torch.Tensor: - """Return a contiguous float32 view containing the exact float8 bytes.""" - - if weight.dtype is not _FP8_DTYPE: - raise TypeError(f"Expected float8_e4m3fn weight, got {weight.dtype}") - if weight.numel() % 4: - raise ValueError(f"FP8 weight has {weight.numel()} elements; byte packing requires a multiple of four") - return weight.contiguous().view(torch.uint8).view(torch.float32) - - -def unpack_float32_as_fp8(packed: torch.Tensor, shape: tuple[int, ...]) -> torch.Tensor: - """Recover a float8 view without numerical conversion.""" - - if packed.dtype is not torch.float32: - raise TypeError( - f"Packed native-FP8 state has dtype {packed.dtype}, expected float32; " - "an FSDP mixed-precision cast may have corrupted it" - ) - expected = 1 - for dim in shape: - expected *= dim - if packed.numel() * 4 != expected: - raise ValueError( - f"Packed native-FP8 state has {packed.numel() * 4} bytes, expected {expected} for shape {shape}" - ) - return packed.contiguous().view(torch.uint8).view(_FP8_DTYPE).reshape(shape) - - -def validate_native_fp8_state_metadata( - module: nn.Module, - metadata: dict[str, tuple[torch.dtype, tuple[int, ...]]], - *, - prefix: str = "", -) -> None: - """Fail before DCP load if serialized dtype/shape metadata can cast bytes. - - DCP callers must build ``metadata`` from the checkpoint reader before - invoking ``set_model_state_dict``. State-dict hooks below cover ordinary - ``load_state_dict``; this preflight covers loaders that copy shards without - calling module hooks. - """ - - expected = { - f"{prefix}{name}": (parameter.dtype, tuple(parameter.shape)) - for name, parameter in module.named_parameters() - if "packed_weight_f32" in name or name.endswith("weight_scale_inv") - } - missing = sorted(set(expected) - set(metadata)) - mismatched = { - name: (metadata[name], contract) - for name, contract in expected.items() - if name in metadata and metadata[name] != contract - } - if missing or mismatched: - raise ValueError(f"Native FP8 DCP metadata mismatch: missing={missing[:8]} mismatched={mismatched}") - - -def validate_native_fp8_dcp_checkpoint( - checkpoint_path: str, - expected_state: dict[str, torch.Tensor], - *, - state_prefix: str = "model.", -) -> None: - """Read DCP metadata and reject any castable native-FP8 payload pre-load.""" - - from torch.distributed.checkpoint import FileSystemReader # noqa: PLC0415 - - state_metadata = FileSystemReader(checkpoint_path).read_metadata().state_dict_metadata - tensor_metadata = {} - for name, metadata in state_metadata.items(): - properties = getattr(metadata, "properties", None) - size = getattr(metadata, "size", None) - if properties is not None and size is not None: - tensor_metadata[name] = (properties.dtype, tuple(size)) - expected_metadata = { - f"{state_prefix}{name}": (tensor.dtype, tuple(tensor.shape)) - for name, tensor in expected_state.items() - if "packed_weight_f32" in name or name.endswith("weight_scale_inv") - } - missing = sorted(set(expected_metadata) - set(tensor_metadata)) - mismatched = { - name: (tensor_metadata[name], contract) - for name, contract in expected_metadata.items() - if name in tensor_metadata and tensor_metadata[name] != contract - } - if missing or mismatched: - raise ValueError(f"Native FP8 DCP metadata mismatch: missing={missing[:8]} mismatched={mismatched}") - - -def _validate_state_dict_contract( - state_dict: dict[str, torch.Tensor], - prefix: str, - expected: dict[str, tuple[torch.dtype, tuple[int, ...]]], -) -> None: - for name, (dtype, shape) in expected.items(): - key = f"{prefix}{name}" - tensor = state_dict.get(key) - if tensor is None: - continue - if tensor.dtype is not dtype or tuple(tensor.shape) != shape: - raise TypeError( - f"Native FP8 state {key} must be {dtype} {shape}, got {tensor.dtype} {tuple(tensor.shape)}; " - "refusing a load_state_dict cast" - ) - - -class _NativeBlockFP8FrozenDgradFunction(torch.autograd.Function): - """Exact-value forward with the frozen-trunk activation backward. - - Forward: the UNCHANGED SGLang W8A8 dispatch on the module's frozen bytes - (byte-identical to the scoring-only path — the value program is not - touched). Backward: dgrad only, ``grad_output @ dequant(cache)`` in the - declared BF16 linear program — the same base-branch treatment the - trainable full-param composites and the exact QLoRA surrogate apply to - quantized bytes. There is deliberately NO wgrad and NO master/cache - mutation: the trunk is frozen, so the only gradient this boundary may - produce is the activation gradient that lets upstream trainable - parameters learn. - """ - - @staticmethod - def forward(ctx, input: torch.Tensor, module, out_start: int, out_end: int, in_start: int, in_end: int): - weight = module.fp8_weight()[out_start:out_end, in_start:in_end].contiguous() - scale = module.weight_scale_inv[ - out_start // 128 : (out_end + 127) // 128, - in_start // 128 : (in_end + 127) // 128, - ].contiguous() - output = _sglang_native_block_fp8_linear_value( - input, - weight, - scale, - block_size=module.block_size, - ) - ctx.module = module - ctx.ranges = (out_start, out_end, in_start, in_end) - # dgrad needs only the frozen weight slice (recovered in backward from - # the frozen module); the input is deliberately NOT saved — there is - # no wgrad at a frozen boundary. - return output - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - out_start, out_end, in_start, in_end = ctx.ranges - grad_input = ctx.module._frozen_activation_dgrad( - grad_output, - output_range=(out_start, out_end), - input_range=(in_start, in_end), - ) - return grad_input, None, None, None, None, None - - -class NativeBlockFP8Linear(nn.Module): - """Frozen W8A8 block-FP8 linear whose base weight is never dequantized. - - SGLang is imported only on an engaged CUDA forward. Construction, config - inspection, meta initialization, DCP planning, and ordinary non-FP8 model - imports therefore do not depend on SGLang. - """ - - fsdp_requires_full_precision = True - contract_version = NATIVE_BLOCK_FP8_CONTRACT_VERSION - frozen_dgrad_contract_version = NATIVE_BLOCK_FP8_FROZEN_DGRAD_CONTRACT_VERSION - # Fail-closed default: the scoring-only lanes never admit an activation - # backward. A training admission that OWNS the trainable-set semantics - # (the GLM-5.2 full-param admission) must opt in per module via - # :meth:`enable_frozen_activation_dgrad`. - _frozen_dgrad_admitted = False - _frozen_dgrad_engagement_logged = False - - def __init__( - self, - in_features: int, - out_features: int, - *, - block_size: tuple[int, int] = (128, 128), - device: torch.device | str | None = None, - ) -> None: - super().__init__() - if in_features <= 0 or out_features <= 0: - raise ValueError("Native block-FP8 dimensions must be positive") - if in_features % 4: - raise ValueError("Native block-FP8 in_features must be divisible by four for byte packing") - if tuple(block_size) != (128, 128): - raise ValueError(f"Only the official GLM block shape (128, 128) is supported, got {block_size}") - - self.in_features = int(in_features) - self.out_features = int(out_features) - self.block_size = tuple(block_size) - self.packed_weight_f32 = nn.Parameter( - torch.empty((out_features, in_features // 4), dtype=torch.float32, device=device), - requires_grad=False, - ) - self.weight_scale_inv = nn.Parameter( - torch.empty( - ((out_features + 127) // 128, (in_features + 127) // 128), - dtype=torch.float32, - device=device, - ), - requires_grad=False, - ) - - @classmethod - def from_linear(cls, module: nn.Linear) -> "NativeBlockFP8Linear": - if module.bias is not None: - raise ValueError("Native block-FP8 linear does not support bias") - return cls( - module.in_features, - module.out_features, - device=module.weight.device, - ) - - def _apply(self, fn, recurse: bool = True): - # Preserve packed bytes and scale values across model.to(dtype=...). - probe = fn(torch.empty(0, dtype=torch.float32, device=self.packed_weight_f32.device)) - protected = {name: self._parameters.pop(name) for name in ("packed_weight_f32", "weight_scale_inv")} - try: - result = super()._apply(fn, recurse=recurse) - replacements = {} - for name, parameter in protected.items(): - if parameter.is_meta: - value = torch.empty_like(parameter, dtype=torch.float32, device=probe.device) - else: - value = parameter.to(device=probe.device, dtype=torch.float32) - replacements[name] = nn.Parameter(value, requires_grad=False) - except Exception: - self._parameters.update(protected) - raise - self._parameters.update(replacements) - return result - - def _load_from_state_dict( - self, - state_dict, - prefix, - local_metadata, - strict, - missing_keys, - unexpected_keys, - error_msgs, - ): - _validate_state_dict_contract( - state_dict, - prefix, - { - "packed_weight_f32": (torch.float32, tuple(self.packed_weight_f32.shape)), - "weight_scale_inv": (torch.float32, tuple(self.weight_scale_inv.shape)), - }, - ) - return super()._load_from_state_dict( - state_dict, - prefix, - local_metadata, - strict, - missing_keys, - unexpected_keys, - error_msgs, - ) - - def load_prequantized(self, weight: torch.Tensor, weight_scale_inv: torch.Tensor) -> None: - """Copy one official weight/scale pair with strict byte and shape checks.""" - - if tuple(weight.shape) != (self.out_features, self.in_features): - raise ValueError( - f"FP8 weight shape {tuple(weight.shape)} does not match ({self.out_features}, {self.in_features})" - ) - if weight_scale_inv.dtype is not torch.float32: - raise TypeError(f"weight_scale_inv must remain FP32, got {weight_scale_inv.dtype}") - if tuple(weight_scale_inv.shape) != tuple(self.weight_scale_inv.shape): - raise ValueError( - f"FP8 scale shape {tuple(weight_scale_inv.shape)} does not match {tuple(self.weight_scale_inv.shape)}" - ) - if not bool(torch.all(torch.isfinite(weight_scale_inv))): - raise ValueError("weight_scale_inv contains non-finite values") - packed = pack_fp8_as_float32(weight) - with torch.no_grad(): - self.packed_weight_f32.copy_(packed.to(self.packed_weight_f32.device)) - self.weight_scale_inv.copy_(weight_scale_inv.to(self.weight_scale_inv.device)) - - def fp8_weight(self) -> torch.Tensor: - return unpack_float32_as_fp8( - self.packed_weight_f32, - (self.out_features, self.in_features), - ) - - def enable_frozen_activation_dgrad(self) -> None: - """Admit the validated frozen-trunk activation backward. - - Idempotent per module. The forward VALUE program is byte-unchanged; - only the refusal on grad-requiring inputs is replaced by the checked - BF16 dequant-program dgrad (no wgrad, no master/cache mutation — - frozen means frozen). Only a trainable-set admission that must - backpropagate THROUGH this frozen module may call this; scoring-only - lanes keep the fail-closed refusal. - """ - - self._frozen_dgrad_admitted = True - cls = NativeBlockFP8Linear - if not cls._frozen_dgrad_engagement_logged: - cls._frozen_dgrad_engagement_logged = True - logger.info( - "Native block-FP8 frozen-trunk activation dgrad engaged: contract=%s " - "(forward bytes unchanged; dgrad = grad_output @ dequant(cache) in the " - "declared BF16 program; wgrad none; frozen bytes immutable)", - self.frozen_dgrad_contract_version, - ) - - def _frozen_activation_dgrad( - self, - grad_output: torch.Tensor, - *, - output_range: tuple[int, int], - input_range: tuple[int, int], - ) -> torch.Tensor: - """dgrad through the dequantized frozen bytes in the BF16 program. - - Dequantization is the same program point the trainable composites' - surrogate uses (``block_fp8_dequantize_gkn`` -> FP32 -> one explicit - BF16 rounding); the GEMM is the declared BF16 linear backward - ``grad_output @ W``. Gated bitwise against the reference - dequant-matmul autograd (tests/ops/test_block_fp8_frozen_dgrad.py). - """ - - from xorl.ops.quantize import block_fp8_dequantize_gkn # noqa: PLC0415 - - out_start, out_end = output_range - in_start, in_end = input_range - weight = self.fp8_weight()[out_start:out_end, in_start:in_end].contiguous() - scale = self.weight_scale_inv[ - out_start // 128 : (out_end + 127) // 128, - in_start // 128 : (in_end + 127) // 128, - ].contiguous() - dequantized = block_fp8_dequantize_gkn(weight, scale, 128).to(torch.float32).to(torch.bfloat16) - grad_2d = grad_output.reshape(-1, out_end - out_start).to(torch.bfloat16) - return grad_2d.matmul(dequantized).reshape(*grad_output.shape[:-1], in_end - in_start) - - @staticmethod - def _validate_partition(value: tuple[int, int] | None, size: int, name: str) -> tuple[int, int]: - if value is None: - return 0, size - start, end = value - if start < 0 or end > size or start >= end: - raise ValueError(f"Invalid native block-FP8 {name} range {value} for size {size}") - if start % 128 or (end != size and end % 128): - raise ValueError(f"Native block-FP8 {name} range {value} must follow 128-element block boundaries") - return start, end - - def forward_partition( - self, - input: torch.Tensor, - *, - output_range: tuple[int, int] | None = None, - input_range: tuple[int, int] | None = None, - ) -> torch.Tensor: - """Run an aligned output-row/input-column partition through SGLang. - - This supports callers that preserve a model-specific reduction policy - by assigning disjoint, block-aligned projection slices to ranks. - """ - - if self.packed_weight_f32.requires_grad or self.weight_scale_inv.requires_grad: - raise RuntimeError("Native block-FP8 base weights and scales must remain frozen") - grad_engaged = torch.is_grad_enabled() and input.requires_grad - if grad_engaged and not self._frozen_dgrad_admitted: - raise RuntimeError( - "Native block-FP8 phase-one forward is scoring-only; activation backward requires a validated kernel" - ) - out_start, out_end = self._validate_partition(output_range, self.out_features, "output") - in_start, in_end = self._validate_partition(input_range, self.in_features, "input") - if input.shape[-1] != in_end - in_start: - raise ValueError( - f"Native block-FP8 input width {input.shape[-1]} does not match selected range {in_start}:{in_end}" - ) - if grad_engaged: - # Same slicing + same kernel inside the autograd boundary; the - # value bytes are identical to the scoring-only path below. - return _NativeBlockFP8FrozenDgradFunction.apply(input, self, out_start, out_end, in_start, in_end) - weight = self.fp8_weight()[out_start:out_end, in_start:in_end].contiguous() - scale = self.weight_scale_inv[ - out_start // 128 : (out_end + 127) // 128, - in_start // 128 : (in_end + 127) // 128, - ].contiguous() - output = _sglang_native_block_fp8_linear_value( - input, - weight, - scale, - block_size=self.block_size, - ) - return output - - def forward( - self, - input: torch.Tensor | None = None, - *, - return_dequantized_weight: bool = False, - output_range: tuple[int, int] | None = None, - input_range: tuple[int, int] | None = None, - ) -> torch.Tensor: - if return_dequantized_weight: - if input is not None or output_range is not None or input_range is not None: - raise ValueError("Native block-FP8 weight materialization does not accept activation or range inputs") - if self.packed_weight_f32.device.type != "cuda": - raise RuntimeError("Native block-FP8 weight materialization requires CUDA") - from sglang.srt.layers.quantization.fp8_utils import block_quant_dequant # noqa: PLC0415 - - return block_quant_dequant( - self.fp8_weight(), - self.weight_scale_inv, - list(self.block_size), - torch.bfloat16, - ) - if input is None: - raise ValueError("Native block-FP8 linear forward requires an activation input") - return self.forward_partition(input, output_range=output_range, input_range=input_range) - - -__all__ = [ - "NATIVE_BLOCK_FP8_CONTRACT_VERSION", - "NATIVE_BLOCK_FP8_FROZEN_DGRAD_CONTRACT_VERSION", - "NativeBlockFP8Linear", - "pack_fp8_as_float32", - "unpack_float32_as_fp8", - "validate_native_fp8_dcp_checkpoint", - "validate_native_fp8_state_metadata", -] +from xorl.ops.exact.block_fp8_native import * # noqa: F401,F403 diff --git a/src/xorl/ops/canonical_moe_cast.py b/src/xorl/ops/canonical_moe_cast.py index f58e8b56..4302cf79 100644 --- a/src/xorl/ops/canonical_moe_cast.py +++ b/src/xorl/ops/canonical_moe_cast.py @@ -1,123 +1,3 @@ -"""Exact final cast for CUDA canonical-MoE FP64 accumulators.""" +"""Moved to ``xorl.ops.exact.canonical_moe_cast`` (issue #78 phase 3); compat stub.""" -from __future__ import annotations - -import torch - - -try: # Triton is optional in CPU-only development environments. - import triton - import triton.language as tl - - _TRITON_AVAILABLE = True -except ImportError: # pragma: no cover - exercised only without Triton - _TRITON_AVAILABLE = False - - -def validate_canonical_moe_fp64_cast_input( - value: torch.Tensor, - output_dtype: torch.dtype, -) -> None: - if value.dtype is not torch.float64: - raise TypeError(f"Canonical MoE direct cast requires FP64 input, got {value.dtype}") - if output_dtype not in (torch.bfloat16, torch.float16): - raise TypeError(f"Canonical MoE direct cast requires BF16 or FP16 output, got {output_dtype}") - if not value.is_cuda: - raise ValueError("Canonical MoE direct cast requires a CUDA tensor") - if not value.is_contiguous(): - raise ValueError("Canonical MoE direct cast requires contiguous input") - - -if _TRITON_AVAILABLE: - - @triton.jit - def _canonical_moe_fp64_to_lowp_rne_kernel( - input_ptr, - output_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, - ): - offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - value = tl.load(input_ptr + offsets, mask=mask, other=0.0) - # Store the FP64 register directly through the BF16 pointer. An - # explicit FP32 conversion here would introduce a double-rounding - # boundary and can select the opposite BF16 neighbor at a midpoint. - tl.store(output_ptr + offsets, value, mask=mask) - - -def _canonical_moe_fp64_to_lowp_cuda( - value: torch.Tensor, - output_dtype: torch.dtype, -) -> torch.Tensor: - if not _TRITON_AVAILABLE: # pragma: no cover - guarded by CUDA environments - raise RuntimeError("Canonical MoE CUDA direct cast requires Triton") - output = torch.empty_like(value, dtype=output_dtype) - if output.numel() == 0: - return output - block_size = 256 - _canonical_moe_fp64_to_lowp_rne_kernel[(triton.cdiv(output.numel(), block_size),)]( - value, - output, - output.numel(), - BLOCK_SIZE=block_size, - num_warps=4, - ) - return output - - -@torch.library.custom_op("xorl_k3::_canonical_moe_fp64_to_bf16_rne", mutates_args=()) -def _canonical_moe_fp64_to_bf16_rne_op(value: torch.Tensor) -> torch.Tensor: - """Round FP64 directly to BF16.""" - validate_canonical_moe_fp64_cast_input(value, torch.bfloat16) - return _canonical_moe_fp64_to_lowp_cuda(value, torch.bfloat16) - - -@_canonical_moe_fp64_to_bf16_rne_op.register_fake -def _canonical_moe_fp64_to_bf16_rne_fake(value: torch.Tensor) -> torch.Tensor: - return torch.empty_like(value, dtype=torch.bfloat16) - - -@torch.library.custom_op("xorl_k3::_canonical_moe_fp64_to_fp16_rne", mutates_args=()) -def _canonical_moe_fp64_to_fp16_rne_op(value: torch.Tensor) -> torch.Tensor: - """Round FP64 directly to FP16.""" - validate_canonical_moe_fp64_cast_input(value, torch.float16) - return _canonical_moe_fp64_to_lowp_cuda(value, torch.float16) - - -@_canonical_moe_fp64_to_fp16_rne_op.register_fake -def _canonical_moe_fp64_to_fp16_rne_fake(value: torch.Tensor) -> torch.Tensor: - return torch.empty_like(value, dtype=torch.float16) - - -def _canonical_moe_fp64_to_lowp_backward(ctx, grad_output: torch.Tensor): - del ctx - # Match torch's straight-through cast derivative. The surrounding FP64 - # tree then distributes this FP64 gradient to every contributor before its - # original BF16 leaf cast handles the final backward conversion. - return grad_output.to(torch.float64) - - -_canonical_moe_fp64_to_bf16_rne_op.register_autograd( - _canonical_moe_fp64_to_lowp_backward, -) -_canonical_moe_fp64_to_fp16_rne_op.register_autograd( - _canonical_moe_fp64_to_lowp_backward, -) - - -def canonical_moe_fp64_to_lowp_rne( - value: torch.Tensor, - output_dtype: torch.dtype, -) -> torch.Tensor: - """Validate and directly round FP64 to the canonical transport dtype.""" - validate_canonical_moe_fp64_cast_input(value, output_dtype) - if output_dtype is torch.bfloat16: - return _canonical_moe_fp64_to_bf16_rne_op(value) - return _canonical_moe_fp64_to_fp16_rne_op(value) - - -__all__ = [ - "canonical_moe_fp64_to_lowp_rne", - "validate_canonical_moe_fp64_cast_input", -] +from xorl.ops.exact.canonical_moe_cast import * # noqa: F401,F403 diff --git a/src/xorl/ops/canonical_moe_leaf.py b/src/xorl/ops/canonical_moe_leaf.py index 1165471a..5024c8fd 100644 --- a/src/xorl/ops/canonical_moe_leaf.py +++ b/src/xorl/ops/canonical_moe_leaf.py @@ -1,163 +1,3 @@ -"""Compile-stable one-round arithmetic for canonical MoE contributor leaves.""" +"""Moved to ``xorl.ops.exact.canonical_moe_leaf`` (issue #78 phase 3); compat stub.""" -from __future__ import annotations - -import torch - - -try: # Triton is optional in CPU-only development environments. - import triton - import triton.language as tl - - _TRITON_AVAILABLE = True -except ImportError: # pragma: no cover - exercised only without Triton - _TRITON_AVAILABLE = False - - -def validate_canonical_moe_leaf_operands( - shared: torch.Tensor, - routed: torch.Tensor, -) -> None: - if shared.shape != routed.shape: - raise ValueError( - f"Canonical MoE leaf operands must have equal shapes, got {tuple(shared.shape)} and {tuple(routed.shape)}" - ) - if shared.dtype != routed.dtype or shared.dtype not in ( - torch.bfloat16, - torch.float16, - ): - raise TypeError( - f"Canonical MoE leaf operands must have the same BF16 or FP16 dtype, got {shared.dtype} and {routed.dtype}" - ) - if shared.device != routed.device: - raise ValueError(f"Canonical MoE leaf operands must share a device, got {shared.device} and {routed.device}") - if shared.device.type not in ("cpu", "cuda"): - raise ValueError(f"Canonical MoE leaf operands must be CPU or CUDA tensors, got {shared.device.type}") - if not shared.is_contiguous() or not routed.is_contiguous(): - raise ValueError("Canonical MoE leaf operands must be contiguous") - - -if _TRITON_AVAILABLE: - - @triton.jit - def _fp32_fma_rn(multiplicand, multiplier, addend): - return tl.inline_asm_elementwise( - asm="fma.rn.f32 $0, $1, $2, $3;", - constraints="=f,f,f,f", - args=(multiplicand, multiplier, addend), - dtype=tl.float32, - is_pure=True, - pack=1, - ) - - @triton.jit - def _canonical_moe_leaf_fp32_kernel( - shared_ptr, - routed_ptr, - output_ptr, - n_elements, - routed_scale, - BLOCK_SIZE: tl.constexpr, - ): - offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - shared = tl.load(shared_ptr + offsets, mask=mask, other=0.0).to(tl.float32) - routed = tl.load(routed_ptr + offsets, mask=mask, other=0.0).to(tl.float32) - leaf = _fp32_fma_rn(routed, routed_scale.to(tl.float32), shared) - # The accumulator stays in a register. This store is the sole cast to - # the declared BF16/FP16 transport dtype. - tl.store(output_ptr + offsets, leaf, mask=mask) - - -def _canonical_moe_leaf_cuda( - shared: torch.Tensor, - routed: torch.Tensor, - routed_scale: float, -) -> torch.Tensor: - if not _TRITON_AVAILABLE: # pragma: no cover - guarded by CUDA environments - raise RuntimeError("Canonical MoE CUDA leaf construction requires Triton") - output = torch.empty_like(shared) - if output.numel() == 0: - return output - block_size = 256 - _canonical_moe_leaf_fp32_kernel[(triton.cdiv(output.numel(), block_size),)]( - shared, - routed, - output, - output.numel(), - float(routed_scale), - BLOCK_SIZE=block_size, - num_warps=4, - ) - return output - - -def _canonical_moe_leaf_cpu( - shared: torch.Tensor, - routed: torch.Tensor, - routed_scale: float, -) -> torch.Tensor: - # BF16/FP16 operands and an FP32 scale have at most 35 product bits, so the - # product is exact in FP64. The FP64 add followed by one FP32 round is the - # independent CPU reference (apart from a theoretical double-rounding - # corner); the CUDA identity tests pin the production FMA bytes. The final - # conversion is the sole transport cast. - scale_fp32 = torch.tensor(routed_scale, dtype=torch.float32) - pre_cast = shared.double() + routed.double() * scale_fp32.double() - return pre_cast.float().to(shared.dtype) - - -@torch.library.custom_op("xorl_k3::_canonical_moe_leaf_fp32_v1", mutates_args=()) -def _canonical_moe_leaf_fp32_v1_op( - shared: torch.Tensor, - routed: torch.Tensor, - routed_scale: float, -) -> torch.Tensor: - """Opaque one-FMA leaf primitive; output is the transport dtype.""" - validate_canonical_moe_leaf_operands(shared, routed) - if shared.is_cuda: - return _canonical_moe_leaf_cuda(shared, routed, routed_scale) - return _canonical_moe_leaf_cpu(shared, routed, routed_scale) - - -@_canonical_moe_leaf_fp32_v1_op.register_fake -def _canonical_moe_leaf_fp32_v1_fake( - shared: torch.Tensor, - routed: torch.Tensor, - routed_scale: float, -) -> torch.Tensor: - del routed, routed_scale - return torch.empty_like(shared) - - -def _canonical_moe_leaf_setup_context(ctx, inputs, output) -> None: - del output - shared, _routed, routed_scale = inputs - ctx.input_dtype = shared.dtype - ctx.routed_scale = float(routed_scale) - - -def _canonical_moe_leaf_backward(ctx, grad_output: torch.Tensor): - grad_fp32 = grad_output.float() - grad_shared = grad_fp32.to(ctx.input_dtype) - grad_routed = (grad_fp32 * ctx.routed_scale).to(ctx.input_dtype) - return grad_shared, grad_routed, None - - -_canonical_moe_leaf_fp32_v1_op.register_autograd( - _canonical_moe_leaf_backward, - setup_context=_canonical_moe_leaf_setup_context, -) - - -def canonical_moe_leaf_fp32_v1_op( - shared: torch.Tensor, - routed: torch.Tensor, - routed_scale: float, -) -> torch.Tensor: - """Validate and execute the compile-opaque canonical contributor leaf.""" - validate_canonical_moe_leaf_operands(shared, routed) - return _canonical_moe_leaf_fp32_v1_op(shared, routed, float(routed_scale)) - - -__all__ = ["canonical_moe_leaf_fp32_v1_op", "validate_canonical_moe_leaf_operands"] +from xorl.ops.exact.canonical_moe_leaf import * # noqa: F401,F403 diff --git a/src/xorl/ops/exact/__init__.py b/src/xorl/ops/exact/__init__.py new file mode 100644 index 00000000..c77ddae0 --- /dev/null +++ b/src/xorl/ops/exact/__init__.py @@ -0,0 +1,21 @@ +"""The serving-parity (exact) contract programs (issue #78 phase 3). + +One home and one name for what the historical ``bi_`` / ``exact_`` / +``canonical_`` / ``class_b`` prefixes all meant: byte-pinned programs shared +with the serving engine. + +Physically here: ``rope_class_b``, ``canonical_moe_leaf``, +``canonical_moe_cast``, ``kernel_config_pin``, ``bi_gemm_configs``, +``block_fp8_native``, ``fused_silu_and_mul``. + +Aliased here but deliberately NOT moved: + +- ``families_v2`` -> :mod:`xorl.ops.bi_families_v2` — vendored byte-identical + into the serving engine and sha256-gated; the file cannot move. +- ``batch_invariant`` -> :mod:`xorl.ops.batch_invariant_ops` — vendored- + adapted from SGLang's ``srt/batch_invariant_ops``; it stays a single file + at its path so it remains diffable against the serving twin. +- ``sampling_transforms`` -> :mod:`xorl.ops.exact_sampling_transforms` — the + replay contract; kept in place while in-flight work (#74) rewrites it, to + be flipped to canonical here afterwards. +""" diff --git a/src/xorl/ops/exact/batch_invariant.py b/src/xorl/ops/exact/batch_invariant.py new file mode 100644 index 00000000..943e6769 --- /dev/null +++ b/src/xorl/ops/exact/batch_invariant.py @@ -0,0 +1,3 @@ +"""Alias for :mod:`xorl.ops.batch_invariant_ops` (parity twin of SGLang's copy).""" + +from xorl.ops.batch_invariant_ops import * # noqa: F401,F403 diff --git a/src/xorl/ops/exact/bi_gemm_configs.py b/src/xorl/ops/exact/bi_gemm_configs.py new file mode 100644 index 00000000..f366058f --- /dev/null +++ b/src/xorl/ops/exact/bi_gemm_configs.py @@ -0,0 +1,147 @@ +"""Shape-keyed tile configs for the batch-invariant persistent matmul (gap-map R1). + +Only BLOCK_SIZE_K is bit-relevant in ``matmul_kernel_persistent``: it sets the +per-element K-reduction order. BLOCK_SIZE_M/N, GROUP_SIZE_M, num_stages and +num_warps are bit-neutral tuning axes (warps split M/N, never K; no split-K), +so they may vary per shape without moving the contract bits. Every entry below +was admitted only after torch.equal against the pinned baseline on multiple +seeds plus a cross-M row-invariance check (xorl +experiments/k3_tests/tools/tune_bi_gemm.py, H100, 2026-07-07). + +The shape-keyed table is the production configuration. The pinned baseline is +retained only as an internal launch-failure fallback; it is not a user-selectable +mode. +""" + +# The bit-relevant axis: pinned per dtype, NEVER shape-keyed. +PINNED_BLOCK_K = { + "torch.bfloat16": 64, + "torch.float16": 64, + "torch.float32": 32, +} + +# The pre-R1 one-config-per-dtype baselines (also the launch-failure fallback). +BASELINE_CONFIG = { + "torch.bfloat16": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}, + "torch.float16": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}, + "torch.float32": {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}, +} + +# Class defaults on (dtype, M-bucket) for shapes without an exact (K, N) entry. +# Buckets: decode/graph shapes (M<=256) want tiny BLOCK_M; batch shapes want +# the wide 128x256 tile. Measured on delphi-class shapes; bit-neutral either way. +CLASS_DEFAULTS = { + "torch.bfloat16": ( + (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), + (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), + ), + "torch.float32": ( + (256, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 4, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), + ), +} + +# Exact-shape entries: {(dtype, K, N): ((max_M, cfg), ..., (None, cfg))}. +# Generated from offline tuner results. Do not hand-edit these values; rerun +# the tuning and generation pipeline instead. +# BEGIN GENERATED TABLE +TABLE = { + ("torch.bfloat16", 3840, 3840): ( + (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 8, "num_stages": 12, "num_warps": 4}), + (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "GROUP_SIZE_M": 1, "num_stages": 12, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), + ), + ("torch.bfloat16", 3840, 5376): ( + (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 1, "num_stages": 12, "num_warps": 4}), + (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 1, "num_stages": 8, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), + ), + ("torch.bfloat16", 3840, 8192): ( + (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 8, "num_stages": 8, "num_warps": 4}), + (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 1, "num_stages": 8, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), + ), + ("torch.bfloat16", 3840, 11520): ( + (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 5, "num_warps": 4}), + (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), + ), + ("torch.bfloat16", 3840, 30720): ( + (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), + (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 1, "num_stages": 5, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), + ), + ("torch.bfloat16", 3840, 128256): ( + (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 5, "num_warps": 4}), + (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), + ), + ("torch.bfloat16", 15360, 3840): ( + (16, {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "GROUP_SIZE_M": 1, "num_stages": 20, "num_warps": 4}), + (256, {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "GROUP_SIZE_M": 1, "num_stages": 12, "num_warps": 4}), + (8192, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 8, "num_stages": 4, "num_warps": 8}), + (None, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "GROUP_SIZE_M": 1, "num_stages": 4, "num_warps": 8}), + ), + ("torch.float32", 3840, 8192): ( + (8192, {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 5, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), + ), +} +# END GENERATED TABLE + + +# Conservative H100 dynamic-smem budget. The persistent kernel needs +# num_stages * (BM*BK + BK*BN) * in_elt for the pipeline plus BM*BN*out_elt for +# the epilogue store staging when the output is wider than the inputs (the BI +# lm-head's fp32-out form) — a config that fits bf16-out can OOM fp32-out. +_SMEM_BUDGET = 220 * 1024 + +_ELT = {"torch.bfloat16": 2, "torch.float16": 2, "torch.float32": 4} + + +def _fits(cfg, block_k: int, in_elt: int, out_elt: int) -> bool: + pipe = cfg["num_stages"] * (cfg["BLOCK_SIZE_M"] * block_k + block_k * cfg["BLOCK_SIZE_N"]) * in_elt + # same-width outputs reuse pipeline smem; wider outputs stage a full C tile + epilogue = cfg["BLOCK_SIZE_M"] * cfg["BLOCK_SIZE_N"] * out_elt if out_elt > in_elt else 0 + return pipe + epilogue <= _SMEM_BUDGET + + +def lookup_mm_config(dtype, M: int, N: int, K: int, out_itemsize: int | None = None): + """Tile config for the persistent matmul at (dtype, M, N, K). + + Returns the full launch config including the pinned BLOCK_SIZE_K. Exact + (K, N) entries win, then (dtype, M-bucket) class defaults, then the pinned + baseline. All candidates are bit-identical; this only picks speed. Pass + ``out_itemsize`` when the output dtype is wider than the inputs so the + epilogue smem staging is budgeted (oversized configs fall back to baseline). + """ + key = str(dtype) + block_k = PINNED_BLOCK_K.get(key) + if block_k is None: + raise ValueError(f"Unsupported dtype {dtype} for batch-invariant matmul") + in_elt = _ELT[key] + out_elt = out_itemsize if out_itemsize is not None else in_elt + cfg = None + buckets = TABLE.get((key, K, N)) + exact = buckets is not None + if buckets is None: + buckets = CLASS_DEFAULTS.get(key) + if buckets is not None: + for max_m, bucket_cfg in buckets: + if max_m is None or M <= max_m: + # exact entries were measured/launch-validated at their out + # dtype; class defaults get the conservative smem model (the + # launch sites also fall back to baseline on OutOfResources) + if exact or _fits(bucket_cfg, block_k, in_elt, out_elt): + cfg = bucket_cfg + break + if cfg is None: + cfg = BASELINE_CONFIG[key] + return dict(cfg, BLOCK_SIZE_K=block_k) + + +def baseline_mm_config(dtype): + """The pinned pre-R1 config for dtype — the launch-failure fallback.""" + key = str(dtype) + return dict(BASELINE_CONFIG[key], BLOCK_SIZE_K=PINNED_BLOCK_K[key]) diff --git a/src/xorl/ops/exact/block_fp8_native.py b/src/xorl/ops/exact/block_fp8_native.py new file mode 100644 index 00000000..f0c78158 --- /dev/null +++ b/src/xorl/ops/exact/block_fp8_native.py @@ -0,0 +1,488 @@ +"""Frozen native block-FP8 state and exact SGLang dense dispatch. + +The packed parameter uses the same public storage convention as XoRL QLoRA: +four float8 bytes are viewed as one float32 element. This makes the payload a +normal reshardable parameter while retaining every checkpoint byte. It must +be wrapped without an FSDP mixed-precision policy; casting the packed float32 +values numerically would corrupt the embedded bytes. +""" + +from __future__ import annotations + +import logging + +import torch +from torch import nn + + +logger = logging.getLogger(__name__) + +NATIVE_BLOCK_FP8_CONTRACT_VERSION = "xorl_native_block_fp8_sglang_v1" +NATIVE_BLOCK_FP8_FROZEN_DGRAD_CONTRACT_VERSION = "xorl_native_block_fp8_frozen_dgrad_v1" +_FP8_DTYPE = torch.float8_e4m3fn + + +def _sglang_native_block_fp8_linear_value( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + *, + block_size: tuple[int, int] = (128, 128), +) -> torch.Tensor: + """Run the shared exact-value SGLang W8A8 dispatch. + + Gradient policy belongs to the caller: :class:`NativeBlockFP8Linear` + remains scoring-only, while the exact QLoRA wrapper supplies its validated + surrogate through a custom autograd boundary. + """ + + if input.device.type != "cuda": + raise RuntimeError("Native block-FP8 forward requires CUDA and the pinned SGLang Triton kernel") + if input.dtype is not torch.bfloat16: + raise TypeError(f"Native block-FP8 forward requires BF16 activations, got {input.dtype}") + if weight.dtype is not _FP8_DTYPE: + raise TypeError(f"Native block-FP8 weight must remain float8_e4m3fn, got {weight.dtype}") + if weight_scale_inv.dtype is not torch.float32: + raise TypeError(f"Native block-FP8 scales must remain FP32, got {weight_scale_inv.dtype}") + if input.device != weight.device or input.device != weight_scale_inv.device: + raise RuntimeError("Native block-FP8 weight, scales, and input must be on the same CUDA device") + + try: + from sglang.srt.layers.quantization.fp8_utils import ( # noqa: PLC0415 + triton_w8a8_block_fp8_linear, + ) + except Exception as exc: + raise RuntimeError("Pinned public SGLang block-FP8 kernel is required for native FP8") from exc + + return triton_w8a8_block_fp8_linear( + input, + weight, + list(block_size), + weight_scale_inv, + ) + + +def pack_fp8_as_float32(weight: torch.Tensor) -> torch.Tensor: + """Return a contiguous float32 view containing the exact float8 bytes.""" + + if weight.dtype is not _FP8_DTYPE: + raise TypeError(f"Expected float8_e4m3fn weight, got {weight.dtype}") + if weight.numel() % 4: + raise ValueError(f"FP8 weight has {weight.numel()} elements; byte packing requires a multiple of four") + return weight.contiguous().view(torch.uint8).view(torch.float32) + + +def unpack_float32_as_fp8(packed: torch.Tensor, shape: tuple[int, ...]) -> torch.Tensor: + """Recover a float8 view without numerical conversion.""" + + if packed.dtype is not torch.float32: + raise TypeError( + f"Packed native-FP8 state has dtype {packed.dtype}, expected float32; " + "an FSDP mixed-precision cast may have corrupted it" + ) + expected = 1 + for dim in shape: + expected *= dim + if packed.numel() * 4 != expected: + raise ValueError( + f"Packed native-FP8 state has {packed.numel() * 4} bytes, expected {expected} for shape {shape}" + ) + return packed.contiguous().view(torch.uint8).view(_FP8_DTYPE).reshape(shape) + + +def validate_native_fp8_state_metadata( + module: nn.Module, + metadata: dict[str, tuple[torch.dtype, tuple[int, ...]]], + *, + prefix: str = "", +) -> None: + """Fail before DCP load if serialized dtype/shape metadata can cast bytes. + + DCP callers must build ``metadata`` from the checkpoint reader before + invoking ``set_model_state_dict``. State-dict hooks below cover ordinary + ``load_state_dict``; this preflight covers loaders that copy shards without + calling module hooks. + """ + + expected = { + f"{prefix}{name}": (parameter.dtype, tuple(parameter.shape)) + for name, parameter in module.named_parameters() + if "packed_weight_f32" in name or name.endswith("weight_scale_inv") + } + missing = sorted(set(expected) - set(metadata)) + mismatched = { + name: (metadata[name], contract) + for name, contract in expected.items() + if name in metadata and metadata[name] != contract + } + if missing or mismatched: + raise ValueError(f"Native FP8 DCP metadata mismatch: missing={missing[:8]} mismatched={mismatched}") + + +def validate_native_fp8_dcp_checkpoint( + checkpoint_path: str, + expected_state: dict[str, torch.Tensor], + *, + state_prefix: str = "model.", +) -> None: + """Read DCP metadata and reject any castable native-FP8 payload pre-load.""" + + from torch.distributed.checkpoint import FileSystemReader # noqa: PLC0415 + + state_metadata = FileSystemReader(checkpoint_path).read_metadata().state_dict_metadata + tensor_metadata = {} + for name, metadata in state_metadata.items(): + properties = getattr(metadata, "properties", None) + size = getattr(metadata, "size", None) + if properties is not None and size is not None: + tensor_metadata[name] = (properties.dtype, tuple(size)) + expected_metadata = { + f"{state_prefix}{name}": (tensor.dtype, tuple(tensor.shape)) + for name, tensor in expected_state.items() + if "packed_weight_f32" in name or name.endswith("weight_scale_inv") + } + missing = sorted(set(expected_metadata) - set(tensor_metadata)) + mismatched = { + name: (tensor_metadata[name], contract) + for name, contract in expected_metadata.items() + if name in tensor_metadata and tensor_metadata[name] != contract + } + if missing or mismatched: + raise ValueError(f"Native FP8 DCP metadata mismatch: missing={missing[:8]} mismatched={mismatched}") + + +def _validate_state_dict_contract( + state_dict: dict[str, torch.Tensor], + prefix: str, + expected: dict[str, tuple[torch.dtype, tuple[int, ...]]], +) -> None: + for name, (dtype, shape) in expected.items(): + key = f"{prefix}{name}" + tensor = state_dict.get(key) + if tensor is None: + continue + if tensor.dtype is not dtype or tuple(tensor.shape) != shape: + raise TypeError( + f"Native FP8 state {key} must be {dtype} {shape}, got {tensor.dtype} {tuple(tensor.shape)}; " + "refusing a load_state_dict cast" + ) + + +class _NativeBlockFP8FrozenDgradFunction(torch.autograd.Function): + """Exact-value forward with the frozen-trunk activation backward. + + Forward: the UNCHANGED SGLang W8A8 dispatch on the module's frozen bytes + (byte-identical to the scoring-only path — the value program is not + touched). Backward: dgrad only, ``grad_output @ dequant(cache)`` in the + declared BF16 linear program — the same base-branch treatment the + trainable full-param composites and the exact QLoRA surrogate apply to + quantized bytes. There is deliberately NO wgrad and NO master/cache + mutation: the trunk is frozen, so the only gradient this boundary may + produce is the activation gradient that lets upstream trainable + parameters learn. + """ + + @staticmethod + def forward(ctx, input: torch.Tensor, module, out_start: int, out_end: int, in_start: int, in_end: int): + weight = module.fp8_weight()[out_start:out_end, in_start:in_end].contiguous() + scale = module.weight_scale_inv[ + out_start // 128 : (out_end + 127) // 128, + in_start // 128 : (in_end + 127) // 128, + ].contiguous() + output = _sglang_native_block_fp8_linear_value( + input, + weight, + scale, + block_size=module.block_size, + ) + ctx.module = module + ctx.ranges = (out_start, out_end, in_start, in_end) + # dgrad needs only the frozen weight slice (recovered in backward from + # the frozen module); the input is deliberately NOT saved — there is + # no wgrad at a frozen boundary. + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + out_start, out_end, in_start, in_end = ctx.ranges + grad_input = ctx.module._frozen_activation_dgrad( + grad_output, + output_range=(out_start, out_end), + input_range=(in_start, in_end), + ) + return grad_input, None, None, None, None, None + + +class NativeBlockFP8Linear(nn.Module): + """Frozen W8A8 block-FP8 linear whose base weight is never dequantized. + + SGLang is imported only on an engaged CUDA forward. Construction, config + inspection, meta initialization, DCP planning, and ordinary non-FP8 model + imports therefore do not depend on SGLang. + """ + + fsdp_requires_full_precision = True + contract_version = NATIVE_BLOCK_FP8_CONTRACT_VERSION + frozen_dgrad_contract_version = NATIVE_BLOCK_FP8_FROZEN_DGRAD_CONTRACT_VERSION + # Fail-closed default: the scoring-only lanes never admit an activation + # backward. A training admission that OWNS the trainable-set semantics + # (the GLM-5.2 full-param admission) must opt in per module via + # :meth:`enable_frozen_activation_dgrad`. + _frozen_dgrad_admitted = False + _frozen_dgrad_engagement_logged = False + + def __init__( + self, + in_features: int, + out_features: int, + *, + block_size: tuple[int, int] = (128, 128), + device: torch.device | str | None = None, + ) -> None: + super().__init__() + if in_features <= 0 or out_features <= 0: + raise ValueError("Native block-FP8 dimensions must be positive") + if in_features % 4: + raise ValueError("Native block-FP8 in_features must be divisible by four for byte packing") + if tuple(block_size) != (128, 128): + raise ValueError(f"Only the official GLM block shape (128, 128) is supported, got {block_size}") + + self.in_features = int(in_features) + self.out_features = int(out_features) + self.block_size = tuple(block_size) + self.packed_weight_f32 = nn.Parameter( + torch.empty((out_features, in_features // 4), dtype=torch.float32, device=device), + requires_grad=False, + ) + self.weight_scale_inv = nn.Parameter( + torch.empty( + ((out_features + 127) // 128, (in_features + 127) // 128), + dtype=torch.float32, + device=device, + ), + requires_grad=False, + ) + + @classmethod + def from_linear(cls, module: nn.Linear) -> "NativeBlockFP8Linear": + if module.bias is not None: + raise ValueError("Native block-FP8 linear does not support bias") + return cls( + module.in_features, + module.out_features, + device=module.weight.device, + ) + + def _apply(self, fn, recurse: bool = True): + # Preserve packed bytes and scale values across model.to(dtype=...). + probe = fn(torch.empty(0, dtype=torch.float32, device=self.packed_weight_f32.device)) + protected = {name: self._parameters.pop(name) for name in ("packed_weight_f32", "weight_scale_inv")} + try: + result = super()._apply(fn, recurse=recurse) + replacements = {} + for name, parameter in protected.items(): + if parameter.is_meta: + value = torch.empty_like(parameter, dtype=torch.float32, device=probe.device) + else: + value = parameter.to(device=probe.device, dtype=torch.float32) + replacements[name] = nn.Parameter(value, requires_grad=False) + except Exception: + self._parameters.update(protected) + raise + self._parameters.update(replacements) + return result + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + _validate_state_dict_contract( + state_dict, + prefix, + { + "packed_weight_f32": (torch.float32, tuple(self.packed_weight_f32.shape)), + "weight_scale_inv": (torch.float32, tuple(self.weight_scale_inv.shape)), + }, + ) + return super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def load_prequantized(self, weight: torch.Tensor, weight_scale_inv: torch.Tensor) -> None: + """Copy one official weight/scale pair with strict byte and shape checks.""" + + if tuple(weight.shape) != (self.out_features, self.in_features): + raise ValueError( + f"FP8 weight shape {tuple(weight.shape)} does not match ({self.out_features}, {self.in_features})" + ) + if weight_scale_inv.dtype is not torch.float32: + raise TypeError(f"weight_scale_inv must remain FP32, got {weight_scale_inv.dtype}") + if tuple(weight_scale_inv.shape) != tuple(self.weight_scale_inv.shape): + raise ValueError( + f"FP8 scale shape {tuple(weight_scale_inv.shape)} does not match {tuple(self.weight_scale_inv.shape)}" + ) + if not bool(torch.all(torch.isfinite(weight_scale_inv))): + raise ValueError("weight_scale_inv contains non-finite values") + packed = pack_fp8_as_float32(weight) + with torch.no_grad(): + self.packed_weight_f32.copy_(packed.to(self.packed_weight_f32.device)) + self.weight_scale_inv.copy_(weight_scale_inv.to(self.weight_scale_inv.device)) + + def fp8_weight(self) -> torch.Tensor: + return unpack_float32_as_fp8( + self.packed_weight_f32, + (self.out_features, self.in_features), + ) + + def enable_frozen_activation_dgrad(self) -> None: + """Admit the validated frozen-trunk activation backward. + + Idempotent per module. The forward VALUE program is byte-unchanged; + only the refusal on grad-requiring inputs is replaced by the checked + BF16 dequant-program dgrad (no wgrad, no master/cache mutation — + frozen means frozen). Only a trainable-set admission that must + backpropagate THROUGH this frozen module may call this; scoring-only + lanes keep the fail-closed refusal. + """ + + self._frozen_dgrad_admitted = True + cls = NativeBlockFP8Linear + if not cls._frozen_dgrad_engagement_logged: + cls._frozen_dgrad_engagement_logged = True + logger.info( + "Native block-FP8 frozen-trunk activation dgrad engaged: contract=%s " + "(forward bytes unchanged; dgrad = grad_output @ dequant(cache) in the " + "declared BF16 program; wgrad none; frozen bytes immutable)", + self.frozen_dgrad_contract_version, + ) + + def _frozen_activation_dgrad( + self, + grad_output: torch.Tensor, + *, + output_range: tuple[int, int], + input_range: tuple[int, int], + ) -> torch.Tensor: + """dgrad through the dequantized frozen bytes in the BF16 program. + + Dequantization is the same program point the trainable composites' + surrogate uses (``block_fp8_dequantize_gkn`` -> FP32 -> one explicit + BF16 rounding); the GEMM is the declared BF16 linear backward + ``grad_output @ W``. Gated bitwise against the reference + dequant-matmul autograd (tests/ops/test_block_fp8_frozen_dgrad.py). + """ + + from xorl.ops.quantize import block_fp8_dequantize_gkn # noqa: PLC0415 + + out_start, out_end = output_range + in_start, in_end = input_range + weight = self.fp8_weight()[out_start:out_end, in_start:in_end].contiguous() + scale = self.weight_scale_inv[ + out_start // 128 : (out_end + 127) // 128, + in_start // 128 : (in_end + 127) // 128, + ].contiguous() + dequantized = block_fp8_dequantize_gkn(weight, scale, 128).to(torch.float32).to(torch.bfloat16) + grad_2d = grad_output.reshape(-1, out_end - out_start).to(torch.bfloat16) + return grad_2d.matmul(dequantized).reshape(*grad_output.shape[:-1], in_end - in_start) + + @staticmethod + def _validate_partition(value: tuple[int, int] | None, size: int, name: str) -> tuple[int, int]: + if value is None: + return 0, size + start, end = value + if start < 0 or end > size or start >= end: + raise ValueError(f"Invalid native block-FP8 {name} range {value} for size {size}") + if start % 128 or (end != size and end % 128): + raise ValueError(f"Native block-FP8 {name} range {value} must follow 128-element block boundaries") + return start, end + + def forward_partition( + self, + input: torch.Tensor, + *, + output_range: tuple[int, int] | None = None, + input_range: tuple[int, int] | None = None, + ) -> torch.Tensor: + """Run an aligned output-row/input-column partition through SGLang. + + This supports callers that preserve a model-specific reduction policy + by assigning disjoint, block-aligned projection slices to ranks. + """ + + if self.packed_weight_f32.requires_grad or self.weight_scale_inv.requires_grad: + raise RuntimeError("Native block-FP8 base weights and scales must remain frozen") + grad_engaged = torch.is_grad_enabled() and input.requires_grad + if grad_engaged and not self._frozen_dgrad_admitted: + raise RuntimeError( + "Native block-FP8 phase-one forward is scoring-only; activation backward requires a validated kernel" + ) + out_start, out_end = self._validate_partition(output_range, self.out_features, "output") + in_start, in_end = self._validate_partition(input_range, self.in_features, "input") + if input.shape[-1] != in_end - in_start: + raise ValueError( + f"Native block-FP8 input width {input.shape[-1]} does not match selected range {in_start}:{in_end}" + ) + if grad_engaged: + # Same slicing + same kernel inside the autograd boundary; the + # value bytes are identical to the scoring-only path below. + return _NativeBlockFP8FrozenDgradFunction.apply(input, self, out_start, out_end, in_start, in_end) + weight = self.fp8_weight()[out_start:out_end, in_start:in_end].contiguous() + scale = self.weight_scale_inv[ + out_start // 128 : (out_end + 127) // 128, + in_start // 128 : (in_end + 127) // 128, + ].contiguous() + output = _sglang_native_block_fp8_linear_value( + input, + weight, + scale, + block_size=self.block_size, + ) + return output + + def forward( + self, + input: torch.Tensor | None = None, + *, + return_dequantized_weight: bool = False, + output_range: tuple[int, int] | None = None, + input_range: tuple[int, int] | None = None, + ) -> torch.Tensor: + if return_dequantized_weight: + if input is not None or output_range is not None or input_range is not None: + raise ValueError("Native block-FP8 weight materialization does not accept activation or range inputs") + if self.packed_weight_f32.device.type != "cuda": + raise RuntimeError("Native block-FP8 weight materialization requires CUDA") + from sglang.srt.layers.quantization.fp8_utils import block_quant_dequant # noqa: PLC0415 + + return block_quant_dequant( + self.fp8_weight(), + self.weight_scale_inv, + list(self.block_size), + torch.bfloat16, + ) + if input is None: + raise ValueError("Native block-FP8 linear forward requires an activation input") + return self.forward_partition(input, output_range=output_range, input_range=input_range) + + +__all__ = [ + "NATIVE_BLOCK_FP8_CONTRACT_VERSION", + "NATIVE_BLOCK_FP8_FROZEN_DGRAD_CONTRACT_VERSION", + "NativeBlockFP8Linear", + "pack_fp8_as_float32", + "unpack_float32_as_fp8", + "validate_native_fp8_dcp_checkpoint", + "validate_native_fp8_state_metadata", +] diff --git a/src/xorl/ops/exact/canonical_moe_cast.py b/src/xorl/ops/exact/canonical_moe_cast.py new file mode 100644 index 00000000..f58e8b56 --- /dev/null +++ b/src/xorl/ops/exact/canonical_moe_cast.py @@ -0,0 +1,123 @@ +"""Exact final cast for CUDA canonical-MoE FP64 accumulators.""" + +from __future__ import annotations + +import torch + + +try: # Triton is optional in CPU-only development environments. + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: # pragma: no cover - exercised only without Triton + _TRITON_AVAILABLE = False + + +def validate_canonical_moe_fp64_cast_input( + value: torch.Tensor, + output_dtype: torch.dtype, +) -> None: + if value.dtype is not torch.float64: + raise TypeError(f"Canonical MoE direct cast requires FP64 input, got {value.dtype}") + if output_dtype not in (torch.bfloat16, torch.float16): + raise TypeError(f"Canonical MoE direct cast requires BF16 or FP16 output, got {output_dtype}") + if not value.is_cuda: + raise ValueError("Canonical MoE direct cast requires a CUDA tensor") + if not value.is_contiguous(): + raise ValueError("Canonical MoE direct cast requires contiguous input") + + +if _TRITON_AVAILABLE: + + @triton.jit + def _canonical_moe_fp64_to_lowp_rne_kernel( + input_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, + ): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + value = tl.load(input_ptr + offsets, mask=mask, other=0.0) + # Store the FP64 register directly through the BF16 pointer. An + # explicit FP32 conversion here would introduce a double-rounding + # boundary and can select the opposite BF16 neighbor at a midpoint. + tl.store(output_ptr + offsets, value, mask=mask) + + +def _canonical_moe_fp64_to_lowp_cuda( + value: torch.Tensor, + output_dtype: torch.dtype, +) -> torch.Tensor: + if not _TRITON_AVAILABLE: # pragma: no cover - guarded by CUDA environments + raise RuntimeError("Canonical MoE CUDA direct cast requires Triton") + output = torch.empty_like(value, dtype=output_dtype) + if output.numel() == 0: + return output + block_size = 256 + _canonical_moe_fp64_to_lowp_rne_kernel[(triton.cdiv(output.numel(), block_size),)]( + value, + output, + output.numel(), + BLOCK_SIZE=block_size, + num_warps=4, + ) + return output + + +@torch.library.custom_op("xorl_k3::_canonical_moe_fp64_to_bf16_rne", mutates_args=()) +def _canonical_moe_fp64_to_bf16_rne_op(value: torch.Tensor) -> torch.Tensor: + """Round FP64 directly to BF16.""" + validate_canonical_moe_fp64_cast_input(value, torch.bfloat16) + return _canonical_moe_fp64_to_lowp_cuda(value, torch.bfloat16) + + +@_canonical_moe_fp64_to_bf16_rne_op.register_fake +def _canonical_moe_fp64_to_bf16_rne_fake(value: torch.Tensor) -> torch.Tensor: + return torch.empty_like(value, dtype=torch.bfloat16) + + +@torch.library.custom_op("xorl_k3::_canonical_moe_fp64_to_fp16_rne", mutates_args=()) +def _canonical_moe_fp64_to_fp16_rne_op(value: torch.Tensor) -> torch.Tensor: + """Round FP64 directly to FP16.""" + validate_canonical_moe_fp64_cast_input(value, torch.float16) + return _canonical_moe_fp64_to_lowp_cuda(value, torch.float16) + + +@_canonical_moe_fp64_to_fp16_rne_op.register_fake +def _canonical_moe_fp64_to_fp16_rne_fake(value: torch.Tensor) -> torch.Tensor: + return torch.empty_like(value, dtype=torch.float16) + + +def _canonical_moe_fp64_to_lowp_backward(ctx, grad_output: torch.Tensor): + del ctx + # Match torch's straight-through cast derivative. The surrounding FP64 + # tree then distributes this FP64 gradient to every contributor before its + # original BF16 leaf cast handles the final backward conversion. + return grad_output.to(torch.float64) + + +_canonical_moe_fp64_to_bf16_rne_op.register_autograd( + _canonical_moe_fp64_to_lowp_backward, +) +_canonical_moe_fp64_to_fp16_rne_op.register_autograd( + _canonical_moe_fp64_to_lowp_backward, +) + + +def canonical_moe_fp64_to_lowp_rne( + value: torch.Tensor, + output_dtype: torch.dtype, +) -> torch.Tensor: + """Validate and directly round FP64 to the canonical transport dtype.""" + validate_canonical_moe_fp64_cast_input(value, output_dtype) + if output_dtype is torch.bfloat16: + return _canonical_moe_fp64_to_bf16_rne_op(value) + return _canonical_moe_fp64_to_fp16_rne_op(value) + + +__all__ = [ + "canonical_moe_fp64_to_lowp_rne", + "validate_canonical_moe_fp64_cast_input", +] diff --git a/src/xorl/ops/exact/canonical_moe_leaf.py b/src/xorl/ops/exact/canonical_moe_leaf.py new file mode 100644 index 00000000..1165471a --- /dev/null +++ b/src/xorl/ops/exact/canonical_moe_leaf.py @@ -0,0 +1,163 @@ +"""Compile-stable one-round arithmetic for canonical MoE contributor leaves.""" + +from __future__ import annotations + +import torch + + +try: # Triton is optional in CPU-only development environments. + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: # pragma: no cover - exercised only without Triton + _TRITON_AVAILABLE = False + + +def validate_canonical_moe_leaf_operands( + shared: torch.Tensor, + routed: torch.Tensor, +) -> None: + if shared.shape != routed.shape: + raise ValueError( + f"Canonical MoE leaf operands must have equal shapes, got {tuple(shared.shape)} and {tuple(routed.shape)}" + ) + if shared.dtype != routed.dtype or shared.dtype not in ( + torch.bfloat16, + torch.float16, + ): + raise TypeError( + f"Canonical MoE leaf operands must have the same BF16 or FP16 dtype, got {shared.dtype} and {routed.dtype}" + ) + if shared.device != routed.device: + raise ValueError(f"Canonical MoE leaf operands must share a device, got {shared.device} and {routed.device}") + if shared.device.type not in ("cpu", "cuda"): + raise ValueError(f"Canonical MoE leaf operands must be CPU or CUDA tensors, got {shared.device.type}") + if not shared.is_contiguous() or not routed.is_contiguous(): + raise ValueError("Canonical MoE leaf operands must be contiguous") + + +if _TRITON_AVAILABLE: + + @triton.jit + def _fp32_fma_rn(multiplicand, multiplier, addend): + return tl.inline_asm_elementwise( + asm="fma.rn.f32 $0, $1, $2, $3;", + constraints="=f,f,f,f", + args=(multiplicand, multiplier, addend), + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + @triton.jit + def _canonical_moe_leaf_fp32_kernel( + shared_ptr, + routed_ptr, + output_ptr, + n_elements, + routed_scale, + BLOCK_SIZE: tl.constexpr, + ): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + shared = tl.load(shared_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + routed = tl.load(routed_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + leaf = _fp32_fma_rn(routed, routed_scale.to(tl.float32), shared) + # The accumulator stays in a register. This store is the sole cast to + # the declared BF16/FP16 transport dtype. + tl.store(output_ptr + offsets, leaf, mask=mask) + + +def _canonical_moe_leaf_cuda( + shared: torch.Tensor, + routed: torch.Tensor, + routed_scale: float, +) -> torch.Tensor: + if not _TRITON_AVAILABLE: # pragma: no cover - guarded by CUDA environments + raise RuntimeError("Canonical MoE CUDA leaf construction requires Triton") + output = torch.empty_like(shared) + if output.numel() == 0: + return output + block_size = 256 + _canonical_moe_leaf_fp32_kernel[(triton.cdiv(output.numel(), block_size),)]( + shared, + routed, + output, + output.numel(), + float(routed_scale), + BLOCK_SIZE=block_size, + num_warps=4, + ) + return output + + +def _canonical_moe_leaf_cpu( + shared: torch.Tensor, + routed: torch.Tensor, + routed_scale: float, +) -> torch.Tensor: + # BF16/FP16 operands and an FP32 scale have at most 35 product bits, so the + # product is exact in FP64. The FP64 add followed by one FP32 round is the + # independent CPU reference (apart from a theoretical double-rounding + # corner); the CUDA identity tests pin the production FMA bytes. The final + # conversion is the sole transport cast. + scale_fp32 = torch.tensor(routed_scale, dtype=torch.float32) + pre_cast = shared.double() + routed.double() * scale_fp32.double() + return pre_cast.float().to(shared.dtype) + + +@torch.library.custom_op("xorl_k3::_canonical_moe_leaf_fp32_v1", mutates_args=()) +def _canonical_moe_leaf_fp32_v1_op( + shared: torch.Tensor, + routed: torch.Tensor, + routed_scale: float, +) -> torch.Tensor: + """Opaque one-FMA leaf primitive; output is the transport dtype.""" + validate_canonical_moe_leaf_operands(shared, routed) + if shared.is_cuda: + return _canonical_moe_leaf_cuda(shared, routed, routed_scale) + return _canonical_moe_leaf_cpu(shared, routed, routed_scale) + + +@_canonical_moe_leaf_fp32_v1_op.register_fake +def _canonical_moe_leaf_fp32_v1_fake( + shared: torch.Tensor, + routed: torch.Tensor, + routed_scale: float, +) -> torch.Tensor: + del routed, routed_scale + return torch.empty_like(shared) + + +def _canonical_moe_leaf_setup_context(ctx, inputs, output) -> None: + del output + shared, _routed, routed_scale = inputs + ctx.input_dtype = shared.dtype + ctx.routed_scale = float(routed_scale) + + +def _canonical_moe_leaf_backward(ctx, grad_output: torch.Tensor): + grad_fp32 = grad_output.float() + grad_shared = grad_fp32.to(ctx.input_dtype) + grad_routed = (grad_fp32 * ctx.routed_scale).to(ctx.input_dtype) + return grad_shared, grad_routed, None + + +_canonical_moe_leaf_fp32_v1_op.register_autograd( + _canonical_moe_leaf_backward, + setup_context=_canonical_moe_leaf_setup_context, +) + + +def canonical_moe_leaf_fp32_v1_op( + shared: torch.Tensor, + routed: torch.Tensor, + routed_scale: float, +) -> torch.Tensor: + """Validate and execute the compile-opaque canonical contributor leaf.""" + validate_canonical_moe_leaf_operands(shared, routed) + return _canonical_moe_leaf_fp32_v1_op(shared, routed, float(routed_scale)) + + +__all__ = ["canonical_moe_leaf_fp32_v1_op", "validate_canonical_moe_leaf_operands"] diff --git a/src/xorl/ops/exact/families_v2.py b/src/xorl/ops/exact/families_v2.py new file mode 100644 index 00000000..82cadc9f --- /dev/null +++ b/src/xorl/ops/exact/families_v2.py @@ -0,0 +1,3 @@ +"""Alias for :mod:`xorl.ops.bi_families_v2` (sha256-gated; cannot move).""" + +from xorl.ops.bi_families_v2 import * # noqa: F401,F403 diff --git a/src/xorl/ops/exact/fused_silu_and_mul.py b/src/xorl/ops/exact/fused_silu_and_mul.py new file mode 100644 index 00000000..6122c025 --- /dev/null +++ b/src/xorl/ops/exact/fused_silu_and_mul.py @@ -0,0 +1,359 @@ +"""Fused SiLU-and-multiply (SwiGLU activation) using Triton kernels. + +Computes: output = SiLU(input[:, :N]) * input[:, N:] +Used by both dense MLP (SwiGLU) and MoE expert layers. + +Two programs live here, deliberately: + +- ``fused_silu_and_mul`` — the historical TWO-ROUND program (SiLU result + rounded to the input dtype, then multiplied). Every pre-existing caller + keeps these exact bytes; nothing off the exact-contract path changes. +- ``exact_fp32_silu_and_mul`` — the ONE-ROUND FP32 program (SiLU and + multiply in fp32, single rounding), byte-paired with serving's + ``fp32_silu_and_mul`` (xorl-sglang f10b907d8). Selected ONLY by the + exact-contract dispatch (``_exact_one_round_swiglu``). +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _silu_and_mul_kernel( + input_ptr, + output_ptr, + N: tl.constexpr, # intermediate_size (half of input dim) + BLOCK_SIZE: tl.constexpr, +): + """Fused SiLU activation and element-wise multiplication. + + Computes: output = SiLU(input[:, :N]) * input[:, N:] + + Args: + input_ptr: Input tensor of shape [num_tokens, 2*N] + output_ptr: Output tensor of shape [num_tokens, N] + N: intermediate_size (half of input dimension) + BLOCK_SIZE: Block size for processing + """ + row_idx = tl.program_id(0) + + # Process in blocks along the N dimension + for block_start in range(0, N, BLOCK_SIZE): + col_offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < N + + # Load gate (first half) and up (second half) + gate_ptr = input_ptr + row_idx * 2 * N + col_offsets + up_ptr = input_ptr + row_idx * 2 * N + N + col_offsets + + gate = tl.load(gate_ptr, mask=mask, other=0.0) + up = tl.load(up_ptr, mask=mask, other=0.0) + + # Compute SiLU(gate) * up + gate_f32 = gate.to(tl.float32) + silu_gate = gate_f32 * tl.sigmoid(gate_f32) + result = silu_gate.to(gate.dtype) * up + + # Store result + out_ptr = output_ptr + row_idx * N + col_offsets + tl.store(out_ptr, result, mask=mask) + + +def silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: + """Fused SiLU activation and element-wise multiplication. + + Computes: output = SiLU(input[:, :N]) * input[:, N:] + where N = input.shape[-1] // 2 + + Args: + input_tensor: Input tensor of shape [..., 2*N] + + Returns: + Output tensor of shape [..., N] + """ + assert input_tensor.shape[-1] % 2 == 0, "Last dimension must be even" + original_shape = input_tensor.shape + input_2d = input_tensor.view(-1, original_shape[-1]) + + num_tokens = input_2d.shape[0] + N = input_2d.shape[1] // 2 + + output = torch.empty( + (num_tokens, N), + dtype=input_tensor.dtype, + device=input_tensor.device, + ) + + BLOCK_SIZE = 1024 + grid = (num_tokens,) + + _silu_and_mul_kernel[grid]( + input_2d, + output, + N, + BLOCK_SIZE, + ) + + # Reshape to match input shape (except last dim is halved) + output_shape = list(original_shape) + output_shape[-1] = N + return output.view(output_shape) + + +@triton.jit +def _silu_and_mul_backward_kernel( + grad_output_ptr, + input_ptr, + grad_input_ptr, + N: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Backward pass for fused SiLU and multiply. + + Given: y = SiLU(gate) * up + Computes: + d_gate = grad_output * up * SiLU_grad(gate) + d_up = grad_output * SiLU(gate) + """ + row_idx = tl.program_id(0) + + for block_start in range(0, N, BLOCK_SIZE): + col_offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < N + + # Load grad_output + grad_out = tl.load(grad_output_ptr + row_idx * N + col_offsets, mask=mask, other=0.0) + + # Load gate and up + gate = tl.load(input_ptr + row_idx * 2 * N + col_offsets, mask=mask, other=0.0) + up = tl.load(input_ptr + row_idx * 2 * N + N + col_offsets, mask=mask, other=0.0) + + # Compute SiLU and its gradient + gate_f32 = gate.to(tl.float32) + sigmoid_gate = tl.sigmoid(gate_f32) + silu_gate = gate_f32 * sigmoid_gate + silu_grad = sigmoid_gate + gate_f32 * sigmoid_gate * (1.0 - sigmoid_gate) + + # Compute gradients + grad_out_f32 = grad_out.to(tl.float32) + up_f32 = up.to(tl.float32) + + d_gate = grad_out_f32 * up_f32 * silu_grad + d_up = grad_out_f32 * silu_gate + + # Store gradients + tl.store(grad_input_ptr + row_idx * 2 * N + col_offsets, d_gate.to(gate.dtype), mask=mask) + tl.store(grad_input_ptr + row_idx * 2 * N + N + col_offsets, d_up.to(up.dtype), mask=mask) + + +def silu_and_mul_backward(grad_output: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + """Backward pass for fused SiLU and multiply. + + Args: + grad_output: Gradient of output, shape [..., N] + input_tensor: Original input, shape [..., 2*N] + + Returns: + Gradient of input, shape [..., 2*N] + """ + original_shape = input_tensor.shape + input_2d = input_tensor.view(-1, original_shape[-1]) + grad_output_2d = grad_output.view(-1, grad_output.shape[-1]) + + num_tokens = input_2d.shape[0] + N = input_2d.shape[1] // 2 + + grad_input = torch.empty_like(input_2d) + + BLOCK_SIZE = 1024 + grid = (num_tokens,) + + _silu_and_mul_backward_kernel[grid]( + grad_output_2d, + input_2d, + grad_input, + N, + BLOCK_SIZE, + ) + + return grad_input.view(original_shape) + + +class SiluAndMulFunction(torch.autograd.Function): + """Autograd function for fused SiLU and multiply.""" + + @staticmethod + def forward(ctx, input_tensor: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(input_tensor) + return silu_and_mul(input_tensor) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: + (input_tensor,) = ctx.saved_tensors + return silu_and_mul_backward(grad_output, input_tensor) + + +def fused_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: + """Fused SiLU and multiply with autograd support. + + Args: + input_tensor: Input tensor of shape [..., 2*N] + + Returns: + Output tensor of shape [..., N] + """ + return SiluAndMulFunction.apply(input_tensor) + + +# --------------------------------------------------------------------------- +# Exact-contract one-round FP32 SwiGLU (serving-paired program). +# Byte-contract: bitwise-equal to serving's fp32_silu_and_mul +# (xorl-sglang f10b907d8). Forward bits are the contract; backward is stock +# numerics (analytic Triton backward / native autograd on the fallback). +# --------------------------------------------------------------------------- + + +def _use_fp32_fused_swiglu(input_tensor: torch.Tensor) -> bool: + """Admit the shared tiled realization on contiguous Hopper inputs.""" + if not input_tensor.is_cuda or not input_tensor.is_contiguous(): + return False + if input_tensor.dtype not in (torch.bfloat16, torch.float16): + return False + major, minor = torch.cuda.get_device_capability(input_tensor.device) + if (major, minor) != (9, 0): + return False + return True + + +def _fp32_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: + """Keep the SiLU result in FP32 through the multiply and round once. + + Native-op realization: differentiable, used off-Hopper/CPU.""" + assert input_tensor.shape[-1] % 2 == 0, "Last dimension must be even" + split = input_tensor.shape[-1] // 2 + gate = input_tensor[..., :split].float() + up = input_tensor[..., split:].float() + return (torch.nn.functional.silu(gate) * up).to(input_tensor.dtype) + + +@triton.jit +def _exact_fp32_silu_and_mul_kernel( + input_ptr, + output_ptr, + N: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row_idx = tl.program_id(0) + block_start = tl.program_id(1) * BLOCK_SIZE + col_offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < N + + gate_ptr = input_ptr + row_idx * 2 * N + col_offsets + up_ptr = input_ptr + row_idx * 2 * N + N + col_offsets + gate = tl.load(gate_ptr, mask=mask, other=0.0) + up = tl.load(up_ptr, mask=mask, other=0.0) + + gate_f32 = gate.to(tl.float32) + silu_gate = gate_f32 * tl.sigmoid(gate_f32) + result = silu_gate * up.to(tl.float32) + tl.store(output_ptr + row_idx * N + col_offsets, result, mask=mask) + + +def _exact_fp32_silu_and_mul_forward(input_tensor: torch.Tensor) -> torch.Tensor: + original_shape = input_tensor.shape + input_2d = input_tensor.view(-1, original_shape[-1]) + num_tokens = input_2d.shape[0] + N = input_2d.shape[1] // 2 + output = torch.empty( + (num_tokens, N), + dtype=input_tensor.dtype, + device=input_tensor.device, + ) + BLOCK_SIZE = min(1024, triton.next_power_of_2(N)) + grid = (num_tokens, triton.cdiv(N, BLOCK_SIZE)) + _exact_fp32_silu_and_mul_kernel[grid]( + input_2d, + output, + N, + BLOCK_SIZE, + ) + output_shape = list(original_shape) + output_shape[-1] = N + return output.view(output_shape) + + +@triton.jit +def _exact_fp32_silu_and_mul_backward_kernel( + grad_output_ptr, + input_ptr, + grad_input_ptr, + N: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Analytic backward for y = SiLU(gate) * up (fp32 intermediates).""" + row_idx = tl.program_id(0) + block_start = tl.program_id(1) * BLOCK_SIZE + col_offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < N + + grad_out = tl.load(grad_output_ptr + row_idx * N + col_offsets, mask=mask, other=0.0) + gate = tl.load(input_ptr + row_idx * 2 * N + col_offsets, mask=mask, other=0.0) + up = tl.load(input_ptr + row_idx * 2 * N + N + col_offsets, mask=mask, other=0.0) + + gate_f32 = gate.to(tl.float32) + sigmoid_gate = tl.sigmoid(gate_f32) + silu_gate = gate_f32 * sigmoid_gate + silu_grad = sigmoid_gate + gate_f32 * sigmoid_gate * (1.0 - sigmoid_gate) + grad_out_f32 = grad_out.to(tl.float32) + up_f32 = up.to(tl.float32) + + d_gate = grad_out_f32 * up_f32 * silu_grad + d_up = grad_out_f32 * silu_gate + tl.store(grad_input_ptr + row_idx * 2 * N + col_offsets, d_gate.to(gate.dtype), mask=mask) + tl.store(grad_input_ptr + row_idx * 2 * N + N + col_offsets, d_up.to(up.dtype), mask=mask) + + +def _exact_fp32_silu_and_mul_backward(grad_output: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + original_shape = input_tensor.shape + input_2d = input_tensor.view(-1, original_shape[-1]) + grad_output_2d = grad_output.view(-1, grad_output.shape[-1]) + num_tokens = input_2d.shape[0] + N = input_2d.shape[1] // 2 + grad_input = torch.empty_like(input_2d) + BLOCK_SIZE = min(1024, triton.next_power_of_2(N)) + grid = (num_tokens, triton.cdiv(N, BLOCK_SIZE)) + _exact_fp32_silu_and_mul_backward_kernel[grid]( + grad_output_2d, + input_2d, + grad_input, + N, + BLOCK_SIZE, + ) + return grad_input.view(original_shape) + + +class ExactFp32SiluAndMulFunction(torch.autograd.Function): + """Autograd function for the exact one-round FP32 SwiGLU.""" + + @staticmethod + def forward(ctx, input_tensor: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(input_tensor) + return _exact_fp32_silu_and_mul_forward(input_tensor) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: + (input_tensor,) = ctx.saved_tensors + return _exact_fp32_silu_and_mul_backward(grad_output, input_tensor) + + +def exact_fp32_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: + """One-round FP32 SwiGLU with autograd support (exact-contract callers only). + + Byte-paired with serving's fp32_silu_and_mul; see the module docstring + for the scoping rationale.""" + if input_tensor.shape[-1] % 2: + raise ValueError("SwiGLU input width must be even") + if not _use_fp32_fused_swiglu(input_tensor): + return _fp32_silu_and_mul(input_tensor) + return ExactFp32SiluAndMulFunction.apply(input_tensor) diff --git a/src/xorl/ops/exact/kernel_config_pin.py b/src/xorl/ops/exact/kernel_config_pin.py new file mode 100644 index 00000000..4734a935 --- /dev/null +++ b/src/xorl/ops/exact/kernel_config_pin.py @@ -0,0 +1,182 @@ +"""First-class kernel/toolchain pinning for exact contracts. + +Byte-exact programs are TOOLCHAIN-SCOPED claims: different Triton or FA4 +builds can compile the same source into different arithmetic. Autotune configs +are likewise per-process unless pinned: Triton's ``cache_results`` replays +tuned configs from the cache directory, so all ranks and the qualification +oracle must share one seeded cache. Each rank receives a separate clone to +avoid concurrent writes to the shared seed. + +The admission contract is mechanical: + +- A qualification run SEEDS a pin directory: its triton cache plus a + toolchain manifest (torch/triton/flash-attn versions). +- Admission (``pin_exact_kernel_configs``) FAILS CLOSED unless + ``XORL_EXACT_KERNEL_CONFIG_DIR`` names a seeded pin directory whose + manifest matches the running toolchain, then points this process's + ``TRITON_CACHE_DIR`` at a per-rank clone of the seeded cache. It must run + BEFORE the first kernel compilation (admission time satisfies this). +""" + +from __future__ import annotations + +import json +import logging +import os +import shutil + +import torch +import triton + + +logger = logging.getLogger("xorl.kernel_config_pin") + +PIN_DIR_ENV = "XORL_EXACT_KERNEL_CONFIG_DIR" +MANIFEST_NAME = "toolchain_manifest.json" +CACHE_SUBDIR = "triton-cache" + + +class KernelConfigPinError(RuntimeError): + """The kernel/toolchain pin is missing or violated. Fail closed.""" + + +# Ownership sentinel: this module only ever deletes directories it created +# itself (the marker IS the authorization). Env-var-fed paths never reach +# rmtree without it. +OWNED_SENTINEL = ".xorl-kernel-pin-owned" + + +def _mark_owned(path: str) -> None: + with open(os.path.join(path, OWNED_SENTINEL), "w") as f: + f.write("created by xorl.ops.exact.kernel_config_pin; safe for it to replace\n") + + +def _rmtree_owned(path: str) -> None: + """Delete `path` only if this module created it (sentinel present).""" + if not os.path.isdir(path): + return + if not os.path.isfile(os.path.join(path, OWNED_SENTINEL)): + raise KernelConfigPinError( + f"refusing to delete {path!r}: it lacks the ownership sentinel " + f"{OWNED_SENTINEL!r} and was not created by this module. Remove or " + "relocate it manually if it is stale.", + ) + shutil.rmtree(path) + + +def _runtime_fingerprint() -> dict: + # Use the distribution version rather than only the module attribute: + # distinct flash-attn wheel builds can share a torch/triton fingerprint, + # while some builds do not expose a useful ``flash_attn.__version__``. + from importlib import metadata # noqa: PLC0415 + + fa = "unavailable" + for dist in ("flash-attn-4", "flash_attn_4", "flash-attn", "flash_attn"): + try: + fa = metadata.version(dist) + break + except metadata.PackageNotFoundError: + continue + if fa == "unavailable": + try: + import flash_attn # noqa: PLC0415 + + fa = getattr(flash_attn, "__version__", "unavailable") + except Exception: # pragma: no cover - build dependent + pass + return { + "torch": torch.__version__, + "triton": triton.__version__, + "flash_attn": fa, + "cuda": torch.version.cuda or "none", + } + + +def seed_exact_kernel_config_pin(pin_dir: str, *, source_cache: str | None = None) -> dict: + """Create/refresh a pin directory from the CURRENT runtime. + + Called by qualification runs (e.g. the fixture oracle phase) after their + kernels have been tuned. Copies `source_cache` (default: the active + TRITON_CACHE_DIR or ~/.triton/cache) into the pin and writes the + toolchain manifest. + """ + fingerprint = _runtime_fingerprint() + pin_dir = os.path.realpath(pin_dir) + parent = os.path.dirname(pin_dir) + if not os.path.isdir(parent): + raise KernelConfigPinError( + f"pin directory parent {parent!r} does not exist; refusing to create a pin at an implausible location", + ) + os.makedirs(pin_dir, exist_ok=True) + cache_src = os.path.realpath( + source_cache or os.environ.get("TRITON_CACHE_DIR", os.path.expanduser("~/.triton/cache")) + ) + cache_dst = os.path.join(pin_dir, CACHE_SUBDIR) + if os.path.commonpath([cache_src, cache_dst]) == cache_src: + raise KernelConfigPinError( + f"seed source cache {cache_src!r} contains the pin destination {cache_dst!r}; " + "copying would recurse into itself", + ) + _rmtree_owned(cache_dst) + if os.path.isdir(cache_src): + shutil.copytree(cache_src, cache_dst) + else: + os.makedirs(cache_dst, exist_ok=True) + _mark_owned(cache_dst) + with open(os.path.join(pin_dir, MANIFEST_NAME), "w") as f: + json.dump(fingerprint, f, indent=2) + logger.info("exact kernel-config pin seeded at %s: %s", pin_dir, fingerprint) + return fingerprint + + +def pin_exact_kernel_configs(*, rank: int | None = None) -> str: + """Admission-time pin. Returns the per-rank TRITON_CACHE_DIR it installed. + + Fail-closed on: env unset, pin dir or manifest missing, toolchain + fingerprint mismatch. Engagement-logged once per process. + """ + pin_dir = os.environ.get(PIN_DIR_ENV) + if not pin_dir: + raise KernelConfigPinError( + f"{PIN_DIR_ENV} is not set. Exact hybrid-Ulysses admission requires a seeded " + "kernel/toolchain pin directory (seed_exact_kernel_config_pin from the " + "qualification run). Byte claims are toolchain-scoped; refusing to run unpinned.", + ) + # The env var feeds filesystem mutations below: resolve it and refuse + # anything that is not an existing, seeded pin directory. + pin_dir = os.path.realpath(pin_dir) + if not os.path.isdir(pin_dir): + raise KernelConfigPinError( + f"{PIN_DIR_ENV}={pin_dir!r} is not an existing directory. Fail closed.", + ) + manifest_path = os.path.join(pin_dir, MANIFEST_NAME) + if not os.path.isfile(manifest_path): + raise KernelConfigPinError( + f"{PIN_DIR_ENV}={pin_dir!r} has no {MANIFEST_NAME}; the pin directory was never " + "seeded by a qualification run. Fail closed.", + ) + with open(manifest_path) as f: + pinned = json.load(f) + running = _runtime_fingerprint() + mismatches = {k: (pinned.get(k), running[k]) for k in running if pinned.get(k) != running[k]} + if mismatches: + raise KernelConfigPinError( + "Toolchain fingerprint mismatch against the kernel-config pin — the byte " + f"qualification does not transfer: {mismatches}. Re-qualify or restore the " + "pinned environment.", + ) + if rank is None: + rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) + clone = os.path.join(pin_dir, "clones", f"rank{rank}") + _rmtree_owned(clone) + shutil.copytree(os.path.join(pin_dir, CACHE_SUBDIR), clone) + _mark_owned(clone) + os.environ["TRITON_CACHE_DIR"] = clone + logger.info( + "exact kernel-config pin engaged: %s (rank %d clone %s, toolchain %s)", + pin_dir, + rank, + clone, + running, + ) + return clone diff --git a/src/xorl/ops/exact/rope_class_b.py b/src/xorl/ops/exact/rope_class_b.py new file mode 100644 index 00000000..7471737d --- /dev/null +++ b/src/xorl/ops/exact/rope_class_b.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +# Class-B RoPE numerics, vendored from the matching sampler implementation so the trainer +# takes no runtime dependency on SGLang. +# +# Two RoPE numerics classes exist across the trainer/sampler pair: +# Class A -- per-op bf16 rounding (8 rounding points). The trainer's +# ``_naive_apply_rotary_pos_emb``, SGLang's eager ``apply_rotary_emb`` and its +# ``bi_fused_native`` triton kernel all agree bitwise here. +# Class B -- one fp32 chain with a single final round. SGLang's stock fused CUDA rope +# (``jit_kernel/csrc/elementwise/rope.cuh``) and the ``torch.compile`` RoPE path +# SGLang's RL lane selects agree bitwise here. +# +# ``_rotary_emb`` is the body of SGLang's ``rotary_embedding/utils.py::apply_rotary_emb``, +# vendored verbatim. What selects Class B is not the function but its input: callers hand it +# the fp32 cos/sin table serving builds, rather than the bf16 one the trainer's eager path +# rounds first. Compiled, it is bitwise equal to the stock CUDA kernel across every variant and +# shape the target models use -- that equality is a gate, not an assumption; see +# ``gate_phase1.py``. (The ``.to(x.dtype)`` below is load-bearing for the output dtype and is +# round-tripped away by inductor for the arithmetic; dropping it yields fp32 outputs.) +# +# A hand-written triton reimplementation of the CUDA kernel was tried first and rejected: it +# lands within one bf16 ULP on ~7e-4 of elements but is not bitwise equal, and no FMA +# association (plain, ``fma(x,cos,-(y*sin))``, ``fma(-y,sin,x*cos)``, or fully uncontracted +# ``mul.rn``) closes the gap. Bitwise is the requirement, so the compiled fp32 expression is +# what ships. +import logging +import os + +import torch + + +logger = logging.getLogger(__name__) + +# torch.compile is load-bearing here, not an optimisation: the EAGER form of this same +# expression is Class A. Dynamo's default recompile budget is 8, far under what a real model +# needs (separate q and k shapes x rollout lengths x forward and backward), and on exhausting +# it dynamo silently falls back to eager -- reverting those layers to Class A with no error and +# no log line. That is indistinguishable from a numerics bug downstream, so pin the budget and +# make exhaustion loud. +_ROPE_CLASS_B_RECOMPILE_LIMIT = 2048 +_ROPE_CLASS_B_ACCUMULATED_LIMIT = 8192 + + +def _pin_compile_budget() -> None: + cfg = torch._dynamo.config + cfg.recompile_limit = max(getattr(cfg, "recompile_limit", 0), _ROPE_CLASS_B_RECOMPILE_LIMIT) + cfg.accumulated_recompile_limit = max( + getattr(cfg, "accumulated_recompile_limit", 0), _ROPE_CLASS_B_ACCUMULATED_LIMIT + ) + # A silent eager fallback breaks the zero-K3 contract, so fail instead. Opt out with + # XORL_ROPE_CLASS_B_ALLOW_FALLBACK=1 for throughput experiments outside the contract. + if os.environ.get("XORL_ROPE_CLASS_B_ALLOW_FALLBACK") != "1" and hasattr(cfg, "fail_on_recompile_limit_hit"): + cfg.fail_on_recompile_limit_hit = True + logger.info( + "rope Class B: recompile_limit=%s accumulated=%s fail_on_limit=%s", + cfg.recompile_limit, + cfg.accumulated_recompile_limit, + getattr(cfg, "fail_on_recompile_limit_hit", None), + ) + + +_pin_compile_budget() + + +def _rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, is_neox: bool) -> torch.Tensor: + """SGLang's ``apply_rotary_emb``, vendored. + + x: ``[num_tokens, num_heads, rotary_dim]``; cos/sin: ``[num_tokens, rotary_dim // 2]`` fp32. + """ + cos = cos.unsqueeze(-2).to(x.dtype) + sin = sin.unsqueeze(-2).to(x.dtype) + if is_neox: + x1, x2 = torch.chunk(x, 2, dim=-1) + else: + x1 = x[..., ::2] + x2 = x[..., 1::2] + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + if is_neox: + return torch.cat((o1, o2), dim=-1) + return torch.stack((o1, o2), dim=-1).flatten(-2) + + +_rotary_emb_compiled = torch.compile(dynamic=True)(_rotary_emb) + + +class _ClassBRoPE(torch.autograd.Function): + """Class-B rope with a hand-written backward. + + RoPE is a block rotation, so its adjoint is the same rotation driven by a negated sin. + Supplying it explicitly keeps the backward one fused pass and lands it closer to the fp64 + adjoint than differentiating the forward expression. + """ + + @staticmethod + def forward(ctx, q_rot, k_rot, cos, sin, is_neox): + ctx.save_for_backward(cos, sin) + ctx.is_neox = is_neox + return ( + _rotary_emb_compiled(q_rot, cos, sin, is_neox), + _rotary_emb_compiled(k_rot, cos, sin, is_neox), + ) + + @staticmethod + def backward(ctx, grad_q, grad_k): + cos, sin = ctx.saved_tensors + neg_sin = -sin + return ( + _rotary_emb_compiled(grad_q.contiguous(), cos, neg_sin, ctx.is_neox), + _rotary_emb_compiled(grad_k.contiguous(), cos, neg_sin, ctx.is_neox), + None, + None, + None, + ) + + +def build_class_b_cos_sin(cos, sin, *, doubled: bool = True): + """Flatten the trainer's cos/sin into the fp32 ``[num_tokens, rotary_dim // 2]`` serving layout. + + With ``doubled`` (the default) cos/sin arrive in the ``[..., rotary_dim]`` layout the shared + ``RotaryEmbedding`` emits, where only the leading half carries distinct frequencies. GPT-OSS + emits them already halved, at ``[..., rotary_dim // 2]``. + """ + half = cos.shape[-1] // 2 if doubled else cos.shape[-1] + return cos[..., :half].reshape(-1, half).float(), sin[..., :half].reshape(-1, half).float() + + +def class_b_apply_rotary_pos_emb(q, k, cos, sin, *, interleaved: bool = False, doubled: bool = True): + """Class-B RoPE over ``[B, S, H, D]`` q/k. + + ``interleaved`` selects GPT-J pairing; the default is the half-split Neox pairing. Partial + rotary rotates the leading rotary_dim features and passes the tail through untouched. + """ + if cos.dtype != torch.float32 or sin.dtype != torch.float32: + raise RuntimeError( + "Class-B RoPE requires fp32 cos/sin at the apply boundary; " + f"got cos={cos.dtype}, sin={sin.dtype}. A mixed-precision wrapper " + "likely downcast the table in transit." + ) + cos_f, sin_f = build_class_b_cos_sin(cos, sin, doubled=doubled) + rotary_dim = 2 * cos_f.shape[-1] + batch, seq_len = q.shape[0], q.shape[1] + num_tokens = batch * seq_len + + q_rot = q[..., :rotary_dim].reshape(num_tokens, -1, rotary_dim) + k_rot = k[..., :rotary_dim].reshape(num_tokens, -1, rotary_dim) + q_out, k_out = _ClassBRoPE.apply(q_rot, k_rot, cos_f, sin_f, not interleaved) + q_out = q_out.view(batch, seq_len, -1, rotary_dim) + k_out = k_out.view(batch, seq_len, -1, rotary_dim) + + if q.shape[-1] > rotary_dim: + q_out = torch.cat((q_out, q[..., rotary_dim:]), dim=-1) + k_out = torch.cat((k_out, k[..., rotary_dim:]), dim=-1) + return q_out, k_out + + +__all__ = [ + "build_class_b_cos_sin", + "class_b_apply_rotary_pos_emb", +] diff --git a/src/xorl/ops/exact/sampling_transforms.py b/src/xorl/ops/exact/sampling_transforms.py new file mode 100644 index 00000000..65ae89c6 --- /dev/null +++ b/src/xorl/ops/exact/sampling_transforms.py @@ -0,0 +1,3 @@ +"""Alias for :mod:`xorl.ops.exact_sampling_transforms` (kept in place during #74).""" + +from xorl.ops.exact_sampling_transforms import * # noqa: F401,F403 diff --git a/src/xorl/ops/fused_silu_and_mul.py b/src/xorl/ops/fused_silu_and_mul.py index 6122c025..5ef6b591 100644 --- a/src/xorl/ops/fused_silu_and_mul.py +++ b/src/xorl/ops/fused_silu_and_mul.py @@ -1,359 +1,3 @@ -"""Fused SiLU-and-multiply (SwiGLU activation) using Triton kernels. +"""Moved to ``xorl.ops.exact.fused_silu_and_mul`` (issue #78 phase 3); compat stub.""" -Computes: output = SiLU(input[:, :N]) * input[:, N:] -Used by both dense MLP (SwiGLU) and MoE expert layers. - -Two programs live here, deliberately: - -- ``fused_silu_and_mul`` — the historical TWO-ROUND program (SiLU result - rounded to the input dtype, then multiplied). Every pre-existing caller - keeps these exact bytes; nothing off the exact-contract path changes. -- ``exact_fp32_silu_and_mul`` — the ONE-ROUND FP32 program (SiLU and - multiply in fp32, single rounding), byte-paired with serving's - ``fp32_silu_and_mul`` (xorl-sglang f10b907d8). Selected ONLY by the - exact-contract dispatch (``_exact_one_round_swiglu``). -""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _silu_and_mul_kernel( - input_ptr, - output_ptr, - N: tl.constexpr, # intermediate_size (half of input dim) - BLOCK_SIZE: tl.constexpr, -): - """Fused SiLU activation and element-wise multiplication. - - Computes: output = SiLU(input[:, :N]) * input[:, N:] - - Args: - input_ptr: Input tensor of shape [num_tokens, 2*N] - output_ptr: Output tensor of shape [num_tokens, N] - N: intermediate_size (half of input dimension) - BLOCK_SIZE: Block size for processing - """ - row_idx = tl.program_id(0) - - # Process in blocks along the N dimension - for block_start in range(0, N, BLOCK_SIZE): - col_offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = col_offsets < N - - # Load gate (first half) and up (second half) - gate_ptr = input_ptr + row_idx * 2 * N + col_offsets - up_ptr = input_ptr + row_idx * 2 * N + N + col_offsets - - gate = tl.load(gate_ptr, mask=mask, other=0.0) - up = tl.load(up_ptr, mask=mask, other=0.0) - - # Compute SiLU(gate) * up - gate_f32 = gate.to(tl.float32) - silu_gate = gate_f32 * tl.sigmoid(gate_f32) - result = silu_gate.to(gate.dtype) * up - - # Store result - out_ptr = output_ptr + row_idx * N + col_offsets - tl.store(out_ptr, result, mask=mask) - - -def silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: - """Fused SiLU activation and element-wise multiplication. - - Computes: output = SiLU(input[:, :N]) * input[:, N:] - where N = input.shape[-1] // 2 - - Args: - input_tensor: Input tensor of shape [..., 2*N] - - Returns: - Output tensor of shape [..., N] - """ - assert input_tensor.shape[-1] % 2 == 0, "Last dimension must be even" - original_shape = input_tensor.shape - input_2d = input_tensor.view(-1, original_shape[-1]) - - num_tokens = input_2d.shape[0] - N = input_2d.shape[1] // 2 - - output = torch.empty( - (num_tokens, N), - dtype=input_tensor.dtype, - device=input_tensor.device, - ) - - BLOCK_SIZE = 1024 - grid = (num_tokens,) - - _silu_and_mul_kernel[grid]( - input_2d, - output, - N, - BLOCK_SIZE, - ) - - # Reshape to match input shape (except last dim is halved) - output_shape = list(original_shape) - output_shape[-1] = N - return output.view(output_shape) - - -@triton.jit -def _silu_and_mul_backward_kernel( - grad_output_ptr, - input_ptr, - grad_input_ptr, - N: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - """Backward pass for fused SiLU and multiply. - - Given: y = SiLU(gate) * up - Computes: - d_gate = grad_output * up * SiLU_grad(gate) - d_up = grad_output * SiLU(gate) - """ - row_idx = tl.program_id(0) - - for block_start in range(0, N, BLOCK_SIZE): - col_offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = col_offsets < N - - # Load grad_output - grad_out = tl.load(grad_output_ptr + row_idx * N + col_offsets, mask=mask, other=0.0) - - # Load gate and up - gate = tl.load(input_ptr + row_idx * 2 * N + col_offsets, mask=mask, other=0.0) - up = tl.load(input_ptr + row_idx * 2 * N + N + col_offsets, mask=mask, other=0.0) - - # Compute SiLU and its gradient - gate_f32 = gate.to(tl.float32) - sigmoid_gate = tl.sigmoid(gate_f32) - silu_gate = gate_f32 * sigmoid_gate - silu_grad = sigmoid_gate + gate_f32 * sigmoid_gate * (1.0 - sigmoid_gate) - - # Compute gradients - grad_out_f32 = grad_out.to(tl.float32) - up_f32 = up.to(tl.float32) - - d_gate = grad_out_f32 * up_f32 * silu_grad - d_up = grad_out_f32 * silu_gate - - # Store gradients - tl.store(grad_input_ptr + row_idx * 2 * N + col_offsets, d_gate.to(gate.dtype), mask=mask) - tl.store(grad_input_ptr + row_idx * 2 * N + N + col_offsets, d_up.to(up.dtype), mask=mask) - - -def silu_and_mul_backward(grad_output: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: - """Backward pass for fused SiLU and multiply. - - Args: - grad_output: Gradient of output, shape [..., N] - input_tensor: Original input, shape [..., 2*N] - - Returns: - Gradient of input, shape [..., 2*N] - """ - original_shape = input_tensor.shape - input_2d = input_tensor.view(-1, original_shape[-1]) - grad_output_2d = grad_output.view(-1, grad_output.shape[-1]) - - num_tokens = input_2d.shape[0] - N = input_2d.shape[1] // 2 - - grad_input = torch.empty_like(input_2d) - - BLOCK_SIZE = 1024 - grid = (num_tokens,) - - _silu_and_mul_backward_kernel[grid]( - grad_output_2d, - input_2d, - grad_input, - N, - BLOCK_SIZE, - ) - - return grad_input.view(original_shape) - - -class SiluAndMulFunction(torch.autograd.Function): - """Autograd function for fused SiLU and multiply.""" - - @staticmethod - def forward(ctx, input_tensor: torch.Tensor) -> torch.Tensor: - ctx.save_for_backward(input_tensor) - return silu_and_mul(input_tensor) - - @staticmethod - def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: - (input_tensor,) = ctx.saved_tensors - return silu_and_mul_backward(grad_output, input_tensor) - - -def fused_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: - """Fused SiLU and multiply with autograd support. - - Args: - input_tensor: Input tensor of shape [..., 2*N] - - Returns: - Output tensor of shape [..., N] - """ - return SiluAndMulFunction.apply(input_tensor) - - -# --------------------------------------------------------------------------- -# Exact-contract one-round FP32 SwiGLU (serving-paired program). -# Byte-contract: bitwise-equal to serving's fp32_silu_and_mul -# (xorl-sglang f10b907d8). Forward bits are the contract; backward is stock -# numerics (analytic Triton backward / native autograd on the fallback). -# --------------------------------------------------------------------------- - - -def _use_fp32_fused_swiglu(input_tensor: torch.Tensor) -> bool: - """Admit the shared tiled realization on contiguous Hopper inputs.""" - if not input_tensor.is_cuda or not input_tensor.is_contiguous(): - return False - if input_tensor.dtype not in (torch.bfloat16, torch.float16): - return False - major, minor = torch.cuda.get_device_capability(input_tensor.device) - if (major, minor) != (9, 0): - return False - return True - - -def _fp32_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: - """Keep the SiLU result in FP32 through the multiply and round once. - - Native-op realization: differentiable, used off-Hopper/CPU.""" - assert input_tensor.shape[-1] % 2 == 0, "Last dimension must be even" - split = input_tensor.shape[-1] // 2 - gate = input_tensor[..., :split].float() - up = input_tensor[..., split:].float() - return (torch.nn.functional.silu(gate) * up).to(input_tensor.dtype) - - -@triton.jit -def _exact_fp32_silu_and_mul_kernel( - input_ptr, - output_ptr, - N: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - row_idx = tl.program_id(0) - block_start = tl.program_id(1) * BLOCK_SIZE - col_offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = col_offsets < N - - gate_ptr = input_ptr + row_idx * 2 * N + col_offsets - up_ptr = input_ptr + row_idx * 2 * N + N + col_offsets - gate = tl.load(gate_ptr, mask=mask, other=0.0) - up = tl.load(up_ptr, mask=mask, other=0.0) - - gate_f32 = gate.to(tl.float32) - silu_gate = gate_f32 * tl.sigmoid(gate_f32) - result = silu_gate * up.to(tl.float32) - tl.store(output_ptr + row_idx * N + col_offsets, result, mask=mask) - - -def _exact_fp32_silu_and_mul_forward(input_tensor: torch.Tensor) -> torch.Tensor: - original_shape = input_tensor.shape - input_2d = input_tensor.view(-1, original_shape[-1]) - num_tokens = input_2d.shape[0] - N = input_2d.shape[1] // 2 - output = torch.empty( - (num_tokens, N), - dtype=input_tensor.dtype, - device=input_tensor.device, - ) - BLOCK_SIZE = min(1024, triton.next_power_of_2(N)) - grid = (num_tokens, triton.cdiv(N, BLOCK_SIZE)) - _exact_fp32_silu_and_mul_kernel[grid]( - input_2d, - output, - N, - BLOCK_SIZE, - ) - output_shape = list(original_shape) - output_shape[-1] = N - return output.view(output_shape) - - -@triton.jit -def _exact_fp32_silu_and_mul_backward_kernel( - grad_output_ptr, - input_ptr, - grad_input_ptr, - N: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - """Analytic backward for y = SiLU(gate) * up (fp32 intermediates).""" - row_idx = tl.program_id(0) - block_start = tl.program_id(1) * BLOCK_SIZE - col_offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = col_offsets < N - - grad_out = tl.load(grad_output_ptr + row_idx * N + col_offsets, mask=mask, other=0.0) - gate = tl.load(input_ptr + row_idx * 2 * N + col_offsets, mask=mask, other=0.0) - up = tl.load(input_ptr + row_idx * 2 * N + N + col_offsets, mask=mask, other=0.0) - - gate_f32 = gate.to(tl.float32) - sigmoid_gate = tl.sigmoid(gate_f32) - silu_gate = gate_f32 * sigmoid_gate - silu_grad = sigmoid_gate + gate_f32 * sigmoid_gate * (1.0 - sigmoid_gate) - grad_out_f32 = grad_out.to(tl.float32) - up_f32 = up.to(tl.float32) - - d_gate = grad_out_f32 * up_f32 * silu_grad - d_up = grad_out_f32 * silu_gate - tl.store(grad_input_ptr + row_idx * 2 * N + col_offsets, d_gate.to(gate.dtype), mask=mask) - tl.store(grad_input_ptr + row_idx * 2 * N + N + col_offsets, d_up.to(up.dtype), mask=mask) - - -def _exact_fp32_silu_and_mul_backward(grad_output: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: - original_shape = input_tensor.shape - input_2d = input_tensor.view(-1, original_shape[-1]) - grad_output_2d = grad_output.view(-1, grad_output.shape[-1]) - num_tokens = input_2d.shape[0] - N = input_2d.shape[1] // 2 - grad_input = torch.empty_like(input_2d) - BLOCK_SIZE = min(1024, triton.next_power_of_2(N)) - grid = (num_tokens, triton.cdiv(N, BLOCK_SIZE)) - _exact_fp32_silu_and_mul_backward_kernel[grid]( - grad_output_2d, - input_2d, - grad_input, - N, - BLOCK_SIZE, - ) - return grad_input.view(original_shape) - - -class ExactFp32SiluAndMulFunction(torch.autograd.Function): - """Autograd function for the exact one-round FP32 SwiGLU.""" - - @staticmethod - def forward(ctx, input_tensor: torch.Tensor) -> torch.Tensor: - ctx.save_for_backward(input_tensor) - return _exact_fp32_silu_and_mul_forward(input_tensor) - - @staticmethod - def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: - (input_tensor,) = ctx.saved_tensors - return _exact_fp32_silu_and_mul_backward(grad_output, input_tensor) - - -def exact_fp32_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: - """One-round FP32 SwiGLU with autograd support (exact-contract callers only). - - Byte-paired with serving's fp32_silu_and_mul; see the module docstring - for the scoping rationale.""" - if input_tensor.shape[-1] % 2: - raise ValueError("SwiGLU input width must be even") - if not _use_fp32_fused_swiglu(input_tensor): - return _fp32_silu_and_mul(input_tensor) - return ExactFp32SiluAndMulFunction.apply(input_tensor) +from xorl.ops.exact.fused_silu_and_mul import * # noqa: F401,F403 diff --git a/src/xorl/ops/kernel_config_pin.py b/src/xorl/ops/kernel_config_pin.py index 3b512174..b83b850f 100644 --- a/src/xorl/ops/kernel_config_pin.py +++ b/src/xorl/ops/kernel_config_pin.py @@ -1,182 +1,3 @@ -"""First-class kernel/toolchain pinning for exact contracts. +"""Moved to ``xorl.ops.exact.kernel_config_pin`` (issue #78 phase 3); compat stub.""" -Byte-exact programs are TOOLCHAIN-SCOPED claims: different Triton or FA4 -builds can compile the same source into different arithmetic. Autotune configs -are likewise per-process unless pinned: Triton's ``cache_results`` replays -tuned configs from the cache directory, so all ranks and the qualification -oracle must share one seeded cache. Each rank receives a separate clone to -avoid concurrent writes to the shared seed. - -The admission contract is mechanical: - -- A qualification run SEEDS a pin directory: its triton cache plus a - toolchain manifest (torch/triton/flash-attn versions). -- Admission (``pin_exact_kernel_configs``) FAILS CLOSED unless - ``XORL_EXACT_KERNEL_CONFIG_DIR`` names a seeded pin directory whose - manifest matches the running toolchain, then points this process's - ``TRITON_CACHE_DIR`` at a per-rank clone of the seeded cache. It must run - BEFORE the first kernel compilation (admission time satisfies this). -""" - -from __future__ import annotations - -import json -import logging -import os -import shutil - -import torch -import triton - - -logger = logging.getLogger("xorl.kernel_config_pin") - -PIN_DIR_ENV = "XORL_EXACT_KERNEL_CONFIG_DIR" -MANIFEST_NAME = "toolchain_manifest.json" -CACHE_SUBDIR = "triton-cache" - - -class KernelConfigPinError(RuntimeError): - """The kernel/toolchain pin is missing or violated. Fail closed.""" - - -# Ownership sentinel: this module only ever deletes directories it created -# itself (the marker IS the authorization). Env-var-fed paths never reach -# rmtree without it. -OWNED_SENTINEL = ".xorl-kernel-pin-owned" - - -def _mark_owned(path: str) -> None: - with open(os.path.join(path, OWNED_SENTINEL), "w") as f: - f.write("created by xorl.ops.kernel_config_pin; safe for it to replace\n") - - -def _rmtree_owned(path: str) -> None: - """Delete `path` only if this module created it (sentinel present).""" - if not os.path.isdir(path): - return - if not os.path.isfile(os.path.join(path, OWNED_SENTINEL)): - raise KernelConfigPinError( - f"refusing to delete {path!r}: it lacks the ownership sentinel " - f"{OWNED_SENTINEL!r} and was not created by this module. Remove or " - "relocate it manually if it is stale.", - ) - shutil.rmtree(path) - - -def _runtime_fingerprint() -> dict: - # Use the distribution version rather than only the module attribute: - # distinct flash-attn wheel builds can share a torch/triton fingerprint, - # while some builds do not expose a useful ``flash_attn.__version__``. - from importlib import metadata # noqa: PLC0415 - - fa = "unavailable" - for dist in ("flash-attn-4", "flash_attn_4", "flash-attn", "flash_attn"): - try: - fa = metadata.version(dist) - break - except metadata.PackageNotFoundError: - continue - if fa == "unavailable": - try: - import flash_attn # noqa: PLC0415 - - fa = getattr(flash_attn, "__version__", "unavailable") - except Exception: # pragma: no cover - build dependent - pass - return { - "torch": torch.__version__, - "triton": triton.__version__, - "flash_attn": fa, - "cuda": torch.version.cuda or "none", - } - - -def seed_exact_kernel_config_pin(pin_dir: str, *, source_cache: str | None = None) -> dict: - """Create/refresh a pin directory from the CURRENT runtime. - - Called by qualification runs (e.g. the fixture oracle phase) after their - kernels have been tuned. Copies `source_cache` (default: the active - TRITON_CACHE_DIR or ~/.triton/cache) into the pin and writes the - toolchain manifest. - """ - fingerprint = _runtime_fingerprint() - pin_dir = os.path.realpath(pin_dir) - parent = os.path.dirname(pin_dir) - if not os.path.isdir(parent): - raise KernelConfigPinError( - f"pin directory parent {parent!r} does not exist; refusing to create a pin at an implausible location", - ) - os.makedirs(pin_dir, exist_ok=True) - cache_src = os.path.realpath( - source_cache or os.environ.get("TRITON_CACHE_DIR", os.path.expanduser("~/.triton/cache")) - ) - cache_dst = os.path.join(pin_dir, CACHE_SUBDIR) - if os.path.commonpath([cache_src, cache_dst]) == cache_src: - raise KernelConfigPinError( - f"seed source cache {cache_src!r} contains the pin destination {cache_dst!r}; " - "copying would recurse into itself", - ) - _rmtree_owned(cache_dst) - if os.path.isdir(cache_src): - shutil.copytree(cache_src, cache_dst) - else: - os.makedirs(cache_dst, exist_ok=True) - _mark_owned(cache_dst) - with open(os.path.join(pin_dir, MANIFEST_NAME), "w") as f: - json.dump(fingerprint, f, indent=2) - logger.info("exact kernel-config pin seeded at %s: %s", pin_dir, fingerprint) - return fingerprint - - -def pin_exact_kernel_configs(*, rank: int | None = None) -> str: - """Admission-time pin. Returns the per-rank TRITON_CACHE_DIR it installed. - - Fail-closed on: env unset, pin dir or manifest missing, toolchain - fingerprint mismatch. Engagement-logged once per process. - """ - pin_dir = os.environ.get(PIN_DIR_ENV) - if not pin_dir: - raise KernelConfigPinError( - f"{PIN_DIR_ENV} is not set. Exact hybrid-Ulysses admission requires a seeded " - "kernel/toolchain pin directory (seed_exact_kernel_config_pin from the " - "qualification run). Byte claims are toolchain-scoped; refusing to run unpinned.", - ) - # The env var feeds filesystem mutations below: resolve it and refuse - # anything that is not an existing, seeded pin directory. - pin_dir = os.path.realpath(pin_dir) - if not os.path.isdir(pin_dir): - raise KernelConfigPinError( - f"{PIN_DIR_ENV}={pin_dir!r} is not an existing directory. Fail closed.", - ) - manifest_path = os.path.join(pin_dir, MANIFEST_NAME) - if not os.path.isfile(manifest_path): - raise KernelConfigPinError( - f"{PIN_DIR_ENV}={pin_dir!r} has no {MANIFEST_NAME}; the pin directory was never " - "seeded by a qualification run. Fail closed.", - ) - with open(manifest_path) as f: - pinned = json.load(f) - running = _runtime_fingerprint() - mismatches = {k: (pinned.get(k), running[k]) for k in running if pinned.get(k) != running[k]} - if mismatches: - raise KernelConfigPinError( - "Toolchain fingerprint mismatch against the kernel-config pin — the byte " - f"qualification does not transfer: {mismatches}. Re-qualify or restore the " - "pinned environment.", - ) - if rank is None: - rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) - clone = os.path.join(pin_dir, "clones", f"rank{rank}") - _rmtree_owned(clone) - shutil.copytree(os.path.join(pin_dir, CACHE_SUBDIR), clone) - _mark_owned(clone) - os.environ["TRITON_CACHE_DIR"] = clone - logger.info( - "exact kernel-config pin engaged: %s (rank %d clone %s, toolchain %s)", - pin_dir, - rank, - clone, - running, - ) - return clone +from xorl.ops.exact.kernel_config_pin import * # noqa: F401,F403 diff --git a/src/xorl/ops/rope_class_b.py b/src/xorl/ops/rope_class_b.py index 7471737d..c68883f9 100644 --- a/src/xorl/ops/rope_class_b.py +++ b/src/xorl/ops/rope_class_b.py @@ -1,161 +1,3 @@ -from __future__ import annotations +"""Moved to ``xorl.ops.exact.rope_class_b`` (issue #78 phase 3); compat stub.""" -# Class-B RoPE numerics, vendored from the matching sampler implementation so the trainer -# takes no runtime dependency on SGLang. -# -# Two RoPE numerics classes exist across the trainer/sampler pair: -# Class A -- per-op bf16 rounding (8 rounding points). The trainer's -# ``_naive_apply_rotary_pos_emb``, SGLang's eager ``apply_rotary_emb`` and its -# ``bi_fused_native`` triton kernel all agree bitwise here. -# Class B -- one fp32 chain with a single final round. SGLang's stock fused CUDA rope -# (``jit_kernel/csrc/elementwise/rope.cuh``) and the ``torch.compile`` RoPE path -# SGLang's RL lane selects agree bitwise here. -# -# ``_rotary_emb`` is the body of SGLang's ``rotary_embedding/utils.py::apply_rotary_emb``, -# vendored verbatim. What selects Class B is not the function but its input: callers hand it -# the fp32 cos/sin table serving builds, rather than the bf16 one the trainer's eager path -# rounds first. Compiled, it is bitwise equal to the stock CUDA kernel across every variant and -# shape the target models use -- that equality is a gate, not an assumption; see -# ``gate_phase1.py``. (The ``.to(x.dtype)`` below is load-bearing for the output dtype and is -# round-tripped away by inductor for the arithmetic; dropping it yields fp32 outputs.) -# -# A hand-written triton reimplementation of the CUDA kernel was tried first and rejected: it -# lands within one bf16 ULP on ~7e-4 of elements but is not bitwise equal, and no FMA -# association (plain, ``fma(x,cos,-(y*sin))``, ``fma(-y,sin,x*cos)``, or fully uncontracted -# ``mul.rn``) closes the gap. Bitwise is the requirement, so the compiled fp32 expression is -# what ships. -import logging -import os - -import torch - - -logger = logging.getLogger(__name__) - -# torch.compile is load-bearing here, not an optimisation: the EAGER form of this same -# expression is Class A. Dynamo's default recompile budget is 8, far under what a real model -# needs (separate q and k shapes x rollout lengths x forward and backward), and on exhausting -# it dynamo silently falls back to eager -- reverting those layers to Class A with no error and -# no log line. That is indistinguishable from a numerics bug downstream, so pin the budget and -# make exhaustion loud. -_ROPE_CLASS_B_RECOMPILE_LIMIT = 2048 -_ROPE_CLASS_B_ACCUMULATED_LIMIT = 8192 - - -def _pin_compile_budget() -> None: - cfg = torch._dynamo.config - cfg.recompile_limit = max(getattr(cfg, "recompile_limit", 0), _ROPE_CLASS_B_RECOMPILE_LIMIT) - cfg.accumulated_recompile_limit = max( - getattr(cfg, "accumulated_recompile_limit", 0), _ROPE_CLASS_B_ACCUMULATED_LIMIT - ) - # A silent eager fallback breaks the zero-K3 contract, so fail instead. Opt out with - # XORL_ROPE_CLASS_B_ALLOW_FALLBACK=1 for throughput experiments outside the contract. - if os.environ.get("XORL_ROPE_CLASS_B_ALLOW_FALLBACK") != "1" and hasattr(cfg, "fail_on_recompile_limit_hit"): - cfg.fail_on_recompile_limit_hit = True - logger.info( - "rope Class B: recompile_limit=%s accumulated=%s fail_on_limit=%s", - cfg.recompile_limit, - cfg.accumulated_recompile_limit, - getattr(cfg, "fail_on_recompile_limit_hit", None), - ) - - -_pin_compile_budget() - - -def _rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, is_neox: bool) -> torch.Tensor: - """SGLang's ``apply_rotary_emb``, vendored. - - x: ``[num_tokens, num_heads, rotary_dim]``; cos/sin: ``[num_tokens, rotary_dim // 2]`` fp32. - """ - cos = cos.unsqueeze(-2).to(x.dtype) - sin = sin.unsqueeze(-2).to(x.dtype) - if is_neox: - x1, x2 = torch.chunk(x, 2, dim=-1) - else: - x1 = x[..., ::2] - x2 = x[..., 1::2] - o1 = x1 * cos - x2 * sin - o2 = x2 * cos + x1 * sin - if is_neox: - return torch.cat((o1, o2), dim=-1) - return torch.stack((o1, o2), dim=-1).flatten(-2) - - -_rotary_emb_compiled = torch.compile(dynamic=True)(_rotary_emb) - - -class _ClassBRoPE(torch.autograd.Function): - """Class-B rope with a hand-written backward. - - RoPE is a block rotation, so its adjoint is the same rotation driven by a negated sin. - Supplying it explicitly keeps the backward one fused pass and lands it closer to the fp64 - adjoint than differentiating the forward expression. - """ - - @staticmethod - def forward(ctx, q_rot, k_rot, cos, sin, is_neox): - ctx.save_for_backward(cos, sin) - ctx.is_neox = is_neox - return ( - _rotary_emb_compiled(q_rot, cos, sin, is_neox), - _rotary_emb_compiled(k_rot, cos, sin, is_neox), - ) - - @staticmethod - def backward(ctx, grad_q, grad_k): - cos, sin = ctx.saved_tensors - neg_sin = -sin - return ( - _rotary_emb_compiled(grad_q.contiguous(), cos, neg_sin, ctx.is_neox), - _rotary_emb_compiled(grad_k.contiguous(), cos, neg_sin, ctx.is_neox), - None, - None, - None, - ) - - -def build_class_b_cos_sin(cos, sin, *, doubled: bool = True): - """Flatten the trainer's cos/sin into the fp32 ``[num_tokens, rotary_dim // 2]`` serving layout. - - With ``doubled`` (the default) cos/sin arrive in the ``[..., rotary_dim]`` layout the shared - ``RotaryEmbedding`` emits, where only the leading half carries distinct frequencies. GPT-OSS - emits them already halved, at ``[..., rotary_dim // 2]``. - """ - half = cos.shape[-1] // 2 if doubled else cos.shape[-1] - return cos[..., :half].reshape(-1, half).float(), sin[..., :half].reshape(-1, half).float() - - -def class_b_apply_rotary_pos_emb(q, k, cos, sin, *, interleaved: bool = False, doubled: bool = True): - """Class-B RoPE over ``[B, S, H, D]`` q/k. - - ``interleaved`` selects GPT-J pairing; the default is the half-split Neox pairing. Partial - rotary rotates the leading rotary_dim features and passes the tail through untouched. - """ - if cos.dtype != torch.float32 or sin.dtype != torch.float32: - raise RuntimeError( - "Class-B RoPE requires fp32 cos/sin at the apply boundary; " - f"got cos={cos.dtype}, sin={sin.dtype}. A mixed-precision wrapper " - "likely downcast the table in transit." - ) - cos_f, sin_f = build_class_b_cos_sin(cos, sin, doubled=doubled) - rotary_dim = 2 * cos_f.shape[-1] - batch, seq_len = q.shape[0], q.shape[1] - num_tokens = batch * seq_len - - q_rot = q[..., :rotary_dim].reshape(num_tokens, -1, rotary_dim) - k_rot = k[..., :rotary_dim].reshape(num_tokens, -1, rotary_dim) - q_out, k_out = _ClassBRoPE.apply(q_rot, k_rot, cos_f, sin_f, not interleaved) - q_out = q_out.view(batch, seq_len, -1, rotary_dim) - k_out = k_out.view(batch, seq_len, -1, rotary_dim) - - if q.shape[-1] > rotary_dim: - q_out = torch.cat((q_out, q[..., rotary_dim:]), dim=-1) - k_out = torch.cat((k_out, k[..., rotary_dim:]), dim=-1) - return q_out, k_out - - -__all__ = [ - "build_class_b_cos_sin", - "class_b_apply_rotary_pos_emb", -] +from xorl.ops.exact.rope_class_b import * # noqa: F401,F403 diff --git a/tests/distributed/test_glm52_fullparam_construction_gate.py b/tests/distributed/test_glm52_fullparam_construction_gate.py index 15e9dacf..11a089b0 100644 --- a/tests/distributed/test_glm52_fullparam_construction_gate.py +++ b/tests/distributed/test_glm52_fullparam_construction_gate.py @@ -146,7 +146,7 @@ def _fabricate_snapshot(config, snapshot_dir: str) -> None: import safetensors.torch as st from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8Experts - from xorl.ops.block_fp8_native import NativeBlockFP8Linear + from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear model = _build_meta_model(config) state: dict[str, torch.Tensor] = {} @@ -286,7 +286,7 @@ def _run_construction_gate() -> None: ) from xorl.models.transformers.glm5.modeling_glm5 import Glm5ForCausalLM from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8Experts - from xorl.ops.block_fp8_native import NativeBlockFP8Linear + from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear snapshot_dir = os.environ[_SNAPSHOT_ENV] dist.init_process_group("nccl") diff --git a/tests/distributed/test_glm52_fullparam_ep16_combine.py b/tests/distributed/test_glm52_fullparam_ep16_combine.py index c15c56ce..cfa00009 100644 --- a/tests/distributed/test_glm52_fullparam_ep16_combine.py +++ b/tests/distributed/test_glm52_fullparam_ep16_combine.py @@ -57,7 +57,7 @@ def _seed_bank(rank: int, device: torch.device = torch.device("cpu"), kind: str Glm52FullParamBlockFP8RoutedExperts, ) from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8Experts - from xorl.ops.block_fp8_native import pack_fp8_as_float32 + from xorl.ops.exact.block_fp8_native import pack_fp8_as_float32 if kind == "frozen": bank = Glm52NativeBlockFP8Experts(_LOCAL_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device) diff --git a/tests/distributed/test_native_fp8_fsdp2_materialization.py b/tests/distributed/test_native_fp8_fsdp2_materialization.py index dd57329c..29ce91da 100644 --- a/tests/distributed/test_native_fp8_fsdp2_materialization.py +++ b/tests/distributed/test_native_fp8_fsdp2_materialization.py @@ -19,7 +19,7 @@ from torch.distributed.tensor import DTensor, Shard from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8Experts -from xorl.ops.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear THIS_DIR = Path(__file__).resolve().parent diff --git a/tests/distributed/test_torch_parallelize_policies.py b/tests/distributed/test_torch_parallelize_policies.py index 312b7c54..8159841a 100644 --- a/tests/distributed/test_torch_parallelize_policies.py +++ b/tests/distributed/test_torch_parallelize_policies.py @@ -16,7 +16,7 @@ ) from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear from xorl.models.transformers.glm5.exact_shared_expert_qlora import Glm52ExactTP16SharedExpertBlockFP8QLoRA -from xorl.ops.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear class _FakeBlock: diff --git a/tests/models/test_glm52_exact_absorbed_kv_b_qlora.py b/tests/models/test_glm52_exact_absorbed_kv_b_qlora.py index ccccfd13..9619110f 100644 --- a/tests/models/test_glm52_exact_absorbed_kv_b_qlora.py +++ b/tests/models/test_glm52_exact_absorbed_kv_b_qlora.py @@ -10,7 +10,7 @@ GLM52_EXACT_TP1_ABSORBED_KV_B_QLORA_CONTRACT_VERSION, Glm52ExactTP1AbsorbedKvBBlockFP8QLoRA, ) -from xorl.ops.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear _NUM_HEADS = 64 diff --git a/tests/models/test_glm52_exact_dense_mlp_composition.py b/tests/models/test_glm52_exact_dense_mlp_composition.py index e3c6cfbc..296a7c96 100644 --- a/tests/models/test_glm52_exact_dense_mlp_composition.py +++ b/tests/models/test_glm52_exact_dense_mlp_composition.py @@ -8,7 +8,7 @@ Glm52ExactTP1FusedGateUpBlockFP8QLoRA, ) from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul def _pattern( diff --git a/tests/models/test_glm52_exact_fullparam_fp8.py b/tests/models/test_glm52_exact_fullparam_fp8.py index 12c3ee02..19463dae 100644 --- a/tests/models/test_glm52_exact_fullparam_fp8.py +++ b/tests/models/test_glm52_exact_fullparam_fp8.py @@ -26,8 +26,8 @@ quantize_master_to_serving_bytes, ) from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8DenseMLP -from xorl.ops.block_fp8_native import NativeBlockFP8Linear -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul def _linear(in_features: int = 8, out_features: int = 16) -> Glm52ExactTP1BlockFP8FullParamLinear: diff --git a/tests/models/test_glm52_exact_gate_up_qlora.py b/tests/models/test_glm52_exact_gate_up_qlora.py index 9ca76d42..231dfd9c 100644 --- a/tests/models/test_glm52_exact_gate_up_qlora.py +++ b/tests/models/test_glm52_exact_gate_up_qlora.py @@ -12,8 +12,8 @@ Glm52ExactTP1FusedGateUpBlockFP8QLoRA, ) from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear -from xorl.ops.block_fp8_native import NativeBlockFP8Linear -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul def _module() -> Glm52ExactTP1FusedGateUpBlockFP8QLoRA: diff --git a/tests/models/test_glm52_exact_qlora.py b/tests/models/test_glm52_exact_qlora.py index 45e20ed6..9034c953 100644 --- a/tests/models/test_glm52_exact_qlora.py +++ b/tests/models/test_glm52_exact_qlora.py @@ -10,7 +10,7 @@ GLM52_EXACT_TP1_QLORA_CONTRACT_VERSION, Glm52ExactTP1BlockFP8QLoRALinear, ) -from xorl.ops.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear def _module() -> Glm52ExactTP1BlockFP8QLoRALinear: diff --git a/tests/models/test_glm52_exact_shared_expert_qlora.py b/tests/models/test_glm52_exact_shared_expert_qlora.py index d6ba82f2..2b0914cf 100644 --- a/tests/models/test_glm52_exact_shared_expert_qlora.py +++ b/tests/models/test_glm52_exact_shared_expert_qlora.py @@ -12,7 +12,7 @@ GLM52_EXACT_TP16_SHARED_EXPERT_QLORA_CONTRACT_VERSION, Glm52ExactTP16SharedExpertBlockFP8QLoRA, ) -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul def _canonical_moe_reference(partials: torch.Tensor, metadata: CanonicalMoEGraphMetadata) -> torch.Tensor: diff --git a/tests/models/test_glm52_fullparam_admission.py b/tests/models/test_glm52_fullparam_admission.py index 3606f79e..67cbae99 100644 --- a/tests/models/test_glm52_fullparam_admission.py +++ b/tests/models/test_glm52_fullparam_admission.py @@ -164,7 +164,7 @@ def _seed_native_bytes(model: nn.Module) -> None: from xorl.models.transformers.glm5.modeling_glm5 import Glm5TopkRouter from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8Experts - from xorl.ops.block_fp8_native import NativeBlockFP8Linear + from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear def _pattern(*shape: int, offset: int) -> torch.Tensor: values = torch.arange(int(torch.tensor(shape).prod()), dtype=torch.float32) diff --git a/tests/models/test_glm52_native_fp8.py b/tests/models/test_glm52_native_fp8.py index 031e90b3..e6fddbc7 100644 --- a/tests/models/test_glm52_native_fp8.py +++ b/tests/models/test_glm52_native_fp8.py @@ -15,7 +15,7 @@ native_fp8_dense_source_map, validate_glm52_native_fp8_config, ) -from xorl.ops.block_fp8_native import NativeBlockFP8Linear, unpack_float32_as_fp8 +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear, unpack_float32_as_fp8 OFFICIAL_QUANT_CONFIG = { diff --git a/tests/models/test_glm52_qlora.py b/tests/models/test_glm52_qlora.py index 78d8896b..8f3a5e79 100644 --- a/tests/models/test_glm52_qlora.py +++ b/tests/models/test_glm52_qlora.py @@ -20,7 +20,7 @@ prepare_glm52_block_fp8_qlora, ) from xorl.models.transformers.glm5.support import validate_glm5_training_mode -from xorl.ops.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear from xorl.qlora.modules.moe_experts import BlockFP8QLoRAMoeExperts diff --git a/tests/models/test_glm5_one_round_swiglu.py b/tests/models/test_glm5_one_round_swiglu.py index b2d9f900..f4b573f7 100644 --- a/tests/models/test_glm5_one_round_swiglu.py +++ b/tests/models/test_glm5_one_round_swiglu.py @@ -48,7 +48,7 @@ ) from xorl.models.transformers.glm5.configuration_glm5 import Glm5Config from xorl.models.transformers.glm5.modeling_glm5 import GLM52_LOCAL_PARTIAL_POLICY, Glm5MLP -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul # GLM-5.2 geometries with byte evidence: the 128-wide TP16 shared-expert diff --git a/tests/models/test_module_utils_broadcast.py b/tests/models/test_module_utils_broadcast.py index 6f4b940e..e3912151 100644 --- a/tests/models/test_module_utils_broadcast.py +++ b/tests/models/test_module_utils_broadcast.py @@ -14,7 +14,7 @@ from xorl.models import module_utils from xorl.models.transformers.glm5.native_fp8 import NativeBlockFP8PairBuffer -from xorl.ops.block_fp8_native import NativeBlockFP8Linear +from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear pytestmark = [pytest.mark.cpu] diff --git a/tests/models/test_op_parity_dense.py b/tests/models/test_op_parity_dense.py index c55e31df..68d3dd88 100644 --- a/tests/models/test_op_parity_dense.py +++ b/tests/models/test_op_parity_dense.py @@ -16,7 +16,7 @@ import torch.nn.functional as F from xorl.models.layers import rope as xrope -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/models/test_qwen35_hybrid_ulysses_admission.py b/tests/models/test_qwen35_hybrid_ulysses_admission.py index 7082da36..bc639a4d 100644 --- a/tests/models/test_qwen35_hybrid_ulysses_admission.py +++ b/tests/models/test_qwen35_hybrid_ulysses_admission.py @@ -8,7 +8,7 @@ import torch import triton -from xorl.ops.kernel_config_pin import ( +from xorl.ops.exact.kernel_config_pin import ( KernelConfigPinError, pin_exact_kernel_configs, seed_exact_kernel_config_pin, @@ -67,7 +67,7 @@ def test_pin_replaces_its_own_clone(self, tmp_path, monkeypatch): assert first == second and os.path.isdir(second) def test_seed_refuses_to_delete_unowned_cache(self, tmp_path): - from xorl.ops.kernel_config_pin import CACHE_SUBDIR, OWNED_SENTINEL + from xorl.ops.exact.kernel_config_pin import CACHE_SUBDIR, OWNED_SENTINEL pin = self._seeded(tmp_path) (pin / CACHE_SUBDIR / OWNED_SENTINEL).unlink() # simulate a foreign dir diff --git a/tests/ops/test_bi_gemm_config_table.py b/tests/ops/test_bi_gemm_config_table.py index 2a61a5b1..78c5c694 100644 --- a/tests/ops/test_bi_gemm_config_table.py +++ b/tests/ops/test_bi_gemm_config_table.py @@ -17,7 +17,7 @@ matmul_kernel_persistent, set_batch_invariant_mode, ) -from xorl.ops.bi_gemm_configs import BASELINE_CONFIG, PINNED_BLOCK_K, lookup_mm_config +from xorl.ops.exact.bi_gemm_configs import BASELINE_CONFIG, PINNED_BLOCK_K, lookup_mm_config requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/ops/test_block_fp8_frozen_dgrad.py b/tests/ops/test_block_fp8_frozen_dgrad.py index dcac5173..e25160a3 100644 --- a/tests/ops/test_block_fp8_frozen_dgrad.py +++ b/tests/ops/test_block_fp8_frozen_dgrad.py @@ -26,7 +26,7 @@ import torch import torch.nn.functional as F -from xorl.ops.block_fp8_native import ( +from xorl.ops.exact.block_fp8_native import ( NATIVE_BLOCK_FP8_FROZEN_DGRAD_CONTRACT_VERSION, NativeBlockFP8Linear, ) @@ -126,7 +126,7 @@ def test_cuda_dgrad_matches_reference_and_composite_and_mutates_nothing(caplog) packed_before = module.packed_weight_f32.detach().view(torch.uint8).clone() scales_before = module.weight_scale_inv.detach().clone() - with caplog.at_level(logging.INFO, logger="xorl.ops.block_fp8_native"): + with caplog.at_level(logging.INFO, logger="xorl.ops.exact.block_fp8_native"): module.enable_frozen_activation_dgrad() module.enable_frozen_activation_dgrad() # idempotent engagement = [record for record in caplog.records if "frozen-trunk activation dgrad engaged" in record.message] diff --git a/tests/ops/test_block_fp8_native.py b/tests/ops/test_block_fp8_native.py index ac2ef594..c0c594ff 100644 --- a/tests/ops/test_block_fp8_native.py +++ b/tests/ops/test_block_fp8_native.py @@ -4,7 +4,7 @@ import torch import torch.distributed.checkpoint as dcp -from xorl.ops.block_fp8_native import ( +from xorl.ops.exact.block_fp8_native import ( NativeBlockFP8Linear, pack_fp8_as_float32, unpack_float32_as_fp8, diff --git a/tests/ops/test_canonical_moe_cast.py b/tests/ops/test_canonical_moe_cast.py index 38f555bf..39511e5c 100644 --- a/tests/ops/test_canonical_moe_cast.py +++ b/tests/ops/test_canonical_moe_cast.py @@ -3,7 +3,7 @@ import pytest import torch -from xorl.ops.canonical_moe_cast import canonical_moe_fp64_to_lowp_rne +from xorl.ops.exact.canonical_moe_cast import canonical_moe_fp64_to_lowp_rne @pytest.mark.gpu diff --git a/tests/ops/test_fused_silu_and_mul_scoping.py b/tests/ops/test_fused_silu_and_mul_scoping.py index 53b27fa4..fade45a5 100644 --- a/tests/ops/test_fused_silu_and_mul_scoping.py +++ b/tests/ops/test_fused_silu_and_mul_scoping.py @@ -20,7 +20,7 @@ import torch.nn.functional as F from xorl.models.transformers.qwen3_5.modeling_qwen3_5 import Qwen3_5MLP -from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul def _two_round_reference(x: torch.Tensor) -> torch.Tensor: diff --git a/tests/ops/test_rope_class_b_contract.py b/tests/ops/test_rope_class_b_contract.py index a8f72cf5..d5e027e4 100644 --- a/tests/ops/test_rope_class_b_contract.py +++ b/tests/ops/test_rope_class_b_contract.py @@ -3,7 +3,7 @@ import pytest import torch -from xorl.ops.rope_class_b import ( +from xorl.ops.exact.rope_class_b import ( build_class_b_cos_sin, class_b_apply_rotary_pos_emb, ) diff --git a/tests/server/weight_sync/test_glm52_fullparam_step_publish.py b/tests/server/weight_sync/test_glm52_fullparam_step_publish.py index 0c10322e..db8c106e 100644 --- a/tests/server/weight_sync/test_glm52_fullparam_step_publish.py +++ b/tests/server/weight_sync/test_glm52_fullparam_step_publish.py @@ -42,7 +42,7 @@ def _seeded_bank(expert_start: int, device: torch.device) -> Glm52FullParamBlockFP8RoutedExperts: - from xorl.ops.block_fp8_native import pack_fp8_as_float32 + from xorl.ops.exact.block_fp8_native import pack_fp8_as_float32 bank = Glm52FullParamBlockFP8RoutedExperts(_LOCAL_EXPERTS, _HIDDEN, _INTERMEDIATE, device=device) gate_up_fp8 = ( From cf680fb01787b2b8fc3651e9634493cbc33e1428 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 03:54:04 +0000 Subject: [PATCH 04/14] One home per family kernel tree; layer classes out of ops (#78 phase 4) - ops/glm5_kernels -> ops/families/glm5 and ops/dsv4 -> ops/families/dsv4. Deliberate deviation from the #78 sketch: the kernels stay under ops/ rather than moving into models/transformers// because those packages have import side effects (DSV4 mutates the HF Auto registries on import) and kernels must be importable without them. Old paths are sys.modules alias stubs. - The GatedDeltaNet and Mamba2Mixer nn.Module layer classes move from ops/linear_attention/layers and ops/ssm/layers to xorl.models.layers; ops keeps only kernels. The old package paths re-export the classes lazily (PEP 562) for one deprecation cycle. - In-repo imports rewritten; the kernel-adjacent contract modules (linear_attention/modules, ssm/modules) stay in ops by design. --- .../glm52/benchmark_sparse_mla_backward.py | 2 +- src/xorl/models/layers/gated_deltanet.py | 467 +++++++++++++++++ .../ssm => models}/layers/mamba2_mixer.py | 0 .../deepseek_v4/modeling_deepseek_v4.py | 32 +- src/xorl/models/transformers/glm5/indexer.py | 2 +- .../models/transformers/glm5/sparse_mla.py | 4 +- .../transformers/qwen3_5/modeling_qwen3_5.py | 2 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 2 +- src/xorl/ops/README.md | 15 +- src/xorl/ops/__init__.py | 12 +- src/xorl/ops/dsv4/__init__.py | 12 +- src/xorl/ops/families/__init__.py | 7 + src/xorl/ops/families/dsv4/__init__.py | 7 + .../ops/{ => families}/dsv4/attention_core.py | 0 .../ops/{ => families}/dsv4/compressor.py | 0 src/xorl/ops/{ => families}/dsv4/cp_utils.py | 0 .../{ => families}/dsv4/exact_attention.py | 4 +- .../{ => families}/dsv4/hyper_connection.py | 0 .../{ => families}/dsv4/kernel/__init__.py | 0 .../{ => families}/dsv4/kernel/act_quant.py | 0 .../{ => families}/dsv4/kernel/sinkhorn.py | 0 .../dsv4/kernel/tilelang_indexer.py | 0 .../dsv4/kernel/tilelang_indexer_bwd.py | 0 .../dsv4/kernel/tilelang_indexer_fwd.py | 0 .../dsv4/kernel/tilelang_sparse_mla_bwd.py | 0 .../dsv4/kernel/tilelang_sparse_mla_fwd.py | 0 src/xorl/ops/{ => families}/dsv4/qat.py | 0 src/xorl/ops/{ => families}/dsv4/rope.py | 0 src/xorl/ops/{ => families}/dsv4/utils.py | 0 .../ops/{ => families}/dsv4/v4_indexer.py | 0 src/xorl/ops/families/glm5/__init__.py | 13 + .../glm5}/flashmla_sparse_mla.py | 0 .../glm5}/indexer.py | 0 .../glm5}/sparse_mla.py | 0 .../glm5}/tilelang_indexer_bwd.py | 0 .../glm5}/tilelang_indexer_fwd.py | 0 .../glm5}/tilelang_sparse_mla_bwd.py | 0 .../glm5}/tilelang_sparse_mla_fwd.py | 0 src/xorl/ops/glm5_kernels/__init__.py | 16 +- src/xorl/ops/linear_attention/__init__.py | 16 +- .../ops/linear_attention/layers/__init__.py | 10 +- .../linear_attention/layers/gated_deltanet.py | 469 +----------------- src/xorl/ops/ssm/__init__.py | 16 +- src/xorl/ops/ssm/layers/__init__.py | 14 +- .../test_dsv4_exact_cp_attention_layout.py | 2 +- .../test_linear_attention_cp_equivalence.py | 2 +- tests/models/test_dsv4_attention.py | 14 +- tests/models/test_dsv4_exact_contract.py | 2 +- tests/models/test_dsv4_hf_to_dcp.py | 2 +- tests/models/test_dsv4_loader.py | 2 +- tests/models/test_dsv4_lora.py | 2 +- tests/models/test_dsv4_model.py | 2 +- tests/models/test_fused_gdn_lora.py | 2 +- tests/models/test_glm5_flashmla_sparse_mla.py | 4 +- tests/models/test_glm5_support.py | 4 +- .../test_qwen35_lora_projection_topology.py | 2 +- tests/ops/dsv4/test_compressor.py | 8 +- tests/ops/dsv4/test_exact_attention.py | 2 +- tests/ops/dsv4/test_v4_tilelang_indexer.py | 6 +- tests/ops/dsv4/test_v4_tilelang_sparse_mla.py | 16 +- tests/ops/test_bi_gdn_contract.py | 2 +- tests/ops/test_flashqla_contract_pin.py | 4 +- tests/ops/test_gdn_conv_contract.py | 4 +- tests/ops/test_sparse_mla_kernels.py | 8 +- tests/ops/test_ssm_mamba2.py | 3 +- 65 files changed, 645 insertions(+), 570 deletions(-) create mode 100644 src/xorl/models/layers/gated_deltanet.py rename src/xorl/{ops/ssm => models}/layers/mamba2_mixer.py (100%) create mode 100644 src/xorl/ops/families/__init__.py create mode 100644 src/xorl/ops/families/dsv4/__init__.py rename src/xorl/ops/{ => families}/dsv4/attention_core.py (100%) rename src/xorl/ops/{ => families}/dsv4/compressor.py (100%) rename src/xorl/ops/{ => families}/dsv4/cp_utils.py (100%) rename src/xorl/ops/{ => families}/dsv4/exact_attention.py (99%) rename src/xorl/ops/{ => families}/dsv4/hyper_connection.py (100%) rename src/xorl/ops/{ => families}/dsv4/kernel/__init__.py (100%) rename src/xorl/ops/{ => families}/dsv4/kernel/act_quant.py (100%) rename src/xorl/ops/{ => families}/dsv4/kernel/sinkhorn.py (100%) rename src/xorl/ops/{ => families}/dsv4/kernel/tilelang_indexer.py (100%) rename src/xorl/ops/{ => families}/dsv4/kernel/tilelang_indexer_bwd.py (100%) rename src/xorl/ops/{ => families}/dsv4/kernel/tilelang_indexer_fwd.py (100%) rename src/xorl/ops/{ => families}/dsv4/kernel/tilelang_sparse_mla_bwd.py (100%) rename src/xorl/ops/{ => families}/dsv4/kernel/tilelang_sparse_mla_fwd.py (100%) rename src/xorl/ops/{ => families}/dsv4/qat.py (100%) rename src/xorl/ops/{ => families}/dsv4/rope.py (100%) rename src/xorl/ops/{ => families}/dsv4/utils.py (100%) rename src/xorl/ops/{ => families}/dsv4/v4_indexer.py (100%) create mode 100644 src/xorl/ops/families/glm5/__init__.py rename src/xorl/ops/{glm5_kernels => families/glm5}/flashmla_sparse_mla.py (100%) rename src/xorl/ops/{glm5_kernels => families/glm5}/indexer.py (100%) rename src/xorl/ops/{glm5_kernels => families/glm5}/sparse_mla.py (100%) rename src/xorl/ops/{glm5_kernels => families/glm5}/tilelang_indexer_bwd.py (100%) rename src/xorl/ops/{glm5_kernels => families/glm5}/tilelang_indexer_fwd.py (100%) rename src/xorl/ops/{glm5_kernels => families/glm5}/tilelang_sparse_mla_bwd.py (100%) rename src/xorl/ops/{glm5_kernels => families/glm5}/tilelang_sparse_mla_fwd.py (100%) diff --git a/certification/glm52/benchmark_sparse_mla_backward.py b/certification/glm52/benchmark_sparse_mla_backward.py index 28ad6696..9e8eed46 100755 --- a/certification/glm52/benchmark_sparse_mla_backward.py +++ b/certification/glm52/benchmark_sparse_mla_backward.py @@ -36,7 +36,7 @@ def main() -> None: except ImportError as exc: raise SystemExit("This certification requires TileLang") from exc - from xorl.ops.glm5_kernels.sparse_mla import SparseMLA # noqa: PLC0415 + from xorl.ops.families.glm5.sparse_mla import SparseMLA # noqa: PLC0415 sequence, kv_sequence, heads, rank, tail, topk = 2048, 32768, 64, 512, 64, 2048 scale = (rank + tail) ** -0.5 diff --git a/src/xorl/models/layers/gated_deltanet.py b/src/xorl/models/layers/gated_deltanet.py new file mode 100644 index 00000000..5201a803 --- /dev/null +++ b/src/xorl/models/layers/gated_deltanet.py @@ -0,0 +1,467 @@ +from __future__ import annotations + +# Adapted from flash-linear-attention/fla/layers/gated_deltanet.py. +# Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. +import math +import warnings +from typing import Any + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from torch.nn import functional as F + +from xorl.lora.fold import lora_merged_forward_enabled +from xorl.ops.linear_attention.backend import ( + flashqla_chunk_gated_delta_rule, + flashqla_chunk_gated_delta_rule_cp, + get_gdn_backend, + warn_cp_fallback_once, +) +from xorl.ops.linear_attention.layers.utils import get_unpad_data, index_first_axis, pad_input +from xorl.ops.linear_attention.modules import ( + FusedRMSNormGated, + RMSNorm, + ShortConvolution, + causal_conv1d_qkv_contract, +) +from xorl.ops.linear_attention.modules.bi_contract import ( + _is_gdn_contract_enabled, + bi_fused_gdn_gating, + gdn_contract, +) +from xorl.ops.linear_attention.ops.gated_delta_rule import ( + chunk_gated_delta_rule, + fused_recurrent_gated_delta_rule, +) + + +def _sglang_compatible_beta_gate(b_input: torch.Tensor) -> torch.Tensor: + beta = b_input.float().sigmoid() + if beta.dtype != b_input.dtype: + beta = beta.to(dtype=b_input.dtype).float() + return beta + + +class GatedDeltaNet(nn.Module): + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 2, + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int | None = None, + mode: str = "chunk", + use_gate: bool = True, + use_short_conv: bool = True, + allow_neg_eigval: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int | None = None, + norm_eps: float = 1e-5, + exact_contract: bool = False, + **kwargs: Any, + ) -> None: + del kwargs + super().__init__() + + self.mode = mode + self.allow_neg_eigval = allow_neg_eigval + self.hidden_size = hidden_size + self.expand_v = expand_v + self.use_gate = use_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads + self.exact_contract = exact_contract + self.exact_merged_forward = exact_contract + + self.head_k_dim = head_dim + self.head_v_dim = int(self.head_dim * self.expand_v) + self.key_dim = int(self.num_heads * self.head_k_dim) + self.value_dim = int(self.num_v_heads * self.head_v_dim) + self.layer_idx = layer_idx + + if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by " + f"num_v_heads * head_dim={self.num_v_heads * self.head_dim}." + ) + if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: + raise ValueError( + f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.", + ) + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}.", + ) + if mode not in {"chunk", "fused_recurrent"}: + raise ValueError(f"Unsupported GatedDeltaNet mode: {mode}") + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) + self.b_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) + + A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp(torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min)) + dt = torch.clamp(dt, min=dt_init_floor) + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + self.dt_bias._no_weight_decay = True + + if use_short_conv: + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation="silu", + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation="silu", + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation="silu", + ) + else: + warnings.warn( + "ShortConvolution is usually important for GatedDeltaNet quality; " + "leave `use_short_conv=True` unless you know you want it disabled.", + stacklevel=2, + ) + + if use_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def _apply(self, fn, recurse: bool = True): + module = super()._apply(fn, recurse=recurse) + # Match SGLang's Qwen3.5/GDN contract: decay/time-step parameters stay fp32 + # even when the rest of the model is moved to bf16/fp16. + for name in ("A_log", "dt_bias"): + param = getattr(self, name, None) + if isinstance(param, nn.Parameter) and param.is_floating_point() and param.dtype != torch.float32: + with torch.no_grad(): + param.data = param.data.float() + if param.grad is not None: + param.grad.data = param.grad.data.float() + return module + + def _fused_qkvz_lora_delta( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Apply River's fused q/k/v/z LoRA delta and split its output.""" + adapter = getattr(self, "in_proj_qkvz", None) + if adapter is None: + return None + delta = adapter(hidden_states) + expected = 2 * self.key_dim + 2 * self.value_dim + if delta.shape[-1] != expected: + raise RuntimeError(f"Fused GDN LoRA produced {delta.shape[-1]} features, expected {expected}") + return delta.split( + (self.key_dim, self.key_dim, self.value_dim, self.value_dim), + dim=-1, + ) + + def _add_output_lora(self, inputs: torch.Tensor, base_output: torch.Tensor) -> torch.Tensor: + """Add River's fused GDN output LoRA delta when one is injected.""" + adapter = getattr(self, "out_proj", None) + if adapter is None: + return base_output + return base_output + adapter(inputs).to(base_output.dtype) + + @staticmethod + def _linear_with_contract( + module: nn.Linear, + inputs: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + if getattr(module, "_xorl_bi_trunk_wrapped", False): + from xorl.ops.batch_invariant_ops import batch_invariant_trunk_linear # noqa: PLC0415 + + return batch_invariant_trunk_linear(inputs, weight, module.bias) + return F.linear(inputs, weight, module.bias) + + def _project_output_linear(self, o: torch.Tensor) -> torch.Tensor: + output_adapter = getattr(self, "out_proj", None) + if output_adapter is not None and lora_merged_forward_enabled(self): + folded = output_adapter.merged_weight_for_forward(self.o_proj.weight) + return self._linear_with_contract(self.o_proj, o, folded) + + return self._add_output_lora(o, self.o_proj(o)) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Any | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor | None, Any | None]: + with gdn_contract(self.exact_contract): + return self._forward_impl( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + def _forward_impl( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Any | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor | None, Any | None]: + del output_attentions + if attention_mask is not None and len(attention_mask.shape) != 2: + raise ValueError( + "Expected `attention_mask` with shape [batch_size, seq_len] where 0 marks padding.", + ) + + batch_size, q_len, _ = hidden_states.shape + cp_context = kwargs.get("cp_context") + mode = ( + self.mode + if cp_context is not None + else ("fused_recurrent" if (q_len <= 64 and not self.training) else self.mode) + ) + if self.training and mode != "chunk": + raise AssertionError("Only chunk mode is supported in training.") + + last_state = None + if past_key_values is not None and self.layer_idx is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get("cu_seqlens") + indices = None + if cp_context is not None: + if attention_mask is not None: + raise ValueError( + "Ulysses linear attention currently requires packed inputs without a 2D attention_mask.", + ) + if cp_context.cu_seqlens is None: + raise ValueError( + "Ulysses linear attention requires cu_seqlens metadata from the collator.", + ) + if use_cache: + raise ValueError( + "Ulysses native FLA CP does not yet support KV/conv cache updates.", + ) + if mode != "chunk": + raise ValueError("Ulysses native FLA CP currently supports chunk mode only.") + cu_seqlens = cp_context.cu_seqlens + elif attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + merged_input_adapter = getattr(self, "in_proj_qkvz", None) + if merged_input_adapter is not None and lora_merged_forward_enabled(self): + offsets = [0] + for projection in (self.q_proj, self.k_proj, self.v_proj, self.g_proj): + offsets.append(offsets[-1] + projection.out_features) + q_weight = merged_input_adapter.merged_weight_for_forward( + self.q_proj.weight, output_start=offsets[0], output_end=offsets[1] + ) + k_weight = merged_input_adapter.merged_weight_for_forward( + self.k_proj.weight, output_start=offsets[1], output_end=offsets[2] + ) + v_weight = merged_input_adapter.merged_weight_for_forward( + self.v_proj.weight, output_start=offsets[2], output_end=offsets[3] + ) + gate_weight = merged_input_adapter.merged_weight_for_forward( + self.g_proj.weight, output_start=offsets[3], output_end=offsets[4] + ) + q_input = self._linear_with_contract(self.q_proj, hidden_states, q_weight) + k_input = self._linear_with_contract(self.k_proj, hidden_states, k_weight) + v_input = self._linear_with_contract(self.v_proj, hidden_states, v_weight) + gate_input = self._linear_with_contract(self.g_proj, hidden_states, gate_weight) + a_input = self.a_proj(hidden_states).float() + b_input = self.b_proj(hidden_states) + else: + q_input = self.q_proj(hidden_states) + k_input = self.k_proj(hidden_states) + v_input = self.v_proj(hidden_states) + a_input = self.a_proj(hidden_states).float() + b_input = self.b_proj(hidden_states) + gate_input = self.g_proj(hidden_states) if self.use_gate else None + + fused_lora_delta = None if lora_merged_forward_enabled(self) else self._fused_qkvz_lora_delta(hidden_states) + if fused_lora_delta is not None: + q_delta, k_delta, v_delta, gate_delta = fused_lora_delta + q_input = q_input + q_delta.to(q_input.dtype) + k_input = k_input + k_delta.to(k_input.dtype) + v_input = v_input + v_delta.to(v_input.dtype) + if gate_input is None: + raise RuntimeError("Fused qkvz LoRA requires a gated GDN layer") + gate_input = gate_input + gate_delta.to(gate_input.dtype) + + if _is_gdn_contract_enabled() and not self.use_short_conv: + raise RuntimeError("Exact Qwen3.5 GDN requires short convolution") + + if self.use_short_conv and _is_gdn_contract_enabled(): + if use_cache or last_state is not None: + raise RuntimeError( + "Exact Qwen3.5 trainer GDN supports packed prefill only, not recurrent cache updates" + ) + if cp_context is None: + # Fail-closed backstop against the ring/silent-skip fail-open: + # build_linear_attention_cp_context returns None for ring>1 + # and for missing metadata — an exact GDN layer must never + # run U1 math on sequence-sharded rows. + from xorl.distributed.parallel_state import get_parallel_state # noqa: PLC0415 + + ps = get_parallel_state() + if ps.ulysses_size > 1 or ps.ringattn_size > 1: + raise RuntimeError( + "Exact Qwen3.5 GDN: sequence parallelism is active " + f"(ulysses={ps.ulysses_size}, ring={ps.ringattn_size}) but no " + "cp_context reached the layer. Ring is unsupported (contract C4) " + "and a missing Ulysses context would silently break bytes. " + "Fail closed.", + ) + q, k, v = causal_conv1d_qkv_contract( + q_input, + k_input, + v_input, + self.q_conv1d, + self.k_conv1d, + self.v_conv1d, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + ) + conv_state_q = conv_state_k = conv_state_v = None + elif self.use_short_conv: + conv_state_q = conv_state_k = conv_state_v = None + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state["conv_state"] + q, conv_state_q = self.q_conv1d( + x=q_input, + cache=conv_state_q, + output_final_state=bool(use_cache), + cu_seqlens=cu_seqlens, + cp_context=cp_context, + ) + k, conv_state_k = self.k_conv1d( + x=k_input, + cache=conv_state_k, + output_final_state=bool(use_cache), + cu_seqlens=cu_seqlens, + cp_context=cp_context, + ) + v, conv_state_v = self.v_conv1d( + x=v_input, + cache=conv_state_v, + output_final_state=bool(use_cache), + cu_seqlens=cu_seqlens, + cp_context=cp_context, + ) + else: + q = F.silu(q_input) + k = F.silu(k_input) + v = F.silu(v_input) + conv_state_q = conv_state_k = conv_state_v = None + + q, k = (rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim) for x in (q, k)) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_v_dim) + + if self.num_v_heads > self.num_heads: + repeat_factor = self.num_v_heads // self.num_heads + q, k = (repeat(x, "... h d -> ... (h g) d", g=repeat_factor) for x in (q, k)) + + if _is_gdn_contract_enabled(): + # The serving fused-GDN gating kernel removes the + # 1-ULP g term (torch softplus vs tl.log(1+tl.exp)); beta is bitwise + # either way. + g, beta = bi_fused_gdn_gating(self.A_log, a_input, b_input, self.dt_bias) + else: + beta = _sglang_compatible_beta_gate(b_input) + g = -self.A_log.float().exp() * F.softplus(a_input + self.dt_bias) + if self.allow_neg_eigval: + beta = beta * 2.0 + recurrent_state = last_state["recurrent_state"] if last_state is not None else None + + if mode == "chunk": + backend = get_gdn_backend() + chunk_fn = chunk_gated_delta_rule + if backend == "flashqla": + if self.head_k_dim != 128 or self.head_v_dim != 128: + warn_cp_fallback_once() + elif cp_context is not None: + chunk_fn = flashqla_chunk_gated_delta_rule_cp + else: + chunk_fn = flashqla_chunk_gated_delta_rule + + chunk_kwargs = dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=bool(use_cache), + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + if cp_context is not None: + chunk_kwargs["cp_context"] = cp_context + o, recurrent_state = chunk_fn(**chunk_kwargs) + + elif mode == "fused_recurrent": + o, recurrent_state = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=bool(use_cache), + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + else: + raise NotImplementedError(f"Unsupported mode `{mode}`.") + + if past_key_values is not None and self.layer_idx is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q_len, + ) + + if self.use_gate: + gate = rearrange(gate_input, "... (h d) -> ... h d", d=self.head_v_dim) + o = self.o_norm(o, gate) + else: + o = self.o_norm(o) + o = rearrange(o, "b t h d -> b t (h d)") + o = self._project_output_linear(o) + if attention_mask is not None and indices is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/src/xorl/ops/ssm/layers/mamba2_mixer.py b/src/xorl/models/layers/mamba2_mixer.py similarity index 100% rename from src/xorl/ops/ssm/layers/mamba2_mixer.py rename to src/xorl/models/layers/mamba2_mixer.py diff --git a/src/xorl/models/transformers/deepseek_v4/modeling_deepseek_v4.py b/src/xorl/models/transformers/deepseek_v4/modeling_deepseek_v4.py index e4ba4858..6258d3ee 100644 --- a/src/xorl/models/transformers/deepseek_v4/modeling_deepseek_v4.py +++ b/src/xorl/models/transformers/deepseek_v4/modeling_deepseek_v4.py @@ -42,9 +42,9 @@ from xorl.models.layers.normalization import RMSNorm from xorl.models.module_utils import DEFAULT_GRADIENT_CHECKPOINTING_METHOD from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput -from xorl.ops.dsv4.attention_core import dense_attn_torch, sparse_attn_tilelang, sparse_attn_torch -from xorl.ops.dsv4.compressor import DeepSeekV4Compressor -from xorl.ops.dsv4.cp_utils import ( +from xorl.ops.families.dsv4.attention_core import dense_attn_torch, sparse_attn_tilelang, sparse_attn_torch +from xorl.ops.families.dsv4.compressor import DeepSeekV4Compressor +from xorl.ops.families.dsv4.cp_utils import ( Dsv4ExactCPLayout, all_gather_cp, build_dsv4_exact_cp_layout, @@ -54,11 +54,11 @@ get_q_positions_for_cp, get_window_topk_idxs_cp, ) -from xorl.ops.dsv4.hyper_connection import DeepSeekV4HyperConnectionUtil -from xorl.ops.dsv4.qat import fp8_simulate_qat -from xorl.ops.dsv4.rope import apply_rotary_emb, wrapped_precompute_freqs_cis -from xorl.ops.dsv4.utils import dsv4_kv_qat_enabled -from xorl.ops.dsv4.v4_indexer import V4Indexer +from xorl.ops.families.dsv4.hyper_connection import DeepSeekV4HyperConnectionUtil +from xorl.ops.families.dsv4.qat import fp8_simulate_qat +from xorl.ops.families.dsv4.rope import apply_rotary_emb, wrapped_precompute_freqs_cis +from xorl.ops.families.dsv4.utils import dsv4_kv_qat_enabled +from xorl.ops.families.dsv4.v4_indexer import V4Indexer # Defaults: when ``XORL_DSV4_SPARSE_ATTN_IMPL`` is unset, autodetect — @@ -393,7 +393,7 @@ def validate_lora_metadata(where: str) -> None: raise RuntimeError("DSV4 exact decode-cache carry admits one request without context parallelism") carry_state = self.__dict__.get("_dsv4_decode_state") if carry_state is None: - from xorl.ops.dsv4.exact_attention import Dsv4DecodeCarryState # noqa: PLC0415 + from xorl.ops.families.dsv4.exact_attention import Dsv4DecodeCarryState # noqa: PLC0415 carry_state = Dsv4DecodeCarryState() self._dsv4_decode_state = carry_state @@ -464,14 +464,14 @@ def validate_lora_metadata(where: str) -> None: q = self.wq_b(q_lora) # [B, S, n_heads * head_dim] q = q.unflatten(-1, (self.n_local_heads, self.head_dim)) # [B, S, H, D] if self._exact_attention: - from xorl.ops.dsv4.exact_attention import exact_q_norm_rope # noqa: PLC0415 + from xorl.ops.families.dsv4.exact_attention import exact_q_norm_rope # noqa: PLC0415 q = exact_q_norm_rope(q, freqs_cis, self.eps, position_offset=carry_offset or 0) else: q_dtype = q.dtype q = q.float() q = (q * torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps)).to(q_dtype) - # ``apply_rotary_emb`` (xorl.ops.dsv4.rope) writes into ``q[..., -rd:]`` + # ``apply_rotary_emb`` (xorl.ops.families.dsv4.rope) writes into ``q[..., -rd:]`` # in place, so the slice's storage must be exclusively ours. Without # ``q.clone()``, the in-place rotary would mutate the upstream # ``self.wq_b(...)`` activation that's still held by autograd, which @@ -529,7 +529,7 @@ def validate_lora_metadata(where: str) -> None: request_x = x.index_select(1, local_rows) if ratio else None if ratio == 0: - from xorl.ops.dsv4.exact_attention import exact_c0_attention # noqa: PLC0415 + from xorl.ops.families.dsv4.exact_attention import exact_c0_attention # noqa: PLC0415 request_o = exact_c0_attention( request_q, @@ -543,7 +543,7 @@ def validate_lora_metadata(where: str) -> None: kv_preprocessed=False, ) else: - from xorl.ops.dsv4.exact_attention import exact_compressed_attention # noqa: PLC0415 + from xorl.ops.families.dsv4.exact_attention import exact_compressed_attention # noqa: PLC0415 request_o = exact_compressed_attention( request_q, @@ -571,7 +571,7 @@ def validate_lora_metadata(where: str) -> None: if ratio: exact_x = all_gather_cp(exact_x, dim=1, cp_group=self.cp_group) if ratio == 0: - from xorl.ops.dsv4.exact_attention import exact_c0_attention # noqa: PLC0415 + from xorl.ops.families.dsv4.exact_attention import exact_c0_attention # noqa: PLC0415 o = exact_c0_attention( q, @@ -587,7 +587,7 @@ def validate_lora_metadata(where: str) -> None: kv_preprocessed=False, ) else: - from xorl.ops.dsv4.exact_attention import exact_compressed_attention # noqa: PLC0415 + from xorl.ops.families.dsv4.exact_attention import exact_compressed_attention # noqa: PLC0415 o = exact_compressed_attention( q, @@ -669,7 +669,7 @@ def validate_lora_metadata(where: str) -> None: # Inverse RoPE on the rope slice of the output. if self._exact_attention: - from xorl.ops.dsv4.exact_attention import exact_inverse_rope # noqa: PLC0415 + from xorl.ops.families.dsv4.exact_attention import exact_inverse_rope # noqa: PLC0415 o = exact_inverse_rope(o, freqs_cis, position_offset=carry_offset or 0) else: diff --git a/src/xorl/models/transformers/glm5/indexer.py b/src/xorl/models/transformers/glm5/indexer.py index 7bb11464..b522ece9 100644 --- a/src/xorl/models/transformers/glm5/indexer.py +++ b/src/xorl/models/transformers/glm5/indexer.py @@ -478,7 +478,7 @@ def _try_tilelang_select_topk( return None try: - from xorl.ops.glm5_kernels.tilelang_indexer_fwd import ( # noqa: PLC0415 + from xorl.ops.families.glm5.tilelang_indexer_fwd import ( # noqa: PLC0415 clean_logits_, tl_indexer_fwd_impl, ) diff --git a/src/xorl/models/transformers/glm5/sparse_mla.py b/src/xorl/models/transformers/glm5/sparse_mla.py index 264812f2..bec7c0a5 100644 --- a/src/xorl/models/transformers/glm5/sparse_mla.py +++ b/src/xorl/models/transformers/glm5/sparse_mla.py @@ -143,7 +143,7 @@ def _sparse_mla_tilelang( miles' 3D ``[B*S_q, H, D]`` and ``[B*S_kv, 1, D]`` forms before calling `SparseMLA.apply`, then reshapes back. """ - from xorl.ops.glm5_kernels.sparse_mla import SparseMLA # noqa: PLC0415 (lazy: tilelang is optional) + from xorl.ops.families.glm5.sparse_mla import SparseMLA # noqa: PLC0415 (lazy: tilelang is optional) del query_offset B, S_q, H, D = q.shape @@ -239,7 +239,7 @@ def _sparse_mla_flashmla( if error is not None: raise RuntimeError(f"backend='flashmla' requested outside its certified GLM-5.2 envelope: {error}") - from xorl.ops.glm5_kernels.flashmla_sparse_mla import ( # noqa: PLC0415 + from xorl.ops.families.glm5.flashmla_sparse_mla import ( # noqa: PLC0415 FlashMLASparseWithTileLangBackward, ) diff --git a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py index b5aa20a6..a0dafb40 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -19,6 +19,7 @@ ) from xorl.models.layers.attention.backend import get_attention_fn from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora +from xorl.models.layers.gated_deltanet import GatedDeltaNet from xorl.models.layers.normalization import ( compiled_zero_centered_rms_norm, eager_zero_centered_rms_norm, @@ -43,7 +44,6 @@ qwen3_5_apply_rotary_pos_emb, ) from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul -from xorl.ops.linear_attention import GatedDeltaNet from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils import logging diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index bc4380bc..00d6cc91 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -18,6 +18,7 @@ from xorl.models.layers.attention import AttentionKwargs, update_causal_mask from xorl.models.layers.attention.backend import ATTENTION_FUNCTIONS from xorl.models.layers.attention.backend.eager import eager_attention_forward +from xorl.models.layers.gated_deltanet import GatedDeltaNet from xorl.models.layers.moe import MoEBlock, MoEExperts from xorl.models.layers.moe.ep_native_combine import validate_native_ep_combine_size from xorl.models.layers.normalization import ( @@ -44,7 +45,6 @@ qwen3_5_apply_rotary_pos_emb, ) from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul -from xorl.ops.linear_attention import GatedDeltaNet from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils import logging diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index 79ace5a9..a38533de 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -27,13 +27,16 @@ Three kinds of code live here with three different rules: paths here are compat stubs for one deprecation cycle. - `moe/` — MoE expert compute backends (triton/quack/native, LoRA variants). `ep_kernels/` (DeepEP sort/scatter) merges in here (#78 phase 5). -- `linear_attention/` — GDN/linear-attention kernels; also currently hosts - the `GatedDeltaNet` layer class (moves to `models/layers/`, #78 phase 4). -- `ssm/` — Mamba-2 kernels; also currently hosts the `Mamba2Mixer` layer - class (same plan as above). +- `linear_attention/` — GDN/linear-attention kernels. The `GatedDeltaNet` + layer class lives in `models/layers/gated_deltanet.py` (#78 phase 4); the + old paths re-export it lazily for one deprecation cycle. +- `ssm/` — Mamba-2 kernels; `Mamba2Mixer` likewise lives in + `models/layers/mamba2_mixer.py`. - `quantize/` — NF4/INT4/FP4/FP8 quantization codecs and fake-quant ops. -- `glm5_kernels/`, `dsv4/` — model-family-specific kernels; planned home: - `models/transformers/{glm5,deepseek_v4}/kernels/` (#78 phase 4). +- `families/` — model-family-specific kernels (`glm5/`, `dsv4/`, #78 + phase 4). Kept under `ops/` (not inside `models/transformers//`) + because the model packages have import side effects (DSV4 registers with + the HF Auto registries on import); old paths are alias stubs. - `_vendored/` — vendored trees (see above); old paths (`ops/quack`, `ops/linear_attention/flashqla`) are alias stubs for one deprecation cycle. ## What does NOT belong here diff --git a/src/xorl/ops/__init__.py b/src/xorl/ops/__init__.py index 36d8f446..5afc1ae8 100644 --- a/src/xorl/ops/__init__.py +++ b/src/xorl/ops/__init__.py @@ -1,5 +1,4 @@ from .linear_attention import ( - GatedDeltaNet, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule, ) @@ -9,7 +8,7 @@ TritonMoeExpertsLoRAFunction, triton_moe_lora_forward, ) -from .ssm import Mamba2Mixer, ssd_chunked +from .ssm import ssd_chunked def __getattr__(name: str): @@ -20,6 +19,15 @@ def __getattr__(name: str): import xorl.objectives as _objectives return getattr(_objectives, name) + # Layer classes moved to xorl.models.layers (#78 phase 4). + if name == "GatedDeltaNet": + from xorl.models.layers.gated_deltanet import GatedDeltaNet + + return GatedDeltaNet + if name == "Mamba2Mixer": + from xorl.models.layers.mamba2_mixer import Mamba2Mixer + + return Mamba2Mixer raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/xorl/ops/dsv4/__init__.py b/src/xorl/ops/dsv4/__init__.py index 027b24c4..1e353edc 100644 --- a/src/xorl/ops/dsv4/__init__.py +++ b/src/xorl/ops/dsv4/__init__.py @@ -1,7 +1,7 @@ -"""DeepSeek V4 operator kernels and helpers. +"""Moved to ``xorl.ops.families.dsv4`` (issue #78 phase 4); alias stub.""" -The ``kernel`` submodule contains the TileLang-based kernels (sparse MLA fwd/bwd, -DSA indexer fwd/bwd, FP8 act-quant, hyper-connection sinkhorn). Higher-level -layer wrappers (compressor, indexer, attention) live alongside this package -in subsequent phases. -""" +import importlib as _importlib +import sys as _sys + + +_sys.modules[__name__] = _importlib.import_module("xorl.ops.families.dsv4") diff --git a/src/xorl/ops/families/__init__.py b/src/xorl/ops/families/__init__.py new file mode 100644 index 00000000..977926b7 --- /dev/null +++ b/src/xorl/ops/families/__init__.py @@ -0,0 +1,7 @@ +"""Model-family-specific kernels (issue #78 phase 4). + +One home per family. These stay under ``ops/`` rather than inside +``models/transformers//`` deliberately: the model packages have +import side effects (e.g. DeepSeek-V4 registers itself with the HF Auto +registries on import), and kernels must be importable without them. +""" diff --git a/src/xorl/ops/families/dsv4/__init__.py b/src/xorl/ops/families/dsv4/__init__.py new file mode 100644 index 00000000..027b24c4 --- /dev/null +++ b/src/xorl/ops/families/dsv4/__init__.py @@ -0,0 +1,7 @@ +"""DeepSeek V4 operator kernels and helpers. + +The ``kernel`` submodule contains the TileLang-based kernels (sparse MLA fwd/bwd, +DSA indexer fwd/bwd, FP8 act-quant, hyper-connection sinkhorn). Higher-level +layer wrappers (compressor, indexer, attention) live alongside this package +in subsequent phases. +""" diff --git a/src/xorl/ops/dsv4/attention_core.py b/src/xorl/ops/families/dsv4/attention_core.py similarity index 100% rename from src/xorl/ops/dsv4/attention_core.py rename to src/xorl/ops/families/dsv4/attention_core.py diff --git a/src/xorl/ops/dsv4/compressor.py b/src/xorl/ops/families/dsv4/compressor.py similarity index 100% rename from src/xorl/ops/dsv4/compressor.py rename to src/xorl/ops/families/dsv4/compressor.py diff --git a/src/xorl/ops/dsv4/cp_utils.py b/src/xorl/ops/families/dsv4/cp_utils.py similarity index 100% rename from src/xorl/ops/dsv4/cp_utils.py rename to src/xorl/ops/families/dsv4/cp_utils.py diff --git a/src/xorl/ops/dsv4/exact_attention.py b/src/xorl/ops/families/dsv4/exact_attention.py similarity index 99% rename from src/xorl/ops/dsv4/exact_attention.py rename to src/xorl/ops/families/dsv4/exact_attention.py index 7d53a2e6..18ff24ea 100644 --- a/src/xorl/ops/dsv4/exact_attention.py +++ b/src/xorl/ops/families/dsv4/exact_attention.py @@ -935,7 +935,7 @@ def forward( @staticmethod def backward(ctx, grad_output: Tensor): - from xorl.ops.dsv4.attention_core import sparse_attn_torch # noqa: PLC0415 + from xorl.ops.families.dsv4.attention_core import sparse_attn_torch # noqa: PLC0415 if ctx.carry_offset is not None: ( @@ -1225,7 +1225,7 @@ def forward( @staticmethod def backward(ctx, grad_output: Tensor): - from xorl.ops.dsv4.attention_core import sparse_attn_torch # noqa: PLC0415 + from xorl.ops.families.dsv4.attention_core import sparse_attn_torch # noqa: PLC0415 if ctx.carry_offset is not None: ( diff --git a/src/xorl/ops/dsv4/hyper_connection.py b/src/xorl/ops/families/dsv4/hyper_connection.py similarity index 100% rename from src/xorl/ops/dsv4/hyper_connection.py rename to src/xorl/ops/families/dsv4/hyper_connection.py diff --git a/src/xorl/ops/dsv4/kernel/__init__.py b/src/xorl/ops/families/dsv4/kernel/__init__.py similarity index 100% rename from src/xorl/ops/dsv4/kernel/__init__.py rename to src/xorl/ops/families/dsv4/kernel/__init__.py diff --git a/src/xorl/ops/dsv4/kernel/act_quant.py b/src/xorl/ops/families/dsv4/kernel/act_quant.py similarity index 100% rename from src/xorl/ops/dsv4/kernel/act_quant.py rename to src/xorl/ops/families/dsv4/kernel/act_quant.py diff --git a/src/xorl/ops/dsv4/kernel/sinkhorn.py b/src/xorl/ops/families/dsv4/kernel/sinkhorn.py similarity index 100% rename from src/xorl/ops/dsv4/kernel/sinkhorn.py rename to src/xorl/ops/families/dsv4/kernel/sinkhorn.py diff --git a/src/xorl/ops/dsv4/kernel/tilelang_indexer.py b/src/xorl/ops/families/dsv4/kernel/tilelang_indexer.py similarity index 100% rename from src/xorl/ops/dsv4/kernel/tilelang_indexer.py rename to src/xorl/ops/families/dsv4/kernel/tilelang_indexer.py diff --git a/src/xorl/ops/dsv4/kernel/tilelang_indexer_bwd.py b/src/xorl/ops/families/dsv4/kernel/tilelang_indexer_bwd.py similarity index 100% rename from src/xorl/ops/dsv4/kernel/tilelang_indexer_bwd.py rename to src/xorl/ops/families/dsv4/kernel/tilelang_indexer_bwd.py diff --git a/src/xorl/ops/dsv4/kernel/tilelang_indexer_fwd.py b/src/xorl/ops/families/dsv4/kernel/tilelang_indexer_fwd.py similarity index 100% rename from src/xorl/ops/dsv4/kernel/tilelang_indexer_fwd.py rename to src/xorl/ops/families/dsv4/kernel/tilelang_indexer_fwd.py diff --git a/src/xorl/ops/dsv4/kernel/tilelang_sparse_mla_bwd.py b/src/xorl/ops/families/dsv4/kernel/tilelang_sparse_mla_bwd.py similarity index 100% rename from src/xorl/ops/dsv4/kernel/tilelang_sparse_mla_bwd.py rename to src/xorl/ops/families/dsv4/kernel/tilelang_sparse_mla_bwd.py diff --git a/src/xorl/ops/dsv4/kernel/tilelang_sparse_mla_fwd.py b/src/xorl/ops/families/dsv4/kernel/tilelang_sparse_mla_fwd.py similarity index 100% rename from src/xorl/ops/dsv4/kernel/tilelang_sparse_mla_fwd.py rename to src/xorl/ops/families/dsv4/kernel/tilelang_sparse_mla_fwd.py diff --git a/src/xorl/ops/dsv4/qat.py b/src/xorl/ops/families/dsv4/qat.py similarity index 100% rename from src/xorl/ops/dsv4/qat.py rename to src/xorl/ops/families/dsv4/qat.py diff --git a/src/xorl/ops/dsv4/rope.py b/src/xorl/ops/families/dsv4/rope.py similarity index 100% rename from src/xorl/ops/dsv4/rope.py rename to src/xorl/ops/families/dsv4/rope.py diff --git a/src/xorl/ops/dsv4/utils.py b/src/xorl/ops/families/dsv4/utils.py similarity index 100% rename from src/xorl/ops/dsv4/utils.py rename to src/xorl/ops/families/dsv4/utils.py diff --git a/src/xorl/ops/dsv4/v4_indexer.py b/src/xorl/ops/families/dsv4/v4_indexer.py similarity index 100% rename from src/xorl/ops/dsv4/v4_indexer.py rename to src/xorl/ops/families/dsv4/v4_indexer.py diff --git a/src/xorl/ops/families/glm5/__init__.py b/src/xorl/ops/families/glm5/__init__.py new file mode 100644 index 00000000..7ddff7c1 --- /dev/null +++ b/src/xorl/ops/families/glm5/__init__.py @@ -0,0 +1,13 @@ +"""GLM-5 sparse-MLA + DSA-indexer kernels. + +Vendored from the radixark `miles` glm5 plugin +(`miles_plugins/models/glm5/ops/`). All four kernel files are tilelang-only +— the dispatch layer in ``xorl.models.transformers.glm5.sparse_mla`` and +``...glm5.indexer`` lazy-imports tilelang and falls back to the pure-torch +reference paths when tilelang isn't installed or the input is on CPU. + +We carry the bwd kernels even though current usage is fwd-only because +miles' `SparseMLA` / `IndexerFunction` autograd wrappers reference them at +class-definition time; vendoring fwd alone would force a runtime import +error on `import miles_plugins.models.glm5.ops.sparse_mla`. +""" diff --git a/src/xorl/ops/glm5_kernels/flashmla_sparse_mla.py b/src/xorl/ops/families/glm5/flashmla_sparse_mla.py similarity index 100% rename from src/xorl/ops/glm5_kernels/flashmla_sparse_mla.py rename to src/xorl/ops/families/glm5/flashmla_sparse_mla.py diff --git a/src/xorl/ops/glm5_kernels/indexer.py b/src/xorl/ops/families/glm5/indexer.py similarity index 100% rename from src/xorl/ops/glm5_kernels/indexer.py rename to src/xorl/ops/families/glm5/indexer.py diff --git a/src/xorl/ops/glm5_kernels/sparse_mla.py b/src/xorl/ops/families/glm5/sparse_mla.py similarity index 100% rename from src/xorl/ops/glm5_kernels/sparse_mla.py rename to src/xorl/ops/families/glm5/sparse_mla.py diff --git a/src/xorl/ops/glm5_kernels/tilelang_indexer_bwd.py b/src/xorl/ops/families/glm5/tilelang_indexer_bwd.py similarity index 100% rename from src/xorl/ops/glm5_kernels/tilelang_indexer_bwd.py rename to src/xorl/ops/families/glm5/tilelang_indexer_bwd.py diff --git a/src/xorl/ops/glm5_kernels/tilelang_indexer_fwd.py b/src/xorl/ops/families/glm5/tilelang_indexer_fwd.py similarity index 100% rename from src/xorl/ops/glm5_kernels/tilelang_indexer_fwd.py rename to src/xorl/ops/families/glm5/tilelang_indexer_fwd.py diff --git a/src/xorl/ops/glm5_kernels/tilelang_sparse_mla_bwd.py b/src/xorl/ops/families/glm5/tilelang_sparse_mla_bwd.py similarity index 100% rename from src/xorl/ops/glm5_kernels/tilelang_sparse_mla_bwd.py rename to src/xorl/ops/families/glm5/tilelang_sparse_mla_bwd.py diff --git a/src/xorl/ops/glm5_kernels/tilelang_sparse_mla_fwd.py b/src/xorl/ops/families/glm5/tilelang_sparse_mla_fwd.py similarity index 100% rename from src/xorl/ops/glm5_kernels/tilelang_sparse_mla_fwd.py rename to src/xorl/ops/families/glm5/tilelang_sparse_mla_fwd.py diff --git a/src/xorl/ops/glm5_kernels/__init__.py b/src/xorl/ops/glm5_kernels/__init__.py index 7ddff7c1..60d4376f 100644 --- a/src/xorl/ops/glm5_kernels/__init__.py +++ b/src/xorl/ops/glm5_kernels/__init__.py @@ -1,13 +1,7 @@ -"""GLM-5 sparse-MLA + DSA-indexer kernels. +"""Moved to ``xorl.ops.families.glm5`` (issue #78 phase 4); alias stub.""" -Vendored from the radixark `miles` glm5 plugin -(`miles_plugins/models/glm5/ops/`). All four kernel files are tilelang-only -— the dispatch layer in ``xorl.models.transformers.glm5.sparse_mla`` and -``...glm5.indexer`` lazy-imports tilelang and falls back to the pure-torch -reference paths when tilelang isn't installed or the input is on CPU. +import importlib as _importlib +import sys as _sys -We carry the bwd kernels even though current usage is fwd-only because -miles' `SparseMLA` / `IndexerFunction` autograd wrappers reference them at -class-definition time; vendoring fwd alone would force a runtime import -error on `import miles_plugins.models.glm5.ops.sparse_mla`. -""" + +_sys.modules[__name__] = _importlib.import_module("xorl.ops.families.glm5") diff --git a/src/xorl/ops/linear_attention/__init__.py b/src/xorl/ops/linear_attention/__init__.py index a10dcd3a..1a5453c6 100644 --- a/src/xorl/ops/linear_attention/__init__.py +++ b/src/xorl/ops/linear_attention/__init__.py @@ -1,9 +1,21 @@ -"""Linear attention ops and layers used by Qwen3.5.""" +"""Linear attention kernels used by Qwen3.5. + +The ``GatedDeltaNet`` layer class moved to +:mod:`xorl.models.layers.gated_deltanet` (issue #78 phase 4); it is +re-exported here lazily for one deprecation cycle. +""" -from .layers.gated_deltanet import GatedDeltaNet from .ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule +def __getattr__(name: str): + if name == "GatedDeltaNet": + from xorl.models.layers.gated_deltanet import GatedDeltaNet + + return GatedDeltaNet + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = [ "GatedDeltaNet", "chunk_gated_delta_rule", diff --git a/src/xorl/ops/linear_attention/layers/__init__.py b/src/xorl/ops/linear_attention/layers/__init__.py index b91b02dd..01709de3 100644 --- a/src/xorl/ops/linear_attention/layers/__init__.py +++ b/src/xorl/ops/linear_attention/layers/__init__.py @@ -1,4 +1,12 @@ -from .gated_deltanet import GatedDeltaNet +"""Moved: ``GatedDeltaNet`` lives in :mod:`xorl.models.layers.gated_deltanet` (issue #78 phase 4).""" + + +def __getattr__(name: str): + if name == "GatedDeltaNet": + from xorl.models.layers.gated_deltanet import GatedDeltaNet + + return GatedDeltaNet + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") __all__ = ["GatedDeltaNet"] diff --git a/src/xorl/ops/linear_attention/layers/gated_deltanet.py b/src/xorl/ops/linear_attention/layers/gated_deltanet.py index 5201a803..0eeb2f84 100644 --- a/src/xorl/ops/linear_attention/layers/gated_deltanet.py +++ b/src/xorl/ops/linear_attention/layers/gated_deltanet.py @@ -1,467 +1,4 @@ -from __future__ import annotations +"""Moved to ``xorl.models.layers.gated_deltanet`` (issue #78 phase 4); compat stub.""" -# Adapted from flash-linear-attention/fla/layers/gated_deltanet.py. -# Portions of this file are adapted from flash-linear-attention, Copyright (c) 2023-2025 Songlin Yang, licensed under the MIT License. -import math -import warnings -from typing import Any - -import torch -import torch.nn as nn -from einops import rearrange, repeat -from torch.nn import functional as F - -from xorl.lora.fold import lora_merged_forward_enabled -from xorl.ops.linear_attention.backend import ( - flashqla_chunk_gated_delta_rule, - flashqla_chunk_gated_delta_rule_cp, - get_gdn_backend, - warn_cp_fallback_once, -) -from xorl.ops.linear_attention.layers.utils import get_unpad_data, index_first_axis, pad_input -from xorl.ops.linear_attention.modules import ( - FusedRMSNormGated, - RMSNorm, - ShortConvolution, - causal_conv1d_qkv_contract, -) -from xorl.ops.linear_attention.modules.bi_contract import ( - _is_gdn_contract_enabled, - bi_fused_gdn_gating, - gdn_contract, -) -from xorl.ops.linear_attention.ops.gated_delta_rule import ( - chunk_gated_delta_rule, - fused_recurrent_gated_delta_rule, -) - - -def _sglang_compatible_beta_gate(b_input: torch.Tensor) -> torch.Tensor: - beta = b_input.float().sigmoid() - if beta.dtype != b_input.dtype: - beta = beta.to(dtype=b_input.dtype).float() - return beta - - -class GatedDeltaNet(nn.Module): - def __init__( - self, - hidden_size: int = 2048, - expand_v: float = 2, - head_dim: int = 256, - num_heads: int = 6, - num_v_heads: int | None = None, - mode: str = "chunk", - use_gate: bool = True, - use_short_conv: bool = True, - allow_neg_eigval: bool = False, - conv_size: int = 4, - conv_bias: bool = False, - layer_idx: int | None = None, - norm_eps: float = 1e-5, - exact_contract: bool = False, - **kwargs: Any, - ) -> None: - del kwargs - super().__init__() - - self.mode = mode - self.allow_neg_eigval = allow_neg_eigval - self.hidden_size = hidden_size - self.expand_v = expand_v - self.use_gate = use_gate - self.use_short_conv = use_short_conv - self.conv_size = conv_size - self.conv_bias = conv_bias - self.head_dim = head_dim - self.num_heads = num_heads - self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads - self.exact_contract = exact_contract - self.exact_merged_forward = exact_contract - - self.head_k_dim = head_dim - self.head_v_dim = int(self.head_dim * self.expand_v) - self.key_dim = int(self.num_heads * self.head_k_dim) - self.value_dim = int(self.num_v_heads * self.head_v_dim) - self.layer_idx = layer_idx - - if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): - raise ValueError( - f"expand_v={expand_v} does not produce an integer value when multiplied by " - f"num_v_heads * head_dim={self.num_v_heads * self.head_dim}." - ) - if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: - raise ValueError( - f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.", - ) - if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): - raise ValueError( - f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}.", - ) - if mode not in {"chunk", "fused_recurrent"}: - raise ValueError(f"Unsupported GatedDeltaNet mode: {mode}") - - self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) - self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) - self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) - self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) - self.b_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) - - A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16) - self.A_log = nn.Parameter(torch.log(A)) - self.A_log._no_weight_decay = True - - dt_min = 0.001 - dt_max = 0.1 - dt_init_floor = 1e-4 - dt = torch.exp(torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min)) - dt = torch.clamp(dt, min=dt_init_floor) - inv_dt = dt + torch.log(-torch.expm1(-dt)) - self.dt_bias = nn.Parameter(inv_dt) - self.dt_bias._no_weight_decay = True - - if use_short_conv: - self.q_conv1d = ShortConvolution( - hidden_size=self.key_dim, - kernel_size=conv_size, - bias=conv_bias, - activation="silu", - ) - self.k_conv1d = ShortConvolution( - hidden_size=self.key_dim, - kernel_size=conv_size, - bias=conv_bias, - activation="silu", - ) - self.v_conv1d = ShortConvolution( - hidden_size=self.value_dim, - kernel_size=conv_size, - bias=conv_bias, - activation="silu", - ) - else: - warnings.warn( - "ShortConvolution is usually important for GatedDeltaNet quality; " - "leave `use_short_conv=True` unless you know you want it disabled.", - stacklevel=2, - ) - - if use_gate: - self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) - self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) - else: - self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32) - self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) - - def _apply(self, fn, recurse: bool = True): - module = super()._apply(fn, recurse=recurse) - # Match SGLang's Qwen3.5/GDN contract: decay/time-step parameters stay fp32 - # even when the rest of the model is moved to bf16/fp16. - for name in ("A_log", "dt_bias"): - param = getattr(self, name, None) - if isinstance(param, nn.Parameter) and param.is_floating_point() and param.dtype != torch.float32: - with torch.no_grad(): - param.data = param.data.float() - if param.grad is not None: - param.grad.data = param.grad.data.float() - return module - - def _fused_qkvz_lora_delta( - self, hidden_states: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None: - """Apply River's fused q/k/v/z LoRA delta and split its output.""" - adapter = getattr(self, "in_proj_qkvz", None) - if adapter is None: - return None - delta = adapter(hidden_states) - expected = 2 * self.key_dim + 2 * self.value_dim - if delta.shape[-1] != expected: - raise RuntimeError(f"Fused GDN LoRA produced {delta.shape[-1]} features, expected {expected}") - return delta.split( - (self.key_dim, self.key_dim, self.value_dim, self.value_dim), - dim=-1, - ) - - def _add_output_lora(self, inputs: torch.Tensor, base_output: torch.Tensor) -> torch.Tensor: - """Add River's fused GDN output LoRA delta when one is injected.""" - adapter = getattr(self, "out_proj", None) - if adapter is None: - return base_output - return base_output + adapter(inputs).to(base_output.dtype) - - @staticmethod - def _linear_with_contract( - module: nn.Linear, - inputs: torch.Tensor, - weight: torch.Tensor, - ) -> torch.Tensor: - if getattr(module, "_xorl_bi_trunk_wrapped", False): - from xorl.ops.batch_invariant_ops import batch_invariant_trunk_linear # noqa: PLC0415 - - return batch_invariant_trunk_linear(inputs, weight, module.bias) - return F.linear(inputs, weight, module.bias) - - def _project_output_linear(self, o: torch.Tensor) -> torch.Tensor: - output_adapter = getattr(self, "out_proj", None) - if output_adapter is not None and lora_merged_forward_enabled(self): - folded = output_adapter.merged_weight_for_forward(self.o_proj.weight) - return self._linear_with_contract(self.o_proj, o, folded) - - return self._add_output_lora(o, self.o_proj(o)) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - past_key_values: Any | None = None, - use_cache: bool | None = False, - output_attentions: bool | None = False, - **kwargs: Any, - ) -> tuple[torch.Tensor, torch.Tensor | None, Any | None]: - with gdn_contract(self.exact_contract): - return self._forward_impl( - hidden_states, - attention_mask=attention_mask, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - **kwargs, - ) - - def _forward_impl( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - past_key_values: Any | None = None, - use_cache: bool | None = False, - output_attentions: bool | None = False, - **kwargs: Any, - ) -> tuple[torch.Tensor, torch.Tensor | None, Any | None]: - del output_attentions - if attention_mask is not None and len(attention_mask.shape) != 2: - raise ValueError( - "Expected `attention_mask` with shape [batch_size, seq_len] where 0 marks padding.", - ) - - batch_size, q_len, _ = hidden_states.shape - cp_context = kwargs.get("cp_context") - mode = ( - self.mode - if cp_context is not None - else ("fused_recurrent" if (q_len <= 64 and not self.training) else self.mode) - ) - if self.training and mode != "chunk": - raise AssertionError("Only chunk mode is supported in training.") - - last_state = None - if past_key_values is not None and self.layer_idx is not None and len(past_key_values) > self.layer_idx: - last_state = past_key_values[self.layer_idx] - - cu_seqlens = kwargs.get("cu_seqlens") - indices = None - if cp_context is not None: - if attention_mask is not None: - raise ValueError( - "Ulysses linear attention currently requires packed inputs without a 2D attention_mask.", - ) - if cp_context.cu_seqlens is None: - raise ValueError( - "Ulysses linear attention requires cu_seqlens metadata from the collator.", - ) - if use_cache: - raise ValueError( - "Ulysses native FLA CP does not yet support KV/conv cache updates.", - ) - if mode != "chunk": - raise ValueError("Ulysses native FLA CP currently supports chunk mode only.") - cu_seqlens = cp_context.cu_seqlens - elif attention_mask is not None: - indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) - hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) - - merged_input_adapter = getattr(self, "in_proj_qkvz", None) - if merged_input_adapter is not None and lora_merged_forward_enabled(self): - offsets = [0] - for projection in (self.q_proj, self.k_proj, self.v_proj, self.g_proj): - offsets.append(offsets[-1] + projection.out_features) - q_weight = merged_input_adapter.merged_weight_for_forward( - self.q_proj.weight, output_start=offsets[0], output_end=offsets[1] - ) - k_weight = merged_input_adapter.merged_weight_for_forward( - self.k_proj.weight, output_start=offsets[1], output_end=offsets[2] - ) - v_weight = merged_input_adapter.merged_weight_for_forward( - self.v_proj.weight, output_start=offsets[2], output_end=offsets[3] - ) - gate_weight = merged_input_adapter.merged_weight_for_forward( - self.g_proj.weight, output_start=offsets[3], output_end=offsets[4] - ) - q_input = self._linear_with_contract(self.q_proj, hidden_states, q_weight) - k_input = self._linear_with_contract(self.k_proj, hidden_states, k_weight) - v_input = self._linear_with_contract(self.v_proj, hidden_states, v_weight) - gate_input = self._linear_with_contract(self.g_proj, hidden_states, gate_weight) - a_input = self.a_proj(hidden_states).float() - b_input = self.b_proj(hidden_states) - else: - q_input = self.q_proj(hidden_states) - k_input = self.k_proj(hidden_states) - v_input = self.v_proj(hidden_states) - a_input = self.a_proj(hidden_states).float() - b_input = self.b_proj(hidden_states) - gate_input = self.g_proj(hidden_states) if self.use_gate else None - - fused_lora_delta = None if lora_merged_forward_enabled(self) else self._fused_qkvz_lora_delta(hidden_states) - if fused_lora_delta is not None: - q_delta, k_delta, v_delta, gate_delta = fused_lora_delta - q_input = q_input + q_delta.to(q_input.dtype) - k_input = k_input + k_delta.to(k_input.dtype) - v_input = v_input + v_delta.to(v_input.dtype) - if gate_input is None: - raise RuntimeError("Fused qkvz LoRA requires a gated GDN layer") - gate_input = gate_input + gate_delta.to(gate_input.dtype) - - if _is_gdn_contract_enabled() and not self.use_short_conv: - raise RuntimeError("Exact Qwen3.5 GDN requires short convolution") - - if self.use_short_conv and _is_gdn_contract_enabled(): - if use_cache or last_state is not None: - raise RuntimeError( - "Exact Qwen3.5 trainer GDN supports packed prefill only, not recurrent cache updates" - ) - if cp_context is None: - # Fail-closed backstop against the ring/silent-skip fail-open: - # build_linear_attention_cp_context returns None for ring>1 - # and for missing metadata — an exact GDN layer must never - # run U1 math on sequence-sharded rows. - from xorl.distributed.parallel_state import get_parallel_state # noqa: PLC0415 - - ps = get_parallel_state() - if ps.ulysses_size > 1 or ps.ringattn_size > 1: - raise RuntimeError( - "Exact Qwen3.5 GDN: sequence parallelism is active " - f"(ulysses={ps.ulysses_size}, ring={ps.ringattn_size}) but no " - "cp_context reached the layer. Ring is unsupported (contract C4) " - "and a missing Ulysses context would silently break bytes. " - "Fail closed.", - ) - q, k, v = causal_conv1d_qkv_contract( - q_input, - k_input, - v_input, - self.q_conv1d, - self.k_conv1d, - self.v_conv1d, - cu_seqlens=cu_seqlens, - cp_context=cp_context, - ) - conv_state_q = conv_state_k = conv_state_v = None - elif self.use_short_conv: - conv_state_q = conv_state_k = conv_state_v = None - if last_state is not None: - conv_state_q, conv_state_k, conv_state_v = last_state["conv_state"] - q, conv_state_q = self.q_conv1d( - x=q_input, - cache=conv_state_q, - output_final_state=bool(use_cache), - cu_seqlens=cu_seqlens, - cp_context=cp_context, - ) - k, conv_state_k = self.k_conv1d( - x=k_input, - cache=conv_state_k, - output_final_state=bool(use_cache), - cu_seqlens=cu_seqlens, - cp_context=cp_context, - ) - v, conv_state_v = self.v_conv1d( - x=v_input, - cache=conv_state_v, - output_final_state=bool(use_cache), - cu_seqlens=cu_seqlens, - cp_context=cp_context, - ) - else: - q = F.silu(q_input) - k = F.silu(k_input) - v = F.silu(v_input) - conv_state_q = conv_state_k = conv_state_v = None - - q, k = (rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim) for x in (q, k)) - v = rearrange(v, "... (h d) -> ... h d", d=self.head_v_dim) - - if self.num_v_heads > self.num_heads: - repeat_factor = self.num_v_heads // self.num_heads - q, k = (repeat(x, "... h d -> ... (h g) d", g=repeat_factor) for x in (q, k)) - - if _is_gdn_contract_enabled(): - # The serving fused-GDN gating kernel removes the - # 1-ULP g term (torch softplus vs tl.log(1+tl.exp)); beta is bitwise - # either way. - g, beta = bi_fused_gdn_gating(self.A_log, a_input, b_input, self.dt_bias) - else: - beta = _sglang_compatible_beta_gate(b_input) - g = -self.A_log.float().exp() * F.softplus(a_input + self.dt_bias) - if self.allow_neg_eigval: - beta = beta * 2.0 - recurrent_state = last_state["recurrent_state"] if last_state is not None else None - - if mode == "chunk": - backend = get_gdn_backend() - chunk_fn = chunk_gated_delta_rule - if backend == "flashqla": - if self.head_k_dim != 128 or self.head_v_dim != 128: - warn_cp_fallback_once() - elif cp_context is not None: - chunk_fn = flashqla_chunk_gated_delta_rule_cp - else: - chunk_fn = flashqla_chunk_gated_delta_rule - - chunk_kwargs = dict( - q=q, - k=k, - v=v, - g=g, - beta=beta, - initial_state=recurrent_state, - output_final_state=bool(use_cache), - cu_seqlens=cu_seqlens, - use_qk_l2norm_in_kernel=True, - ) - if cp_context is not None: - chunk_kwargs["cp_context"] = cp_context - o, recurrent_state = chunk_fn(**chunk_kwargs) - - elif mode == "fused_recurrent": - o, recurrent_state = fused_recurrent_gated_delta_rule( - q=q, - k=k, - v=v, - g=g, - beta=beta, - initial_state=recurrent_state, - output_final_state=bool(use_cache), - cu_seqlens=cu_seqlens, - use_qk_l2norm_in_kernel=True, - ) - else: - raise NotImplementedError(f"Unsupported mode `{mode}`.") - - if past_key_values is not None and self.layer_idx is not None: - past_key_values.update( - recurrent_state=recurrent_state, - conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, - layer_idx=self.layer_idx, - offset=q_len, - ) - - if self.use_gate: - gate = rearrange(gate_input, "... (h d) -> ... h d", d=self.head_v_dim) - o = self.o_norm(o, gate) - else: - o = self.o_norm(o) - o = rearrange(o, "b t h d -> b t (h d)") - o = self._project_output_linear(o) - if attention_mask is not None and indices is not None: - o = pad_input(o.squeeze(0), indices, batch_size, q_len) - - return o, None, past_key_values +from xorl.models.layers.gated_deltanet import * # noqa: F401,F403 +from xorl.models.layers.gated_deltanet import GatedDeltaNet # noqa: F401 diff --git a/src/xorl/ops/ssm/__init__.py b/src/xorl/ops/ssm/__init__.py index 9aedcbd5..19485736 100644 --- a/src/xorl/ops/ssm/__init__.py +++ b/src/xorl/ops/ssm/__init__.py @@ -1,11 +1,23 @@ -"""Mamba2 (SSD) ops and layers used by Nemotron-H.""" +"""Mamba2 (SSD) kernels used by Nemotron-H. + +The ``Mamba2Mixer`` layer class moved to +:mod:`xorl.models.layers.mamba2_mixer` (issue #78 phase 4); it is re-exported +here lazily for one deprecation cycle. +""" -from .layers.mamba2_mixer import Mamba2Mixer from .modules.gated_norm import GroupRMSNormGated from .ops.conv import causal_depthwise_conv1d from .ops.ssd import ssd_chunked +def __getattr__(name: str): + if name == "Mamba2Mixer": + from xorl.models.layers.mamba2_mixer import Mamba2Mixer + + return Mamba2Mixer + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = [ "GroupRMSNormGated", "Mamba2Mixer", diff --git a/src/xorl/ops/ssm/layers/__init__.py b/src/xorl/ops/ssm/layers/__init__.py index 5bf08160..8eadc5e3 100644 --- a/src/xorl/ops/ssm/layers/__init__.py +++ b/src/xorl/ops/ssm/layers/__init__.py @@ -1,6 +1,12 @@ -from .mamba2_mixer import Mamba2Mixer +"""Moved: ``Mamba2Mixer`` lives in :mod:`xorl.models.layers.mamba2_mixer` (issue #78 phase 4).""" -__all__ = [ - "Mamba2Mixer", -] +def __getattr__(name: str): + if name == "Mamba2Mixer": + from xorl.models.layers.mamba2_mixer import Mamba2Mixer + + return Mamba2Mixer + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = ["Mamba2Mixer"] diff --git a/tests/distributed/test_dsv4_exact_cp_attention_layout.py b/tests/distributed/test_dsv4_exact_cp_attention_layout.py index 266aec5f..8e158965 100644 --- a/tests/distributed/test_dsv4_exact_cp_attention_layout.py +++ b/tests/distributed/test_dsv4_exact_cp_attention_layout.py @@ -21,7 +21,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src")) from xorl.data.collators.sequence_shard_collator import zigzag_reorder_packed_sequence # noqa: E402 -from xorl.ops.dsv4.cp_utils import ( # noqa: E402 +from xorl.ops.families.dsv4.cp_utils import ( # noqa: E402 build_dsv4_exact_cp_layout, gather_dsv4_exact_cp_rows, ) diff --git a/tests/distributed/test_linear_attention_cp_equivalence.py b/tests/distributed/test_linear_attention_cp_equivalence.py index d3dbdca2..cf8f669e 100644 --- a/tests/distributed/test_linear_attention_cp_equivalence.py +++ b/tests/distributed/test_linear_attention_cp_equivalence.py @@ -11,7 +11,7 @@ import torch.distributed as dist from xorl.distributed.parallel_state import init_parallel_state -from xorl.ops.linear_attention import GatedDeltaNet +from xorl.models.layers.gated_deltanet import GatedDeltaNet from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils.device import get_nccl_backend diff --git a/tests/models/test_dsv4_attention.py b/tests/models/test_dsv4_attention.py index 4921718c..42802ff4 100644 --- a/tests/models/test_dsv4_attention.py +++ b/tests/models/test_dsv4_attention.py @@ -26,7 +26,7 @@ def _small_rope_buffer(monkeypatch): monkeypatch.setenv("XORL_DSV4_SPARSE_ATTN_IMPL", "sparse") # pure-torch ref # Clear the @lru_cache on precompute_freqs_cis so cross-test state with # different (dim, seqlen, factor, ...) keys doesn't leak. - from xorl.ops.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 + from xorl.ops.families.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 precompute_freqs_cis.cache_clear() yield @@ -135,7 +135,7 @@ def test_attn_sink_is_fp32_and_keep_fp32_marked(): def test_kv_qat_enabled_from_quantization_config(): """Xorl mirrors Miles' config-driven FP8-QAT gate instead of an env toggle.""" from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepSeekV4Attention # noqa: PLC0415 - from xorl.ops.dsv4.utils import dsv4_kv_qat_enabled # noqa: PLC0415 + from xorl.ops.families.dsv4.utils import dsv4_kv_qat_enabled # noqa: PLC0415 cfg = _tiny_config(compress_ratios=[0]) assert dsv4_kv_qat_enabled(cfg) is False @@ -215,7 +215,7 @@ def test_exact_cp_selects_the_serving_kv_boundary(monkeypatch, compress_ratio, c from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import ( # noqa: PLC0415 DeepSeekV4Attention, ) - from xorl.ops.dsv4 import exact_attention # noqa: PLC0415 + from xorl.ops.families.dsv4 import exact_attention # noqa: PLC0415 class _FakeCPGroup: def size(self): @@ -303,8 +303,8 @@ def test_exact_ring_cp_restores_gathered_rows_and_uses_local_rope_positions(monk from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import ( # noqa: PLC0415 DeepSeekV4Attention, ) - from xorl.ops.dsv4 import exact_attention # noqa: PLC0415 - from xorl.ops.dsv4.cp_utils import Dsv4ExactCPLayout # noqa: PLC0415 + from xorl.ops.families.dsv4 import exact_attention # noqa: PLC0415 + from xorl.ops.families.dsv4.cp_utils import Dsv4ExactCPLayout # noqa: PLC0415 class _FakeCPGroup: @staticmethod @@ -404,8 +404,8 @@ def test_exact_packed_requests_reset_c0_c4_c128_state(monkeypatch, compress_rati from xorl.models.transformers.deepseek_v4 import modeling_deepseek_v4 # noqa: PLC0415 from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepSeekV4Attention # noqa: PLC0415 - from xorl.ops.dsv4 import exact_attention # noqa: PLC0415 - from xorl.ops.dsv4.cp_utils import build_dsv4_exact_cp_layout # noqa: PLC0415 + from xorl.ops.families.dsv4 import exact_attention # noqa: PLC0415 + from xorl.ops.families.dsv4.cp_utils import build_dsv4_exact_cp_layout # noqa: PLC0415 cfg = _tiny_config(compress_ratios=[compress_ratio]) cfg._dsv4_flash_exact_mode = True diff --git a/tests/models/test_dsv4_exact_contract.py b/tests/models/test_dsv4_exact_contract.py index 11216c6c..735e9f48 100644 --- a/tests/models/test_dsv4_exact_contract.py +++ b/tests/models/test_dsv4_exact_contract.py @@ -28,7 +28,7 @@ bind_dsv4_exact_lm_head, dsv4_lm_head_shard, ) -from xorl.ops.dsv4.exact_attention import ( +from xorl.ops.families.dsv4.exact_attention import ( _causal_window_indices, _hybrid_indices_for_positions, _hybrid_prefill_indices, diff --git a/tests/models/test_dsv4_hf_to_dcp.py b/tests/models/test_dsv4_hf_to_dcp.py index 411381db..a517d959 100644 --- a/tests/models/test_dsv4_hf_to_dcp.py +++ b/tests/models/test_dsv4_hf_to_dcp.py @@ -25,7 +25,7 @@ @pytest.fixture(autouse=True) def _clean_pg_and_cache(): - from xorl.ops.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 + from xorl.ops.families.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 precompute_freqs_cis.cache_clear() yield diff --git a/tests/models/test_dsv4_loader.py b/tests/models/test_dsv4_loader.py index 59a0cb15..b5aee7a0 100644 --- a/tests/models/test_dsv4_loader.py +++ b/tests/models/test_dsv4_loader.py @@ -28,7 +28,7 @@ @pytest.fixture(autouse=True) def _clear_rope_cache(): - from xorl.ops.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 + from xorl.ops.families.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 precompute_freqs_cis.cache_clear() yield diff --git a/tests/models/test_dsv4_lora.py b/tests/models/test_dsv4_lora.py index 76ca652e..96bb7b57 100644 --- a/tests/models/test_dsv4_lora.py +++ b/tests/models/test_dsv4_lora.py @@ -18,7 +18,7 @@ def _cpu_env(monkeypatch): monkeypatch.setenv("XORL_DSV4_ROPE_MAX_SEQ_LEN", "256") monkeypatch.setenv("XORL_DSV4_SPARSE_ATTN_IMPL", "sparse") - from xorl.ops.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 + from xorl.ops.families.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 precompute_freqs_cis.cache_clear() yield diff --git a/tests/models/test_dsv4_model.py b/tests/models/test_dsv4_model.py index 1019d6bd..4389410c 100644 --- a/tests/models/test_dsv4_model.py +++ b/tests/models/test_dsv4_model.py @@ -24,7 +24,7 @@ def _cpu_env(monkeypatch): monkeypatch.setenv("XORL_DSV4_ROPE_MAX_SEQ_LEN", "256") monkeypatch.setenv("XORL_DSV4_SPARSE_ATTN_IMPL", "sparse") - from xorl.ops.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 + from xorl.ops.families.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 precompute_freqs_cis.cache_clear() yield diff --git a/tests/models/test_fused_gdn_lora.py b/tests/models/test_fused_gdn_lora.py index e2676a1d..04c0f5ae 100644 --- a/tests/models/test_fused_gdn_lora.py +++ b/tests/models/test_fused_gdn_lora.py @@ -24,7 +24,7 @@ load_lora_checkpoint, save_lora_checkpoint, ) -from xorl.ops.linear_attention.layers.gated_deltanet import GatedDeltaNet +from xorl.models.layers.gated_deltanet import GatedDeltaNet class _SharedExpert(nn.Module): diff --git a/tests/models/test_glm5_flashmla_sparse_mla.py b/tests/models/test_glm5_flashmla_sparse_mla.py index fe02eb6c..51e15185 100644 --- a/tests/models/test_glm5_flashmla_sparse_mla.py +++ b/tests/models/test_glm5_flashmla_sparse_mla.py @@ -11,8 +11,8 @@ _flatten_sparse_mla_inputs, sparse_mla_dispatch, ) -from xorl.ops.glm5_kernels import flashmla_sparse_mla -from xorl.ops.glm5_kernels.flashmla_sparse_mla import FlashMLASparseWithTileLangBackward +from xorl.ops.families.glm5 import flashmla_sparse_mla +from xorl.ops.families.glm5.flashmla_sparse_mla import FlashMLASparseWithTileLangBackward pytestmark = [pytest.mark.cpu] diff --git a/tests/models/test_glm5_support.py b/tests/models/test_glm5_support.py index 85948c25..09db54b2 100644 --- a/tests/models/test_glm5_support.py +++ b/tests/models/test_glm5_support.py @@ -366,7 +366,7 @@ def _assert_indexer_padding_mask_policy(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_indexer_tilelang_fast_path_accepts_causal_mask_and_matches_blocked_torch(): """Both fast-path mask forms match an independently forced torch path.""" - from xorl.ops.glm5_kernels.tilelang_indexer_fwd import tl_indexer_fwd_impl # noqa: F401 + from xorl.ops.families.glm5.tilelang_indexer_fwd import tl_indexer_fwd_impl # noqa: F401 torch.manual_seed(0) config = _tiny_config(index_topk=8, index_n_heads=4) @@ -624,7 +624,7 @@ def apply(q_flat, kv_flat, indices_flat, scaling): monkeypatch.setitem( sys.modules, - "xorl.ops.glm5_kernels.sparse_mla", + "xorl.ops.families.glm5.sparse_mla", SimpleNamespace(SparseMLA=FakeSparseMLA), ) diff --git a/tests/models/test_qwen35_lora_projection_topology.py b/tests/models/test_qwen35_lora_projection_topology.py index cc088fcc..07720047 100644 --- a/tests/models/test_qwen35_lora_projection_topology.py +++ b/tests/models/test_qwen35_lora_projection_topology.py @@ -17,9 +17,9 @@ load_lora_checkpoint, save_lora_checkpoint, ) +from xorl.models.layers.gated_deltanet import GatedDeltaNet from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeMLP from xorl.ops.batch_invariant_ops import set_trunk_linear_contract, wrap_trunk_linears_batch_invariant -from xorl.ops.linear_attention.layers.gated_deltanet import GatedDeltaNet RANK = 16 diff --git a/tests/ops/dsv4/test_compressor.py b/tests/ops/dsv4/test_compressor.py index 433deb7e..ebd2be3c 100644 --- a/tests/ops/dsv4/test_compressor.py +++ b/tests/ops/dsv4/test_compressor.py @@ -5,10 +5,10 @@ import pytest import torch -from xorl.ops.dsv4 import utils -from xorl.ops.dsv4.compressor import DeepSeekV4Compressor -from xorl.ops.dsv4.exact_attention import _legacy_compressor_state_pages -from xorl.ops.dsv4.rope import precompute_freqs_cis +from xorl.ops.families.dsv4 import utils +from xorl.ops.families.dsv4.compressor import DeepSeekV4Compressor +from xorl.ops.families.dsv4.exact_attention import _legacy_compressor_state_pages +from xorl.ops.families.dsv4.rope import precompute_freqs_cis pytestmark = pytest.mark.cpu diff --git a/tests/ops/dsv4/test_exact_attention.py b/tests/ops/dsv4/test_exact_attention.py index 7e043b72..fea38efd 100644 --- a/tests/ops/dsv4/test_exact_attention.py +++ b/tests/ops/dsv4/test_exact_attention.py @@ -10,7 +10,7 @@ def _exact_attention_module(): - from xorl.ops.dsv4 import exact_attention # noqa: PLC0415 + from xorl.ops.families.dsv4 import exact_attention # noqa: PLC0415 return exact_attention diff --git a/tests/ops/dsv4/test_v4_tilelang_indexer.py b/tests/ops/dsv4/test_v4_tilelang_indexer.py index 06cac426..b190911c 100644 --- a/tests/ops/dsv4/test_v4_tilelang_indexer.py +++ b/tests/ops/dsv4/test_v4_tilelang_indexer.py @@ -161,7 +161,7 @@ def make_inputs(seqlen_q, batch, heads, dim, compress_ratio, device="cuda"): @pytest.mark.parametrize("seqlen_q,batch,heads,dim,compress_ratio,topk", FORWARD_CONFIGS, ids=FORWARD_CONFIG_IDS) def test_indexer_forward_scores(seqlen_q, batch, heads, dim, compress_ratio, topk): """Compare tilelang forward logits against PyTorch reference.""" - from xorl.ops.dsv4.kernel.tilelang_indexer_fwd import ( + from xorl.ops.families.dsv4.kernel.tilelang_indexer_fwd import ( _make_causal_cu_seqlens, batched_indexer_fwd, ) @@ -204,7 +204,7 @@ def test_indexer_forward_scores(seqlen_q, batch, heads, dim, compress_ratio, top # --------------------------------------------------------------------------- def _assert_large_values(): """Test with large input values to check for overflow/underflow.""" - from xorl.ops.dsv4.kernel.tilelang_indexer_fwd import ( + from xorl.ops.families.dsv4.kernel.tilelang_indexer_fwd import ( _make_causal_cu_seqlens, batched_indexer_fwd, ) @@ -238,7 +238,7 @@ def _assert_large_values(): # --------------------------------------------------------------------------- def _assert_zero_inputs(): """Test that zero inputs produce zero scores.""" - from xorl.ops.dsv4.kernel.tilelang_indexer_fwd import ( + from xorl.ops.families.dsv4.kernel.tilelang_indexer_fwd import ( _make_causal_cu_seqlens, batched_indexer_fwd, ) diff --git a/tests/ops/dsv4/test_v4_tilelang_sparse_mla.py b/tests/ops/dsv4/test_v4_tilelang_sparse_mla.py index 3bd7d896..978bd502 100644 --- a/tests/ops/dsv4/test_v4_tilelang_sparse_mla.py +++ b/tests/ops/dsv4/test_v4_tilelang_sparse_mla.py @@ -181,7 +181,7 @@ def make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk, device="cuda", sink_ @pytest.mark.parametrize("batch,seqlen,heads,dim,seqlen_kv,topk", FORWARD_CONFIGS, ids=FORWARD_IDS) def test_sparse_mla_forward_policy(batch, seqlen, heads, dim, seqlen_kv, topk): """Compare tilelang sparse MLA forward against PyTorch reference.""" - from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 + from xorl.ops.families.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 q, kv, attn_sink, topk_idxs = make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk) sm_scale = (1.0 / dim) ** 0.5 @@ -205,7 +205,7 @@ def test_sparse_mla_forward_policy(batch, seqlen, heads, dim, seqlen_kv, topk): # --------------------------------------------------------------------------- def _assert_attn_sink_policy(): """Test attn_sink reference parity and prove that the kernel does not ignore it.""" - from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 + from xorl.ops.families.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 # Zero is the boundary; random contains both signs. Positive-only and # negative-only tensors use the identical arithmetic, while the separate @@ -233,7 +233,7 @@ def _assert_attn_sink_policy(): def _assert_attn_sink_changes_output(sparse_mqa_fwd_interface): """Verify attn_sink actually changes output (not ignored).""" - from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 + from xorl.ops.families.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 batch, seqlen, heads, dim, seqlen_kv, topk = 1, 128, 8, 512, 160, 64 q, kv, _, topk_idxs = make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk) @@ -309,7 +309,7 @@ def ref_dense_attn_with_grad(q, kv, attn_sink, topk_idxs, sm_scale): @pytest.mark.parametrize("batch,seqlen,heads,dim,seqlen_kv,topk", BACKWARD_CONFIGS, ids=BACKWARD_IDS) def test_sparse_mla_backward_policy(batch, seqlen, heads, dim, seqlen_kv, topk): """Compare tilelang backward gradients against PyTorch autograd reference.""" - from xorl.ops.dsv4.attention_core import sparse_attn_tilelang # noqa: PLC0415 + from xorl.ops.families.dsv4.attention_core import sparse_attn_tilelang # noqa: PLC0415 q_base, kv_base, attn_sink_base, topk_idxs = make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk) sm_scale = (1.0 / dim) ** 0.5 @@ -359,8 +359,8 @@ def test_sparse_mla_backward_deterministic_dkv(): atomic path within bf16 GEMM tolerance, with no NaN.""" import os as _os # noqa: PLC0415 - from xorl.ops.dsv4.kernel.tilelang_sparse_mla_bwd import sparse_mqa_bwd_interface # noqa: PLC0415 - from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 + from xorl.ops.families.dsv4.kernel.tilelang_sparse_mla_bwd import sparse_mqa_bwd_interface # noqa: PLC0415 + from xorl.ops.families.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 batch, seqlen, heads, dim, seqlen_kv, topk = 1, 256, 64, 512, 320, 128 q, kv, attn_sink, topk_idxs = make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk) @@ -392,7 +392,7 @@ def test_sparse_mla_backward_deterministic_dkv(): @requires_tilelang() def test_sparse_mla_partial_invalid_indices_forward_backward_policy(): """Forward and backward must ignore -1 sparse slots without touching invalid dKV rows.""" - from xorl.ops.dsv4.attention_core import sparse_attn_tilelang # noqa: PLC0415 + from xorl.ops.families.dsv4.attention_core import sparse_attn_tilelang # noqa: PLC0415 batch, seqlen, heads, dim, seqlen_kv, topk = 1, 128, 8, 512, 160, 64 q_base, kv_base, attn_sink_base, topk_idxs = make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk) @@ -427,7 +427,7 @@ def test_sparse_mla_partial_invalid_indices_forward_backward_policy(): # --------------------------------------------------------------------------- def _assert_partial_invalid_forward_interface(): """Test with some indices set to -1 (invalid).""" - from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 + from xorl.ops.families.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 batch, seqlen, heads, dim, seqlen_kv, topk = 1, 256, 8, 512, 320, 128 q, kv, attn_sink, topk_idxs = make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk) diff --git a/tests/ops/test_bi_gdn_contract.py b/tests/ops/test_bi_gdn_contract.py index 2e5f8e45..23d068d7 100644 --- a/tests/ops/test_bi_gdn_contract.py +++ b/tests/ops/test_bi_gdn_contract.py @@ -137,7 +137,7 @@ def _assert_fused_rms_norm_gated_module_routes_under_exact_model_program(): def _assert_gated_deltanet_gating_routes_under_exact_model_program(): - from xorl.ops.linear_attention.layers.gated_deltanet import GatedDeltaNet # noqa: PLC0415 + from xorl.models.layers.gated_deltanet import GatedDeltaNet # noqa: PLC0415 ordinary = ( GatedDeltaNet( diff --git a/tests/ops/test_flashqla_contract_pin.py b/tests/ops/test_flashqla_contract_pin.py index b8aac332..e3475688 100644 --- a/tests/ops/test_flashqla_contract_pin.py +++ b/tests/ops/test_flashqla_contract_pin.py @@ -18,8 +18,8 @@ import pytest import torch -import xorl.ops.linear_attention.layers.gated_deltanet as gated_deltanet -from xorl.ops.linear_attention import GatedDeltaNet +import xorl.models.layers.gated_deltanet as gated_deltanet +from xorl.models.layers.gated_deltanet import GatedDeltaNet from xorl.ops.linear_attention.backend import FLASHQLA_AUTOCP_ENV, resolve_flashqla_auto_cp from xorl.ops.linear_attention.modules.bi_contract import gdn_contract diff --git a/tests/ops/test_gdn_conv_contract.py b/tests/ops/test_gdn_conv_contract.py index ec32c936..a1187318 100644 --- a/tests/ops/test_gdn_conv_contract.py +++ b/tests/ops/test_gdn_conv_contract.py @@ -6,8 +6,8 @@ import torch.nn.functional as F from torch.utils.checkpoint import checkpoint -import xorl.ops.linear_attention.layers.gated_deltanet as gated_deltanet -from xorl.ops.linear_attention import GatedDeltaNet +import xorl.models.layers.gated_deltanet as gated_deltanet +from xorl.models.layers.gated_deltanet import GatedDeltaNet from xorl.ops.linear_attention.modules import ShortConvolution, causal_conv1d_qkv_contract from xorl.ops.linear_attention.modules.bi_contract import _is_gdn_contract_enabled from xorl.ops.linear_attention.modules.conv_contract import _pack_conv_weight diff --git a/tests/ops/test_sparse_mla_kernels.py b/tests/ops/test_sparse_mla_kernels.py index ca3e4a6b..3ca20d94 100644 --- a/tests/ops/test_sparse_mla_kernels.py +++ b/tests/ops/test_sparse_mla_kernels.py @@ -57,7 +57,7 @@ def _make_inputs(S, S_kv, H, D, tail, topk, *, device, seed=1234): @pytest.mark.skipif(not _have_cuda_h100(), reason="needs H100+") def test_sparse_mla_fwd_kernel_matches_torch_reference(): - from xorl.ops.glm5_kernels.tilelang_sparse_mla_fwd import sparse_mla_fwd_interface + from xorl.ops.families.glm5.tilelang_sparse_mla_fwd import sparse_mla_fwd_interface S, S_kv, H, D, tail, topk = 256, 1024, 64, 512, 64, 128 sm_scale = (D + tail) ** -0.5 @@ -82,7 +82,7 @@ def _assert_sparse_mla_bwd_kernel_produces_finite_gradients(): (b) dq matches torch to bf16 attention tolerance (c) dkv max-abs error < 1× dkv_ref_max (signal-to-noise > 1) """ - from xorl.ops.glm5_kernels.sparse_mla import SparseMLA + from xorl.ops.families.glm5.sparse_mla import SparseMLA S, S_kv, H, D, tail, topk = 256, 1024, 64, 512, 64, 128 sm_scale = (D + tail) ** -0.5 @@ -130,8 +130,8 @@ def test_sparse_mla_bwd_reference_and_deterministic_policy(): import os as _os # noqa: PLC0415 - from xorl.ops.glm5_kernels.tilelang_sparse_mla_bwd import sparse_mla_bwd # noqa: PLC0415 - from xorl.ops.glm5_kernels.tilelang_sparse_mla_fwd import sparse_mla_fwd_interface # noqa: PLC0415 + from xorl.ops.families.glm5.tilelang_sparse_mla_bwd import sparse_mla_bwd # noqa: PLC0415 + from xorl.ops.families.glm5.tilelang_sparse_mla_fwd import sparse_mla_fwd_interface # noqa: PLC0415 S, S_kv, H, D, tail, topk = 256, 1024, 64, 512, 64, 128 sm_scale = (D + tail) ** -0.5 diff --git a/tests/ops/test_ssm_mamba2.py b/tests/ops/test_ssm_mamba2.py index 1e3a3a6b..361cbd72 100644 --- a/tests/ops/test_ssm_mamba2.py +++ b/tests/ops/test_ssm_mamba2.py @@ -15,7 +15,8 @@ from transformers.models.nemotron_h.configuration_nemotron_h import NemotronHConfig from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHMamba2Mixer -from xorl.ops.ssm import Mamba2Mixer, causal_depthwise_conv1d, ssd_chunked +from xorl.models.layers.mamba2_mixer import Mamba2Mixer +from xorl.ops.ssm import causal_depthwise_conv1d, ssd_chunked pytestmark = pytest.mark.cpu From 2a505ce994f56773791441cd1f87854870bf87fa Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 03:54:48 +0000 Subject: [PATCH 05/14] Delete the dead ops/ep_kernels package (#78 phase 5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every export (build_sorted_indices, group_tokens_by_expert{,_v2}, DeepEPScatter, DeepEPWeightedGather, deepep_scatter, deepep_weighted_gather) has zero importers in src/, tests/, benchmarks/, or certification/ — the DeepEP dispatch paths in ops/moe grew their own implementations. The #78 plan said 'merge into moe/'; deleting dead code beats moving it. --- src/xorl/ops/README.md | 1 - src/xorl/ops/ep_kernels/__init__.py | 22 --- .../ops/ep_kernels/deepep_counting_sort.py | 109 -------------- .../ops/ep_kernels/deepep_scatter_gather.py | 142 ------------------ 4 files changed, 274 deletions(-) delete mode 100644 src/xorl/ops/ep_kernels/__init__.py delete mode 100644 src/xorl/ops/ep_kernels/deepep_counting_sort.py delete mode 100644 src/xorl/ops/ep_kernels/deepep_scatter_gather.py diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index a38533de..f5ecacbd 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -26,7 +26,6 @@ Three kinds of code live here with three different rules: objective functions live in `xorl/objectives/` (#78 phase 2); old module paths here are compat stubs for one deprecation cycle. - `moe/` — MoE expert compute backends (triton/quack/native, LoRA variants). - `ep_kernels/` (DeepEP sort/scatter) merges in here (#78 phase 5). - `linear_attention/` — GDN/linear-attention kernels. The `GatedDeltaNet` layer class lives in `models/layers/gated_deltanet.py` (#78 phase 4); the old paths re-export it lazily for one deprecation cycle. diff --git a/src/xorl/ops/ep_kernels/__init__.py b/src/xorl/ops/ep_kernels/__init__.py deleted file mode 100644 index ba136a32..00000000 --- a/src/xorl/ops/ep_kernels/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -from .deepep_counting_sort import ( - build_sorted_indices, - group_tokens_by_expert, - group_tokens_by_expert_v2, -) -from .deepep_scatter_gather import ( - DeepEPScatter, - DeepEPWeightedGather, - deepep_scatter, - deepep_weighted_gather, -) - - -__all__ = [ - "deepep_scatter", - "deepep_weighted_gather", - "DeepEPScatter", - "DeepEPWeightedGather", - "group_tokens_by_expert", - "group_tokens_by_expert_v2", - "build_sorted_indices", -] diff --git a/src/xorl/ops/ep_kernels/deepep_counting_sort.py b/src/xorl/ops/ep_kernels/deepep_counting_sort.py deleted file mode 100644 index aefaf9d2..00000000 --- a/src/xorl/ops/ep_kernels/deepep_counting_sort.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Optimized sorting and grouping operations for DeepEP.""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _build_sorted_indices_kernel( - sorted_token_idx_ptr, - sorted_k_idx_ptr, - sorted_flat_indices_ptr, - num_valid: tl.constexpr, - topk: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - idx = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = idx < num_valid - flat_idx = tl.load(sorted_flat_indices_ptr + idx, mask=mask, other=0) - token_idx = flat_idx // topk - k_idx = flat_idx % topk - tl.store(sorted_token_idx_ptr + idx, token_idx, mask=mask) - tl.store(sorted_k_idx_ptr + idx, k_idx, mask=mask) - - -def build_sorted_indices( - sorted_flat_indices: torch.Tensor, - topk: int, -) -> tuple[torch.Tensor, torch.Tensor]: - num_valid = sorted_flat_indices.shape[0] - device = sorted_flat_indices.device - if num_valid == 0: - return ( - torch.empty(0, dtype=torch.int64, device=device), - torch.empty(0, dtype=torch.int64, device=device), - ) - - sorted_token_idx = torch.empty(num_valid, dtype=torch.int64, device=device) - sorted_k_idx = torch.empty(num_valid, dtype=torch.int64, device=device) - BLOCK_SIZE = 1024 - grid = ((num_valid + BLOCK_SIZE - 1) // BLOCK_SIZE,) - _build_sorted_indices_kernel[grid]( - sorted_token_idx, - sorted_k_idx, - sorted_flat_indices, - num_valid, - topk, - BLOCK_SIZE=BLOCK_SIZE, - ) - return sorted_token_idx, sorted_k_idx - - -def group_tokens_by_expert_v2( - expert_ids: torch.Tensor, - topk: int, - num_experts: int, - num_valid: int = None, -) -> tuple[torch.Tensor, torch.Tensor, int]: - del num_experts - device = expert_ids.device - n = expert_ids.numel() - flat_expert_ids = expert_ids.view(-1).long() - flat_idx = torch.arange(n, device=device, dtype=torch.int64) - composite_keys = torch.where( - flat_expert_ids >= 0, - flat_expert_ids * n + flat_idx, - torch.iinfo(torch.int64).max, - ) - _, sorted_indices = torch.sort(composite_keys, stable=False) - if num_valid is None: - # Fallback: requires CPU-GPU sync via .item() - valid_mask = flat_expert_ids >= 0 - num_valid = valid_mask.sum().item() - if num_valid == 0: - return ( - torch.empty(0, dtype=torch.int64, device=device), - torch.empty(0, dtype=torch.int64, device=device), - 0, - ) - sorted_flat_indices = sorted_indices[:num_valid] - sorted_token_idx = sorted_flat_indices // topk - sorted_k_idx = sorted_flat_indices % topk - return sorted_token_idx, sorted_k_idx, num_valid - - -def group_tokens_by_expert( - expert_ids: torch.Tensor, - topk: int, - num_valid: int = None, -) -> tuple[torch.Tensor, torch.Tensor, int]: - device = expert_ids.device - flat_expert_ids = expert_ids.view(-1).long() - valid_mask = flat_expert_ids >= 0 - if num_valid is None: - # Fallback: requires CPU-GPU sync via .item() - num_valid = valid_mask.sum().item() - if num_valid == 0: - return ( - torch.empty(0, dtype=torch.int64, device=device), - torch.empty(0, dtype=torch.int64, device=device), - 0, - ) - valid_expert_ids = flat_expert_ids[valid_mask] - valid_flat_indices = valid_mask.nonzero(as_tuple=True)[0] - sorted_order = torch.argsort(valid_expert_ids, stable=True) - sorted_flat_indices = valid_flat_indices[sorted_order] - sorted_token_idx, sorted_k_idx = build_sorted_indices(sorted_flat_indices, topk) - return sorted_token_idx, sorted_k_idx, num_valid diff --git a/src/xorl/ops/ep_kernels/deepep_scatter_gather.py b/src/xorl/ops/ep_kernels/deepep_scatter_gather.py deleted file mode 100644 index 9c50550b..00000000 --- a/src/xorl/ops/ep_kernels/deepep_scatter_gather.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Optimized Triton kernels for DeepEP scatter/gather operations.""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _gather_by_index_kernel( - output_ptr, - input_ptr, - indices_ptr, - num_gather: tl.constexpr, - hidden_dim: tl.constexpr, - BLOCK_H: tl.constexpr, -): - pid = tl.program_id(0) - if pid >= num_gather: - return - src_idx = tl.load(indices_ptr + pid) - for h_start in range(0, hidden_dim, BLOCK_H): - h_offs = h_start + tl.arange(0, BLOCK_H) - mask = h_offs < hidden_dim - val = tl.load(input_ptr + src_idx * hidden_dim + h_offs, mask=mask, other=0.0) - tl.store(output_ptr + pid * hidden_dim + h_offs, val, mask=mask) - - -@triton.jit -def _weighted_scatter_add_no_atomic_kernel( - output_ptr, - expert_output_ptr, - sorted_token_idx_ptr, - weights_ptr, - num_expert_outputs: tl.constexpr, - num_tokens: tl.constexpr, - hidden_dim: tl.constexpr, - BLOCK_H: tl.constexpr, -): - pid = tl.program_id(0) - if pid >= num_expert_outputs: - return - token_idx = tl.load(sorted_token_idx_ptr + pid) - weight = tl.load(weights_ptr + pid).to(tl.float32) - for h_start in range(0, hidden_dim, BLOCK_H): - h_offs = h_start + tl.arange(0, BLOCK_H) - mask = h_offs < hidden_dim - val = tl.load(expert_output_ptr + pid * hidden_dim + h_offs, mask=mask, other=0.0).to(tl.float32) - weighted_val = val * weight - out_ptr = output_ptr + token_idx * hidden_dim + h_offs - tl.atomic_add(out_ptr, weighted_val, mask=mask) - - -class DeepEPScatter(torch.autograd.Function): - @staticmethod - def forward(ctx, tokens, sorted_indices): - num_gather = sorted_indices.shape[0] - hidden_dim = tokens.shape[1] - output = torch.empty(num_gather, hidden_dim, dtype=tokens.dtype, device=tokens.device) - if num_gather == 0: - ctx.save_for_backward(sorted_indices) - ctx.num_tokens = tokens.shape[0] - ctx.hidden_dim = hidden_dim - return output - BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024) - grid = (num_gather,) - _gather_by_index_kernel[grid](output, tokens, sorted_indices, num_gather, hidden_dim, BLOCK_H=BLOCK_H) - ctx.save_for_backward(sorted_indices) - ctx.num_tokens = tokens.shape[0] - ctx.hidden_dim = hidden_dim - return output - - @staticmethod - def backward(ctx, grad_output): - (sorted_indices,) = ctx.saved_tensors - num_tokens = ctx.num_tokens - hidden_dim = ctx.hidden_dim - grad_tokens = torch.zeros(num_tokens, hidden_dim, dtype=torch.float32, device=grad_output.device) - if sorted_indices.shape[0] > 0: - expanded_indices = sorted_indices.unsqueeze(1).expand(-1, hidden_dim) - grad_tokens.scatter_add_(0, expanded_indices, grad_output.float()) - if grad_output.dtype != torch.float32: - grad_tokens = grad_tokens.to(grad_output.dtype) - return grad_tokens, None - - -class DeepEPWeightedGather(torch.autograd.Function): - @staticmethod - def forward(ctx, expert_output, sorted_token_idx, sorted_k_idx, weights, num_tokens): - num_expert_tokens = expert_output.shape[0] - hidden_dim = expert_output.shape[1] - dtype = expert_output.dtype - device = expert_output.device - output = torch.zeros(num_tokens, hidden_dim, dtype=torch.float32, device=device) - if num_expert_tokens == 0: - ctx.save_for_backward(expert_output, sorted_token_idx, sorted_k_idx, weights) - ctx.num_tokens = num_tokens - return output.to(dtype) - gathered_weights = weights[sorted_token_idx, sorted_k_idx] - BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024) - grid = (num_expert_tokens,) - _weighted_scatter_add_no_atomic_kernel[grid]( - output, - expert_output, - sorted_token_idx, - gathered_weights, - num_expert_tokens, - num_tokens, - hidden_dim, - BLOCK_H=BLOCK_H, - ) - ctx.save_for_backward(expert_output, sorted_token_idx, sorted_k_idx, weights) - ctx.num_tokens = num_tokens - return output.to(dtype) - - @staticmethod - def backward(ctx, grad_output): - expert_output, sorted_token_idx, sorted_k_idx, weights = ctx.saved_tensors - num_expert_tokens = expert_output.shape[0] - topk = weights.shape[1] - grad_output = grad_output.contiguous() - gathered_weights = weights[sorted_token_idx, sorted_k_idx] - grad_expert_output = grad_output[sorted_token_idx] * gathered_weights.unsqueeze(1) - grad_weights = torch.zeros_like(weights) - if num_expert_tokens > 0: - dot_products = (grad_output[sorted_token_idx] * expert_output).sum(dim=1) - flat_idx = sorted_token_idx * topk + sorted_k_idx - grad_weights.view(-1).scatter_add_(0, flat_idx, dot_products) - return grad_expert_output, None, None, grad_weights, None - - -def deepep_scatter(tokens: torch.Tensor, sorted_indices: torch.Tensor) -> torch.Tensor: - return DeepEPScatter.apply(tokens, sorted_indices) - - -def deepep_weighted_gather( - expert_output: torch.Tensor, - sorted_token_idx: torch.Tensor, - sorted_k_idx: torch.Tensor, - weights: torch.Tensor, - num_tokens: int, -) -> torch.Tensor: - return DeepEPWeightedGather.apply(expert_output, sorted_token_idx, sorted_k_idx, weights, num_tokens) From ff5d1a75fb37574ab29947ca3ce3f2acfa58aded Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 04:15:21 +0000 Subject: [PATCH 06/14] Fix flashqla_cp over-replacement from the phase-1 codemod The vendored-path rewrite matched 'xorl.ops.linear_attention.flashqla' as a prefix of the first-party module 'flashqla_cp', mispointing backend.py and the CP equivalence test at a nonexistent _vendored.flashqla_cp. --- src/xorl/ops/linear_attention/backend.py | 6 +++--- tests/distributed/test_flashqla_cp_equivalence.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/xorl/ops/linear_attention/backend.py b/src/xorl/ops/linear_attention/backend.py index a9013f53..379394e6 100644 --- a/src/xorl/ops/linear_attention/backend.py +++ b/src/xorl/ops/linear_attention/backend.py @@ -9,7 +9,7 @@ (:mod:`xorl.ops._vendored.flashqla`). Hopper (SM90) only, requires ``tilelang``. Faster fwd/bwd. The single-GPU kernel is CP-unaware; under Ulysses CP it is driven by xorl's native CP - orchestration via :mod:`xorl.ops._vendored.flashqla_cp` + orchestration via :mod:`xorl.ops.linear_attention.flashqla_cp` (see :func:`flashqla_chunk_gated_delta_rule_cp`). Select the backend with the ``XORL_GDN_BACKEND`` environment variable @@ -127,7 +127,7 @@ def flashqla_chunk_gated_delta_rule_cp(**kwargs: Any) -> Any: Like :func:`flashqla_chunk_gated_delta_rule`, the import is deferred (FlashQLA pulls in ``tilelang`` and validates SM90 at import time). This path drives the FlashQLA interior with xorl's Ulysses/sequence-parallel orchestration; see - :mod:`xorl.ops._vendored.flashqla_cp`. + :mod:`xorl.ops.linear_attention.flashqla_cp`. """ global _flashqla_chunk_cp if _flashqla_chunk_cp is None: @@ -135,7 +135,7 @@ def flashqla_chunk_gated_delta_rule_cp(**kwargs: Any) -> Any: from xorl.ops.linear_attention import tilelang_gemm_v1 # noqa: PLC0415 tilelang_gemm_v1.patch() - from xorl.ops._vendored.flashqla_cp import ( # noqa: PLC0415 + from xorl.ops.linear_attention.flashqla_cp import ( # noqa: PLC0415 flashqla_chunk_gated_delta_rule_cp as _chunk_cp, ) except Exception as exc: # pragma: no cover - environment dependent diff --git a/tests/distributed/test_flashqla_cp_equivalence.py b/tests/distributed/test_flashqla_cp_equivalence.py index 3eda4e57..4abf97f3 100644 --- a/tests/distributed/test_flashqla_cp_equivalence.py +++ b/tests/distributed/test_flashqla_cp_equivalence.py @@ -19,7 +19,7 @@ pytest.skip("FlashQLA CP parity requires CUDA", allow_module_level=True) from xorl.distributed.parallel_state import init_parallel_state # noqa: E402 -from xorl.ops._vendored.flashqla_cp import flashqla_chunk_gated_delta_rule_cp # noqa: E402 +from xorl.ops.linear_attention.flashqla_cp import flashqla_chunk_gated_delta_rule_cp # noqa: E402 from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context # noqa: E402 from xorl.ops.linear_attention.ops.gated_delta_rule import chunk_gated_delta_rule # noqa: E402 from xorl.utils.device import get_nccl_backend # noqa: E402 From cc8f635e03796f3b2d669ff5563d130434b0262a Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 06:02:47 +0000 Subject: [PATCH 07/14] Address review: move the serving twins to ops/sglang; no aliases left MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per kiddyboots216 on #78/#81: the three 'aliased in place' modules are now physically moved. - New ops/sglang/ package for the literal serving-engine twins: bi_families_v2.py (moved byte-identical — same git blob as main; the sha256 gate hashes content, not paths) and batch_invariant_ops.py (vendored-adapted from sglang/srt/batch_invariant_ops; kept a single diffable file). The package docstring states the paired-edit policy; the format/rewrite excludes follow the new path. - exact_sampling_transforms.py -> ops/exact/sampling_transforms.py, making the replay contract physically canonical under ops/exact. - Old module paths are sys.modules alias stubs (full-fidelity, including underscore names) for one deprecation cycle; in-repo imports codemodded. --- .pre-commit-config.yaml | 2 +- pyproject.toml | 2 +- src/xorl/models/auto.py | 4 +- src/xorl/models/layers/gated_deltanet.py | 2 +- src/xorl/models/layers/moe/moe_block.py | 2 +- src/xorl/models/layers/moe/router.py | 2 +- src/xorl/models/layers/normalization.py | 6 +- .../transformers/deepseek_v4/exact_lm_head.py | 4 +- .../transformers/glm5/exact_fullparam_fp8.py | 2 +- .../transformers/glm5/exact_lm_head_qlora.py | 4 +- src/xorl/models/transformers/glm5/indexer.py | 2 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 4 +- .../models/transformers/qwen3_5_shared.py | 4 +- src/xorl/objectives/causallm_loss.py | 2 +- src/xorl/objectives/cispo_loss.py | 2 +- src/xorl/objectives/grpo_loss.py | 2 +- .../objectives/importance_sampling_loss.py | 2 +- src/xorl/objectives/policy_loss.py | 2 +- src/xorl/ops/README.md | 13 +- src/xorl/ops/batch_invariant_ops.py | 2189 +---------------- src/xorl/ops/bi_families_v2.py | 1064 +------- src/xorl/ops/exact/__init__.py | 19 +- src/xorl/ops/exact/batch_invariant.py | 3 - src/xorl/ops/exact/families_v2.py | 3 - src/xorl/ops/exact/sampling_transforms.py | 314 ++- src/xorl/ops/exact_sampling_transforms.py | 314 +-- src/xorl/ops/loss/bi_fused_lm_head.py | 20 +- src/xorl/ops/loss/per_token_ce.py | 2 +- src/xorl/ops/sglang/__init__.py | 17 + src/xorl/ops/sglang/batch_invariant_ops.py | 2188 ++++++++++++++++ src/xorl/ops/sglang/bi_families_v2.py | 1063 ++++++++ src/xorl/server/runner/model_runner.py | 4 +- src/xorl/trainers/model_builder.py | 8 +- src/xorl/trainers/trainer.py | 2 +- .../distributed/test_bi_trunk_linear_fsdp.py | 2 +- tests/models/test_batch_invariance_dense.py | 2 +- .../models/test_glm52_exact_fullparam_fp8.py | 2 +- .../models/test_glm52_exact_lm_head_qlora.py | 2 +- tests/models/test_lora_merged_forward.py | 4 +- tests/models/test_moe_ep_native_combine.py | 2 +- .../test_qwen35_lora_projection_topology.py | 2 +- tests/models/test_qwen3_5_rmsnorm.py | 4 +- tests/models/test_qwen3_5_trunk_wrap.py | 2 +- tests/models/test_rmsnorm_family_contract.py | 4 +- .../test_rmsnorm_family_cross_engine.py | 4 +- tests/models/test_rmsnorm_sglang_fused.py | 4 +- tests/ops/test_bi_families_v2_norm.py | 12 +- tests/ops/test_bi_fused_lm_head.py | 8 +- tests/ops/test_bi_gemm_config_table.py | 4 +- tests/ops/test_bi_router_gemm.py | 2 +- tests/ops/test_bi_trunk_linear.py | 2 +- tests/ops/test_exact_sampling_transforms.py | 2 +- .../test_bi_trunk_linear_server_mode.py | 6 +- 53 files changed, 3683 insertions(+), 3660 deletions(-) delete mode 100644 src/xorl/ops/exact/batch_invariant.py delete mode 100644 src/xorl/ops/exact/families_v2.py create mode 100644 src/xorl/ops/sglang/__init__.py create mode 100644 src/xorl/ops/sglang/batch_invariant_ops.py create mode 100644 src/xorl/ops/sglang/bi_families_v2.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 250dd41c..a60a77b1 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ default_stages: [pre-commit, pre-push, manual] # and bi_families_v2.py is vendored byte-identical into the serving engine # with both copies sha256-gated, so it keeps that engine's formatting # (black, 88 columns). See src/xorl/ops/README.md. -exclude: ^(submodules/|\.venv/|src/xorl/ops/_vendored/|src/xorl/ops/bi_families_v2\.py$) +exclude: ^(submodules/|\.venv/|src/xorl/ops/_vendored/|src/xorl/ops/sglang/bi_families_v2\.py$) repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 diff --git a/pyproject.toml b/pyproject.toml index bcdcccdf..87a6cdb6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -149,7 +149,7 @@ line-length = 120 force-exclude = true extend-exclude = [ "src/xorl/ops/_vendored", - "src/xorl/ops/bi_families_v2.py", + "src/xorl/ops/sglang/bi_families_v2.py", ] [tool.ruff.lint] diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 1085225a..0ee5c673 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -1128,10 +1128,10 @@ def build_foundation_model( model.config._attn_implementation = attn_implementation if config._qwen3_dense_exact_contract: - from xorl.ops.batch_invariant_ops import ( # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import ( # noqa: PLC0415 wrap_trunk_linears_batch_invariant, ) - from xorl.ops.bi_families_v2 import ( # noqa: PLC0415 + from xorl.ops.sglang.bi_families_v2 import ( # noqa: PLC0415 _select_qwen3_dense_families_v2, ) diff --git a/src/xorl/models/layers/gated_deltanet.py b/src/xorl/models/layers/gated_deltanet.py index 5201a803..56158880 100644 --- a/src/xorl/models/layers/gated_deltanet.py +++ b/src/xorl/models/layers/gated_deltanet.py @@ -196,7 +196,7 @@ def _linear_with_contract( weight: torch.Tensor, ) -> torch.Tensor: if getattr(module, "_xorl_bi_trunk_wrapped", False): - from xorl.ops.batch_invariant_ops import batch_invariant_trunk_linear # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import batch_invariant_trunk_linear # noqa: PLC0415 return batch_invariant_trunk_linear(inputs, weight, module.bias) return F.linear(inputs, weight, module.bias) diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index 3d0f95e9..4bad2f16 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -54,7 +54,7 @@ class _BIRouterGemm(torch.autograd.Function): @staticmethod def forward(ctx, hidden, weight): - from xorl.ops.batch_invariant_ops import bi_router_gemm # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import bi_router_gemm # noqa: PLC0415 ctx.save_for_backward(hidden, weight) return bi_router_gemm(hidden, weight) diff --git a/src/xorl/models/layers/moe/router.py b/src/xorl/models/layers/moe/router.py index b10bc38f..7a5a5571 100644 --- a/src/xorl/models/layers/moe/router.py +++ b/src/xorl/models/layers/moe/router.py @@ -249,7 +249,7 @@ def _forward_softmax(self, router_logits: torch.Tensor, input_dtype: torch.dtype # Exact router contract: fixed-order renorm + cast so the top-k weights # are bit-identical to SGLang's batch-invariant path (the stock # sum(dim=-1) reduction is build-dependent; see bi_router_topk_weights). - from xorl.ops.batch_invariant_ops import bi_router_topk_weights # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import bi_router_topk_weights # noqa: PLC0415 routing_weights = bi_router_topk_weights(routing_weights, self.norm_topk_prob, input_dtype) else: diff --git a/src/xorl/models/layers/normalization.py b/src/xorl/models/layers/normalization.py index a9a67e63..d2e2ef11 100644 --- a/src/xorl/models/layers/normalization.py +++ b/src/xorl/models/layers/normalization.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from torch import nn -from xorl.ops.batch_invariant_ops import ( +from xorl.ops.sglang.batch_invariant_ops import ( RMS_NORM_FAMILIES, RMS_NORM_FAMILY_NO_RESIDUAL, RMS_NORM_FAMILY_RESIDUAL_TREE, @@ -21,7 +21,7 @@ mean_dim, set_batch_invariant_mode, ) -from xorl.ops.bi_families_v2 import families_v2_enabled, rms_norm_v2 +from xorl.ops.sglang.bi_families_v2 import families_v2_enabled, rms_norm_v2 RMSNormMode = Literal["eager", "native", "compile", "sglang", "sglang_fused", "sglang_jit", "sglang_kernel"] @@ -629,7 +629,7 @@ class RMSNorm(nn.Module): """Root Mean Square Layer Normalization. ``family`` declares which batch-invariant kernel family the site's serving - counterpart executes (see ``xorl.ops.batch_invariant_ops``): + counterpart executes (see ``xorl.ops.sglang.batch_invariant_ops``): ``"serving_no_residual"`` for qk-norms and layer-0 input layernorms, ``"serving_residual_tree"`` for residual-carrying norms (input layernorm at layer>0, post-attention layernorm, final norm). The declaration replaces the diff --git a/src/xorl/models/transformers/deepseek_v4/exact_lm_head.py b/src/xorl/models/transformers/deepseek_v4/exact_lm_head.py index 7e6878ec..b630a203 100644 --- a/src/xorl/models/transformers/deepseek_v4/exact_lm_head.py +++ b/src/xorl/models/transformers/deepseek_v4/exact_lm_head.py @@ -20,14 +20,14 @@ from xorl.distributed.canonical_moe import LogicalRowOwnership from xorl.lora.modules.linear import LoraLinear -from xorl.ops.bi_families_v2 import exact_temperature_scale_bf16_logits -from xorl.ops.exact_sampling_transforms import ( +from xorl.ops.exact.sampling_transforms import ( EXACT_FILTER_ROW_CHUNK, exact_sampling_identity_rows, exact_sampling_support, exact_selected_logprob_from_support, exact_selected_logprob_partitioned_from_support, ) +from xorl.ops.sglang.bi_families_v2 import exact_temperature_scale_bf16_logits DSV4_EXACT_TP8_LM_HEAD_CONTRACT_VERSION = "dsv4_flash_exact_tp8_lm_head_rank1_lora_v1" diff --git a/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py b/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py index 6cf1f694..83da20a5 100644 --- a/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py +++ b/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py @@ -838,7 +838,7 @@ def _router_value(self, hidden_states: Tensor) -> Tensor: if self._effective_weight.device != hidden_states.device: raise RuntimeError("GLM-5.2 full-param router view and activations must share one CUDA device") - from xorl.ops.batch_invariant_ops import bi_router_gemm # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import bi_router_gemm # noqa: PLC0415 return bi_router_gemm(hidden_states, self._effective_weight) diff --git a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py index 2eb74b5c..e90feb38 100644 --- a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py @@ -32,14 +32,14 @@ from xorl.lora.modules.linear import LoraLinear from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling -from xorl.ops.bi_families_v2 import exact_temperature_scale_fp32_logits -from xorl.ops.exact_sampling_transforms import ( +from xorl.ops.exact.sampling_transforms import ( EXACT_FILTER_ROW_CHUNK, exact_sampling_identity_rows, exact_sampling_support, exact_selected_logprob_from_support, exact_selected_logprob_partitioned_from_support, ) +from xorl.ops.sglang.bi_families_v2 import exact_temperature_scale_fp32_logits GLM52_EXACT_TP16_LM_HEAD_CONTRACT_VERSION = "glm52_exact_tp16_lm_head_qlora_v2" diff --git a/src/xorl/models/transformers/glm5/indexer.py b/src/xorl/models/transformers/glm5/indexer.py index b522ece9..f900005b 100644 --- a/src/xorl/models/transformers/glm5/indexer.py +++ b/src/xorl/models/transformers/glm5/indexer.py @@ -25,8 +25,8 @@ GLM52_SELECTOR_VERSION, select_glm52_logical_indices, ) -from xorl.ops.batch_invariant_ops import bi_bf16_fp32_linear, matmul_persistent from xorl.ops.exact.rope_class_b import build_class_b_cos_sin +from xorl.ops.sglang.batch_invariant_ops import bi_bf16_fp32_linear, matmul_persistent GLM52_FUSED_INDEX_K_IMPORT = "sglang.kernels.ops.quantization.dsv32.elementwise.fused_k_indexer_norm_rope" diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index 00d6cc91..2cd9278f 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -118,7 +118,7 @@ def unfuse_for_tp(self): @staticmethod def _linear_with_contract(module: nn.Linear, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: if getattr(module, "_xorl_bi_trunk_wrapped", False): - from xorl.ops.batch_invariant_ops import _BatchInvariantTrunkLinearFn # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import _BatchInvariantTrunkLinearFn # noqa: PLC0415 return _BatchInvariantTrunkLinearFn.apply(x, weight, module.bias) return F.linear(x, weight, module.bias) @@ -448,7 +448,7 @@ def _ep_combine_native( max_rows_for_ep_combine, sglang_fused_gate_sigmoid_mul_add, ) - from xorl.ops.batch_invariant_ops import _BatchInvariantTrunkLinearFn # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import _BatchInvariantTrunkLinearFn # noqa: PLC0415 ps = get_parallel_state() if not ps.ep_enabled: diff --git a/src/xorl/models/transformers/qwen3_5_shared.py b/src/xorl/models/transformers/qwen3_5_shared.py index 6c168564..b1d7a4e1 100644 --- a/src/xorl/models/transformers/qwen3_5_shared.py +++ b/src/xorl/models/transformers/qwen3_5_shared.py @@ -64,8 +64,8 @@ def _apply_qwen35_gdn_exact(model: torch.nn.Module) -> dict[str, int]: validate_native_ep_combine_size(ps.ep_size) from xorl.lora.modules.base import LoraModule # noqa: PLC0415 - from xorl.ops.batch_invariant_ops import wrap_trunk_linears_batch_invariant # noqa: PLC0415 - from xorl.ops.bi_families_v2 import _select_qwen35_families_v1 # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import wrap_trunk_linears_batch_invariant # noqa: PLC0415 + from xorl.ops.sglang.bi_families_v2 import _select_qwen35_families_v1 # noqa: PLC0415 # RMSNorm uses the qualified v2 tree. The LM-head/LSE remains on its # separately qualified v1 program; that selector does not control norms. diff --git a/src/xorl/objectives/causallm_loss.py b/src/xorl/objectives/causallm_loss.py index ff74d1cc..a627f55b 100644 --- a/src/xorl/objectives/causallm_loss.py +++ b/src/xorl/objectives/causallm_loss.py @@ -8,7 +8,7 @@ from xorl.objectives.loss_output import LossOutput from xorl.objectives.reducers import Reducer, TokenPartial -from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.exact.sampling_transforms import TOP_K_ALL from xorl.ops.loss.compiled_cross_entropy import ( compiled_ce_and_lse_sq_function, compiled_cross_entropy_function, diff --git a/src/xorl/objectives/cispo_loss.py b/src/xorl/objectives/cispo_loss.py index acf49dfc..2e9e0d88 100644 --- a/src/xorl/objectives/cispo_loss.py +++ b/src/xorl/objectives/cispo_loss.py @@ -8,7 +8,7 @@ from xorl.objectives.importance_sampling_loss import K3_DEBUG_THRESHOLDS from xorl.objectives.loss_output import LossOutput from xorl.objectives.reducers import Reducer, TokenPartial -from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.exact.sampling_transforms import TOP_K_ALL from xorl.ops.loss.per_token_ce import compute_per_token_ce diff --git a/src/xorl/objectives/grpo_loss.py b/src/xorl/objectives/grpo_loss.py index 9692cf4b..b9bf7b86 100644 --- a/src/xorl/objectives/grpo_loss.py +++ b/src/xorl/objectives/grpo_loss.py @@ -12,7 +12,7 @@ from xorl.objectives.loss_output import LossOutput from xorl.objectives.reducers import Reducer, TokenPartial -from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.exact.sampling_transforms import TOP_K_ALL from xorl.ops.loss.per_token_ce import compute_per_token_ce diff --git a/src/xorl/objectives/importance_sampling_loss.py b/src/xorl/objectives/importance_sampling_loss.py index 7af4d275..42422e59 100644 --- a/src/xorl/objectives/importance_sampling_loss.py +++ b/src/xorl/objectives/importance_sampling_loss.py @@ -7,7 +7,7 @@ from xorl.objectives.loss_output import LossOutput from xorl.objectives.reducers import Reducer, TokenPartial -from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.exact.sampling_transforms import TOP_K_ALL from xorl.ops.loss.per_token_ce import compute_per_token_ce diff --git a/src/xorl/objectives/policy_loss.py b/src/xorl/objectives/policy_loss.py index 59ff3dc6..165646ea 100644 --- a/src/xorl/objectives/policy_loss.py +++ b/src/xorl/objectives/policy_loss.py @@ -17,7 +17,7 @@ from xorl.objectives.loss_output import LossOutput from xorl.objectives.reducers import Reducer, TokenPartial -from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.exact.sampling_transforms import TOP_K_ALL from xorl.ops.loss.per_token_ce import compute_per_token_ce diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index f5ecacbd..375565c9 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -11,17 +11,18 @@ Three kinds of code live here with three different rules: | kind | rule | | --- | --- | | **Vendored** — `_vendored/` (`quack/`, `flashqla/`) | Never hand-edit, lint, or reformat. Each tree carries a `VENDORED.md` with provenance and the local-patch ledger. First-party tooling skips them (`[tool.ruff]` excludes in `pyproject.toml`; top-level `exclude:` in `.pre-commit-config.yaml`). | -| **Byte-contract-gated** — `bi_families_v2.py` (sha256-gated), `batch_invariant_ops.py` (parity-diffable twin of SGLang's copy; edits must consider the serving side) | Vendored byte-identical into the serving engine; both copies are sha256-gated. Any edit here without the paired serving-side edit breaks the gate. It keeps the engine's formatting (black, 88 columns) and is excluded from all rewriting hooks. | +| **Serving-parity twins** — `sglang/` (`bi_families_v2.py`, `batch_invariant_ops.py`) | Modules mirrored into/from the serving engine. `bi_families_v2.py` is sha256-gated byte-identical (keeps the engine's black-88 formatting, excluded from all rewriting hooks); `batch_invariant_ops.py` is vendored-adapted and stays a single diffable file. Edits require considering the paired serving-side copy. | | **First-party** — everything else | Normal rules. | ## Map (current) - `exact/` — the **serving-parity (exact) contract family**: byte-pinned - programs shared with the serving engine (#78 phase 3). Three members are - aliased rather than moved: `bi_families_v2.py` (sha256-gated), - `batch_invariant_ops.py` (diffable parity twin of SGLang's copy), and - `exact_sampling_transforms.py` (in-flight in #74). Old root-level module - paths are compat stubs for one deprecation cycle. + programs shared with the serving engine (#78 phase 3), including the + replay contract (`sampling_transforms.py`). +- `sglang/` — the literal serving-engine twins (`bi_families_v2.py`, + `batch_invariant_ops.py`); see the edit-policy table. Old root-level + module paths for both packages are compat stubs for one deprecation + cycle. - `loss/` — the CE/selected-logprob kernel stack. The RL/supervised objective functions live in `xorl/objectives/` (#78 phase 2); old module paths here are compat stubs for one deprecation cycle. diff --git a/src/xorl/ops/batch_invariant_ops.py b/src/xorl/ops/batch_invariant_ops.py index 79f4d56b..df710fad 100644 --- a/src/xorl/ops/batch_invariant_ops.py +++ b/src/xorl/ops/batch_invariant_ops.py @@ -1,2188 +1,7 @@ -# Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/batch_invariant_ops/batch_invariant_ops.py -# -# Vendored into xorl from SGLang's -# python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py so that the -# xorl trainer/server forward can use the SAME batch-invariant Triton matmul -# kernels SGLang uses. This makes the linear layers, gate matmul and lm_head -# bit-for-bit match SGLang's reduction order, collapsing the cross-engine K3 -# logprob tail. -# -# SGLang-internal helpers are stubbed/inlined so this module is self-contained. -# DeepGEMM is discovered lazily and used only on a supported Hopper bf16 route; -# otherwise the vendored Triton path remains the fallback. Environment parsing -# and calc_diff are inlined below. +"""Moved to ``xorl.ops.sglang.batch_invariant_ops`` (issue #78); alias stub.""" -import contextlib -import os -from collections import namedtuple -from collections.abc import Callable -from typing import Any, Dict, Literal, Optional, Tuple +import importlib as _importlib +import sys as _sys -import torch -import triton -import triton.language as tl -from triton.runtime.errors import OutOfResources -from xorl.ops.exact.bi_gemm_configs import baseline_mm_config, lookup_mm_config - - -# --- Stubs for SGLang-internal imports --------------------------------------- -# DeepGEMM's bf16 NN GEMM is bitwise-identical to the batch-invariant Triton -# persistent kernel (verified per-shape by tune_bi_gemm.py and gated by -# tests/ops/test_bi_gemm_config_table.py). Serving's mm route already relies on -# that equality — SGLang sends every contiguous bf16 N>=16 mm to DeepGEMM while -# the trainer ran pure Triton — so routing the trainer the same way cannot move -# cross-engine bits, and it is markedly faster at trainer token counts. -# Requires Hopper+ and an importable deep_gemm. Detection is lazy (first mm), -# so importing this module never initializes CUDA. -_DEEPGEMM_READY: bool | None = None - - -def _deepgemm_ready() -> bool: - global _DEEPGEMM_READY - if _DEEPGEMM_READY is None: - try: - import deep_gemm # noqa: F401, PLC0415 - - _DEEPGEMM_READY = ( - torch.cuda.is_available() - and torch.cuda.get_device_capability()[0] >= 9 - and hasattr(deep_gemm, "bf16_gemm_nn") - ) - except Exception: # noqa: BLE001 - _DEEPGEMM_READY = False - return _DEEPGEMM_READY - - -ENABLE_JIT_DEEPGEMM = True # gated per-call by _deepgemm_ready() - -# Shapes whose table config exceeded shared memory at launch (Triton's epilogue -# staging for wide-output tiles is version-dependent): remembered so the hot -# path re-launches straight on the pinned baseline without re-raising. -_MM_CONFIG_OOM_SHAPES: set[tuple] = set() - - -def _launch_with_config_fallback(launch, dtype, M, N, K, out_itemsize=None): - key = (str(dtype), M, N, K, out_itemsize) - if key in _MM_CONFIG_OOM_SHAPES: - launch(baseline_mm_config(dtype)) - return - try: - launch(lookup_mm_config(dtype, M, N, K, out_itemsize=out_itemsize)) - except OutOfResources: - _MM_CONFIG_OOM_SHAPES.add(key) - launch(baseline_mm_config(dtype)) - - -# ----------------------------------------------------------------------------- - - -# Exact-model activation always uses the admitted production route. Ambient -# process variables cannot substitute comparison or order-variant fallbacks. -_ENABLE_MM_DEEPGEMM = True - -__all__ = [ - "set_batch_invariant_mode", - "is_batch_invariant_mode_enabled", - "is_batch_invariant_op_enabled", - "get_batch_invariant_ops", - "disable_batch_invariant_mode", - "enable_batch_invariant_mode", - "rms_norm_batch_invariant", - "fused_add_rms_norm_batch_invariant", - "sglang_rms_norm_batch_invariant", - "fused_rms_norm_backward", - "wrap_trunk_linears_batch_invariant", - "is_trunk_linear_contract_enabled", - "batch_invariant_trunk_linear", - "bi_bf16_fp32_linear", - "set_trunk_linear_contract", - "RMSNormFamily", - "RMS_NORM_FAMILY_NO_RESIDUAL", - "RMS_NORM_FAMILY_RESIDUAL_TREE", - "RMS_NORM_FAMILIES", - "bi_rms_norm", - "bi_fused_add_rms_norm", -] - - -_BATCH_INVARIANT_ALL_OPS = { - "mm", - "addmm", - "mm_dtype", - "log_softmax", - "mean", - "rms_norm", - "bmm", -} -_BATCH_INVARIANT_ALIASES = { - "matmul": "mm", - "logsoftmax": "log_softmax", - "log-softmax": "log_softmax", - "rmsnorm": "rms_norm", - "rms-norm": "rms_norm", -} - - -def _parse_batch_invariant_ops() -> set[str]: - raw = os.environ.get("XORL_BATCH_INVARIANT_OPS", "all").strip().lower() - if raw in ("", "1", "true", "yes", "all"): - return set(_BATCH_INVARIANT_ALL_OPS) - if raw in ("0", "false", "no", "none"): - return set() - - ops = set() - for part in raw.replace(";", ",").split(","): - op = part.strip().lower().replace("-", "_") - if not op: - continue - op = _BATCH_INVARIANT_ALIASES.get(op, op) - if op == "mm": - ops.add("mm_dtype") - if op not in _BATCH_INVARIANT_ALL_OPS: - raise ValueError( - f"Unsupported XORL_BATCH_INVARIANT_OPS entry {part!r}; " - f"supported values are: {sorted(_BATCH_INVARIANT_ALL_OPS)}" - ) - ops.add(op) - return ops - - -def _matmul_launch_metadata(grid: Callable[..., Any], kernel: Any, args: Dict[str, Any]) -> Dict[str, Any]: - ret = {} - m, n, k = args["M"], args["N"], args["K"] - ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}]" - if "tiles_per_update" in args: - ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}, tiles_per_update={args['tiles_per_update']:02}]" - if "c_ptr" in args: - bytes_per_elem = args["c_ptr"].element_size() - else: - bytes_per_elem = 1 if args["FP8_OUTPUT"] else 2 - ret[f"flops{bytes_per_elem * 8}"] = 2.0 * m * n * k - ret["bytes"] = bytes_per_elem * (m * k + n * k + m * n) - return ret - - -@triton.jit -def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS): - group_id = tile_id // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + (tile_id % group_size_m) - pid_n = (tile_id % num_pid_in_group) // group_size_m - return pid_m, pid_n - - -@triton.jit(launch_metadata=_matmul_launch_metadata) -def matmul_kernel_persistent( - a_ptr, - b_ptr, - c_ptr, # - bias_ptr, - M, - N, - K, # - stride_am, - stride_ak, - stride_bk, - stride_bn, - stride_cm, - stride_cn, - BLOCK_SIZE_M: tl.constexpr, # - BLOCK_SIZE_N: tl.constexpr, # - BLOCK_SIZE_K: tl.constexpr, # - GROUP_SIZE_M: tl.constexpr, # - NUM_SMS: tl.constexpr, # - A_LARGE: tl.constexpr, - B_LARGE: tl.constexpr, - C_LARGE: tl.constexpr, - HAS_BIAS: tl.constexpr, -): - start_pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - k_tiles = tl.cdiv(K, BLOCK_SIZE_K) - num_tiles = num_pid_m * num_pid_n - - offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - - for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): - pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS) - start_m = pid_m * BLOCK_SIZE_M - start_n = pid_n * BLOCK_SIZE_N - offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) - offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) - if A_LARGE: - offs_am = offs_am.to(tl.int64) - if B_LARGE: - offs_bn = offs_bn.to(tl.int64) - offs_am = tl.where(offs_am < M, offs_am, 0) - offs_bn = tl.where(offs_bn < N, offs_bn, 0) - offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) - offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for ki in range(k_tiles): - if A_LARGE or B_LARGE: - offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) - else: - offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) - a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - - a = tl.load(a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0) - b = tl.load(b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0) - accumulator = tl.dot(a, b, accumulator) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - if C_LARGE: - offs_cm = offs_cm.to(tl.int64) - offs_cn = offs_cn.to(tl.int64) - c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - if HAS_BIAS: - bias_ptrs = bias_ptr + offs_cn - bias = tl.load(bias_ptrs, mask=offs_cn < N, other=0.0).to(tl.float32) - accumulator += bias - if c_ptr.dtype.element_ty == tl.float8e4nv: - c = accumulator.to(tl.float8e4nv) - elif c_ptr.dtype.element_ty == tl.bfloat16: - c = accumulator.to(tl.bfloat16) - elif c_ptr.dtype.element_ty == tl.float32: - c = accumulator.to(tl.float32) - else: - c = accumulator.to(tl.float16) - tl.store(c_ptrs, c, mask=c_mask) - - -def _matmul_persistent_triton(a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None): - # Check constraints. - assert a.shape[1] == b.shape[0], "Incompatible dimensions" - assert a.dtype == b.dtype, "Incompatible dtypes" - assert bias is None or bias.dim() == 1, "Currently assuming bias is 1D, let Horace know if you run into this" - NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count - M, K = a.shape - K, N = b.shape - dtype = a.dtype - # Allocates output. - c = torch.empty((M, N), device=a.device, dtype=dtype) - - # 1D launch kernel where each block gets its own program. - def grid(META): - return ( - min( - NUM_SMS, - triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), - ), - ) - - # Shape-keyed on the bit-neutral axes only; BLOCK_SIZE_K stays pinned per - # dtype (bi_gemm_configs — the R1 config table, identical in both engines). - def _launch(config): - matmul_kernel_persistent[grid]( - a, - b, - c, # - bias, - M, - N, - K, # - a.stride(0), - a.stride(1), # - b.stride(0), - b.stride(1), # - c.stride(0), - c.stride(1), # - NUM_SMS=NUM_SMS, # - A_LARGE=a.numel() > 2**31, - B_LARGE=b.numel() > 2**31, - C_LARGE=c.numel() > 2**31, - HAS_BIAS=bias is not None, - **config, - ) - - _launch_with_config_fallback(_launch, dtype, M, N, K) - return c - - -def _matmul_persistent_deepgemm(a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None): - M, K = a.shape - K, N = b.shape - dtype = a.dtype - out = torch.empty((M, N), device=a.device, dtype=dtype) - - try: - import deep_gemm # noqa: PLC0415 - - deep_gemm.bf16_gemm_nn(a, b, out) - except RuntimeError: - return None - - # TODO can this be put in DeepGEMM's `c`? - if bias is not None: - out += bias - - return out - - -def matmul_persistent(a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None): - K, N = b.shape - - # DeepGEMM has minimum dimension requirements for TMA descriptors - MIN_DEEPGEMM_DIM = 16 - - if ( - _ENABLE_MM_DEEPGEMM - and (a.dtype == torch.bfloat16) - and (b.dtype == torch.bfloat16) - and a.is_contiguous() - and b.transpose(0, 1).is_contiguous() - and N >= MIN_DEEPGEMM_DIM - and _deepgemm_ready() - ): - result = _matmul_persistent_deepgemm(a=a, b=b, bias=bias) - if result is not None: - return result - # DeepGEMM failed (e.g. dimensions too small for TMA descriptors), - # fall through to batch-invariant Triton persistent kernel - - return _matmul_persistent_triton(a=a, b=b, bias=bias) - - -@triton.jit -def _log_softmax_kernel( - input_ptr, - output_ptr, - input_row_stride: tl.constexpr, - output_row_stride: tl.constexpr, - n_cols: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - """ - Compute log_softmax along the last dimension of a 2D tensor. - Each block handles one row of the input tensor. - """ - # Get the row index for this block - row_idx = tl.program_id(0).to(tl.int64) - - # Compute base pointers for input and output rows - row_start_ptr = input_ptr + row_idx * input_row_stride - output_row_start_ptr = output_ptr + row_idx * output_row_stride - - # Step 1: Find maximum value in the row for numerical stability - # Load first block to infer dtype and initialize max_val with correct type - col_idx_init = tl.arange(0, BLOCK_SIZE) - mask_init = col_idx_init < n_cols - vals_init = tl.load(row_start_ptr + col_idx_init, mask=mask_init, other=-float("inf")) - max_val = tl.max(vals_init) - - # Continue with remaining blocks - for col_offset in range(BLOCK_SIZE, n_cols, BLOCK_SIZE): - col_idx = col_offset + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - - # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=-float("inf")) - - # Update maximum - max_val = tl.max(tl.maximum(vals, max_val)) - - # Step 2: Compute sum of exp(x - max_val) - # Initialize sum_exp with correct dtype by using tl.sum on a zero vector - sum_exp = tl.sum(tl.zeros([1], dtype=max_val.dtype)) - - for col_offset in range(0, n_cols, BLOCK_SIZE): - col_idx = col_offset + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - - # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) - - # Compute exp(x - max_val) and accumulate - exp_vals = tl.exp(vals - max_val) - sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0)) - - # Compute log(sum_exp) - log_sum_exp = tl.log(sum_exp) - - # Step 3: Compute final log_softmax values: x - max_val - log_sum_exp - for col_offset in range(0, n_cols, BLOCK_SIZE): - col_idx = col_offset + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - - # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask) - - # Compute log_softmax - output = vals - max_val - log_sum_exp - - # Store results - tl.store(output_row_start_ptr + col_idx, output, mask=mask) - - -def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: - """ - Compute log_softmax using Triton kernel. - - Args: - input: Input tensor - dim: Dimension along which to compute log_softmax (only -1 or last dim supported) - - Returns: - Tensor with log_softmax applied along the specified dimension - """ - if dim != -1 and dim != input.ndim - 1: - raise ValueError("This implementation only supports log_softmax along the last dimension") - - # Flatten all dimensions except the last one - original_shape = input.shape - input_2d = input.reshape(-1, input.shape[-1]) - input_2d = input_2d.contiguous() - - n_rows, n_cols = input_2d.shape - - # Allocate output tensor - output = torch.empty_like(input_2d) - - # Choose block size based on the number of columns - BLOCK_SIZE = 1024 - - # Launch kernel with one block per row - grid = (n_rows,) - _log_softmax_kernel[grid]( - input_2d, - output, - input_2d.stride(0), - output.stride(0), - n_cols, - BLOCK_SIZE=BLOCK_SIZE, - ) - # Reshape output back to original shape - return output.reshape(original_shape) - - -@triton.jit -def mean_kernel( - input_ptr, - output_ptr, - input_stride0, - input_stride1, - input_stride2, - output_stride0, - output_stride1, - M, # size before reduction dim - N, # size of reduction dim - K, # size after reduction dim - BLOCK_SIZE: tl.constexpr, -): - """ - Kernel for computing mean along a single dimension. - Input is viewed as (M, N, K) where N is the dimension being reduced. - """ - # Program ID gives us which output element we're computing - pid = tl.program_id(0) - - # Compute output indices - m_idx = pid // K - k_idx = pid % K - - # Bounds check - if m_idx >= M or k_idx >= K: - return - - # Accumulate sum across reduction dimension - acc = 0.0 - for n_start in range(0, N, BLOCK_SIZE): - n_offsets = n_start + tl.arange(0, BLOCK_SIZE) - mask = n_offsets < N - - # Calculate input indices - input_idx = m_idx * input_stride0 + n_offsets * input_stride1 + k_idx * input_stride2 - - # Load and accumulate - vals = tl.load(input_ptr + input_idx, mask=mask, other=0.0) - acc += tl.sum(vals) - - # Compute mean and store - mean_val = acc / N - output_idx = m_idx * output_stride0 + k_idx * output_stride1 - tl.store(output_ptr + output_idx, mean_val) - - -def mean_dim( - input: torch.Tensor, - dim: int, - keepdim: bool = False, - dtype: torch.dtype | None = None, -) -> torch.Tensor: - """ - Triton implementation of torch.mean with single dimension reduction. - - Args: - input: Input tensor - dim: Single dimension along which to compute mean - keepdim: Whether to keep the reduced dimension - dtype: Output dtype. If None, uses input dtype (or float32 for integer inputs) - - Returns: - Tensor with mean values along specified dimension - """ - # Validate inputs - assert input.is_cuda, "Input must be a CUDA tensor" - assert -input.ndim <= dim < input.ndim, f"Invalid dimension {dim} for tensor with {input.ndim} dimensions" - - # Handle negative dim - if dim < 0: - dim = dim + input.ndim - - # Handle dtype - if dtype is None: - if input.dtype in [torch.int8, torch.int16, torch.int32, torch.int64]: - dtype = torch.float32 - else: - dtype = input.dtype - - # Convert input to appropriate dtype if needed - if input.dtype != dtype: - input = input.to(dtype) - - # Get input shape and strides - shape = list(input.shape) - - # Calculate dimensions for kernel - M = 1 - for i in range(dim): - M *= shape[i] - - N = shape[dim] - - K = 1 - for i in range(dim + 1, len(shape)): - K *= shape[i] - - # Reshape input to 3D view (M, N, K) - input_3d = input.reshape(M, N, K) - - # Create output shape - if keepdim: - output_shape = shape.copy() - output_shape[dim] = 1 - else: - output_shape = shape[:dim] + shape[dim + 1 :] - - # Create output tensor - output = torch.empty(output_shape, dtype=dtype, device=input.device) - - # Reshape output for kernel - if keepdim: - output_2d = output.reshape(M, 1, K).squeeze(1) - else: - output_2d = output.reshape(M, K) - - # Launch kernel - grid = (M * K,) - BLOCK_SIZE = 1024 - - mean_kernel[grid]( - input_3d, - output_2d, - input_3d.stride(0), - input_3d.stride(1), - input_3d.stride(2), - output_2d.stride(0), - output_2d.stride(1) if output_2d.ndim > 1 else 0, - M, - N, - K, - BLOCK_SIZE, - ) - - return output - - -_INTERPOSE_GRAD_ERROR = ( - "Batch-invariant interposed op '{op}' received a grad-requiring input while grad is enabled. " - "The global interpose (XORL_BATCH_INVARIANT_MATMUL / enable_batch_invariant_mode) is " - "inference/verification-only: the aten::rms_norm override records no autograd graph (q/k-norm " - "gradients silently vanish) and the torch.bmm monkeypatch detaches the graph. For a training " - "forward on the batch-invariant contract use the module-scoped XORL_BI_TRUNK_LINEAR=1 lane " - "instead." -) - - -def _guard_interpose_no_grad(op: str, *tensors) -> None: - """Loud-fail: the global interpose must never see a training forward.""" - if torch.is_grad_enabled() and any(isinstance(t, torch.Tensor) and t.requires_grad for t in tensors): - raise RuntimeError(_INTERPOSE_GRAD_ERROR.format(op=op)) - - -def mm_batch_invariant(a, b): - _guard_interpose_no_grad("aten::mm", a, b) - return matmul_persistent(a, b) - - -def addmm_batch_invariant(bias, a, b): - _guard_interpose_no_grad("aten::addmm", bias, a, b) - return matmul_persistent(a, b, bias=bias) - - -def _log_softmax_batch_invariant(input, dim, _half_to_float): - assert not _half_to_float, "not implemented" - _guard_interpose_no_grad("aten::_log_softmax", input) - return log_softmax(input, dim=dim) - - -def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = None): - _guard_interpose_no_grad("aten::mean.dim", input) - assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}" - if dim is None or len(dim) == 0: - # aten::mean full-reduce dispatches here with dim=[]; the empty n_elems product returned the SUM - dim = list(range(input.ndim)) - if len(dim) == 1: - return mean_dim(input, dim[0], keepdim=keepdim) - else: - assert input.dtype in { - torch.float16, - torch.bfloat16, - torch.float32, - }, "only float types supported for now" - n_elems = 1 - for d in dim: - n_elems *= input.shape[d] - return torch.sum(input, dim=dim, keepdim=keepdim, dtype=torch.float32) / n_elems - - -@triton.jit -def bmm_kernel_persistent( - a_ptr, - b_ptr, - c_ptr, # - B, - M, - N, - K, # - stride_ab, - stride_am, - stride_ak, - stride_bb, - stride_bk, - stride_bn, - stride_cb, - stride_cm, - stride_cn, - BLOCK_SIZE_M: tl.constexpr, # - BLOCK_SIZE_N: tl.constexpr, # - BLOCK_SIZE_K: tl.constexpr, # - GROUP_SIZE_M: tl.constexpr, # - NUM_SMS: tl.constexpr, # - A_LARGE: tl.constexpr, - B_LARGE: tl.constexpr, - C_LARGE: tl.constexpr, -): - """ - Batched matrix multiplication kernel that processes batches in parallel. - Each tile processes a (BLOCK_SIZE_M, BLOCK_SIZE_N) output block for a specific batch. - """ - start_pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - k_tiles = tl.cdiv(K, BLOCK_SIZE_K) - num_tiles_per_batch = num_pid_m * num_pid_n - num_tiles_total = B * num_tiles_per_batch - - offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - - # Process tiles in a deterministic order: batch-major ordering - for tile_id in tl.range(start_pid, num_tiles_total, NUM_SMS, flatten=True): - # Decompose tile_id into batch and within-batch tile - batch_idx = tile_id // num_tiles_per_batch - tile_in_batch = tile_id % num_tiles_per_batch - - pid_m, pid_n = _compute_pid(tile_in_batch, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS) - start_m = pid_m * BLOCK_SIZE_M - start_n = pid_n * BLOCK_SIZE_N - offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) - offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) - if A_LARGE: - offs_am = offs_am.to(tl.int64) - if B_LARGE: - offs_bn = offs_bn.to(tl.int64) - offs_am = tl.where(offs_am < M, offs_am, 0) - offs_bn = tl.where(offs_bn < N, offs_bn, 0) - offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) - offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) - - # Add batch offset - if A_LARGE or B_LARGE: - batch_idx_typed = batch_idx.to(tl.int64) - else: - batch_idx_typed = batch_idx - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for ki in range(k_tiles): - if A_LARGE or B_LARGE: - offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) - else: - offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) - - a_ptrs = a_ptr + (batch_idx_typed * stride_ab + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - b_ptrs = b_ptr + (batch_idx_typed * stride_bb + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - - a = tl.load(a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0) - b = tl.load(b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0) - accumulator = tl.dot(a, b, accumulator) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - if C_LARGE: - offs_cm = offs_cm.to(tl.int64) - offs_cn = offs_cn.to(tl.int64) - c_ptrs = c_ptr + batch_idx_typed * stride_cb + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - - if c_ptr.dtype.element_ty == tl.float8e4nv: - c = accumulator.to(tl.float8e4nv) - elif c_ptr.dtype.element_ty == tl.bfloat16: - c = accumulator.to(tl.bfloat16) - elif c_ptr.dtype.element_ty == tl.float32: - c = accumulator.to(tl.float32) - else: - c = accumulator.to(tl.float16) - tl.store(c_ptrs, c, mask=c_mask) - - -def bmm_batch_invariant(a, b, *, out=None): - _guard_interpose_no_grad("aten::bmm/torch.bmm", a, b) - # Batched matrix multiply: (B, M, K) x (B, K, N) -> (B, M, N) - # Process batches in parallel with our persistent kernel - if a.ndim == 3 and b.ndim == 3: - # Check constraints - assert a.shape[0] == b.shape[0], "Batch sizes must match" - assert a.shape[2] == b.shape[1], "Incompatible dimensions" - assert a.dtype == b.dtype, "Incompatible dtypes" - - B = a.shape[0] - M = a.shape[1] - K = a.shape[2] - N = b.shape[2] - dtype = a.dtype - - # Allocate output - if out is None: - c = torch.empty((B, M, N), device=a.device, dtype=dtype) - else: - c = out - - NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count - - # Use fixed kernel configuration for determinism - configs = { - torch.bfloat16: { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "num_stages": 3, - "num_warps": 8, - }, - torch.float16: { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 256, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "num_stages": 3, - "num_warps": 8, - }, - torch.float32: { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 8, - "num_stages": 3, - "num_warps": 8, - }, - } - - config = configs.get(dtype) - if config is None: - raise ValueError( - f"Unsupported dtype {dtype} for bmm_batch_invariant. Supported dtypes are: {list(configs.keys())}" - ) - - # Grid: limit by NUM_SMS for persistent kernel approach - num_tiles_per_batch = triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]) - num_tiles_total = B * num_tiles_per_batch - grid = (min(NUM_SMS, num_tiles_total),) - - bmm_kernel_persistent[grid]( - a, - b, - c, # - B, - M, - N, - K, # - a.stride(0), - a.stride(1), - a.stride(2), # - b.stride(0), - b.stride(1), - b.stride(2), # - c.stride(0), - c.stride(1), - c.stride(2), # - NUM_SMS=NUM_SMS, # - A_LARGE=a.numel() > 2**31, - B_LARGE=b.numel() > 2**31, - C_LARGE=c.numel() > 2**31, - **config, - ) - - return c - else: - raise ValueError(f"bmm_batch_invariant expects 3D tensors, got shapes {a.shape} and {b.shape}") - - -@triton.jit -def _rms_norm_kernel( - input_ptr, - weight_ptr, - output_ptr, - input_row_stride: tl.constexpr, - output_row_stride: tl.constexpr, - n_cols: tl.constexpr, - eps, - BLOCK_SIZE: tl.constexpr, -): - """ - Compute RMS normalization along the last dimension of a 2D tensor. - RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight - Each block handles one row of the input tensor. - """ - row_idx = tl.program_id(0).to(tl.int64) - row_start_ptr = input_ptr + row_idx * input_row_stride - output_row_start_ptr = output_ptr + row_idx * output_row_stride - - # Step 1: Compute sum of squares in float32 to avoid overflow - sum_sq = tl.zeros([1], dtype=tl.float32) - for col_offset in range(0, n_cols, BLOCK_SIZE): - col_idx = col_offset + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) - # Convert to float32 for accumulation to prevent overflow - vals_f32 = vals.to(tl.float32) - sq_vals = vals_f32 * vals_f32 - sum_sq += tl.sum(tl.where(mask, sq_vals, 0.0)) - - # Step 2: Compute RMS (root mean square) in float32 - mean_sq = sum_sq / n_cols - rms = tl.sqrt(mean_sq + eps) - inv_rms = 1.0 / rms - - # Step 3: Normalize and apply weight - for col_offset in range(0, n_cols, BLOCK_SIZE): - col_idx = col_offset + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) - weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) - # Compute in float32 then convert back to input dtype - vals_f32 = vals.to(tl.float32) - weight_f32 = weight.to(tl.float32) - output_f32 = vals_f32 * inv_rms * weight_f32 - output = output_f32.to(vals.dtype) - tl.store(output_row_start_ptr + col_idx, output, mask=mask) - - -def rms_norm(input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: - """ - Compute RMS normalization using Triton kernel. - - RMS Norm normalizes the input by the root mean square and scales by weight: - output = input / sqrt(mean(input^2) + eps) * weight - - Args: - input: Input tensor of shape (..., hidden_size) - weight: Weight tensor of shape (hidden_size,) - eps: Small constant for numerical stability - - Returns: - Tensor with RMS normalization applied along the last dimension - """ - assert weight.dim() == 1, "Weight must be 1-dimensional" - assert input.shape[-1] == weight.shape[0], ( - f"Input last dimension ({input.shape[-1]}) must match weight dimension ({weight.shape[0]})" - ) - - # Flatten all dimensions except the last one - original_shape = input.shape - input_2d = input.reshape(-1, input.shape[-1]) - input_2d = input_2d.contiguous() - weight = weight.contiguous() - - n_rows, n_cols = input_2d.shape - - output = torch.empty_like(input_2d) - BLOCK_SIZE = 1024 - grid = (n_rows,) - _rms_norm_kernel[grid]( - input_2d, - weight, - output, - input_2d.stride(0), - output.stride(0), - n_cols, - eps, - BLOCK_SIZE=BLOCK_SIZE, - ) - return output.reshape(original_shape) - - -def rms_norm_batch_invariant(input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: - """ - Batch-invariant wrapper for RMS normalization. - - This function provides a deterministic, batch-invariant implementation - of RMS normalization for use with the batch_invariant mode. - - Adapted from https://github.com/vllm-project/vllm/blob/66a168a197ba214a5b70a74fa2e713c9eeb3251a/vllm/model_executor/layers/batch_invariant.py#L649 - - Args: - input: Input tensor of shape (..., hidden_size) - weight: Weight tensor of shape (hidden_size,) - eps: Small constant for numerical stability - - Returns: - RMS normalized tensor - """ - return rms_norm(input, weight, eps=eps) - - -# --------------------------------------------------------------------------- # -# Fused batch-invariant "sglang" RMSNorm (residual + no-residual) -# -# Vendored from SGLang's fused batch-invariant residual RMSNorm -# (python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py, branch -# feature/k3-train-serve-numerics). These reproduce, bit-for-bit, the eager -# ``normalization.sglang_residual_rms_norm`` path (fp32 upcast, ``mean_dim`` -# variance, ``rsqrt``, fp32 weight multiply, cast last) while replacing its -# ~6-launch eager forward with three kernels. Bit-exactness keys: -# - reuse the batch-invariant ``mean_dim`` for the variance (a hand-rolled -# ``tl.sum`` reduction does NOT match it), -# - ``tl.rsqrt`` matches ``torch.rsqrt(var + eps)`` (``1.0 / tl.sqrt`` does not), -# - residual add is upcast -> add -> round back to the input dtype (matches -# torch's fp32-accumulated bf16 elementwise add). -# The forward is order-identical to the eager path, so the static K3 forward is -# preserved exactly. These are forward-only; the trainer wraps them in an -# ``autograd.Function`` with a closed-form backward (see normalization.py). -# --------------------------------------------------------------------------- # -@triton.jit -def _add_residual_square_kernel( - input_ptr, - residual_ptr, - residual_out_ptr, - sq_ptr, - input_row_stride: tl.constexpr, - residual_row_stride: tl.constexpr, - residual_out_row_stride: tl.constexpr, - sq_row_stride: tl.constexpr, - n_cols: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - """ - Stage 1 of the fused residual-add RMSNorm: residual add in the original - dtype and the per-element square in float32. - - residual_out = (x + residual).to(orig_dtype) - sq = residual_out.float() ** 2 # float32 - - ``sq`` is then reduced by the existing batch-invariant ``mean_dim`` kernel, - so the variance reduction order is bit-identical to the eager - ``x.pow(2).mean(-1)`` path this replaces. - """ - row_idx = tl.program_id(0).to(tl.int64) - in_row = input_ptr + row_idx * input_row_stride - res_row = residual_ptr + row_idx * residual_row_stride - res_out_row = residual_out_ptr + row_idx * residual_out_row_stride - sq_row = sq_ptr + row_idx * sq_row_stride - for col_offset in range(0, n_cols, BLOCK_SIZE): - col_idx = col_offset + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - x = tl.load(in_row + col_idx, mask=mask, other=0.0) - r = tl.load(res_row + col_idx, mask=mask, other=0.0) - # Match torch's elementwise add for low-precision dtypes: upcast to - # float32, add, round the result back to the original dtype. The - # normalization then operates on this rounded value. - s = (x.to(tl.float32) + r.to(tl.float32)).to(x.dtype) - tl.store(res_out_row + col_idx, s, mask=mask) - s_f32 = s.to(tl.float32) - tl.store(sq_row + col_idx, s_f32 * s_f32, mask=mask) - - -@triton.jit -def _square_kernel( - input_ptr, - sq_ptr, - input_row_stride: tl.constexpr, - sq_row_stride: tl.constexpr, - n_cols: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - """No-residual analog of stage 1: per-element square in float32. - - sq = x.float() ** 2 - - Reduced by ``mean_dim`` for a variance bit-identical to the eager - ``x.float().pow(2).mean(-1)`` path (fp32 ``s * s`` == ``pow(x, 2)``). - """ - row_idx = tl.program_id(0).to(tl.int64) - in_row = input_ptr + row_idx * input_row_stride - sq_row = sq_ptr + row_idx * sq_row_stride - for col_offset in range(0, n_cols, BLOCK_SIZE): - col_idx = col_offset + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - x = tl.load(in_row + col_idx, mask=mask, other=0.0) - x_f32 = x.to(tl.float32) - tl.store(sq_row + col_idx, x_f32 * x_f32, mask=mask) - - -@triton.jit -def _rms_normalize_with_var_kernel( - input_ptr, - var_ptr, - weight_ptr, - output_ptr, - input_row_stride: tl.constexpr, - output_row_stride: tl.constexpr, - n_cols: tl.constexpr, - eps, - BLOCK_SIZE: tl.constexpr, -): - """ - Stage 2 of the fused RMSNorm: normalize by a precomputed per-row variance - and multiply weight in float32, casting last. - - out = (x.float() * rsqrt(var + eps) * weight.float()).to(orig_dtype) - - ``tl.rsqrt`` bit-matches ``torch.rsqrt(var + eps)`` used by the eager path - (``1.0 / tl.sqrt`` does not). - """ - row_idx = tl.program_id(0).to(tl.int64) - in_row = input_ptr + row_idx * input_row_stride - out_row = output_ptr + row_idx * output_row_stride - var = tl.load(var_ptr + row_idx) - inv_rms = tl.rsqrt(var + eps) - for col_offset in range(0, n_cols, BLOCK_SIZE): - col_idx = col_offset + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - x = tl.load(in_row + col_idx, mask=mask, other=0.0) - weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) - output_f32 = x.to(tl.float32) * inv_rms * weight.to(tl.float32) - tl.store(out_row + col_idx, output_f32.to(x.dtype), mask=mask) - - -def fused_add_rms_norm_batch_invariant( - input: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float = 1e-6, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Batch-invariant fused residual-add + RMS normalization. - - Returns ``(output, residual_out)`` bit-matching the eager - ``residual_out = input + residual`` followed by - ``normalization.sglang_residual_rms_norm(residual_out, weight, eps)`` for the - closed dense Qwen3 recipe (fp32 upcast, ``mean_dim`` variance, ``rsqrt``, - fp32 weight multiply, cast last). - - The eager path is ~6 small launches per call; here it is three: a fused - residual-add+square, the batch-invariant ``mean_dim`` reduction (reused - verbatim so the variance is bit-identical to ``x.pow(2).mean(-1)``), and a - fused normalize. Forward-only; wrap in an autograd.Function for training. - """ - assert weight.dim() == 1, "Weight must be 1-dimensional" - assert input.shape == residual.shape, "Input and residual must share a shape" - assert input.shape[-1] == weight.shape[0], ( - f"Input last dimension ({input.shape[-1]}) must match weight dimension ({weight.shape[0]})" - ) - - original_shape = input.shape - input_2d = input.reshape(-1, input.shape[-1]).contiguous() - residual_2d = residual.reshape(-1, residual.shape[-1]).contiguous() - weight = weight.contiguous() - - n_rows, n_cols = input_2d.shape - residual_out = torch.empty_like(input_2d) - sq = torch.empty((n_rows, n_cols), dtype=torch.float32, device=input.device) - - BLOCK_SIZE = 1024 - grid = (n_rows,) - _add_residual_square_kernel[grid]( - input_2d, - residual_2d, - residual_out, - sq, - input_2d.stride(0), - residual_2d.stride(0), - residual_out.stride(0), - sq.stride(0), - n_cols, - BLOCK_SIZE=BLOCK_SIZE, - ) - - # Reuse the batch-invariant mean reduction verbatim: variance is then - # bit-identical to the eager path's x.pow(2).mean(-1). - var = mean_dim(sq, -1, keepdim=True).reshape(-1).contiguous() - - output = torch.empty_like(input_2d) - _rms_normalize_with_var_kernel[grid]( - residual_out, - var, - weight, - output, - residual_out.stride(0), - output.stride(0), - n_cols, - eps, - BLOCK_SIZE=BLOCK_SIZE, - ) - return output.reshape(original_shape), residual_out.reshape(original_shape) - - -@triton.jit -def _rms_norm_backward_kernel( - normed_ptr, - grad_out_ptr, - grad_ro_ptr, - weight_ptr, - grad_in_ptr, - gw_partial_ptr, - normed_row_stride: tl.constexpr, - grad_out_row_stride: tl.constexpr, - grad_ro_row_stride: tl.constexpr, - grad_in_row_stride: tl.constexpr, - gw_row_stride: tl.constexpr, - n_cols: tl.constexpr, - eps, - HAS_RESIDUAL_GRAD: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - """Fused RMSNorm backward for the ``sglang``-style norm. - - Given ``normed_input`` (the tensor that was normalized), ``grad_output`` and - optionally ``grad_residual_out``, computes per row:: - - var = mean(rf^2); inv = rsqrt(var + eps) - S = sum_j grad_out_j * w_j * rf_j - grad_normed_k = inv * grad_out_k * w_k - (inv^3 * S / D) * rf_k - gw_partial_k = grad_out_k * rf_k * inv # summed over rows -> grad_weight - - When ``grad_residual_out`` is present it is added to ``grad_normed`` (the - residual stream feeds both the norm and the next layer). This is the - gradient path only; it does not enter the forward K3, so ``var`` is - recomputed with a plain reduction (not ``mean_dim``). - """ - row = tl.program_id(0).to(tl.int64) - n_row = normed_ptr + row * normed_row_stride - go_row = grad_out_ptr + row * grad_out_row_stride - gi_row = grad_in_ptr + row * grad_in_row_stride - gw_row = gw_partial_ptr + row * gw_row_stride - - sum_sq = tl.zeros([1], dtype=tl.float32) - s_acc = tl.zeros([1], dtype=tl.float32) - for off in range(0, n_cols, BLOCK_SIZE): - idx = off + tl.arange(0, BLOCK_SIZE) - mask = idx < n_cols - rf = tl.load(n_row + idx, mask=mask, other=0.0).to(tl.float32) - go = tl.load(go_row + idx, mask=mask, other=0.0).to(tl.float32) - w = tl.load(weight_ptr + idx, mask=mask, other=0.0).to(tl.float32) - sum_sq += tl.sum(tl.where(mask, rf * rf, 0.0)) - s_acc += tl.sum(tl.where(mask, go * w * rf, 0.0)) - var = sum_sq / n_cols - inv = tl.rsqrt(var + eps) - c = inv * inv * inv * s_acc / n_cols - - for off in range(0, n_cols, BLOCK_SIZE): - idx = off + tl.arange(0, BLOCK_SIZE) - mask = idx < n_cols - rf = tl.load(n_row + idx, mask=mask, other=0.0).to(tl.float32) - go = tl.load(go_row + idx, mask=mask, other=0.0).to(tl.float32) - w = tl.load(weight_ptr + idx, mask=mask, other=0.0).to(tl.float32) - grad_normed = inv * go * w - c * rf - if HAS_RESIDUAL_GRAD: - gro = tl.load(grad_ro_ptr + row * grad_ro_row_stride + idx, mask=mask, other=0.0).to(tl.float32) - grad_normed = grad_normed + gro - tl.store(gi_row + idx, grad_normed, mask=mask) - tl.store(gw_row + idx, go * rf * inv, mask=mask) - - -def fused_rms_norm_backward( - normed_input: torch.Tensor, - weight: torch.Tensor, - eps: float, - grad_output: torch.Tensor, - grad_residual_out: torch.Tensor | None = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Fused backward for :func:`fused_add_rms_norm_batch_invariant` / - :func:`sglang_rms_norm_batch_invariant`. - - Returns ``(grad_normed_input_fp32, grad_weight_fp32)``. ``grad_normed_input`` - already includes ``grad_residual_out`` when supplied. Not order-sensitive: - gradients do not enter the forward K3. - """ - original_shape = normed_input.shape - ni = normed_input.reshape(-1, original_shape[-1]).contiguous() - go = grad_output.reshape(-1, original_shape[-1]).contiguous() - wf = weight.float().contiguous() - - n_rows, n_cols = ni.shape - grad_in = torch.empty((n_rows, n_cols), dtype=torch.float32, device=ni.device) - gw_partial = torch.empty((n_rows, n_cols), dtype=torch.float32, device=ni.device) - - has_residual = grad_residual_out is not None - if has_residual: - gro = grad_residual_out.reshape(-1, original_shape[-1]).contiguous() - else: - gro = ni # unused; kernel does not read it when HAS_RESIDUAL_GRAD=False - - BLOCK_SIZE = 1024 - grid = (n_rows,) - _rms_norm_backward_kernel[grid]( - ni, - go, - gro, - wf, - grad_in, - gw_partial, - ni.stride(0), - go.stride(0), - gro.stride(0), - grad_in.stride(0), - gw_partial.stride(0), - n_cols, - eps, - HAS_RESIDUAL_GRAD=has_residual, - BLOCK_SIZE=BLOCK_SIZE, - ) - grad_weight = gw_partial.sum(0) - return grad_in.reshape(original_shape), grad_weight - - -def sglang_rms_norm_batch_invariant( - input: torch.Tensor, - weight: torch.Tensor, - eps: float = 1e-6, -) -> torch.Tensor: - """ - Batch-invariant RMS normalization bit-matching the eager - ``normalization.sglang_residual_rms_norm`` (no residual add). - - This is the no-residual analog of :func:`fused_add_rms_norm_batch_invariant` - for the ``force_sglang_residual`` call sites (input layernorm at layer>0 and - the final norm), which apply the residual-style fp32 normalization to a - single input. Two kernels + ``mean_dim`` replace the eager ~6 launches. - Forward-only; wrap in an autograd.Function for training. - """ - assert weight.dim() == 1, "Weight must be 1-dimensional" - assert input.shape[-1] == weight.shape[0], ( - f"Input last dimension ({input.shape[-1]}) must match weight dimension ({weight.shape[0]})" - ) - - original_shape = input.shape - input_2d = input.reshape(-1, input.shape[-1]).contiguous() - weight = weight.contiguous() - - n_rows, n_cols = input_2d.shape - sq = torch.empty((n_rows, n_cols), dtype=torch.float32, device=input.device) - - BLOCK_SIZE = 1024 - grid = (n_rows,) - _square_kernel[grid]( - input_2d, - sq, - input_2d.stride(0), - sq.stride(0), - n_cols, - BLOCK_SIZE=BLOCK_SIZE, - ) - - var = mean_dim(sq, -1, keepdim=True).reshape(-1).contiguous() - - output = torch.empty_like(input_2d) - _rms_normalize_with_var_kernel[grid]( - input_2d, - var, - weight, - output, - input_2d.stride(0), - output.stride(0), - n_cols, - eps, - BLOCK_SIZE=BLOCK_SIZE, - ) - return output.reshape(original_shape) - - -# --------------------------------------------------------------------------- # -# RMSNorm kernel-family contract -# -# Two batch-invariant RMSNorm kernel families coexist, and they disagree at -# 1 ulp on rare bf16 boundary values (~2/524288 at [4096, 128]), so silently -# swapping one for the other seeds K3 divergence that amplifies downstream. -# Each family is pinned to the serving site-class that executes it: -# - "serving_no_residual" (family-1): the looped ``tl.sum`` + ``1.0/tl.sqrt`` -# kernel (``rms_norm_batch_invariant``), what SGLang dispatches when -# ``residual is None`` under batch-invariant mode and what the -# ``aten::rms_norm`` interpose runs. Site-classes: qk-norm, layer-0 input -# layernorm. -# - "serving_residual_tree" (family-2): the ``mean_dim`` + ``tl.rsqrt`` -# fused residual-tree kernels (``fused_add_rms_norm_batch_invariant`` / -# ``sglang_rms_norm_batch_invariant``), what SGLang dispatches for -# residual calls under the rl-on-policy lane. Site-classes: input -# layernorm at layer>0, post-attention layernorm, final norm. -# Every call site must name its family through ``bi_rms_norm`` / -# ``bi_fused_add_rms_norm``; never call the family kernels directly. -# --------------------------------------------------------------------------- # -RMS_NORM_FAMILY_NO_RESIDUAL = "serving_no_residual" -RMS_NORM_FAMILY_RESIDUAL_TREE = "serving_residual_tree" -RMS_NORM_FAMILIES = (RMS_NORM_FAMILY_NO_RESIDUAL, RMS_NORM_FAMILY_RESIDUAL_TREE) -RMSNormFamily = Literal["serving_no_residual", "serving_residual_tree"] - - -def bi_rms_norm( - input: torch.Tensor, - weight: torch.Tensor, - eps: float = 1e-6, - *, - family: RMSNormFamily, - zero_centered: bool = False, -) -> torch.Tensor: - """Single-tensor batch-invariant RMSNorm with an explicit kernel family. - - The single sanctioned entry point for both no-residual family kernels; the - ``family`` keyword must name the serving counterpart of the call site. - - ``zero_centered`` is the Gemma-style Qwen3.5 form: an fp32 upcast with the - ``1 + weight`` scale folded in fp32, cast back last, exactly as - ``normalization.native_zero_centered_rms_norm`` composes it. It is an affine - fold around the SAME family-1 reduction tree — not a third family — and only - exists in no-residual form (Qwen3.5 residual-tree norms run the eager native - path, never a batch-invariant kernel). - """ - if zero_centered: - if family != RMS_NORM_FAMILY_NO_RESIDUAL: - raise ValueError( - "zero-centered RMSNorm only exists in the 'serving_no_residual' family; " - "Qwen3.5 residual-tree norms run the native path, not a batch-invariant kernel" - ) - return rms_norm_batch_invariant(input.float(), 1.0 + weight.float(), eps=eps).type_as(input) - if family == RMS_NORM_FAMILY_NO_RESIDUAL: - return rms_norm_batch_invariant(input, weight, eps=eps) - if family == RMS_NORM_FAMILY_RESIDUAL_TREE: - return sglang_rms_norm_batch_invariant(input, weight, eps=eps) - raise ValueError(f"Unknown RMSNorm family {family!r}; expected one of {RMS_NORM_FAMILIES}") - - -def bi_fused_add_rms_norm( - input: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float = 1e-6, - *, - family: RMSNormFamily, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Fused residual-add batch-invariant RMSNorm with an explicit kernel family. - - Only the serving residual tree has a fused-add kernel; requesting the - no-residual family with a residual stream is a contract violation (serving - never runs family-1 on a residual site) and raises. - """ - if family == RMS_NORM_FAMILY_RESIDUAL_TREE: - return fused_add_rms_norm_batch_invariant(input, residual, weight, eps=eps) - if family == RMS_NORM_FAMILY_NO_RESIDUAL: - raise ValueError( - "RMSNorm family 'serving_no_residual' has no fused-add kernel: residual " - "site-classes are 'serving_residual_tree' by the cross-engine contract" - ) - raise ValueError(f"Unknown RMSNorm family {family!r}; expected one of {RMS_NORM_FAMILIES}") - - -# --------------------------------------------------------------------------- # -# Batch-invariant fused LM-head selected-token log-probability -# -# The K3 lm-head contract, vendored identically in xorl and SGLang -# (python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py). Both engines -# compute per-token logprobs of given token ids from bit-exact bf16 hidden -# states and the bf16 lm-head weight through the SAME reduction trees, so the -# results are bitwise identical cross-engine: -# 1. chunk GEMM — ``matmul_kernel_persistent`` with the family's fixed bf16 -# tile config and an fp32 output buffer. bf16xbf16 products are exact in -# fp32, so reading the weight in bf16 with tensor-core fp32 accumulation -# equals a GEMM over materialized fp32 upcasts with the same tree (and -# deletes the fp32 weight copy the eager paths materialize). -# 2. chunk stats — per-row max and sum(exp(x - chunk_max)) in a fixed -# sequential BLOCK loop (same discipline as ``_rms_norm_kernel``), plus -# the selected-token logit gather. -# 3. merge — global max over chunk maxima (exact), then the rescaled sumexp -# accumulated in pinned chunk order; lse = gmax + log(acc). -# All transcendentals stay inside these kernels: tl.exp/tl.log measured -# bit-identical across triton 3.5.1 (serving venv) and 3.7.1 (trainer venv), -# as is the fixed-tile tl.dot fp32 accumulator. VOCAB_CHUNK and STATS_BLOCK are -# contract constants — changing either changes the bits (the LSE reduction -# tree). The chunk GEMM's tile config is shape-keyed via bi_gemm_configs: only -# its BLOCK_SIZE_K (pinned there) is bit-relevant. -# Forward-only; the trainer wraps it in an autograd.Function (ops/loss). -# --------------------------------------------------------------------------- # - -BI_LM_HEAD_VOCAB_CHUNK = 8192 -_BI_LM_HEAD_STATS_BLOCK = 1024 - - -@triton.jit -def _lm_head_chunk_stats_kernel( - logits_ptr, - token_ids_ptr, - sel_ptr, - m_ptr, - s_ptr, - temp_ptr, - logits_row_stride, - n_cols, - col_offset, - chunk_idx, - n_chunks, - BLOCK_SIZE: tl.constexpr, - HAS_TEMP: tl.constexpr, -): - """Per-row chunk statistics over an fp32 logits tile [N, n_cols]: - chunk max, sum(exp(x - chunk_max)) in a fixed sequential block loop, and - the selected-token logit when ``token_ids[row]`` falls in this chunk. - With HAS_TEMP, logits are scaled by 1/temp[row] before the statistics and - the selected logit; the fp32 divide runs in-kernel so every engine - computes the identical scale (elementwise, so batch-invariance holds).""" - row = tl.program_id(0).to(tl.int64) - row_ptr = logits_ptr + row * logits_row_stride - if HAS_TEMP: - inv_t = 1.0 / tl.load(temp_ptr + row) - else: - inv_t = 1.0 - - row_max = float("-inf") - for col_start in range(0, n_cols, BLOCK_SIZE): - col_idx = col_start + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - vals = tl.load(row_ptr + col_idx, mask=mask, other=float("-inf")) - if HAS_TEMP: - vals = vals * inv_t - row_max = tl.maximum(row_max, tl.max(vals)) - - sum_exp = 0.0 - for col_start in range(0, n_cols, BLOCK_SIZE): - col_idx = col_start + tl.arange(0, BLOCK_SIZE) - mask = col_idx < n_cols - vals = tl.load(row_ptr + col_idx, mask=mask, other=float("-inf")) - if HAS_TEMP: - vals = vals * inv_t - e = tl.exp(vals - row_max) - sum_exp += tl.sum(tl.where(mask, e, 0.0)) - - tl.store(m_ptr + row * n_chunks + chunk_idx, row_max) - tl.store(s_ptr + row * n_chunks + chunk_idx, sum_exp) - - tok = tl.load(token_ids_ptr + row) - local = tok - col_offset - in_chunk = (local >= 0) & (local < n_cols) - sel = tl.load(row_ptr + local, mask=in_chunk, other=0.0) - if HAS_TEMP: - sel = sel * inv_t - tl.store(sel_ptr + row, sel, mask=in_chunk) - - -@triton.jit -def _lm_head_lse_merge_kernel( - m_ptr, - s_ptr, - lse_ptr, - n_chunks, -): - """lse[row] = gmax + log(sum_c s_c * exp(m_c - gmax)), chunks in pinned order.""" - row = tl.program_id(0).to(tl.int64) - base = row * n_chunks - gmax = float("-inf") - for c in range(n_chunks): - gmax = tl.maximum(gmax, tl.load(m_ptr + base + c)) - acc = 0.0 - for c in range(n_chunks): - acc += tl.load(s_ptr + base + c) * tl.exp(tl.load(m_ptr + base + c) - gmax) - tl.store(lse_ptr + row, gmax + tl.log(acc)) - - -def _bi_lm_head_chunk_gemm_fp32(a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> None: - """Launch the family's persistent matmul with the shape-keyed bf16 config and - an fp32 output buffer (the fp32 store path keeps the raw accumulator bits).""" - NUM_SMS = torch.cuda.get_device_properties(a.device).multi_processor_count - M, K = a.shape - _, N = b.shape - - def grid(META): - return (min(NUM_SMS, triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"])),) - - def _launch(config): - matmul_kernel_persistent[grid]( - a, - b, - out, - None, - M, - N, - K, - a.stride(0), - a.stride(1), - b.stride(0), - b.stride(1), - out.stride(0), - out.stride(1), - NUM_SMS=NUM_SMS, - A_LARGE=a.numel() > 2**31, - B_LARGE=b.numel() > 2**31, - C_LARGE=out.numel() > 2**31, - HAS_BIAS=False, - **config, - ) - - _launch_with_config_fallback(_launch, a.dtype, M, N, K, out_itemsize=out.element_size()) - - -def bi_lm_head_selected_logprob( - hidden: torch.Tensor, - weight: torch.Tensor, - token_ids: torch.Tensor, - temperature: Optional[torch.Tensor] = None, - vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Per-token ``log p(token_ids)`` for the LM head, batch-invariant and - cross-engine bit-exact (the K3 lm-head contract). - - Args: - hidden: ``[N, H]`` bf16 hidden states (pre lm-head). - weight: ``[V, H]`` bf16 lm-head weight (kept resident in bf16; no fp32 - copy is materialized). - token_ids: ``[N]`` integer token ids to score (callers must pre-clamp - ignored positions to a valid id and mask outputs downstream). - temperature: optional ``[N]`` fp32 per-row temperatures (> 0). Logits - are scaled by ``1/temperature[row]`` inside the stats kernel (the - divide runs in-kernel, so engines sharing the contract compute the - identical scale). ``None`` is the exact temperature-1.0 path. - - Returns: - ``(logprob, lse, selected)`` — all ``[N]`` fp32; ``logprob = selected - lse`` - (temperature-scaled when ``temperature`` is given). - """ - assert hidden.ndim == 2 and weight.ndim == 2, "hidden and weight must be 2D" - assert hidden.shape[1] == weight.shape[1], "hidden dim mismatch" - assert hidden.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( - "the lm-head contract takes bf16 hidden/weight (fp32 upcast is exact inside the GEMM)" - ) - assert hidden.is_cuda, "CUDA only" - - hidden = hidden.contiguous() - token_ids = token_ids.contiguous().to(device=hidden.device, dtype=torch.int64) - n_tokens = hidden.shape[0] - vocab = weight.shape[0] - n_chunks = (vocab + vocab_chunk - 1) // vocab_chunk - if temperature is not None: - temperature = temperature.reshape(-1).to(device=hidden.device, dtype=torch.float32).contiguous() - assert temperature.shape[0] == n_tokens, "temperature must be per-row [N]" - assert bool((temperature > 0).all()), "temperature must be > 0" - - chunk_max = torch.empty((n_tokens, n_chunks), dtype=torch.float32, device=hidden.device) - chunk_sumexp = torch.empty_like(chunk_max) - selected = torch.zeros(n_tokens, dtype=torch.float32, device=hidden.device) - lse = torch.empty(n_tokens, dtype=torch.float32, device=hidden.device) - logits_buf = torch.empty((n_tokens, vocab_chunk), dtype=torch.float32, device=hidden.device) - - for chunk_idx, col_start in enumerate(range(0, vocab, vocab_chunk)): - col_end = min(col_start + vocab_chunk, vocab) - n_cols = col_end - col_start - logits_c = logits_buf[:, :n_cols] - # [H, C] transposed view of the resident bf16 weight — the persistent - # kernel takes explicit strides, so no copy is made. - _bi_lm_head_chunk_gemm_fp32(hidden, weight[col_start:col_end].t(), logits_c) - _lm_head_chunk_stats_kernel[(n_tokens,)]( - logits_c, - token_ids, - selected, - chunk_max, - chunk_sumexp, - temperature, - logits_c.stride(0), - n_cols, - col_start, - chunk_idx, - n_chunks, - BLOCK_SIZE=_BI_LM_HEAD_STATS_BLOCK, - HAS_TEMP=temperature is not None, - ) - - _lm_head_lse_merge_kernel[(n_tokens,)](chunk_max, chunk_sumexp, lse, n_chunks) - # In exact math the selected logit never exceeds the LSE; clamp the one-ulp - # fp boundary case (p~1 tokens) so contract logprobs are provably <= 0. - return torch.clamp_max(selected - lse, 0.0), lse, selected - - -def bi_lm_head_full_logits( - hidden: torch.Tensor, - weight: torch.Tensor, - vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, -) -> torch.Tensor: - """Materialize full FP32 logits through the v1 contract GEMM.""" - - assert hidden.ndim == 2 and weight.ndim == 2, "hidden and weight must be 2D" - assert hidden.shape[1] == weight.shape[1], "hidden dim mismatch" - assert hidden.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( - "the lm-head contract takes bf16 hidden/weight" - ) - assert hidden.is_cuda, "CUDA only" - - hidden = hidden.contiguous() - logits = torch.empty( - (hidden.shape[0], weight.shape[0]), - dtype=torch.float32, - device=hidden.device, - ) - if hidden.shape[0] != 0: - _bi_lm_head_chunk_gemm_fp32(hidden, weight.t(), logits) - return logits - - -def bi_lm_head_selected_logprob_from_logits( - logits: torch.Tensor, - token_ids: torch.Tensor, - temperature: Optional[torch.Tensor] = None, - vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Run the v1 stats and merge tree over materialized FP32 logits.""" - - assert logits.ndim == 2 and logits.dtype == torch.float32, "logits must be FP32 [N, V]" - assert logits.is_cuda and logits.stride(1) == 1, "logits must be contiguous CUDA rows" - n_tokens, vocab = logits.shape - token_ids = token_ids.contiguous().to(device=logits.device, dtype=torch.int64) - assert token_ids.shape == (n_tokens,), "token_ids must be per-row [N]" - n_chunks = (vocab + vocab_chunk - 1) // vocab_chunk - if temperature is not None: - temperature = temperature.reshape(-1).to(device=logits.device, dtype=torch.float32).contiguous() - assert temperature.shape == (n_tokens,), "temperature must be per-row [N]" - torch._assert_async((temperature > 0).all(), "temperature must be > 0") - - chunk_max = torch.empty((n_tokens, n_chunks), dtype=torch.float32, device=logits.device) - chunk_sumexp = torch.empty_like(chunk_max) - selected = torch.zeros(n_tokens, dtype=torch.float32, device=logits.device) - lse = torch.empty(n_tokens, dtype=torch.float32, device=logits.device) - for chunk_idx, col_start in enumerate(range(0, vocab, vocab_chunk)): - col_end = min(col_start + vocab_chunk, vocab) - logits_c = logits[:, col_start:col_end] - _lm_head_chunk_stats_kernel[(n_tokens,)]( - logits_c, - token_ids, - selected, - chunk_max, - chunk_sumexp, - temperature, - logits_c.stride(0), - col_end - col_start, - col_start, - chunk_idx, - n_chunks, - BLOCK_SIZE=_BI_LM_HEAD_STATS_BLOCK, - HAS_TEMP=temperature is not None, - ) - - _lm_head_lse_merge_kernel[(n_tokens,)](chunk_max, chunk_sumexp, lse, n_chunks) - return torch.clamp_max(selected - lse, 0.0), lse, selected - - -# --------------------------------------------------------------------------- # -# Batch-invariant MoE router GEMM (the K3 router contract) -# -# Vendored identically in xorl and SGLang so the MoE gate/router logits are -# computed through ONE reduction tree cross-engine. Unlike the capture/replay -# lane, live training routes independently of serving (no routing replay), so a -# ~1e-10..1e-4 router-logit reduction-order diff between the two engines' GEMMs -# can flip the top-k expert selection on razor-edge tokens and cause large, -# rare-token logprob divergence. This kernel removes that last term: -# - bf16 hidden [N, H] @ bf16 gate weight [E, H]^T -> fp32 logits [N, E] -# - ``matmul_kernel_persistent`` with a pinned tile config and an fp32 output -# buffer (same discipline as the lm-head contract). bf16xbf16 products are -# exact in fp32, so reading both operands in bf16 with tensor-core fp32 -# accumulation equals an fp32 GEMM over their (exact) fp32 upcasts, but with -# the reduction order pinned identically in both engines — and without the -# fp32 weight/activation copies the eager fp32-router paths materialize. -# num_experts is small (a single BLOCK_SIZE_N tile for the common E <= 128), so -# the whole GEMM is one persistent launch. The config below is part of the -# contract; changing any constant changes the bits. -# Exact model programs call this kernel directly on both trainer and sampler. -# Forward-only; the trainer wraps it in an autograd.Function with a closed-form -# (order-insensitive) backward — gradients do not enter the forward K3. -# --------------------------------------------------------------------------- # - -_BI_ROUTER_GEMM_CONFIG = { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "num_stages": 3, - "num_warps": 8, -} - - -def bi_router_gemm(hidden: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - """fp32 MoE router logits ``[N, E]`` from bf16 hidden ``[N, H]`` and bf16 gate - weight ``[E, H]`` (the K3 router contract). - - One persistent bf16-in / fp32-accumulate / fp32-out GEMM with a pinned launch - config, vendored identically in xorl and SGLang so the router logits (and - therefore the top-k expert selection) are bitwise identical cross-engine. - bf16xbf16 products are exact in fp32, so this equals an fp32 GEMM over the - upcast operands the eager fp32-router paths materialize — minus the fp32 - weight/activation copies and with a reduction order that no longer depends on - the backend GEMM. - - Args: - hidden: ``[N, H]`` bf16 hidden states (pre-gate). - weight: ``[E, H]`` bf16 gate weight (kept resident in bf16; no fp32 copy). - - Returns: - ``[N, E]`` fp32 router logits. - """ - assert hidden.ndim == 2 and weight.ndim == 2, "hidden and weight must be 2D" - assert hidden.shape[1] == weight.shape[1], "hidden dim mismatch" - assert hidden.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( - "the router contract takes bf16 hidden/weight (fp32 upcast is exact inside the GEMM)" - ) - assert hidden.is_cuda, "CUDA only" - - hidden = hidden.contiguous() - weight = weight.contiguous() - n_tokens = hidden.shape[0] - num_experts = weight.shape[0] - logits = torch.empty((n_tokens, num_experts), dtype=torch.float32, device=hidden.device) - if n_tokens == 0: - return logits - - # [H, E] transposed view of the resident bf16 gate weight — the persistent - # kernel takes explicit strides, so no copy is made. - b = weight.t() - NUM_SMS = torch.cuda.get_device_properties(hidden.device).multi_processor_count - M, K = hidden.shape - N = num_experts - - def grid(META): - return (min(NUM_SMS, triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"])),) - - matmul_kernel_persistent[grid]( - hidden, - b, - logits, - None, - M, - N, - K, - hidden.stride(0), - hidden.stride(1), - b.stride(0), - b.stride(1), - logits.stride(0), - logits.stride(1), - NUM_SMS=NUM_SMS, - A_LARGE=hidden.numel() > 2**31, - B_LARGE=weight.numel() > 2**31, - C_LARGE=logits.numel() > 2**31, - HAS_BIAS=False, - **_BI_ROUTER_GEMM_CONFIG, - ) - return logits - - -class _BIBf16Fp32LinearFn(torch.autograd.Function): - """Trainable wrapper for the shared BF16-input, FP32-output GEMM. - - The pinned forward reduction enters the trainer-sampler numerical - contract. The ordinary linear backward does not enter K3, so it may use - the native matmul reduction while still propagating gradients to both - operands. - """ - - @staticmethod - def forward(ctx, input: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - ctx.save_for_backward(input, weight) - input_2d = input.reshape(-1, input.shape[-1]) - output = bi_router_gemm(input_2d, weight) - return output.reshape(*input.shape[:-1], weight.shape[0]) - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - input, weight = ctx.saved_tensors - grad_output_2d = grad_output.reshape(-1, grad_output.shape[-1]).float() - input_2d = input.reshape(-1, input.shape[-1]) - grad_input = grad_weight = None - if ctx.needs_input_grad[0]: - grad_input = (grad_output_2d @ weight.float()).to(input.dtype).reshape_as(input) - if ctx.needs_input_grad[1]: - grad_weight = (grad_output_2d.t() @ input_2d.float()).to(weight.dtype) - return grad_input, grad_weight - - -def bi_bf16_fp32_linear(input: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - """Apply the batch-invariant BF16-input, FP32-output linear contract. - - This is the differentiable form of :func:`bi_router_gemm`. It is suitable - for small-output projections whose serving counterpart uses DeepGEMM's - ``bf16_gemm_nt`` FP32 output, including GLM-5.2's indexer head weights. - """ - - if input.ndim < 2 or weight.ndim != 2: - raise ValueError(f"expected input.ndim >= 2 and weight.ndim == 2, got {input.shape=} and {weight.shape=}") - if input.shape[-1] != weight.shape[-1]: - raise ValueError(f"input/weight contraction mismatch: {input.shape[-1]} != {weight.shape[-1]}") - return _BIBf16Fp32LinearFn.apply(input, weight) - - -def bi_router_topk_weights( - topk_vals: torch.Tensor, - norm_topk_prob: bool = True, - out_dtype: torch.dtype = torch.bfloat16, -) -> torch.Tensor: - """Renormalize the gathered top-k router scores in a fixed reduction order, - then cast — the second half of the K3 router contract. - - ``bi_router_gemm`` makes the router logits (and therefore ``torch.topk`` - selection and the gathered top-k softmax scores) bitwise identical - cross-engine, but the stock renorm ``vals / vals.sum(dim=-1, keepdim=True)`` - is not: the small last-dim reduction ``sum(dim=-1)`` uses a build-dependent - tree order, so the divisor can differ by ~1 fp32 ulp between the trainer's - and the server's torch/triton, occasionally flipping the final bf16 weight - on rare tokens. Elementwise ops (add, divide, round-to-bf16) are IEEE - correctly-rounded and build-invariant, so accumulating the divisor with a - pinned left-to-right order makes the top-k weights bit-identical too. - - Args: - topk_vals: ``[N, top_k]`` fp32 gathered top-k router scores (softmax - probabilities on the softmax path). - norm_topk_prob: renormalize the top-k slice to sum to 1 (Qwen3 MoE - default). When False the scores are only cast (already bit-identical - cross-engine, since the softmax/top-k that produced them are). - out_dtype: routing-weight dtype (the model activation dtype, bf16). - - Returns: - ``[N, top_k]`` ``out_dtype`` routing weights. - """ - assert topk_vals.dtype == torch.float32, "the router contract renorms fp32 top-k scores" - if norm_topk_prob: - denom = topk_vals[..., 0] - for k in range(1, topk_vals.shape[-1]): - denom = denom + topk_vals[..., k] - topk_vals = topk_vals / denom.unsqueeze(-1) - return topk_vals.to(out_dtype) - - -_ONES_CACHE: Dict[Tuple[str, int | None, torch.dtype, int], torch.Tensor] = {} - - -def _get_or_make_ones(input: torch.Tensor, normalized_shape: list[int]) -> torch.Tensor: - assert len(normalized_shape) == 1, "Only last-dimension RMSNorm is supported" - hidden_size = int(normalized_shape[0]) - key = (input.device.type, input.device.index, input.dtype, hidden_size) - weight = _ONES_CACHE.get(key) - if weight is None or weight.device != input.device: - weight = torch.ones(hidden_size, device=input.device, dtype=input.dtype) - _ONES_CACHE[key] = weight - return weight - - -def _rms_norm_aten_compat(input, normalized_shape, weight=None, eps=None): - _guard_interpose_no_grad("aten::rms_norm", input, weight) - normalized_shape = [int(dim) for dim in normalized_shape] - if len(normalized_shape) != 1 or input.shape[-1] != normalized_shape[0]: - raise NotImplementedError("Batch-invariant RMSNorm only supports last dimension") - if weight is None: - weight = _get_or_make_ones(input, normalized_shape) - if eps is None: - eps = torch.finfo(input.dtype).eps - # The interpose IS the no-residual family: every F.rms_norm that reaches it - # is a no-residual site (qk-norm, layer-0 input norm) by the family contract. - return bi_rms_norm(input, weight, eps=eps, family=RMS_NORM_FAMILY_NO_RESIDUAL) - - -def _mm_dtype_compat(a, b, out_dtype): - out = mm_batch_invariant(a, b) - if out.dtype != out_dtype: - out = out.to(out_dtype) - return out - - -_batch_invariant_MODE = False -_batch_invariant_LIB = None -_batch_invariant_OPS: set[str] = set() -_original_torch_bmm = None - - -def is_batch_invariant_mode_enabled(): - return _batch_invariant_MODE - - -def get_batch_invariant_ops() -> tuple[str, ...]: - return tuple(sorted(_batch_invariant_OPS)) if _batch_invariant_MODE else () - - -def is_batch_invariant_op_enabled(op: str) -> bool: - op = _BATCH_INVARIANT_ALIASES.get(op, op) - return _batch_invariant_MODE and op in _batch_invariant_OPS - - -def enable_batch_invariant_mode( - enable_bmm: bool = True, -): - global _batch_invariant_MODE, _batch_invariant_LIB, _batch_invariant_OPS, _original_torch_bmm - if _batch_invariant_MODE: - return - - _batch_invariant_OPS = _parse_batch_invariant_ops() - _batch_invariant_MODE = True - _batch_invariant_LIB = torch.library.Library("aten", "IMPL") - if "mm" in _batch_invariant_OPS: - _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA") - if "addmm" in _batch_invariant_OPS: - _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA") - if "log_softmax" in _batch_invariant_OPS: - _batch_invariant_LIB.impl("aten::_log_softmax", _log_softmax_batch_invariant, "CUDA") - if "mean" in _batch_invariant_OPS: - _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA") - if "rms_norm" in _batch_invariant_OPS: - _batch_invariant_LIB.impl("aten::rms_norm", _rms_norm_aten_compat, "CUDA") - if "mm_dtype" in _batch_invariant_OPS: - _batch_invariant_LIB.impl("aten::mm.dtype", _mm_dtype_compat, "CUDA") - - if enable_bmm and "bmm" in _batch_invariant_OPS: - _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA") - - # Also monkeypatch torch.bmm directly as a fallback - _original_torch_bmm = torch.bmm - torch.bmm = bmm_batch_invariant - - -def disable_batch_invariant_mode(): - global _batch_invariant_MODE, _batch_invariant_LIB, _batch_invariant_OPS, _original_torch_bmm - if _batch_invariant_LIB is not None: - _batch_invariant_LIB._destroy() - if _original_torch_bmm is not None: - torch.bmm = _original_torch_bmm - _original_torch_bmm = None - _batch_invariant_MODE = False - _batch_invariant_LIB = None - _batch_invariant_OPS = set() - - -@contextlib.contextmanager -def set_batch_invariant_mode(enabled: bool = True): - was_enabled = _batch_invariant_MODE - if enabled == was_enabled: - yield - return - - if enabled: - enable_batch_invariant_mode() - else: - disable_batch_invariant_mode() - try: - yield - finally: - if was_enabled: - enable_batch_invariant_mode() - else: - disable_batch_invariant_mode() - - -_TRUNK_LINEAR_NAMES = ( - "qkv_proj", - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_up_proj", - "gate_proj", - "up_proj", - "down_proj", - # Qwen2/3.5-MoE shared-expert sigmoid gate: serving contracts it through the - # global interpose (plain nn.Linear -> aten::mm), so the trunk lane must too. - "shared_expert_gate", -) - -_TRUNK_LINEAR_CONTRACT_ACTIVE = False - - -def is_trunk_linear_contract_enabled() -> bool: - """True once :func:`wrap_trunk_linears_batch_invariant` armed the contract lane. - - RMSNorm dispatch keys off this to route no-residual (family-1) norms through - the serving batch-invariant kernel (the qk-norm term of the K3 contract). - """ - return _TRUNK_LINEAR_CONTRACT_ACTIVE - - -def set_trunk_linear_contract(enabled: bool) -> None: - global _TRUNK_LINEAR_CONTRACT_ACTIVE - _TRUNK_LINEAR_CONTRACT_ACTIVE = enabled - - -class _BatchInvariantTrunkLinearFn(torch.autograd.Function): - """Forward through the batch-invariant persistent GEMM; backward stays cuBLAS.""" - - @staticmethod - def forward(ctx, input, weight, bias): - if input.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: - raise RuntimeError( - f"XORL_BI_TRUNK_LINEAR contract is bf16-only; got input={input.dtype}, weight={weight.dtype}." - ) - ctx.save_for_backward(input, weight) - ctx.has_bias = bias is not None - x2d = input.reshape(-1, input.shape[-1]) - out = matmul_persistent(x2d, weight.t(), bias=bias) - return out.reshape(*input.shape[:-1], weight.shape[0]) - - @staticmethod - def backward(ctx, grad_output): - input, weight = ctx.saved_tensors - go2d = grad_output.reshape(-1, grad_output.shape[-1]) - grad_input = grad_weight = grad_bias = None - if ctx.needs_input_grad[0]: - grad_input = (go2d @ weight).reshape_as(input) - if ctx.needs_input_grad[1]: - grad_weight = go2d.t() @ input.reshape(-1, input.shape[-1]) - if ctx.has_bias and ctx.needs_input_grad[2]: - grad_bias = go2d.sum(dim=0) - return grad_input, grad_weight, grad_bias - - -def batch_invariant_trunk_linear( - input: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None = None, -) -> torch.Tensor: - """Run the public trunk-linear contract on an explicitly folded weight.""" - return _BatchInvariantTrunkLinearFn.apply(input, weight, bias) - - -def wrap_trunk_linears_batch_invariant( - model: torch.nn.Module, names: Tuple[str, ...] = _TRUNK_LINEAR_NAMES -) -> Dict[str, int]: - """Route ONLY transformer-trunk nn.Linear forwards through the batch-invariant GEMM. - - Module-scoped alternative to enable_batch_invariant_mode(): forward bits match the - serving-side matmul_persistent contract while backward/optimizer/loss stay on cuBLAS. - - Selection is explicit: a module is wrapped iff its leaf name is in ``names`` AND it is - exactly ``torch.nn.Linear``. lm_head/embeddings never match the name set; routed MoE - experts (FQN containing ``.experts.``) are skipped — they are contracted through the - fused sglang expert path. LoRA/QLoRA-wrapped modules and FP8/TE/custom Linear - subclasses RAISE: silently skipping them would void the bitwise contract the flag - promises. Idempotent (already-wrapped modules are left alone). - - Returns ``{leaf_name: wrapped_count}`` and arms the contract lane - (see :func:`is_trunk_linear_contract_enabled`). - """ - import types # noqa: PLC0415 - - from xorl.lora.fold import lora_merged_forward_enabled # noqa: PLC0415 - from xorl.lora.modules.base import LoraModule # noqa: PLC0415 - from xorl.lora.modules.delta_linear import LoraDeltaLinear # noqa: PLC0415 - from xorl.lora.modules.linear import LoraLinear # noqa: PLC0415 - - if is_batch_invariant_mode_enabled(): - raise RuntimeError( - "XORL_BI_TRUNK_LINEAR cannot be combined with the global batch-invariant interpose " - "(XORL_BATCH_INVARIANT_MATMUL): the wrapped backward would silently ride the interposed " - "aten::mm instead of cuBLAS. Pick one lane." - ) - - def _forward(self, input): - return _BatchInvariantTrunkLinearFn.apply(input, self.weight, self.bias) - - def _forward_lora_merged(self, input): - # Merged-forward LoRA lane: the BI GEMM runs on the canonically folded - # weight (identical bits to a serving engine that received W'); backward - # flows dW' through the straight-through fold into the LoRA factors. - return _BatchInvariantTrunkLinearFn.apply(input, self.merged_weight_for_forward(), self.bias) - - wrapped: Dict[str, int] = {} - already_wrapped = 0 - for module_name, module in model.named_modules(): - leaf = module_name.rsplit(".", 1)[-1] - if leaf not in names: - continue - if ".experts." in f".{module_name}.": - continue - if isinstance(module, LoraDeltaLinear) and module_name.endswith( - (".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj") - ): - # These are factor-only children of the fused gate_up_proj. The - # parent base GEMM is wrapped and Qwen3_5MoeMLP folds both logical - # projections into that one contracted call. - continue - if type(module) is LoraLinear and lora_merged_forward_enabled(module): - # Merged-forward contract lane: the adapted linear serves and trains - # through the folded weight, so the trunk contract composes. - if module.weight.dtype not in (torch.bfloat16, torch.float32): - raise RuntimeError( - f"XORL_BI_TRUNK_LINEAR: {module_name} weight is {module.weight.dtype}; the trunk " - "contract is bf16-only." - ) - if getattr(module, "_xorl_bi_trunk_wrapped", False): - already_wrapped += 1 - continue - module.forward = types.MethodType(_forward_lora_merged, module) - module._xorl_bi_trunk_wrapped = True - wrapped[leaf] = wrapped.get(leaf, 0) + 1 - continue - if isinstance(module, LoraModule): - raise NotImplementedError( - f"XORL_BI_TRUNK_LINEAR: {module_name} is adapter-wrapped ({type(module).__qualname__}); " - "the canonical merged-LoRA trunk contract composes only with a plain LoraLinear whose " - "model-owned exact_merged_forward property is true — enable it on that module or exclude the adapter." - ) - if type(module) is not torch.nn.Linear: - raise NotImplementedError( - f"XORL_BI_TRUNK_LINEAR: {module_name} is {type(module).__qualname__}, not a plain " - "nn.Linear; fp8/te/custom linears are outside the bf16 trunk contract." - ) - if module.weight.dtype not in (torch.bfloat16, torch.float32): - # fp32 is allowed here only as a mixed-precision master: FSDP2's mp_policy - # casts it to bf16 before forward, and the runtime guard in - # _BatchInvariantTrunkLinearFn enforces bf16 on the actual GEMM operands. - raise RuntimeError( - f"XORL_BI_TRUNK_LINEAR: {module_name} weight is {module.weight.dtype}; the trunk contract is bf16-only." - ) - if getattr(module, "_xorl_bi_trunk_wrapped", False): - already_wrapped += 1 - continue - module.forward = types.MethodType(_forward, module) - module._xorl_bi_trunk_wrapped = True - wrapped[leaf] = wrapped.get(leaf, 0) + 1 - - if not wrapped and not already_wrapped: - raise RuntimeError( - "XORL_BI_TRUNK_LINEAR=1 matched no trunk linears; expected leaf names " - f"{sorted(names)} — wire the model's projections or drop the flag." - ) - set_trunk_linear_contract(True) - return wrapped - - -AttentionBlockSize = namedtuple("AttentionBlockSize", ["block_m", "block_n"]) - - -def get_batch_invariant_attention_block_size() -> AttentionBlockSize: - return AttentionBlockSize(block_m=16, block_n=16) +_sys.modules[__name__] = _importlib.import_module("xorl.ops.sglang.batch_invariant_ops") diff --git a/src/xorl/ops/bi_families_v2.py b/src/xorl/ops/bi_families_v2.py index 525e6c87..627ae3e2 100644 --- a/src/xorl/ops/bi_families_v2.py +++ b/src/xorl/ops/bi_families_v2.py @@ -1,1063 +1,7 @@ -# Families v2 — the redefined frozen reduction trees for the batch-invariance -# contract (hidden-dim RMSNorm, qk-norm, and the final projection). -# -# The trainer and serving engine carry equivalent implementations because they -# must evaluate the same reduction trees. This module remains self-contained -# (torch + triton only), and cross-engine agreement is established by the -# behavior-logprob replay rather than source-file identity. -# -# Rule A: every bit-relevant reduction is written explicitly — an -# adjacent-pairwise balanced binary tree within a block (tl.split + one add -# per level; a 2-element reduction has exactly one association, so the -# compiler owns no tree choice) and a sequential scalar chain across chunks -# in index order. tl.sum over >2 elements is banned in bit-relevant positions. -# Rule B: golden-value gates pin the bits under both engines' venvs. -# -# These trees are DEFAULT ON inside an engaged contract lane. The exact -# Qwen3.5-family model setup selects the already-certified v1 family directly; -# callers do not coordinate the selection with an environment variable. -# -# Migration: v1 and v2 are two different trees, and both hold the trainer and -# the sampler bitwise equal — that is the contract, and v2 satisfies it. The -# reported values do move between the two, so goldens and frozen anchors -# recorded under v1 must be re-taken under v2, and both engines must flip -# together. +"""Moved to ``xorl.ops.sglang.bi_families_v2`` (issue #78); alias stub.""" -import os +import importlib as _importlib +import sys as _sys -import torch -import triton -import triton.language as tl -from triton.runtime.errors import OutOfResources - -_EXACT_FAMILIES_VERSION: str | None = None - - -def _select_qwen35_families_v1() -> None: - """Pin exact Qwen's fused LM-head loss to its qualified v1 program.""" - global _EXACT_FAMILIES_VERSION - _EXACT_FAMILIES_VERSION = "v1" - - -def _select_glm52_families_v2() -> None: - """Pin the GLM-5.2 process to its certified v2 reduction trees.""" - global _EXACT_FAMILIES_VERSION - _EXACT_FAMILIES_VERSION = "v2" - - -def _select_qwen3_dense_families_v2() -> None: - """Pin exact dense Qwen3 to the shared v2 norm and lm-head trees.""" - global _EXACT_FAMILIES_VERSION - _EXACT_FAMILIES_VERSION = "v2" - - -def _select_nonexact_families() -> None: - """Restore the pre-existing family selection for an ordinary model.""" - global _EXACT_FAMILIES_VERSION - _EXACT_FAMILIES_VERSION = None - - -def families_v2_enabled() -> bool: - """Return the selected reduction family for the current process. - - Exact model programs select their family structurally and ignore the - legacy rollback variables. Without an exact model selection, preserve the - pre-existing non-exact behavior for compatibility. - """ - if _EXACT_FAMILIES_VERSION is not None: - return _EXACT_FAMILIES_VERSION == "v2" - return not any(os.getenv(v, "1").lower() in _V2_OFF for v in FAMILIES_V2_ENV_VARS) - - -# Contract constants (bit-relevant; never tuning axes). -V2_NORM_BLOCK_H = 4096 # per-chunk tree width for hidden-dim norms -V2_QK_MAX_HEAD_DIM = 256 - -FAMILIES_V2_ENV_VARS = ("XORL_FAMILIES_V2", "SGLANG_FAMILIES_V2") -_V2_OFF = ("0", "false", "no") - - -@triton.jit -def _rtne_bf16(x): - # Round-to-nearest-even to bf16 via integer bitcast (triton folds - # f32->bf16->f32 cast pairs; this cannot be folded). NaNs quieted. - bits = x.to(tl.int32, bitcast=True) - rounded = (bits + 0x7FFF + ((bits >> 16) & 1)) & -65536 - nan_bits = (bits & -65536) | 0x00400000 - out = tl.where(x != x, nan_bits, rounded) - return out.to(tl.float32, bitcast=True) - - -@triton.jit -def _pairwise_tree_sum(vec, BLOCK: tl.constexpr): - # Adjacent-pairwise balanced binary tree: log2(BLOCK) levels, one explicit - # fp32 add per level. The reduction ORDER is written in the IR — no - # compiler-owned tl.sum lowering (the v1 1-ulp drift class). Statically - # unrolled; the constexpr `if` prunes levels past log2(BLOCK). - tl.static_assert(BLOCK & (BLOCK - 1) == 0, "BLOCK must be a power of 2") - tl.static_assert(BLOCK <= 4096, "12 unrolled levels cover BLOCK <= 4096") - for _ in tl.static_range(0, 12): - if vec.shape[0] > 1: - lo, hi = tl.split(tl.reshape(vec, (vec.shape[0] // 2, 2))) - vec = lo + hi - return vec # shape (1,) fp32 - - -V2_NORM_TILE = 512 # register tile; BITS-NEUTRAL by the tree factorization below - - -@triton.jit -def _rms_norm_v2_kernel( - x_ptr, - res_ptr, - w_ptr, - out_ptr, - res_out_ptr, - n_cols, - stride_x, - stride_res, - stride_out, - stride_res_out, - eps, - HAS_RESIDUAL: tl.constexpr, - ZERO_CENTERED: tl.constexpr, - BLOCK_H: tl.constexpr, -): - """Family-2'/1' unified row norm: ONE launch per call site. - - Pinned order: (residual add -> RTNE bf16 round ->) fp32 square -> - per-chunk pairwise tree -> sequential chunk chain -> /H -> tl.rsqrt -> - x * inv_rms * w (left-to-right, fp32 weight-mul) -> single cast at store. - Tree is a function of n_cols alone => batch-invariant by construction. - - For the residual specialization, this fused one-launch form serves nearly - every shipped shape. Few rows over many tiles, and the compiler-spilling - no-residual specialization, dispatch to the bit-identical split realization - below (see ``_v2_norm_use_split``), which factorizes the same tree into - per-TILE trees + a partials tree (adjacent pairing preserves contiguity at - every level => identical bits; cross-structure gates + frozen goldens). - """ - row = tl.program_id(0) - n_chunks = tl.cdiv(n_cols, BLOCK_H) - - total = tl.zeros((1,), dtype=tl.float32) - for c in range(n_chunks): - cols = c * BLOCK_H + tl.arange(0, BLOCK_H) - mask = cols < n_cols - x = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_RESIDUAL: - r = tl.load(res_ptr + row * stride_res + cols, mask=mask, other=0.0).to(tl.float32) - s = _rtne_bf16(x + r) # v1 semantic kept: add rounds to input dtype - tl.store( - res_out_ptr + row * stride_res_out + cols, - s.to(res_out_ptr.dtype.element_ty), - mask=mask, - ) - else: - s = x - sq = s * s # masked lanes contribute +0.0 (exact for sums of squares) - total = total + _pairwise_tree_sum(sq, BLOCK_H) - - var = total / n_cols.to(tl.float32) - inv_rms = tl.rsqrt(var + eps) - - for c in range(n_chunks): - cols = c * BLOCK_H + tl.arange(0, BLOCK_H) - mask = cols < n_cols - if HAS_RESIDUAL: - s = tl.load(res_out_ptr + row * stride_res_out + cols, mask=mask, other=0.0).to( - tl.float32 - ) # the rounded sum: bit-identical to what was squared - else: - s = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32) - w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32) - if ZERO_CENTERED: - w = 1.0 + w - y = s * inv_rms * w - tl.store(out_ptr + row * stride_out + cols, y.to(out_ptr.dtype.element_ty), mask=mask) - - -QK_V2_ROWS_PER_PROG = 16 # head-rows per program (perf-only, NOT bit-relevant) - - -@triton.jit -def _qk_norm_v2_kernel( - x_ptr, - w_ptr, - out_ptr, - n_rows, - n_heads, - head_dim, - stride_x_tok, - stride_x_head, - stride_out_tok, - stride_out_head, - eps, - ZERO_CENTERED: tl.constexpr, - BLOCK_D: tl.constexpr, - ROWS: tl.constexpr, -): - """Family-1' strided qk-norm: reads head rows straight out of the packed - qkv projection (no reshape/contiguous copies). Same tree as - _rms_norm_v2_kernel (single chunk: head_dim <= V2_QK_MAX_HEAD_DIM). - Each program handles ROWS independent head-rows (row batching is grid - shape only — per-row math identical, like BLOCK_M in the GEMM). - In-place safe per row: the full head row is loaded before any store. - """ - pid = tl.program_id(0) - rows = pid * ROWS + tl.arange(0, ROWS) - row_mask = rows < n_rows - rows_safe = tl.where(row_mask, rows, 0) - tok = rows_safe // n_heads - head = rows_safe % n_heads - d = tl.arange(0, BLOCK_D) - col_mask = d < head_dim - mask = row_mask[:, None] & col_mask[None, :] - base = tok * stride_x_tok + head * stride_x_head - x = tl.load(x_ptr + base[:, None] + d[None, :], mask=mask, other=0.0).to(tl.float32) - total = _pairwise_tree_sum_rows(x * x, BLOCK_D) - var = total / head_dim.to(tl.float32) - inv_rms = tl.rsqrt(var + eps) - w = tl.load(w_ptr + d, mask=col_mask, other=0.0).to(tl.float32) - if ZERO_CENTERED: - w = 1.0 + w - y = x * inv_rms[:, None] * w[None, :] - out_base = tok * stride_out_tok + head * stride_out_head - tl.store( - out_ptr + out_base[:, None] + d[None, :], - y.to(out_ptr.dtype.element_ty), - mask=mask, - ) - - -def rms_norm_v2( - x: torch.Tensor, - weight: torch.Tensor, - eps: float = 1e-6, - *, - residual: torch.Tensor | None = None, - zero_centered: bool = False, -): - """v2 hidden-dim RMSNorm (family-2' residual / no-residual, family-1' - site-classes unified on one tree). Returns ``out`` or ``(out, residual_out)``. - - bf16-only (the contract dtype); 2D ``[rows, H]`` input, unit column stride. - """ - assert x.ndim == 2 and x.stride(1) == 1, "x must be [rows, H], unit col stride" - assert x.dtype == torch.bfloat16, "families v2 is bf16-only (contract dtype)" - assert weight.ndim == 1 and weight.shape[0] == x.shape[1] - assert x.is_cuda - rows, H = x.shape - weight = weight.contiguous() - if residual is not None: - assert residual.shape == x.shape and residual.stride(1) == 1 - assert residual.dtype == torch.bfloat16 - if _v2_norm_use_split( - rows, - triton.cdiv(H, V2_NORM_TILE), - has_residual=residual is not None, - is_hopper=torch.cuda.get_device_capability(x.device) == (9, 0), - ): - return _rms_norm_v2_split(x, weight, eps, residual, zero_centered) - return _rms_norm_v2_fused(x, weight, eps, residual, zero_centered) - - -def _rms_norm_v2_fused(x, weight, eps, residual, zero_centered): - """One-launch realization of the v2 norm tree (see ``_rms_norm_v2_kernel``).""" - rows, H = x.shape - out = torch.empty_like(x) - if residual is not None: - res_out = torch.empty_like(x) - _rms_norm_v2_kernel[(rows,)]( - x, - residual, - weight, - out, - res_out, - H, - x.stride(0), - residual.stride(0), - out.stride(0), - res_out.stride(0), - eps, - HAS_RESIDUAL=True, - ZERO_CENTERED=zero_centered, - BLOCK_H=V2_NORM_BLOCK_H, - ) - return out, res_out - _rms_norm_v2_kernel[(rows,)]( - x, - x, - weight, - out, - out, - H, - x.stride(0), - 0, - out.stride(0), - 0, - eps, - HAS_RESIDUAL=False, - ZERO_CENTERED=zero_centered, - BLOCK_H=V2_NORM_BLOCK_H, - ) - return out - - -def qk_norm_v2( - x: torch.Tensor, - weight: torch.Tensor, - eps: float = 1e-6, - *, - head_dim: int, - out: torch.Tensor | None = None, - zero_centered: bool = False, -): - """v2 per-head qk-norm over strided input. - - ``x``: ``[T, n_heads * head_dim]`` view into the packed qkv output — any - row stride, unit element stride, heads contiguous within a row. ``out`` - defaults to a fresh tensor (trainer); pass ``out=x`` for in-place (serving). - """ - assert x.ndim == 2 and x.stride(1) == 1 - assert x.dtype == torch.bfloat16, "families v2 is bf16-only (contract dtype)" - assert x.shape[1] % head_dim == 0 - assert head_dim <= V2_QK_MAX_HEAD_DIM - assert weight.ndim == 1 and weight.shape[0] == head_dim - assert x.is_cuda - T = x.shape[0] - n_heads = x.shape[1] // head_dim - weight = weight.contiguous() - if out is None: - out = torch.empty_like(x) - else: - assert out.shape == x.shape and out.stride(1) == 1 and out.dtype == x.dtype - if T > 0: - n_rows = T * n_heads - _qk_norm_v2_kernel[(triton.cdiv(n_rows, QK_V2_ROWS_PER_PROG),)]( - x, - weight, - out, - n_rows, - n_heads, - head_dim, - x.stride(0), - head_dim, - out.stride(0), - head_dim, - eps, - ZERO_CENTERED=zero_centered, - BLOCK_D=triton.next_power_of_2(head_dim), - ROWS=QK_V2_ROWS_PER_PROG, - ) - return out - - -# --------------------------------------------------------------------------- -# Head v2 — online-LSE lm-head (component 3, design note §5) -# -# The GEMM K-chain is v1's matmul_kernel_persistent VERBATIM (pinned -# BLOCK_SIZE_K; logits bitwise == the v1 head GEMM — gated). New: per-tile -# (m, l) stats computed in the epilogue from the fp32 accumulator, one pinned -# merge kernel. BLOCK_SIZE_N tiles a stats reduction, so it is a CONTRACT -# CONSTANT here (HEAD_STATS_TILE_N) — never shape/M-keyed, or decode and -# trainer scoring rows would grow different LSE trees. The OOM fallback chain -# varies only bit-neutral axes (BLOCK_M / GROUP / stages / warps). -# --------------------------------------------------------------------------- - -HEAD_V2_STATS_TILE_N = 256 # contract constant (bit-relevant) -HEAD_V2_BLOCK_K = 64 # == v1 PINNED_BLOCK_K[bf16] (bit-relevant) - -# Bit-neutral launch axes, M-bucketed (perf only; all candidates share N/K). -_HEAD_V2_LAUNCH = ( - (16, {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 4}), - (256, {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 4}), - (None, {"BLOCK_SIZE_M": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), -) -_HEAD_V2_FALLBACK = { - "BLOCK_SIZE_M": 64, - "GROUP_SIZE_M": 8, - "num_stages": 2, - "num_warps": 4, -} - - -@triton.jit -def _pairwise_tree_sum_rows(mat, BLOCK_N: tl.constexpr): - # Row-wise adjacent-pairwise tree over axis 1 (same construction as - # _pairwise_tree_sum; each level is one elementwise add — no compiler - # tree choice, warp count is bit-neutral by construction). - tl.static_assert(BLOCK_N & (BLOCK_N - 1) == 0, "BLOCK_N must be a power of 2") - tl.static_assert(BLOCK_N <= 4096, "12 unrolled levels cover BLOCK_N <= 4096") - for _ in tl.static_range(0, 12): - if mat.shape[1] > 1: - lo, hi = tl.split(tl.reshape(mat, (mat.shape[0], mat.shape[1] // 2, 2))) - mat = lo + hi - return tl.reshape(mat, (mat.shape[0],)) - - -@triton.jit -def _head_v2_gemm_stats_kernel( - a_ptr, - b_ptr, - c_ptr, - m_out_ptr, - l_out_ptr, - sel_ptr, - tok_ptr, - temp_ptr, - M, - N, - K, - stride_am, - stride_ak, - stride_bk, - stride_bn, - stride_cm, - stride_cn, - stride_mm, - stride_lm, - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr, - NUM_SMS: tl.constexpr, - A_LARGE: tl.constexpr, - B_LARGE: tl.constexpr, - C_LARGE: tl.constexpr, - STORE_LOGITS: tl.constexpr, - HAS_TOKS: tl.constexpr, - HAS_TEMP: tl.constexpr, -): - """v1 persistent GEMM (identical K-chain) + per-tile (m, l) epilogue stats. - - Stats are computed from the fp32 accumulator in registers: m = row max of - the tile (order-free), l = row sum of exp(x - m) via the pairwise tree. - Temperature scales stats (and the gathered selected logit) but NOT the - stored logits — matching v1, where the stats kernel applies 1/T while the - decode logits buffer stays unscaled for the sampler. - """ - start_pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - k_tiles = tl.cdiv(K, BLOCK_SIZE_K) - num_tiles = num_pid_m * num_pid_n - - offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - - for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): - group_id = tile_id // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + (tile_id % group_size_m) - pid_n = (tile_id % num_pid_in_group) // group_size_m - - start_m = pid_m * BLOCK_SIZE_M - start_n = pid_n * BLOCK_SIZE_N - offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) - offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) - if A_LARGE: - offs_am = offs_am.to(tl.int64) - if B_LARGE: - offs_bn = offs_bn.to(tl.int64) - offs_am = tl.where(offs_am < M, offs_am, 0) - offs_bn = tl.where(offs_bn < N, offs_bn, 0) - offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) - offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for ki in range(k_tiles): - if A_LARGE or B_LARGE: - offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) - else: - offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) - a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - a = tl.load(a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0) - b = tl.load(b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0) - accumulator = tl.dot(a, b, accumulator) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - if C_LARGE: - offs_cm = offs_cm.to(tl.int64) - offs_cn = offs_cn.to(tl.int64) - row_mask = offs_cm < M - col_mask = offs_cn < N - - if STORE_LOGITS: - c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - tl.store( - c_ptrs, - accumulator.to(tl.float32), - mask=row_mask[:, None] & col_mask[None, :], - ) - - # --- epilogue stats (the new frozen tree) --- - x_stats = accumulator - if HAS_TEMP: - temp = tl.load(temp_ptr + offs_cm, mask=row_mask, other=1.0) - inv_t = 1.0 / temp - x_stats = x_stats * inv_t[:, None] - neg_inf = float("-inf") - m_tile = tl.max(tl.where(col_mask[None, :], x_stats, neg_inf), axis=1) - terms = tl.where(col_mask[None, :], tl.exp(x_stats - m_tile[:, None]), 0.0) - l_tile = _pairwise_tree_sum_rows(terms, BLOCK_SIZE_N) - tl.store(m_out_ptr + offs_cm * stride_mm + pid_n, m_tile, mask=row_mask) - tl.store(l_out_ptr + offs_cm * stride_lm + pid_n, l_tile, mask=row_mask) - - if HAS_TOKS: - toks = tl.load(tok_ptr + offs_cm, mask=row_mask, other=-1) - tok_here = (toks >= start_n) & (toks < start_n + BLOCK_SIZE_N) & row_mask - # -inf-padded max gather: exactly one lane matches, sign of zero preserved - sel = tl.max(tl.where(offs_cn[None, :] == toks[:, None], x_stats, neg_inf), axis=1) - tl.store(sel_ptr + offs_cm, sel, mask=tok_here) - - -@triton.jit -def _head_v2_lse_merge_kernel( - m_ptr, - l_ptr, - sel_ptr, - lse_ptr, - lp_ptr, - n_tiles, - stride_mm, - stride_lm, - HAS_SEL: tl.constexpr, - BLOCK_T: tl.constexpr, -): - """Pinned merge: exact global max over tile maxima (order-free), rescaled - sumexp via the pairwise tree over tiles (padded l=0 / m=-inf), then - lse = gmax + log(acc); logprob = min(sel - lse, 0) (the p~1 boundary - clamp, kept from v1).""" - row = tl.program_id(0) - t = tl.arange(0, BLOCK_T) - mask = t < n_tiles - neg_inf = float("-inf") - m_t = tl.load(m_ptr + row * stride_mm + t, mask=mask, other=neg_inf) - l_t = tl.load(l_ptr + row * stride_lm + t, mask=mask, other=0.0) - gmax = tl.max(m_t) - terms = tl.where(mask, l_t * tl.exp(m_t - gmax), 0.0) - acc = _pairwise_tree_sum(terms, BLOCK_T) - lse = gmax + tl.log(acc) - tl.store(lse_ptr + row + tl.arange(0, 1), lse) - if HAS_SEL: - sel = tl.load(sel_ptr + row) - lp = tl.minimum(sel - lse, 0.0) - tl.store(lp_ptr + row + tl.arange(0, 1), lp) - - -def _head_v2_launch(hidden, weight_t, logits, m_buf, l_buf, sel, toks, temp): - M, K = hidden.shape - N = weight_t.shape[1] - NUM_SMS = torch.cuda.get_device_properties(hidden.device).multi_processor_count - a_large = hidden.numel() > 2**31 - b_large = weight_t.numel() > 2**31 - c_large = logits is not None and (logits.shape[0] * logits.stride(0) + logits.shape[1] * logits.stride(1) > 2**31) - launch = _HEAD_V2_FALLBACK - for bound, cfg in _HEAD_V2_LAUNCH: - if bound is None or M <= bound: - launch = cfg - break - - def _run(cfg): - grid = min( - NUM_SMS, - triton.cdiv(M, cfg["BLOCK_SIZE_M"]) * triton.cdiv(N, HEAD_V2_STATS_TILE_N), - ) - _head_v2_gemm_stats_kernel[(grid,)]( - hidden, - weight_t, - logits if logits is not None else hidden, - m_buf, - l_buf, - sel if sel is not None else m_buf, - toks if toks is not None else m_buf, - temp if temp is not None else m_buf, - M, - N, - K, - hidden.stride(0), - hidden.stride(1), - weight_t.stride(0), - weight_t.stride(1), - logits.stride(0) if logits is not None else 0, - logits.stride(1) if logits is not None else 0, - m_buf.stride(0), - l_buf.stride(0), - BLOCK_SIZE_M=cfg["BLOCK_SIZE_M"], - BLOCK_SIZE_N=HEAD_V2_STATS_TILE_N, - BLOCK_SIZE_K=HEAD_V2_BLOCK_K, - GROUP_SIZE_M=cfg["GROUP_SIZE_M"], - NUM_SMS=NUM_SMS, - A_LARGE=a_large, - B_LARGE=b_large, - C_LARGE=c_large, - STORE_LOGITS=logits is not None, - HAS_TOKS=toks is not None, - HAS_TEMP=temp is not None, - num_stages=cfg["num_stages"], - num_warps=cfg["num_warps"], - ) - - try: - _run(launch) - except OutOfResources: - _run(_HEAD_V2_FALLBACK) # bit-identical: only bit-neutral axes differ - - -def _head_v2_merge(m_buf, l_buf, sel): - M, n_tiles = m_buf.shape - lse = torch.empty((M,), dtype=torch.float32, device=m_buf.device) - lp = torch.empty((M,), dtype=torch.float32, device=m_buf.device) if sel is not None else None - _head_v2_lse_merge_kernel[(M,)]( - m_buf, - l_buf, - sel if sel is not None else m_buf, - lse, - lp if lp is not None else lse, - n_tiles, - m_buf.stride(0), - l_buf.stride(0), - HAS_SEL=sel is not None, - BLOCK_T=max(triton.next_power_of_2(n_tiles), 2), - ) - return lse, lp - - -def _head_v2_check_inputs(hidden, weight, token_ids, temperature): - assert hidden.ndim == 2 and weight.ndim == 2 and hidden.shape[1] == weight.shape[1] - assert hidden.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( - "the head contract takes bf16 hidden/weight (fp32 upcast is exact inside the GEMM)" - ) - assert hidden.is_cuda - hidden = hidden.contiguous() - M = hidden.shape[0] - if token_ids is not None: - token_ids = token_ids.contiguous().to(device=hidden.device, dtype=torch.int64) - assert token_ids.shape == (M,) - if temperature is not None: - temperature = temperature.reshape(-1).to(device=hidden.device, dtype=torch.float32).contiguous() - assert temperature.shape == (M,) - torch._assert_async((temperature > 0).all(), "temperature must be > 0") - return hidden, token_ids, temperature - - -def head_v2_selected_logprob( - hidden: torch.Tensor, - weight: torch.Tensor, - token_ids: torch.Tensor, - temperature: torch.Tensor | None = None, -): - """Scoring path (trainer bi_fused v2 / serving prefill input logprobs): - logits are NEVER materialized — per-tile stats + selected gather in the - GEMM epilogue, one merge launch. Returns ``(logprob, lse, selected)``, - all ``[N]`` fp32 (temperature-scaled when given), like v1.""" - hidden, token_ids, temperature = _head_v2_check_inputs(hidden, weight, token_ids, temperature) - M = hidden.shape[0] - V = weight.shape[0] - n_tiles = triton.cdiv(V, HEAD_V2_STATS_TILE_N) - if M == 0: - z = torch.empty((0,), dtype=torch.float32, device=hidden.device) - return z, z.clone(), z.clone() - m_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=hidden.device) - l_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=hidden.device) - sel = torch.empty((M,), dtype=torch.float32, device=hidden.device) - _head_v2_launch(hidden, weight.t(), None, m_buf, l_buf, sel, token_ids, temperature) - lse, lp = _head_v2_merge(m_buf, l_buf, sel) - return lp, lse, sel - - -def head_v2_full_logits_with_lse( - hidden: torch.Tensor, - weight: torch.Tensor, - temperature: torch.Tensor | None = None, -): - """Decode path: full ``[N, V]`` fp32 logits for the sampler (bitwise == - the v1 head GEMM — same kernel body, same pinned K-chain) PLUS the lse - from the epilogue stats. The sampled token's logprob afterwards is - ``clamp_max(logits[row, tok] - lse[row], 0.0)`` — elementwise aten, no - second stats pass, no BI log_softmax.""" - hidden, _, temperature = _head_v2_check_inputs(hidden, weight, None, temperature) - M = hidden.shape[0] - V = weight.shape[0] - n_tiles = triton.cdiv(V, HEAD_V2_STATS_TILE_N) - logits = torch.empty((M, V), dtype=torch.float32, device=hidden.device) - if M == 0: - return logits, torch.empty((0,), dtype=torch.float32, device=hidden.device) - m_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=hidden.device) - l_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=hidden.device) - _head_v2_launch(hidden, weight.t(), logits, m_buf, l_buf, None, None, temperature) - lse, _ = _head_v2_merge(m_buf, l_buf, None) - return logits, lse - - -@triton.jit -def _exact_temperature_scale_fp32_kernel( - logits_ptr, - temperature_ptr, - output_ptr, - n_cols, - logits_row_stride, - output_row_stride, - BLOCK_N: tl.constexpr, -): - row = tl.program_id(0) - cols = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) - mask = cols < n_cols - values = tl.load( - logits_ptr + row * logits_row_stride + cols, - mask=mask, - other=0.0, - ) - inv_t = 1.0 / tl.load(temperature_ptr + row) - values = values * inv_t - tl.store(output_ptr + row * output_row_stride + cols, values, mask=mask) - - -def exact_temperature_scale_fp32_logits( - logits: torch.Tensor, - temperature: torch.Tensor, -) -> torch.Tensor: - """Materialize the exact per-row FP32 sampling transform.""" - - assert logits.ndim == 2 and logits.dtype == torch.float32 - assert logits.stride(1) == 1, "logits rows must be unit-stride" - assert temperature.dtype == torch.float32 - assert temperature.device == logits.device - assert temperature.is_contiguous() - temperature = temperature.reshape(-1) - assert temperature.shape == (logits.shape[0],) - torch._assert_async( - (torch.isfinite(temperature) & (temperature > 0)).all(), - "temperature must contain finite values > 0", - ) - - output = torch.empty_like(logits) - if logits.numel() == 0: - return output - if logits.is_cuda: - block_n = 1024 - _exact_temperature_scale_fp32_kernel[(logits.shape[0], triton.cdiv(logits.shape[1], block_n))]( - logits, - temperature, - output, - logits.shape[1], - logits.stride(0), - output.stride(0), - BLOCK_N=block_n, - ) - else: - output.copy_(logits * (1.0 / temperature).unsqueeze(1)) - return output - - -def exact_temperature_scale_bf16_logits( - logits: torch.Tensor, - temperature: torch.Tensor | None, -) -> torch.Tensor: - """Apply DSV4's BF16 divide/store temperature boundary.""" - - assert logits.ndim == 2 and logits.dtype == torch.bfloat16 - assert logits.stride(1) == 1, "logits rows must be unit-stride" - if temperature is None: - return logits - assert temperature.dtype == torch.float32 - assert temperature.device == logits.device - assert temperature.is_contiguous() - temperature = temperature.reshape(-1) - assert temperature.shape == (logits.shape[0],) - torch._assert_async( - (torch.isfinite(temperature) & (temperature > 0)).all(), - "temperature must contain finite values > 0", - ) - return logits.bfloat16().div(temperature.unsqueeze(1)).bfloat16() - - -@triton.jit -def _head_v2_stats_from_logits_kernel( - logits_ptr, - m_out_ptr, - l_out_ptr, - sel_ptr, - tok_ptr, - temp_ptr, - N, - n_tiles, - stride_lm, - stride_mm, - stride_lm_out, - HAS_TOKS: tl.constexpr, - HAS_TEMP: tl.constexpr, - TILE_N: tl.constexpr, -): - """Per-tile (m, l) stats over MATERIALIZED fp32 logits — the same tree as - the epilogue stats (same TILE_N, same max/exp/pairwise-tree expressions on - the same fp32 values), for callers that must sample from the logits they - score (the decode rescore hook). Bitwise == the epilogue path.""" - pid = tl.program_id(0) - row = pid // n_tiles - tile = pid % n_tiles - start_n = tile * TILE_N - offs_n = start_n + tl.arange(0, TILE_N) - col_mask = offs_n < N - x = tl.load(logits_ptr + row * stride_lm + offs_n, mask=col_mask, other=0.0) - if HAS_TEMP: - inv_t = 1.0 / tl.load(temp_ptr + row) - x = x * inv_t - neg_inf = float("-inf") - m_tile = tl.max(tl.where(col_mask, x, neg_inf)) - terms = tl.where(col_mask, tl.exp(x - m_tile), 0.0) - l_tile = _pairwise_tree_sum(terms, TILE_N) - tl.store( - m_out_ptr + row * stride_mm + tile + tl.arange(0, 1), - m_tile + tl.zeros((1,), dtype=tl.float32), - ) - tl.store(l_out_ptr + row * stride_lm_out + tile + tl.arange(0, 1), l_tile) - if HAS_TOKS: - tok = tl.load(tok_ptr + row) - if (tok >= start_n) & (tok < start_n + TILE_N): - sel = tl.max(tl.where(offs_n == tok, x, neg_inf)) - tl.store(sel_ptr + row, sel) - - -def head_v2_selected_logprob_from_logits( - logits: torch.Tensor, - token_ids: torch.Tensor, - temperature: torch.Tensor | None = None, -): - """Head-v2 twin of the v1 from-logits rescore: same (m, l) tree over an - existing fp32 logits tensor, so with logits from - ``head_v2_full_logits_with_lse`` (or the bitwise-equal v1 GEMM) the result - is bitwise identical to ``head_v2_selected_logprob`` on the same - hidden/weight/temperature. Returns ``(logprob, lse, selected)``.""" - assert logits.ndim == 2 and logits.dtype == torch.float32 and logits.is_cuda - assert logits.stride(1) == 1, "logits rows must be unit-stride" - M, V = logits.shape - token_ids = token_ids.contiguous().to(device=logits.device, dtype=torch.int64) - assert token_ids.shape == (M,) - if temperature is not None: - temperature = temperature.reshape(-1).to(device=logits.device, dtype=torch.float32).contiguous() - assert temperature.shape == (M,) - torch._assert_async((temperature > 0).all(), "temperature must be > 0") - n_tiles = triton.cdiv(V, HEAD_V2_STATS_TILE_N) - if M == 0: - z = torch.empty((0,), dtype=torch.float32, device=logits.device) - return z, z.clone(), z.clone() - m_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=logits.device) - l_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=logits.device) - sel = torch.empty((M,), dtype=torch.float32, device=logits.device) - _head_v2_stats_from_logits_kernel[(M * n_tiles,)]( - logits, - m_buf, - l_buf, - sel, - token_ids, - temperature if temperature is not None else m_buf, - V, - n_tiles, - logits.stride(0), - m_buf.stride(0), - l_buf.stride(0), - HAS_TOKS=True, - HAS_TEMP=temperature is not None, - TILE_N=HEAD_V2_STATS_TILE_N, - ) - lse, lp = _head_v2_merge(m_buf, l_buf, sel) - return lp, lse, sel - - -# --- split realization of the SAME norm tree (structure is bit-neutral) ----- -# -# The split realization computes the IDENTICAL tree — per-tile partial trees in -# a 2D-parallel kernel, the pinned combine in a per-row kernel, elementwise -# normalize in a 2D-parallel kernel — with fp32 partials stored/loaded exactly -# and the bf16 residual round-trip identical to the fused kernel's own res_out -# reload. Bit-equality fused==split is gated by the families-v2 norm tests, -# which force each realization explicitly rather than relying on the rule -# below to select one, and both realizations check the same frozen goldens. -# -# It pays three launches and an HBM round-trip for the partials, and buys -# rows*n_tiles-way parallelism where the fused grid=(rows,) has only rows-way. -# For the residual specialization, that trade pays at few rows over many tiles. -# The no-residual fused specialization currently spills on Hopper, so it always -# uses the split form (see _v2_norm_use_split for the measured boundary). - -V2_NORM_SPLIT_MIN_TILES = 10 # no-residual crossover on non-Hopper devices -V2_NORM_RESIDUAL_SPLIT_MIN_TILES = 13 # keeps all shipped hidden sizes (at most 12 tiles) fused - - -def _v2_norm_use_split( - rows: int, - n_tiles: int, - *, - has_residual: bool = True, - is_hopper: bool = False, -) -> bool: - """Structure switch (perf-only, NOT bit-relevant). - - The no-residual fused specialization spills on Hopper, while the split - realization is bit-identical and avoids that register-pressure cliff. - Other architectures and the residual specialization retain the measured - few-rows/many-tiles switch, where the fused ``grid=(rows,)`` cannot fill - the GPU. - - ``n_tiles`` must be the split kernel's 512-wide tile count, not the fused - kernel's 4096-wide chunk count. The threshold is a Hopper performance - policy. Both realizations compute the same tree, so retuning it may change - speed only; the fused-versus-split equality gate must still pass. - """ - if is_hopper and not has_residual: - return True - min_tiles = V2_NORM_RESIDUAL_SPLIT_MIN_TILES if has_residual else V2_NORM_SPLIT_MIN_TILES - return n_tiles >= min_tiles and rows <= n_tiles - - -@triton.jit -def _rms_norm_v2_partials_kernel( - x_ptr, - res_ptr, - res_out_ptr, - part_ptr, - n_cols, - n_tiles, - stride_x, - stride_res, - stride_res_out, - HAS_RESIDUAL: tl.constexpr, - TILE: tl.constexpr, -): - pid = tl.program_id(0) - row = pid // n_tiles - t = pid % n_tiles - cols = t * TILE + tl.arange(0, TILE) - mask = cols < n_cols - x = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32) - if HAS_RESIDUAL: - r = tl.load(res_ptr + row * stride_res + cols, mask=mask, other=0.0).to(tl.float32) - s = _rtne_bf16(x + r) - tl.store( - res_out_ptr + row * stride_res_out + cols, - s.to(res_out_ptr.dtype.element_ty), - mask=mask, - ) - else: - s = x - p = _pairwise_tree_sum(s * s, TILE) - tl.store(part_ptr + row * n_tiles + t + tl.arange(0, 1), p) - - -@triton.jit -def _rms_norm_v2_invrms_kernel( - part_ptr, - invrms_ptr, - n_cols, - n_tiles, - eps, - CHUNK_TILES: tl.constexpr, - BLOCK_T: tl.constexpr, -): - # combine = per-BLOCK_H-chunk pairwise tree over its CHUNK_TILES partials, - # then the sequential cross-chunk chain — identical to the fused kernel. - row = tl.program_id(0) - n_chunks = tl.cdiv(n_tiles, CHUNK_TILES) - total = tl.zeros((1,), dtype=tl.float32) - for c in range(n_chunks): - idx = c * CHUNK_TILES + tl.arange(0, CHUNK_TILES) - p = tl.load(part_ptr + row * n_tiles + idx, mask=idx < n_tiles, other=0.0) - total = total + _pairwise_tree_sum(p, CHUNK_TILES) - var = total / n_cols.to(tl.float32) - inv_rms = tl.rsqrt(var + eps) - tl.store(invrms_ptr + row + tl.arange(0, 1), inv_rms) - - -@triton.jit -def _rms_norm_v2_normalize_kernel( - s_ptr, - w_ptr, - invrms_ptr, - out_ptr, - n_cols, - n_tiles, - stride_s, - stride_out, - ZERO_CENTERED: tl.constexpr, - TILE: tl.constexpr, -): - pid = tl.program_id(0) - row = pid // n_tiles - t = pid % n_tiles - cols = t * TILE + tl.arange(0, TILE) - mask = cols < n_cols - s = tl.load(s_ptr + row * stride_s + cols, mask=mask, other=0.0).to(tl.float32) - w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32) - if ZERO_CENTERED: - w = 1.0 + w - inv_rms = tl.load(invrms_ptr + row) - y = s * inv_rms * w - tl.store(out_ptr + row * stride_out + cols, y.to(out_ptr.dtype.element_ty), mask=mask) - - -def _rms_norm_v2_split(x, weight, eps, residual, zero_centered): - rows, H = x.shape - n_tiles = triton.cdiv(H, V2_NORM_TILE) - out = torch.empty_like(x) - part = torch.empty((rows, n_tiles), dtype=torch.float32, device=x.device) - invrms = torch.empty((rows,), dtype=torch.float32, device=x.device) - if residual is not None: - res_out = torch.empty_like(x) - _rms_norm_v2_partials_kernel[(rows * n_tiles,)]( - x, - residual, - res_out, - part, - H, - n_tiles, - x.stride(0), - residual.stride(0), - res_out.stride(0), - HAS_RESIDUAL=True, - TILE=V2_NORM_TILE, - ) - s_buf = res_out - else: - res_out = None - _rms_norm_v2_partials_kernel[(rows * n_tiles,)]( - x, - x, - x, - part, - H, - n_tiles, - x.stride(0), - 0, - 0, - HAS_RESIDUAL=False, - TILE=V2_NORM_TILE, - ) - s_buf = x - _rms_norm_v2_invrms_kernel[(rows,)]( - part, - invrms, - H, - n_tiles, - eps, - CHUNK_TILES=V2_NORM_BLOCK_H // V2_NORM_TILE, - BLOCK_T=max(triton.next_power_of_2(n_tiles), 2), - ) - _rms_norm_v2_normalize_kernel[(rows * n_tiles,)]( - s_buf, - weight, - invrms, - out, - H, - n_tiles, - s_buf.stride(0), - out.stride(0), - ZERO_CENTERED=zero_centered, - TILE=V2_NORM_TILE, - ) - return (out, res_out) if residual is not None else out +_sys.modules[__name__] = _importlib.import_module("xorl.ops.sglang.bi_families_v2") diff --git a/src/xorl/ops/exact/__init__.py b/src/xorl/ops/exact/__init__.py index c77ddae0..8055b585 100644 --- a/src/xorl/ops/exact/__init__.py +++ b/src/xorl/ops/exact/__init__.py @@ -4,18 +4,11 @@ ``canonical_`` / ``class_b`` prefixes all meant: byte-pinned programs shared with the serving engine. -Physically here: ``rope_class_b``, ``canonical_moe_leaf``, -``canonical_moe_cast``, ``kernel_config_pin``, ``bi_gemm_configs``, -``block_fp8_native``, ``fused_silu_and_mul``. +Here: ``sampling_transforms`` (the replay contract), ``rope_class_b``, +``canonical_moe_leaf``, ``canonical_moe_cast``, ``kernel_config_pin``, +``bi_gemm_configs``, ``block_fp8_native``, ``fused_silu_and_mul``. -Aliased here but deliberately NOT moved: - -- ``families_v2`` -> :mod:`xorl.ops.bi_families_v2` — vendored byte-identical - into the serving engine and sha256-gated; the file cannot move. -- ``batch_invariant`` -> :mod:`xorl.ops.batch_invariant_ops` — vendored- - adapted from SGLang's ``srt/batch_invariant_ops``; it stays a single file - at its path so it remains diffable against the serving twin. -- ``sampling_transforms`` -> :mod:`xorl.ops.exact_sampling_transforms` — the - replay contract; kept in place while in-flight work (#74) rewrites it, to - be flipped to canonical here afterwards. +The modules that exist as literal twins inside the serving engine +(``bi_families_v2``, ``batch_invariant_ops``) live in +:mod:`xorl.ops.sglang`, which carries the paired-edit policy. """ diff --git a/src/xorl/ops/exact/batch_invariant.py b/src/xorl/ops/exact/batch_invariant.py deleted file mode 100644 index 943e6769..00000000 --- a/src/xorl/ops/exact/batch_invariant.py +++ /dev/null @@ -1,3 +0,0 @@ -"""Alias for :mod:`xorl.ops.batch_invariant_ops` (parity twin of SGLang's copy).""" - -from xorl.ops.batch_invariant_ops import * # noqa: F401,F403 diff --git a/src/xorl/ops/exact/families_v2.py b/src/xorl/ops/exact/families_v2.py deleted file mode 100644 index 82cadc9f..00000000 --- a/src/xorl/ops/exact/families_v2.py +++ /dev/null @@ -1,3 +0,0 @@ -"""Alias for :mod:`xorl.ops.bi_families_v2` (sha256-gated; cannot move).""" - -from xorl.ops.bi_families_v2 import * # noqa: F401,F403 diff --git a/src/xorl/ops/exact/sampling_transforms.py b/src/xorl/ops/exact/sampling_transforms.py index 65ae89c6..3904c644 100644 --- a/src/xorl/ops/exact/sampling_transforms.py +++ b/src/xorl/ops/exact/sampling_transforms.py @@ -1,3 +1,313 @@ -"""Alias for :mod:`xorl.ops.exact_sampling_transforms` (kept in place during #74).""" +"""The backend-independent sampling-transform program for exact RL lanes. -from xorl.ops.exact_sampling_transforms import * # noqa: F401,F403 +Qwen, GLM, and DSV4 first apply their declared temperature dtype/store +boundary, then call this module for one shared program: stable descending +probability order (token ID breaks ties), joint top-k plus inclusive-crossing +top-p plus min-p relative to the original row maximum, and normalization on +exactly that support. Serving performs per-row seeded Gumbel-max on the masked +logits. This is an exact-lane contract; it intentionally makes no claim about +generic SGLang or FlashInfer filter semantics. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable + +import torch + + +TOP_K_ALL = 1 << 30 +EXACT_FILTER_ROW_CHUNK = 32 +EXACT_SAMPLING_TRANSFORM_PROGRAM = ( + "temperature_then_stable_token_id_topk_inclusive_topp_original_max_minp_seeded_gumbel_v1" +) +SamplingTransformRows = tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None] +NativeSelectedScore = Callable[ + [torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor, torch.Tensor], +] + + +def _normalize_row_metadata( + value: int | float | torch.Tensor, + *, + rows: int, + device: torch.device, + dtype: torch.dtype, + name: str, +) -> torch.Tensor: + if isinstance(value, torch.Tensor): + if value.dtype is not dtype: + raise TypeError(f"per-row {name} must be {dtype}, got {value.dtype}") + if value.device != device: + raise ValueError(f"per-row {name} must share the logits device") + if tuple(value.shape) != (rows,): + raise ValueError(f"per-row {name} must have shape ({rows},), got {tuple(value.shape)}") + if not value.is_contiguous() or value.requires_grad: + raise ValueError(f"per-row {name} must be contiguous, non-differentiable sampling metadata") + return value + return torch.full((rows,), value, dtype=dtype, device=device) + + +def normalize_exact_sampling_transforms( + top_k: int | torch.Tensor = TOP_K_ALL, + top_p: float | torch.Tensor = 1.0, + min_p: float | torch.Tensor = 0.0, + *, + rows: int, + device: torch.device, +) -> SamplingTransformRows: + """Validate transforms and return row tensors, or three ``None`` values. + + Returning ``None`` for the all-identity scalar case is the explicit switch + that keeps established no-filter kernels and bytes untouched. + """ + + scalar_identity = ( + not isinstance(top_k, torch.Tensor) + and not isinstance(top_p, torch.Tensor) + and not isinstance(min_p, torch.Tensor) + and int(top_k) >= TOP_K_ALL + and float(top_p) == 1.0 + and float(min_p) == 0.0 + ) + if scalar_identity: + return None, None, None + + top_ks = _normalize_row_metadata( + top_k, + rows=rows, + device=device, + dtype=torch.int64, + name="logprob_top_ks", + ) + top_ps = _normalize_row_metadata( + top_p, + rows=rows, + device=device, + dtype=torch.float32, + name="logprob_top_ps", + ) + min_ps = _normalize_row_metadata( + min_p, + rows=rows, + device=device, + dtype=torch.float32, + name="logprob_min_ps", + ) + torch._assert_async((top_ks >= 1).all(), "logprob_top_ks must contain integers >= 1") + torch._assert_async( + (torch.isfinite(top_ps) & (top_ps > 0.0) & (top_ps <= 1.0)).all(), + "logprob_top_ps must contain finite values in (0, 1]", + ) + torch._assert_async( + (torch.isfinite(min_ps) & (min_ps >= 0.0) & (min_ps <= 1.0)).all(), + "logprob_min_ps must contain finite values in [0, 1]", + ) + if bool(((top_ks >= TOP_K_ALL) & (top_ps == 1.0) & (min_ps == 0.0)).all().item()): + return None, None, None + return top_ks, top_ps, min_ps + + +def exact_sampling_support( + logits: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: torch.Tensor, +) -> torch.Tensor: + """Return the joint top-k/top-p/min-p support for current logits. + + Temperature must already have been applied. Sorting is stable, so equal + probabilities are ordered by their original vocabulary index (token ID). + Top-p keeps the first token that crosses the threshold via + ``cumulative_probability_before <= top_p``. Min-p is relative to the + unfiltered row maximum. The three conditions are applied jointly. + """ + + if logits.ndim != 2 or not logits.is_floating_point(): + raise ValueError("exact sampling transforms require floating [rows, vocab] logits") + rows, vocab = logits.shape + if vocab < 1: + raise ValueError("exact sampling transforms require a non-empty vocabulary") + if tuple(top_ks.shape) != (rows,) or tuple(top_ps.shape) != (rows,) or tuple(min_ps.shape) != (rows,): + raise ValueError("exact sampling transform metadata must align one-to-one with logit rows") + + identity_rows = exact_sampling_identity_rows( + top_ks, + top_ps, + min_ps, + vocab_size=vocab, + ) + # The support is discrete metadata. Detaching is essential: current logits + # choose the support every forward, but gradients do not pass through sort + # indices or threshold comparisons. Identity rows are overwritten to full + # support after the fixed-shape program: a rounded FP32 cumulative sum may + # exceed one, but top-p=1 is mathematically unconditional. Keeping the + # fixed row shape also preserves CUDA-graph capture in serving. + probabilities = torch.softmax(logits.detach(), dim=-1) + sorted_probs, sorted_indices = torch.sort(probabilities, dim=-1, descending=True, stable=True) + ranks = torch.arange(vocab, device=logits.device, dtype=top_ks.dtype).unsqueeze(0) + keep_sorted = ranks < top_ks.clamp(max=vocab).unsqueeze(1) + cumulative_before = sorted_probs.cumsum(dim=-1) - sorted_probs + keep_sorted &= cumulative_before <= top_ps.unsqueeze(1) + keep_sorted &= sorted_probs >= sorted_probs[:, :1] * min_ps.unsqueeze(1) + + support = torch.zeros((rows, vocab), dtype=torch.bool, device=logits.device) + support.scatter_(1, sorted_indices, keep_sorted) + support |= identity_rows.unsqueeze(1) + return support + + +def exact_sampling_identity_rows( + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: torch.Tensor, + *, + vocab_size: int, +) -> torch.Tensor: + """Return rows whose transforms are the exact mathematical identity.""" + + if vocab_size < 1: + raise ValueError("vocab_size must be >= 1") + if top_ks.ndim != 1 or top_ps.shape != top_ks.shape or min_ps.shape != top_ks.shape: + raise ValueError("exact sampling transform metadata must have aligned one-dimensional rows") + return (top_ks >= vocab_size) & (top_ps == 1.0) & (min_ps == 0.0) + + +def exact_masked_logits( + logits: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + support = exact_sampling_support(logits, top_ks, top_ps, min_ps) + return logits.masked_fill(~support, -math.inf), support + + +def exact_selected_logprob_from_support( + logits: torch.Tensor, + token_ids: torch.Tensor, + support: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Score selected tokens under an already-derived exact support. + + A token outside current support has logprob ``-inf`` and zero gradient. + This is the required current-policy meaning when a historical action falls + out of top-k/top-p/min-p support after a weight update. + """ + + rows = logits.shape[0] + if token_ids.dtype is not torch.int64 or tuple(token_ids.shape) != (rows,): + raise ValueError("selected token IDs must be int64 and row-aligned") + masked_logits = logits.masked_fill(~support, -math.inf) + lse = torch.logsumexp(masked_logits, dim=-1) + selected = logits.gather(1, token_ids.unsqueeze(1)).squeeze(1) + selected_support = support.gather(1, token_ids.unsqueeze(1)).squeeze(1) + finite_logprob = torch.minimum(selected - lse, torch.zeros_like(selected)) + logprob = torch.where(selected_support, finite_logprob, torch.full_like(finite_logprob, -math.inf)) + return logprob, lse, selected_support + + +def exact_selected_logprob_partitioned_from_support( + logits: torch.Tensor, + token_ids: torch.Tensor, + support: torch.Tensor, + identity_rows: torch.Tensor, + native_selected_score: NativeSelectedScore, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Score identity rows natively and filtered rows on their exact support.""" + + rows = logits.shape[0] + if identity_rows.dtype is not torch.bool or identity_rows.shape != (rows,): + raise ValueError("identity_rows must be a row-aligned bool tensor") + + native_logprob, native_lse, _ = native_selected_score(logits, token_ids) + filtered_logprob, filtered_lse, filtered_selected_support = exact_selected_logprob_from_support( + logits, + token_ids, + support, + ) + return ( + torch.where(identity_rows, native_logprob, filtered_logprob), + torch.where(identity_rows, native_lse, filtered_lse), + torch.where(identity_rows, torch.ones_like(filtered_selected_support), filtered_selected_support), + ) + + +def exact_selected_logprob( + logits: torch.Tensor, + token_ids: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + support = exact_sampling_support(logits, top_ks, top_ps, min_ps) + logprob, lse, selected_support = exact_selected_logprob_from_support(logits, token_ids, support) + return logprob, lse, selected_support, support + + +def exact_selected_logprob_chunked( + logits: torch.Tensor, + token_ids: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: torch.Tensor, + *, + row_chunk_size: int = EXACT_FILTER_ROW_CHUNK, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Score rows without retaining a dense ``[tokens, vocab]`` support mask.""" + + if row_chunk_size < 1: + raise ValueError("row_chunk_size must be >= 1") + logprob_chunks = [] + lse_chunks = [] + selected_support_chunks = [] + for start in range(0, logits.shape[0], row_chunk_size): + end = min(start + row_chunk_size, logits.shape[0]) + logprob, lse, selected_support, _support = exact_selected_logprob( + logits[start:end], + token_ids[start:end], + top_ks[start:end], + top_ps[start:end], + min_ps[start:end], + ) + logprob_chunks.append(logprob) + lse_chunks.append(lse) + selected_support_chunks.append(selected_support) + if not logprob_chunks: + empty_float = logits.new_empty((0,)) + return empty_float, empty_float.clone(), torch.empty((0,), dtype=torch.bool, device=logits.device) + return ( + torch.cat(logprob_chunks), + torch.cat(lse_chunks), + torch.cat(selected_support_chunks), + ) + + +def exact_support_workspace_bytes(vocab_size: int, row_chunk_size: int = EXACT_FILTER_ROW_CHUNK) -> int: + """Upper bound for dense bool workspace in one support row chunk. + + At most the sorted keep mask and token-order support mask coexist. Float + probabilities and sort indices are separate value-program workspaces. + """ + + if vocab_size < 1 or row_chunk_size < 1: + raise ValueError("vocab_size and row_chunk_size must be >= 1") + return 2 * vocab_size * row_chunk_size + + +__all__ = [ + "EXACT_FILTER_ROW_CHUNK", + "EXACT_SAMPLING_TRANSFORM_PROGRAM", + "TOP_K_ALL", + "exact_masked_logits", + "exact_sampling_identity_rows", + "exact_sampling_support", + "exact_selected_logprob", + "exact_selected_logprob_chunked", + "exact_selected_logprob_from_support", + "exact_selected_logprob_partitioned_from_support", + "exact_support_workspace_bytes", + "normalize_exact_sampling_transforms", +] diff --git a/src/xorl/ops/exact_sampling_transforms.py b/src/xorl/ops/exact_sampling_transforms.py index 3904c644..0b95570a 100644 --- a/src/xorl/ops/exact_sampling_transforms.py +++ b/src/xorl/ops/exact_sampling_transforms.py @@ -1,313 +1,7 @@ -"""The backend-independent sampling-transform program for exact RL lanes. +"""Moved to ``xorl.ops.exact.sampling_transforms`` (issue #78); alias stub.""" -Qwen, GLM, and DSV4 first apply their declared temperature dtype/store -boundary, then call this module for one shared program: stable descending -probability order (token ID breaks ties), joint top-k plus inclusive-crossing -top-p plus min-p relative to the original row maximum, and normalization on -exactly that support. Serving performs per-row seeded Gumbel-max on the masked -logits. This is an exact-lane contract; it intentionally makes no claim about -generic SGLang or FlashInfer filter semantics. -""" +import importlib as _importlib +import sys as _sys -from __future__ import annotations -import math -from collections.abc import Callable - -import torch - - -TOP_K_ALL = 1 << 30 -EXACT_FILTER_ROW_CHUNK = 32 -EXACT_SAMPLING_TRANSFORM_PROGRAM = ( - "temperature_then_stable_token_id_topk_inclusive_topp_original_max_minp_seeded_gumbel_v1" -) -SamplingTransformRows = tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None] -NativeSelectedScore = Callable[ - [torch.Tensor, torch.Tensor], - tuple[torch.Tensor, torch.Tensor, torch.Tensor], -] - - -def _normalize_row_metadata( - value: int | float | torch.Tensor, - *, - rows: int, - device: torch.device, - dtype: torch.dtype, - name: str, -) -> torch.Tensor: - if isinstance(value, torch.Tensor): - if value.dtype is not dtype: - raise TypeError(f"per-row {name} must be {dtype}, got {value.dtype}") - if value.device != device: - raise ValueError(f"per-row {name} must share the logits device") - if tuple(value.shape) != (rows,): - raise ValueError(f"per-row {name} must have shape ({rows},), got {tuple(value.shape)}") - if not value.is_contiguous() or value.requires_grad: - raise ValueError(f"per-row {name} must be contiguous, non-differentiable sampling metadata") - return value - return torch.full((rows,), value, dtype=dtype, device=device) - - -def normalize_exact_sampling_transforms( - top_k: int | torch.Tensor = TOP_K_ALL, - top_p: float | torch.Tensor = 1.0, - min_p: float | torch.Tensor = 0.0, - *, - rows: int, - device: torch.device, -) -> SamplingTransformRows: - """Validate transforms and return row tensors, or three ``None`` values. - - Returning ``None`` for the all-identity scalar case is the explicit switch - that keeps established no-filter kernels and bytes untouched. - """ - - scalar_identity = ( - not isinstance(top_k, torch.Tensor) - and not isinstance(top_p, torch.Tensor) - and not isinstance(min_p, torch.Tensor) - and int(top_k) >= TOP_K_ALL - and float(top_p) == 1.0 - and float(min_p) == 0.0 - ) - if scalar_identity: - return None, None, None - - top_ks = _normalize_row_metadata( - top_k, - rows=rows, - device=device, - dtype=torch.int64, - name="logprob_top_ks", - ) - top_ps = _normalize_row_metadata( - top_p, - rows=rows, - device=device, - dtype=torch.float32, - name="logprob_top_ps", - ) - min_ps = _normalize_row_metadata( - min_p, - rows=rows, - device=device, - dtype=torch.float32, - name="logprob_min_ps", - ) - torch._assert_async((top_ks >= 1).all(), "logprob_top_ks must contain integers >= 1") - torch._assert_async( - (torch.isfinite(top_ps) & (top_ps > 0.0) & (top_ps <= 1.0)).all(), - "logprob_top_ps must contain finite values in (0, 1]", - ) - torch._assert_async( - (torch.isfinite(min_ps) & (min_ps >= 0.0) & (min_ps <= 1.0)).all(), - "logprob_min_ps must contain finite values in [0, 1]", - ) - if bool(((top_ks >= TOP_K_ALL) & (top_ps == 1.0) & (min_ps == 0.0)).all().item()): - return None, None, None - return top_ks, top_ps, min_ps - - -def exact_sampling_support( - logits: torch.Tensor, - top_ks: torch.Tensor, - top_ps: torch.Tensor, - min_ps: torch.Tensor, -) -> torch.Tensor: - """Return the joint top-k/top-p/min-p support for current logits. - - Temperature must already have been applied. Sorting is stable, so equal - probabilities are ordered by their original vocabulary index (token ID). - Top-p keeps the first token that crosses the threshold via - ``cumulative_probability_before <= top_p``. Min-p is relative to the - unfiltered row maximum. The three conditions are applied jointly. - """ - - if logits.ndim != 2 or not logits.is_floating_point(): - raise ValueError("exact sampling transforms require floating [rows, vocab] logits") - rows, vocab = logits.shape - if vocab < 1: - raise ValueError("exact sampling transforms require a non-empty vocabulary") - if tuple(top_ks.shape) != (rows,) or tuple(top_ps.shape) != (rows,) or tuple(min_ps.shape) != (rows,): - raise ValueError("exact sampling transform metadata must align one-to-one with logit rows") - - identity_rows = exact_sampling_identity_rows( - top_ks, - top_ps, - min_ps, - vocab_size=vocab, - ) - # The support is discrete metadata. Detaching is essential: current logits - # choose the support every forward, but gradients do not pass through sort - # indices or threshold comparisons. Identity rows are overwritten to full - # support after the fixed-shape program: a rounded FP32 cumulative sum may - # exceed one, but top-p=1 is mathematically unconditional. Keeping the - # fixed row shape also preserves CUDA-graph capture in serving. - probabilities = torch.softmax(logits.detach(), dim=-1) - sorted_probs, sorted_indices = torch.sort(probabilities, dim=-1, descending=True, stable=True) - ranks = torch.arange(vocab, device=logits.device, dtype=top_ks.dtype).unsqueeze(0) - keep_sorted = ranks < top_ks.clamp(max=vocab).unsqueeze(1) - cumulative_before = sorted_probs.cumsum(dim=-1) - sorted_probs - keep_sorted &= cumulative_before <= top_ps.unsqueeze(1) - keep_sorted &= sorted_probs >= sorted_probs[:, :1] * min_ps.unsqueeze(1) - - support = torch.zeros((rows, vocab), dtype=torch.bool, device=logits.device) - support.scatter_(1, sorted_indices, keep_sorted) - support |= identity_rows.unsqueeze(1) - return support - - -def exact_sampling_identity_rows( - top_ks: torch.Tensor, - top_ps: torch.Tensor, - min_ps: torch.Tensor, - *, - vocab_size: int, -) -> torch.Tensor: - """Return rows whose transforms are the exact mathematical identity.""" - - if vocab_size < 1: - raise ValueError("vocab_size must be >= 1") - if top_ks.ndim != 1 or top_ps.shape != top_ks.shape or min_ps.shape != top_ks.shape: - raise ValueError("exact sampling transform metadata must have aligned one-dimensional rows") - return (top_ks >= vocab_size) & (top_ps == 1.0) & (min_ps == 0.0) - - -def exact_masked_logits( - logits: torch.Tensor, - top_ks: torch.Tensor, - top_ps: torch.Tensor, - min_ps: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - support = exact_sampling_support(logits, top_ks, top_ps, min_ps) - return logits.masked_fill(~support, -math.inf), support - - -def exact_selected_logprob_from_support( - logits: torch.Tensor, - token_ids: torch.Tensor, - support: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Score selected tokens under an already-derived exact support. - - A token outside current support has logprob ``-inf`` and zero gradient. - This is the required current-policy meaning when a historical action falls - out of top-k/top-p/min-p support after a weight update. - """ - - rows = logits.shape[0] - if token_ids.dtype is not torch.int64 or tuple(token_ids.shape) != (rows,): - raise ValueError("selected token IDs must be int64 and row-aligned") - masked_logits = logits.masked_fill(~support, -math.inf) - lse = torch.logsumexp(masked_logits, dim=-1) - selected = logits.gather(1, token_ids.unsqueeze(1)).squeeze(1) - selected_support = support.gather(1, token_ids.unsqueeze(1)).squeeze(1) - finite_logprob = torch.minimum(selected - lse, torch.zeros_like(selected)) - logprob = torch.where(selected_support, finite_logprob, torch.full_like(finite_logprob, -math.inf)) - return logprob, lse, selected_support - - -def exact_selected_logprob_partitioned_from_support( - logits: torch.Tensor, - token_ids: torch.Tensor, - support: torch.Tensor, - identity_rows: torch.Tensor, - native_selected_score: NativeSelectedScore, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Score identity rows natively and filtered rows on their exact support.""" - - rows = logits.shape[0] - if identity_rows.dtype is not torch.bool or identity_rows.shape != (rows,): - raise ValueError("identity_rows must be a row-aligned bool tensor") - - native_logprob, native_lse, _ = native_selected_score(logits, token_ids) - filtered_logprob, filtered_lse, filtered_selected_support = exact_selected_logprob_from_support( - logits, - token_ids, - support, - ) - return ( - torch.where(identity_rows, native_logprob, filtered_logprob), - torch.where(identity_rows, native_lse, filtered_lse), - torch.where(identity_rows, torch.ones_like(filtered_selected_support), filtered_selected_support), - ) - - -def exact_selected_logprob( - logits: torch.Tensor, - token_ids: torch.Tensor, - top_ks: torch.Tensor, - top_ps: torch.Tensor, - min_ps: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - support = exact_sampling_support(logits, top_ks, top_ps, min_ps) - logprob, lse, selected_support = exact_selected_logprob_from_support(logits, token_ids, support) - return logprob, lse, selected_support, support - - -def exact_selected_logprob_chunked( - logits: torch.Tensor, - token_ids: torch.Tensor, - top_ks: torch.Tensor, - top_ps: torch.Tensor, - min_ps: torch.Tensor, - *, - row_chunk_size: int = EXACT_FILTER_ROW_CHUNK, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Score rows without retaining a dense ``[tokens, vocab]`` support mask.""" - - if row_chunk_size < 1: - raise ValueError("row_chunk_size must be >= 1") - logprob_chunks = [] - lse_chunks = [] - selected_support_chunks = [] - for start in range(0, logits.shape[0], row_chunk_size): - end = min(start + row_chunk_size, logits.shape[0]) - logprob, lse, selected_support, _support = exact_selected_logprob( - logits[start:end], - token_ids[start:end], - top_ks[start:end], - top_ps[start:end], - min_ps[start:end], - ) - logprob_chunks.append(logprob) - lse_chunks.append(lse) - selected_support_chunks.append(selected_support) - if not logprob_chunks: - empty_float = logits.new_empty((0,)) - return empty_float, empty_float.clone(), torch.empty((0,), dtype=torch.bool, device=logits.device) - return ( - torch.cat(logprob_chunks), - torch.cat(lse_chunks), - torch.cat(selected_support_chunks), - ) - - -def exact_support_workspace_bytes(vocab_size: int, row_chunk_size: int = EXACT_FILTER_ROW_CHUNK) -> int: - """Upper bound for dense bool workspace in one support row chunk. - - At most the sorted keep mask and token-order support mask coexist. Float - probabilities and sort indices are separate value-program workspaces. - """ - - if vocab_size < 1 or row_chunk_size < 1: - raise ValueError("vocab_size and row_chunk_size must be >= 1") - return 2 * vocab_size * row_chunk_size - - -__all__ = [ - "EXACT_FILTER_ROW_CHUNK", - "EXACT_SAMPLING_TRANSFORM_PROGRAM", - "TOP_K_ALL", - "exact_masked_logits", - "exact_sampling_identity_rows", - "exact_sampling_support", - "exact_selected_logprob", - "exact_selected_logprob_chunked", - "exact_selected_logprob_from_support", - "exact_selected_logprob_partitioned_from_support", - "exact_support_workspace_bytes", - "normalize_exact_sampling_transforms", -] +_sys.modules[__name__] = _importlib.import_module("xorl.ops.exact.sampling_transforms") diff --git a/src/xorl/ops/loss/bi_fused_lm_head.py b/src/xorl/ops/loss/bi_fused_lm_head.py index 35249b21..292eff69 100644 --- a/src/xorl/ops/loss/bi_fused_lm_head.py +++ b/src/xorl/ops/loss/bi_fused_lm_head.py @@ -1,7 +1,7 @@ """Trainable wrapper for the batch-invariant fused LM-head logprob contract. Forward scores per-token cross-entropy through -:func:`xorl.ops.batch_invariant_ops.bi_lm_head_selected_logprob` — the K3 +:func:`xorl.ops.sglang.batch_invariant_ops.bi_lm_head_selected_logprob` — the K3 lm-head contract vendored identically in SGLang, so trainer and serving logprobs are bitwise identical from bit-exact hidden states. The bf16 weight stays resident (no fp32 lm-head copy). Per-row temperature materializes the @@ -21,26 +21,26 @@ import torch import torch.distributed as dist -from xorl.ops.batch_invariant_ops import ( +from xorl.ops.exact.sampling_transforms import ( + EXACT_FILTER_ROW_CHUNK, + TOP_K_ALL, + exact_sampling_identity_rows, + exact_sampling_support, + exact_selected_logprob_partitioned_from_support, +) +from xorl.ops.sglang.batch_invariant_ops import ( BI_LM_HEAD_VOCAB_CHUNK, bi_lm_head_full_logits, bi_lm_head_selected_logprob, bi_lm_head_selected_logprob_from_logits, ) -from xorl.ops.bi_families_v2 import ( +from xorl.ops.sglang.bi_families_v2 import ( exact_temperature_scale_fp32_logits, families_v2_enabled, head_v2_full_logits_with_lse, head_v2_selected_logprob, head_v2_selected_logprob_from_logits, ) -from xorl.ops.exact_sampling_transforms import ( - EXACT_FILTER_ROW_CHUNK, - TOP_K_ALL, - exact_sampling_identity_rows, - exact_sampling_support, - exact_selected_logprob_partitioned_from_support, -) _TEMPERATURE_MATERIALIZE_ROW_CHUNK = EXACT_FILTER_ROW_CHUNK diff --git a/src/xorl/ops/loss/per_token_ce.py b/src/xorl/ops/loss/per_token_ce.py index a14183bb..84d4af4b 100644 --- a/src/xorl/ops/loss/per_token_ce.py +++ b/src/xorl/ops/loss/per_token_ce.py @@ -10,7 +10,7 @@ import torch.nn.functional as F from xorl.distributed.parallel_state import get_parallel_state -from xorl.ops.exact_sampling_transforms import ( +from xorl.ops.exact.sampling_transforms import ( TOP_K_ALL, normalize_exact_sampling_transforms, ) diff --git a/src/xorl/ops/sglang/__init__.py b/src/xorl/ops/sglang/__init__.py new file mode 100644 index 00000000..ae06412b --- /dev/null +++ b/src/xorl/ops/sglang/__init__.py @@ -0,0 +1,17 @@ +"""Serving-parity twins vendored from / mirrored into SGLang (issue #78). + +These modules exist in both engines so trainer and serving arithmetic share +one implementation: + +- ``bi_families_v2`` — vendored byte-identical into the serving engine and + sha256-gated; it keeps that engine's formatting (black, 88 columns) and is + excluded from all rewriting hooks. Any edit requires the paired + serving-side edit. +- ``batch_invariant_ops`` — vendored-adapted from SGLang's + ``srt/batch_invariant_ops`` (SGLang-internal helpers stubbed, DeepGEMM + routing added). Kept as a single file so it stays diffable against the + serving twin; edits must consider the serving side. + +Unlike ``ops/_vendored/`` (untouchable third-party snapshots), these are +first-party-maintained with a paired-edit discipline. +""" diff --git a/src/xorl/ops/sglang/batch_invariant_ops.py b/src/xorl/ops/sglang/batch_invariant_ops.py new file mode 100644 index 00000000..79f4d56b --- /dev/null +++ b/src/xorl/ops/sglang/batch_invariant_ops.py @@ -0,0 +1,2188 @@ +# Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/batch_invariant_ops/batch_invariant_ops.py +# +# Vendored into xorl from SGLang's +# python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py so that the +# xorl trainer/server forward can use the SAME batch-invariant Triton matmul +# kernels SGLang uses. This makes the linear layers, gate matmul and lm_head +# bit-for-bit match SGLang's reduction order, collapsing the cross-engine K3 +# logprob tail. +# +# SGLang-internal helpers are stubbed/inlined so this module is self-contained. +# DeepGEMM is discovered lazily and used only on a supported Hopper bf16 route; +# otherwise the vendored Triton path remains the fallback. Environment parsing +# and calc_diff are inlined below. + +import contextlib +import os +from collections import namedtuple +from collections.abc import Callable +from typing import Any, Dict, Literal, Optional, Tuple + +import torch +import triton +import triton.language as tl +from triton.runtime.errors import OutOfResources + +from xorl.ops.exact.bi_gemm_configs import baseline_mm_config, lookup_mm_config + + +# --- Stubs for SGLang-internal imports --------------------------------------- +# DeepGEMM's bf16 NN GEMM is bitwise-identical to the batch-invariant Triton +# persistent kernel (verified per-shape by tune_bi_gemm.py and gated by +# tests/ops/test_bi_gemm_config_table.py). Serving's mm route already relies on +# that equality — SGLang sends every contiguous bf16 N>=16 mm to DeepGEMM while +# the trainer ran pure Triton — so routing the trainer the same way cannot move +# cross-engine bits, and it is markedly faster at trainer token counts. +# Requires Hopper+ and an importable deep_gemm. Detection is lazy (first mm), +# so importing this module never initializes CUDA. +_DEEPGEMM_READY: bool | None = None + + +def _deepgemm_ready() -> bool: + global _DEEPGEMM_READY + if _DEEPGEMM_READY is None: + try: + import deep_gemm # noqa: F401, PLC0415 + + _DEEPGEMM_READY = ( + torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] >= 9 + and hasattr(deep_gemm, "bf16_gemm_nn") + ) + except Exception: # noqa: BLE001 + _DEEPGEMM_READY = False + return _DEEPGEMM_READY + + +ENABLE_JIT_DEEPGEMM = True # gated per-call by _deepgemm_ready() + +# Shapes whose table config exceeded shared memory at launch (Triton's epilogue +# staging for wide-output tiles is version-dependent): remembered so the hot +# path re-launches straight on the pinned baseline without re-raising. +_MM_CONFIG_OOM_SHAPES: set[tuple] = set() + + +def _launch_with_config_fallback(launch, dtype, M, N, K, out_itemsize=None): + key = (str(dtype), M, N, K, out_itemsize) + if key in _MM_CONFIG_OOM_SHAPES: + launch(baseline_mm_config(dtype)) + return + try: + launch(lookup_mm_config(dtype, M, N, K, out_itemsize=out_itemsize)) + except OutOfResources: + _MM_CONFIG_OOM_SHAPES.add(key) + launch(baseline_mm_config(dtype)) + + +# ----------------------------------------------------------------------------- + + +# Exact-model activation always uses the admitted production route. Ambient +# process variables cannot substitute comparison or order-variant fallbacks. +_ENABLE_MM_DEEPGEMM = True + +__all__ = [ + "set_batch_invariant_mode", + "is_batch_invariant_mode_enabled", + "is_batch_invariant_op_enabled", + "get_batch_invariant_ops", + "disable_batch_invariant_mode", + "enable_batch_invariant_mode", + "rms_norm_batch_invariant", + "fused_add_rms_norm_batch_invariant", + "sglang_rms_norm_batch_invariant", + "fused_rms_norm_backward", + "wrap_trunk_linears_batch_invariant", + "is_trunk_linear_contract_enabled", + "batch_invariant_trunk_linear", + "bi_bf16_fp32_linear", + "set_trunk_linear_contract", + "RMSNormFamily", + "RMS_NORM_FAMILY_NO_RESIDUAL", + "RMS_NORM_FAMILY_RESIDUAL_TREE", + "RMS_NORM_FAMILIES", + "bi_rms_norm", + "bi_fused_add_rms_norm", +] + + +_BATCH_INVARIANT_ALL_OPS = { + "mm", + "addmm", + "mm_dtype", + "log_softmax", + "mean", + "rms_norm", + "bmm", +} +_BATCH_INVARIANT_ALIASES = { + "matmul": "mm", + "logsoftmax": "log_softmax", + "log-softmax": "log_softmax", + "rmsnorm": "rms_norm", + "rms-norm": "rms_norm", +} + + +def _parse_batch_invariant_ops() -> set[str]: + raw = os.environ.get("XORL_BATCH_INVARIANT_OPS", "all").strip().lower() + if raw in ("", "1", "true", "yes", "all"): + return set(_BATCH_INVARIANT_ALL_OPS) + if raw in ("0", "false", "no", "none"): + return set() + + ops = set() + for part in raw.replace(";", ",").split(","): + op = part.strip().lower().replace("-", "_") + if not op: + continue + op = _BATCH_INVARIANT_ALIASES.get(op, op) + if op == "mm": + ops.add("mm_dtype") + if op not in _BATCH_INVARIANT_ALL_OPS: + raise ValueError( + f"Unsupported XORL_BATCH_INVARIANT_OPS entry {part!r}; " + f"supported values are: {sorted(_BATCH_INVARIANT_ALL_OPS)}" + ) + ops.add(op) + return ops + + +def _matmul_launch_metadata(grid: Callable[..., Any], kernel: Any, args: Dict[str, Any]) -> Dict[str, Any]: + ret = {} + m, n, k = args["M"], args["N"], args["K"] + ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}]" + if "tiles_per_update" in args: + ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}, tiles_per_update={args['tiles_per_update']:02}]" + if "c_ptr" in args: + bytes_per_elem = args["c_ptr"].element_size() + else: + bytes_per_elem = 1 if args["FP8_OUTPUT"] else 2 + ret[f"flops{bytes_per_elem * 8}"] = 2.0 * m * n * k + ret["bytes"] = bytes_per_elem * (m * k + n * k + m * n) + return ret + + +@triton.jit +def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS): + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + return pid_m, pid_n + + +@triton.jit(launch_metadata=_matmul_launch_metadata) +def matmul_kernel_persistent( + a_ptr, + b_ptr, + c_ptr, # + bias_ptr, + M, + N, + K, # + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, # + BLOCK_SIZE_N: tl.constexpr, # + BLOCK_SIZE_K: tl.constexpr, # + GROUP_SIZE_M: tl.constexpr, # + NUM_SMS: tl.constexpr, # + A_LARGE: tl.constexpr, + B_LARGE: tl.constexpr, + C_LARGE: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + start_pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + num_tiles = num_pid_m * num_pid_n + + offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): + pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS) + start_m = pid_m * BLOCK_SIZE_M + start_n = pid_n * BLOCK_SIZE_N + offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) + offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) + if A_LARGE: + offs_am = offs_am.to(tl.int64) + if B_LARGE: + offs_bn = offs_bn.to(tl.int64) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for ki in range(k_tiles): + if A_LARGE or B_LARGE: + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + else: + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + a = tl.load(a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0) + accumulator = tl.dot(a, b, accumulator) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + if C_LARGE: + offs_cm = offs_cm.to(tl.int64) + offs_cn = offs_cn.to(tl.int64) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + if HAS_BIAS: + bias_ptrs = bias_ptr + offs_cn + bias = tl.load(bias_ptrs, mask=offs_cn < N, other=0.0).to(tl.float32) + accumulator += bias + if c_ptr.dtype.element_ty == tl.float8e4nv: + c = accumulator.to(tl.float8e4nv) + elif c_ptr.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif c_ptr.dtype.element_ty == tl.float32: + c = accumulator.to(tl.float32) + else: + c = accumulator.to(tl.float16) + tl.store(c_ptrs, c, mask=c_mask) + + +def _matmul_persistent_triton(a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None): + # Check constraints. + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.dtype == b.dtype, "Incompatible dtypes" + assert bias is None or bias.dim() == 1, "Currently assuming bias is 1D, let Horace know if you run into this" + NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count + M, K = a.shape + K, N = b.shape + dtype = a.dtype + # Allocates output. + c = torch.empty((M, N), device=a.device, dtype=dtype) + + # 1D launch kernel where each block gets its own program. + def grid(META): + return ( + min( + NUM_SMS, + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ), + ) + + # Shape-keyed on the bit-neutral axes only; BLOCK_SIZE_K stays pinned per + # dtype (bi_gemm_configs — the R1 config table, identical in both engines). + def _launch(config): + matmul_kernel_persistent[grid]( + a, + b, + c, # + bias, + M, + N, + K, # + a.stride(0), + a.stride(1), # + b.stride(0), + b.stride(1), # + c.stride(0), + c.stride(1), # + NUM_SMS=NUM_SMS, # + A_LARGE=a.numel() > 2**31, + B_LARGE=b.numel() > 2**31, + C_LARGE=c.numel() > 2**31, + HAS_BIAS=bias is not None, + **config, + ) + + _launch_with_config_fallback(_launch, dtype, M, N, K) + return c + + +def _matmul_persistent_deepgemm(a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None): + M, K = a.shape + K, N = b.shape + dtype = a.dtype + out = torch.empty((M, N), device=a.device, dtype=dtype) + + try: + import deep_gemm # noqa: PLC0415 + + deep_gemm.bf16_gemm_nn(a, b, out) + except RuntimeError: + return None + + # TODO can this be put in DeepGEMM's `c`? + if bias is not None: + out += bias + + return out + + +def matmul_persistent(a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None): + K, N = b.shape + + # DeepGEMM has minimum dimension requirements for TMA descriptors + MIN_DEEPGEMM_DIM = 16 + + if ( + _ENABLE_MM_DEEPGEMM + and (a.dtype == torch.bfloat16) + and (b.dtype == torch.bfloat16) + and a.is_contiguous() + and b.transpose(0, 1).is_contiguous() + and N >= MIN_DEEPGEMM_DIM + and _deepgemm_ready() + ): + result = _matmul_persistent_deepgemm(a=a, b=b, bias=bias) + if result is not None: + return result + # DeepGEMM failed (e.g. dimensions too small for TMA descriptors), + # fall through to batch-invariant Triton persistent kernel + + return _matmul_persistent_triton(a=a, b=b, bias=bias) + + +@triton.jit +def _log_softmax_kernel( + input_ptr, + output_ptr, + input_row_stride: tl.constexpr, + output_row_stride: tl.constexpr, + n_cols: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """ + Compute log_softmax along the last dimension of a 2D tensor. + Each block handles one row of the input tensor. + """ + # Get the row index for this block + row_idx = tl.program_id(0).to(tl.int64) + + # Compute base pointers for input and output rows + row_start_ptr = input_ptr + row_idx * input_row_stride + output_row_start_ptr = output_ptr + row_idx * output_row_stride + + # Step 1: Find maximum value in the row for numerical stability + # Load first block to infer dtype and initialize max_val with correct type + col_idx_init = tl.arange(0, BLOCK_SIZE) + mask_init = col_idx_init < n_cols + vals_init = tl.load(row_start_ptr + col_idx_init, mask=mask_init, other=-float("inf")) + max_val = tl.max(vals_init) + + # Continue with remaining blocks + for col_offset in range(BLOCK_SIZE, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + + # Load values + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=-float("inf")) + + # Update maximum + max_val = tl.max(tl.maximum(vals, max_val)) + + # Step 2: Compute sum of exp(x - max_val) + # Initialize sum_exp with correct dtype by using tl.sum on a zero vector + sum_exp = tl.sum(tl.zeros([1], dtype=max_val.dtype)) + + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + + # Load values + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + + # Compute exp(x - max_val) and accumulate + exp_vals = tl.exp(vals - max_val) + sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0)) + + # Compute log(sum_exp) + log_sum_exp = tl.log(sum_exp) + + # Step 3: Compute final log_softmax values: x - max_val - log_sum_exp + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + + # Load values + vals = tl.load(row_start_ptr + col_idx, mask=mask) + + # Compute log_softmax + output = vals - max_val - log_sum_exp + + # Store results + tl.store(output_row_start_ptr + col_idx, output, mask=mask) + + +def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: + """ + Compute log_softmax using Triton kernel. + + Args: + input: Input tensor + dim: Dimension along which to compute log_softmax (only -1 or last dim supported) + + Returns: + Tensor with log_softmax applied along the specified dimension + """ + if dim != -1 and dim != input.ndim - 1: + raise ValueError("This implementation only supports log_softmax along the last dimension") + + # Flatten all dimensions except the last one + original_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + input_2d = input_2d.contiguous() + + n_rows, n_cols = input_2d.shape + + # Allocate output tensor + output = torch.empty_like(input_2d) + + # Choose block size based on the number of columns + BLOCK_SIZE = 1024 + + # Launch kernel with one block per row + grid = (n_rows,) + _log_softmax_kernel[grid]( + input_2d, + output, + input_2d.stride(0), + output.stride(0), + n_cols, + BLOCK_SIZE=BLOCK_SIZE, + ) + # Reshape output back to original shape + return output.reshape(original_shape) + + +@triton.jit +def mean_kernel( + input_ptr, + output_ptr, + input_stride0, + input_stride1, + input_stride2, + output_stride0, + output_stride1, + M, # size before reduction dim + N, # size of reduction dim + K, # size after reduction dim + BLOCK_SIZE: tl.constexpr, +): + """ + Kernel for computing mean along a single dimension. + Input is viewed as (M, N, K) where N is the dimension being reduced. + """ + # Program ID gives us which output element we're computing + pid = tl.program_id(0) + + # Compute output indices + m_idx = pid // K + k_idx = pid % K + + # Bounds check + if m_idx >= M or k_idx >= K: + return + + # Accumulate sum across reduction dimension + acc = 0.0 + for n_start in range(0, N, BLOCK_SIZE): + n_offsets = n_start + tl.arange(0, BLOCK_SIZE) + mask = n_offsets < N + + # Calculate input indices + input_idx = m_idx * input_stride0 + n_offsets * input_stride1 + k_idx * input_stride2 + + # Load and accumulate + vals = tl.load(input_ptr + input_idx, mask=mask, other=0.0) + acc += tl.sum(vals) + + # Compute mean and store + mean_val = acc / N + output_idx = m_idx * output_stride0 + k_idx * output_stride1 + tl.store(output_ptr + output_idx, mean_val) + + +def mean_dim( + input: torch.Tensor, + dim: int, + keepdim: bool = False, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """ + Triton implementation of torch.mean with single dimension reduction. + + Args: + input: Input tensor + dim: Single dimension along which to compute mean + keepdim: Whether to keep the reduced dimension + dtype: Output dtype. If None, uses input dtype (or float32 for integer inputs) + + Returns: + Tensor with mean values along specified dimension + """ + # Validate inputs + assert input.is_cuda, "Input must be a CUDA tensor" + assert -input.ndim <= dim < input.ndim, f"Invalid dimension {dim} for tensor with {input.ndim} dimensions" + + # Handle negative dim + if dim < 0: + dim = dim + input.ndim + + # Handle dtype + if dtype is None: + if input.dtype in [torch.int8, torch.int16, torch.int32, torch.int64]: + dtype = torch.float32 + else: + dtype = input.dtype + + # Convert input to appropriate dtype if needed + if input.dtype != dtype: + input = input.to(dtype) + + # Get input shape and strides + shape = list(input.shape) + + # Calculate dimensions for kernel + M = 1 + for i in range(dim): + M *= shape[i] + + N = shape[dim] + + K = 1 + for i in range(dim + 1, len(shape)): + K *= shape[i] + + # Reshape input to 3D view (M, N, K) + input_3d = input.reshape(M, N, K) + + # Create output shape + if keepdim: + output_shape = shape.copy() + output_shape[dim] = 1 + else: + output_shape = shape[:dim] + shape[dim + 1 :] + + # Create output tensor + output = torch.empty(output_shape, dtype=dtype, device=input.device) + + # Reshape output for kernel + if keepdim: + output_2d = output.reshape(M, 1, K).squeeze(1) + else: + output_2d = output.reshape(M, K) + + # Launch kernel + grid = (M * K,) + BLOCK_SIZE = 1024 + + mean_kernel[grid]( + input_3d, + output_2d, + input_3d.stride(0), + input_3d.stride(1), + input_3d.stride(2), + output_2d.stride(0), + output_2d.stride(1) if output_2d.ndim > 1 else 0, + M, + N, + K, + BLOCK_SIZE, + ) + + return output + + +_INTERPOSE_GRAD_ERROR = ( + "Batch-invariant interposed op '{op}' received a grad-requiring input while grad is enabled. " + "The global interpose (XORL_BATCH_INVARIANT_MATMUL / enable_batch_invariant_mode) is " + "inference/verification-only: the aten::rms_norm override records no autograd graph (q/k-norm " + "gradients silently vanish) and the torch.bmm monkeypatch detaches the graph. For a training " + "forward on the batch-invariant contract use the module-scoped XORL_BI_TRUNK_LINEAR=1 lane " + "instead." +) + + +def _guard_interpose_no_grad(op: str, *tensors) -> None: + """Loud-fail: the global interpose must never see a training forward.""" + if torch.is_grad_enabled() and any(isinstance(t, torch.Tensor) and t.requires_grad for t in tensors): + raise RuntimeError(_INTERPOSE_GRAD_ERROR.format(op=op)) + + +def mm_batch_invariant(a, b): + _guard_interpose_no_grad("aten::mm", a, b) + return matmul_persistent(a, b) + + +def addmm_batch_invariant(bias, a, b): + _guard_interpose_no_grad("aten::addmm", bias, a, b) + return matmul_persistent(a, b, bias=bias) + + +def _log_softmax_batch_invariant(input, dim, _half_to_float): + assert not _half_to_float, "not implemented" + _guard_interpose_no_grad("aten::_log_softmax", input) + return log_softmax(input, dim=dim) + + +def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = None): + _guard_interpose_no_grad("aten::mean.dim", input) + assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}" + if dim is None or len(dim) == 0: + # aten::mean full-reduce dispatches here with dim=[]; the empty n_elems product returned the SUM + dim = list(range(input.ndim)) + if len(dim) == 1: + return mean_dim(input, dim[0], keepdim=keepdim) + else: + assert input.dtype in { + torch.float16, + torch.bfloat16, + torch.float32, + }, "only float types supported for now" + n_elems = 1 + for d in dim: + n_elems *= input.shape[d] + return torch.sum(input, dim=dim, keepdim=keepdim, dtype=torch.float32) / n_elems + + +@triton.jit +def bmm_kernel_persistent( + a_ptr, + b_ptr, + c_ptr, # + B, + M, + N, + K, # + stride_ab, + stride_am, + stride_ak, + stride_bb, + stride_bk, + stride_bn, + stride_cb, + stride_cm, + stride_cn, + BLOCK_SIZE_M: tl.constexpr, # + BLOCK_SIZE_N: tl.constexpr, # + BLOCK_SIZE_K: tl.constexpr, # + GROUP_SIZE_M: tl.constexpr, # + NUM_SMS: tl.constexpr, # + A_LARGE: tl.constexpr, + B_LARGE: tl.constexpr, + C_LARGE: tl.constexpr, +): + """ + Batched matrix multiplication kernel that processes batches in parallel. + Each tile processes a (BLOCK_SIZE_M, BLOCK_SIZE_N) output block for a specific batch. + """ + start_pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + num_tiles_per_batch = num_pid_m * num_pid_n + num_tiles_total = B * num_tiles_per_batch + + offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + # Process tiles in a deterministic order: batch-major ordering + for tile_id in tl.range(start_pid, num_tiles_total, NUM_SMS, flatten=True): + # Decompose tile_id into batch and within-batch tile + batch_idx = tile_id // num_tiles_per_batch + tile_in_batch = tile_id % num_tiles_per_batch + + pid_m, pid_n = _compute_pid(tile_in_batch, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS) + start_m = pid_m * BLOCK_SIZE_M + start_n = pid_n * BLOCK_SIZE_N + offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) + offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) + if A_LARGE: + offs_am = offs_am.to(tl.int64) + if B_LARGE: + offs_bn = offs_bn.to(tl.int64) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + # Add batch offset + if A_LARGE or B_LARGE: + batch_idx_typed = batch_idx.to(tl.int64) + else: + batch_idx_typed = batch_idx + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for ki in range(k_tiles): + if A_LARGE or B_LARGE: + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + else: + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + + a_ptrs = a_ptr + (batch_idx_typed * stride_ab + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (batch_idx_typed * stride_bb + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + a = tl.load(a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0) + accumulator = tl.dot(a, b, accumulator) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + if C_LARGE: + offs_cm = offs_cm.to(tl.int64) + offs_cn = offs_cn.to(tl.int64) + c_ptrs = c_ptr + batch_idx_typed * stride_cb + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + + if c_ptr.dtype.element_ty == tl.float8e4nv: + c = accumulator.to(tl.float8e4nv) + elif c_ptr.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif c_ptr.dtype.element_ty == tl.float32: + c = accumulator.to(tl.float32) + else: + c = accumulator.to(tl.float16) + tl.store(c_ptrs, c, mask=c_mask) + + +def bmm_batch_invariant(a, b, *, out=None): + _guard_interpose_no_grad("aten::bmm/torch.bmm", a, b) + # Batched matrix multiply: (B, M, K) x (B, K, N) -> (B, M, N) + # Process batches in parallel with our persistent kernel + if a.ndim == 3 and b.ndim == 3: + # Check constraints + assert a.shape[0] == b.shape[0], "Batch sizes must match" + assert a.shape[2] == b.shape[1], "Incompatible dimensions" + assert a.dtype == b.dtype, "Incompatible dtypes" + + B = a.shape[0] + M = a.shape[1] + K = a.shape[2] + N = b.shape[2] + dtype = a.dtype + + # Allocate output + if out is None: + c = torch.empty((B, M, N), device=a.device, dtype=dtype) + else: + c = out + + NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count + + # Use fixed kernel configuration for determinism + configs = { + torch.bfloat16: { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_stages": 3, + "num_warps": 8, + }, + torch.float16: { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_stages": 3, + "num_warps": 8, + }, + torch.float32: { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_stages": 3, + "num_warps": 8, + }, + } + + config = configs.get(dtype) + if config is None: + raise ValueError( + f"Unsupported dtype {dtype} for bmm_batch_invariant. Supported dtypes are: {list(configs.keys())}" + ) + + # Grid: limit by NUM_SMS for persistent kernel approach + num_tiles_per_batch = triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]) + num_tiles_total = B * num_tiles_per_batch + grid = (min(NUM_SMS, num_tiles_total),) + + bmm_kernel_persistent[grid]( + a, + b, + c, # + B, + M, + N, + K, # + a.stride(0), + a.stride(1), + a.stride(2), # + b.stride(0), + b.stride(1), + b.stride(2), # + c.stride(0), + c.stride(1), + c.stride(2), # + NUM_SMS=NUM_SMS, # + A_LARGE=a.numel() > 2**31, + B_LARGE=b.numel() > 2**31, + C_LARGE=c.numel() > 2**31, + **config, + ) + + return c + else: + raise ValueError(f"bmm_batch_invariant expects 3D tensors, got shapes {a.shape} and {b.shape}") + + +@triton.jit +def _rms_norm_kernel( + input_ptr, + weight_ptr, + output_ptr, + input_row_stride: tl.constexpr, + output_row_stride: tl.constexpr, + n_cols: tl.constexpr, + eps, + BLOCK_SIZE: tl.constexpr, +): + """ + Compute RMS normalization along the last dimension of a 2D tensor. + RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight + Each block handles one row of the input tensor. + """ + row_idx = tl.program_id(0).to(tl.int64) + row_start_ptr = input_ptr + row_idx * input_row_stride + output_row_start_ptr = output_ptr + row_idx * output_row_stride + + # Step 1: Compute sum of squares in float32 to avoid overflow + sum_sq = tl.zeros([1], dtype=tl.float32) + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + # Convert to float32 for accumulation to prevent overflow + vals_f32 = vals.to(tl.float32) + sq_vals = vals_f32 * vals_f32 + sum_sq += tl.sum(tl.where(mask, sq_vals, 0.0)) + + # Step 2: Compute RMS (root mean square) in float32 + mean_sq = sum_sq / n_cols + rms = tl.sqrt(mean_sq + eps) + inv_rms = 1.0 / rms + + # Step 3: Normalize and apply weight + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) + # Compute in float32 then convert back to input dtype + vals_f32 = vals.to(tl.float32) + weight_f32 = weight.to(tl.float32) + output_f32 = vals_f32 * inv_rms * weight_f32 + output = output_f32.to(vals.dtype) + tl.store(output_row_start_ptr + col_idx, output, mask=mask) + + +def rms_norm(input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + """ + Compute RMS normalization using Triton kernel. + + RMS Norm normalizes the input by the root mean square and scales by weight: + output = input / sqrt(mean(input^2) + eps) * weight + + Args: + input: Input tensor of shape (..., hidden_size) + weight: Weight tensor of shape (hidden_size,) + eps: Small constant for numerical stability + + Returns: + Tensor with RMS normalization applied along the last dimension + """ + assert weight.dim() == 1, "Weight must be 1-dimensional" + assert input.shape[-1] == weight.shape[0], ( + f"Input last dimension ({input.shape[-1]}) must match weight dimension ({weight.shape[0]})" + ) + + # Flatten all dimensions except the last one + original_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + input_2d = input_2d.contiguous() + weight = weight.contiguous() + + n_rows, n_cols = input_2d.shape + + output = torch.empty_like(input_2d) + BLOCK_SIZE = 1024 + grid = (n_rows,) + _rms_norm_kernel[grid]( + input_2d, + weight, + output, + input_2d.stride(0), + output.stride(0), + n_cols, + eps, + BLOCK_SIZE=BLOCK_SIZE, + ) + return output.reshape(original_shape) + + +def rms_norm_batch_invariant(input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + """ + Batch-invariant wrapper for RMS normalization. + + This function provides a deterministic, batch-invariant implementation + of RMS normalization for use with the batch_invariant mode. + + Adapted from https://github.com/vllm-project/vllm/blob/66a168a197ba214a5b70a74fa2e713c9eeb3251a/vllm/model_executor/layers/batch_invariant.py#L649 + + Args: + input: Input tensor of shape (..., hidden_size) + weight: Weight tensor of shape (hidden_size,) + eps: Small constant for numerical stability + + Returns: + RMS normalized tensor + """ + return rms_norm(input, weight, eps=eps) + + +# --------------------------------------------------------------------------- # +# Fused batch-invariant "sglang" RMSNorm (residual + no-residual) +# +# Vendored from SGLang's fused batch-invariant residual RMSNorm +# (python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py, branch +# feature/k3-train-serve-numerics). These reproduce, bit-for-bit, the eager +# ``normalization.sglang_residual_rms_norm`` path (fp32 upcast, ``mean_dim`` +# variance, ``rsqrt``, fp32 weight multiply, cast last) while replacing its +# ~6-launch eager forward with three kernels. Bit-exactness keys: +# - reuse the batch-invariant ``mean_dim`` for the variance (a hand-rolled +# ``tl.sum`` reduction does NOT match it), +# - ``tl.rsqrt`` matches ``torch.rsqrt(var + eps)`` (``1.0 / tl.sqrt`` does not), +# - residual add is upcast -> add -> round back to the input dtype (matches +# torch's fp32-accumulated bf16 elementwise add). +# The forward is order-identical to the eager path, so the static K3 forward is +# preserved exactly. These are forward-only; the trainer wraps them in an +# ``autograd.Function`` with a closed-form backward (see normalization.py). +# --------------------------------------------------------------------------- # +@triton.jit +def _add_residual_square_kernel( + input_ptr, + residual_ptr, + residual_out_ptr, + sq_ptr, + input_row_stride: tl.constexpr, + residual_row_stride: tl.constexpr, + residual_out_row_stride: tl.constexpr, + sq_row_stride: tl.constexpr, + n_cols: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """ + Stage 1 of the fused residual-add RMSNorm: residual add in the original + dtype and the per-element square in float32. + + residual_out = (x + residual).to(orig_dtype) + sq = residual_out.float() ** 2 # float32 + + ``sq`` is then reduced by the existing batch-invariant ``mean_dim`` kernel, + so the variance reduction order is bit-identical to the eager + ``x.pow(2).mean(-1)`` path this replaces. + """ + row_idx = tl.program_id(0).to(tl.int64) + in_row = input_ptr + row_idx * input_row_stride + res_row = residual_ptr + row_idx * residual_row_stride + res_out_row = residual_out_ptr + row_idx * residual_out_row_stride + sq_row = sq_ptr + row_idx * sq_row_stride + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + x = tl.load(in_row + col_idx, mask=mask, other=0.0) + r = tl.load(res_row + col_idx, mask=mask, other=0.0) + # Match torch's elementwise add for low-precision dtypes: upcast to + # float32, add, round the result back to the original dtype. The + # normalization then operates on this rounded value. + s = (x.to(tl.float32) + r.to(tl.float32)).to(x.dtype) + tl.store(res_out_row + col_idx, s, mask=mask) + s_f32 = s.to(tl.float32) + tl.store(sq_row + col_idx, s_f32 * s_f32, mask=mask) + + +@triton.jit +def _square_kernel( + input_ptr, + sq_ptr, + input_row_stride: tl.constexpr, + sq_row_stride: tl.constexpr, + n_cols: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """No-residual analog of stage 1: per-element square in float32. + + sq = x.float() ** 2 + + Reduced by ``mean_dim`` for a variance bit-identical to the eager + ``x.float().pow(2).mean(-1)`` path (fp32 ``s * s`` == ``pow(x, 2)``). + """ + row_idx = tl.program_id(0).to(tl.int64) + in_row = input_ptr + row_idx * input_row_stride + sq_row = sq_ptr + row_idx * sq_row_stride + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + x = tl.load(in_row + col_idx, mask=mask, other=0.0) + x_f32 = x.to(tl.float32) + tl.store(sq_row + col_idx, x_f32 * x_f32, mask=mask) + + +@triton.jit +def _rms_normalize_with_var_kernel( + input_ptr, + var_ptr, + weight_ptr, + output_ptr, + input_row_stride: tl.constexpr, + output_row_stride: tl.constexpr, + n_cols: tl.constexpr, + eps, + BLOCK_SIZE: tl.constexpr, +): + """ + Stage 2 of the fused RMSNorm: normalize by a precomputed per-row variance + and multiply weight in float32, casting last. + + out = (x.float() * rsqrt(var + eps) * weight.float()).to(orig_dtype) + + ``tl.rsqrt`` bit-matches ``torch.rsqrt(var + eps)`` used by the eager path + (``1.0 / tl.sqrt`` does not). + """ + row_idx = tl.program_id(0).to(tl.int64) + in_row = input_ptr + row_idx * input_row_stride + out_row = output_ptr + row_idx * output_row_stride + var = tl.load(var_ptr + row_idx) + inv_rms = tl.rsqrt(var + eps) + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + x = tl.load(in_row + col_idx, mask=mask, other=0.0) + weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) + output_f32 = x.to(tl.float32) * inv_rms * weight.to(tl.float32) + tl.store(out_row + col_idx, output_f32.to(x.dtype), mask=mask) + + +def fused_add_rms_norm_batch_invariant( + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Batch-invariant fused residual-add + RMS normalization. + + Returns ``(output, residual_out)`` bit-matching the eager + ``residual_out = input + residual`` followed by + ``normalization.sglang_residual_rms_norm(residual_out, weight, eps)`` for the + closed dense Qwen3 recipe (fp32 upcast, ``mean_dim`` variance, ``rsqrt``, + fp32 weight multiply, cast last). + + The eager path is ~6 small launches per call; here it is three: a fused + residual-add+square, the batch-invariant ``mean_dim`` reduction (reused + verbatim so the variance is bit-identical to ``x.pow(2).mean(-1)``), and a + fused normalize. Forward-only; wrap in an autograd.Function for training. + """ + assert weight.dim() == 1, "Weight must be 1-dimensional" + assert input.shape == residual.shape, "Input and residual must share a shape" + assert input.shape[-1] == weight.shape[0], ( + f"Input last dimension ({input.shape[-1]}) must match weight dimension ({weight.shape[0]})" + ) + + original_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]).contiguous() + residual_2d = residual.reshape(-1, residual.shape[-1]).contiguous() + weight = weight.contiguous() + + n_rows, n_cols = input_2d.shape + residual_out = torch.empty_like(input_2d) + sq = torch.empty((n_rows, n_cols), dtype=torch.float32, device=input.device) + + BLOCK_SIZE = 1024 + grid = (n_rows,) + _add_residual_square_kernel[grid]( + input_2d, + residual_2d, + residual_out, + sq, + input_2d.stride(0), + residual_2d.stride(0), + residual_out.stride(0), + sq.stride(0), + n_cols, + BLOCK_SIZE=BLOCK_SIZE, + ) + + # Reuse the batch-invariant mean reduction verbatim: variance is then + # bit-identical to the eager path's x.pow(2).mean(-1). + var = mean_dim(sq, -1, keepdim=True).reshape(-1).contiguous() + + output = torch.empty_like(input_2d) + _rms_normalize_with_var_kernel[grid]( + residual_out, + var, + weight, + output, + residual_out.stride(0), + output.stride(0), + n_cols, + eps, + BLOCK_SIZE=BLOCK_SIZE, + ) + return output.reshape(original_shape), residual_out.reshape(original_shape) + + +@triton.jit +def _rms_norm_backward_kernel( + normed_ptr, + grad_out_ptr, + grad_ro_ptr, + weight_ptr, + grad_in_ptr, + gw_partial_ptr, + normed_row_stride: tl.constexpr, + grad_out_row_stride: tl.constexpr, + grad_ro_row_stride: tl.constexpr, + grad_in_row_stride: tl.constexpr, + gw_row_stride: tl.constexpr, + n_cols: tl.constexpr, + eps, + HAS_RESIDUAL_GRAD: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Fused RMSNorm backward for the ``sglang``-style norm. + + Given ``normed_input`` (the tensor that was normalized), ``grad_output`` and + optionally ``grad_residual_out``, computes per row:: + + var = mean(rf^2); inv = rsqrt(var + eps) + S = sum_j grad_out_j * w_j * rf_j + grad_normed_k = inv * grad_out_k * w_k - (inv^3 * S / D) * rf_k + gw_partial_k = grad_out_k * rf_k * inv # summed over rows -> grad_weight + + When ``grad_residual_out`` is present it is added to ``grad_normed`` (the + residual stream feeds both the norm and the next layer). This is the + gradient path only; it does not enter the forward K3, so ``var`` is + recomputed with a plain reduction (not ``mean_dim``). + """ + row = tl.program_id(0).to(tl.int64) + n_row = normed_ptr + row * normed_row_stride + go_row = grad_out_ptr + row * grad_out_row_stride + gi_row = grad_in_ptr + row * grad_in_row_stride + gw_row = gw_partial_ptr + row * gw_row_stride + + sum_sq = tl.zeros([1], dtype=tl.float32) + s_acc = tl.zeros([1], dtype=tl.float32) + for off in range(0, n_cols, BLOCK_SIZE): + idx = off + tl.arange(0, BLOCK_SIZE) + mask = idx < n_cols + rf = tl.load(n_row + idx, mask=mask, other=0.0).to(tl.float32) + go = tl.load(go_row + idx, mask=mask, other=0.0).to(tl.float32) + w = tl.load(weight_ptr + idx, mask=mask, other=0.0).to(tl.float32) + sum_sq += tl.sum(tl.where(mask, rf * rf, 0.0)) + s_acc += tl.sum(tl.where(mask, go * w * rf, 0.0)) + var = sum_sq / n_cols + inv = tl.rsqrt(var + eps) + c = inv * inv * inv * s_acc / n_cols + + for off in range(0, n_cols, BLOCK_SIZE): + idx = off + tl.arange(0, BLOCK_SIZE) + mask = idx < n_cols + rf = tl.load(n_row + idx, mask=mask, other=0.0).to(tl.float32) + go = tl.load(go_row + idx, mask=mask, other=0.0).to(tl.float32) + w = tl.load(weight_ptr + idx, mask=mask, other=0.0).to(tl.float32) + grad_normed = inv * go * w - c * rf + if HAS_RESIDUAL_GRAD: + gro = tl.load(grad_ro_ptr + row * grad_ro_row_stride + idx, mask=mask, other=0.0).to(tl.float32) + grad_normed = grad_normed + gro + tl.store(gi_row + idx, grad_normed, mask=mask) + tl.store(gw_row + idx, go * rf * inv, mask=mask) + + +def fused_rms_norm_backward( + normed_input: torch.Tensor, + weight: torch.Tensor, + eps: float, + grad_output: torch.Tensor, + grad_residual_out: torch.Tensor | None = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fused backward for :func:`fused_add_rms_norm_batch_invariant` / + :func:`sglang_rms_norm_batch_invariant`. + + Returns ``(grad_normed_input_fp32, grad_weight_fp32)``. ``grad_normed_input`` + already includes ``grad_residual_out`` when supplied. Not order-sensitive: + gradients do not enter the forward K3. + """ + original_shape = normed_input.shape + ni = normed_input.reshape(-1, original_shape[-1]).contiguous() + go = grad_output.reshape(-1, original_shape[-1]).contiguous() + wf = weight.float().contiguous() + + n_rows, n_cols = ni.shape + grad_in = torch.empty((n_rows, n_cols), dtype=torch.float32, device=ni.device) + gw_partial = torch.empty((n_rows, n_cols), dtype=torch.float32, device=ni.device) + + has_residual = grad_residual_out is not None + if has_residual: + gro = grad_residual_out.reshape(-1, original_shape[-1]).contiguous() + else: + gro = ni # unused; kernel does not read it when HAS_RESIDUAL_GRAD=False + + BLOCK_SIZE = 1024 + grid = (n_rows,) + _rms_norm_backward_kernel[grid]( + ni, + go, + gro, + wf, + grad_in, + gw_partial, + ni.stride(0), + go.stride(0), + gro.stride(0), + grad_in.stride(0), + gw_partial.stride(0), + n_cols, + eps, + HAS_RESIDUAL_GRAD=has_residual, + BLOCK_SIZE=BLOCK_SIZE, + ) + grad_weight = gw_partial.sum(0) + return grad_in.reshape(original_shape), grad_weight + + +def sglang_rms_norm_batch_invariant( + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, +) -> torch.Tensor: + """ + Batch-invariant RMS normalization bit-matching the eager + ``normalization.sglang_residual_rms_norm`` (no residual add). + + This is the no-residual analog of :func:`fused_add_rms_norm_batch_invariant` + for the ``force_sglang_residual`` call sites (input layernorm at layer>0 and + the final norm), which apply the residual-style fp32 normalization to a + single input. Two kernels + ``mean_dim`` replace the eager ~6 launches. + Forward-only; wrap in an autograd.Function for training. + """ + assert weight.dim() == 1, "Weight must be 1-dimensional" + assert input.shape[-1] == weight.shape[0], ( + f"Input last dimension ({input.shape[-1]}) must match weight dimension ({weight.shape[0]})" + ) + + original_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]).contiguous() + weight = weight.contiguous() + + n_rows, n_cols = input_2d.shape + sq = torch.empty((n_rows, n_cols), dtype=torch.float32, device=input.device) + + BLOCK_SIZE = 1024 + grid = (n_rows,) + _square_kernel[grid]( + input_2d, + sq, + input_2d.stride(0), + sq.stride(0), + n_cols, + BLOCK_SIZE=BLOCK_SIZE, + ) + + var = mean_dim(sq, -1, keepdim=True).reshape(-1).contiguous() + + output = torch.empty_like(input_2d) + _rms_normalize_with_var_kernel[grid]( + input_2d, + var, + weight, + output, + input_2d.stride(0), + output.stride(0), + n_cols, + eps, + BLOCK_SIZE=BLOCK_SIZE, + ) + return output.reshape(original_shape) + + +# --------------------------------------------------------------------------- # +# RMSNorm kernel-family contract +# +# Two batch-invariant RMSNorm kernel families coexist, and they disagree at +# 1 ulp on rare bf16 boundary values (~2/524288 at [4096, 128]), so silently +# swapping one for the other seeds K3 divergence that amplifies downstream. +# Each family is pinned to the serving site-class that executes it: +# - "serving_no_residual" (family-1): the looped ``tl.sum`` + ``1.0/tl.sqrt`` +# kernel (``rms_norm_batch_invariant``), what SGLang dispatches when +# ``residual is None`` under batch-invariant mode and what the +# ``aten::rms_norm`` interpose runs. Site-classes: qk-norm, layer-0 input +# layernorm. +# - "serving_residual_tree" (family-2): the ``mean_dim`` + ``tl.rsqrt`` +# fused residual-tree kernels (``fused_add_rms_norm_batch_invariant`` / +# ``sglang_rms_norm_batch_invariant``), what SGLang dispatches for +# residual calls under the rl-on-policy lane. Site-classes: input +# layernorm at layer>0, post-attention layernorm, final norm. +# Every call site must name its family through ``bi_rms_norm`` / +# ``bi_fused_add_rms_norm``; never call the family kernels directly. +# --------------------------------------------------------------------------- # +RMS_NORM_FAMILY_NO_RESIDUAL = "serving_no_residual" +RMS_NORM_FAMILY_RESIDUAL_TREE = "serving_residual_tree" +RMS_NORM_FAMILIES = (RMS_NORM_FAMILY_NO_RESIDUAL, RMS_NORM_FAMILY_RESIDUAL_TREE) +RMSNormFamily = Literal["serving_no_residual", "serving_residual_tree"] + + +def bi_rms_norm( + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + *, + family: RMSNormFamily, + zero_centered: bool = False, +) -> torch.Tensor: + """Single-tensor batch-invariant RMSNorm with an explicit kernel family. + + The single sanctioned entry point for both no-residual family kernels; the + ``family`` keyword must name the serving counterpart of the call site. + + ``zero_centered`` is the Gemma-style Qwen3.5 form: an fp32 upcast with the + ``1 + weight`` scale folded in fp32, cast back last, exactly as + ``normalization.native_zero_centered_rms_norm`` composes it. It is an affine + fold around the SAME family-1 reduction tree — not a third family — and only + exists in no-residual form (Qwen3.5 residual-tree norms run the eager native + path, never a batch-invariant kernel). + """ + if zero_centered: + if family != RMS_NORM_FAMILY_NO_RESIDUAL: + raise ValueError( + "zero-centered RMSNorm only exists in the 'serving_no_residual' family; " + "Qwen3.5 residual-tree norms run the native path, not a batch-invariant kernel" + ) + return rms_norm_batch_invariant(input.float(), 1.0 + weight.float(), eps=eps).type_as(input) + if family == RMS_NORM_FAMILY_NO_RESIDUAL: + return rms_norm_batch_invariant(input, weight, eps=eps) + if family == RMS_NORM_FAMILY_RESIDUAL_TREE: + return sglang_rms_norm_batch_invariant(input, weight, eps=eps) + raise ValueError(f"Unknown RMSNorm family {family!r}; expected one of {RMS_NORM_FAMILIES}") + + +def bi_fused_add_rms_norm( + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + *, + family: RMSNormFamily, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fused residual-add batch-invariant RMSNorm with an explicit kernel family. + + Only the serving residual tree has a fused-add kernel; requesting the + no-residual family with a residual stream is a contract violation (serving + never runs family-1 on a residual site) and raises. + """ + if family == RMS_NORM_FAMILY_RESIDUAL_TREE: + return fused_add_rms_norm_batch_invariant(input, residual, weight, eps=eps) + if family == RMS_NORM_FAMILY_NO_RESIDUAL: + raise ValueError( + "RMSNorm family 'serving_no_residual' has no fused-add kernel: residual " + "site-classes are 'serving_residual_tree' by the cross-engine contract" + ) + raise ValueError(f"Unknown RMSNorm family {family!r}; expected one of {RMS_NORM_FAMILIES}") + + +# --------------------------------------------------------------------------- # +# Batch-invariant fused LM-head selected-token log-probability +# +# The K3 lm-head contract, vendored identically in xorl and SGLang +# (python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py). Both engines +# compute per-token logprobs of given token ids from bit-exact bf16 hidden +# states and the bf16 lm-head weight through the SAME reduction trees, so the +# results are bitwise identical cross-engine: +# 1. chunk GEMM — ``matmul_kernel_persistent`` with the family's fixed bf16 +# tile config and an fp32 output buffer. bf16xbf16 products are exact in +# fp32, so reading the weight in bf16 with tensor-core fp32 accumulation +# equals a GEMM over materialized fp32 upcasts with the same tree (and +# deletes the fp32 weight copy the eager paths materialize). +# 2. chunk stats — per-row max and sum(exp(x - chunk_max)) in a fixed +# sequential BLOCK loop (same discipline as ``_rms_norm_kernel``), plus +# the selected-token logit gather. +# 3. merge — global max over chunk maxima (exact), then the rescaled sumexp +# accumulated in pinned chunk order; lse = gmax + log(acc). +# All transcendentals stay inside these kernels: tl.exp/tl.log measured +# bit-identical across triton 3.5.1 (serving venv) and 3.7.1 (trainer venv), +# as is the fixed-tile tl.dot fp32 accumulator. VOCAB_CHUNK and STATS_BLOCK are +# contract constants — changing either changes the bits (the LSE reduction +# tree). The chunk GEMM's tile config is shape-keyed via bi_gemm_configs: only +# its BLOCK_SIZE_K (pinned there) is bit-relevant. +# Forward-only; the trainer wraps it in an autograd.Function (ops/loss). +# --------------------------------------------------------------------------- # + +BI_LM_HEAD_VOCAB_CHUNK = 8192 +_BI_LM_HEAD_STATS_BLOCK = 1024 + + +@triton.jit +def _lm_head_chunk_stats_kernel( + logits_ptr, + token_ids_ptr, + sel_ptr, + m_ptr, + s_ptr, + temp_ptr, + logits_row_stride, + n_cols, + col_offset, + chunk_idx, + n_chunks, + BLOCK_SIZE: tl.constexpr, + HAS_TEMP: tl.constexpr, +): + """Per-row chunk statistics over an fp32 logits tile [N, n_cols]: + chunk max, sum(exp(x - chunk_max)) in a fixed sequential block loop, and + the selected-token logit when ``token_ids[row]`` falls in this chunk. + With HAS_TEMP, logits are scaled by 1/temp[row] before the statistics and + the selected logit; the fp32 divide runs in-kernel so every engine + computes the identical scale (elementwise, so batch-invariance holds).""" + row = tl.program_id(0).to(tl.int64) + row_ptr = logits_ptr + row * logits_row_stride + if HAS_TEMP: + inv_t = 1.0 / tl.load(temp_ptr + row) + else: + inv_t = 1.0 + + row_max = float("-inf") + for col_start in range(0, n_cols, BLOCK_SIZE): + col_idx = col_start + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + vals = tl.load(row_ptr + col_idx, mask=mask, other=float("-inf")) + if HAS_TEMP: + vals = vals * inv_t + row_max = tl.maximum(row_max, tl.max(vals)) + + sum_exp = 0.0 + for col_start in range(0, n_cols, BLOCK_SIZE): + col_idx = col_start + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + vals = tl.load(row_ptr + col_idx, mask=mask, other=float("-inf")) + if HAS_TEMP: + vals = vals * inv_t + e = tl.exp(vals - row_max) + sum_exp += tl.sum(tl.where(mask, e, 0.0)) + + tl.store(m_ptr + row * n_chunks + chunk_idx, row_max) + tl.store(s_ptr + row * n_chunks + chunk_idx, sum_exp) + + tok = tl.load(token_ids_ptr + row) + local = tok - col_offset + in_chunk = (local >= 0) & (local < n_cols) + sel = tl.load(row_ptr + local, mask=in_chunk, other=0.0) + if HAS_TEMP: + sel = sel * inv_t + tl.store(sel_ptr + row, sel, mask=in_chunk) + + +@triton.jit +def _lm_head_lse_merge_kernel( + m_ptr, + s_ptr, + lse_ptr, + n_chunks, +): + """lse[row] = gmax + log(sum_c s_c * exp(m_c - gmax)), chunks in pinned order.""" + row = tl.program_id(0).to(tl.int64) + base = row * n_chunks + gmax = float("-inf") + for c in range(n_chunks): + gmax = tl.maximum(gmax, tl.load(m_ptr + base + c)) + acc = 0.0 + for c in range(n_chunks): + acc += tl.load(s_ptr + base + c) * tl.exp(tl.load(m_ptr + base + c) - gmax) + tl.store(lse_ptr + row, gmax + tl.log(acc)) + + +def _bi_lm_head_chunk_gemm_fp32(a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> None: + """Launch the family's persistent matmul with the shape-keyed bf16 config and + an fp32 output buffer (the fp32 store path keeps the raw accumulator bits).""" + NUM_SMS = torch.cuda.get_device_properties(a.device).multi_processor_count + M, K = a.shape + _, N = b.shape + + def grid(META): + return (min(NUM_SMS, triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"])),) + + def _launch(config): + matmul_kernel_persistent[grid]( + a, + b, + out, + None, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + out.stride(0), + out.stride(1), + NUM_SMS=NUM_SMS, + A_LARGE=a.numel() > 2**31, + B_LARGE=b.numel() > 2**31, + C_LARGE=out.numel() > 2**31, + HAS_BIAS=False, + **config, + ) + + _launch_with_config_fallback(_launch, a.dtype, M, N, K, out_itemsize=out.element_size()) + + +def bi_lm_head_selected_logprob( + hidden: torch.Tensor, + weight: torch.Tensor, + token_ids: torch.Tensor, + temperature: Optional[torch.Tensor] = None, + vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Per-token ``log p(token_ids)`` for the LM head, batch-invariant and + cross-engine bit-exact (the K3 lm-head contract). + + Args: + hidden: ``[N, H]`` bf16 hidden states (pre lm-head). + weight: ``[V, H]`` bf16 lm-head weight (kept resident in bf16; no fp32 + copy is materialized). + token_ids: ``[N]`` integer token ids to score (callers must pre-clamp + ignored positions to a valid id and mask outputs downstream). + temperature: optional ``[N]`` fp32 per-row temperatures (> 0). Logits + are scaled by ``1/temperature[row]`` inside the stats kernel (the + divide runs in-kernel, so engines sharing the contract compute the + identical scale). ``None`` is the exact temperature-1.0 path. + + Returns: + ``(logprob, lse, selected)`` — all ``[N]`` fp32; ``logprob = selected - lse`` + (temperature-scaled when ``temperature`` is given). + """ + assert hidden.ndim == 2 and weight.ndim == 2, "hidden and weight must be 2D" + assert hidden.shape[1] == weight.shape[1], "hidden dim mismatch" + assert hidden.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( + "the lm-head contract takes bf16 hidden/weight (fp32 upcast is exact inside the GEMM)" + ) + assert hidden.is_cuda, "CUDA only" + + hidden = hidden.contiguous() + token_ids = token_ids.contiguous().to(device=hidden.device, dtype=torch.int64) + n_tokens = hidden.shape[0] + vocab = weight.shape[0] + n_chunks = (vocab + vocab_chunk - 1) // vocab_chunk + if temperature is not None: + temperature = temperature.reshape(-1).to(device=hidden.device, dtype=torch.float32).contiguous() + assert temperature.shape[0] == n_tokens, "temperature must be per-row [N]" + assert bool((temperature > 0).all()), "temperature must be > 0" + + chunk_max = torch.empty((n_tokens, n_chunks), dtype=torch.float32, device=hidden.device) + chunk_sumexp = torch.empty_like(chunk_max) + selected = torch.zeros(n_tokens, dtype=torch.float32, device=hidden.device) + lse = torch.empty(n_tokens, dtype=torch.float32, device=hidden.device) + logits_buf = torch.empty((n_tokens, vocab_chunk), dtype=torch.float32, device=hidden.device) + + for chunk_idx, col_start in enumerate(range(0, vocab, vocab_chunk)): + col_end = min(col_start + vocab_chunk, vocab) + n_cols = col_end - col_start + logits_c = logits_buf[:, :n_cols] + # [H, C] transposed view of the resident bf16 weight — the persistent + # kernel takes explicit strides, so no copy is made. + _bi_lm_head_chunk_gemm_fp32(hidden, weight[col_start:col_end].t(), logits_c) + _lm_head_chunk_stats_kernel[(n_tokens,)]( + logits_c, + token_ids, + selected, + chunk_max, + chunk_sumexp, + temperature, + logits_c.stride(0), + n_cols, + col_start, + chunk_idx, + n_chunks, + BLOCK_SIZE=_BI_LM_HEAD_STATS_BLOCK, + HAS_TEMP=temperature is not None, + ) + + _lm_head_lse_merge_kernel[(n_tokens,)](chunk_max, chunk_sumexp, lse, n_chunks) + # In exact math the selected logit never exceeds the LSE; clamp the one-ulp + # fp boundary case (p~1 tokens) so contract logprobs are provably <= 0. + return torch.clamp_max(selected - lse, 0.0), lse, selected + + +def bi_lm_head_full_logits( + hidden: torch.Tensor, + weight: torch.Tensor, + vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, +) -> torch.Tensor: + """Materialize full FP32 logits through the v1 contract GEMM.""" + + assert hidden.ndim == 2 and weight.ndim == 2, "hidden and weight must be 2D" + assert hidden.shape[1] == weight.shape[1], "hidden dim mismatch" + assert hidden.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( + "the lm-head contract takes bf16 hidden/weight" + ) + assert hidden.is_cuda, "CUDA only" + + hidden = hidden.contiguous() + logits = torch.empty( + (hidden.shape[0], weight.shape[0]), + dtype=torch.float32, + device=hidden.device, + ) + if hidden.shape[0] != 0: + _bi_lm_head_chunk_gemm_fp32(hidden, weight.t(), logits) + return logits + + +def bi_lm_head_selected_logprob_from_logits( + logits: torch.Tensor, + token_ids: torch.Tensor, + temperature: Optional[torch.Tensor] = None, + vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the v1 stats and merge tree over materialized FP32 logits.""" + + assert logits.ndim == 2 and logits.dtype == torch.float32, "logits must be FP32 [N, V]" + assert logits.is_cuda and logits.stride(1) == 1, "logits must be contiguous CUDA rows" + n_tokens, vocab = logits.shape + token_ids = token_ids.contiguous().to(device=logits.device, dtype=torch.int64) + assert token_ids.shape == (n_tokens,), "token_ids must be per-row [N]" + n_chunks = (vocab + vocab_chunk - 1) // vocab_chunk + if temperature is not None: + temperature = temperature.reshape(-1).to(device=logits.device, dtype=torch.float32).contiguous() + assert temperature.shape == (n_tokens,), "temperature must be per-row [N]" + torch._assert_async((temperature > 0).all(), "temperature must be > 0") + + chunk_max = torch.empty((n_tokens, n_chunks), dtype=torch.float32, device=logits.device) + chunk_sumexp = torch.empty_like(chunk_max) + selected = torch.zeros(n_tokens, dtype=torch.float32, device=logits.device) + lse = torch.empty(n_tokens, dtype=torch.float32, device=logits.device) + for chunk_idx, col_start in enumerate(range(0, vocab, vocab_chunk)): + col_end = min(col_start + vocab_chunk, vocab) + logits_c = logits[:, col_start:col_end] + _lm_head_chunk_stats_kernel[(n_tokens,)]( + logits_c, + token_ids, + selected, + chunk_max, + chunk_sumexp, + temperature, + logits_c.stride(0), + col_end - col_start, + col_start, + chunk_idx, + n_chunks, + BLOCK_SIZE=_BI_LM_HEAD_STATS_BLOCK, + HAS_TEMP=temperature is not None, + ) + + _lm_head_lse_merge_kernel[(n_tokens,)](chunk_max, chunk_sumexp, lse, n_chunks) + return torch.clamp_max(selected - lse, 0.0), lse, selected + + +# --------------------------------------------------------------------------- # +# Batch-invariant MoE router GEMM (the K3 router contract) +# +# Vendored identically in xorl and SGLang so the MoE gate/router logits are +# computed through ONE reduction tree cross-engine. Unlike the capture/replay +# lane, live training routes independently of serving (no routing replay), so a +# ~1e-10..1e-4 router-logit reduction-order diff between the two engines' GEMMs +# can flip the top-k expert selection on razor-edge tokens and cause large, +# rare-token logprob divergence. This kernel removes that last term: +# - bf16 hidden [N, H] @ bf16 gate weight [E, H]^T -> fp32 logits [N, E] +# - ``matmul_kernel_persistent`` with a pinned tile config and an fp32 output +# buffer (same discipline as the lm-head contract). bf16xbf16 products are +# exact in fp32, so reading both operands in bf16 with tensor-core fp32 +# accumulation equals an fp32 GEMM over their (exact) fp32 upcasts, but with +# the reduction order pinned identically in both engines — and without the +# fp32 weight/activation copies the eager fp32-router paths materialize. +# num_experts is small (a single BLOCK_SIZE_N tile for the common E <= 128), so +# the whole GEMM is one persistent launch. The config below is part of the +# contract; changing any constant changes the bits. +# Exact model programs call this kernel directly on both trainer and sampler. +# Forward-only; the trainer wraps it in an autograd.Function with a closed-form +# (order-insensitive) backward — gradients do not enter the forward K3. +# --------------------------------------------------------------------------- # + +_BI_ROUTER_GEMM_CONFIG = { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_stages": 3, + "num_warps": 8, +} + + +def bi_router_gemm(hidden: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """fp32 MoE router logits ``[N, E]`` from bf16 hidden ``[N, H]`` and bf16 gate + weight ``[E, H]`` (the K3 router contract). + + One persistent bf16-in / fp32-accumulate / fp32-out GEMM with a pinned launch + config, vendored identically in xorl and SGLang so the router logits (and + therefore the top-k expert selection) are bitwise identical cross-engine. + bf16xbf16 products are exact in fp32, so this equals an fp32 GEMM over the + upcast operands the eager fp32-router paths materialize — minus the fp32 + weight/activation copies and with a reduction order that no longer depends on + the backend GEMM. + + Args: + hidden: ``[N, H]`` bf16 hidden states (pre-gate). + weight: ``[E, H]`` bf16 gate weight (kept resident in bf16; no fp32 copy). + + Returns: + ``[N, E]`` fp32 router logits. + """ + assert hidden.ndim == 2 and weight.ndim == 2, "hidden and weight must be 2D" + assert hidden.shape[1] == weight.shape[1], "hidden dim mismatch" + assert hidden.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( + "the router contract takes bf16 hidden/weight (fp32 upcast is exact inside the GEMM)" + ) + assert hidden.is_cuda, "CUDA only" + + hidden = hidden.contiguous() + weight = weight.contiguous() + n_tokens = hidden.shape[0] + num_experts = weight.shape[0] + logits = torch.empty((n_tokens, num_experts), dtype=torch.float32, device=hidden.device) + if n_tokens == 0: + return logits + + # [H, E] transposed view of the resident bf16 gate weight — the persistent + # kernel takes explicit strides, so no copy is made. + b = weight.t() + NUM_SMS = torch.cuda.get_device_properties(hidden.device).multi_processor_count + M, K = hidden.shape + N = num_experts + + def grid(META): + return (min(NUM_SMS, triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"])),) + + matmul_kernel_persistent[grid]( + hidden, + b, + logits, + None, + M, + N, + K, + hidden.stride(0), + hidden.stride(1), + b.stride(0), + b.stride(1), + logits.stride(0), + logits.stride(1), + NUM_SMS=NUM_SMS, + A_LARGE=hidden.numel() > 2**31, + B_LARGE=weight.numel() > 2**31, + C_LARGE=logits.numel() > 2**31, + HAS_BIAS=False, + **_BI_ROUTER_GEMM_CONFIG, + ) + return logits + + +class _BIBf16Fp32LinearFn(torch.autograd.Function): + """Trainable wrapper for the shared BF16-input, FP32-output GEMM. + + The pinned forward reduction enters the trainer-sampler numerical + contract. The ordinary linear backward does not enter K3, so it may use + the native matmul reduction while still propagating gradients to both + operands. + """ + + @staticmethod + def forward(ctx, input: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(input, weight) + input_2d = input.reshape(-1, input.shape[-1]) + output = bi_router_gemm(input_2d, weight) + return output.reshape(*input.shape[:-1], weight.shape[0]) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + input, weight = ctx.saved_tensors + grad_output_2d = grad_output.reshape(-1, grad_output.shape[-1]).float() + input_2d = input.reshape(-1, input.shape[-1]) + grad_input = grad_weight = None + if ctx.needs_input_grad[0]: + grad_input = (grad_output_2d @ weight.float()).to(input.dtype).reshape_as(input) + if ctx.needs_input_grad[1]: + grad_weight = (grad_output_2d.t() @ input_2d.float()).to(weight.dtype) + return grad_input, grad_weight + + +def bi_bf16_fp32_linear(input: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Apply the batch-invariant BF16-input, FP32-output linear contract. + + This is the differentiable form of :func:`bi_router_gemm`. It is suitable + for small-output projections whose serving counterpart uses DeepGEMM's + ``bf16_gemm_nt`` FP32 output, including GLM-5.2's indexer head weights. + """ + + if input.ndim < 2 or weight.ndim != 2: + raise ValueError(f"expected input.ndim >= 2 and weight.ndim == 2, got {input.shape=} and {weight.shape=}") + if input.shape[-1] != weight.shape[-1]: + raise ValueError(f"input/weight contraction mismatch: {input.shape[-1]} != {weight.shape[-1]}") + return _BIBf16Fp32LinearFn.apply(input, weight) + + +def bi_router_topk_weights( + topk_vals: torch.Tensor, + norm_topk_prob: bool = True, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Renormalize the gathered top-k router scores in a fixed reduction order, + then cast — the second half of the K3 router contract. + + ``bi_router_gemm`` makes the router logits (and therefore ``torch.topk`` + selection and the gathered top-k softmax scores) bitwise identical + cross-engine, but the stock renorm ``vals / vals.sum(dim=-1, keepdim=True)`` + is not: the small last-dim reduction ``sum(dim=-1)`` uses a build-dependent + tree order, so the divisor can differ by ~1 fp32 ulp between the trainer's + and the server's torch/triton, occasionally flipping the final bf16 weight + on rare tokens. Elementwise ops (add, divide, round-to-bf16) are IEEE + correctly-rounded and build-invariant, so accumulating the divisor with a + pinned left-to-right order makes the top-k weights bit-identical too. + + Args: + topk_vals: ``[N, top_k]`` fp32 gathered top-k router scores (softmax + probabilities on the softmax path). + norm_topk_prob: renormalize the top-k slice to sum to 1 (Qwen3 MoE + default). When False the scores are only cast (already bit-identical + cross-engine, since the softmax/top-k that produced them are). + out_dtype: routing-weight dtype (the model activation dtype, bf16). + + Returns: + ``[N, top_k]`` ``out_dtype`` routing weights. + """ + assert topk_vals.dtype == torch.float32, "the router contract renorms fp32 top-k scores" + if norm_topk_prob: + denom = topk_vals[..., 0] + for k in range(1, topk_vals.shape[-1]): + denom = denom + topk_vals[..., k] + topk_vals = topk_vals / denom.unsqueeze(-1) + return topk_vals.to(out_dtype) + + +_ONES_CACHE: Dict[Tuple[str, int | None, torch.dtype, int], torch.Tensor] = {} + + +def _get_or_make_ones(input: torch.Tensor, normalized_shape: list[int]) -> torch.Tensor: + assert len(normalized_shape) == 1, "Only last-dimension RMSNorm is supported" + hidden_size = int(normalized_shape[0]) + key = (input.device.type, input.device.index, input.dtype, hidden_size) + weight = _ONES_CACHE.get(key) + if weight is None or weight.device != input.device: + weight = torch.ones(hidden_size, device=input.device, dtype=input.dtype) + _ONES_CACHE[key] = weight + return weight + + +def _rms_norm_aten_compat(input, normalized_shape, weight=None, eps=None): + _guard_interpose_no_grad("aten::rms_norm", input, weight) + normalized_shape = [int(dim) for dim in normalized_shape] + if len(normalized_shape) != 1 or input.shape[-1] != normalized_shape[0]: + raise NotImplementedError("Batch-invariant RMSNorm only supports last dimension") + if weight is None: + weight = _get_or_make_ones(input, normalized_shape) + if eps is None: + eps = torch.finfo(input.dtype).eps + # The interpose IS the no-residual family: every F.rms_norm that reaches it + # is a no-residual site (qk-norm, layer-0 input norm) by the family contract. + return bi_rms_norm(input, weight, eps=eps, family=RMS_NORM_FAMILY_NO_RESIDUAL) + + +def _mm_dtype_compat(a, b, out_dtype): + out = mm_batch_invariant(a, b) + if out.dtype != out_dtype: + out = out.to(out_dtype) + return out + + +_batch_invariant_MODE = False +_batch_invariant_LIB = None +_batch_invariant_OPS: set[str] = set() +_original_torch_bmm = None + + +def is_batch_invariant_mode_enabled(): + return _batch_invariant_MODE + + +def get_batch_invariant_ops() -> tuple[str, ...]: + return tuple(sorted(_batch_invariant_OPS)) if _batch_invariant_MODE else () + + +def is_batch_invariant_op_enabled(op: str) -> bool: + op = _BATCH_INVARIANT_ALIASES.get(op, op) + return _batch_invariant_MODE and op in _batch_invariant_OPS + + +def enable_batch_invariant_mode( + enable_bmm: bool = True, +): + global _batch_invariant_MODE, _batch_invariant_LIB, _batch_invariant_OPS, _original_torch_bmm + if _batch_invariant_MODE: + return + + _batch_invariant_OPS = _parse_batch_invariant_ops() + _batch_invariant_MODE = True + _batch_invariant_LIB = torch.library.Library("aten", "IMPL") + if "mm" in _batch_invariant_OPS: + _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA") + if "addmm" in _batch_invariant_OPS: + _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA") + if "log_softmax" in _batch_invariant_OPS: + _batch_invariant_LIB.impl("aten::_log_softmax", _log_softmax_batch_invariant, "CUDA") + if "mean" in _batch_invariant_OPS: + _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA") + if "rms_norm" in _batch_invariant_OPS: + _batch_invariant_LIB.impl("aten::rms_norm", _rms_norm_aten_compat, "CUDA") + if "mm_dtype" in _batch_invariant_OPS: + _batch_invariant_LIB.impl("aten::mm.dtype", _mm_dtype_compat, "CUDA") + + if enable_bmm and "bmm" in _batch_invariant_OPS: + _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA") + + # Also monkeypatch torch.bmm directly as a fallback + _original_torch_bmm = torch.bmm + torch.bmm = bmm_batch_invariant + + +def disable_batch_invariant_mode(): + global _batch_invariant_MODE, _batch_invariant_LIB, _batch_invariant_OPS, _original_torch_bmm + if _batch_invariant_LIB is not None: + _batch_invariant_LIB._destroy() + if _original_torch_bmm is not None: + torch.bmm = _original_torch_bmm + _original_torch_bmm = None + _batch_invariant_MODE = False + _batch_invariant_LIB = None + _batch_invariant_OPS = set() + + +@contextlib.contextmanager +def set_batch_invariant_mode(enabled: bool = True): + was_enabled = _batch_invariant_MODE + if enabled == was_enabled: + yield + return + + if enabled: + enable_batch_invariant_mode() + else: + disable_batch_invariant_mode() + try: + yield + finally: + if was_enabled: + enable_batch_invariant_mode() + else: + disable_batch_invariant_mode() + + +_TRUNK_LINEAR_NAMES = ( + "qkv_proj", + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_up_proj", + "gate_proj", + "up_proj", + "down_proj", + # Qwen2/3.5-MoE shared-expert sigmoid gate: serving contracts it through the + # global interpose (plain nn.Linear -> aten::mm), so the trunk lane must too. + "shared_expert_gate", +) + +_TRUNK_LINEAR_CONTRACT_ACTIVE = False + + +def is_trunk_linear_contract_enabled() -> bool: + """True once :func:`wrap_trunk_linears_batch_invariant` armed the contract lane. + + RMSNorm dispatch keys off this to route no-residual (family-1) norms through + the serving batch-invariant kernel (the qk-norm term of the K3 contract). + """ + return _TRUNK_LINEAR_CONTRACT_ACTIVE + + +def set_trunk_linear_contract(enabled: bool) -> None: + global _TRUNK_LINEAR_CONTRACT_ACTIVE + _TRUNK_LINEAR_CONTRACT_ACTIVE = enabled + + +class _BatchInvariantTrunkLinearFn(torch.autograd.Function): + """Forward through the batch-invariant persistent GEMM; backward stays cuBLAS.""" + + @staticmethod + def forward(ctx, input, weight, bias): + if input.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + raise RuntimeError( + f"XORL_BI_TRUNK_LINEAR contract is bf16-only; got input={input.dtype}, weight={weight.dtype}." + ) + ctx.save_for_backward(input, weight) + ctx.has_bias = bias is not None + x2d = input.reshape(-1, input.shape[-1]) + out = matmul_persistent(x2d, weight.t(), bias=bias) + return out.reshape(*input.shape[:-1], weight.shape[0]) + + @staticmethod + def backward(ctx, grad_output): + input, weight = ctx.saved_tensors + go2d = grad_output.reshape(-1, grad_output.shape[-1]) + grad_input = grad_weight = grad_bias = None + if ctx.needs_input_grad[0]: + grad_input = (go2d @ weight).reshape_as(input) + if ctx.needs_input_grad[1]: + grad_weight = go2d.t() @ input.reshape(-1, input.shape[-1]) + if ctx.has_bias and ctx.needs_input_grad[2]: + grad_bias = go2d.sum(dim=0) + return grad_input, grad_weight, grad_bias + + +def batch_invariant_trunk_linear( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + """Run the public trunk-linear contract on an explicitly folded weight.""" + return _BatchInvariantTrunkLinearFn.apply(input, weight, bias) + + +def wrap_trunk_linears_batch_invariant( + model: torch.nn.Module, names: Tuple[str, ...] = _TRUNK_LINEAR_NAMES +) -> Dict[str, int]: + """Route ONLY transformer-trunk nn.Linear forwards through the batch-invariant GEMM. + + Module-scoped alternative to enable_batch_invariant_mode(): forward bits match the + serving-side matmul_persistent contract while backward/optimizer/loss stay on cuBLAS. + + Selection is explicit: a module is wrapped iff its leaf name is in ``names`` AND it is + exactly ``torch.nn.Linear``. lm_head/embeddings never match the name set; routed MoE + experts (FQN containing ``.experts.``) are skipped — they are contracted through the + fused sglang expert path. LoRA/QLoRA-wrapped modules and FP8/TE/custom Linear + subclasses RAISE: silently skipping them would void the bitwise contract the flag + promises. Idempotent (already-wrapped modules are left alone). + + Returns ``{leaf_name: wrapped_count}`` and arms the contract lane + (see :func:`is_trunk_linear_contract_enabled`). + """ + import types # noqa: PLC0415 + + from xorl.lora.fold import lora_merged_forward_enabled # noqa: PLC0415 + from xorl.lora.modules.base import LoraModule # noqa: PLC0415 + from xorl.lora.modules.delta_linear import LoraDeltaLinear # noqa: PLC0415 + from xorl.lora.modules.linear import LoraLinear # noqa: PLC0415 + + if is_batch_invariant_mode_enabled(): + raise RuntimeError( + "XORL_BI_TRUNK_LINEAR cannot be combined with the global batch-invariant interpose " + "(XORL_BATCH_INVARIANT_MATMUL): the wrapped backward would silently ride the interposed " + "aten::mm instead of cuBLAS. Pick one lane." + ) + + def _forward(self, input): + return _BatchInvariantTrunkLinearFn.apply(input, self.weight, self.bias) + + def _forward_lora_merged(self, input): + # Merged-forward LoRA lane: the BI GEMM runs on the canonically folded + # weight (identical bits to a serving engine that received W'); backward + # flows dW' through the straight-through fold into the LoRA factors. + return _BatchInvariantTrunkLinearFn.apply(input, self.merged_weight_for_forward(), self.bias) + + wrapped: Dict[str, int] = {} + already_wrapped = 0 + for module_name, module in model.named_modules(): + leaf = module_name.rsplit(".", 1)[-1] + if leaf not in names: + continue + if ".experts." in f".{module_name}.": + continue + if isinstance(module, LoraDeltaLinear) and module_name.endswith( + (".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj") + ): + # These are factor-only children of the fused gate_up_proj. The + # parent base GEMM is wrapped and Qwen3_5MoeMLP folds both logical + # projections into that one contracted call. + continue + if type(module) is LoraLinear and lora_merged_forward_enabled(module): + # Merged-forward contract lane: the adapted linear serves and trains + # through the folded weight, so the trunk contract composes. + if module.weight.dtype not in (torch.bfloat16, torch.float32): + raise RuntimeError( + f"XORL_BI_TRUNK_LINEAR: {module_name} weight is {module.weight.dtype}; the trunk " + "contract is bf16-only." + ) + if getattr(module, "_xorl_bi_trunk_wrapped", False): + already_wrapped += 1 + continue + module.forward = types.MethodType(_forward_lora_merged, module) + module._xorl_bi_trunk_wrapped = True + wrapped[leaf] = wrapped.get(leaf, 0) + 1 + continue + if isinstance(module, LoraModule): + raise NotImplementedError( + f"XORL_BI_TRUNK_LINEAR: {module_name} is adapter-wrapped ({type(module).__qualname__}); " + "the canonical merged-LoRA trunk contract composes only with a plain LoraLinear whose " + "model-owned exact_merged_forward property is true — enable it on that module or exclude the adapter." + ) + if type(module) is not torch.nn.Linear: + raise NotImplementedError( + f"XORL_BI_TRUNK_LINEAR: {module_name} is {type(module).__qualname__}, not a plain " + "nn.Linear; fp8/te/custom linears are outside the bf16 trunk contract." + ) + if module.weight.dtype not in (torch.bfloat16, torch.float32): + # fp32 is allowed here only as a mixed-precision master: FSDP2's mp_policy + # casts it to bf16 before forward, and the runtime guard in + # _BatchInvariantTrunkLinearFn enforces bf16 on the actual GEMM operands. + raise RuntimeError( + f"XORL_BI_TRUNK_LINEAR: {module_name} weight is {module.weight.dtype}; the trunk contract is bf16-only." + ) + if getattr(module, "_xorl_bi_trunk_wrapped", False): + already_wrapped += 1 + continue + module.forward = types.MethodType(_forward, module) + module._xorl_bi_trunk_wrapped = True + wrapped[leaf] = wrapped.get(leaf, 0) + 1 + + if not wrapped and not already_wrapped: + raise RuntimeError( + "XORL_BI_TRUNK_LINEAR=1 matched no trunk linears; expected leaf names " + f"{sorted(names)} — wire the model's projections or drop the flag." + ) + set_trunk_linear_contract(True) + return wrapped + + +AttentionBlockSize = namedtuple("AttentionBlockSize", ["block_m", "block_n"]) + + +def get_batch_invariant_attention_block_size() -> AttentionBlockSize: + return AttentionBlockSize(block_m=16, block_n=16) diff --git a/src/xorl/ops/sglang/bi_families_v2.py b/src/xorl/ops/sglang/bi_families_v2.py new file mode 100644 index 00000000..525e6c87 --- /dev/null +++ b/src/xorl/ops/sglang/bi_families_v2.py @@ -0,0 +1,1063 @@ +# Families v2 — the redefined frozen reduction trees for the batch-invariance +# contract (hidden-dim RMSNorm, qk-norm, and the final projection). +# +# The trainer and serving engine carry equivalent implementations because they +# must evaluate the same reduction trees. This module remains self-contained +# (torch + triton only), and cross-engine agreement is established by the +# behavior-logprob replay rather than source-file identity. +# +# Rule A: every bit-relevant reduction is written explicitly — an +# adjacent-pairwise balanced binary tree within a block (tl.split + one add +# per level; a 2-element reduction has exactly one association, so the +# compiler owns no tree choice) and a sequential scalar chain across chunks +# in index order. tl.sum over >2 elements is banned in bit-relevant positions. +# Rule B: golden-value gates pin the bits under both engines' venvs. +# +# These trees are DEFAULT ON inside an engaged contract lane. The exact +# Qwen3.5-family model setup selects the already-certified v1 family directly; +# callers do not coordinate the selection with an environment variable. +# +# Migration: v1 and v2 are two different trees, and both hold the trainer and +# the sampler bitwise equal — that is the contract, and v2 satisfies it. The +# reported values do move between the two, so goldens and frozen anchors +# recorded under v1 must be re-taken under v2, and both engines must flip +# together. + +import os + +import torch +import triton +import triton.language as tl +from triton.runtime.errors import OutOfResources + + +_EXACT_FAMILIES_VERSION: str | None = None + + +def _select_qwen35_families_v1() -> None: + """Pin exact Qwen's fused LM-head loss to its qualified v1 program.""" + global _EXACT_FAMILIES_VERSION + _EXACT_FAMILIES_VERSION = "v1" + + +def _select_glm52_families_v2() -> None: + """Pin the GLM-5.2 process to its certified v2 reduction trees.""" + global _EXACT_FAMILIES_VERSION + _EXACT_FAMILIES_VERSION = "v2" + + +def _select_qwen3_dense_families_v2() -> None: + """Pin exact dense Qwen3 to the shared v2 norm and lm-head trees.""" + global _EXACT_FAMILIES_VERSION + _EXACT_FAMILIES_VERSION = "v2" + + +def _select_nonexact_families() -> None: + """Restore the pre-existing family selection for an ordinary model.""" + global _EXACT_FAMILIES_VERSION + _EXACT_FAMILIES_VERSION = None + + +def families_v2_enabled() -> bool: + """Return the selected reduction family for the current process. + + Exact model programs select their family structurally and ignore the + legacy rollback variables. Without an exact model selection, preserve the + pre-existing non-exact behavior for compatibility. + """ + if _EXACT_FAMILIES_VERSION is not None: + return _EXACT_FAMILIES_VERSION == "v2" + return not any(os.getenv(v, "1").lower() in _V2_OFF for v in FAMILIES_V2_ENV_VARS) + + +# Contract constants (bit-relevant; never tuning axes). +V2_NORM_BLOCK_H = 4096 # per-chunk tree width for hidden-dim norms +V2_QK_MAX_HEAD_DIM = 256 + +FAMILIES_V2_ENV_VARS = ("XORL_FAMILIES_V2", "SGLANG_FAMILIES_V2") +_V2_OFF = ("0", "false", "no") + + +@triton.jit +def _rtne_bf16(x): + # Round-to-nearest-even to bf16 via integer bitcast (triton folds + # f32->bf16->f32 cast pairs; this cannot be folded). NaNs quieted. + bits = x.to(tl.int32, bitcast=True) + rounded = (bits + 0x7FFF + ((bits >> 16) & 1)) & -65536 + nan_bits = (bits & -65536) | 0x00400000 + out = tl.where(x != x, nan_bits, rounded) + return out.to(tl.float32, bitcast=True) + + +@triton.jit +def _pairwise_tree_sum(vec, BLOCK: tl.constexpr): + # Adjacent-pairwise balanced binary tree: log2(BLOCK) levels, one explicit + # fp32 add per level. The reduction ORDER is written in the IR — no + # compiler-owned tl.sum lowering (the v1 1-ulp drift class). Statically + # unrolled; the constexpr `if` prunes levels past log2(BLOCK). + tl.static_assert(BLOCK & (BLOCK - 1) == 0, "BLOCK must be a power of 2") + tl.static_assert(BLOCK <= 4096, "12 unrolled levels cover BLOCK <= 4096") + for _ in tl.static_range(0, 12): + if vec.shape[0] > 1: + lo, hi = tl.split(tl.reshape(vec, (vec.shape[0] // 2, 2))) + vec = lo + hi + return vec # shape (1,) fp32 + + +V2_NORM_TILE = 512 # register tile; BITS-NEUTRAL by the tree factorization below + + +@triton.jit +def _rms_norm_v2_kernel( + x_ptr, + res_ptr, + w_ptr, + out_ptr, + res_out_ptr, + n_cols, + stride_x, + stride_res, + stride_out, + stride_res_out, + eps, + HAS_RESIDUAL: tl.constexpr, + ZERO_CENTERED: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Family-2'/1' unified row norm: ONE launch per call site. + + Pinned order: (residual add -> RTNE bf16 round ->) fp32 square -> + per-chunk pairwise tree -> sequential chunk chain -> /H -> tl.rsqrt -> + x * inv_rms * w (left-to-right, fp32 weight-mul) -> single cast at store. + Tree is a function of n_cols alone => batch-invariant by construction. + + For the residual specialization, this fused one-launch form serves nearly + every shipped shape. Few rows over many tiles, and the compiler-spilling + no-residual specialization, dispatch to the bit-identical split realization + below (see ``_v2_norm_use_split``), which factorizes the same tree into + per-TILE trees + a partials tree (adjacent pairing preserves contiguity at + every level => identical bits; cross-structure gates + frozen goldens). + """ + row = tl.program_id(0) + n_chunks = tl.cdiv(n_cols, BLOCK_H) + + total = tl.zeros((1,), dtype=tl.float32) + for c in range(n_chunks): + cols = c * BLOCK_H + tl.arange(0, BLOCK_H) + mask = cols < n_cols + x = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + r = tl.load(res_ptr + row * stride_res + cols, mask=mask, other=0.0).to(tl.float32) + s = _rtne_bf16(x + r) # v1 semantic kept: add rounds to input dtype + tl.store( + res_out_ptr + row * stride_res_out + cols, + s.to(res_out_ptr.dtype.element_ty), + mask=mask, + ) + else: + s = x + sq = s * s # masked lanes contribute +0.0 (exact for sums of squares) + total = total + _pairwise_tree_sum(sq, BLOCK_H) + + var = total / n_cols.to(tl.float32) + inv_rms = tl.rsqrt(var + eps) + + for c in range(n_chunks): + cols = c * BLOCK_H + tl.arange(0, BLOCK_H) + mask = cols < n_cols + if HAS_RESIDUAL: + s = tl.load(res_out_ptr + row * stride_res_out + cols, mask=mask, other=0.0).to( + tl.float32 + ) # the rounded sum: bit-identical to what was squared + else: + s = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32) + w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32) + if ZERO_CENTERED: + w = 1.0 + w + y = s * inv_rms * w + tl.store(out_ptr + row * stride_out + cols, y.to(out_ptr.dtype.element_ty), mask=mask) + + +QK_V2_ROWS_PER_PROG = 16 # head-rows per program (perf-only, NOT bit-relevant) + + +@triton.jit +def _qk_norm_v2_kernel( + x_ptr, + w_ptr, + out_ptr, + n_rows, + n_heads, + head_dim, + stride_x_tok, + stride_x_head, + stride_out_tok, + stride_out_head, + eps, + ZERO_CENTERED: tl.constexpr, + BLOCK_D: tl.constexpr, + ROWS: tl.constexpr, +): + """Family-1' strided qk-norm: reads head rows straight out of the packed + qkv projection (no reshape/contiguous copies). Same tree as + _rms_norm_v2_kernel (single chunk: head_dim <= V2_QK_MAX_HEAD_DIM). + Each program handles ROWS independent head-rows (row batching is grid + shape only — per-row math identical, like BLOCK_M in the GEMM). + In-place safe per row: the full head row is loaded before any store. + """ + pid = tl.program_id(0) + rows = pid * ROWS + tl.arange(0, ROWS) + row_mask = rows < n_rows + rows_safe = tl.where(row_mask, rows, 0) + tok = rows_safe // n_heads + head = rows_safe % n_heads + d = tl.arange(0, BLOCK_D) + col_mask = d < head_dim + mask = row_mask[:, None] & col_mask[None, :] + base = tok * stride_x_tok + head * stride_x_head + x = tl.load(x_ptr + base[:, None] + d[None, :], mask=mask, other=0.0).to(tl.float32) + total = _pairwise_tree_sum_rows(x * x, BLOCK_D) + var = total / head_dim.to(tl.float32) + inv_rms = tl.rsqrt(var + eps) + w = tl.load(w_ptr + d, mask=col_mask, other=0.0).to(tl.float32) + if ZERO_CENTERED: + w = 1.0 + w + y = x * inv_rms[:, None] * w[None, :] + out_base = tok * stride_out_tok + head * stride_out_head + tl.store( + out_ptr + out_base[:, None] + d[None, :], + y.to(out_ptr.dtype.element_ty), + mask=mask, + ) + + +def rms_norm_v2( + x: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + *, + residual: torch.Tensor | None = None, + zero_centered: bool = False, +): + """v2 hidden-dim RMSNorm (family-2' residual / no-residual, family-1' + site-classes unified on one tree). Returns ``out`` or ``(out, residual_out)``. + + bf16-only (the contract dtype); 2D ``[rows, H]`` input, unit column stride. + """ + assert x.ndim == 2 and x.stride(1) == 1, "x must be [rows, H], unit col stride" + assert x.dtype == torch.bfloat16, "families v2 is bf16-only (contract dtype)" + assert weight.ndim == 1 and weight.shape[0] == x.shape[1] + assert x.is_cuda + rows, H = x.shape + weight = weight.contiguous() + if residual is not None: + assert residual.shape == x.shape and residual.stride(1) == 1 + assert residual.dtype == torch.bfloat16 + if _v2_norm_use_split( + rows, + triton.cdiv(H, V2_NORM_TILE), + has_residual=residual is not None, + is_hopper=torch.cuda.get_device_capability(x.device) == (9, 0), + ): + return _rms_norm_v2_split(x, weight, eps, residual, zero_centered) + return _rms_norm_v2_fused(x, weight, eps, residual, zero_centered) + + +def _rms_norm_v2_fused(x, weight, eps, residual, zero_centered): + """One-launch realization of the v2 norm tree (see ``_rms_norm_v2_kernel``).""" + rows, H = x.shape + out = torch.empty_like(x) + if residual is not None: + res_out = torch.empty_like(x) + _rms_norm_v2_kernel[(rows,)]( + x, + residual, + weight, + out, + res_out, + H, + x.stride(0), + residual.stride(0), + out.stride(0), + res_out.stride(0), + eps, + HAS_RESIDUAL=True, + ZERO_CENTERED=zero_centered, + BLOCK_H=V2_NORM_BLOCK_H, + ) + return out, res_out + _rms_norm_v2_kernel[(rows,)]( + x, + x, + weight, + out, + out, + H, + x.stride(0), + 0, + out.stride(0), + 0, + eps, + HAS_RESIDUAL=False, + ZERO_CENTERED=zero_centered, + BLOCK_H=V2_NORM_BLOCK_H, + ) + return out + + +def qk_norm_v2( + x: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + *, + head_dim: int, + out: torch.Tensor | None = None, + zero_centered: bool = False, +): + """v2 per-head qk-norm over strided input. + + ``x``: ``[T, n_heads * head_dim]`` view into the packed qkv output — any + row stride, unit element stride, heads contiguous within a row. ``out`` + defaults to a fresh tensor (trainer); pass ``out=x`` for in-place (serving). + """ + assert x.ndim == 2 and x.stride(1) == 1 + assert x.dtype == torch.bfloat16, "families v2 is bf16-only (contract dtype)" + assert x.shape[1] % head_dim == 0 + assert head_dim <= V2_QK_MAX_HEAD_DIM + assert weight.ndim == 1 and weight.shape[0] == head_dim + assert x.is_cuda + T = x.shape[0] + n_heads = x.shape[1] // head_dim + weight = weight.contiguous() + if out is None: + out = torch.empty_like(x) + else: + assert out.shape == x.shape and out.stride(1) == 1 and out.dtype == x.dtype + if T > 0: + n_rows = T * n_heads + _qk_norm_v2_kernel[(triton.cdiv(n_rows, QK_V2_ROWS_PER_PROG),)]( + x, + weight, + out, + n_rows, + n_heads, + head_dim, + x.stride(0), + head_dim, + out.stride(0), + head_dim, + eps, + ZERO_CENTERED=zero_centered, + BLOCK_D=triton.next_power_of_2(head_dim), + ROWS=QK_V2_ROWS_PER_PROG, + ) + return out + + +# --------------------------------------------------------------------------- +# Head v2 — online-LSE lm-head (component 3, design note §5) +# +# The GEMM K-chain is v1's matmul_kernel_persistent VERBATIM (pinned +# BLOCK_SIZE_K; logits bitwise == the v1 head GEMM — gated). New: per-tile +# (m, l) stats computed in the epilogue from the fp32 accumulator, one pinned +# merge kernel. BLOCK_SIZE_N tiles a stats reduction, so it is a CONTRACT +# CONSTANT here (HEAD_STATS_TILE_N) — never shape/M-keyed, or decode and +# trainer scoring rows would grow different LSE trees. The OOM fallback chain +# varies only bit-neutral axes (BLOCK_M / GROUP / stages / warps). +# --------------------------------------------------------------------------- + +HEAD_V2_STATS_TILE_N = 256 # contract constant (bit-relevant) +HEAD_V2_BLOCK_K = 64 # == v1 PINNED_BLOCK_K[bf16] (bit-relevant) + +# Bit-neutral launch axes, M-bucketed (perf only; all candidates share N/K). +_HEAD_V2_LAUNCH = ( + (16, {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 4}), + (256, {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 4}), + (None, {"BLOCK_SIZE_M": 128, "GROUP_SIZE_M": 8, "num_stages": 3, "num_warps": 8}), +) +_HEAD_V2_FALLBACK = { + "BLOCK_SIZE_M": 64, + "GROUP_SIZE_M": 8, + "num_stages": 2, + "num_warps": 4, +} + + +@triton.jit +def _pairwise_tree_sum_rows(mat, BLOCK_N: tl.constexpr): + # Row-wise adjacent-pairwise tree over axis 1 (same construction as + # _pairwise_tree_sum; each level is one elementwise add — no compiler + # tree choice, warp count is bit-neutral by construction). + tl.static_assert(BLOCK_N & (BLOCK_N - 1) == 0, "BLOCK_N must be a power of 2") + tl.static_assert(BLOCK_N <= 4096, "12 unrolled levels cover BLOCK_N <= 4096") + for _ in tl.static_range(0, 12): + if mat.shape[1] > 1: + lo, hi = tl.split(tl.reshape(mat, (mat.shape[0], mat.shape[1] // 2, 2))) + mat = lo + hi + return tl.reshape(mat, (mat.shape[0],)) + + +@triton.jit +def _head_v2_gemm_stats_kernel( + a_ptr, + b_ptr, + c_ptr, + m_out_ptr, + l_out_ptr, + sel_ptr, + tok_ptr, + temp_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_mm, + stride_lm, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, + A_LARGE: tl.constexpr, + B_LARGE: tl.constexpr, + C_LARGE: tl.constexpr, + STORE_LOGITS: tl.constexpr, + HAS_TOKS: tl.constexpr, + HAS_TEMP: tl.constexpr, +): + """v1 persistent GEMM (identical K-chain) + per-tile (m, l) epilogue stats. + + Stats are computed from the fp32 accumulator in registers: m = row max of + the tile (order-free), l = row sum of exp(x - m) via the pairwise tree. + Temperature scales stats (and the gathered selected logit) but NOT the + stored logits — matching v1, where the stats kernel applies 1/T while the + decode logits buffer stays unscaled for the sampler. + """ + start_pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + num_tiles = num_pid_m * num_pid_n + + offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + + start_m = pid_m * BLOCK_SIZE_M + start_n = pid_n * BLOCK_SIZE_N + offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) + offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) + if A_LARGE: + offs_am = offs_am.to(tl.int64) + if B_LARGE: + offs_bn = offs_bn.to(tl.int64) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for ki in range(k_tiles): + if A_LARGE or B_LARGE: + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + else: + offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + a = tl.load(a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0) + accumulator = tl.dot(a, b, accumulator) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + if C_LARGE: + offs_cm = offs_cm.to(tl.int64) + offs_cn = offs_cn.to(tl.int64) + row_mask = offs_cm < M + col_mask = offs_cn < N + + if STORE_LOGITS: + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + tl.store( + c_ptrs, + accumulator.to(tl.float32), + mask=row_mask[:, None] & col_mask[None, :], + ) + + # --- epilogue stats (the new frozen tree) --- + x_stats = accumulator + if HAS_TEMP: + temp = tl.load(temp_ptr + offs_cm, mask=row_mask, other=1.0) + inv_t = 1.0 / temp + x_stats = x_stats * inv_t[:, None] + neg_inf = float("-inf") + m_tile = tl.max(tl.where(col_mask[None, :], x_stats, neg_inf), axis=1) + terms = tl.where(col_mask[None, :], tl.exp(x_stats - m_tile[:, None]), 0.0) + l_tile = _pairwise_tree_sum_rows(terms, BLOCK_SIZE_N) + tl.store(m_out_ptr + offs_cm * stride_mm + pid_n, m_tile, mask=row_mask) + tl.store(l_out_ptr + offs_cm * stride_lm + pid_n, l_tile, mask=row_mask) + + if HAS_TOKS: + toks = tl.load(tok_ptr + offs_cm, mask=row_mask, other=-1) + tok_here = (toks >= start_n) & (toks < start_n + BLOCK_SIZE_N) & row_mask + # -inf-padded max gather: exactly one lane matches, sign of zero preserved + sel = tl.max(tl.where(offs_cn[None, :] == toks[:, None], x_stats, neg_inf), axis=1) + tl.store(sel_ptr + offs_cm, sel, mask=tok_here) + + +@triton.jit +def _head_v2_lse_merge_kernel( + m_ptr, + l_ptr, + sel_ptr, + lse_ptr, + lp_ptr, + n_tiles, + stride_mm, + stride_lm, + HAS_SEL: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """Pinned merge: exact global max over tile maxima (order-free), rescaled + sumexp via the pairwise tree over tiles (padded l=0 / m=-inf), then + lse = gmax + log(acc); logprob = min(sel - lse, 0) (the p~1 boundary + clamp, kept from v1).""" + row = tl.program_id(0) + t = tl.arange(0, BLOCK_T) + mask = t < n_tiles + neg_inf = float("-inf") + m_t = tl.load(m_ptr + row * stride_mm + t, mask=mask, other=neg_inf) + l_t = tl.load(l_ptr + row * stride_lm + t, mask=mask, other=0.0) + gmax = tl.max(m_t) + terms = tl.where(mask, l_t * tl.exp(m_t - gmax), 0.0) + acc = _pairwise_tree_sum(terms, BLOCK_T) + lse = gmax + tl.log(acc) + tl.store(lse_ptr + row + tl.arange(0, 1), lse) + if HAS_SEL: + sel = tl.load(sel_ptr + row) + lp = tl.minimum(sel - lse, 0.0) + tl.store(lp_ptr + row + tl.arange(0, 1), lp) + + +def _head_v2_launch(hidden, weight_t, logits, m_buf, l_buf, sel, toks, temp): + M, K = hidden.shape + N = weight_t.shape[1] + NUM_SMS = torch.cuda.get_device_properties(hidden.device).multi_processor_count + a_large = hidden.numel() > 2**31 + b_large = weight_t.numel() > 2**31 + c_large = logits is not None and (logits.shape[0] * logits.stride(0) + logits.shape[1] * logits.stride(1) > 2**31) + launch = _HEAD_V2_FALLBACK + for bound, cfg in _HEAD_V2_LAUNCH: + if bound is None or M <= bound: + launch = cfg + break + + def _run(cfg): + grid = min( + NUM_SMS, + triton.cdiv(M, cfg["BLOCK_SIZE_M"]) * triton.cdiv(N, HEAD_V2_STATS_TILE_N), + ) + _head_v2_gemm_stats_kernel[(grid,)]( + hidden, + weight_t, + logits if logits is not None else hidden, + m_buf, + l_buf, + sel if sel is not None else m_buf, + toks if toks is not None else m_buf, + temp if temp is not None else m_buf, + M, + N, + K, + hidden.stride(0), + hidden.stride(1), + weight_t.stride(0), + weight_t.stride(1), + logits.stride(0) if logits is not None else 0, + logits.stride(1) if logits is not None else 0, + m_buf.stride(0), + l_buf.stride(0), + BLOCK_SIZE_M=cfg["BLOCK_SIZE_M"], + BLOCK_SIZE_N=HEAD_V2_STATS_TILE_N, + BLOCK_SIZE_K=HEAD_V2_BLOCK_K, + GROUP_SIZE_M=cfg["GROUP_SIZE_M"], + NUM_SMS=NUM_SMS, + A_LARGE=a_large, + B_LARGE=b_large, + C_LARGE=c_large, + STORE_LOGITS=logits is not None, + HAS_TOKS=toks is not None, + HAS_TEMP=temp is not None, + num_stages=cfg["num_stages"], + num_warps=cfg["num_warps"], + ) + + try: + _run(launch) + except OutOfResources: + _run(_HEAD_V2_FALLBACK) # bit-identical: only bit-neutral axes differ + + +def _head_v2_merge(m_buf, l_buf, sel): + M, n_tiles = m_buf.shape + lse = torch.empty((M,), dtype=torch.float32, device=m_buf.device) + lp = torch.empty((M,), dtype=torch.float32, device=m_buf.device) if sel is not None else None + _head_v2_lse_merge_kernel[(M,)]( + m_buf, + l_buf, + sel if sel is not None else m_buf, + lse, + lp if lp is not None else lse, + n_tiles, + m_buf.stride(0), + l_buf.stride(0), + HAS_SEL=sel is not None, + BLOCK_T=max(triton.next_power_of_2(n_tiles), 2), + ) + return lse, lp + + +def _head_v2_check_inputs(hidden, weight, token_ids, temperature): + assert hidden.ndim == 2 and weight.ndim == 2 and hidden.shape[1] == weight.shape[1] + assert hidden.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( + "the head contract takes bf16 hidden/weight (fp32 upcast is exact inside the GEMM)" + ) + assert hidden.is_cuda + hidden = hidden.contiguous() + M = hidden.shape[0] + if token_ids is not None: + token_ids = token_ids.contiguous().to(device=hidden.device, dtype=torch.int64) + assert token_ids.shape == (M,) + if temperature is not None: + temperature = temperature.reshape(-1).to(device=hidden.device, dtype=torch.float32).contiguous() + assert temperature.shape == (M,) + torch._assert_async((temperature > 0).all(), "temperature must be > 0") + return hidden, token_ids, temperature + + +def head_v2_selected_logprob( + hidden: torch.Tensor, + weight: torch.Tensor, + token_ids: torch.Tensor, + temperature: torch.Tensor | None = None, +): + """Scoring path (trainer bi_fused v2 / serving prefill input logprobs): + logits are NEVER materialized — per-tile stats + selected gather in the + GEMM epilogue, one merge launch. Returns ``(logprob, lse, selected)``, + all ``[N]`` fp32 (temperature-scaled when given), like v1.""" + hidden, token_ids, temperature = _head_v2_check_inputs(hidden, weight, token_ids, temperature) + M = hidden.shape[0] + V = weight.shape[0] + n_tiles = triton.cdiv(V, HEAD_V2_STATS_TILE_N) + if M == 0: + z = torch.empty((0,), dtype=torch.float32, device=hidden.device) + return z, z.clone(), z.clone() + m_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=hidden.device) + l_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=hidden.device) + sel = torch.empty((M,), dtype=torch.float32, device=hidden.device) + _head_v2_launch(hidden, weight.t(), None, m_buf, l_buf, sel, token_ids, temperature) + lse, lp = _head_v2_merge(m_buf, l_buf, sel) + return lp, lse, sel + + +def head_v2_full_logits_with_lse( + hidden: torch.Tensor, + weight: torch.Tensor, + temperature: torch.Tensor | None = None, +): + """Decode path: full ``[N, V]`` fp32 logits for the sampler (bitwise == + the v1 head GEMM — same kernel body, same pinned K-chain) PLUS the lse + from the epilogue stats. The sampled token's logprob afterwards is + ``clamp_max(logits[row, tok] - lse[row], 0.0)`` — elementwise aten, no + second stats pass, no BI log_softmax.""" + hidden, _, temperature = _head_v2_check_inputs(hidden, weight, None, temperature) + M = hidden.shape[0] + V = weight.shape[0] + n_tiles = triton.cdiv(V, HEAD_V2_STATS_TILE_N) + logits = torch.empty((M, V), dtype=torch.float32, device=hidden.device) + if M == 0: + return logits, torch.empty((0,), dtype=torch.float32, device=hidden.device) + m_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=hidden.device) + l_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=hidden.device) + _head_v2_launch(hidden, weight.t(), logits, m_buf, l_buf, None, None, temperature) + lse, _ = _head_v2_merge(m_buf, l_buf, None) + return logits, lse + + +@triton.jit +def _exact_temperature_scale_fp32_kernel( + logits_ptr, + temperature_ptr, + output_ptr, + n_cols, + logits_row_stride, + output_row_stride, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + mask = cols < n_cols + values = tl.load( + logits_ptr + row * logits_row_stride + cols, + mask=mask, + other=0.0, + ) + inv_t = 1.0 / tl.load(temperature_ptr + row) + values = values * inv_t + tl.store(output_ptr + row * output_row_stride + cols, values, mask=mask) + + +def exact_temperature_scale_fp32_logits( + logits: torch.Tensor, + temperature: torch.Tensor, +) -> torch.Tensor: + """Materialize the exact per-row FP32 sampling transform.""" + + assert logits.ndim == 2 and logits.dtype == torch.float32 + assert logits.stride(1) == 1, "logits rows must be unit-stride" + assert temperature.dtype == torch.float32 + assert temperature.device == logits.device + assert temperature.is_contiguous() + temperature = temperature.reshape(-1) + assert temperature.shape == (logits.shape[0],) + torch._assert_async( + (torch.isfinite(temperature) & (temperature > 0)).all(), + "temperature must contain finite values > 0", + ) + + output = torch.empty_like(logits) + if logits.numel() == 0: + return output + if logits.is_cuda: + block_n = 1024 + _exact_temperature_scale_fp32_kernel[(logits.shape[0], triton.cdiv(logits.shape[1], block_n))]( + logits, + temperature, + output, + logits.shape[1], + logits.stride(0), + output.stride(0), + BLOCK_N=block_n, + ) + else: + output.copy_(logits * (1.0 / temperature).unsqueeze(1)) + return output + + +def exact_temperature_scale_bf16_logits( + logits: torch.Tensor, + temperature: torch.Tensor | None, +) -> torch.Tensor: + """Apply DSV4's BF16 divide/store temperature boundary.""" + + assert logits.ndim == 2 and logits.dtype == torch.bfloat16 + assert logits.stride(1) == 1, "logits rows must be unit-stride" + if temperature is None: + return logits + assert temperature.dtype == torch.float32 + assert temperature.device == logits.device + assert temperature.is_contiguous() + temperature = temperature.reshape(-1) + assert temperature.shape == (logits.shape[0],) + torch._assert_async( + (torch.isfinite(temperature) & (temperature > 0)).all(), + "temperature must contain finite values > 0", + ) + return logits.bfloat16().div(temperature.unsqueeze(1)).bfloat16() + + +@triton.jit +def _head_v2_stats_from_logits_kernel( + logits_ptr, + m_out_ptr, + l_out_ptr, + sel_ptr, + tok_ptr, + temp_ptr, + N, + n_tiles, + stride_lm, + stride_mm, + stride_lm_out, + HAS_TOKS: tl.constexpr, + HAS_TEMP: tl.constexpr, + TILE_N: tl.constexpr, +): + """Per-tile (m, l) stats over MATERIALIZED fp32 logits — the same tree as + the epilogue stats (same TILE_N, same max/exp/pairwise-tree expressions on + the same fp32 values), for callers that must sample from the logits they + score (the decode rescore hook). Bitwise == the epilogue path.""" + pid = tl.program_id(0) + row = pid // n_tiles + tile = pid % n_tiles + start_n = tile * TILE_N + offs_n = start_n + tl.arange(0, TILE_N) + col_mask = offs_n < N + x = tl.load(logits_ptr + row * stride_lm + offs_n, mask=col_mask, other=0.0) + if HAS_TEMP: + inv_t = 1.0 / tl.load(temp_ptr + row) + x = x * inv_t + neg_inf = float("-inf") + m_tile = tl.max(tl.where(col_mask, x, neg_inf)) + terms = tl.where(col_mask, tl.exp(x - m_tile), 0.0) + l_tile = _pairwise_tree_sum(terms, TILE_N) + tl.store( + m_out_ptr + row * stride_mm + tile + tl.arange(0, 1), + m_tile + tl.zeros((1,), dtype=tl.float32), + ) + tl.store(l_out_ptr + row * stride_lm_out + tile + tl.arange(0, 1), l_tile) + if HAS_TOKS: + tok = tl.load(tok_ptr + row) + if (tok >= start_n) & (tok < start_n + TILE_N): + sel = tl.max(tl.where(offs_n == tok, x, neg_inf)) + tl.store(sel_ptr + row, sel) + + +def head_v2_selected_logprob_from_logits( + logits: torch.Tensor, + token_ids: torch.Tensor, + temperature: torch.Tensor | None = None, +): + """Head-v2 twin of the v1 from-logits rescore: same (m, l) tree over an + existing fp32 logits tensor, so with logits from + ``head_v2_full_logits_with_lse`` (or the bitwise-equal v1 GEMM) the result + is bitwise identical to ``head_v2_selected_logprob`` on the same + hidden/weight/temperature. Returns ``(logprob, lse, selected)``.""" + assert logits.ndim == 2 and logits.dtype == torch.float32 and logits.is_cuda + assert logits.stride(1) == 1, "logits rows must be unit-stride" + M, V = logits.shape + token_ids = token_ids.contiguous().to(device=logits.device, dtype=torch.int64) + assert token_ids.shape == (M,) + if temperature is not None: + temperature = temperature.reshape(-1).to(device=logits.device, dtype=torch.float32).contiguous() + assert temperature.shape == (M,) + torch._assert_async((temperature > 0).all(), "temperature must be > 0") + n_tiles = triton.cdiv(V, HEAD_V2_STATS_TILE_N) + if M == 0: + z = torch.empty((0,), dtype=torch.float32, device=logits.device) + return z, z.clone(), z.clone() + m_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=logits.device) + l_buf = torch.empty((M, n_tiles), dtype=torch.float32, device=logits.device) + sel = torch.empty((M,), dtype=torch.float32, device=logits.device) + _head_v2_stats_from_logits_kernel[(M * n_tiles,)]( + logits, + m_buf, + l_buf, + sel, + token_ids, + temperature if temperature is not None else m_buf, + V, + n_tiles, + logits.stride(0), + m_buf.stride(0), + l_buf.stride(0), + HAS_TOKS=True, + HAS_TEMP=temperature is not None, + TILE_N=HEAD_V2_STATS_TILE_N, + ) + lse, lp = _head_v2_merge(m_buf, l_buf, sel) + return lp, lse, sel + + +# --- split realization of the SAME norm tree (structure is bit-neutral) ----- +# +# The split realization computes the IDENTICAL tree — per-tile partial trees in +# a 2D-parallel kernel, the pinned combine in a per-row kernel, elementwise +# normalize in a 2D-parallel kernel — with fp32 partials stored/loaded exactly +# and the bf16 residual round-trip identical to the fused kernel's own res_out +# reload. Bit-equality fused==split is gated by the families-v2 norm tests, +# which force each realization explicitly rather than relying on the rule +# below to select one, and both realizations check the same frozen goldens. +# +# It pays three launches and an HBM round-trip for the partials, and buys +# rows*n_tiles-way parallelism where the fused grid=(rows,) has only rows-way. +# For the residual specialization, that trade pays at few rows over many tiles. +# The no-residual fused specialization currently spills on Hopper, so it always +# uses the split form (see _v2_norm_use_split for the measured boundary). + +V2_NORM_SPLIT_MIN_TILES = 10 # no-residual crossover on non-Hopper devices +V2_NORM_RESIDUAL_SPLIT_MIN_TILES = 13 # keeps all shipped hidden sizes (at most 12 tiles) fused + + +def _v2_norm_use_split( + rows: int, + n_tiles: int, + *, + has_residual: bool = True, + is_hopper: bool = False, +) -> bool: + """Structure switch (perf-only, NOT bit-relevant). + + The no-residual fused specialization spills on Hopper, while the split + realization is bit-identical and avoids that register-pressure cliff. + Other architectures and the residual specialization retain the measured + few-rows/many-tiles switch, where the fused ``grid=(rows,)`` cannot fill + the GPU. + + ``n_tiles`` must be the split kernel's 512-wide tile count, not the fused + kernel's 4096-wide chunk count. The threshold is a Hopper performance + policy. Both realizations compute the same tree, so retuning it may change + speed only; the fused-versus-split equality gate must still pass. + """ + if is_hopper and not has_residual: + return True + min_tiles = V2_NORM_RESIDUAL_SPLIT_MIN_TILES if has_residual else V2_NORM_SPLIT_MIN_TILES + return n_tiles >= min_tiles and rows <= n_tiles + + +@triton.jit +def _rms_norm_v2_partials_kernel( + x_ptr, + res_ptr, + res_out_ptr, + part_ptr, + n_cols, + n_tiles, + stride_x, + stride_res, + stride_res_out, + HAS_RESIDUAL: tl.constexpr, + TILE: tl.constexpr, +): + pid = tl.program_id(0) + row = pid // n_tiles + t = pid % n_tiles + cols = t * TILE + tl.arange(0, TILE) + mask = cols < n_cols + x = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + r = tl.load(res_ptr + row * stride_res + cols, mask=mask, other=0.0).to(tl.float32) + s = _rtne_bf16(x + r) + tl.store( + res_out_ptr + row * stride_res_out + cols, + s.to(res_out_ptr.dtype.element_ty), + mask=mask, + ) + else: + s = x + p = _pairwise_tree_sum(s * s, TILE) + tl.store(part_ptr + row * n_tiles + t + tl.arange(0, 1), p) + + +@triton.jit +def _rms_norm_v2_invrms_kernel( + part_ptr, + invrms_ptr, + n_cols, + n_tiles, + eps, + CHUNK_TILES: tl.constexpr, + BLOCK_T: tl.constexpr, +): + # combine = per-BLOCK_H-chunk pairwise tree over its CHUNK_TILES partials, + # then the sequential cross-chunk chain — identical to the fused kernel. + row = tl.program_id(0) + n_chunks = tl.cdiv(n_tiles, CHUNK_TILES) + total = tl.zeros((1,), dtype=tl.float32) + for c in range(n_chunks): + idx = c * CHUNK_TILES + tl.arange(0, CHUNK_TILES) + p = tl.load(part_ptr + row * n_tiles + idx, mask=idx < n_tiles, other=0.0) + total = total + _pairwise_tree_sum(p, CHUNK_TILES) + var = total / n_cols.to(tl.float32) + inv_rms = tl.rsqrt(var + eps) + tl.store(invrms_ptr + row + tl.arange(0, 1), inv_rms) + + +@triton.jit +def _rms_norm_v2_normalize_kernel( + s_ptr, + w_ptr, + invrms_ptr, + out_ptr, + n_cols, + n_tiles, + stride_s, + stride_out, + ZERO_CENTERED: tl.constexpr, + TILE: tl.constexpr, +): + pid = tl.program_id(0) + row = pid // n_tiles + t = pid % n_tiles + cols = t * TILE + tl.arange(0, TILE) + mask = cols < n_cols + s = tl.load(s_ptr + row * stride_s + cols, mask=mask, other=0.0).to(tl.float32) + w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32) + if ZERO_CENTERED: + w = 1.0 + w + inv_rms = tl.load(invrms_ptr + row) + y = s * inv_rms * w + tl.store(out_ptr + row * stride_out + cols, y.to(out_ptr.dtype.element_ty), mask=mask) + + +def _rms_norm_v2_split(x, weight, eps, residual, zero_centered): + rows, H = x.shape + n_tiles = triton.cdiv(H, V2_NORM_TILE) + out = torch.empty_like(x) + part = torch.empty((rows, n_tiles), dtype=torch.float32, device=x.device) + invrms = torch.empty((rows,), dtype=torch.float32, device=x.device) + if residual is not None: + res_out = torch.empty_like(x) + _rms_norm_v2_partials_kernel[(rows * n_tiles,)]( + x, + residual, + res_out, + part, + H, + n_tiles, + x.stride(0), + residual.stride(0), + res_out.stride(0), + HAS_RESIDUAL=True, + TILE=V2_NORM_TILE, + ) + s_buf = res_out + else: + res_out = None + _rms_norm_v2_partials_kernel[(rows * n_tiles,)]( + x, + x, + x, + part, + H, + n_tiles, + x.stride(0), + 0, + 0, + HAS_RESIDUAL=False, + TILE=V2_NORM_TILE, + ) + s_buf = x + _rms_norm_v2_invrms_kernel[(rows,)]( + part, + invrms, + H, + n_tiles, + eps, + CHUNK_TILES=V2_NORM_BLOCK_H // V2_NORM_TILE, + BLOCK_T=max(triton.next_power_of_2(n_tiles), 2), + ) + _rms_norm_v2_normalize_kernel[(rows * n_tiles,)]( + s_buf, + weight, + invrms, + out, + H, + n_tiles, + s_buf.stride(0), + out.stride(0), + ZERO_CENTERED=zero_centered, + TILE=V2_NORM_TILE, + ) + return (out, res_out) if residual is not None else out diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index d34654c5..72b3365e 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -74,8 +74,7 @@ from xorl.models.transformers.deepseek_v4.exact_contract import DSV4_FLASH_REQUIRED_TARGET_MODULES from xorl.models.transformers.glm5.index_share import IndexShareMode from xorl.models.transformers.glm5.support import glm5_default_lora_targets -from xorl.ops.batch_invariant_ops import enable_batch_invariant_mode, get_batch_invariant_ops -from xorl.ops.exact_sampling_transforms import TOP_K_ALL +from xorl.ops.exact.sampling_transforms import TOP_K_ALL from xorl.ops.loss import ( LossOutput, OPDLossMetrics, @@ -88,6 +87,7 @@ opd_vocab_parallel_loss_function, policy_loss_function, ) +from xorl.ops.sglang.batch_invariant_ops import enable_batch_invariant_mode, get_batch_invariant_ops from xorl.optim import build_optimizer from xorl.server.runner.adapters import LoRAAdapterManager from xorl.server.runner.adapters.gradient_finalizer import AdapterGradientMutationFailure diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index cf224676..268f2444 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -609,17 +609,17 @@ def build_training_model( # ------------------------------------------------------------------ # 5. Exact model contract / legacy scoped trunk contract (must precede FSDP2) # ------------------------------------------------------------------ - from xorl.ops.bi_families_v2 import _select_nonexact_families # noqa: PLC0415 + from xorl.ops.sglang.bi_families_v2 import _select_nonexact_families # noqa: PLC0415 _select_nonexact_families() if glm52_exact_forward_enabled(model.config): - from xorl.ops.bi_families_v2 import _select_glm52_families_v2 # noqa: PLC0415 + from xorl.ops.sglang.bi_families_v2 import _select_glm52_families_v2 # noqa: PLC0415 _select_glm52_families_v2() elif getattr(model.config, "_qwen3_dense_exact_contract", False): from xorl.lora.modules.base import LoraModule # noqa: PLC0415 - from xorl.ops.batch_invariant_ops import wrap_trunk_linears_batch_invariant # noqa: PLC0415 - from xorl.ops.bi_families_v2 import _select_qwen3_dense_families_v2 # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import wrap_trunk_linears_batch_invariant # noqa: PLC0415 + from xorl.ops.sglang.bi_families_v2 import _select_qwen3_dense_families_v2 # noqa: PLC0415 # Foundation construction initially wraps plain projections, but TP # unfusing and LoRA/QLoRA injection replace those module objects. The diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index a07879e9..d29a5c64 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -401,7 +401,7 @@ def _bootstrap(self) -> None: helper.enable_third_party_logging() if os.environ.get("XORL_BATCH_INVARIANT_MATMUL", "0") == "1": - from xorl.ops.batch_invariant_ops import enable_batch_invariant_mode # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import enable_batch_invariant_mode # noqa: PLC0415 enable_batch_invariant_mode() logger.info_rank0("XORL_BATCH_INVARIANT_MATMUL=1: enabled SGLang-compatible batch-invariant ops") diff --git a/tests/distributed/test_bi_trunk_linear_fsdp.py b/tests/distributed/test_bi_trunk_linear_fsdp.py index 882ba11e..22d006aa 100644 --- a/tests/distributed/test_bi_trunk_linear_fsdp.py +++ b/tests/distributed/test_bi_trunk_linear_fsdp.py @@ -21,7 +21,7 @@ import torch.nn.functional as F from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard -from xorl.ops.batch_invariant_ops import wrap_trunk_linears_batch_invariant +from xorl.ops.sglang.batch_invariant_ops import wrap_trunk_linears_batch_invariant from xorl.utils.device import get_nccl_backend diff --git a/tests/models/test_batch_invariance_dense.py b/tests/models/test_batch_invariance_dense.py index 50587064..4cbe4a32 100644 --- a/tests/models/test_batch_invariance_dense.py +++ b/tests/models/test_batch_invariance_dense.py @@ -16,7 +16,7 @@ from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3ForCausalLM -from xorl.ops.batch_invariant_ops import set_batch_invariant_mode +from xorl.ops.sglang.batch_invariant_ops import set_batch_invariant_mode requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/models/test_glm52_exact_fullparam_fp8.py b/tests/models/test_glm52_exact_fullparam_fp8.py index 19463dae..c51ac72b 100644 --- a/tests/models/test_glm52_exact_fullparam_fp8.py +++ b/tests/models/test_glm52_exact_fullparam_fp8.py @@ -587,7 +587,7 @@ def test_cuda_expert_quantization_preserves_expert_boundaries_and_fuse_after_qua @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") def test_cuda_router_forward_bytes_match_frozen_router_program_on_published_bytes() -> None: device = _hopper_or_skip() - from xorl.ops.batch_invariant_ops import bi_router_gemm + from xorl.ops.sglang.batch_invariant_ops import bi_router_gemm torch.manual_seed(4) num_experts, hidden_size = 256, 512 diff --git a/tests/models/test_glm52_exact_lm_head_qlora.py b/tests/models/test_glm52_exact_lm_head_qlora.py index bcc9e274..97dec686 100644 --- a/tests/models/test_glm52_exact_lm_head_qlora.py +++ b/tests/models/test_glm52_exact_lm_head_qlora.py @@ -23,7 +23,7 @@ _selected_logprob_reference_grad_partitioned, glm52_lm_head_shard, ) -from xorl.ops.bi_families_v2 import exact_temperature_scale_fp32_logits +from xorl.ops.sglang.bi_families_v2 import exact_temperature_scale_fp32_logits def _component(tp_rank: int = 0, tp_group=None) -> Glm52ExactTP16LmHeadSelectedLogprob: diff --git a/tests/models/test_lora_merged_forward.py b/tests/models/test_lora_merged_forward.py index 97262d45..c0b7acc2 100644 --- a/tests/models/test_lora_merged_forward.py +++ b/tests/models/test_lora_merged_forward.py @@ -482,7 +482,7 @@ def _model(self): return model def test_trunk_wrap_composition_policy(self, monkeypatch): - from xorl.ops.batch_invariant_ops import wrap_trunk_linears_batch_invariant # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import wrap_trunk_linears_batch_invariant # noqa: PLC0415 with pytest.raises(NotImplementedError, match="canonical merged-LoRA"): wrap_trunk_linears_batch_invariant(self._model()) @@ -490,7 +490,7 @@ def test_trunk_wrap_composition_policy(self, monkeypatch): self._assert_wrap_composes_with_merged_flag(monkeypatch) def _assert_wrap_composes_with_merged_flag(self, monkeypatch): - from xorl.ops.batch_invariant_ops import ( # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import ( # noqa: PLC0415 set_trunk_linear_contract, wrap_trunk_linears_batch_invariant, ) diff --git a/tests/models/test_moe_ep_native_combine.py b/tests/models/test_moe_ep_native_combine.py index e5c423d4..ed035aa8 100644 --- a/tests/models/test_moe_ep_native_combine.py +++ b/tests/models/test_moe_ep_native_combine.py @@ -270,7 +270,7 @@ def test_native_combine_captures_actual_operands(monkeypatch): """The layer-selected diagnostic hook exposes every exact-combine boundary.""" import xorl.distributed.parallel_state as parallel_state # noqa: PLC0415 import xorl.models.layers.moe.ep_native_combine as combine # noqa: PLC0415 - import xorl.ops.batch_invariant_ops as batch_invariant_ops # noqa: PLC0415 + import xorl.ops.sglang.batch_invariant_ops as batch_invariant_ops # noqa: PLC0415 class DummyParallelState: ep_enabled = True diff --git a/tests/models/test_qwen35_lora_projection_topology.py b/tests/models/test_qwen35_lora_projection_topology.py index 07720047..fb492be4 100644 --- a/tests/models/test_qwen35_lora_projection_topology.py +++ b/tests/models/test_qwen35_lora_projection_topology.py @@ -19,7 +19,7 @@ ) from xorl.models.layers.gated_deltanet import GatedDeltaNet from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeMLP -from xorl.ops.batch_invariant_ops import set_trunk_linear_contract, wrap_trunk_linears_batch_invariant +from xorl.ops.sglang.batch_invariant_ops import set_trunk_linear_contract, wrap_trunk_linears_batch_invariant RANK = 16 diff --git a/tests/models/test_qwen3_5_rmsnorm.py b/tests/models/test_qwen3_5_rmsnorm.py index 3e7d7735..8d578402 100644 --- a/tests/models/test_qwen3_5_rmsnorm.py +++ b/tests/models/test_qwen3_5_rmsnorm.py @@ -26,7 +26,7 @@ Qwen3_5MoeModel, Qwen3_5MoeRMSNorm, ) -from xorl.ops.batch_invariant_ops import rms_norm_batch_invariant, set_batch_invariant_mode +from xorl.ops.sglang.batch_invariant_ops import rms_norm_batch_invariant, set_batch_invariant_mode HIDDEN = 2048 @@ -447,7 +447,7 @@ def _assert_family2_residual_matches_serving_tree(): fast_zero_centered_batch_invariant_residual_rms_norm, native_zero_centered_rms_norm, ) - from xorl.ops.batch_invariant_ops import mean_dim # noqa: PLC0415 + from xorl.ops.sglang.batch_invariant_ops import mean_dim # noqa: PLC0415 torch.manual_seed(11) x = torch.randn(513, HIDDEN, device="cuda", dtype=torch.bfloat16) diff --git a/tests/models/test_qwen3_5_trunk_wrap.py b/tests/models/test_qwen3_5_trunk_wrap.py index bf101e8c..ef240774 100644 --- a/tests/models/test_qwen3_5_trunk_wrap.py +++ b/tests/models/test_qwen3_5_trunk_wrap.py @@ -19,7 +19,7 @@ from xorl.models.transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeForCausalLM from xorl.models.transformers.qwen3_5_shared import _apply_qwen35_gdn_exact -from xorl.ops.batch_invariant_ops import ( +from xorl.ops.sglang.batch_invariant_ops import ( is_trunk_linear_contract_enabled, set_trunk_linear_contract, wrap_trunk_linears_batch_invariant, diff --git a/tests/models/test_rmsnorm_family_contract.py b/tests/models/test_rmsnorm_family_contract.py index e4a1e2a4..e341389f 100644 --- a/tests/models/test_rmsnorm_family_contract.py +++ b/tests/models/test_rmsnorm_family_contract.py @@ -34,7 +34,7 @@ from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3DecoderLayer from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeModel -from xorl.ops.batch_invariant_ops import ( +from xorl.ops.sglang.batch_invariant_ops import ( bi_fused_add_rms_norm, bi_rms_norm, fused_add_rms_norm_batch_invariant, @@ -57,7 +57,7 @@ @pytest.fixture(autouse=True) def _pin_qualified_v1_family(): """Keep the v1-family contract independent of test collection order.""" - from xorl.ops.bi_families_v2 import _select_nonexact_families, _select_qwen35_families_v1 + from xorl.ops.sglang.bi_families_v2 import _select_nonexact_families, _select_qwen35_families_v1 _select_qwen35_families_v1() try: diff --git a/tests/models/test_rmsnorm_family_cross_engine.py b/tests/models/test_rmsnorm_family_cross_engine.py index f3953b7d..c280fee9 100644 --- a/tests/models/test_rmsnorm_family_cross_engine.py +++ b/tests/models/test_rmsnorm_family_cross_engine.py @@ -26,11 +26,11 @@ RMSNorm, fast_zero_centered_batch_invariant_rms_norm, ) -from xorl.ops.batch_invariant_ops import ( # noqa: E402 +from xorl.ops.sglang.batch_invariant_ops import ( # noqa: E402 bi_rms_norm, set_batch_invariant_mode, ) -from xorl.ops.bi_families_v2 import rms_norm_v2 as xorl_rms_norm_v2 # noqa: E402 +from xorl.ops.sglang.bi_families_v2 import rms_norm_v2 as xorl_rms_norm_v2 # noqa: E402 requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/models/test_rmsnorm_sglang_fused.py b/tests/models/test_rmsnorm_sglang_fused.py index 31ae9108..930403f7 100644 --- a/tests/models/test_rmsnorm_sglang_fused.py +++ b/tests/models/test_rmsnorm_sglang_fused.py @@ -21,7 +21,7 @@ ) from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3DecoderLayer -from xorl.ops.batch_invariant_ops import ( +from xorl.ops.sglang.batch_invariant_ops import ( fused_add_rms_norm_batch_invariant, rms_norm_batch_invariant, set_batch_invariant_mode, @@ -62,7 +62,7 @@ def _assert_sglang_fused_cpu_residual_matches_eager(): @pytest.fixture(autouse=True) def _pin_qualified_v1_family(): """This suite owns the qualified v1 fast-dispatch path.""" - from xorl.ops.bi_families_v2 import _select_nonexact_families, _select_qwen35_families_v1 + from xorl.ops.sglang.bi_families_v2 import _select_nonexact_families, _select_qwen35_families_v1 _select_qwen35_families_v1() try: diff --git a/tests/ops/test_bi_families_v2_norm.py b/tests/ops/test_bi_families_v2_norm.py index 8cd65b67..ae3f37bd 100644 --- a/tests/ops/test_bi_families_v2_norm.py +++ b/tests/ops/test_bi_families_v2_norm.py @@ -23,7 +23,7 @@ import pytest import torch -from xorl.ops.bi_families_v2 import ( +from xorl.ops.sglang.bi_families_v2 import ( V2_NORM_SPLIT_MIN_TILES, V2_NORM_TILE, qk_norm_v2, @@ -61,7 +61,7 @@ def _reference(x, weight, eps=EPS, residual=None, zero_centered=False): def _force(monkeypatch, realization): """Pin the realization so the bit gates never depend on the dispatch rule.""" - import xorl.ops.bi_families_v2 as module + import xorl.ops.sglang.bi_families_v2 as module monkeypatch.setattr(module, "_v2_norm_use_split", lambda *_a, **_k: realization == "split") @@ -168,7 +168,7 @@ def test_split_realization_is_actually_exercised(monkeypatch): Without this, a refactor that dropped the split realization entirely would leave the equality tests above passing vacuously. """ - import xorl.ops.bi_families_v2 as module + import xorl.ops.sglang.bi_families_v2 as module calls = [] original = module._rms_norm_v2_split @@ -191,7 +191,7 @@ def counting_split(*args, **kwargs): @pytest.mark.parametrize("rows", [1, 64, 256, 512, 2048]) def test_shipped_hidden_sizes_always_take_the_fused_realization(hidden_size, rows): """Common shipped hidden sizes stay below the measured split boundary.""" - import xorl.ops.bi_families_v2 as module + import xorl.ops.sglang.bi_families_v2 as module n_tiles = -(-hidden_size // V2_NORM_TILE) assert n_tiles < V2_NORM_SPLIT_MIN_TILES @@ -206,7 +206,7 @@ def test_dispatch_rule_needs_a_deep_tile_chain_and_few_rows(): ``has_residual=False`` — the residual specialization has its own, higher threshold in ``V2_NORM_RESIDUAL_SPLIT_MIN_TILES``. """ - import xorl.ops.bi_families_v2 as module + import xorl.ops.sglang.bi_families_v2 as module shallow = V2_NORM_SPLIT_MIN_TILES - 1 assert module._v2_norm_use_split(1, shallow, has_residual=False) is False @@ -220,7 +220,7 @@ def test_dispatch_rule_needs_a_deep_tile_chain_and_few_rows(): def test_dispatch_uses_the_split_kernels_tile_basis(): """The rejected rule was passed the fused kernel's 4096-wide chunk count, understating split parallelism by exactly 8x.""" - import xorl.ops.bi_families_v2 as module + import xorl.ops.sglang.bi_families_v2 as module hidden_size = 5120 split_tiles = -(-hidden_size // V2_NORM_TILE) diff --git a/tests/ops/test_bi_fused_lm_head.py b/tests/ops/test_bi_fused_lm_head.py index a4af5844..6224a703 100644 --- a/tests/ops/test_bi_fused_lm_head.py +++ b/tests/ops/test_bi_fused_lm_head.py @@ -177,8 +177,8 @@ def test_bi_fused_temperature_matches_serving_materialize_then_score(family): ) from xorl.ops import bi_families_v2 - from xorl.ops.batch_invariant_ops import bi_lm_head_full_logits - from xorl.ops.bi_families_v2 import head_v2_full_logits_with_lse + from xorl.ops.sglang.batch_invariant_ops import bi_lm_head_full_logits + from xorl.ops.sglang.bi_families_v2 import head_v2_full_logits_with_lse torch.manual_seed(53) hidden = torch.randn((4, 128), dtype=torch.bfloat16, device="cuda") @@ -210,7 +210,7 @@ def test_bi_fused_temperature_matches_serving_materialize_then_score(family): @requires_cuda @pytest.mark.gpu def test_bi_kernel_unit_temperature_is_exact_identity(): - from xorl.ops.batch_invariant_ops import bi_lm_head_selected_logprob + from xorl.ops.sglang.batch_invariant_ops import bi_lm_head_selected_logprob hidden, weight, _ = _inputs(6) hidden = hidden.view(-1, hidden.shape[-1]) @@ -228,7 +228,7 @@ def test_bi_kernel_unit_temperature_is_exact_identity(): @requires_cuda @pytest.mark.gpu def test_bi_kernel_p1_tokens_clamp_to_exact_zero(): - from xorl.ops.batch_invariant_ops import bi_lm_head_selected_logprob + from xorl.ops.sglang.batch_invariant_ops import bi_lm_head_selected_logprob torch.manual_seed(7) N, H, V = 8192, 1024, 12800 diff --git a/tests/ops/test_bi_gemm_config_table.py b/tests/ops/test_bi_gemm_config_table.py index 78c5c694..cb8f8a45 100644 --- a/tests/ops/test_bi_gemm_config_table.py +++ b/tests/ops/test_bi_gemm_config_table.py @@ -11,13 +11,13 @@ import torch import triton -from xorl.ops.batch_invariant_ops import ( +from xorl.ops.exact.bi_gemm_configs import BASELINE_CONFIG, PINNED_BLOCK_K, lookup_mm_config +from xorl.ops.sglang.batch_invariant_ops import ( _deepgemm_ready, _matmul_persistent_deepgemm, matmul_kernel_persistent, set_batch_invariant_mode, ) -from xorl.ops.exact.bi_gemm_configs import BASELINE_CONFIG, PINNED_BLOCK_K, lookup_mm_config requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/ops/test_bi_router_gemm.py b/tests/ops/test_bi_router_gemm.py index d8cf0c30..bb5ce369 100644 --- a/tests/ops/test_bi_router_gemm.py +++ b/tests/ops/test_bi_router_gemm.py @@ -2,7 +2,7 @@ import torch from xorl.models.layers.moe.moe_block import MoEBlock, _BIRouterGemm -from xorl.ops.batch_invariant_ops import bi_bf16_fp32_linear, bi_router_gemm, bi_router_topk_weights +from xorl.ops.sglang.batch_invariant_ops import bi_bf16_fp32_linear, bi_router_gemm, bi_router_topk_weights requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/tests/ops/test_bi_trunk_linear.py b/tests/ops/test_bi_trunk_linear.py index e5c844f0..24303786 100644 --- a/tests/ops/test_bi_trunk_linear.py +++ b/tests/ops/test_bi_trunk_linear.py @@ -17,7 +17,7 @@ import torch.nn.functional as F from xorl.lora.modules.linear import LoraLinear -from xorl.ops.batch_invariant_ops import ( +from xorl.ops.sglang.batch_invariant_ops import ( is_trunk_linear_contract_enabled, matmul_persistent, rms_norm_batch_invariant, diff --git a/tests/ops/test_exact_sampling_transforms.py b/tests/ops/test_exact_sampling_transforms.py index ae4fd10d..2d5833ba 100644 --- a/tests/ops/test_exact_sampling_transforms.py +++ b/tests/ops/test_exact_sampling_transforms.py @@ -5,7 +5,7 @@ import pytest import torch -from xorl.ops.exact_sampling_transforms import ( +from xorl.ops.exact.sampling_transforms import ( EXACT_FILTER_ROW_CHUNK, EXACT_SAMPLING_TRANSFORM_PROGRAM, TOP_K_ALL, diff --git a/tests/trainers/test_bi_trunk_linear_server_mode.py b/tests/trainers/test_bi_trunk_linear_server_mode.py index cc4b527e..2faaf305 100644 --- a/tests/trainers/test_bi_trunk_linear_server_mode.py +++ b/tests/trainers/test_bi_trunk_linear_server_mode.py @@ -7,7 +7,7 @@ import torch.nn as nn from xorl.ops import bi_families_v2 -from xorl.ops.batch_invariant_ops import is_trunk_linear_contract_enabled, set_trunk_linear_contract +from xorl.ops.sglang.batch_invariant_ops import is_trunk_linear_contract_enabled, set_trunk_linear_contract from xorl.trainers.model_builder import build_training_model @@ -37,7 +37,7 @@ def __init__(self): class ExactTinyTrunkModel(TinyTrunkModel): def _apply_qwen35_gdn_exact(self): - from xorl.ops.batch_invariant_ops import wrap_trunk_linears_batch_invariant + from xorl.ops.sglang.batch_invariant_ops import wrap_trunk_linears_batch_invariant return wrap_trunk_linears_batch_invariant(self) @@ -134,7 +134,7 @@ def test_glm52_selects_v2_family_structurally(monkeypatch): def test_dense_qwen_exact_program_is_reinstalled_after_lora_replacement(monkeypatch): from xorl.lora.modules.linear import LoraLinear - from xorl.ops.batch_invariant_ops import wrap_trunk_linears_batch_invariant + from xorl.ops.sglang.batch_invariant_ops import wrap_trunk_linears_batch_invariant captured = {} model = TinyTrunkModel() From abed122b62e8384fafcffd291ebda56480cd20fb Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 16:17:39 +0000 Subject: [PATCH 08/14] Add a name glossary to ops/README (#81 review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Decoder ring for the historical jargon (bi_, K3, Class-B, canonical MoE reduce, one-round SwiGLU, families v1/v2, GKN, exact) pending the rename pass discussed on #78 — several terms are user-facing config fields or vocabulary shared with the serving engine and need coordinated renames. --- src/xorl/ops/README.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index 375565c9..a4c6eefc 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -44,3 +44,20 @@ Three kinds of code live here with three different rules: New `nn.Module` layer classes (→ `models/layers/`), RL objectives (→ `xorl/objectives/`), and orchestration logic. `ops/` is for kernels and the autograd boundaries directly over them. + +## Name glossary + +Decoder ring for the historical jargon (candidates for renaming are tracked +on #78; several terms are user-facing config or shared vocabulary with the +serving engine and cannot be renamed unilaterally): + +| name | meaning | +| --- | --- | +| **batch-invariant / `bi_`** | a kernel whose per-element reduction order does not depend on batch composition, so the same token produces the same bits in any batch — the property that makes trainer/sampler logprobs comparable bitwise. | +| **K3 / "zero-K3"** | the k3 KL-divergence estimator between trainer and sampler logprobs for the same tokens; "zero-K3" = bit-identical train/serve forward, the goal of the exact contracts. | +| **Class-A / Class-B RoPE** | the two RoPE numerics classes across the trainer/sampler pair: Class A rounds to bf16 per op (8 rounding points); Class B computes one fp32 chain with a single final round (SGLang's fused CUDA rope and its compiled RL-lane path). `rope_class_b: true` selects Class B. | +| **canonical MoE reduce** | the pinned fixed-order expert-contribution reduction (contributor-leaf arithmetic + final FP64-accumulator cast) shared with serving, versioned by `CANONICAL_MOE_REDUCE_VERSION`. | +| **one-round SwiGLU** | the exact-contract SwiGLU with a single FP32 rounding point (`exact_fp32_silu_and_mul`), vs the generic fused SwiGLU. | +| **families v1 / v2** | versioned batch-invariant kernel families (norms, LM head); v2 is the epilogue-stats generation. `bi_families_v2` is the serving twin module. | +| **GKN layout** | grouped expert-weight layout `[G=experts, K=in_features, N=out_features]`. | +| **exact** | shorthand for "serving-parity byte contract": the forward reproduces the serving engine's bits, not just its math. | From 743af782e68f211e4ef8ce6748cb9817edd8eb03 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 16:26:20 +0000 Subject: [PATCH 09/14] Rename rope_class_b to rope_fp32_single_round (#81 review, row 1) 'Class-B' was blog-post nomenclature; the descriptive name says what it is: one fp32 chain with a single final round. Module, functions, internal config stamps, and the YAML field are renamed; 'rope_class_b:' remains accepted as a deprecated alias in both entry points (DeprecationWarning), and XORL_ROPE_CLASS_B stays an env alias. The ops module docstring keeps the Class-A/Class-B mapping as historical context. --- .github/workflows/pr-test-cpu.yml | 2 +- .../content/docs/config-reference/server.md | 2 +- src/xorl/arguments.py | 21 +++++- src/xorl/distributed/torch_parallelize.py | 6 +- src/xorl/models/auto.py | 50 +++++++-------- src/xorl/models/layers/rope.py | 32 +++++----- .../deepseek_v3/modeling_deepseek_v3.py | 2 +- src/xorl/models/transformers/glm5/indexer.py | 6 +- .../models/transformers/glm5/modeling_glm5.py | 4 +- src/xorl/models/transformers/glm5/rotary.py | 8 +-- .../transformers/gpt_oss/modeling_gpt_oss.py | 6 +- .../minimax_m3/modeling_minimax_m3.py | 4 +- .../transformers/qwen3_5/modeling_qwen3_5.py | 2 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 2 +- .../models/transformers/qwen3_5_shared.py | 4 +- src/xorl/ops/README.md | 2 +- src/xorl/ops/exact/__init__.py | 2 +- ...e_class_b.py => rope_fp32_single_round.py} | 26 ++++---- src/xorl/ops/rope_class_b.py | 4 +- src/xorl/server/runner/model_runner.py | 4 +- src/xorl/server/server_arguments.py | 24 +++++-- src/xorl/trainers/model_builder.py | 10 +-- src/xorl/trainers/trainer.py | 4 +- ..._rope_fp32_single_round_fsdp_transport.py} | 6 +- tests/models/test_qwen3_5_apply_rotary.py | 12 ++-- .../test_rmsnorm_family_cross_engine.py | 8 +-- tests/models/test_rope_inv_freq_fp32.py | 10 +-- tests/ops/test_rope_class_b_contract.py | 18 +++--- .../runner/test_model_runner_fp8_training.py | 2 +- tests/server/test_server_arguments.py | 4 +- .../trainers/test_qwen3_dense_exact_config.py | 8 +-- ... => test_rope_fp32_single_round_config.py} | 64 +++++++++---------- .../test_trainer_model_alignment_flags.py | 6 +- tests/trainers/test_unfuse_for_lora.py | 4 +- 34 files changed, 201 insertions(+), 168 deletions(-) rename src/xorl/ops/exact/{rope_class_b.py => rope_fp32_single_round.py} (87%) rename tests/distributed/{test_rope_class_b_fsdp_transport.py => test_rope_fp32_single_round_fsdp_transport.py} (94%) rename tests/trainers/{test_rope_class_b_config.py => test_rope_fp32_single_round_config.py} (89%) diff --git a/.github/workflows/pr-test-cpu.yml b/.github/workflows/pr-test-cpu.yml index f02515f5..ad424921 100644 --- a/.github/workflows/pr-test-cpu.yml +++ b/.github/workflows/pr-test-cpu.yml @@ -75,7 +75,7 @@ jobs: tests/distributed/test_native_fp8_fsdp2_materialization.py tests/distributed/test_bi_trunk_linear_fsdp.py tests/distributed/test_qwen35_lora_projection_fsdp2.py - tests/distributed/test_rope_class_b_fsdp_transport.py + tests/distributed/test_rope_fp32_single_round_fsdp_transport.py tests/distributed/test_torch_parallelize_policies.py tests/distributed/test_parallel_plan_meta_slice.py diff --git a/docs/src/content/docs/config-reference/server.md b/docs/src/content/docs/config-reference/server.md index 4ee926a7..006a398f 100644 --- a/docs/src/content/docs/config-reference/server.md +++ b/docs/src/content/docs/config-reference/server.md @@ -53,7 +53,7 @@ These stored defaults are resolved after the model architecture is known. Ordina | `rmsnorm_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `native`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs require `sglang_fused`. Other explicit diagnostic modes are also accepted by the argument type. | | `activation_native` | `false` (resolved) | Use native SiLU instead of the fused Triton kernel. Exact Qwen3.5-family programs resolve this to `true`; the other exact programs retain their architecture-owned fused arithmetic. | | `rope_native` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `false`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `true`. | -| `rope_class_b` | `null` (resolved) | Select the compiled Class-B RoPE FP32-chain path. It is enabled for exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs; DSV4 owns a separate RoPE program. | +| `rope_fp32_single_round` | `null` (resolved) | Select the compiled Class-B RoPE FP32-chain path. It is enabled for exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs; DSV4 owns a separate RoPE program. | | `attention_cast_bf16` | `false` (resolved) | Explicitly cast Q/K to BF16 after RoPE. Exact Qwen3.5-family programs resolve this to `true`; dense Qwen3, GLM-5.2, and DSV4-Flash exact programs require `false`. | | `qwen35_rmsnorm_family` | `null` (resolved) | Exact Qwen3.5/3.6 programs require the qualified `v2` arithmetic; other architectures reject an override. | | `sparse_mla_enabled` | `null` (resolved) | Canonical GLM-5.2 enables the sparse-MLA path; ordinary models resolve to `false`. | diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index a4d9e046..06804fb9 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -585,13 +585,18 @@ class ModelArguments: default=None, metadata={"help": "Use native RoPE. Auto-enables for canonical GLM-5.2 when omitted."}, ) - rope_class_b: Optional[bool] = field( + rope_fp32_single_round: Optional[bool] = field( default=None, metadata={ - "help": "Use compiled Class-B RoPE fp32-chain numerics. Auto-enables for canonical GLM-5.2 and " - "exact Qwen3.5-family training when omitted." + "help": "Use the compiled fp32 single-round RoPE numerics (one fp32 chain, one final round; " + "formerly 'Class-B'). Auto-enables for canonical GLM-5.2 and exact Qwen3.5-family training " + "when omitted." }, ) + rope_class_b: Optional[bool] = field( + default=None, + metadata={"help": "DEPRECATED alias for rope_fp32_single_round."}, + ) attention_cast_bf16: bool = field( default=False, metadata={"help": "Explicitly cast Q/K to bfloat16 after RoPE for SGLang alignment."} ) @@ -624,6 +629,16 @@ class ModelArguments: ) def __post_init__(self): + if self.rope_class_b is not None: + import warnings + + warnings.warn( + "rope_class_b is deprecated; use rope_fp32_single_round", + DeprecationWarning, + stacklevel=2, + ) + if self.rope_fp32_single_round is None: + self.rope_fp32_single_round = self.rope_class_b if self.config_path is None and self.model_path is None: raise ValueError("`config_path` must be specified when `model_path` is None.") diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index 5211cd7c..3a58ee44 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -135,12 +135,12 @@ def _bf16_mixed_precision_policy( def _decoder_bf16_mixed_precision_policy( reduce_dtype: torch.dtype = torch.float32, *, - class_b: bool = False, + fp32_single_round: bool = False, ): """Preserve the fp32 RoPE table when the Class-B contract is active.""" return _bf16_mixed_precision_policy( reduce_dtype=reduce_dtype, - cast_forward_inputs=not class_b, + cast_forward_inputs=not fp32_single_round, ) @@ -443,7 +443,7 @@ def parallelize_model_fsdp2( fsdp_kwargs["mp_policy"] = mp_policy decoder_mp_policy = _decoder_bf16_mixed_precision_policy( reduce_dtype=reduce_dtype, - class_b=bool(getattr(model.config, "_rope_class_b", False)), + fp32_single_round=bool(getattr(model.config, "_rope_fp32_single_round", False)), ) if not decoder_mp_policy.cast_forward_inputs: logger.info_rank0( diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 0ee5c673..14055c61 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -25,7 +25,7 @@ ) from .layers.attention import get_attention_fn from .layers.normalization import set_rmsnorm_mode -from .layers.rope import set_rope_class_b, set_rope_native +from .layers.rope import set_rope_fp32_single_round, set_rope_native from .loader import ModelLoader, get_loader from .transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config from .transformers.deepseek_v3.support import validate_deepseek_v3_router_settings @@ -522,7 +522,7 @@ class ResolvedModelNumericalProgram: qwen35_rmsnorm_family: Optional[str] activation_native: bool rope_native: bool - rope_class_b: bool + rope_fp32_single_round: bool attention_cast_bf16: bool sparse_mla_enabled: bool sparse_mla_backend: str @@ -532,42 +532,42 @@ def _resolve_rope_modes( config: PretrainedConfig, *, rope_native: Optional[bool], - rope_class_b: Optional[bool], + rope_fp32_single_round: Optional[bool], ) -> tuple[bool, bool]: if _is_exact_glm52(config): - if rope_native is False or rope_class_b is False: + if rope_native is False or rope_fp32_single_round is False: raise ValueError( "Canonical GLM-5.2 requires native Class-B RoPE; explicit rope_native=false " - "or rope_class_b=false is incompatible with the model's numerical contract" + "or rope_fp32_single_round=false is incompatible with the model's numerical contract" ) return True, True if _is_exact_qwen35(config): - if rope_native is False or rope_class_b is False: + if rope_native is False or rope_fp32_single_round is False: raise ValueError( "Exact Qwen3.5-family server training requires native Class-B RoPE; " - "explicit rope_native=false or rope_class_b=false is incompatible " + "explicit rope_native=false or rope_fp32_single_round=false is incompatible " "with the model's numerical contract" ) return True, True if _is_exact_qwen3_dense(config): - if rope_native is False or rope_class_b is False: + if rope_native is False or rope_fp32_single_round is False: raise ValueError( "Exact dense Qwen3 server training requires native Class-B RoPE; " - "explicit rope_native=false or rope_class_b=false is incompatible " + "explicit rope_native=false or rope_fp32_single_round=false is incompatible " "with the model's numerical contract" ) return True, True effective_rope_native = bool(rope_native) - effective_rope_class_b = bool(rope_class_b) - if effective_rope_class_b and not effective_rope_native: + effective_rope_fp32_single_round = bool(rope_fp32_single_round) + if effective_rope_fp32_single_round and not effective_rope_native: raise ValueError( - "rope_class_b=True requires rope_native=True: the Class-B contract uses " + "rope_fp32_single_round=True requires rope_native=True: the Class-B contract uses " "the CPU-built serving-layout cos/sin cache selected by rope_native" ) - return effective_rope_native, effective_rope_class_b + return effective_rope_native, effective_rope_fp32_single_round def resolve_model_numerical_program( @@ -580,7 +580,7 @@ def resolve_model_numerical_program( rmsnorm_mode: Optional[str], activation_native: bool, rope_native: Optional[bool], - rope_class_b: Optional[bool], + rope_fp32_single_round: Optional[bool], attention_cast_bf16: bool, sparse_mla_enabled: Optional[bool], sparse_mla_backend: Optional[str], @@ -593,10 +593,10 @@ def resolve_model_numerical_program( explicit values fail before weights are loaded. """ - effective_rope_native, effective_rope_class_b = _resolve_rope_modes( + effective_rope_native, effective_rope_fp32_single_round = _resolve_rope_modes( config, rope_native=rope_native, - rope_class_b=rope_class_b, + rope_fp32_single_round=rope_fp32_single_round, ) if qwen35_rmsnorm_family not in (None, "v1", "v2"): raise ValueError(f"qwen35_rmsnorm_family must be one of None, 'v1', or 'v2'; got {qwen35_rmsnorm_family!r}") @@ -627,7 +627,7 @@ def resolve_model_numerical_program( qwen35_rmsnorm_family="v2", activation_native=True, rope_native=True, - rope_class_b=effective_rope_class_b, + rope_fp32_single_round=effective_rope_fp32_single_round, attention_cast_bf16=True, sparse_mla_enabled=False, sparse_mla_backend="auto", @@ -657,7 +657,7 @@ def resolve_model_numerical_program( qwen35_rmsnorm_family=None, activation_native=False, rope_native=True, - rope_class_b=effective_rope_class_b, + rope_fp32_single_round=effective_rope_fp32_single_round, attention_cast_bf16=False, sparse_mla_enabled=False, sparse_mla_backend="auto", @@ -694,7 +694,7 @@ def resolve_model_numerical_program( qwen35_rmsnorm_family=None, activation_native=False, rope_native=False, - rope_class_b=False, + rope_fp32_single_round=False, attention_cast_bf16=False, sparse_mla_enabled=False, sparse_mla_backend="auto", @@ -714,7 +714,7 @@ def resolve_model_numerical_program( qwen35_rmsnorm_family=None, activation_native=activation_native, rope_native=effective_rope_native, - rope_class_b=effective_rope_class_b, + rope_fp32_single_round=effective_rope_fp32_single_round, attention_cast_bf16=attention_cast_bf16, sparse_mla_enabled=False if sparse_mla_enabled is None else sparse_mla_enabled, sparse_mla_backend=sparse_mla_backend or "auto", @@ -751,7 +751,7 @@ def resolve_model_numerical_program( qwen35_rmsnorm_family=None, activation_native=False, rope_native=True, - rope_class_b=True, + rope_fp32_single_round=True, attention_cast_bf16=False, sparse_mla_enabled=True, sparse_mla_backend="flashmla", @@ -809,7 +809,7 @@ def build_foundation_model( qwen35_rmsnorm_family: Optional[Literal["v1", "v2"]] = None, activation_native: bool = False, rope_native: Optional[bool] = None, - rope_class_b: Optional[bool] = None, + rope_fp32_single_round: Optional[bool] = None, attention_cast_bf16: bool = False, sparse_mla_enabled: Optional[bool] = None, sparse_mla_backend: Optional[str] = None, @@ -945,7 +945,7 @@ def build_foundation_model( qwen35_rmsnorm_family=qwen35_rmsnorm_family, activation_native=activation_native, rope_native=rope_native, - rope_class_b=rope_class_b, + rope_fp32_single_round=rope_fp32_single_round, attention_cast_bf16=attention_cast_bf16, sparse_mla_enabled=sparse_mla_enabled, sparse_mla_backend=sparse_mla_backend, @@ -965,9 +965,9 @@ def build_foundation_model( # rotary helper. if not (canonical_glm52 or config._qwen35_exact_contract): set_rope_native(numerical_program.rope_native) - set_rope_class_b(numerical_program.rope_class_b) + set_rope_fp32_single_round(numerical_program.rope_fp32_single_round) config._rope_native = numerical_program.rope_native - config._rope_class_b = numerical_program.rope_class_b + config._rope_fp32_single_round = numerical_program.rope_fp32_single_round config._resolved_numerical_program = asdict(numerical_program) if canonical_glm52: logger.info_rank0(f"Canonical GLM-5.2 numerical program: {numerical_program}") diff --git a/src/xorl/models/layers/rope.py b/src/xorl/models/layers/rope.py index b8e83be8..6935f3c1 100644 --- a/src/xorl/models/layers/rope.py +++ b/src/xorl/models/layers/rope.py @@ -462,7 +462,7 @@ def __init__(self, config, device=None): self._set_inv_freq_fp32(self._cpu_fp32_inv_freq()) self._sglang_default_cache = None self._use_sglang_default_cache = bool(getattr(config, "_rope_native", False) and self.rope_type == "default") - self._class_b = bool(getattr(config, "_rope_class_b", False) or glm52_exact_forward_enabled(config)) + self._fp32_single_round = bool(getattr(config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(config)) def _cpu_fp32_inv_freq(self) -> torch.Tensor: """Frequency table computed on CPU in fp32 — the provenance serving's cos/sin cache is built with.""" @@ -545,7 +545,7 @@ def forward(self, x, position_ids): cos_half, sin_half = cos_sin.chunk(2, dim=-1) cos = torch.cat((cos_half, cos_half), dim=-1).view(*position_ids.shape, -1) sin = torch.cat((sin_half, sin_half), dim=-1).view(*position_ids.shape, -1) - out_dtype = torch.float32 if self._class_b else x.dtype + out_dtype = torch.float32 if self._fp32_single_round else x.dtype return cos.to(device=x.device, dtype=out_dtype), sin.to(device=x.device, dtype=out_dtype) inv_freq = self._resolve_inv_freq(x.device) @@ -561,7 +561,7 @@ def forward(self, x, position_ids): # Class B feeds cos/sin straight into an fp32 kernel cache; a bf16 round here # would put the result back in Class A. - out_dtype = torch.float32 if self._class_b else x.dtype + out_dtype = torch.float32 if self._fp32_single_round else x.dtype return cos.to(dtype=out_dtype), sin.to(dtype=out_dtype) @@ -623,24 +623,26 @@ def set_rope_native(enabled: bool): # first would land back in Class A, so the Class-B lane keeps the table in fp32 all # the way to the kernel. -_rope_class_b = os.environ.get("XORL_ROPE_CLASS_B") == "1" +_rope_fp32_single_round = ( + os.environ.get("XORL_ROPE_FP32_SINGLE_ROUND", os.environ.get("XORL_ROPE_CLASS_B", "")) == "1" +) -def set_rope_class_b(enabled: bool) -> None: +def set_rope_fp32_single_round(enabled: bool) -> None: """Select compiled fp32-chain numerics aligned with SGLang's stock fused CUDA kernel.""" - global _rope_class_b - _rope_class_b = enabled + global _rope_fp32_single_round + _rope_fp32_single_round = enabled -def rope_class_b_enabled() -> bool: - return _rope_class_b +def rope_fp32_single_round_enabled() -> bool: + return _rope_fp32_single_round def stock_fused_apply_rotary_pos_emb(q, k, cos, sin, *, interleaved: bool = False, doubled: bool = True): - """Class-B RoPE application backed by the compiled expression in ``xorl.ops.exact.rope_class_b``.""" - from xorl.ops.exact.rope_class_b import class_b_apply_rotary_pos_emb # noqa: PLC0415 + """Class-B RoPE application backed by the compiled expression in ``xorl.ops.exact.rope_fp32_single_round``.""" + from xorl.ops.exact.rope_fp32_single_round import single_round_apply_rotary_pos_emb # noqa: PLC0415 - return class_b_apply_rotary_pos_emb(q, k, cos, sin, interleaved=interleaved, doubled=doubled) + return single_round_apply_rotary_pos_emb(q, k, cos, sin, interleaved=interleaved, doubled=doubled) def apply_rotary_pos_emb(q, k, cos, sin, *, force_native: bool = False): @@ -658,7 +660,7 @@ def apply_rotary_pos_emb(q, k, cos, sin, *, force_native: bool = False): cos: The cosine part from RotaryEmbedding, shape [batch, seq_len, head_dim]. sin: The sine part from RotaryEmbedding, shape [batch, seq_len, head_dim]. """ - if _rope_class_b and q.is_cuda: + if _rope_fp32_single_round and q.is_cuda: return stock_fused_apply_rotary_pos_emb(q, k, cos, sin) if _flash_apply_rotary_emb is not None and q.is_cuda and not (_rope_native or force_native): @@ -708,9 +710,9 @@ def rope_config_validation(config, ignore_keys=None): "RotaryEmbedding", "apply_rotary_pos_emb", "dynamic_rope_update", - "rope_class_b_enabled", + "rope_fp32_single_round_enabled", "rope_config_validation", "rotate_half", - "set_rope_class_b", + "set_rope_fp32_single_round", "stock_fused_apply_rotary_pos_emb", ] diff --git a/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py b/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py index 0233373d..9093fb24 100644 --- a/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py +++ b/src/xorl/models/transformers/deepseek_v3/modeling_deepseek_v3.py @@ -143,7 +143,7 @@ def _project_qkv( cos, sin, interleaved=getattr(self.config, "rope_interleave", True), - class_b=bool(getattr(self.config, "_rope_class_b", False)), + fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False)), ) k_rot = k_rot.expand(*k_pass.shape[:-1], -1) diff --git a/src/xorl/models/transformers/glm5/indexer.py b/src/xorl/models/transformers/glm5/indexer.py index f900005b..ca3efd8e 100644 --- a/src/xorl/models/transformers/glm5/indexer.py +++ b/src/xorl/models/transformers/glm5/indexer.py @@ -25,7 +25,7 @@ GLM52_SELECTOR_VERSION, select_glm52_logical_indices, ) -from xorl.ops.exact.rope_class_b import build_class_b_cos_sin +from xorl.ops.exact.rope_fp32_single_round import build_single_round_cos_sin from xorl.ops.sglang.batch_invariant_ops import bi_bf16_fp32_linear, matmul_persistent @@ -72,7 +72,7 @@ def _fused_sampler_index_k_prepare( f"got key={tuple(raw_key.shape)} cos={tuple(cos.shape)} sin={tuple(sin.shape)}" ) - cos_half, sin_half = build_class_b_cos_sin(cos, sin) + cos_half, sin_half = build_single_round_cos_sin(cos, sin) cos_sin_cache = torch.cat((cos_half, sin_half), dim=-1).contiguous() flat_key = raw_key.reshape(-1, raw_key.shape[-1]) if flat_key.stride(-1) != 1: @@ -332,7 +332,7 @@ def project( cos, sin, interleaved=getattr(self.config, "indexer_rope_interleave", True), - class_b=bool(getattr(self.config, "_rope_class_b", False) or glm52_exact_forward_enabled(self.config)), + fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(self.config)), ) index_q = torch.cat([q_pe, q_no_pe], dim=-1) diff --git a/src/xorl/models/transformers/glm5/modeling_glm5.py b/src/xorl/models/transformers/glm5/modeling_glm5.py index 7bb2e375..0f6341b1 100644 --- a/src/xorl/models/transformers/glm5/modeling_glm5.py +++ b/src/xorl/models/transformers/glm5/modeling_glm5.py @@ -271,7 +271,7 @@ def _project_qkv( cos, sin, interleaved=getattr(self.config, "rope_interleave", True), - class_b=bool(getattr(self.config, "_rope_class_b", False) or glm52_exact_forward_enabled(self.config)), + fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(self.config)), ) k_rot = k_rot.expand(*k_pass.shape[:-1], -1) @@ -461,7 +461,7 @@ def _project_qkv_absorb( cos, sin, interleaved=getattr(self.config, "rope_interleave", True), - class_b=bool(getattr(self.config, "_rope_class_b", False) or glm52_exact_forward_enabled(self.config)), + fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(self.config)), ) k_pe = k_pe.squeeze(2) diff --git a/src/xorl/models/transformers/glm5/rotary.py b/src/xorl/models/transformers/glm5/rotary.py index d18a5df4..841a340a 100644 --- a/src/xorl/models/transformers/glm5/rotary.py +++ b/src/xorl/models/transformers/glm5/rotary.py @@ -4,7 +4,7 @@ import torch -from xorl.models.layers.rope import rope_class_b_enabled, stock_fused_apply_rotary_pos_emb +from xorl.models.layers.rope import rope_fp32_single_round_enabled, stock_fused_apply_rotary_pos_emb def glm5_rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor: @@ -24,10 +24,10 @@ def glm5_apply_rotary_pos_emb( cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False, - class_b: bool | None = None, + fp32_single_round: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - use_class_b = rope_class_b_enabled() if class_b is None else class_b - if use_class_b and q.is_cuda: + use_fp32_single_round = rope_fp32_single_round_enabled() if fp32_single_round is None else fp32_single_round + if use_fp32_single_round and q.is_cuda: return stock_fused_apply_rotary_pos_emb(q, k, cos, sin, interleaved=interleaved) if interleaved: diff --git a/src/xorl/models/transformers/gpt_oss/modeling_gpt_oss.py b/src/xorl/models/transformers/gpt_oss/modeling_gpt_oss.py index d04a9d8c..267451d8 100644 --- a/src/xorl/models/transformers/gpt_oss/modeling_gpt_oss.py +++ b/src/xorl/models/transformers/gpt_oss/modeling_gpt_oss.py @@ -26,7 +26,7 @@ from xorl.models.layers.attention.backend import ATTENTION_FUNCTIONS from xorl.models.layers.moe import MoEBlock, MoEExperts from xorl.models.layers.normalization import RMSNorm -from xorl.models.layers.rope import rope_class_b_enabled, stock_fused_apply_rotary_pos_emb +from xorl.models.layers.rope import rope_fp32_single_round_enabled, stock_fused_apply_rotary_pos_emb from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput from xorl.models.transformers.gpt_oss import parallelize from xorl.models.transformers.gpt_oss.checkpoint_handler import GptOssCheckpointHandler @@ -85,7 +85,7 @@ def gpt_oss_apply_rotary_pos_emb( cos: ``[batch, seq, head_dim // 2]`` sin: ``[batch, seq, head_dim // 2]`` """ - if rope_class_b_enabled() and query.is_cuda: + if rope_fp32_single_round_enabled() and query.is_cuda: return stock_fused_apply_rotary_pos_emb(query, key, cos, sin, doubled=False) # Unsqueeze for head dimension: [batch, seq, 1, head_dim // 2] @@ -160,7 +160,7 @@ def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Te freqs = torch.einsum("bi,j->bij", position_ids.float(), inv_freq) cos = freqs.cos() * concentration sin = freqs.sin() * concentration - if rope_class_b_enabled(): + if rope_fp32_single_round_enabled(): return cos.float(), sin.float() return cos, sin diff --git a/src/xorl/models/transformers/minimax_m3/modeling_minimax_m3.py b/src/xorl/models/transformers/minimax_m3/modeling_minimax_m3.py index 5a7567ef..3dc6f16d 100644 --- a/src/xorl/models/transformers/minimax_m3/modeling_minimax_m3.py +++ b/src/xorl/models/transformers/minimax_m3/modeling_minimax_m3.py @@ -401,7 +401,7 @@ def _project_qkv( key_states, cos, sin, - class_b=bool(getattr(self.config, "_rope_class_b", False)), + fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False)), ) + (value_states,) def _project_index_qk( @@ -424,7 +424,7 @@ def _project_index_qk( index_key, cos, sin, - class_b=bool(getattr(self.config, "_rope_class_b", False)), + fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False)), ) def _project_output(self, attn_output: torch.Tensor) -> torch.Tensor: diff --git a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py index a0dafb40..ca4199af 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -273,7 +273,7 @@ def _project_qkv( key_states, cos, sin, - class_b=bool(getattr(self.config, "_rope_class_b", False)), + fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False)), ) if getattr(self.config, "_attention_cast_bf16", False): query_states = query_states.to(torch.bfloat16) diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index 2cd9278f..f6817c85 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -336,7 +336,7 @@ def _project_qkv( key_states, cos, sin, - class_b=bool(getattr(self.config, "_rope_class_b", False)), + fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False)), ) if getattr(self.config, "_attention_cast_bf16", False): query_states = query_states.to(torch.bfloat16) diff --git a/src/xorl/models/transformers/qwen3_5_shared.py b/src/xorl/models/transformers/qwen3_5_shared.py index b1d7a4e1..a1854313 100644 --- a/src/xorl/models/transformers/qwen3_5_shared.py +++ b/src/xorl/models/transformers/qwen3_5_shared.py @@ -119,9 +119,9 @@ def qwen3_5_apply_rotary_pos_emb( cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False, - class_b: bool = False, + fp32_single_round: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - if class_b: + if fp32_single_round: if not q.is_cuda or not k.is_cuda: raise RuntimeError("Qwen3.5-family Class-B RoPE requires CUDA q/k tensors") if q.dtype is not torch.bfloat16 or k.dtype is not torch.bfloat16: diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index a4c6eefc..077cf802 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -55,7 +55,7 @@ serving engine and cannot be renamed unilaterally): | --- | --- | | **batch-invariant / `bi_`** | a kernel whose per-element reduction order does not depend on batch composition, so the same token produces the same bits in any batch — the property that makes trainer/sampler logprobs comparable bitwise. | | **K3 / "zero-K3"** | the k3 KL-divergence estimator between trainer and sampler logprobs for the same tokens; "zero-K3" = bit-identical train/serve forward, the goal of the exact contracts. | -| **Class-A / Class-B RoPE** | the two RoPE numerics classes across the trainer/sampler pair: Class A rounds to bf16 per op (8 rounding points); Class B computes one fp32 chain with a single final round (SGLang's fused CUDA rope and its compiled RL-lane path). `rope_class_b: true` selects Class B. | +| **Class-A / Class-B RoPE** | the two RoPE numerics classes across the trainer/sampler pair: Class A rounds to bf16 per op (8 rounding points); Class B computes one fp32 chain with a single final round (SGLang's fused CUDA rope and its compiled RL-lane path). `rope_fp32_single_round: true` selects Class B. | | **canonical MoE reduce** | the pinned fixed-order expert-contribution reduction (contributor-leaf arithmetic + final FP64-accumulator cast) shared with serving, versioned by `CANONICAL_MOE_REDUCE_VERSION`. | | **one-round SwiGLU** | the exact-contract SwiGLU with a single FP32 rounding point (`exact_fp32_silu_and_mul`), vs the generic fused SwiGLU. | | **families v1 / v2** | versioned batch-invariant kernel families (norms, LM head); v2 is the epilogue-stats generation. `bi_families_v2` is the serving twin module. | diff --git a/src/xorl/ops/exact/__init__.py b/src/xorl/ops/exact/__init__.py index 8055b585..2ce7ea88 100644 --- a/src/xorl/ops/exact/__init__.py +++ b/src/xorl/ops/exact/__init__.py @@ -4,7 +4,7 @@ ``canonical_`` / ``class_b`` prefixes all meant: byte-pinned programs shared with the serving engine. -Here: ``sampling_transforms`` (the replay contract), ``rope_class_b``, +Here: ``sampling_transforms`` (the replay contract), ``rope_fp32_single_round``, ``canonical_moe_leaf``, ``canonical_moe_cast``, ``kernel_config_pin``, ``bi_gemm_configs``, ``block_fp8_native``, ``fused_silu_and_mul``. diff --git a/src/xorl/ops/exact/rope_class_b.py b/src/xorl/ops/exact/rope_fp32_single_round.py similarity index 87% rename from src/xorl/ops/exact/rope_class_b.py rename to src/xorl/ops/exact/rope_fp32_single_round.py index 7471737d..ce6f6973 100644 --- a/src/xorl/ops/exact/rope_class_b.py +++ b/src/xorl/ops/exact/rope_fp32_single_round.py @@ -38,19 +38,19 @@ # it dynamo silently falls back to eager -- reverting those layers to Class A with no error and # no log line. That is indistinguishable from a numerics bug downstream, so pin the budget and # make exhaustion loud. -_ROPE_CLASS_B_RECOMPILE_LIMIT = 2048 -_ROPE_CLASS_B_ACCUMULATED_LIMIT = 8192 +_ROPE_SINGLE_ROUND_RECOMPILE_LIMIT = 2048 +_ROPE_SINGLE_ROUND_ACCUMULATED_LIMIT = 8192 def _pin_compile_budget() -> None: cfg = torch._dynamo.config - cfg.recompile_limit = max(getattr(cfg, "recompile_limit", 0), _ROPE_CLASS_B_RECOMPILE_LIMIT) + cfg.recompile_limit = max(getattr(cfg, "recompile_limit", 0), _ROPE_SINGLE_ROUND_RECOMPILE_LIMIT) cfg.accumulated_recompile_limit = max( - getattr(cfg, "accumulated_recompile_limit", 0), _ROPE_CLASS_B_ACCUMULATED_LIMIT + getattr(cfg, "accumulated_recompile_limit", 0), _ROPE_SINGLE_ROUND_ACCUMULATED_LIMIT ) # A silent eager fallback breaks the zero-K3 contract, so fail instead. Opt out with - # XORL_ROPE_CLASS_B_ALLOW_FALLBACK=1 for throughput experiments outside the contract. - if os.environ.get("XORL_ROPE_CLASS_B_ALLOW_FALLBACK") != "1" and hasattr(cfg, "fail_on_recompile_limit_hit"): + # XORL_ROPE_SINGLE_ROUND_ALLOW_FALLBACK=1 for throughput experiments outside the contract. + if os.environ.get("XORL_ROPE_SINGLE_ROUND_ALLOW_FALLBACK") != "1" and hasattr(cfg, "fail_on_recompile_limit_hit"): cfg.fail_on_recompile_limit_hit = True logger.info( "rope Class B: recompile_limit=%s accumulated=%s fail_on_limit=%s", @@ -85,7 +85,7 @@ def _rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, is_neox: _rotary_emb_compiled = torch.compile(dynamic=True)(_rotary_emb) -class _ClassBRoPE(torch.autograd.Function): +class _SingleRoundRoPE(torch.autograd.Function): """Class-B rope with a hand-written backward. RoPE is a block rotation, so its adjoint is the same rotation driven by a negated sin. @@ -115,7 +115,7 @@ def backward(ctx, grad_q, grad_k): ) -def build_class_b_cos_sin(cos, sin, *, doubled: bool = True): +def build_single_round_cos_sin(cos, sin, *, doubled: bool = True): """Flatten the trainer's cos/sin into the fp32 ``[num_tokens, rotary_dim // 2]`` serving layout. With ``doubled`` (the default) cos/sin arrive in the ``[..., rotary_dim]`` layout the shared @@ -126,7 +126,7 @@ def build_class_b_cos_sin(cos, sin, *, doubled: bool = True): return cos[..., :half].reshape(-1, half).float(), sin[..., :half].reshape(-1, half).float() -def class_b_apply_rotary_pos_emb(q, k, cos, sin, *, interleaved: bool = False, doubled: bool = True): +def single_round_apply_rotary_pos_emb(q, k, cos, sin, *, interleaved: bool = False, doubled: bool = True): """Class-B RoPE over ``[B, S, H, D]`` q/k. ``interleaved`` selects GPT-J pairing; the default is the half-split Neox pairing. Partial @@ -138,14 +138,14 @@ def class_b_apply_rotary_pos_emb(q, k, cos, sin, *, interleaved: bool = False, d f"got cos={cos.dtype}, sin={sin.dtype}. A mixed-precision wrapper " "likely downcast the table in transit." ) - cos_f, sin_f = build_class_b_cos_sin(cos, sin, doubled=doubled) + cos_f, sin_f = build_single_round_cos_sin(cos, sin, doubled=doubled) rotary_dim = 2 * cos_f.shape[-1] batch, seq_len = q.shape[0], q.shape[1] num_tokens = batch * seq_len q_rot = q[..., :rotary_dim].reshape(num_tokens, -1, rotary_dim) k_rot = k[..., :rotary_dim].reshape(num_tokens, -1, rotary_dim) - q_out, k_out = _ClassBRoPE.apply(q_rot, k_rot, cos_f, sin_f, not interleaved) + q_out, k_out = _SingleRoundRoPE.apply(q_rot, k_rot, cos_f, sin_f, not interleaved) q_out = q_out.view(batch, seq_len, -1, rotary_dim) k_out = k_out.view(batch, seq_len, -1, rotary_dim) @@ -156,6 +156,6 @@ def class_b_apply_rotary_pos_emb(q, k, cos, sin, *, interleaved: bool = False, d __all__ = [ - "build_class_b_cos_sin", - "class_b_apply_rotary_pos_emb", + "build_single_round_cos_sin", + "single_round_apply_rotary_pos_emb", ] diff --git a/src/xorl/ops/rope_class_b.py b/src/xorl/ops/rope_class_b.py index c68883f9..b7b277c8 100644 --- a/src/xorl/ops/rope_class_b.py +++ b/src/xorl/ops/rope_class_b.py @@ -1,3 +1,3 @@ -"""Moved to ``xorl.ops.exact.rope_class_b`` (issue #78 phase 3); compat stub.""" +"""Moved to ``xorl.ops.exact.rope_fp32_single_round`` (issue #78 phase 3); compat stub.""" -from xorl.ops.exact.rope_class_b import * # noqa: F401,F403 +from xorl.ops.exact.rope_fp32_single_round import * # noqa: F401,F403 diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 72b3365e..d6cbcd8b 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -1663,7 +1663,7 @@ def _initialize_model(self): qwen35_rmsnorm_family=self.model_config.get("qwen35_rmsnorm_family"), activation_native=self.model_config.get("activation_native", False), rope_native=self.model_config.get("rope_native"), - rope_class_b=self.model_config.get("rope_class_b"), + rope_fp32_single_round=self.model_config.get("rope_fp32_single_round"), attention_cast_bf16=self.model_config.get("attention_cast_bf16", False), sparse_mla_enabled=self.model_config.get("sparse_mla_enabled"), sparse_mla_backend=self.model_config.get("sparse_mla_backend"), @@ -1689,7 +1689,7 @@ def _initialize_model(self): "rmsnorm_mode": numerical_program["rmsnorm_mode"], "activation_native": numerical_program["activation_native"], "rope_native": numerical_program["rope_native"], - "rope_class_b": numerical_program["rope_class_b"], + "rope_fp32_single_round": numerical_program["rope_fp32_single_round"], "attention_cast_bf16": numerical_program["attention_cast_bf16"], "sparse_mla_enabled": numerical_program["sparse_mla_enabled"], "sparse_mla_backend": numerical_program["sparse_mla_backend"], diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 544b114d..649caec0 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -175,13 +175,18 @@ class ServerArguments: default=None, metadata={"help": "Use native RoPE. Auto-enables for canonical GLM-5.2 when omitted."}, ) - rope_class_b: Optional[bool] = field( + rope_fp32_single_round: Optional[bool] = field( default=None, metadata={ - "help": "Use compiled Class-B RoPE fp32-chain numerics. Auto-enables for canonical GLM-5.2 and " - "exact Qwen3.5-family training when omitted." + "help": "Use the compiled fp32 single-round RoPE numerics (one fp32 chain, one final round; " + "formerly 'Class-B'). Auto-enables for canonical GLM-5.2 and exact Qwen3.5-family training " + "when omitted." }, ) + rope_class_b: Optional[bool] = field( + default=None, + metadata={"help": "DEPRECATED alias for rope_fp32_single_round."}, + ) attention_cast_bf16: bool = field( default=False, metadata={"help": "Explicitly cast Q/K to bfloat16 after RoPE for SGLang alignment."} @@ -1224,6 +1229,17 @@ def __post_init__(self): """Validate and set defaults.""" from xorl.fp8_training.config_compat import normalize_fp8_training_config # noqa: PLC0415 + if self.rope_class_b is not None: + import warnings + + warnings.warn( + "rope_class_b is deprecated; use rope_fp32_single_round", + DeprecationWarning, + stacklevel=2, + ) + if self.rope_fp32_single_round is None: + self.rope_fp32_single_round = self.rope_class_b + if isinstance(self.max_grad_norm, bool): raise ValueError("max_grad_norm must be a finite number; use a value <= 0 to disable clipping") try: @@ -1496,7 +1512,7 @@ def to_config_dict(self) -> Dict[str, Any]: "qwen35_rmsnorm_family": self.qwen35_rmsnorm_family, "activation_native": self.activation_native, "rope_native": self.rope_native, - "rope_class_b": self.rope_class_b, + "rope_fp32_single_round": self.rope_fp32_single_round, "attention_cast_bf16": self.attention_cast_bf16, "flash_attention_deterministic": self.flash_attention_deterministic, "sparse_mla_enabled": self.sparse_mla_enabled, diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 268f2444..4cd2b6a8 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -293,7 +293,7 @@ def build_training_model( qwen35_rmsnorm_family: Optional[str] = None, activation_native: bool = False, rope_native: Optional[bool] = None, - rope_class_b: Optional[bool] = None, + rope_fp32_single_round: Optional[bool] = None, attention_cast_bf16: bool = False, sparse_mla_enabled: Optional[bool] = None, sparse_mla_backend: Optional[str] = None, @@ -325,9 +325,9 @@ def build_training_model( "Use physical pipeline parallelism with one stage per rank " "(pipeline_parallel_virtual_stages=1)." ) - if rope_class_b is True and rope_native is False: + if rope_fp32_single_round is True and rope_native is False: raise ValueError( - "rope_class_b=True requires rope_native=True: the Class-B contract uses " + "rope_fp32_single_round=True requires rope_native=True: the Class-B contract uses " "the CPU-built serving-layout cos/sin cache selected by rope_native" ) if block_fp8_qlora_training and (not enable_lora or not enable_qlora): @@ -398,7 +398,7 @@ def build_training_model( qwen35_rmsnorm_family=qwen35_rmsnorm_family, activation_native=activation_native, rope_native=rope_native, - rope_class_b=rope_class_b, + rope_fp32_single_round=rope_fp32_single_round, attention_cast_bf16=attention_cast_bf16, sparse_mla_enabled=sparse_mla_enabled, sparse_mla_backend=sparse_mla_backend, @@ -421,7 +421,7 @@ def build_training_model( # Set module-level flags for rope and activation if getattr(model.config, "_rope_native", False): logger.info_rank0("Using native RoPE (flash_attn fused kernel disabled)") - if getattr(model.config, "_rope_class_b", False): + if getattr(model.config, "_rope_fp32_single_round", False): logger.info_rank0( "Using compiled Class-B RoPE fp32-chain numerics aligned with SGLang's stock fused CUDA kernel" ) diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index d29a5c64..09fb9ce2 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -713,7 +713,7 @@ def _build_model(self) -> None: qwen35_rmsnorm_family=args.model.qwen35_rmsnorm_family, activation_native=args.model.activation_native, rope_native=args.model.rope_native, - rope_class_b=args.model.rope_class_b, + rope_fp32_single_round=args.model.rope_fp32_single_round, attention_cast_bf16=args.model.attention_cast_bf16, sparse_mla_enabled=args.model.sparse_mla_enabled, sparse_mla_backend=args.model.sparse_mla_backend, @@ -737,7 +737,7 @@ def _build_model(self) -> None: args.model.qwen35_rmsnorm_family = numerical_program["qwen35_rmsnorm_family"] args.model.activation_native = numerical_program["activation_native"] args.model.rope_native = numerical_program["rope_native"] - args.model.rope_class_b = numerical_program["rope_class_b"] + args.model.rope_fp32_single_round = numerical_program["rope_fp32_single_round"] args.model.attention_cast_bf16 = numerical_program["attention_cast_bf16"] args.model.sparse_mla_enabled = numerical_program["sparse_mla_enabled"] args.model.sparse_mla_backend = numerical_program["sparse_mla_backend"] diff --git a/tests/distributed/test_rope_class_b_fsdp_transport.py b/tests/distributed/test_rope_fp32_single_round_fsdp_transport.py similarity index 94% rename from tests/distributed/test_rope_class_b_fsdp_transport.py rename to tests/distributed/test_rope_fp32_single_round_fsdp_transport.py index d28bd117..5e396ca9 100644 --- a/tests/distributed/test_rope_class_b_fsdp_transport.py +++ b/tests/distributed/test_rope_fp32_single_round_fsdp_transport.py @@ -71,7 +71,7 @@ def _run() -> None: mesh = dist.device_mesh.init_device_mesh("cuda", (dist.get_world_size(),), mesh_dim_names=("dp_shard",)) model = _Root().to(device) - fully_shard(model.decoder, mesh=mesh, mp_policy=_decoder_bf16_mixed_precision_policy(class_b=True)) + fully_shard(model.decoder, mesh=mesh, mp_policy=_decoder_bf16_mixed_precision_policy(fp32_single_round=True)) fully_shard(model, mesh=mesh, mp_policy=_bf16_mixed_precision_policy()) x = torch.randn((2, 4, 8), dtype=torch.float32, device=device, requires_grad=True) @@ -83,7 +83,7 @@ def _run() -> None: assert model.decoder.proj.weight.grad is not None assert torch.isfinite(model.decoder.proj.weight.grad.to_local()).all() - control = _decoder_bf16_mixed_precision_policy(class_b=False) + control = _decoder_bf16_mixed_precision_policy(fp32_single_round=False) assert control.cast_forward_inputs is True, "control must retain stock FSDP input casting" dist.destroy_process_group() @@ -92,7 +92,7 @@ def _run() -> None: if __name__ != "__main__": @skip_if_gpu_count_less_than(1) - def test_class_b_table_survives_nested_fsdp_boundary(): + def test_fp32_single_round_table_survives_nested_fsdp_boundary(): result = run_distributed_script(__file__, num_gpus=1, timeout=120) result.assert_success("Class-B RoPE table should remain fp32 across decoder FSDP") diff --git a/tests/models/test_qwen3_5_apply_rotary.py b/tests/models/test_qwen3_5_apply_rotary.py index 8d000421..2dbbcdb9 100644 --- a/tests/models/test_qwen3_5_apply_rotary.py +++ b/tests/models/test_qwen3_5_apply_rotary.py @@ -77,7 +77,7 @@ def _assert_interleaved_matches_pairwise_reference(): torch.testing.assert_close(k_ours, k_ref, atol=1e-6, rtol=1e-6) -def _assert_class_b_rotary_admission_policy(monkeypatch): +def _assert_fp32_single_round_rotary_admission_policy(monkeypatch): sentinel = (object(), object()) calls = [] @@ -96,20 +96,20 @@ def _stock(q, k, cos, sin, *, interleaved): cos = torch.zeros((1, 2, 4), dtype=torch.float32) sin = torch.zeros_like(cos) - assert qwen3_5_shared.qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, class_b=True) is sentinel + assert qwen3_5_shared.qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, fp32_single_round=True) is sentinel assert calls == [(q, k, cos, sin, False)] - _assert_class_b_fails_loudly_outside_cuda_contract() + _assert_fp32_single_round_fails_loudly_outside_cuda_contract() -def _assert_class_b_fails_loudly_outside_cuda_contract(): +def _assert_fp32_single_round_fails_loudly_outside_cuda_contract(): cos = torch.zeros((1, 2, 4), dtype=torch.float32) sin = torch.zeros_like(cos) for dtype in (torch.bfloat16, torch.float32): q = torch.zeros((1, 2, 1, 4), dtype=dtype) k = torch.zeros_like(q) with pytest.raises(RuntimeError, match="requires CUDA"): - qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, class_b=True) + qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, fp32_single_round=True) def _assert_qwen35_attention_keeps_half_rotate_when_mrope_is_interleaved(attention_type, config_type): @@ -151,7 +151,7 @@ def _assert_qwen35_attention_keeps_half_rotate_when_mrope_is_interleaved(attenti def test_qwen35_rotary_numerics_admission_and_attention_policy(monkeypatch): _assert_interleaved_matches_pairwise_reference() - _assert_class_b_rotary_admission_policy(monkeypatch) + _assert_fp32_single_round_rotary_admission_policy(monkeypatch) for attention_type, config_type in ( (modeling_qwen3_5.Qwen3_5Attention, Qwen3_5Config), diff --git a/tests/models/test_rmsnorm_family_cross_engine.py b/tests/models/test_rmsnorm_family_cross_engine.py index c280fee9..c566d912 100644 --- a/tests/models/test_rmsnorm_family_cross_engine.py +++ b/tests/models/test_rmsnorm_family_cross_engine.py @@ -74,13 +74,13 @@ def test_rmsnorm_site_class_cross_engine_bitwise_policy(): family2 = sgl_bio.rms_norm_residual_tree_batch_invariant(x, w, EPS) assert not torch.equal(serving_out, family2), "serving families agree on the seed shape; gate is vacuous" - _assert_presummed_residual_tree_site_class_bitwise() - _assert_post_attention_residual_site_class_bitwise() + _assert_presummed_residual_tree_site_fp32_single_rounditwise() + _assert_post_attention_residual_site_fp32_single_rounditwise() _assert_zero_centered_family1_twin_bitwise() _assert_zero_centered_families_v2_candidate_bitwise() -def _assert_presummed_residual_tree_site_class_bitwise(): +def _assert_presummed_residual_tree_site_fp32_single_rounditwise(): """Input layernorm at layer>0 / final norm: xorl normalizes the pre-summed single tensor through the residual tree; SGLang fuses the add. On the same summed value both must produce identical bits (gate via a zero residual and @@ -100,7 +100,7 @@ def _assert_presummed_residual_tree_site_class_bitwise(): assert torch.equal(xorl_out, serving_funnel), f"pre-summed funnel {shape} diverged" -def _assert_post_attention_residual_site_class_bitwise(): +def _assert_post_attention_residual_site_fp32_single_rounditwise(): """Post-attention layernorm: xorl's fused residual dispatch must bit-match SGLang's fused residual dispatch, on both the normed output and the carried residual stream.""" diff --git a/tests/models/test_rope_inv_freq_fp32.py b/tests/models/test_rope_inv_freq_fp32.py index e68e7013..97ea9fc7 100644 --- a/tests/models/test_rope_inv_freq_fp32.py +++ b/tests/models/test_rope_inv_freq_fp32.py @@ -146,14 +146,14 @@ def _assert_native_default_cache_is_lazy_and_follows_execution_device(): assert torch.equal(cos[..., : HEAD_DIM // 2].reshape_as(cached_cos), cached_cos) assert torch.equal(sin[..., : HEAD_DIM // 2].reshape_as(cached_sin), cached_sin) - _assert_qwen_class_b_cache_growth_preserves_cpu_fp32_recipe() + _assert_qwen_fp32_single_round_cache_growth_preserves_cpu_fp32_recipe() -def _assert_qwen_class_b_cache_growth_preserves_cpu_fp32_recipe(): +def _assert_qwen_fp32_single_round_cache_growth_preserves_cpu_fp32_recipe(): config = _config("default") config.max_position_embeddings = 8 config._rope_native = True - config._rope_class_b = True + config._rope_fp32_single_round = True config._glm52_exact_contract = False config._qwen35_exact_contract = True rotary = RotaryEmbedding(config) @@ -177,14 +177,14 @@ def test_exact_architectures_build_default_rope_tables_on_their_serving_devices( glm_config = _config("default") glm_config._rope_native = True - glm_config._rope_class_b = True + glm_config._rope_fp32_single_round = True glm_config._glm52_exact_contract = True glm_config._qwen35_exact_contract = False glm_rotary = RotaryEmbedding(glm_config) qwen_config = _config("default") qwen_config._rope_native = True - qwen_config._rope_class_b = False + qwen_config._rope_fp32_single_round = False qwen_config._glm52_exact_contract = False qwen_config._qwen35_exact_contract = True qwen_rotary = RotaryEmbedding(qwen_config) diff --git a/tests/ops/test_rope_class_b_contract.py b/tests/ops/test_rope_class_b_contract.py index d5e027e4..fba5d7ff 100644 --- a/tests/ops/test_rope_class_b_contract.py +++ b/tests/ops/test_rope_class_b_contract.py @@ -3,38 +3,38 @@ import pytest import torch -from xorl.ops.exact.rope_class_b import ( - build_class_b_cos_sin, - class_b_apply_rotary_pos_emb, +from xorl.ops.exact.rope_fp32_single_round import ( + build_single_round_cos_sin, + single_round_apply_rotary_pos_emb, ) pytestmark = pytest.mark.cpu -def test_class_b_admission_shape_backward_and_table_layout_contract(): +def test_fp32_single_round_admission_shape_backward_and_table_layout_contract(): q = torch.zeros((1, 1, 1, 4), dtype=torch.bfloat16) cos = torch.ones((1, 1, 4), dtype=torch.bfloat16) sin = torch.zeros_like(cos) with pytest.raises(RuntimeError, match="requires fp32 cos/sin"): - class_b_apply_rotary_pos_emb(q, q, cos, sin) + single_round_apply_rotary_pos_emb(q, q, cos, sin) for shape in ((2, 1, 8, 8), (3, 2, 12, 8), (1, 1, 16, 8)): - _assert_class_b_shape_and_partial_rotary_backward(*shape) + _assert_fp32_single_round_shape_and_partial_rotary_backward(*shape) cos_half = torch.arange(12, dtype=torch.float32).view(1, 3, 4) sin_half = -cos_half cos = torch.cat((cos_half, cos_half), dim=-1) sin = torch.cat((sin_half, sin_half), dim=-1) - cos_flat, sin_flat = build_class_b_cos_sin(cos, sin) + cos_flat, sin_flat = build_single_round_cos_sin(cos, sin) assert cos_flat.shape == sin_flat.shape == (3, 4) assert torch.equal(cos_flat, cos_half.view(3, 4)) assert torch.equal(sin_flat, sin_half.view(3, 4)) -def _assert_class_b_shape_and_partial_rotary_backward(q_heads, k_heads, head_dim, rotary_dim): +def _assert_fp32_single_round_shape_and_partial_rotary_backward(q_heads, k_heads, head_dim, rotary_dim): torch.manual_seed(17) q = torch.randn((1, 3, q_heads, head_dim), dtype=torch.bfloat16, requires_grad=True) k = torch.randn((1, 3, k_heads, head_dim), dtype=torch.bfloat16, requires_grad=True) @@ -44,7 +44,7 @@ def _assert_class_b_shape_and_partial_rotary_backward(q_heads, k_heads, head_dim cos = torch.cat((cos_half, cos_half), dim=-1) sin = torch.cat((sin_half, sin_half), dim=-1) - q_out, k_out = class_b_apply_rotary_pos_emb(q, k, cos, sin) + q_out, k_out = single_round_apply_rotary_pos_emb(q, k, cos, sin) assert q_out.shape == q.shape assert k_out.shape == k.shape if rotary_dim < head_dim: diff --git a/tests/server/runner/test_model_runner_fp8_training.py b/tests/server/runner/test_model_runner_fp8_training.py index ea347157..fe96b5e6 100644 --- a/tests/server/runner/test_model_runner_fp8_training.py +++ b/tests/server/runner/test_model_runner_fp8_training.py @@ -18,7 +18,7 @@ def _fake_model_config() -> SimpleNamespace: "rmsnorm_mode": "eager", "activation_native": False, "rope_native": False, - "rope_class_b": False, + "rope_fp32_single_round": False, "attention_cast_bf16": False, "sparse_mla_enabled": False, "sparse_mla_backend": None, diff --git a/tests/server/test_server_arguments.py b/tests/server/test_server_arguments.py index 80e762bb..9c53c0cc 100644 --- a/tests/server/test_server_arguments.py +++ b/tests/server/test_server_arguments.py @@ -325,7 +325,7 @@ def test_shipped_exact_qwen35_lora_backends_pass_exact_admission(clean_shipped_a rmsnorm_mode=raw.get("rmsnorm_mode"), activation_native=raw.get("activation_native", False), rope_native=raw.get("rope_native"), - rope_class_b=raw.get("rope_class_b"), + rope_fp32_single_round=raw.get("rope_fp32_single_round"), attention_cast_bf16=raw.get("attention_cast_bf16", False), sparse_mla_enabled=raw.get("sparse_mla_enabled"), sparse_mla_backend=raw.get("sparse_mla_backend"), @@ -379,7 +379,7 @@ def test_canonical_moe_and_rope_auto_defaults_serialize(tmp_path): assert model_config["lm_head_fp32"] is None assert model_config["rmsnorm_mode"] is None assert model_config["rope_native"] is None - assert model_config["rope_class_b"] is None + assert model_config["rope_fp32_single_round"] is None assert model_config["sparse_mla_enabled"] is None assert args.to_config_dict()["train"]["ce_mode"] is None diff --git a/tests/trainers/test_qwen3_dense_exact_config.py b/tests/trainers/test_qwen3_dense_exact_config.py index ee3fee3a..41e0b73b 100644 --- a/tests/trainers/test_qwen3_dense_exact_config.py +++ b/tests/trainers/test_qwen3_dense_exact_config.py @@ -49,7 +49,7 @@ def _program(config, **overrides): "rmsnorm_mode": None, "activation_native": False, "rope_native": None, - "rope_class_b": None, + "rope_fp32_single_round": None, "attention_cast_bf16": False, "sparse_mla_enabled": None, "sparse_mla_backend": None, @@ -68,12 +68,12 @@ def test_dense_qwen3_resolves_shared_exact_program(): assert program.rmsnorm_mode == "sglang_fused" assert not program.activation_native assert program.rope_native - assert program.rope_class_b + assert program.rope_fp32_single_round assert resolve_cross_entropy_mode(config, None) == "bi_fused" assert _resolve_rope_modes( config, rope_native=None, - rope_class_b=None, + rope_fp32_single_round=None, ) == (True, True) @@ -85,7 +85,7 @@ def test_dense_qwen3_resolves_shared_exact_program(): ("rmsnorm_mode", "native"), ("activation_native", True), ("rope_native", False), - ("rope_class_b", False), + ("rope_fp32_single_round", False), ], ) def test_dense_qwen3_rejects_numerical_opt_out(name, value): diff --git a/tests/trainers/test_rope_class_b_config.py b/tests/trainers/test_rope_fp32_single_round_config.py similarity index 89% rename from tests/trainers/test_rope_class_b_config.py rename to tests/trainers/test_rope_fp32_single_round_config.py index 51f811e8..d51c5059 100644 --- a/tests/trainers/test_rope_class_b_config.py +++ b/tests/trainers/test_rope_fp32_single_round_config.py @@ -15,7 +15,7 @@ resolve_cross_entropy_mode, resolve_model_numerical_program, ) -from xorl.models.layers.rope import RotaryEmbedding, rope_class_b_enabled, set_rope_class_b +from xorl.models.layers.rope import RotaryEmbedding, rope_fp32_single_round_enabled, set_rope_fp32_single_round from xorl.models.transformers.glm5.configuration_glm5 import Glm5Config from xorl.models.transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config from xorl.models.transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig @@ -72,29 +72,29 @@ def _exact_qwen35_moe_config() -> Qwen3_5MoeConfig: ) -def test_class_b_requires_serving_table_provenance(): - with pytest.raises(ValueError, match="rope_class_b=True requires rope_native=True"): +def test_fp32_single_round_requires_serving_table_provenance(): + with pytest.raises(ValueError, match="rope_fp32_single_round=True requires rope_native=True"): build_training_model( config_path="unused", weights_path="unused", - rope_class_b=True, + rope_fp32_single_round=True, rope_native=False, ) -def test_class_b_selector_can_be_reset(): - set_rope_class_b(True) - assert rope_class_b_enabled() - set_rope_class_b(False) - assert not rope_class_b_enabled() +def test_fp32_single_round_selector_can_be_reset(): + set_rope_fp32_single_round(True) + assert rope_fp32_single_round_enabled() + set_rope_fp32_single_round(False) + assert not rope_fp32_single_round_enabled() -def test_canonical_glm_resolves_native_class_b_without_environment(monkeypatch): +def test_canonical_glm_resolves_native_fp32_single_round_without_environment(monkeypatch): monkeypatch.delenv("XORL_ROPE_CLASS_B", raising=False) config = _exact_glm52_config() - assert _resolve_rope_modes(config, rope_native=None, rope_class_b=None) == (True, True) + assert _resolve_rope_modes(config, rope_native=None, rope_fp32_single_round=None) == (True, True) - set_rope_class_b(False) + set_rope_fp32_single_round(False) rotary = RotaryEmbedding(config) cos, sin = rotary( torch.zeros((1, 1, config.hidden_size), dtype=torch.bfloat16), @@ -103,18 +103,18 @@ def test_canonical_glm_resolves_native_class_b_without_environment(monkeypatch): assert cos.dtype is sin.dtype is torch.float32 -@pytest.mark.parametrize("name", ["rope_native", "rope_class_b"]) -def test_canonical_glm_rejects_explicit_class_b_opt_out(name): +@pytest.mark.parametrize("name", ["rope_native", "rope_fp32_single_round"]) +def test_canonical_glm_rejects_explicit_fp32_single_round_opt_out(name): config = _exact_glm52_config() - kwargs = {"rope_native": None, "rope_class_b": None, name: False} + kwargs = {"rope_native": None, "rope_fp32_single_round": None, name: False} with pytest.raises(ValueError, match="Canonical GLM-5.2 requires native Class-B RoPE"): _resolve_rope_modes(config, **kwargs) def test_non_glm_rope_defaults_and_opt_in_are_unchanged(): config = PretrainedConfig() - assert _resolve_rope_modes(config, rope_native=None, rope_class_b=None) == (False, False) - assert _resolve_rope_modes(config, rope_native=True, rope_class_b=True) == (True, True) + assert _resolve_rope_modes(config, rope_native=None, rope_fp32_single_round=None) == (False, False) + assert _resolve_rope_modes(config, rope_native=True, rope_fp32_single_round=True) == (True, True) def test_canonical_glm_resolves_complete_exact_program(monkeypatch): @@ -137,7 +137,7 @@ def test_canonical_glm_resolves_complete_exact_program(monkeypatch): rmsnorm_mode=None, activation_native=False, rope_native=None, - rope_class_b=None, + rope_fp32_single_round=None, attention_cast_bf16=False, sparse_mla_enabled=None, sparse_mla_backend="auto", @@ -150,7 +150,7 @@ def test_canonical_glm_resolves_complete_exact_program(monkeypatch): qwen35_rmsnorm_family=None, activation_native=False, rope_native=True, - rope_class_b=True, + rope_fp32_single_round=True, attention_cast_bf16=False, sparse_mla_enabled=True, sparse_mla_backend="flashmla", @@ -181,7 +181,7 @@ def test_canonical_glm_rejects_incompatible_numerical_override(override, value): "rmsnorm_mode": None, "activation_native": False, "rope_native": None, - "rope_class_b": None, + "rope_fp32_single_round": None, "attention_cast_bf16": False, "sparse_mla_enabled": None, "sparse_mla_backend": "auto", @@ -208,7 +208,7 @@ def test_non_glm_numerical_defaults_are_preserved(): rmsnorm_mode=None, activation_native=False, rope_native=None, - rope_class_b=None, + rope_fp32_single_round=None, attention_cast_bf16=False, sparse_mla_enabled=None, sparse_mla_backend=None, @@ -221,7 +221,7 @@ def test_non_glm_numerical_defaults_are_preserved(): qwen35_rmsnorm_family=None, activation_native=False, rope_native=False, - rope_class_b=False, + rope_fp32_single_round=False, attention_cast_bf16=False, sparse_mla_enabled=False, sparse_mla_backend="auto", @@ -243,7 +243,7 @@ def test_exact_qwen35_resolves_the_certified_numerical_program(config_factory): rmsnorm_mode=None, activation_native=False, rope_native=None, - rope_class_b=None, + rope_fp32_single_round=None, attention_cast_bf16=False, sparse_mla_enabled=None, sparse_mla_backend=None, @@ -257,7 +257,7 @@ def test_exact_qwen35_resolves_the_certified_numerical_program(config_factory): qwen35_rmsnorm_family="v2", activation_native=True, rope_native=True, - rope_class_b=True, + rope_fp32_single_round=True, attention_cast_bf16=True, sparse_mla_enabled=False, sparse_mla_backend="auto", @@ -266,14 +266,14 @@ def test_exact_qwen35_resolves_the_certified_numerical_program(config_factory): @pytest.mark.parametrize("config_factory", [_exact_qwen35_dense_config, _exact_qwen35_moe_config]) -def test_exact_qwen35_resolves_class_b_and_rejects_opt_out(config_factory): +def test_exact_qwen35_resolves_fp32_single_round_and_rejects_opt_out(config_factory): config = config_factory() config._qwen35_exact_contract = True - default_modes = _resolve_rope_modes(config, rope_native=None, rope_class_b=None) + default_modes = _resolve_rope_modes(config, rope_native=None, rope_fp32_single_round=None) assert default_modes == (True, True) with pytest.raises(ValueError, match="requires native Class-B RoPE"): - _resolve_rope_modes(config, rope_native=None, rope_class_b=False) + _resolve_rope_modes(config, rope_native=None, rope_fp32_single_round=False) program = resolve_model_numerical_program( config, @@ -284,13 +284,13 @@ def test_exact_qwen35_resolves_class_b_and_rejects_opt_out(config_factory): rmsnorm_mode=None, activation_native=False, rope_native=None, - rope_class_b=True, + rope_fp32_single_round=True, attention_cast_bf16=False, sparse_mla_enabled=None, sparse_mla_backend=None, ) assert program.rope_native is True - assert program.rope_class_b is True + assert program.rope_fp32_single_round is True @pytest.mark.parametrize("config_factory", [_exact_qwen35_dense_config, _exact_qwen35_moe_config]) @@ -306,7 +306,7 @@ def test_exact_qwen35_rmsnorm_v2_is_the_architecture_scoped_default(config_facto "qwen35_rmsnorm_family": None, "activation_native": False, "rope_native": None, - "rope_class_b": None, + "rope_fp32_single_round": None, "attention_cast_bf16": False, "sparse_mla_enabled": None, "sparse_mla_backend": None, @@ -330,7 +330,7 @@ def test_non_qwen_rejects_qwen35_rmsnorm_v2(): qwen35_rmsnorm_family="v2", activation_native=False, rope_native=None, - rope_class_b=None, + rope_fp32_single_round=None, attention_cast_bf16=False, sparse_mla_enabled=None, sparse_mla_backend=None, @@ -359,7 +359,7 @@ def test_exact_qwen35_rejects_incompatible_numerical_override(override, value): "rmsnorm_mode": None, "activation_native": False, "rope_native": None, - "rope_class_b": None, + "rope_fp32_single_round": None, "attention_cast_bf16": False, "sparse_mla_enabled": None, "sparse_mla_backend": None, diff --git a/tests/trainers/test_trainer_model_alignment_flags.py b/tests/trainers/test_trainer_model_alignment_flags.py index ed7fca9c..8cf4b76a 100644 --- a/tests/trainers/test_trainer_model_alignment_flags.py +++ b/tests/trainers/test_trainer_model_alignment_flags.py @@ -24,7 +24,7 @@ def __init__(self): "qwen35_rmsnorm_family": "v2", "activation_native": True, "rope_native": True, - "rope_class_b": True, + "rope_fp32_single_round": True, "attention_cast_bf16": True, "sparse_mla_enabled": False, "sparse_mla_backend": "auto", @@ -53,7 +53,7 @@ def _trainer_args(): qwen35_rmsnorm_family="v2", activation_native=True, rope_native=True, - rope_class_b=True, + rope_fp32_single_round=True, attention_cast_bf16=True, sparse_mla_enabled=False, sparse_mla_backend="auto", @@ -111,7 +111,7 @@ def fake_build_foundation_model(**kwargs): assert captured["qwen35_rmsnorm_family"] == "v2" assert captured["activation_native"] is True assert captured["rope_native"] is True - assert captured["rope_class_b"] is True + assert captured["rope_fp32_single_round"] is True assert captured["attention_cast_bf16"] is True assert captured["moe_routing_weights_before_down"] is True assert captured["lora_rank"] == 16 diff --git a/tests/trainers/test_unfuse_for_lora.py b/tests/trainers/test_unfuse_for_lora.py index 74640592..eb2bd761 100644 --- a/tests/trainers/test_unfuse_for_lora.py +++ b/tests/trainers/test_unfuse_for_lora.py @@ -36,7 +36,7 @@ def __init__(self, calls: list[str] | None = None) -> None: "qwen35_rmsnorm_family": None, "activation_native": False, "rope_native": False, - "rope_class_b": False, + "rope_fp32_single_round": False, "attention_cast_bf16": False, "sparse_mla_enabled": False, "sparse_mla_backend": None, @@ -235,7 +235,7 @@ def _trainer(monkeypatch, calls, model): qwen35_rmsnorm_family=None, activation_native=False, rope_native=False, - rope_class_b=False, + rope_fp32_single_round=False, attention_cast_bf16=False, sparse_mla_enabled=False, sparse_mla_backend=None, From a9d5bbee2c0e741a417856de397ccfedc8930e29 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 16:29:18 +0000 Subject: [PATCH 10/14] Rename canonical_moe_{leaf,cast} to moe_fixed_order_reduce_{leaf,cast} (#81 review, row 3) Module and Python-level function names spell out what 'canonical' meant: the pinned fixed-order expert-contribution reduction. The cross-engine contract key CANONICAL_MOE_REDUCE_VERSION and the registered torch.library op names (xorl_k3::_canonical_moe_*) are pinned identifiers and keep their strings. --- src/xorl/distributed/__init__.py | 2 +- src/xorl/distributed/canonical_moe.py | 18 +++---- .../deepseek_v4/native_payload.py | 4 +- .../models/transformers/glm5/modeling_glm5.py | 4 +- src/xorl/ops/canonical_moe_cast.py | 8 ++- src/xorl/ops/canonical_moe_leaf.py | 8 ++- src/xorl/ops/exact/__init__.py | 2 +- ...cast.py => moe_fixed_order_reduce_cast.py} | 50 +++++++++---------- ...leaf.py => moe_fixed_order_reduce_leaf.py} | 40 +++++++-------- .../test_canonical_moe_contract.py | 24 ++++----- tests/ops/test_canonical_moe_cast.py | 4 +- 11 files changed, 86 insertions(+), 78 deletions(-) rename src/xorl/ops/exact/{canonical_moe_cast.py => moe_fixed_order_reduce_cast.py} (66%) rename src/xorl/ops/exact/{canonical_moe_leaf.py => moe_fixed_order_reduce_leaf.py} (80%) diff --git a/src/xorl/distributed/__init__.py b/src/xorl/distributed/__init__.py index 7330b425..be8f47c9 100644 --- a/src/xorl/distributed/__init__.py +++ b/src/xorl/distributed/__init__.py @@ -14,7 +14,7 @@ OutputDistribution, ParallelRole, canonical_moe_fold_fp64_v3, - canonical_moe_leaf_fp32_v1, + moe_fixed_order_leaf_fp32_v1, canonical_moe_reduce_cp_sharded_v3, canonical_moe_reduce_fp64_v3, canonical_moe_reduce_packed_ep16_v2, diff --git a/src/xorl/distributed/canonical_moe.py b/src/xorl/distributed/canonical_moe.py index 63eaa91c..01cde145 100644 --- a/src/xorl/distributed/canonical_moe.py +++ b/src/xorl/distributed/canonical_moe.py @@ -26,12 +26,12 @@ import torch import torch.distributed as dist -from xorl.ops.exact.canonical_moe_leaf import canonical_moe_leaf_fp32_v1_op +from xorl.ops.exact.moe_fixed_order_reduce_leaf import moe_fixed_order_leaf_fp32_v1_op CANONICAL_MOE_FOLD_VERSION = "canonical_moe_fold_fp64_v3" CANONICAL_MOE_REDUCE_VERSION = "canonical_moe_reduce_fp64_v3" -CANONICAL_MOE_LEAF_VERSION = "canonical_moe_leaf_fp32_v1" +CANONICAL_MOE_LEAF_VERSION = "moe_fixed_order_leaf_fp32_v1" CANONICAL_MOE_PACKED_EP16_TRANSPORT_VERSION = "packed_ep16_v2" CANONICAL_MOE_CP_SHARDED_TRANSPORT_VERSION = "cp_sharded_v3" GLM52_NUM_LAYERS = 78 @@ -616,7 +616,7 @@ def _validate_runtime_plan( ) -def canonical_moe_leaf_fp32_v1( +def moe_fixed_order_leaf_fp32_v1( shared: torch.Tensor, routed: torch.Tensor, routed_scale: float = 1.0, @@ -629,7 +629,7 @@ def canonical_moe_leaf_fp32_v1( primitive prevents compiler reassociation without materializing an FP32 output tensor. """ - return canonical_moe_leaf_fp32_v1_op(shared, routed, routed_scale) + return moe_fixed_order_leaf_fp32_v1_op(shared, routed, routed_scale) def _canonical_moe_fold_fp64_tree(partials_by_logical_ordinal: torch.Tensor) -> torch.Tensor: @@ -659,7 +659,7 @@ def _canonical_moe_fold_fp64_tree(partials_by_logical_ordinal: torch.Tensor) -> return level[0] -def _canonical_moe_fp64_fold_chunk_elements(contributor_count: int) -> int: +def _moe_fixed_order_fp64_fold_chunk_elements(contributor_count: int) -> int: """Maximum payload elements whose initial FP64 level fits the bound.""" if contributor_count <= 0: raise ValueError("Canonical MoE fold requires a positive contributor count") @@ -677,9 +677,9 @@ def _canonical_moe_cast_fp64_to_transport( if folded_fp64.dtype is not torch.float64: raise TypeError(f"Canonical MoE final cast requires FP64 input, got {folded_fp64.dtype}") if folded_fp64.is_cuda and transport_dtype in (torch.bfloat16, torch.float16): - from xorl.ops.exact.canonical_moe_cast import canonical_moe_fp64_to_lowp_rne # noqa: PLC0415 + from xorl.ops.exact.moe_fixed_order_reduce_cast import moe_fixed_order_fp64_to_lowp_rne # noqa: PLC0415 - return canonical_moe_fp64_to_lowp_rne(folded_fp64.contiguous(), transport_dtype) + return moe_fixed_order_fp64_to_lowp_rne(folded_fp64.contiguous(), transport_dtype) return folded_fp64.to(transport_dtype) @@ -704,7 +704,7 @@ def canonical_moe_fold_fp64_v3(partials_by_logical_ordinal: torch.Tensor) -> tor payload_shape = partials.shape[1:] payload_elements = partials[0].numel() - chunk_elements = _canonical_moe_fp64_fold_chunk_elements(contributor_count) + chunk_elements = _moe_fixed_order_fp64_fold_chunk_elements(contributor_count) if payload_elements <= chunk_elements: return _canonical_moe_cast_fp64_to_transport( _canonical_moe_fold_fp64_tree(partials), @@ -1223,7 +1223,7 @@ def canonical_moe_reduce_cp_sharded_v3( "ParallelPlan", "ParallelRole", "canonical_moe_fold_fp64_v3", - "canonical_moe_leaf_fp32_v1", + "moe_fixed_order_leaf_fp32_v1", "canonical_moe_reduce_reference", "canonical_moe_reduce_packed_ep16_v2", "canonical_moe_reduce_cp_sharded_v3", diff --git a/src/xorl/models/transformers/deepseek_v4/native_payload.py b/src/xorl/models/transformers/deepseek_v4/native_payload.py index c74bdeb7..6921f2a3 100644 --- a/src/xorl/models/transformers/deepseek_v4/native_payload.py +++ b/src/xorl/models/transformers/deepseek_v4/native_payload.py @@ -1464,12 +1464,12 @@ def dsv4_native_shared_expert_tp_partial( class _Dsv4RoutedSharedJoin(torch.autograd.Function): @staticmethod def forward(ctx, routed: torch.Tensor, shared: torch.Tensor, scale: float): - from xorl.distributed.canonical_moe import canonical_moe_leaf_fp32_v1 # noqa: PLC0415 + from xorl.distributed.canonical_moe import moe_fixed_order_leaf_fp32_v1 # noqa: PLC0415 ctx.scale = scale ctx.routed_dtype = routed.dtype ctx.shared_dtype = shared.dtype - return canonical_moe_leaf_fp32_v1(shared, routed, routed_scale=scale) + return moe_fixed_order_leaf_fp32_v1(shared, routed, routed_scale=scale) @staticmethod def backward(ctx, grad_output): diff --git a/src/xorl/models/transformers/glm5/modeling_glm5.py b/src/xorl/models/transformers/glm5/modeling_glm5.py index 0f6341b1..f806630a 100644 --- a/src/xorl/models/transformers/glm5/modeling_glm5.py +++ b/src/xorl/models/transformers/glm5/modeling_glm5.py @@ -15,7 +15,7 @@ LogicalRowOwnership, OutputDistribution, ParallelPlan, - canonical_moe_leaf_fp32_v1, + moe_fixed_order_leaf_fp32_v1, canonical_moe_reduce_cp_sharded_v3, canonical_moe_reduce_fp64_v3, canonical_moe_reduce_packed_ep16_v2, @@ -1201,7 +1201,7 @@ def _canonical_ep_forward( contributor_ordinal=ep_rank, contributor_count=ps.ep_size, ) - local_partial = canonical_moe_leaf_fp32_v1(shared, routed) + local_partial = moe_fixed_order_leaf_fp32_v1(shared, routed) capacity = int(getattr(self.config, "_glm52_canonical_moe_capacity", local_partial.shape[0])) if local_partial.shape[0] > capacity: diff --git a/src/xorl/ops/canonical_moe_cast.py b/src/xorl/ops/canonical_moe_cast.py index 4302cf79..78eb55c2 100644 --- a/src/xorl/ops/canonical_moe_cast.py +++ b/src/xorl/ops/canonical_moe_cast.py @@ -1,3 +1,7 @@ -"""Moved to ``xorl.ops.exact.canonical_moe_cast`` (issue #78 phase 3); compat stub.""" +"""Moved to ``xorl.ops.exact.moe_fixed_order_reduce_cast`` (issue #78); alias stub.""" -from xorl.ops.exact.canonical_moe_cast import * # noqa: F401,F403 +import importlib as _importlib +import sys as _sys + + +_sys.modules[__name__] = _importlib.import_module("xorl.ops.exact.moe_fixed_order_reduce_cast") diff --git a/src/xorl/ops/canonical_moe_leaf.py b/src/xorl/ops/canonical_moe_leaf.py index 5024c8fd..c512aed4 100644 --- a/src/xorl/ops/canonical_moe_leaf.py +++ b/src/xorl/ops/canonical_moe_leaf.py @@ -1,3 +1,7 @@ -"""Moved to ``xorl.ops.exact.canonical_moe_leaf`` (issue #78 phase 3); compat stub.""" +"""Moved to ``xorl.ops.exact.moe_fixed_order_reduce_leaf`` (issue #78); alias stub.""" -from xorl.ops.exact.canonical_moe_leaf import * # noqa: F401,F403 +import importlib as _importlib +import sys as _sys + + +_sys.modules[__name__] = _importlib.import_module("xorl.ops.exact.moe_fixed_order_reduce_leaf") diff --git a/src/xorl/ops/exact/__init__.py b/src/xorl/ops/exact/__init__.py index 2ce7ea88..11912ff1 100644 --- a/src/xorl/ops/exact/__init__.py +++ b/src/xorl/ops/exact/__init__.py @@ -5,7 +5,7 @@ with the serving engine. Here: ``sampling_transforms`` (the replay contract), ``rope_fp32_single_round``, -``canonical_moe_leaf``, ``canonical_moe_cast``, ``kernel_config_pin``, +``moe_fixed_order_leaf``, ``canonical_moe_cast``, ``kernel_config_pin``, ``bi_gemm_configs``, ``block_fp8_native``, ``fused_silu_and_mul``. The modules that exist as literal twins inside the serving engine diff --git a/src/xorl/ops/exact/canonical_moe_cast.py b/src/xorl/ops/exact/moe_fixed_order_reduce_cast.py similarity index 66% rename from src/xorl/ops/exact/canonical_moe_cast.py rename to src/xorl/ops/exact/moe_fixed_order_reduce_cast.py index f58e8b56..d0b59139 100644 --- a/src/xorl/ops/exact/canonical_moe_cast.py +++ b/src/xorl/ops/exact/moe_fixed_order_reduce_cast.py @@ -14,7 +14,7 @@ _TRITON_AVAILABLE = False -def validate_canonical_moe_fp64_cast_input( +def validate_moe_fixed_order_fp64_cast_input( value: torch.Tensor, output_dtype: torch.dtype, ) -> None: @@ -31,7 +31,7 @@ def validate_canonical_moe_fp64_cast_input( if _TRITON_AVAILABLE: @triton.jit - def _canonical_moe_fp64_to_lowp_rne_kernel( + def _moe_fixed_order_fp64_to_lowp_rne_kernel( input_ptr, output_ptr, n_elements, @@ -46,7 +46,7 @@ def _canonical_moe_fp64_to_lowp_rne_kernel( tl.store(output_ptr + offsets, value, mask=mask) -def _canonical_moe_fp64_to_lowp_cuda( +def _moe_fixed_order_fp64_to_lowp_cuda( value: torch.Tensor, output_dtype: torch.dtype, ) -> torch.Tensor: @@ -56,7 +56,7 @@ def _canonical_moe_fp64_to_lowp_cuda( if output.numel() == 0: return output block_size = 256 - _canonical_moe_fp64_to_lowp_rne_kernel[(triton.cdiv(output.numel(), block_size),)]( + _moe_fixed_order_fp64_to_lowp_rne_kernel[(triton.cdiv(output.numel(), block_size),)]( value, output, output.numel(), @@ -67,30 +67,30 @@ def _canonical_moe_fp64_to_lowp_cuda( @torch.library.custom_op("xorl_k3::_canonical_moe_fp64_to_bf16_rne", mutates_args=()) -def _canonical_moe_fp64_to_bf16_rne_op(value: torch.Tensor) -> torch.Tensor: +def _moe_fixed_order_fp64_to_bf16_rne_op(value: torch.Tensor) -> torch.Tensor: """Round FP64 directly to BF16.""" - validate_canonical_moe_fp64_cast_input(value, torch.bfloat16) - return _canonical_moe_fp64_to_lowp_cuda(value, torch.bfloat16) + validate_moe_fixed_order_fp64_cast_input(value, torch.bfloat16) + return _moe_fixed_order_fp64_to_lowp_cuda(value, torch.bfloat16) -@_canonical_moe_fp64_to_bf16_rne_op.register_fake -def _canonical_moe_fp64_to_bf16_rne_fake(value: torch.Tensor) -> torch.Tensor: +@_moe_fixed_order_fp64_to_bf16_rne_op.register_fake +def _moe_fixed_order_fp64_to_bf16_rne_fake(value: torch.Tensor) -> torch.Tensor: return torch.empty_like(value, dtype=torch.bfloat16) @torch.library.custom_op("xorl_k3::_canonical_moe_fp64_to_fp16_rne", mutates_args=()) -def _canonical_moe_fp64_to_fp16_rne_op(value: torch.Tensor) -> torch.Tensor: +def _moe_fixed_order_fp64_to_fp16_rne_op(value: torch.Tensor) -> torch.Tensor: """Round FP64 directly to FP16.""" - validate_canonical_moe_fp64_cast_input(value, torch.float16) - return _canonical_moe_fp64_to_lowp_cuda(value, torch.float16) + validate_moe_fixed_order_fp64_cast_input(value, torch.float16) + return _moe_fixed_order_fp64_to_lowp_cuda(value, torch.float16) -@_canonical_moe_fp64_to_fp16_rne_op.register_fake -def _canonical_moe_fp64_to_fp16_rne_fake(value: torch.Tensor) -> torch.Tensor: +@_moe_fixed_order_fp64_to_fp16_rne_op.register_fake +def _moe_fixed_order_fp64_to_fp16_rne_fake(value: torch.Tensor) -> torch.Tensor: return torch.empty_like(value, dtype=torch.float16) -def _canonical_moe_fp64_to_lowp_backward(ctx, grad_output: torch.Tensor): +def _moe_fixed_order_fp64_to_lowp_backward(ctx, grad_output: torch.Tensor): del ctx # Match torch's straight-through cast derivative. The surrounding FP64 # tree then distributes this FP64 gradient to every contributor before its @@ -98,26 +98,26 @@ def _canonical_moe_fp64_to_lowp_backward(ctx, grad_output: torch.Tensor): return grad_output.to(torch.float64) -_canonical_moe_fp64_to_bf16_rne_op.register_autograd( - _canonical_moe_fp64_to_lowp_backward, +_moe_fixed_order_fp64_to_bf16_rne_op.register_autograd( + _moe_fixed_order_fp64_to_lowp_backward, ) -_canonical_moe_fp64_to_fp16_rne_op.register_autograd( - _canonical_moe_fp64_to_lowp_backward, +_moe_fixed_order_fp64_to_fp16_rne_op.register_autograd( + _moe_fixed_order_fp64_to_lowp_backward, ) -def canonical_moe_fp64_to_lowp_rne( +def moe_fixed_order_fp64_to_lowp_rne( value: torch.Tensor, output_dtype: torch.dtype, ) -> torch.Tensor: """Validate and directly round FP64 to the canonical transport dtype.""" - validate_canonical_moe_fp64_cast_input(value, output_dtype) + validate_moe_fixed_order_fp64_cast_input(value, output_dtype) if output_dtype is torch.bfloat16: - return _canonical_moe_fp64_to_bf16_rne_op(value) - return _canonical_moe_fp64_to_fp16_rne_op(value) + return _moe_fixed_order_fp64_to_bf16_rne_op(value) + return _moe_fixed_order_fp64_to_fp16_rne_op(value) __all__ = [ - "canonical_moe_fp64_to_lowp_rne", - "validate_canonical_moe_fp64_cast_input", + "moe_fixed_order_fp64_to_lowp_rne", + "validate_moe_fixed_order_fp64_cast_input", ] diff --git a/src/xorl/ops/exact/canonical_moe_leaf.py b/src/xorl/ops/exact/moe_fixed_order_reduce_leaf.py similarity index 80% rename from src/xorl/ops/exact/canonical_moe_leaf.py rename to src/xorl/ops/exact/moe_fixed_order_reduce_leaf.py index 1165471a..4dd59443 100644 --- a/src/xorl/ops/exact/canonical_moe_leaf.py +++ b/src/xorl/ops/exact/moe_fixed_order_reduce_leaf.py @@ -14,7 +14,7 @@ _TRITON_AVAILABLE = False -def validate_canonical_moe_leaf_operands( +def validate_moe_fixed_order_leaf_operands( shared: torch.Tensor, routed: torch.Tensor, ) -> None: @@ -51,7 +51,7 @@ def _fp32_fma_rn(multiplicand, multiplier, addend): ) @triton.jit - def _canonical_moe_leaf_fp32_kernel( + def _moe_fixed_order_leaf_fp32_kernel( shared_ptr, routed_ptr, output_ptr, @@ -69,7 +69,7 @@ def _canonical_moe_leaf_fp32_kernel( tl.store(output_ptr + offsets, leaf, mask=mask) -def _canonical_moe_leaf_cuda( +def _moe_fixed_order_leaf_cuda( shared: torch.Tensor, routed: torch.Tensor, routed_scale: float, @@ -80,7 +80,7 @@ def _canonical_moe_leaf_cuda( if output.numel() == 0: return output block_size = 256 - _canonical_moe_leaf_fp32_kernel[(triton.cdiv(output.numel(), block_size),)]( + _moe_fixed_order_leaf_fp32_kernel[(triton.cdiv(output.numel(), block_size),)]( shared, routed, output, @@ -92,7 +92,7 @@ def _canonical_moe_leaf_cuda( return output -def _canonical_moe_leaf_cpu( +def _moe_fixed_order_leaf_cpu( shared: torch.Tensor, routed: torch.Tensor, routed_scale: float, @@ -108,20 +108,20 @@ def _canonical_moe_leaf_cpu( @torch.library.custom_op("xorl_k3::_canonical_moe_leaf_fp32_v1", mutates_args=()) -def _canonical_moe_leaf_fp32_v1_op( +def _moe_fixed_order_leaf_fp32_v1_op( shared: torch.Tensor, routed: torch.Tensor, routed_scale: float, ) -> torch.Tensor: """Opaque one-FMA leaf primitive; output is the transport dtype.""" - validate_canonical_moe_leaf_operands(shared, routed) + validate_moe_fixed_order_leaf_operands(shared, routed) if shared.is_cuda: - return _canonical_moe_leaf_cuda(shared, routed, routed_scale) - return _canonical_moe_leaf_cpu(shared, routed, routed_scale) + return _moe_fixed_order_leaf_cuda(shared, routed, routed_scale) + return _moe_fixed_order_leaf_cpu(shared, routed, routed_scale) -@_canonical_moe_leaf_fp32_v1_op.register_fake -def _canonical_moe_leaf_fp32_v1_fake( +@_moe_fixed_order_leaf_fp32_v1_op.register_fake +def _moe_fixed_order_leaf_fp32_v1_fake( shared: torch.Tensor, routed: torch.Tensor, routed_scale: float, @@ -130,34 +130,34 @@ def _canonical_moe_leaf_fp32_v1_fake( return torch.empty_like(shared) -def _canonical_moe_leaf_setup_context(ctx, inputs, output) -> None: +def _moe_fixed_order_leaf_setup_context(ctx, inputs, output) -> None: del output shared, _routed, routed_scale = inputs ctx.input_dtype = shared.dtype ctx.routed_scale = float(routed_scale) -def _canonical_moe_leaf_backward(ctx, grad_output: torch.Tensor): +def _moe_fixed_order_leaf_backward(ctx, grad_output: torch.Tensor): grad_fp32 = grad_output.float() grad_shared = grad_fp32.to(ctx.input_dtype) grad_routed = (grad_fp32 * ctx.routed_scale).to(ctx.input_dtype) return grad_shared, grad_routed, None -_canonical_moe_leaf_fp32_v1_op.register_autograd( - _canonical_moe_leaf_backward, - setup_context=_canonical_moe_leaf_setup_context, +_moe_fixed_order_leaf_fp32_v1_op.register_autograd( + _moe_fixed_order_leaf_backward, + setup_context=_moe_fixed_order_leaf_setup_context, ) -def canonical_moe_leaf_fp32_v1_op( +def moe_fixed_order_leaf_fp32_v1_op( shared: torch.Tensor, routed: torch.Tensor, routed_scale: float, ) -> torch.Tensor: """Validate and execute the compile-opaque canonical contributor leaf.""" - validate_canonical_moe_leaf_operands(shared, routed) - return _canonical_moe_leaf_fp32_v1_op(shared, routed, float(routed_scale)) + validate_moe_fixed_order_leaf_operands(shared, routed) + return _moe_fixed_order_leaf_fp32_v1_op(shared, routed, float(routed_scale)) -__all__ = ["canonical_moe_leaf_fp32_v1_op", "validate_canonical_moe_leaf_operands"] +__all__ = ["moe_fixed_order_leaf_fp32_v1_op", "validate_moe_fixed_order_leaf_operands"] diff --git a/tests/distributed/test_canonical_moe_contract.py b/tests/distributed/test_canonical_moe_contract.py index 069843be..bbbcd072 100644 --- a/tests/distributed/test_canonical_moe_contract.py +++ b/tests/distributed/test_canonical_moe_contract.py @@ -23,12 +23,12 @@ ParallelPlan, ParallelRole, _canonical_moe_fold_fp64_tree, - _canonical_moe_fp64_fold_chunk_elements, + _moe_fixed_order_fp64_fold_chunk_elements, _resolve_transport_chunk_rows, _RuntimePlan, _transport_and_fold, canonical_moe_fold_fp64_v3, - canonical_moe_leaf_fp32_v1, + moe_fixed_order_leaf_fp32_v1, canonical_moe_reduce_cp_sharded_v3, canonical_moe_reduce_fp64_v3, canonical_moe_reduce_packed_ep16_v2, @@ -184,7 +184,7 @@ def test_chunked_fp64_tree_is_bitwise_exact_across_payload_boundaries( witness_count = min(4, contributors) partials[:witness_count, 0, :4] = torch.tensor([[4096.0], [1.0], [-4096.0], [1.0]], dtype=dtype)[:witness_count] - chunk_elements = _canonical_moe_fp64_fold_chunk_elements(contributors) + chunk_elements = _moe_fixed_order_fp64_fold_chunk_elements(contributors) assert chunk_elements < rows * payload actual = canonical_moe_fold_fp64_v3(partials) expected = _explicit_tree(partials) @@ -244,7 +244,7 @@ def test_chunked_fp64_tree_strided_payload_is_bitwise_exact( @pytest.mark.cpu @pytest.mark.parametrize("contributors", [1, 3, 16, 17, 64]) def test_fp64_fold_chunk_planner_bounds_the_initial_level(contributors: int): - chunk_elements = _canonical_moe_fp64_fold_chunk_elements(contributors) + chunk_elements = _moe_fixed_order_fp64_fold_chunk_elements(contributors) level_bytes = contributors * chunk_elements * torch.float64.itemsize assert 0 < level_bytes <= _CANONICAL_MOE_FP64_FOLD_MAX_LEVEL_BYTES @@ -345,17 +345,17 @@ def test_leaf_uses_compile_stable_one_round_fp32_fma_before_transport_cast( transport_bits: int, dynamic: bool, ): - assert CANONICAL_MOE_LEAF_VERSION == "canonical_moe_leaf_fp32_v1" + assert CANONICAL_MOE_LEAF_VERSION == "moe_fixed_order_leaf_fp32_v1" shared = torch.tensor([shared_value], dtype=dtype) routed = torch.tensor([routed_value], dtype=dtype) fma_oracle = _one_round_leaf_oracle(shared, routed, scale) scale_fp32 = torch.tensor(scale, dtype=torch.float32) separately_rounded = shared.float() + routed.float() * scale_fp32 - eager = canonical_moe_leaf_fp32_v1(shared, routed, scale) + eager = moe_fixed_order_leaf_fp32_v1(shared, routed, scale) def leaf_fn(shared_arg: torch.Tensor, routed_arg: torch.Tensor) -> torch.Tensor: - return canonical_moe_leaf_fp32_v1(shared_arg, routed_arg, scale) + return moe_fixed_order_leaf_fp32_v1(shared_arg, routed_arg, scale) compiled = torch.compile(leaf_fn, fullgraph=True, dynamic=dynamic)(shared, routed) @@ -380,7 +380,7 @@ def test_leaf_autograd_matches_declared_fp32_scale_under_compile( grad_output = torch.tensor([0.5, -1.25, 2.0], dtype=dtype) def leaf_fn(shared_arg: torch.Tensor, routed_arg: torch.Tensor) -> torch.Tensor: - return canonical_moe_leaf_fp32_v1(shared_arg, routed_arg, scale) + return moe_fixed_order_leaf_fp32_v1(shared_arg, routed_arg, scale) compiled = torch.compile(leaf_fn, fullgraph=True, dynamic=dynamic) output = compiled(shared, routed) @@ -435,7 +435,7 @@ def test_cuda_leaf_matches_one_round_oracle_under_compile( ) def leaf_fn(shared_arg: torch.Tensor, routed_arg: torch.Tensor) -> torch.Tensor: - return canonical_moe_leaf_fp32_v1(shared_arg, routed_arg, scale) + return moe_fixed_order_leaf_fp32_v1(shared_arg, routed_arg, scale) eager = leaf_fn(shared, routed) compiled = torch.compile(leaf_fn, fullgraph=True, dynamic=dynamic)(shared, routed) @@ -461,12 +461,12 @@ def test_cuda_leaf_replays_in_cuda_graph_without_fp32_output(dtype: torch.dtype) shared = torch.tensor([0.25, -0.5, 1.0], device="cuda", dtype=dtype) routed = torch.tensor([-2.0, 0.75, 4.0], device="cuda", dtype=dtype) scale = -1.375 - canonical_moe_leaf_fp32_v1(shared, routed, scale) + moe_fixed_order_leaf_fp32_v1(shared, routed, scale) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - output = canonical_moe_leaf_fp32_v1(shared, routed, scale) + output = moe_fixed_order_leaf_fp32_v1(shared, routed, scale) first = output.clone() routed.copy_(torch.tensor([3.0, -1.0, 0.5], device="cuda", dtype=dtype)) graph.replay() @@ -537,7 +537,7 @@ def test_shared_fold_replays_in_cuda_graph(monkeypatch: pytest.MonkeyPatch, cont # production-sized. The production-shape allocator bound is tested below. monkeypatch.setattr(canonical_moe, "_CANONICAL_MOE_FP64_FOLD_MAX_LEVEL_BYTES", 32 * 1024) partials = torch.randn((contributors, 64, 32), device="cuda", dtype=torch.bfloat16) - assert _canonical_moe_fp64_fold_chunk_elements(contributors) < partials[0].numel() + assert _moe_fixed_order_fp64_fold_chunk_elements(contributors) < partials[0].numel() canonical_moe_fold_fp64_v3(partials) torch.cuda.synchronize() diff --git a/tests/ops/test_canonical_moe_cast.py b/tests/ops/test_canonical_moe_cast.py index 39511e5c..bfc56b58 100644 --- a/tests/ops/test_canonical_moe_cast.py +++ b/tests/ops/test_canonical_moe_cast.py @@ -3,7 +3,7 @@ import pytest import torch -from xorl.ops.exact.canonical_moe_cast import canonical_moe_fp64_to_lowp_rne +from xorl.ops.exact.moe_fixed_order_reduce_cast import moe_fixed_order_fp64_to_lowp_rne @pytest.mark.gpu @@ -23,7 +23,7 @@ def test_cuda_direct_rne_result_and_backward( ): value = torch.tensor([midpoint_witness], device="cuda", dtype=torch.float64, requires_grad=True) - output = canonical_moe_fp64_to_lowp_rne(value, output_dtype) + output = moe_fixed_order_fp64_to_lowp_rne(value, output_dtype) assert int(output.view(torch.uint16).item()) == expected_bits assert int(value.detach().to(output_dtype).view(torch.uint16).item()) == double_rounded_bits From 65b6323c07ea53e239cb81ea5187bbee04a7ff0c Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 16:31:07 +0000 Subject: [PATCH 11/14] Rename exact_fp32_silu_and_mul to one_round_swiglu (#81 review, row 5) Matches the existing _exact_one_round_swiglu resolution stamp vocabulary; a deprecated module-level alias keeps the old callable importable. --- .../transformers/glm5/exact_dense_mlp.py | 4 ++-- .../transformers/glm5/exact_fullparam_fp8.py | 4 ++-- .../glm5/exact_shared_expert_qlora.py | 6 ++--- .../models/transformers/glm5/modeling_glm5.py | 6 ++--- .../models/transformers/glm5/native_fp8.py | 4 ++-- .../transformers/qwen3_5/modeling_qwen3_5.py | 4 ++-- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 6 ++--- src/xorl/ops/README.md | 2 +- src/xorl/ops/exact/fused_silu_and_mul.py | 24 +++++++++++-------- tests/models/test_glm52_exact_dense_mlp.py | 2 +- .../test_glm52_exact_dense_mlp_composition.py | 8 +++---- .../models/test_glm52_exact_fullparam_fp8.py | 6 ++--- .../models/test_glm52_exact_gate_up_qlora.py | 4 ++-- .../test_glm52_exact_shared_expert_qlora.py | 4 ++-- tests/models/test_glm5_one_round_swiglu.py | 16 ++++++------- tests/models/test_op_parity_dense.py | 4 ++-- ...est_qwen35_exact_swiglu_grad_engagement.py | 2 +- tests/ops/test_fused_silu_and_mul_scoping.py | 6 ++--- ...lu_and_mul.py => test_one_round_swiglu.py} | 8 +++---- 19 files changed, 62 insertions(+), 58 deletions(-) rename tests/ops/{test_exact_fp32_silu_and_mul.py => test_one_round_swiglu.py} (89%) diff --git a/src/xorl/models/transformers/glm5/exact_dense_mlp.py b/src/xorl/models/transformers/glm5/exact_dense_mlp.py index bfa33685..ea11aca5 100644 --- a/src/xorl/models/transformers/glm5/exact_dense_mlp.py +++ b/src/xorl/models/transformers/glm5/exact_dense_mlp.py @@ -16,7 +16,7 @@ ) from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu GLM52_EXACT_TP1_DENSE_MLP_CONTRACT_VERSION = "glm52_exact_tp1_dense_mlp_qlora_v2" @@ -130,7 +130,7 @@ def forward(self, input: Tensor) -> Tensor: gate_up = Glm52ExactTP1FusedGateUpBlockFP8QLoRA.forward(self, input) # Serving's exact mode computes the one-round FP32 SwiGLU # (SiluAndMul.forward_exact, xorl-sglang f10b907d8). - activated = exact_fp32_silu_and_mul(gate_up) + activated = one_round_swiglu(gate_up) return self.down_proj(activated) diff --git a/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py b/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py index 83da20a5..e3a45ee0 100644 --- a/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py +++ b/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py @@ -30,7 +30,7 @@ pack_fp8_as_float32, unpack_float32_as_fp8, ) -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu logger = logging.getLogger(__name__) @@ -663,7 +663,7 @@ def forward(self, input: Tensor) -> Tensor: # SiLU and multiply in FP32 with one final rounding. The historical # trainer op rounded between SiLU and multiply, so it is not the same # numerical program. - activated = exact_fp32_silu_and_mul(gate_up) + activated = one_round_swiglu(gate_up) return self.down_proj(activated) diff --git a/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py b/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py index 882de8c2..db37e691 100644 --- a/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py @@ -26,7 +26,7 @@ from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear, _sglang_native_block_fp8_linear_value -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu GLM52_EXACT_TP16_SHARED_EXPERT_QLORA_CONTRACT_VERSION = "glm52_exact_tp16_shared_expert_qlora_v2" @@ -521,7 +521,7 @@ def _exact_forward_value( # Exact SGLang target mode resolves SiluAndMul.forward_exact to the # one-round FP32 SwiGLU (xorl-sglang f10b907d8); this site must # produce those bytes. - activated = exact_fp32_silu_and_mul(gate_up) + activated = one_round_swiglu(gate_up) down_base = _sglang_native_block_fp8_linear_value( activated, @@ -674,7 +674,7 @@ def _surrogate_vjp( gate_up_input = exact_gate_up.detach().requires_grad_(True) # VJP reference differentiates the same one-round program the # forward emits (backward stays trainer-owned numerics). - activation = exact_fp32_silu_and_mul(gate_up_input) + activation = one_round_swiglu(gate_up_input) (gate_up_grad,) = torch.autograd.grad( activation, gate_up_input, diff --git a/src/xorl/models/transformers/glm5/modeling_glm5.py b/src/xorl/models/transformers/glm5/modeling_glm5.py index f806630a..de12a344 100644 --- a/src/xorl/models/transformers/glm5/modeling_glm5.py +++ b/src/xorl/models/transformers/glm5/modeling_glm5.py @@ -70,7 +70,7 @@ from xorl.models.transformers.glm5.sparse_mla import sparse_mla_dispatch from xorl.models.transformers.glm5.support import validate_glm5_sequence_parallel from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu, fused_silu_and_mul from xorl.utils import logging @@ -140,7 +140,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: gate = self.gate_proj(x) up = self.up_proj(x) if self._exact_one_round: - hidden_states = exact_fp32_silu_and_mul(torch.cat([gate, up], dim=-1)) + hidden_states = one_round_swiglu(torch.cat([gate, up], dim=-1)) elif self._use_fused_silu: hidden_states = fused_silu_and_mul(torch.cat([gate, up], dim=-1)) else: @@ -1103,7 +1103,7 @@ def _canonical_shared_local_partial( # Serving computes the exact-mode shared expert with the one-round # FP32 SwiGLU (SiluAndMul.forward_exact since xorl-sglang # f10b907d8); the canonical partial must produce those bytes. - activated = exact_fp32_silu_and_mul(torch.cat([gate, up], dim=-1)) + activated = one_round_swiglu(torch.cat([gate, up], dim=-1)) else: activated = F.silu(gate) * up if isinstance(self.shared_experts.down_proj, NativeBlockFP8Linear): diff --git a/src/xorl/models/transformers/glm5/native_fp8.py b/src/xorl/models/transformers/glm5/native_fp8.py index 128a7dbc..725ce8b5 100644 --- a/src/xorl/models/transformers/glm5/native_fp8.py +++ b/src/xorl/models/transformers/glm5/native_fp8.py @@ -1003,10 +1003,10 @@ def checkpoint_split_bytes(self) -> tuple[torch.Tensor, ...]: ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul # noqa: PLC0415 + from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu # noqa: PLC0415 gate_up = self.gate_up_proj(hidden_states) - activated = exact_fp32_silu_and_mul(gate_up) + activated = one_round_swiglu(gate_up) return self.down_proj(activated) diff --git a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py index ca4199af..0fdd7e21 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -43,7 +43,7 @@ has_linear_attention_layers, qwen3_5_apply_rotary_pos_emb, ) -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu, fused_silu_and_mul from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils import logging @@ -104,7 +104,7 @@ def __init__(self, config): def _fused_act(self, gate_up): if self._exact_one_round: - return exact_fp32_silu_and_mul(gate_up) + return one_round_swiglu(gate_up) return fused_silu_and_mul(gate_up) def unfuse_for_tp(self): diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index f6817c85..1f2dd970 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -44,7 +44,7 @@ has_linear_attention_layers, qwen3_5_apply_rotary_pos_emb, ) -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu, fused_silu_and_mul from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils import logging @@ -105,7 +105,7 @@ def __init__(self, config, intermediate_size=None): def _fused_act(self, gate_up): if self._exact_one_round: - return exact_fp32_silu_and_mul(gate_up) + return one_round_swiglu(gate_up) return fused_silu_and_mul(gate_up) def unfuse_for_tp(self): @@ -528,7 +528,7 @@ def _ep_combine_native( self._capture_diagnostic_component("moe_native_shared_gate_up", gate_up) # Exact serving-value path: one-round FP32, paired with serving's # fp32_silu_and_mul (in-scope for the exact contract by construction). - act = exact_fp32_silu_and_mul(gate_up) + act = one_round_swiglu(gate_up) self._capture_diagnostic_component("moe_native_shared_act", act) down = _BatchInvariantTrunkLinearFn.apply(act, w_down[:, lo_s : lo_s + shard].contiguous(), None) self._capture_diagnostic_component("moe_native_shared_down", down) diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index 077cf802..ddcc04ef 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -57,7 +57,7 @@ serving engine and cannot be renamed unilaterally): | **K3 / "zero-K3"** | the k3 KL-divergence estimator between trainer and sampler logprobs for the same tokens; "zero-K3" = bit-identical train/serve forward, the goal of the exact contracts. | | **Class-A / Class-B RoPE** | the two RoPE numerics classes across the trainer/sampler pair: Class A rounds to bf16 per op (8 rounding points); Class B computes one fp32 chain with a single final round (SGLang's fused CUDA rope and its compiled RL-lane path). `rope_fp32_single_round: true` selects Class B. | | **canonical MoE reduce** | the pinned fixed-order expert-contribution reduction (contributor-leaf arithmetic + final FP64-accumulator cast) shared with serving, versioned by `CANONICAL_MOE_REDUCE_VERSION`. | -| **one-round SwiGLU** | the exact-contract SwiGLU with a single FP32 rounding point (`exact_fp32_silu_and_mul`), vs the generic fused SwiGLU. | +| **one-round SwiGLU** | the exact-contract SwiGLU with a single FP32 rounding point (`one_round_swiglu`), vs the generic fused SwiGLU. | | **families v1 / v2** | versioned batch-invariant kernel families (norms, LM head); v2 is the epilogue-stats generation. `bi_families_v2` is the serving twin module. | | **GKN layout** | grouped expert-weight layout `[G=experts, K=in_features, N=out_features]`. | | **exact** | shorthand for "serving-parity byte contract": the forward reproduces the serving engine's bits, not just its math. | diff --git a/src/xorl/ops/exact/fused_silu_and_mul.py b/src/xorl/ops/exact/fused_silu_and_mul.py index 6122c025..99d9db5f 100644 --- a/src/xorl/ops/exact/fused_silu_and_mul.py +++ b/src/xorl/ops/exact/fused_silu_and_mul.py @@ -8,7 +8,7 @@ - ``fused_silu_and_mul`` — the historical TWO-ROUND program (SiLU result rounded to the input dtype, then multiplied). Every pre-existing caller keeps these exact bytes; nothing off the exact-contract path changes. -- ``exact_fp32_silu_and_mul`` — the ONE-ROUND FP32 program (SiLU and +- ``one_round_swiglu`` — the ONE-ROUND FP32 program (SiLU and multiply in fp32, single rounding), byte-paired with serving's ``fp32_silu_and_mul`` (xorl-sglang f10b907d8). Selected ONLY by the exact-contract dispatch (``_exact_one_round_swiglu``). @@ -238,7 +238,7 @@ def _fp32_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: @triton.jit -def _exact_fp32_silu_and_mul_kernel( +def _one_round_swiglu_kernel( input_ptr, output_ptr, N: tl.constexpr, @@ -260,7 +260,7 @@ def _exact_fp32_silu_and_mul_kernel( tl.store(output_ptr + row_idx * N + col_offsets, result, mask=mask) -def _exact_fp32_silu_and_mul_forward(input_tensor: torch.Tensor) -> torch.Tensor: +def _one_round_swiglu_forward(input_tensor: torch.Tensor) -> torch.Tensor: original_shape = input_tensor.shape input_2d = input_tensor.view(-1, original_shape[-1]) num_tokens = input_2d.shape[0] @@ -272,7 +272,7 @@ def _exact_fp32_silu_and_mul_forward(input_tensor: torch.Tensor) -> torch.Tensor ) BLOCK_SIZE = min(1024, triton.next_power_of_2(N)) grid = (num_tokens, triton.cdiv(N, BLOCK_SIZE)) - _exact_fp32_silu_and_mul_kernel[grid]( + _one_round_swiglu_kernel[grid]( input_2d, output, N, @@ -284,7 +284,7 @@ def _exact_fp32_silu_and_mul_forward(input_tensor: torch.Tensor) -> torch.Tensor @triton.jit -def _exact_fp32_silu_and_mul_backward_kernel( +def _one_round_swiglu_backward_kernel( grad_output_ptr, input_ptr, grad_input_ptr, @@ -314,7 +314,7 @@ def _exact_fp32_silu_and_mul_backward_kernel( tl.store(grad_input_ptr + row_idx * 2 * N + N + col_offsets, d_up.to(up.dtype), mask=mask) -def _exact_fp32_silu_and_mul_backward(grad_output: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: +def _one_round_swiglu_backward(grad_output: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: original_shape = input_tensor.shape input_2d = input_tensor.view(-1, original_shape[-1]) grad_output_2d = grad_output.view(-1, grad_output.shape[-1]) @@ -323,7 +323,7 @@ def _exact_fp32_silu_and_mul_backward(grad_output: torch.Tensor, input_tensor: t grad_input = torch.empty_like(input_2d) BLOCK_SIZE = min(1024, triton.next_power_of_2(N)) grid = (num_tokens, triton.cdiv(N, BLOCK_SIZE)) - _exact_fp32_silu_and_mul_backward_kernel[grid]( + _one_round_swiglu_backward_kernel[grid]( grad_output_2d, input_2d, grad_input, @@ -339,15 +339,15 @@ class ExactFp32SiluAndMulFunction(torch.autograd.Function): @staticmethod def forward(ctx, input_tensor: torch.Tensor) -> torch.Tensor: ctx.save_for_backward(input_tensor) - return _exact_fp32_silu_and_mul_forward(input_tensor) + return _one_round_swiglu_forward(input_tensor) @staticmethod def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: (input_tensor,) = ctx.saved_tensors - return _exact_fp32_silu_and_mul_backward(grad_output, input_tensor) + return _one_round_swiglu_backward(grad_output, input_tensor) -def exact_fp32_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: +def one_round_swiglu(input_tensor: torch.Tensor) -> torch.Tensor: """One-round FP32 SwiGLU with autograd support (exact-contract callers only). Byte-paired with serving's fp32_silu_and_mul; see the module docstring @@ -357,3 +357,7 @@ def exact_fp32_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: if not _use_fp32_fused_swiglu(input_tensor): return _fp32_silu_and_mul(input_tensor) return ExactFp32SiluAndMulFunction.apply(input_tensor) + + +# Deprecated alias (renamed in #78/#81): the exact-contract single-rounding SwiGLU. +exact_fp32_silu_and_mul = one_round_swiglu diff --git a/tests/models/test_glm52_exact_dense_mlp.py b/tests/models/test_glm52_exact_dense_mlp.py index 47c3b5ec..2aedb889 100644 --- a/tests/models/test_glm52_exact_dense_mlp.py +++ b/tests/models/test_glm52_exact_dense_mlp.py @@ -130,7 +130,7 @@ def down_value(input, factor_A, factor_B): return _literal_linear_value(input, down_base, factor_A, factor_B) monkeypatch.setattr(module, "_exact_forward_value", gate_up_value) - monkeypatch.setattr(exact_dense_mlp_module, "exact_fp32_silu_and_mul", activation_value) + monkeypatch.setattr(exact_dense_mlp_module, "one_round_swiglu", activation_value) monkeypatch.setattr(module.down_proj, "_exact_forward_value", down_value) input = torch.arange(24, dtype=torch.float32).reshape(3, 8).sub_(7).div_(53).to(torch.bfloat16) diff --git a/tests/models/test_glm52_exact_dense_mlp_composition.py b/tests/models/test_glm52_exact_dense_mlp_composition.py index 296a7c96..7122ac38 100644 --- a/tests/models/test_glm52_exact_dense_mlp_composition.py +++ b/tests/models/test_glm52_exact_dense_mlp_composition.py @@ -8,7 +8,7 @@ Glm52ExactTP1FusedGateUpBlockFP8QLoRA, ) from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu def _pattern( @@ -167,7 +167,7 @@ def test_official_tp1_dense_vertical_composition_bytes_and_manual_vjp() -> None: trainer_gate_up = gate_up(input) trainer_gate_up.retain_grad() - trainer_activation = exact_fp32_silu_and_mul(trainer_gate_up) + trainer_activation = one_round_swiglu(trainer_gate_up) trainer_activation.retain_grad() trainer_output = down(trainer_activation) @@ -266,11 +266,11 @@ def test_official_tp1_dense_vertical_composition_bytes_and_manual_vjp() -> None: ) # Match the BF16 activation-storage boundary traversed by autograd, then # differentiate the exact one-round activation exactly as the trainer - # does: through exact_fp32_silu_and_mul's own autograd definition. + # does: through one_round_swiglu's own autograd definition. with torch.enable_grad(): manual_gate_up_leaf = trainer_gate_up.detach().requires_grad_(True) (manual_gate_up_grad,) = torch.autograd.grad( - exact_fp32_silu_and_mul(manual_gate_up_leaf), + one_round_swiglu(manual_gate_up_leaf), manual_gate_up_leaf, grad_outputs=manual_activation_grad.to(trainer_activation.dtype), ) diff --git a/tests/models/test_glm52_exact_fullparam_fp8.py b/tests/models/test_glm52_exact_fullparam_fp8.py index c51ac72b..eca16c57 100644 --- a/tests/models/test_glm52_exact_fullparam_fp8.py +++ b/tests/models/test_glm52_exact_fullparam_fp8.py @@ -27,7 +27,7 @@ ) from xorl.models.transformers.glm5.native_fp8 import Glm52NativeBlockFP8DenseMLP from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu def _linear(in_features: int = 8, out_features: int = 16) -> Glm52ExactTP1BlockFP8FullParamLinear: @@ -343,7 +343,7 @@ def test_dense_mlp_routes_to_one_round_and_discriminates_the_old_program(monkeyp generator = torch.Generator().manual_seed(9173) input = (torch.randn(32, 128, generator=generator) * 1.5).to(torch.bfloat16) gate_up = _dense_mlp_gate_up(input) - expected = exact_fp32_silu_and_mul(gate_up) + expected = one_round_swiglu(gate_up) split = gate_up.shape[-1] // 2 retired = F.silu(gate_up[..., :split]) * gate_up[..., split:] @@ -366,7 +366,7 @@ def test_dense_mlp_one_round_gradient_matches_the_selected_program(monkeypatch) module(input).backward(grad_output) reference_input = input.detach().clone().requires_grad_(True) - exact_fp32_silu_and_mul(_dense_mlp_gate_up(reference_input)).backward(grad_output) + one_round_swiglu(_dense_mlp_gate_up(reference_input)).backward(grad_output) assert input.grad is not None assert torch.equal(input.grad, reference_input.grad) assert bool(torch.all(torch.isfinite(input.grad.float()))) diff --git a/tests/models/test_glm52_exact_gate_up_qlora.py b/tests/models/test_glm52_exact_gate_up_qlora.py index 231dfd9c..9074252a 100644 --- a/tests/models/test_glm52_exact_gate_up_qlora.py +++ b/tests/models/test_glm52_exact_gate_up_qlora.py @@ -13,7 +13,7 @@ ) from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu def _module() -> Glm52ExactTP1FusedGateUpBlockFP8QLoRA: @@ -429,7 +429,7 @@ def test_official_fused_gate_up_literal_bytes_graph_metadata_zero_and_gradients( # Serving's exact mode computes the one-round FP32 SwiGLU # (SiluAndMul.forward_exact, xorl-sglang f10b907d8); the trainer op must # match a one-round sampler oracle bitwise. - trainer_activation = exact_fp32_silu_and_mul(cold_actual) + trainer_activation = one_round_swiglu(cold_actual) sampler_activation = fp32_silu_and_mul(expected) assert torch.equal(trainer_activation.view(torch.uint8), sampler_activation.view(torch.uint8)) diff --git a/tests/models/test_glm52_exact_shared_expert_qlora.py b/tests/models/test_glm52_exact_shared_expert_qlora.py index 2b0914cf..ead2abca 100644 --- a/tests/models/test_glm52_exact_shared_expert_qlora.py +++ b/tests/models/test_glm52_exact_shared_expert_qlora.py @@ -12,7 +12,7 @@ GLM52_EXACT_TP16_SHARED_EXPERT_QLORA_CONTRACT_VERSION, Glm52ExactTP16SharedExpertBlockFP8QLoRA, ) -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu def _canonical_moe_reference(partials: torch.Tensor, metadata: CanonicalMoEGraphMetadata) -> torch.Tensor: @@ -352,7 +352,7 @@ def projection_vjp( gate_up_input = exact_gate_up.detach().requires_grad_(True) # Mirror the module's VJP reference: differentiate the one-round FP32 # SwiGLU program the forward now emits. - activation = exact_fp32_silu_and_mul(gate_up_input) + activation = one_round_swiglu(gate_up_input) (gate_up_grad,) = torch.autograd.grad( activation, gate_up_input, diff --git a/tests/models/test_glm5_one_round_swiglu.py b/tests/models/test_glm5_one_round_swiglu.py index f4b573f7..649bd986 100644 --- a/tests/models/test_glm5_one_round_swiglu.py +++ b/tests/models/test_glm5_one_round_swiglu.py @@ -6,7 +6,7 @@ through that op. These gates pin the trainer half of the pairing at GLM geometries: -1. the trainer op (``exact_fp32_silu_and_mul``) matches serving's +1. the trainer op (``one_round_swiglu``) matches serving's ``fp32_silu_and_mul`` bitwise on matched bf16 inputs at the TP16 shared-expert shard ([T, 256] gate_up) and the dense width ([T, 2*12288], the Glm5Config default intermediate size) — direct sglang @@ -15,7 +15,7 @@ from the one-round output on these inputs, so the byte gates above have discriminating power (the deterministic fixture must differ in at least one element); -3. the separate-tensor form ``exact_fp32_silu_and_mul(cat([gate, up], -1))`` +3. the separate-tensor form ``one_round_swiglu(cat([gate, up], -1))`` — the ``_canonical_shared_local_partial`` pattern — equals serving on the concatenated tensor; 4. grad smoke: the op under ``requires_grad`` produces finite grads at the @@ -48,7 +48,7 @@ ) from xorl.models.transformers.glm5.configuration_glm5 import Glm5Config from xorl.models.transformers.glm5.modeling_glm5 import GLM52_LOCAL_PARTIAL_POLICY, Glm5MLP -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu # GLM-5.2 geometries with byte evidence: the 128-wide TP16 shared-expert @@ -80,7 +80,7 @@ def test_trainer_op_matches_serving_one_round_bitwise(rows: int, width: int) -> fp32_silu_and_mul = _serving_one_round_op() matched = _matched_input(rows, width, seed=17) - trainer = exact_fp32_silu_and_mul(matched) + trainer = one_round_swiglu(matched) serving = fp32_silu_and_mul(matched) assert trainer.dtype is torch.bfloat16 and serving.dtype is torch.bfloat16 @@ -97,7 +97,7 @@ def test_old_two_op_program_differs_bitwise(rows: int, width: int) -> None: matched = _matched_input(rows, width, seed=17) split = width // 2 - one_round = exact_fp32_silu_and_mul(matched) + one_round = one_round_swiglu(matched) two_op = F.silu(matched[..., :split]) * matched[..., split:] differing = int((one_round.view(torch.uint16) != two_op.view(torch.uint16)).sum()) @@ -121,7 +121,7 @@ def test_separate_gate_up_concat_form_matches_serving() -> None: up = _matched_input(rows, 2 * shard, seed=29)[..., shard:].contiguous() concatenated = torch.cat([gate, up], dim=-1) - trainer = exact_fp32_silu_and_mul(concatenated) + trainer = one_round_swiglu(concatenated) serving = fp32_silu_and_mul(concatenated) assert torch.equal(trainer.view(torch.uint8), serving.view(torch.uint8)) @@ -134,7 +134,7 @@ def test_grad_smoke_at_shard_geometry() -> None: matched = _matched_input(rows, width, seed=31).requires_grad_(True) grad_output = _matched_input(rows, 2 * (width // 2), seed=37)[..., : width // 2].contiguous() - output = exact_fp32_silu_and_mul(matched) + output = one_round_swiglu(matched) output.backward(grad_output) assert matched.grad is not None @@ -276,7 +276,7 @@ def two_round_capture(gate_up: torch.Tensor) -> torch.Tensor: activated = F.silu(gate_up[..., :split].float()).to(gate_up.dtype) return (activated * gate_up[..., split:]).to(gate_up.dtype) - monkeypatch.setattr(modeling_glm5, "exact_fp32_silu_and_mul", one_round_capture) + monkeypatch.setattr(modeling_glm5, "one_round_swiglu", one_round_capture) monkeypatch.setattr(modeling_glm5, "fused_silu_and_mul", two_round_capture) torch.manual_seed(5) diff --git a/tests/models/test_op_parity_dense.py b/tests/models/test_op_parity_dense.py index 68d3dd88..7887e84a 100644 --- a/tests/models/test_op_parity_dense.py +++ b/tests/models/test_op_parity_dense.py @@ -16,7 +16,7 @@ import torch.nn.functional as F from xorl.models.layers import rope as xrope -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -101,7 +101,7 @@ def test_swiglu_xorl_one_round_bit_exact_vs_sglang_fp32(): inter = 6144 gate_up = torch.randn(SEQ, 2 * inter, device=DEV, dtype=DT) - x_xorl_exact = exact_fp32_silu_and_mul(gate_up) + x_xorl_exact = one_round_swiglu(gate_up) d = gate_up.shape[-1] // 2 x_sg = (F.silu(gate_up[..., :d].float()) * gate_up[..., d:].float()).to(DT) diff --git a/tests/models/test_qwen35_exact_swiglu_grad_engagement.py b/tests/models/test_qwen35_exact_swiglu_grad_engagement.py index 5573c9d7..f7744312 100644 --- a/tests/models/test_qwen35_exact_swiglu_grad_engagement.py +++ b/tests/models/test_qwen35_exact_swiglu_grad_engagement.py @@ -4,7 +4,7 @@ a grad-engagement gate — the module forward/backward must produce non-None, finite gradients for every parameter and for the input, and the gradients must match a reference trajectory. The op-level backward test -(tests/ops/test_exact_fp32_silu_and_mul.py) covers the kernel; this gate +(tests/ops/test_one_round_swiglu.py) covers the kernel; this gate covers engagement through Qwen3_5MLP, where the fused one-round path is selected by the exact-contract policy rather than by a flag. """ diff --git a/tests/ops/test_fused_silu_and_mul_scoping.py b/tests/ops/test_fused_silu_and_mul_scoping.py index fade45a5..f25d561a 100644 --- a/tests/ops/test_fused_silu_and_mul_scoping.py +++ b/tests/ops/test_fused_silu_and_mul_scoping.py @@ -20,7 +20,7 @@ import torch.nn.functional as F from xorl.models.transformers.qwen3_5.modeling_qwen3_5 import Qwen3_5MLP -from xorl.ops.exact.fused_silu_and_mul import exact_fp32_silu_and_mul, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu, fused_silu_and_mul def _two_round_reference(x: torch.Tensor) -> torch.Tensor: @@ -61,7 +61,7 @@ def test_default_op_keeps_two_round_bytes(shape): def test_exact_op_is_one_round(shape=(96, 7168)): torch.manual_seed(3) x = torch.randn(*shape, device="cuda", dtype=torch.bfloat16).contiguous() - assert torch.equal(exact_fp32_silu_and_mul(x), _one_round_reference(x)) + assert torch.equal(one_round_swiglu(x), _one_round_reference(x)) @pytest.mark.gpu @@ -88,5 +88,5 @@ def test_exact_qwen35_mlp_selects_one_round(): x = torch.randn(96, 64, device="cuda", dtype=torch.bfloat16) with torch.no_grad(): out = mlp(x) - reference = mlp.down_proj(exact_fp32_silu_and_mul(mlp.gate_up_proj(x))) + reference = mlp.down_proj(one_round_swiglu(mlp.gate_up_proj(x))) assert torch.equal(out, reference) diff --git a/tests/ops/test_exact_fp32_silu_and_mul.py b/tests/ops/test_one_round_swiglu.py similarity index 89% rename from tests/ops/test_exact_fp32_silu_and_mul.py rename to tests/ops/test_one_round_swiglu.py index cd3e9f10..1fc3d2c9 100644 --- a/tests/ops/test_exact_fp32_silu_and_mul.py +++ b/tests/ops/test_one_round_swiglu.py @@ -7,7 +7,7 @@ _MODULE_PATH = Path(__file__).resolve().parents[2] / "src/xorl/ops/fused_silu_and_mul.py" -_SPEC = importlib.util.spec_from_file_location("xorl_exact_fp32_silu_and_mul", _MODULE_PATH) +_SPEC = importlib.util.spec_from_file_location("xorl_one_round_swiglu", _MODULE_PATH) assert _SPEC is not None and _SPEC.loader is not None _MODULE = importlib.util.module_from_spec(_SPEC) _SPEC.loader.exec_module(_MODULE) @@ -29,7 +29,7 @@ def test_cpu_fallback_uses_one_round_program(): [[0.5, -1.25, 3.0, -0.75], [-2.0, 0.125, 1.5, 8.0]], dtype=torch.bfloat16, ) - actual = _MODULE.exact_fp32_silu_and_mul(values) + actual = _MODULE.one_round_swiglu(values) assert torch.equal(actual, _one_round_reference(values)) assert not torch.equal(actual, _two_round_reference(values)) @@ -41,7 +41,7 @@ def test_forward_is_byte_exact_to_one_round_reference(shape): torch.manual_seed(4) input_tensor = torch.randn(*shape, device="cuda", dtype=torch.bfloat16).contiguous() - actual = _MODULE.exact_fp32_silu_and_mul(input_tensor) + actual = _MODULE.one_round_swiglu(input_tensor) expected = _one_round_reference(input_tensor) assert torch.equal(actual, expected) @@ -55,7 +55,7 @@ def test_backward_matches_one_round_reference(): reference_input = input_tensor.detach().clone().requires_grad_(True) grad_output = torch.randn(512, 2048, device="cuda", dtype=torch.bfloat16) - _MODULE.exact_fp32_silu_and_mul(input_tensor).backward(grad_output) + _MODULE.one_round_swiglu(input_tensor).backward(grad_output) _one_round_reference(reference_input).backward(grad_output) split = input_tensor.shape[-1] // 2 From cd0d568b1b1e59e2420e1758af77b6e19d92da07 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 16:41:43 +0000 Subject: [PATCH 12/14] Spell out batch_invariant; deprecate the bi_fused ce_mode alias (#81 review, row 2) - ce_mode canonical value is now 'batch_invariant'; 'bi_fused' remains an accepted deprecated alias normalized (with a DeprecationWarning) in both argument entry points, and the CrossEntropyMode Literal accepts both. - ops/loss/bi_fused_lm_head.py -> batch_invariant_lm_head.py; bi_fused_per_token_ce / bi_fused_vocab_parallel_per_token_ce / resolve_bi_fused_lm_head_tp_groups spell out batch_invariant_*. - ops/exact/bi_gemm_configs.py -> batch_invariant_gemm_configs.py. - Prose, error messages, examples, and docs updated; test files renamed (incl. the tests/distributed shard claim). - The bi_ prefix survives only inside the serving twins (bi_families_v2, batch_invariant_ops' bi_* kernels, bi_fused_gdn_gating) whose symbol names must match the engine. - Alias regression tests cover rope_class_b and ce_mode='bi_fused'. --- .github/workflows/pr-test-cpu.yml | 2 +- docs/k3/GEMM_CONTRACT.md | 2 +- docs/k3/LM_HEAD_CONTRACT.md | 4 +- .../content/docs/config-reference/local.md | 2 +- .../content/docs/config-reference/server.md | 2 +- docs/src/content/docs/loss-functions.mdx | 4 +- .../full/qwen3_235b_a22b_8node_ep64.yaml | 4 +- .../full/qwen3_5_35b_a3b_ep1_fsdp8.yaml | 4 +- .../configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml | 4 +- .../configs/full/qwen3_5_35b_a3b_full.yaml | 4 +- .../full/qwen3_5_35b_a3b_full_no_deepep.yaml | 4 +- .../configs/full/qwen3_5_397b_a17b_full.yaml | 4 +- .../server/configs/full/qwen3_8b_full.yaml | 4 +- .../full/qwen3_coder_30b_a3b_full.yaml | 4 +- .../configs/lora/qwen3_30b_a3b_lora.yaml | 4 +- .../configs/lora/qwen3_5_35b_a3b_lora.yaml | 4 +- .../configs/lora/qwen3_5_397b_a17b_lora.yaml | 4 +- .../server/configs/lora/qwen3_8b_lora.yaml | 4 +- .../lora/qwen3_coder_30b_a3b_lora.yaml | 4 +- src/xorl/arguments.py | 11 +- src/xorl/distributed/__init__.py | 2 +- src/xorl/models/auto.py | 22 +- src/xorl/models/layers/rope.py | 8 +- src/xorl/models/module_utils.py | 10 +- .../transformers/glm5/exact_lm_head_qlora.py | 4 +- src/xorl/models/transformers/glm5/indexer.py | 4 +- .../models/transformers/glm5/modeling_glm5.py | 12 +- .../transformers/qwen3_5/modeling_qwen3_5.py | 2 +- .../qwen3_5_moe/modeling_qwen3_5_moe.py | 2 +- src/xorl/objectives/causallm_loss.py | 42 +- src/xorl/ops/README.md | 4 +- src/xorl/ops/bi_gemm_configs.py | 8 +- src/xorl/ops/exact/__init__.py | 2 +- ...igs.py => batch_invariant_gemm_configs.py} | 0 src/xorl/ops/loss/__init__.py | 6 +- src/xorl/ops/loss/batch_invariant_lm_head.py | 905 ++++++++++++++++++ src/xorl/ops/loss/bi_fused_lm_head.py | 904 +---------------- src/xorl/ops/loss/per_token_ce.py | 30 +- src/xorl/ops/sglang/batch_invariant_ops.py | 2 +- src/xorl/server/runner/model_runner.py | 2 +- src/xorl/server/server_arguments.py | 12 +- src/xorl/trainers/trainer.py | 2 +- src/xorl/trainers/training_utils.py | 6 +- ....py => test_batch_invariant_lm_head_tp.py} | 20 +- .../test_canonical_moe_contract.py | 2 +- tests/distributed/test_pp_byte_alignment.py | 6 +- tests/distributed/test_pp_byte_contract.py | 4 +- .../test_ulysses_byte_alignment.py | 4 +- ...test_batch_invariant_lm_head_tp_wiring.py} | 34 +- tests/models/test_dsv4_exact_contract.py | 4 +- ...st_glm52_exact_lm_head_loss_integration.py | 22 +- ...test_batch_invariant_gemm_config_table.py} | 2 +- ...ead.py => test_batch_invariant_lm_head.py} | 48 +- tests/ops/test_exact_sampling_transforms.py | 4 +- tests/ops/test_fused_silu_and_mul_scoping.py | 2 +- .../server/runner/test_model_runner_drgrpo.py | 22 +- tests/test_arguments.py | 24 + .../trainers/test_qwen3_dense_exact_config.py | 2 +- .../test_rope_fp32_single_round_config.py | 8 +- 59 files changed, 1175 insertions(+), 1103 deletions(-) rename src/xorl/ops/exact/{bi_gemm_configs.py => batch_invariant_gemm_configs.py} (100%) create mode 100644 src/xorl/ops/loss/batch_invariant_lm_head.py rename tests/distributed/{test_bi_fused_lm_head_tp.py => test_batch_invariant_lm_head_tp.py} (92%) rename tests/models/{test_bi_fused_lm_head_tp_wiring.py => test_batch_invariant_lm_head_tp_wiring.py} (89%) rename tests/ops/{test_bi_gemm_config_table.py => test_batch_invariant_gemm_config_table.py} (97%) rename tests/ops/{test_bi_fused_lm_head.py => test_batch_invariant_lm_head.py} (84%) diff --git a/.github/workflows/pr-test-cpu.yml b/.github/workflows/pr-test-cpu.yml index ad424921..79920f5d 100644 --- a/.github/workflows/pr-test-cpu.yml +++ b/.github/workflows/pr-test-cpu.yml @@ -62,7 +62,7 @@ jobs: tests/distributed/test_lm_head_tp_*.py tests/distributed/test_tensor_parallel.py tests/distributed/test_vocab_parallel_ce.py - tests/distributed/test_bi_fused_lm_head_tp.py + tests/distributed/test_batch_invariant_lm_head_tp.py tests/distributed/test_fused_linear_logprob_tp.py tests/distributed/test_olmo2_*.py diff --git a/docs/k3/GEMM_CONTRACT.md b/docs/k3/GEMM_CONTRACT.md index f7883fd7..5d558bc6 100644 --- a/docs/k3/GEMM_CONTRACT.md +++ b/docs/k3/GEMM_CONTRACT.md @@ -15,7 +15,7 @@ The output tile, group size, pipeline depth, and warp count do not split the K reduction, so they remain performance-tuning axes after passing the bitwise gate. -`bi_gemm_configs.py` contains the shared, shape-keyed table. Each entry keeps +`batch_invariant_gemm_configs.py` contains the shared, shape-keyed table. Each entry keeps the dtype's pinned K tile, compares bitwise with the baseline configuration, and checks that an identical row keeps identical output bits across row-count buckets. Set `XORL_BI_GEMM_CONFIG_TABLE=0` to use the pinned baseline table. diff --git a/docs/k3/LM_HEAD_CONTRACT.md b/docs/k3/LM_HEAD_CONTRACT.md index ac81b07e..dcf0301e 100644 --- a/docs/k3/LM_HEAD_CONTRACT.md +++ b/docs/k3/LM_HEAD_CONTRACT.md @@ -27,7 +27,7 @@ vocabulary chunk at a time, records the same maximum, exponential sum, and selected logit, then merges chunks in pinned order. This rollback is exact but uses more launches. -The trainer exposes this as `ce_mode="bi_fused"`. Its custom autograd function +The trainer exposes this as `ce_mode="batch_invariant"`. Its custom autograd function saves the exact forward LSE and recomputes the conventional closed-form CE gradient by chunks. Only forward values enter the train/serve equality contract; backward remains ordinary checked training numerics. @@ -42,7 +42,7 @@ second fp32 copy of the LM-head weight. Run on a Hopper GPU: ```bash -pytest tests/ops/test_bi_fused_lm_head.py -v +pytest tests/ops/test_batch_invariant_lm_head.py -v ``` The gate compares forward probabilities and loss with an eager fp32 reference, diff --git a/docs/src/content/docs/config-reference/local.md b/docs/src/content/docs/config-reference/local.md index 40cb17e0..77ce8e77 100644 --- a/docs/src/content/docs/config-reference/local.md +++ b/docs/src/content/docs/config-reference/local.md @@ -175,7 +175,7 @@ Each entry in `datasets` (or `test_datasets`) is a dict: | `activation_gpu_limit` | `0.0` | GB of activations to keep on GPU when offloading. `0.0` = offload all. | | `enable_compile` | `false` | `torch.compile` for model forward pass. | | `compile_dynamic_shapes` | `false` | Pass `dynamic=True` to `torch.compile`; keep disabled unless a workload has benchmarked a dynamic-shape win. | -| `ce_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `compiled`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `bi_fused`. Explicit modes also include `eager`, `quack_linear`, and `fused_quack`, subject to loss/topology checks. | +| `ce_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `compiled`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `batch_invariant`. Explicit modes also include `eager`, `quack_linear`, and `fused_quack`, subject to loss/topology checks. | | `ce_num_chunks` | `8` | Number of token chunks for chunked/compiled cross-entropy. | | `enable_fp8_training` | `false` | Experimental full-weight block-FP8 compute with BF16/FP32 master parameters. Mutually exclusive with QARL. | | `enable_qarl` | `false` | Experimental dynamic fake-quant training with full-precision masters and STE gradients. E4M3 applies to dense `nn.Linear` modules; NVFP4 also supports MoE expert containers. Mutually exclusive with full-weight FP8 training. | diff --git a/docs/src/content/docs/config-reference/server.md b/docs/src/content/docs/config-reference/server.md index 006a398f..8faaabba 100644 --- a/docs/src/content/docs/config-reference/server.md +++ b/docs/src/content/docs/config-reference/server.md @@ -96,7 +96,7 @@ These stored defaults are resolved after the model architecture is known. Ordina | `enable_forward_prefetch` | `false` | FSDP forward prefetch. | | `init_device` | `meta` | Model initialization device: `cpu`, `meta`, `cuda`. | | `load_weights_mode` | `grouped` | Weight loading mode: `grouped` (default, with rank-0 fallback), `all_ranks`, or `skip`. | -| `ce_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `compiled`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `bi_fused`. Explicit modes also include `eager`, `quack_linear`, and `fused_quack`, subject to loss/topology checks. | +| `ce_mode` | `null` (resolved) | Ordinary models and exact DSV4-Flash resolve to `compiled`; exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs resolve to `batch_invariant`. Explicit modes also include `eager`, `quack_linear`, and `fused_quack`, subject to loss/topology checks. | | `enable_fp8_training` | `false` | Experimental full-weight block-FP8 compute. Mutually exclusive with LoRA/QLoRA and QARL. | | `enable_qarl` | `false` | Experimental dynamic fake-quant training with full-precision masters and STE gradients. E4M3 applies to dense `nn.Linear` modules; NVFP4 also supports MoE expert containers. Mutually exclusive with LoRA/QLoRA and full-weight FP8 training. | | `qarl_quant_cfg` | `null` | QARL alias or dictionary. `null`/`FP8_DEFAULT_CFG` resolves to dynamic E4M3 W8A8 with `[128, 128]` weight blocks. `nvfp4` resolves to dynamic, weight-only W4 with `group_size: 16`; set `activation: true` for W4A4. NVFP4 covers dense linears and MoE expert containers, while E4M3 is dense-only. | diff --git a/docs/src/content/docs/loss-functions.mdx b/docs/src/content/docs/loss-functions.mdx index 1fdc272f..c51d9925 100644 --- a/docs/src/content/docs/loss-functions.mdx +++ b/docs/src/content/docs/loss-functions.mdx @@ -71,13 +71,13 @@ Per-token output behavior depends on the selected loss and the request's `return ### Compiled cross-entropy (`ce_mode`) -The `ce_mode` setting controls the LM-head and cross-entropy implementation. When omitted, XoRL resolves it from the model: ordinary models and exact DSV4-Flash use `compiled`, while exact dense Qwen3, Qwen3.5-family, and GLM-5.2 numerical programs require `bi_fused`. +The `ce_mode` setting controls the LM-head and cross-entropy implementation. When omitted, XoRL resolves it from the model: ordinary models and exact DSV4-Flash use `compiled`, while exact dense Qwen3, Qwen3.5-family, and GLM-5.2 numerical programs require `batch_invariant`. | Value | Description | Important constraints | |---|---|---| | `compiled` | `torch.compile`-compiled chunked cross-entropy | General production default | | `eager` | Standard eager cross-entropy | Debugging; may materialize the full logits tensor | -| `bi_fused` | Batch-invariant selected-token logprob/CE path | Required by the current exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs; topology and dtype restrictions are checked at runtime | +| `batch_invariant` | Batch-invariant selected-token logprob/CE path | Required by the current exact dense Qwen3, Qwen3.5-family, and GLM-5.2 programs; topology and dtype restrictions are checked at runtime | | `quack_linear` | Quack chunked linear plus scalar cross-entropy | Causal-LM path; pipeline parallelism supports this mode on its last stage | | `fused_quack` | Chunked matmul plus fused selected-token CE | Used by supported per-token loss paths; not supported by every loss/topology combination | diff --git a/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml b/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml index 72185ef9..0c67606e 100644 --- a/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml +++ b/examples/server/configs/full/qwen3_235b_a22b_8node_ep64.yaml @@ -56,6 +56,6 @@ enable_packing: true # Full-weights training enable_lora: false -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant freeze_router: false diff --git a/examples/server/configs/full/qwen3_5_35b_a3b_ep1_fsdp8.yaml b/examples/server/configs/full/qwen3_5_35b_a3b_ep1_fsdp8.yaml index c1a47424..55816e31 100644 --- a/examples/server/configs/full/qwen3_5_35b_a3b_ep1_fsdp8.yaml +++ b/examples/server/configs/full/qwen3_5_35b_a3b_ep1_fsdp8.yaml @@ -30,8 +30,8 @@ sample_packing_sequence_len: 128000 enable_packing: true skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant optimizer: muon optimizer_dtype: bf16 diff --git a/examples/server/configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml b/examples/server/configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml index 7db687e3..893162b2 100644 --- a/examples/server/configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml +++ b/examples/server/configs/full/qwen3_5_35b_a3b_ep8_cp1.yaml @@ -31,8 +31,8 @@ sample_packing_sequence_len: 128000 enable_packing: true skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant optimizer: muon optimizer_dtype: bf16 diff --git a/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml b/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml index 6fac4839..b1f865f6 100644 --- a/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml +++ b/examples/server/configs/full/qwen3_5_35b_a3b_full.yaml @@ -36,8 +36,8 @@ sample_packing_sequence_len: 128000 enable_packing: true skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant optimizer: muon optimizer_dtype: bf16 diff --git a/examples/server/configs/full/qwen3_5_35b_a3b_full_no_deepep.yaml b/examples/server/configs/full/qwen3_5_35b_a3b_full_no_deepep.yaml index d07e535f..06433ea2 100644 --- a/examples/server/configs/full/qwen3_5_35b_a3b_full_no_deepep.yaml +++ b/examples/server/configs/full/qwen3_5_35b_a3b_full_no_deepep.yaml @@ -36,8 +36,8 @@ sample_packing_sequence_len: 128000 enable_packing: true skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant optimizer: muon optimizer_dtype: bf16 diff --git a/examples/server/configs/full/qwen3_5_397b_a17b_full.yaml b/examples/server/configs/full/qwen3_5_397b_a17b_full.yaml index 5e17074f..b5e54089 100644 --- a/examples/server/configs/full/qwen3_5_397b_a17b_full.yaml +++ b/examples/server/configs/full/qwen3_5_397b_a17b_full.yaml @@ -45,8 +45,8 @@ sample_packing_sequence_len: 65536 enable_packing: true skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant optimizer: muon optimizer_dtype: bf16 diff --git a/examples/server/configs/full/qwen3_8b_full.yaml b/examples/server/configs/full/qwen3_8b_full.yaml index 2572daee..01dcab9f 100644 --- a/examples/server/configs/full/qwen3_8b_full.yaml +++ b/examples/server/configs/full/qwen3_8b_full.yaml @@ -36,5 +36,5 @@ enable_packing: true # Full weight training skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant diff --git a/examples/server/configs/full/qwen3_coder_30b_a3b_full.yaml b/examples/server/configs/full/qwen3_coder_30b_a3b_full.yaml index afd8c79c..834e6d0c 100644 --- a/examples/server/configs/full/qwen3_coder_30b_a3b_full.yaml +++ b/examples/server/configs/full/qwen3_coder_30b_a3b_full.yaml @@ -41,8 +41,8 @@ enable_packing: true # Full weight training skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant # Optimizer: Muon with bf16 states on GPU (no offload with 8 GPUs). # 8 GPUs: weights ~7.5GB/GPU + Muon momentum ~7.5GB/GPU = ~15GB/GPU. Fits easily. diff --git a/examples/server/configs/lora/qwen3_30b_a3b_lora.yaml b/examples/server/configs/lora/qwen3_30b_a3b_lora.yaml index 35d9e0ae..70d64626 100644 --- a/examples/server/configs/lora/qwen3_30b_a3b_lora.yaml +++ b/examples/server/configs/lora/qwen3_30b_a3b_lora.yaml @@ -74,5 +74,5 @@ lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_p # Skip initial checkpoint skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss). -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss). +ce_mode: batch_invariant diff --git a/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml b/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml index 134dc777..dadd8342 100644 --- a/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml +++ b/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml @@ -48,5 +48,5 @@ lora_alpha: 16 lora_target_modules: ["q_proj", "k_proj", "v_proj", "g_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant diff --git a/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml b/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml index acd0f2cf..f45cace9 100644 --- a/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml +++ b/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml @@ -46,5 +46,5 @@ lora_alpha: 32 lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "g_proj", "down_proj"] skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant diff --git a/examples/server/configs/lora/qwen3_8b_lora.yaml b/examples/server/configs/lora/qwen3_8b_lora.yaml index 72109d89..e8276633 100644 --- a/examples/server/configs/lora/qwen3_8b_lora.yaml +++ b/examples/server/configs/lora/qwen3_8b_lora.yaml @@ -42,5 +42,5 @@ lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_p # Skip initial checkpoint skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant diff --git a/examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml b/examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml index 7725226a..2d2be934 100644 --- a/examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml +++ b/examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml @@ -65,5 +65,5 @@ lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_p # Skip initial checkpoint skip_initial_checkpoint: true -# bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) -ce_mode: bi_fused +# batch_invariant = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for batch_invariant) +ce_mode: batch_invariant diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index 06804fb9..c234917c 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -927,7 +927,7 @@ def optimizer_kwargs(self) -> Dict[str, Any]: default=None, metadata={ "help": "Cross-entropy computation mode for the local-trainer path. " - "Omitted means compiled for ordinary models and bi_fused for canonical GLM-5.2. " + "Omitted means compiled for ordinary models and batch_invariant for canonical GLM-5.2. " "'compiled': torch.compile + auto_chunker, avoids materializing " "the full [batch*seq, vocab] logits tensor. 'quack_linear': Quack chunked " "linear + cross-entropy scalar training loss; return_per_token routes through " @@ -1662,6 +1662,15 @@ def moe_recomputed(self) -> bool: ) def __post_init__(self): + if self.ce_mode == "bi_fused": + import warnings + + warnings.warn( + "ce_mode='bi_fused' is deprecated; use ce_mode='batch_invariant'", + DeprecationWarning, + stacklevel=2, + ) + self.ce_mode = "batch_invariant" from xorl.fp8_training.config_compat import normalize_fp8_training_config # noqa: PLC0415 from xorl.qarl import normalize_qarl_quant_cfg # noqa: PLC0415 diff --git a/src/xorl/distributed/__init__.py b/src/xorl/distributed/__init__.py index be8f47c9..5e6fcb59 100644 --- a/src/xorl/distributed/__init__.py +++ b/src/xorl/distributed/__init__.py @@ -14,10 +14,10 @@ OutputDistribution, ParallelRole, canonical_moe_fold_fp64_v3, - moe_fixed_order_leaf_fp32_v1, canonical_moe_reduce_cp_sharded_v3, canonical_moe_reduce_fp64_v3, canonical_moe_reduce_packed_ep16_v2, canonical_moe_reduce_reference, + moe_fixed_order_leaf_fp32_v1, resolve_canonical_moe_transport, ) diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 14055c61..7f871bbe 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -768,18 +768,22 @@ def resolve_cross_entropy_mode(config: PretrainedConfig, ce_mode: Optional[str]) ) return "compiled" if _is_exact_qwen35(config): - if ce_mode not in (None, "bi_fused"): - raise ValueError(f"Exact Qwen3.5-family server training requires ce_mode='bi_fused'; received {ce_mode!r}") - return "bi_fused" + if ce_mode not in (None, "batch_invariant"): + raise ValueError( + f"Exact Qwen3.5-family server training requires ce_mode='batch_invariant'; received {ce_mode!r}" + ) + return "batch_invariant" if _is_exact_qwen3_dense(config): - if ce_mode not in (None, "bi_fused"): - raise ValueError(f"Exact dense Qwen3 server training requires ce_mode='bi_fused'; received {ce_mode!r}") - return "bi_fused" + if ce_mode not in (None, "batch_invariant"): + raise ValueError( + f"Exact dense Qwen3 server training requires ce_mode='batch_invariant'; received {ce_mode!r}" + ) + return "batch_invariant" if not _is_exact_glm52(config): return ce_mode or "compiled" - if ce_mode not in (None, "bi_fused"): - raise ValueError(f"Canonical GLM-5.2 exact forward requires ce_mode='bi_fused'; received {ce_mode!r}") - return "bi_fused" + if ce_mode not in (None, "batch_invariant"): + raise ValueError(f"Canonical GLM-5.2 exact forward requires ce_mode='batch_invariant'; received {ce_mode!r}") + return "batch_invariant" def build_foundation_model( diff --git a/src/xorl/models/layers/rope.py b/src/xorl/models/layers/rope.py index 6935f3c1..df8b1ae2 100644 --- a/src/xorl/models/layers/rope.py +++ b/src/xorl/models/layers/rope.py @@ -462,7 +462,9 @@ def __init__(self, config, device=None): self._set_inv_freq_fp32(self._cpu_fp32_inv_freq()) self._sglang_default_cache = None self._use_sglang_default_cache = bool(getattr(config, "_rope_native", False) and self.rope_type == "default") - self._fp32_single_round = bool(getattr(config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(config)) + self._fp32_single_round = bool( + getattr(config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(config) + ) def _cpu_fp32_inv_freq(self) -> torch.Tensor: """Frequency table computed on CPU in fp32 — the provenance serving's cos/sin cache is built with.""" @@ -623,9 +625,7 @@ def set_rope_native(enabled: bool): # first would land back in Class A, so the Class-B lane keeps the table in fp32 all # the way to the kernel. -_rope_fp32_single_round = ( - os.environ.get("XORL_ROPE_FP32_SINGLE_ROUND", os.environ.get("XORL_ROPE_CLASS_B", "")) == "1" -) +_rope_fp32_single_round = os.environ.get("XORL_ROPE_FP32_SINGLE_ROUND", os.environ.get("XORL_ROPE_CLASS_B", "")) == "1" def set_rope_fp32_single_round(enabled: bool) -> None: diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index c8218f2b..2c68c2d2 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -2878,12 +2878,12 @@ def compute_loss( getattr(lm_head, "_glm52_exact_tp16_lm_head", False) or getattr(lm_head, "_dsv4_exact_tp8_lm_head", False) ) requested_ce_mode = (loss_fn_params or {}).get("ce_mode") - bi_fused_lm_head_tp = bool( - requested_ce_mode == "bi_fused" + batch_invariant_lm_head_tp = bool( + requested_ce_mode == "batch_invariant" and getattr(ps, "lm_head_tp_size", 1) > 1 and getattr(ps, "lm_head_tp_group", None) is not None ) - if bi_fused_lm_head_tp: + if batch_invariant_lm_head_tp: loss_fn_params = dict(loss_fn_params or {}) global_valid_tokens = loss_fn_params.pop("fsdp_sharded_lm_head_loss_global_valid_tokens", None) loss_fn_params.pop("fsdp_sharded_lm_head_loss_num_chunks", None) @@ -2895,7 +2895,7 @@ def compute_loss( fsdp_sharded_loss = ( bool(getattr(lm_head, "_xorl_fsdp_sharded_lm_head_loss", False)) and not exact_lm_head - and not bi_fused_lm_head_tp + and not batch_invariant_lm_head_tp ) if fsdp_sharded_loss and fn_name not in {"causallm_loss", "cross_entropy"}: raise NotImplementedError(f"fsdp_sharded_lm_head_loss is not supported for loss function {fn_name!r}.") @@ -2987,7 +2987,7 @@ def compute_loss( raise RuntimeError("The exact GLM-5.2 lm head was not prepared for sharded TP16 loss") loss_kwargs["lm_head"] = lm_head loss_kwargs["tp_group"] = ps.lm_head_tp_group - elif bi_fused_lm_head_tp: + elif batch_invariant_lm_head_tp: loss_kwargs["lm_head"] = lm_head loss_kwargs["tp_group"] = ps.lm_head_tp_group elif ps.tp_enabled: diff --git a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py index e90feb38..aa896099 100644 --- a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py @@ -1315,8 +1315,8 @@ def glm52_exact_lm_head_per_token_ce( if not is_glm52_exact_tp16_lm_head(lm_head): raise TypeError("glm52_exact_lm_head_per_token_ce requires the constructed exact GLM-5.2 lm_head") - if ce_mode != "bi_fused": - raise NotImplementedError("The GLM-5.2 exact active-LoRA lm_head requires ce_mode='bi_fused'") + if ce_mode != "batch_invariant": + raise NotImplementedError("The GLM-5.2 exact active-LoRA lm_head requires ce_mode='batch_invariant'") if not lm_head_fp32: raise NotImplementedError("The GLM-5.2 exact active-LoRA lm_head requires lm_head_fp32=true") diff --git a/src/xorl/models/transformers/glm5/indexer.py b/src/xorl/models/transformers/glm5/indexer.py index ca3efd8e..11a63598 100644 --- a/src/xorl/models/transformers/glm5/indexer.py +++ b/src/xorl/models/transformers/glm5/indexer.py @@ -332,7 +332,9 @@ def project( cos, sin, interleaved=getattr(self.config, "indexer_rope_interleave", True), - fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(self.config)), + fp32_single_round=bool( + getattr(self.config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(self.config) + ), ) index_q = torch.cat([q_pe, q_no_pe], dim=-1) diff --git a/src/xorl/models/transformers/glm5/modeling_glm5.py b/src/xorl/models/transformers/glm5/modeling_glm5.py index de12a344..2dddd02f 100644 --- a/src/xorl/models/transformers/glm5/modeling_glm5.py +++ b/src/xorl/models/transformers/glm5/modeling_glm5.py @@ -15,10 +15,10 @@ LogicalRowOwnership, OutputDistribution, ParallelPlan, - moe_fixed_order_leaf_fp32_v1, canonical_moe_reduce_cp_sharded_v3, canonical_moe_reduce_fp64_v3, canonical_moe_reduce_packed_ep16_v2, + moe_fixed_order_leaf_fp32_v1, resolve_canonical_moe_transport, ) from xorl.distributed.moe.deepep import sync_pending_combine @@ -70,7 +70,7 @@ from xorl.models.transformers.glm5.sparse_mla import sparse_mla_dispatch from xorl.models.transformers.glm5.support import validate_glm5_sequence_parallel from xorl.ops.exact.block_fp8_native import NativeBlockFP8Linear -from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul, one_round_swiglu from xorl.utils import logging @@ -271,7 +271,9 @@ def _project_qkv( cos, sin, interleaved=getattr(self.config, "rope_interleave", True), - fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(self.config)), + fp32_single_round=bool( + getattr(self.config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(self.config) + ), ) k_rot = k_rot.expand(*k_pass.shape[:-1], -1) @@ -461,7 +463,9 @@ def _project_qkv_absorb( cos, sin, interleaved=getattr(self.config, "rope_interleave", True), - fp32_single_round=bool(getattr(self.config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(self.config)), + fp32_single_round=bool( + getattr(self.config, "_rope_fp32_single_round", False) or glm52_exact_forward_enabled(self.config) + ), ) k_pe = k_pe.squeeze(2) diff --git a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py index 0fdd7e21..c73ca916 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -43,7 +43,7 @@ has_linear_attention_layers, qwen3_5_apply_rotary_pos_emb, ) -from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul, one_round_swiglu from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils import logging diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index 1f2dd970..e3509192 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -44,7 +44,7 @@ has_linear_attention_layers, qwen3_5_apply_rotary_pos_emb, ) -from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul, one_round_swiglu from xorl.ops.linear_attention.ops.cp import build_linear_attention_cp_context from xorl.utils import logging diff --git a/src/xorl/objectives/causallm_loss.py b/src/xorl/objectives/causallm_loss.py index a627f55b..3b25c5bb 100644 --- a/src/xorl/objectives/causallm_loss.py +++ b/src/xorl/objectives/causallm_loss.py @@ -19,7 +19,7 @@ LogprobTopK, compute_per_token_ce, normalize_logprob_temperature, - resolve_bi_fused_lm_head_tp_groups, + resolve_batch_invariant_lm_head_tp_groups, ) from xorl.ops.loss.vocab_parallel_cross_entropy import ( _backward_kernel as _vocab_parallel_ce_backward_kernel, @@ -388,7 +388,7 @@ def _fused_quack_per_token_ce( ) -def _bi_fused_per_token_ce_checked( +def _batch_invariant_per_token_ce_checked( hidden_states_flat: torch.Tensor, weight: torch.Tensor, labels_flat: torch.Tensor, @@ -396,19 +396,19 @@ def _bi_fused_per_token_ce_checked( lm_head_fp32: bool, z_loss_enabled: bool, ) -> torch.Tensor: - """Guarded entry for ``ce_mode='bi_fused'`` (the batch-invariant lm-head + """Guarded entry for ``ce_mode='batch_invariant'`` (the batch-invariant lm-head contract). The contract IS the fp32-class lm-head computation, so it requires ``lm_head_fp32`` semantics without materializing the fp32 weight.""" - from xorl.ops.loss.bi_fused_lm_head import bi_fused_per_token_ce + from xorl.ops.loss.batch_invariant_lm_head import batch_invariant_per_token_ce if z_loss_enabled: - raise NotImplementedError("ce_mode='bi_fused' does not support softmax_auxiliary_loss") + raise NotImplementedError("ce_mode='batch_invariant' does not support softmax_auxiliary_loss") if not lm_head_fp32: raise NotImplementedError( - "ce_mode='bi_fused' implements the fp32-class lm-head contract; set lm_head_fp32: true" + "ce_mode='batch_invariant' implements the fp32-class lm-head contract; set lm_head_fp32: true" ) local_weight = weight.to_local() if hasattr(weight, "to_local") else weight - return bi_fused_per_token_ce(hidden_states_flat, local_weight, labels_flat, ignore_index) + return batch_invariant_per_token_ce(hidden_states_flat, local_weight, labels_flat, ignore_index) def _quack_linear_per_token_cross_entropy( @@ -511,13 +511,13 @@ def causallm_loss_function( labels_flat = labels.view(-1) hidden_states_flat = hidden_states.view(-1, hidden_states.size(-1)) valid_mask = labels_flat != ignore_index - bi_fused_tp_groups = resolve_bi_fused_lm_head_tp_groups(ce_mode, tp_group, lm_head) + batch_invariant_tp_groups = resolve_batch_invariant_lm_head_tp_groups(ce_mode, tp_group, lm_head) has_explicit_loss_reducer = loss_reducer is not None if loss_reducer is None: scale = valid_mask.sum().float() - if bi_fused_tp_groups is not None: - dedicated_group, replica_group = bi_fused_tp_groups + if batch_invariant_tp_groups is not None: + dedicated_group, replica_group = batch_invariant_tp_groups dist.all_reduce(scale, op=dist.ReduceOp.SUM, group=dedicated_group) if replica_group is not None: dist.all_reduce(scale, op=dist.ReduceOp.SUM, group=replica_group) @@ -564,13 +564,13 @@ def _flatten_sampling_metadata(value, name: str): lm_head is not None and (getattr(lm_head, "_glm52_exact_tp16_lm_head", False) or getattr(lm_head, "_dsv4_exact_tp8_lm_head", False)) ) - if ce_mode == "bi_fused": - if tp_group is not None and not exact_lm_head and bi_fused_tp_groups is None: + if ce_mode == "batch_invariant": + if tp_group is not None and not exact_lm_head and batch_invariant_tp_groups is None: raise NotImplementedError( - "ce_mode='bi_fused' supports TP only through the dedicated vocabulary-sharded LM-head TP path" + "ce_mode='batch_invariant' supports TP only through the dedicated vocabulary-sharded LM-head TP path" ) if lm_head is not None and not lm_head_fp32 and not exact_lm_head: - raise NotImplementedError("ce_mode='bi_fused' does not support FP8 lm_head modules") + raise NotImplementedError("ce_mode='batch_invariant' does not support FP8 lm_head modules") if exact_lm_head: if z_loss_coef > 0.0: raise NotImplementedError("The exact GLM-5.2 active-LoRA lm head does not support Z-loss") @@ -598,9 +598,9 @@ def _flatten_sampling_metadata(value, name: str): per_token_loss=per_token_ce.view(original_shape), ) return LossOutput(loss=loss) - if bi_fused_tp_groups is not None: + if batch_invariant_tp_groups is not None: if z_loss_coef > 0.0: - raise NotImplementedError("ce_mode='bi_fused' does not support softmax_auxiliary_loss") + raise NotImplementedError("ce_mode='batch_invariant' does not support softmax_auxiliary_loss") per_token_ce = compute_per_token_ce( hidden_states_flat, weight, @@ -631,7 +631,7 @@ def _flatten_sampling_metadata(value, name: str): # on every rank. Explicit reducers instead promise a local partial, and # their caller owns detached reporting aggregation. global_loss = local_loss.detach().clone() - dedicated_group, replica_group = bi_fused_tp_groups + dedicated_group, replica_group = batch_invariant_tp_groups dist.all_reduce(global_loss, op=dist.ReduceOp.SUM, group=dedicated_group) if replica_group is not None: dist.all_reduce(global_loss, op=dist.ReduceOp.SUM, group=replica_group) @@ -772,8 +772,8 @@ def _flatten_sampling_metadata(value, name: str): per_token_ce = _fused_quack_per_token_ce( hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, tp_group, lm_head_fp32 ) - elif ce_mode == "bi_fused": - per_token_ce = _bi_fused_per_token_ce_checked( + elif ce_mode == "batch_invariant": + per_token_ce = _batch_invariant_per_token_ce_checked( hidden_states_flat, weight, labels_flat, ignore_index, lm_head_fp32, z_loss_enabled ) elif ce_mode == "quack_linear": @@ -842,8 +842,8 @@ def _flatten_sampling_metadata(value, name: str): per_token_ce = _fused_quack_per_token_ce( hidden_states_flat, weight, labels_flat, ignore_index, num_chunks, tp_group, lm_head_fp32 ) - elif ce_mode == "bi_fused": - per_token_ce = _bi_fused_per_token_ce_checked( + elif ce_mode == "batch_invariant": + per_token_ce = _batch_invariant_per_token_ce_checked( hidden_states_flat, weight, labels_flat, ignore_index, lm_head_fp32, z_loss_enabled ) else: # eager mode diff --git a/src/xorl/ops/README.md b/src/xorl/ops/README.md index ddcc04ef..bcade3a3 100644 --- a/src/xorl/ops/README.md +++ b/src/xorl/ops/README.md @@ -53,10 +53,10 @@ serving engine and cannot be renamed unilaterally): | name | meaning | | --- | --- | -| **batch-invariant / `bi_`** | a kernel whose per-element reduction order does not depend on batch composition, so the same token produces the same bits in any batch — the property that makes trainer/sampler logprobs comparable bitwise. | +| **batch-invariant** (`bi_` in serving-twin symbols) | a kernel whose per-element reduction order does not depend on batch composition, so the same token produces the same bits in any batch. First-party names now spell it out (`ce_mode: batch_invariant`, with `bi_fused` as a deprecated alias); the `bi_` prefix survives only inside the serving twins, whose symbols must match the engine. | | **K3 / "zero-K3"** | the k3 KL-divergence estimator between trainer and sampler logprobs for the same tokens; "zero-K3" = bit-identical train/serve forward, the goal of the exact contracts. | | **Class-A / Class-B RoPE** | the two RoPE numerics classes across the trainer/sampler pair: Class A rounds to bf16 per op (8 rounding points); Class B computes one fp32 chain with a single final round (SGLang's fused CUDA rope and its compiled RL-lane path). `rope_fp32_single_round: true` selects Class B. | -| **canonical MoE reduce** | the pinned fixed-order expert-contribution reduction (contributor-leaf arithmetic + final FP64-accumulator cast) shared with serving, versioned by `CANONICAL_MOE_REDUCE_VERSION`. | +| **fixed-order MoE reduce** (formerly "canonical") | the pinned fixed-order expert-contribution reduction (contributor-leaf arithmetic + final FP64-accumulator cast) shared with serving; the contract key string stays `CANONICAL_MOE_REDUCE_VERSION`. | | **one-round SwiGLU** | the exact-contract SwiGLU with a single FP32 rounding point (`one_round_swiglu`), vs the generic fused SwiGLU. | | **families v1 / v2** | versioned batch-invariant kernel families (norms, LM head); v2 is the epilogue-stats generation. `bi_families_v2` is the serving twin module. | | **GKN layout** | grouped expert-weight layout `[G=experts, K=in_features, N=out_features]`. | diff --git a/src/xorl/ops/bi_gemm_configs.py b/src/xorl/ops/bi_gemm_configs.py index bc72e1aa..6500e8fc 100644 --- a/src/xorl/ops/bi_gemm_configs.py +++ b/src/xorl/ops/bi_gemm_configs.py @@ -1,3 +1,7 @@ -"""Moved to ``xorl.ops.exact.bi_gemm_configs`` (issue #78 phase 3); compat stub.""" +"""Moved to ``xorl.ops.exact.batch_invariant_gemm_configs`` (issue #78); alias stub.""" -from xorl.ops.exact.bi_gemm_configs import * # noqa: F401,F403 +import importlib as _importlib +import sys as _sys + + +_sys.modules[__name__] = _importlib.import_module("xorl.ops.exact.batch_invariant_gemm_configs") diff --git a/src/xorl/ops/exact/__init__.py b/src/xorl/ops/exact/__init__.py index 11912ff1..b4fbfc89 100644 --- a/src/xorl/ops/exact/__init__.py +++ b/src/xorl/ops/exact/__init__.py @@ -6,7 +6,7 @@ Here: ``sampling_transforms`` (the replay contract), ``rope_fp32_single_round``, ``moe_fixed_order_leaf``, ``canonical_moe_cast``, ``kernel_config_pin``, -``bi_gemm_configs``, ``block_fp8_native``, ``fused_silu_and_mul``. +``batch_invariant_gemm_configs``, ``block_fp8_native``, ``fused_silu_and_mul``. The modules that exist as literal twins inside the serving engine (``bi_families_v2``, ``batch_invariant_ops``) live in diff --git a/src/xorl/ops/exact/bi_gemm_configs.py b/src/xorl/ops/exact/batch_invariant_gemm_configs.py similarity index 100% rename from src/xorl/ops/exact/bi_gemm_configs.py rename to src/xorl/ops/exact/batch_invariant_gemm_configs.py diff --git a/src/xorl/ops/loss/__init__.py b/src/xorl/ops/loss/__init__.py index 5094f03a..5bc32cc7 100644 --- a/src/xorl/ops/loss/__init__.py +++ b/src/xorl/ops/loss/__init__.py @@ -14,8 +14,10 @@ # Cross-entropy computation mode shared by the local-trainer (TrainingArguments) # and server-runner (ServerArguments) entry points so the Literal stays in sync. -# ``bi_fused`` runs the shared batch-invariant projection and fixed-order LSE. -CrossEntropyMode = Literal["eager", "compiled", "bi_fused", "quack_linear", "fused_quack"] +# ``batch_invariant`` runs the shared batch-invariant projection and fixed-order LSE. +# "bi_fused" is the deprecated alias for "batch_invariant"; both argument entry +# points normalize it before resolution. +CrossEntropyMode = Literal["eager", "compiled", "batch_invariant", "bi_fused", "quack_linear", "fused_quack"] _OBJECTIVE_EXPORTS = frozenset( { diff --git a/src/xorl/ops/loss/batch_invariant_lm_head.py b/src/xorl/ops/loss/batch_invariant_lm_head.py new file mode 100644 index 00000000..eeb18464 --- /dev/null +++ b/src/xorl/ops/loss/batch_invariant_lm_head.py @@ -0,0 +1,905 @@ +"""Trainable wrapper for the batch-invariant fused LM-head logprob contract. + +Forward scores per-token cross-entropy through +:func:`xorl.ops.sglang.batch_invariant_ops.bi_lm_head_selected_logprob` — the K3 +lm-head contract vendored identically in SGLang, so trainer and serving +logprobs are bitwise identical from bit-exact hidden states. The bf16 weight +stays resident (no fp32 lm-head copy). Per-row temperature materializes the +same FP32 ``z * (1/T)`` tensor that serving samples and scores; the scalar-one +call keeps the original non-materialized path. + +Backward is the closed-form CE gradient computed against the saved forward +``lse`` with chunked cuBLAS recompute (stock-numerics class, like the other +fused CE backwards — the contract governs the forward bits only). All three +recompute GEMMs run on the resident bf16 tensors with fp32 accumulation +(``out_dtype=float32``), so no fp32 copy of hidden or weight ever +materializes. +""" + +import math + +import torch +import torch.distributed as dist + +from xorl.ops.exact.sampling_transforms import ( + EXACT_FILTER_ROW_CHUNK, + TOP_K_ALL, + exact_sampling_identity_rows, + exact_sampling_support, + exact_selected_logprob_partitioned_from_support, +) +from xorl.ops.sglang.batch_invariant_ops import ( + BI_LM_HEAD_VOCAB_CHUNK, + bi_lm_head_full_logits, + bi_lm_head_selected_logprob, + bi_lm_head_selected_logprob_from_logits, +) +from xorl.ops.sglang.bi_families_v2 import ( + exact_temperature_scale_fp32_logits, + families_v2_enabled, + head_v2_full_logits_with_lse, + head_v2_selected_logprob, + head_v2_selected_logprob_from_logits, +) + + +_TEMPERATURE_MATERIALIZE_ROW_CHUNK = EXACT_FILTER_ROW_CHUNK +_TP_LOCAL_ROW_CHUNK = 8 + + +def _score_exact_sampling_rows( + logits, + token_ids, + top_ks, + top_ps, + min_ps, + native_selected_score, +): + support = exact_sampling_support(logits, top_ks, top_ps, min_ps) + identity_rows = exact_sampling_identity_rows( + top_ks, + top_ps, + min_ps, + vocab_size=logits.shape[1], + ) + return exact_selected_logprob_partitioned_from_support( + logits, + token_ids, + support, + identity_rows, + native_selected_score, + ) + + +def _tp_collective_layout( + hidden: torch.Tensor, + weight: torch.Tensor, + *, + group: dist.ProcessGroup, + has_temperature: bool, + has_sampling_filter: bool, + use_v2: bool, +) -> tuple[tuple[int, ...], tuple[int, ...], bool, bool]: + """Exchange the small shape/program header before TP payload collectives.""" + + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("ce_mode='batch_invariant' tensor parallelism requires initialized torch.distributed") + world_size = dist.get_world_size(group) + flags = int(has_temperature) | (int(has_sampling_filter) << 1) | (int(use_v2) << 2) + local = torch.tensor( + [hidden.shape[0], weight.shape[0], hidden.shape[1], flags], + dtype=torch.int64, + device=hidden.device, + ) + gathered = torch.empty(world_size * local.numel(), dtype=local.dtype, device=local.device) + dist.all_gather_into_tensor(gathered, local, group=group) + layout = gathered.view(world_size, local.numel()) + + hidden_sizes = layout[:, 2] + program_flags = layout[:, 3] + if bool((hidden_sizes != hidden_sizes[0]).any().item()): + raise ValueError(f"batch_invariant TP hidden widths differ across ranks: {hidden_sizes.cpu().tolist()}") + head_families = program_flags >> 2 + if bool((head_families != head_families[0]).any().item()): + raise ValueError( + f"batch_invariant TP ranks must use the same head-family program, got {head_families.cpu().tolist()}" + ) + vocab_sizes = tuple(int(value) for value in layout[:, 1].cpu().tolist()) + if any(size <= 0 for size in vocab_sizes): + raise ValueError(f"batch_invariant TP requires a non-empty vocabulary shard on every rank, got {vocab_sizes}") + row_counts = tuple(int(value) for value in layout[:, 0].cpu().tolist()) + group_has_temperature = bool(((program_flags & 1) != 0).any().item()) + group_has_sampling_filter = bool(((program_flags & 2) != 0).any().item()) + return row_counts, vocab_sizes, group_has_temperature, group_has_sampling_filter + + +def _tp_broadcast_source_rows( + value: torch.Tensor, + *, + source_rank: int, + start: int, + rows: int, + group: dist.ProcessGroup, +) -> torch.Tensor: + """Broadcast one real source-owner chunk without a full row all-gather.""" + + rank = dist.get_rank(group) + if rank == source_rank: + chunk = value.narrow(0, start, rows).contiguous() + else: + chunk = value.new_empty((rows, *value.shape[1:])) + global_source_rank = dist.get_global_rank(group, source_rank) + dist.broadcast(chunk, src=global_source_rank, group=group) + return chunk + + +def _tp_gather_full_logits( + local_logits: torch.Tensor, + *, + vocab_sizes: tuple[int, ...], + group: dist.ProcessGroup, +) -> torch.Tensor: + """Gather possibly ragged vocabulary shards in process-group rank order.""" + + max_vocab = max(vocab_sizes) + padded = local_logits.new_zeros((local_logits.shape[0], max_vocab)) + padded[:, : local_logits.shape[1]].copy_(local_logits) + world_size = dist.get_world_size(group) + gathered = local_logits.new_empty((world_size * local_logits.shape[0], max_vocab)) + dist.all_gather_into_tensor(gathered, padded.contiguous(), group=group) + rank_major = gathered.view(world_size, local_logits.shape[0], max_vocab) + return torch.cat( + [rank_major[rank, :, :vocab_size] for rank, vocab_size in enumerate(vocab_sizes)], + dim=1, + ).contiguous() + + +def _tp_exact_full_logits( + hidden: torch.Tensor, + local_weight: torch.Tensor, + *, + vocab_sizes: tuple[int, ...], + group: dist.ProcessGroup, + use_v2: bool, + vocab_chunk: int, +) -> torch.Tensor: + if use_v2: + local_logits, _ = head_v2_full_logits_with_lse(hidden, local_weight, temperature=None) + else: + local_logits = bi_lm_head_full_logits(hidden, local_weight, vocab_chunk=vocab_chunk) + return _tp_gather_full_logits(local_logits, vocab_sizes=vocab_sizes, group=group) + + +def _tp_score_full_logits( + full_logits: torch.Tensor, + labels: torch.Tensor, + temperature: torch.Tensor | None, + sampling_transforms: tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None], + *, + use_v2: bool, + vocab_chunk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + transformed_logits = ( + full_logits if temperature is None else exact_temperature_scale_fp32_logits(full_logits, temperature) + ) + top_ks, top_ps, min_ps = sampling_transforms + if use_v2: + native_score = lambda logits, token_ids: head_v2_selected_logprob_from_logits( # noqa: E731 + logits, + token_ids, + temperature=None, + ) + else: + native_score = lambda logits, token_ids: bi_lm_head_selected_logprob_from_logits( # noqa: E731 + logits, + token_ids, + temperature=None, + vocab_chunk=vocab_chunk, + ) + if top_ks is None: + logprob, lse, _ = native_score(transformed_logits, labels) + else: + logprob, lse, _ = _score_exact_sampling_rows( + transformed_logits, + labels, + top_ks, + top_ps, + min_ps, + native_score, + ) + return logprob, lse + + +class _BatchInvariantVocabParallelPerTokenCE(torch.autograd.Function): + """Exact TP forward over distinct row owners with a local-shard CE VJP.""" + + @staticmethod + def forward( + ctx, + local_hidden, + local_weight, + local_labels, + local_valid, + local_temperature, + local_top_ks, + local_top_ps, + local_min_ps, + tp_group, + vocab_chunk, + ): + use_v2 = families_v2_enabled() + has_temperature = local_temperature is not None + has_sampling_filter = local_top_ks is not None + row_counts, vocab_sizes, has_temperature, has_sampling_filter = _tp_collective_layout( + local_hidden, + local_weight, + group=tp_group, + has_temperature=has_temperature, + has_sampling_filter=has_sampling_filter, + use_v2=use_v2, + ) + if has_temperature and local_temperature is None: + local_temperature = torch.ones( + local_hidden.shape[0], + dtype=torch.float32, + device=local_hidden.device, + ) + if has_sampling_filter and local_top_ks is None: + local_top_ks = torch.full( + (local_hidden.shape[0],), + TOP_K_ALL, + dtype=torch.int64, + device=local_hidden.device, + ) + local_top_ps = torch.ones(local_hidden.shape[0], dtype=torch.float32, device=local_hidden.device) + local_min_ps = torch.zeros(local_hidden.shape[0], dtype=torch.float32, device=local_hidden.device) + rank = dist.get_rank(tp_group) + local_ce_chunks: list[torch.Tensor] = [] + lse_slots: list[torch.Tensor] = [] + + for source_rank, source_rows in enumerate(row_counts): + for start in range(0, source_rows, _TP_LOCAL_ROW_CHUNK): + rows = min(_TP_LOCAL_ROW_CHUNK, source_rows - start) + hidden = _tp_broadcast_source_rows( + local_hidden, + source_rank=source_rank, + start=start, + rows=rows, + group=tp_group, + ) + labels = _tp_broadcast_source_rows( + local_labels, + source_rank=source_rank, + start=start, + rows=rows, + group=tp_group, + ) + valid = _tp_broadcast_source_rows( + local_valid, + source_rank=source_rank, + start=start, + rows=rows, + group=tp_group, + ) + temperature = None + if has_temperature: + temperature = _tp_broadcast_source_rows( + local_temperature, + source_rank=source_rank, + start=start, + rows=rows, + group=tp_group, + ) + sampling_transforms = (None, None, None) + if has_sampling_filter: + top_ks = _tp_broadcast_source_rows( + local_top_ks, + source_rank=source_rank, + start=start, + rows=rows, + group=tp_group, + ) + top_ps = _tp_broadcast_source_rows( + local_top_ps, + source_rank=source_rank, + start=start, + rows=rows, + group=tp_group, + ) + min_ps = _tp_broadcast_source_rows( + local_min_ps, + source_rank=source_rank, + start=start, + rows=rows, + group=tp_group, + ) + sampling_transforms = (top_ks, top_ps, min_ps) + + full_logits = _tp_exact_full_logits( + hidden, + local_weight, + vocab_sizes=vocab_sizes, + group=tp_group, + use_v2=use_v2, + vocab_chunk=vocab_chunk, + ) + logprob, lse = _tp_score_full_logits( + full_logits, + labels, + temperature, + sampling_transforms, + use_v2=use_v2, + vocab_chunk=vocab_chunk, + ) + if rank == source_rank: + local_ce_chunks.append(torch.where(valid, -logprob, torch.zeros_like(logprob))) + lse_slots.append(lse) + + empty_float = torch.empty((0,), dtype=torch.float32, device=local_hidden.device) + empty_long = torch.empty((0,), dtype=torch.int64, device=local_hidden.device) + lse_tensor = ( + torch.cat(lse_slots) if lse_slots else torch.empty((0,), dtype=torch.float32, device=local_hidden.device) + ) + ctx.tp_group = tp_group + ctx.vocab_chunk = int(vocab_chunk) + ctx.use_v2 = use_v2 + ctx.has_temperature = has_temperature + ctx.has_sampling_filter = has_sampling_filter + ctx.row_counts = row_counts + ctx.vocab_sizes = vocab_sizes + ctx.save_for_backward( + local_hidden, + local_weight, + local_labels, + local_valid, + local_temperature if has_temperature else empty_float, + local_top_ks if has_sampling_filter else empty_long, + local_top_ps if has_sampling_filter else empty_float, + local_min_ps if has_sampling_filter else empty_float, + lse_tensor, + ) + if local_ce_chunks: + return torch.cat(local_ce_chunks) + return empty_float + + @staticmethod + def backward(ctx, grad_local_ce): + ( + local_hidden, + local_weight, + local_labels, + local_valid, + stored_temperature, + stored_top_ks, + stored_top_ps, + stored_min_ps, + lse_slots, + ) = ctx.saved_tensors + if grad_local_ce is None: + grad_local_ce = torch.zeros( + local_hidden.shape[0], + dtype=torch.float32, + device=local_hidden.device, + ) + + group = ctx.tp_group + row_counts = ctx.row_counts + vocab_sizes = ctx.vocab_sizes + rank = dist.get_rank(group) + vocab_offset = sum(vocab_sizes[:rank]) + local_vocab = vocab_sizes[rank] + need_hidden = ctx.needs_input_grad[0] + need_weight = ctx.needs_input_grad[1] + need_hidden_tensor = torch.tensor(int(need_hidden), dtype=torch.int64, device=local_hidden.device) + dist.all_reduce(need_hidden_tensor, op=dist.ReduceOp.MAX, group=group) + compute_hidden = bool(need_hidden_tensor.item()) + grad_hidden = torch.zeros_like(local_hidden) if need_hidden else None + grad_weight = ( + torch.zeros(local_weight.shape, dtype=torch.float32, device=local_weight.device) if need_weight else None + ) + grad_local_ce = grad_local_ce.float().contiguous() + lse_offset = 0 + + for source_rank, source_rows in enumerate(row_counts): + for start in range(0, source_rows, _TP_LOCAL_ROW_CHUNK): + rows_in_chunk = min(_TP_LOCAL_ROW_CHUNK, source_rows - start) + hidden = _tp_broadcast_source_rows( + local_hidden, + source_rank=source_rank, + start=start, + rows=rows_in_chunk, + group=group, + ) + labels = _tp_broadcast_source_rows( + local_labels, + source_rank=source_rank, + start=start, + rows=rows_in_chunk, + group=group, + ) + valid = _tp_broadcast_source_rows( + local_valid, + source_rank=source_rank, + start=start, + rows=rows_in_chunk, + group=group, + ) + grad_ce = _tp_broadcast_source_rows( + grad_local_ce, + source_rank=source_rank, + start=start, + rows=rows_in_chunk, + group=group, + ) + temperature = None + if ctx.has_temperature: + temperature = _tp_broadcast_source_rows( + stored_temperature, + source_rank=source_rank, + start=start, + rows=rows_in_chunk, + group=group, + ) + + local_support = None + selected_support = torch.ones_like(valid) + if ctx.has_sampling_filter: + top_ks = _tp_broadcast_source_rows( + stored_top_ks, + source_rank=source_rank, + start=start, + rows=rows_in_chunk, + group=group, + ) + top_ps = _tp_broadcast_source_rows( + stored_top_ps, + source_rank=source_rank, + start=start, + rows=rows_in_chunk, + group=group, + ) + min_ps = _tp_broadcast_source_rows( + stored_min_ps, + source_rank=source_rank, + start=start, + rows=rows_in_chunk, + group=group, + ) + exact_logits = _tp_exact_full_logits( + hidden, + local_weight, + vocab_sizes=vocab_sizes, + group=group, + use_v2=ctx.use_v2, + vocab_chunk=ctx.vocab_chunk, + ) + transformed_logits = ( + exact_logits + if temperature is None + else exact_temperature_scale_fp32_logits(exact_logits, temperature) + ) + support = exact_sampling_support(transformed_logits, top_ks, top_ps, min_ps) + selected_support = support.gather(1, labels.unsqueeze(1)).squeeze(1) + local_support = support[:, vocab_offset : vocab_offset + local_vocab] + + g = (grad_ce * valid * selected_support).float() + local_logits = torch.mm(hidden, local_weight.t(), out_dtype=torch.float32) + inv_temperature = None + if temperature is not None: + inv_temperature = (1.0 / temperature).unsqueeze(1) + local_logits *= inv_temperature + lse = lse_slots.narrow(0, lse_offset, rows_in_chunk) + lse_offset += rows_in_chunk + grad_logits = local_logits.sub_(lse.unsqueeze(1)).exp_() + grad_logits *= g.unsqueeze(1) + if local_support is not None: + grad_logits *= local_support + target_in_shard = ( + selected_support & valid & (labels >= vocab_offset) & (labels < vocab_offset + local_vocab) + ) + rows = torch.arange(labels.shape[0], device=labels.device) + grad_logits[rows[target_in_shard], labels[target_in_shard] - vocab_offset] -= g[target_in_shard] + if inv_temperature is not None: + grad_logits *= inv_temperature + grad_logits_bf16 = grad_logits.to(local_hidden.dtype) + + if compute_hidden: + grad_hidden_chunk = torch.mm( + grad_logits_bf16, + local_weight, + out_dtype=torch.float32, + ) + dist.all_reduce(grad_hidden_chunk, op=dist.ReduceOp.SUM, group=group) + if need_hidden and rank == source_rank: + grad_hidden.narrow(0, start, rows_in_chunk).copy_(grad_hidden_chunk.to(grad_hidden.dtype)) + if need_weight: + grad_weight.add_(torch.mm(grad_logits_bf16.t(), hidden, out_dtype=torch.float32)) + + return ( + grad_hidden, + grad_weight.to(local_weight.dtype) if need_weight else None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +class _BatchInvariantLmHeadPerTokenCE(torch.autograd.Function): + @staticmethod + def forward( + ctx, + hidden, + weight, + labels_safe, + valid_mask, + temp_row, + top_ks, + top_ps, + min_ps, + vocab_chunk, + ): + use_v2 = families_v2_enabled() + has_sampling_filter = top_ks is not None + if temp_row is None and not has_sampling_filter and use_v2: + # head v2 (families-v2 migration): same GEMM K-chain, epilogue-stats + # online LSE; logits never materialize; backward only consumes lse. + logprob, lse, _ = head_v2_selected_logprob(hidden, weight, labels_safe, temperature=temp_row) + elif temp_row is None and not has_sampling_filter: + logprob, lse, _ = bi_lm_head_selected_logprob( + hidden, weight, labels_safe, temperature=temp_row, vocab_chunk=vocab_chunk + ) + else: + logprob_chunks = [] + lse_chunks = [] + for start in range(0, hidden.shape[0], _TEMPERATURE_MATERIALIZE_ROW_CHUNK): + end = min(start + _TEMPERATURE_MATERIALIZE_ROW_CHUNK, hidden.shape[0]) + hidden_chunk = hidden[start:end] + labels_chunk = labels_safe[start:end] + temperature_chunk = None if temp_row is None else temp_row[start:end] + if use_v2: + logits, _ = head_v2_full_logits_with_lse(hidden_chunk, weight, temperature=None) + transformed_logits = ( + logits + if temperature_chunk is None + else exact_temperature_scale_fp32_logits(logits, temperature_chunk) + ) + if has_sampling_filter: + logprob_chunk, lse_chunk, _ = _score_exact_sampling_rows( + transformed_logits, + labels_chunk, + top_ks[start:end], + top_ps[start:end], + min_ps[start:end], + lambda score_logits, score_ids: head_v2_selected_logprob_from_logits( + score_logits, + score_ids, + temperature=None, + ), + ) + else: + logprob_chunk, lse_chunk, _ = head_v2_selected_logprob_from_logits( + transformed_logits, + labels_chunk, + temperature=None, + ) + else: + logits = bi_lm_head_full_logits(hidden_chunk, weight, vocab_chunk=vocab_chunk) + transformed_logits = ( + logits + if temperature_chunk is None + else exact_temperature_scale_fp32_logits(logits, temperature_chunk) + ) + if has_sampling_filter: + logprob_chunk, lse_chunk, _ = _score_exact_sampling_rows( + transformed_logits, + labels_chunk, + top_ks[start:end], + top_ps[start:end], + min_ps[start:end], + lambda score_logits, score_ids: bi_lm_head_selected_logprob_from_logits( + score_logits, + score_ids, + temperature=None, + vocab_chunk=vocab_chunk, + ), + ) + else: + logprob_chunk, lse_chunk, _ = bi_lm_head_selected_logprob_from_logits( + transformed_logits, + labels_chunk, + temperature=None, + vocab_chunk=vocab_chunk, + ) + logprob_chunks.append(logprob_chunk) + lse_chunks.append(lse_chunk) + if logprob_chunks: + logprob = torch.cat(logprob_chunks, dim=0) + lse = torch.cat(lse_chunks, dim=0) + else: + logprob = torch.empty((0,), dtype=torch.float32, device=hidden.device) + lse = logprob.clone() + ctx.save_for_backward( + hidden, + weight, + labels_safe, + valid_mask, + lse, + temp_row, + top_ks, + top_ps, + min_ps, + ) + ctx.vocab_chunk = vocab_chunk + ctx.use_v2 = use_v2 + return torch.where(valid_mask, -logprob, torch.zeros_like(logprob)) + + @staticmethod + def backward(ctx, grad_ce): + hidden, weight, labels, valid_mask, lse, temp_row, top_ks, top_ps, min_ps = ctx.saved_tensors + vocab_chunk = ctx.vocab_chunk + n_tokens = hidden.shape[0] + vocab = weight.shape[0] + need_h = ctx.needs_input_grad[0] + need_w = ctx.needs_input_grad[1] + + if top_ks is not None: + grad_h = torch.zeros(hidden.shape, dtype=torch.float32, device=hidden.device) if need_h else None + grad_w = torch.zeros_like(weight) if need_w else None + for row_start in range(0, n_tokens, _TEMPERATURE_MATERIALIZE_ROW_CHUNK): + row_end = min(row_start + _TEMPERATURE_MATERIALIZE_ROW_CHUNK, n_tokens) + hidden_chunk = hidden[row_start:row_end] + labels_chunk = labels[row_start:row_end] + valid_chunk = valid_mask[row_start:row_end] + temperature_chunk = None if temp_row is None else temp_row[row_start:row_end] + if ctx.use_v2: + exact_logits, _ = head_v2_full_logits_with_lse(hidden_chunk, weight, temperature=None) + else: + exact_logits = bi_lm_head_full_logits(hidden_chunk, weight, vocab_chunk=vocab_chunk) + transformed_logits = ( + exact_logits + if temperature_chunk is None + else exact_temperature_scale_fp32_logits(exact_logits, temperature_chunk) + ) + support = exact_sampling_support( + transformed_logits, + top_ks[row_start:row_end], + top_ps[row_start:row_end], + min_ps[row_start:row_end], + ) + selected_support = support.gather(1, labels_chunk.unsqueeze(1)).squeeze(1) + g = (grad_ce[row_start:row_end] * valid_chunk * selected_support).float() + g_col = g.unsqueeze(1) + lse_col = lse[row_start:row_end].unsqueeze(1) + inv_t = None if temperature_chunk is None else (1.0 / temperature_chunk).unsqueeze(1) + grad_h_chunk = ( + torch.zeros(hidden_chunk.shape, dtype=torch.float32, device=hidden.device) if need_h else None + ) + rows = torch.arange(row_end - row_start, device=hidden.device) + + for col_start in range(0, vocab, vocab_chunk): + col_end = min(col_start + vocab_chunk, vocab) + w_c = weight[col_start:col_end] + logits_c = torch.mm(hidden_chunk, w_c.t(), out_dtype=torch.float32) + if inv_t is not None: + logits_c *= inv_t + grad_z = logits_c.sub_(lse_col).exp_().mul_(g_col) + grad_z *= support[:, col_start:col_end] + in_chunk = selected_support & (labels_chunk >= col_start) & (labels_chunk < col_end) + grad_z[rows[in_chunk], labels_chunk[in_chunk] - col_start] -= g[in_chunk] + if inv_t is not None: + grad_z *= inv_t + grad_z16 = grad_z.to(hidden.dtype) + if need_h: + torch.addmm( + grad_h_chunk, + grad_z16, + w_c, + out_dtype=torch.float32, + out=grad_h_chunk, + ) + if need_w: + grad_w[col_start:col_end].add_( + torch.mm(grad_z16.t(), hidden_chunk, out_dtype=torch.float32).to(weight.dtype) + ) + if need_h: + grad_h[row_start:row_end] = grad_h_chunk + + return ( + grad_h.to(hidden.dtype) if need_h else None, + grad_w if need_w else None, + None, + None, + None, + None, + None, + None, + None, + ) + + selected_support = torch.ones_like(valid_mask) + g = (grad_ce * valid_mask).float() + g_col = g.unsqueeze(1) + lse_col = lse.unsqueeze(1) + inv_t = None if temp_row is None else (1.0 / temp_row).unsqueeze(1) + grad_h = torch.zeros(hidden.shape, dtype=torch.float32, device=hidden.device) if need_h else None + grad_w = torch.empty_like(weight) if need_w else None + rows = torch.arange(n_tokens, device=hidden.device) + + for col_start in range(0, vocab, vocab_chunk): + col_end = min(col_start + vocab_chunk, vocab) + w_c = weight[col_start:col_end] + # bf16 tensor-core GEMM, fp32 accumulate + fp32 out (no fp32 copies) + logits_c = torch.mm(hidden, w_c.t(), out_dtype=torch.float32) + if inv_t is not None: + logits_c *= inv_t + grad_z = logits_c.sub_(lse_col).exp_().mul_(g_col) + in_chunk = selected_support & (labels >= col_start) & (labels < col_end) + grad_z[rows[in_chunk], labels[in_chunk] - col_start] -= g[in_chunk] + if inv_t is not None: + grad_z *= inv_t # dy/dz = 1/T + grad_z16 = grad_z.to(hidden.dtype) + if need_h: + torch.addmm(grad_h, grad_z16, w_c, out_dtype=torch.float32, out=grad_h) + if need_w: + grad_w[col_start:col_end] = torch.mm(grad_z16.t(), hidden, out_dtype=torch.float32).to(weight.dtype) + + return ( + grad_h.to(hidden.dtype) if need_h else None, + grad_w if need_w else None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def batch_invariant_per_token_ce( + hidden_states: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + ignore_index: int = -100, + temperature: float | torch.Tensor = 1.0, + top_ks: torch.Tensor | None = None, + top_ps: torch.Tensor | None = None, + min_ps: torch.Tensor | None = None, + vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, +) -> torch.Tensor: + """Per-token CE (``-log p(labels)``; 0 at ignored positions) through the + batch-invariant lm-head contract. Requires CUDA bf16 hidden/weight; the + fp32-class numerics come from the contract itself, so ``lm_head_fp32`` is + implied rather than materialized. Per-row temperature scores the same + materialized FP32 ``z * (1/T)`` tensor as serving.""" + if not hidden_states.is_cuda: + raise ValueError("ce_mode='batch_invariant' requires CUDA tensors") + if hidden_states.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + raise ValueError("ce_mode='batch_invariant' requires bf16 hidden states and lm-head weight") + valid_mask = labels != ignore_index + labels_safe = torch.where(valid_mask, labels, torch.zeros_like(labels)) + if isinstance(temperature, torch.Tensor): + if temperature.dtype is not torch.float32: + raise TypeError("ce_mode='batch_invariant' requires per-row FP32 temperature") + if temperature.device != hidden_states.device or tuple(temperature.shape) != (hidden_states.shape[0],): + raise ValueError("ce_mode='batch_invariant' requires temperature aligned one-to-one with hidden-state rows") + if not temperature.is_contiguous() or temperature.requires_grad: + raise ValueError("ce_mode='batch_invariant' requires contiguous, non-differentiable temperature metadata") + torch._assert_async( + (torch.isfinite(temperature) & (temperature > 0)).all(), + "ce_mode='batch_invariant' requires finite temperature > 0", + ) + temp_row = temperature + elif temperature == 1.0: + temp_row = None + else: + temperature = float(temperature) + if not math.isfinite(temperature) or temperature <= 0: + raise ValueError("ce_mode='batch_invariant' requires finite temperature > 0") + temp_row = torch.full((hidden_states.shape[0],), temperature, dtype=torch.float32, device=hidden_states.device) + if (top_ks is None, top_ps is None, min_ps is None).count(True) not in (0, 3): + raise ValueError("ce_mode='batch_invariant' requires all or none of top-k/top-p/min-p row metadata") + return _BatchInvariantLmHeadPerTokenCE.apply( + hidden_states, + weight, + labels_safe, + valid_mask, + temp_row, + top_ks, + top_ps, + min_ps, + vocab_chunk, + ) + + +def batch_invariant_vocab_parallel_per_token_ce( + hidden_states: torch.Tensor, + local_weight: torch.Tensor, + labels: torch.Tensor, + tp_group: dist.ProcessGroup, + ignore_index: int = -100, + temperature: float | torch.Tensor = 1.0, + top_ks: torch.Tensor | None = None, + top_ps: torch.Tensor | None = None, + min_ps: torch.Tensor | None = None, + vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, +) -> torch.Tensor: + """Run the batch-invariant LM head over rank-local vocab and token shards. + + LM-head-only TP is composed from DP/CP row owners, so ranks can own + different (or zero) token rows while also owning disjoint vocab rows. The + forward broadcasts bounded chunks from each row owner, computes the exact + local FP32 logits, gathers vocabulary shards in group-rank order, and + scores through the same from-logits tail as serving. Only the source + owner's rows are returned. The custom backward reduces hidden gradients + over vocab ranks and leaves each weight-shard gradient local for the + existing replica sync. + """ + + if not hidden_states.is_cuda or not local_weight.is_cuda: + raise ValueError("ce_mode='batch_invariant' tensor parallelism requires CUDA tensors") + if hidden_states.ndim != 2 or local_weight.ndim != 2: + raise ValueError("ce_mode='batch_invariant' tensor parallelism requires two-dimensional hidden and weight") + if hidden_states.shape[1] != local_weight.shape[1]: + raise ValueError("ce_mode='batch_invariant' tensor parallelism received mismatched hidden dimensions") + if hidden_states.dtype is not torch.bfloat16 or local_weight.dtype is not torch.bfloat16: + raise ValueError("ce_mode='batch_invariant' tensor parallelism requires bf16 hidden states and lm-head weight") + if labels.ndim != 1 or labels.shape[0] != hidden_states.shape[0]: + raise ValueError("ce_mode='batch_invariant' tensor parallelism requires one label per local hidden row") + hidden_states = hidden_states.contiguous() + local_weight = local_weight.contiguous() + labels = labels.contiguous() + valid_mask = labels != ignore_index + labels_safe = torch.where(valid_mask, labels, torch.zeros_like(labels)) + + if isinstance(temperature, torch.Tensor): + if temperature.dtype is not torch.float32: + raise TypeError("ce_mode='batch_invariant' requires per-row FP32 temperature") + if temperature.device != hidden_states.device or tuple(temperature.shape) != (hidden_states.shape[0],): + raise ValueError("ce_mode='batch_invariant' requires temperature aligned one-to-one with local hidden rows") + if not temperature.is_contiguous() or temperature.requires_grad: + raise ValueError("ce_mode='batch_invariant' requires contiguous, non-differentiable temperature metadata") + torch._assert_async( + (torch.isfinite(temperature) & (temperature > 0)).all(), + "ce_mode='batch_invariant' requires finite temperature > 0", + ) + temp_row = temperature + elif temperature == 1.0: + temp_row = None + else: + temperature = float(temperature) + if not math.isfinite(temperature) or temperature <= 0: + raise ValueError("ce_mode='batch_invariant' requires finite temperature > 0") + temp_row = torch.full( + (hidden_states.shape[0],), + temperature, + dtype=torch.float32, + device=hidden_states.device, + ) + if (top_ks is None, top_ps is None, min_ps is None).count(True) not in (0, 3): + raise ValueError("ce_mode='batch_invariant' requires all or none of top-k/top-p/min-p row metadata") + for name, value in (("top_ks", top_ks), ("top_ps", top_ps), ("min_ps", min_ps)): + if value is not None and ( + value.device != hidden_states.device + or tuple(value.shape) != (hidden_states.shape[0],) + or not value.is_contiguous() + ): + raise ValueError(f"ce_mode='batch_invariant' requires contiguous local-row {name} metadata") + + return _BatchInvariantVocabParallelPerTokenCE.apply( + hidden_states, + local_weight, + labels_safe, + valid_mask, + temp_row, + top_ks, + top_ps, + min_ps, + tp_group, + vocab_chunk, + ) diff --git a/src/xorl/ops/loss/bi_fused_lm_head.py b/src/xorl/ops/loss/bi_fused_lm_head.py index 292eff69..69d46e6e 100644 --- a/src/xorl/ops/loss/bi_fused_lm_head.py +++ b/src/xorl/ops/loss/bi_fused_lm_head.py @@ -1,903 +1,7 @@ -"""Trainable wrapper for the batch-invariant fused LM-head logprob contract. +"""Moved to ``xorl.ops.loss.batch_invariant_lm_head`` (issue #78); alias stub.""" -Forward scores per-token cross-entropy through -:func:`xorl.ops.sglang.batch_invariant_ops.bi_lm_head_selected_logprob` — the K3 -lm-head contract vendored identically in SGLang, so trainer and serving -logprobs are bitwise identical from bit-exact hidden states. The bf16 weight -stays resident (no fp32 lm-head copy). Per-row temperature materializes the -same FP32 ``z * (1/T)`` tensor that serving samples and scores; the scalar-one -call keeps the original non-materialized path. +import importlib as _importlib +import sys as _sys -Backward is the closed-form CE gradient computed against the saved forward -``lse`` with chunked cuBLAS recompute (stock-numerics class, like the other -fused CE backwards — the contract governs the forward bits only). All three -recompute GEMMs run on the resident bf16 tensors with fp32 accumulation -(``out_dtype=float32``), so no fp32 copy of hidden or weight ever -materializes. -""" -import math - -import torch -import torch.distributed as dist - -from xorl.ops.exact.sampling_transforms import ( - EXACT_FILTER_ROW_CHUNK, - TOP_K_ALL, - exact_sampling_identity_rows, - exact_sampling_support, - exact_selected_logprob_partitioned_from_support, -) -from xorl.ops.sglang.batch_invariant_ops import ( - BI_LM_HEAD_VOCAB_CHUNK, - bi_lm_head_full_logits, - bi_lm_head_selected_logprob, - bi_lm_head_selected_logprob_from_logits, -) -from xorl.ops.sglang.bi_families_v2 import ( - exact_temperature_scale_fp32_logits, - families_v2_enabled, - head_v2_full_logits_with_lse, - head_v2_selected_logprob, - head_v2_selected_logprob_from_logits, -) - - -_TEMPERATURE_MATERIALIZE_ROW_CHUNK = EXACT_FILTER_ROW_CHUNK -_TP_LOCAL_ROW_CHUNK = 8 - - -def _score_exact_sampling_rows( - logits, - token_ids, - top_ks, - top_ps, - min_ps, - native_selected_score, -): - support = exact_sampling_support(logits, top_ks, top_ps, min_ps) - identity_rows = exact_sampling_identity_rows( - top_ks, - top_ps, - min_ps, - vocab_size=logits.shape[1], - ) - return exact_selected_logprob_partitioned_from_support( - logits, - token_ids, - support, - identity_rows, - native_selected_score, - ) - - -def _tp_collective_layout( - hidden: torch.Tensor, - weight: torch.Tensor, - *, - group: dist.ProcessGroup, - has_temperature: bool, - has_sampling_filter: bool, - use_v2: bool, -) -> tuple[tuple[int, ...], tuple[int, ...], bool, bool]: - """Exchange the small shape/program header before TP payload collectives.""" - - if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError("ce_mode='bi_fused' tensor parallelism requires initialized torch.distributed") - world_size = dist.get_world_size(group) - flags = int(has_temperature) | (int(has_sampling_filter) << 1) | (int(use_v2) << 2) - local = torch.tensor( - [hidden.shape[0], weight.shape[0], hidden.shape[1], flags], - dtype=torch.int64, - device=hidden.device, - ) - gathered = torch.empty(world_size * local.numel(), dtype=local.dtype, device=local.device) - dist.all_gather_into_tensor(gathered, local, group=group) - layout = gathered.view(world_size, local.numel()) - - hidden_sizes = layout[:, 2] - program_flags = layout[:, 3] - if bool((hidden_sizes != hidden_sizes[0]).any().item()): - raise ValueError(f"bi_fused TP hidden widths differ across ranks: {hidden_sizes.cpu().tolist()}") - head_families = program_flags >> 2 - if bool((head_families != head_families[0]).any().item()): - raise ValueError(f"bi_fused TP ranks must use the same head-family program, got {head_families.cpu().tolist()}") - vocab_sizes = tuple(int(value) for value in layout[:, 1].cpu().tolist()) - if any(size <= 0 for size in vocab_sizes): - raise ValueError(f"bi_fused TP requires a non-empty vocabulary shard on every rank, got {vocab_sizes}") - row_counts = tuple(int(value) for value in layout[:, 0].cpu().tolist()) - group_has_temperature = bool(((program_flags & 1) != 0).any().item()) - group_has_sampling_filter = bool(((program_flags & 2) != 0).any().item()) - return row_counts, vocab_sizes, group_has_temperature, group_has_sampling_filter - - -def _tp_broadcast_source_rows( - value: torch.Tensor, - *, - source_rank: int, - start: int, - rows: int, - group: dist.ProcessGroup, -) -> torch.Tensor: - """Broadcast one real source-owner chunk without a full row all-gather.""" - - rank = dist.get_rank(group) - if rank == source_rank: - chunk = value.narrow(0, start, rows).contiguous() - else: - chunk = value.new_empty((rows, *value.shape[1:])) - global_source_rank = dist.get_global_rank(group, source_rank) - dist.broadcast(chunk, src=global_source_rank, group=group) - return chunk - - -def _tp_gather_full_logits( - local_logits: torch.Tensor, - *, - vocab_sizes: tuple[int, ...], - group: dist.ProcessGroup, -) -> torch.Tensor: - """Gather possibly ragged vocabulary shards in process-group rank order.""" - - max_vocab = max(vocab_sizes) - padded = local_logits.new_zeros((local_logits.shape[0], max_vocab)) - padded[:, : local_logits.shape[1]].copy_(local_logits) - world_size = dist.get_world_size(group) - gathered = local_logits.new_empty((world_size * local_logits.shape[0], max_vocab)) - dist.all_gather_into_tensor(gathered, padded.contiguous(), group=group) - rank_major = gathered.view(world_size, local_logits.shape[0], max_vocab) - return torch.cat( - [rank_major[rank, :, :vocab_size] for rank, vocab_size in enumerate(vocab_sizes)], - dim=1, - ).contiguous() - - -def _tp_exact_full_logits( - hidden: torch.Tensor, - local_weight: torch.Tensor, - *, - vocab_sizes: tuple[int, ...], - group: dist.ProcessGroup, - use_v2: bool, - vocab_chunk: int, -) -> torch.Tensor: - if use_v2: - local_logits, _ = head_v2_full_logits_with_lse(hidden, local_weight, temperature=None) - else: - local_logits = bi_lm_head_full_logits(hidden, local_weight, vocab_chunk=vocab_chunk) - return _tp_gather_full_logits(local_logits, vocab_sizes=vocab_sizes, group=group) - - -def _tp_score_full_logits( - full_logits: torch.Tensor, - labels: torch.Tensor, - temperature: torch.Tensor | None, - sampling_transforms: tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None], - *, - use_v2: bool, - vocab_chunk: int, -) -> tuple[torch.Tensor, torch.Tensor]: - transformed_logits = ( - full_logits if temperature is None else exact_temperature_scale_fp32_logits(full_logits, temperature) - ) - top_ks, top_ps, min_ps = sampling_transforms - if use_v2: - native_score = lambda logits, token_ids: head_v2_selected_logprob_from_logits( # noqa: E731 - logits, - token_ids, - temperature=None, - ) - else: - native_score = lambda logits, token_ids: bi_lm_head_selected_logprob_from_logits( # noqa: E731 - logits, - token_ids, - temperature=None, - vocab_chunk=vocab_chunk, - ) - if top_ks is None: - logprob, lse, _ = native_score(transformed_logits, labels) - else: - logprob, lse, _ = _score_exact_sampling_rows( - transformed_logits, - labels, - top_ks, - top_ps, - min_ps, - native_score, - ) - return logprob, lse - - -class _BiFusedVocabParallelPerTokenCE(torch.autograd.Function): - """Exact TP forward over distinct row owners with a local-shard CE VJP.""" - - @staticmethod - def forward( - ctx, - local_hidden, - local_weight, - local_labels, - local_valid, - local_temperature, - local_top_ks, - local_top_ps, - local_min_ps, - tp_group, - vocab_chunk, - ): - use_v2 = families_v2_enabled() - has_temperature = local_temperature is not None - has_sampling_filter = local_top_ks is not None - row_counts, vocab_sizes, has_temperature, has_sampling_filter = _tp_collective_layout( - local_hidden, - local_weight, - group=tp_group, - has_temperature=has_temperature, - has_sampling_filter=has_sampling_filter, - use_v2=use_v2, - ) - if has_temperature and local_temperature is None: - local_temperature = torch.ones( - local_hidden.shape[0], - dtype=torch.float32, - device=local_hidden.device, - ) - if has_sampling_filter and local_top_ks is None: - local_top_ks = torch.full( - (local_hidden.shape[0],), - TOP_K_ALL, - dtype=torch.int64, - device=local_hidden.device, - ) - local_top_ps = torch.ones(local_hidden.shape[0], dtype=torch.float32, device=local_hidden.device) - local_min_ps = torch.zeros(local_hidden.shape[0], dtype=torch.float32, device=local_hidden.device) - rank = dist.get_rank(tp_group) - local_ce_chunks: list[torch.Tensor] = [] - lse_slots: list[torch.Tensor] = [] - - for source_rank, source_rows in enumerate(row_counts): - for start in range(0, source_rows, _TP_LOCAL_ROW_CHUNK): - rows = min(_TP_LOCAL_ROW_CHUNK, source_rows - start) - hidden = _tp_broadcast_source_rows( - local_hidden, - source_rank=source_rank, - start=start, - rows=rows, - group=tp_group, - ) - labels = _tp_broadcast_source_rows( - local_labels, - source_rank=source_rank, - start=start, - rows=rows, - group=tp_group, - ) - valid = _tp_broadcast_source_rows( - local_valid, - source_rank=source_rank, - start=start, - rows=rows, - group=tp_group, - ) - temperature = None - if has_temperature: - temperature = _tp_broadcast_source_rows( - local_temperature, - source_rank=source_rank, - start=start, - rows=rows, - group=tp_group, - ) - sampling_transforms = (None, None, None) - if has_sampling_filter: - top_ks = _tp_broadcast_source_rows( - local_top_ks, - source_rank=source_rank, - start=start, - rows=rows, - group=tp_group, - ) - top_ps = _tp_broadcast_source_rows( - local_top_ps, - source_rank=source_rank, - start=start, - rows=rows, - group=tp_group, - ) - min_ps = _tp_broadcast_source_rows( - local_min_ps, - source_rank=source_rank, - start=start, - rows=rows, - group=tp_group, - ) - sampling_transforms = (top_ks, top_ps, min_ps) - - full_logits = _tp_exact_full_logits( - hidden, - local_weight, - vocab_sizes=vocab_sizes, - group=tp_group, - use_v2=use_v2, - vocab_chunk=vocab_chunk, - ) - logprob, lse = _tp_score_full_logits( - full_logits, - labels, - temperature, - sampling_transforms, - use_v2=use_v2, - vocab_chunk=vocab_chunk, - ) - if rank == source_rank: - local_ce_chunks.append(torch.where(valid, -logprob, torch.zeros_like(logprob))) - lse_slots.append(lse) - - empty_float = torch.empty((0,), dtype=torch.float32, device=local_hidden.device) - empty_long = torch.empty((0,), dtype=torch.int64, device=local_hidden.device) - lse_tensor = ( - torch.cat(lse_slots) if lse_slots else torch.empty((0,), dtype=torch.float32, device=local_hidden.device) - ) - ctx.tp_group = tp_group - ctx.vocab_chunk = int(vocab_chunk) - ctx.use_v2 = use_v2 - ctx.has_temperature = has_temperature - ctx.has_sampling_filter = has_sampling_filter - ctx.row_counts = row_counts - ctx.vocab_sizes = vocab_sizes - ctx.save_for_backward( - local_hidden, - local_weight, - local_labels, - local_valid, - local_temperature if has_temperature else empty_float, - local_top_ks if has_sampling_filter else empty_long, - local_top_ps if has_sampling_filter else empty_float, - local_min_ps if has_sampling_filter else empty_float, - lse_tensor, - ) - if local_ce_chunks: - return torch.cat(local_ce_chunks) - return empty_float - - @staticmethod - def backward(ctx, grad_local_ce): - ( - local_hidden, - local_weight, - local_labels, - local_valid, - stored_temperature, - stored_top_ks, - stored_top_ps, - stored_min_ps, - lse_slots, - ) = ctx.saved_tensors - if grad_local_ce is None: - grad_local_ce = torch.zeros( - local_hidden.shape[0], - dtype=torch.float32, - device=local_hidden.device, - ) - - group = ctx.tp_group - row_counts = ctx.row_counts - vocab_sizes = ctx.vocab_sizes - rank = dist.get_rank(group) - vocab_offset = sum(vocab_sizes[:rank]) - local_vocab = vocab_sizes[rank] - need_hidden = ctx.needs_input_grad[0] - need_weight = ctx.needs_input_grad[1] - need_hidden_tensor = torch.tensor(int(need_hidden), dtype=torch.int64, device=local_hidden.device) - dist.all_reduce(need_hidden_tensor, op=dist.ReduceOp.MAX, group=group) - compute_hidden = bool(need_hidden_tensor.item()) - grad_hidden = torch.zeros_like(local_hidden) if need_hidden else None - grad_weight = ( - torch.zeros(local_weight.shape, dtype=torch.float32, device=local_weight.device) if need_weight else None - ) - grad_local_ce = grad_local_ce.float().contiguous() - lse_offset = 0 - - for source_rank, source_rows in enumerate(row_counts): - for start in range(0, source_rows, _TP_LOCAL_ROW_CHUNK): - rows_in_chunk = min(_TP_LOCAL_ROW_CHUNK, source_rows - start) - hidden = _tp_broadcast_source_rows( - local_hidden, - source_rank=source_rank, - start=start, - rows=rows_in_chunk, - group=group, - ) - labels = _tp_broadcast_source_rows( - local_labels, - source_rank=source_rank, - start=start, - rows=rows_in_chunk, - group=group, - ) - valid = _tp_broadcast_source_rows( - local_valid, - source_rank=source_rank, - start=start, - rows=rows_in_chunk, - group=group, - ) - grad_ce = _tp_broadcast_source_rows( - grad_local_ce, - source_rank=source_rank, - start=start, - rows=rows_in_chunk, - group=group, - ) - temperature = None - if ctx.has_temperature: - temperature = _tp_broadcast_source_rows( - stored_temperature, - source_rank=source_rank, - start=start, - rows=rows_in_chunk, - group=group, - ) - - local_support = None - selected_support = torch.ones_like(valid) - if ctx.has_sampling_filter: - top_ks = _tp_broadcast_source_rows( - stored_top_ks, - source_rank=source_rank, - start=start, - rows=rows_in_chunk, - group=group, - ) - top_ps = _tp_broadcast_source_rows( - stored_top_ps, - source_rank=source_rank, - start=start, - rows=rows_in_chunk, - group=group, - ) - min_ps = _tp_broadcast_source_rows( - stored_min_ps, - source_rank=source_rank, - start=start, - rows=rows_in_chunk, - group=group, - ) - exact_logits = _tp_exact_full_logits( - hidden, - local_weight, - vocab_sizes=vocab_sizes, - group=group, - use_v2=ctx.use_v2, - vocab_chunk=ctx.vocab_chunk, - ) - transformed_logits = ( - exact_logits - if temperature is None - else exact_temperature_scale_fp32_logits(exact_logits, temperature) - ) - support = exact_sampling_support(transformed_logits, top_ks, top_ps, min_ps) - selected_support = support.gather(1, labels.unsqueeze(1)).squeeze(1) - local_support = support[:, vocab_offset : vocab_offset + local_vocab] - - g = (grad_ce * valid * selected_support).float() - local_logits = torch.mm(hidden, local_weight.t(), out_dtype=torch.float32) - inv_temperature = None - if temperature is not None: - inv_temperature = (1.0 / temperature).unsqueeze(1) - local_logits *= inv_temperature - lse = lse_slots.narrow(0, lse_offset, rows_in_chunk) - lse_offset += rows_in_chunk - grad_logits = local_logits.sub_(lse.unsqueeze(1)).exp_() - grad_logits *= g.unsqueeze(1) - if local_support is not None: - grad_logits *= local_support - target_in_shard = ( - selected_support & valid & (labels >= vocab_offset) & (labels < vocab_offset + local_vocab) - ) - rows = torch.arange(labels.shape[0], device=labels.device) - grad_logits[rows[target_in_shard], labels[target_in_shard] - vocab_offset] -= g[target_in_shard] - if inv_temperature is not None: - grad_logits *= inv_temperature - grad_logits_bf16 = grad_logits.to(local_hidden.dtype) - - if compute_hidden: - grad_hidden_chunk = torch.mm( - grad_logits_bf16, - local_weight, - out_dtype=torch.float32, - ) - dist.all_reduce(grad_hidden_chunk, op=dist.ReduceOp.SUM, group=group) - if need_hidden and rank == source_rank: - grad_hidden.narrow(0, start, rows_in_chunk).copy_(grad_hidden_chunk.to(grad_hidden.dtype)) - if need_weight: - grad_weight.add_(torch.mm(grad_logits_bf16.t(), hidden, out_dtype=torch.float32)) - - return ( - grad_hidden, - grad_weight.to(local_weight.dtype) if need_weight else None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class _BiFusedLmHeadPerTokenCE(torch.autograd.Function): - @staticmethod - def forward( - ctx, - hidden, - weight, - labels_safe, - valid_mask, - temp_row, - top_ks, - top_ps, - min_ps, - vocab_chunk, - ): - use_v2 = families_v2_enabled() - has_sampling_filter = top_ks is not None - if temp_row is None and not has_sampling_filter and use_v2: - # head v2 (families-v2 migration): same GEMM K-chain, epilogue-stats - # online LSE; logits never materialize; backward only consumes lse. - logprob, lse, _ = head_v2_selected_logprob(hidden, weight, labels_safe, temperature=temp_row) - elif temp_row is None and not has_sampling_filter: - logprob, lse, _ = bi_lm_head_selected_logprob( - hidden, weight, labels_safe, temperature=temp_row, vocab_chunk=vocab_chunk - ) - else: - logprob_chunks = [] - lse_chunks = [] - for start in range(0, hidden.shape[0], _TEMPERATURE_MATERIALIZE_ROW_CHUNK): - end = min(start + _TEMPERATURE_MATERIALIZE_ROW_CHUNK, hidden.shape[0]) - hidden_chunk = hidden[start:end] - labels_chunk = labels_safe[start:end] - temperature_chunk = None if temp_row is None else temp_row[start:end] - if use_v2: - logits, _ = head_v2_full_logits_with_lse(hidden_chunk, weight, temperature=None) - transformed_logits = ( - logits - if temperature_chunk is None - else exact_temperature_scale_fp32_logits(logits, temperature_chunk) - ) - if has_sampling_filter: - logprob_chunk, lse_chunk, _ = _score_exact_sampling_rows( - transformed_logits, - labels_chunk, - top_ks[start:end], - top_ps[start:end], - min_ps[start:end], - lambda score_logits, score_ids: head_v2_selected_logprob_from_logits( - score_logits, - score_ids, - temperature=None, - ), - ) - else: - logprob_chunk, lse_chunk, _ = head_v2_selected_logprob_from_logits( - transformed_logits, - labels_chunk, - temperature=None, - ) - else: - logits = bi_lm_head_full_logits(hidden_chunk, weight, vocab_chunk=vocab_chunk) - transformed_logits = ( - logits - if temperature_chunk is None - else exact_temperature_scale_fp32_logits(logits, temperature_chunk) - ) - if has_sampling_filter: - logprob_chunk, lse_chunk, _ = _score_exact_sampling_rows( - transformed_logits, - labels_chunk, - top_ks[start:end], - top_ps[start:end], - min_ps[start:end], - lambda score_logits, score_ids: bi_lm_head_selected_logprob_from_logits( - score_logits, - score_ids, - temperature=None, - vocab_chunk=vocab_chunk, - ), - ) - else: - logprob_chunk, lse_chunk, _ = bi_lm_head_selected_logprob_from_logits( - transformed_logits, - labels_chunk, - temperature=None, - vocab_chunk=vocab_chunk, - ) - logprob_chunks.append(logprob_chunk) - lse_chunks.append(lse_chunk) - if logprob_chunks: - logprob = torch.cat(logprob_chunks, dim=0) - lse = torch.cat(lse_chunks, dim=0) - else: - logprob = torch.empty((0,), dtype=torch.float32, device=hidden.device) - lse = logprob.clone() - ctx.save_for_backward( - hidden, - weight, - labels_safe, - valid_mask, - lse, - temp_row, - top_ks, - top_ps, - min_ps, - ) - ctx.vocab_chunk = vocab_chunk - ctx.use_v2 = use_v2 - return torch.where(valid_mask, -logprob, torch.zeros_like(logprob)) - - @staticmethod - def backward(ctx, grad_ce): - hidden, weight, labels, valid_mask, lse, temp_row, top_ks, top_ps, min_ps = ctx.saved_tensors - vocab_chunk = ctx.vocab_chunk - n_tokens = hidden.shape[0] - vocab = weight.shape[0] - need_h = ctx.needs_input_grad[0] - need_w = ctx.needs_input_grad[1] - - if top_ks is not None: - grad_h = torch.zeros(hidden.shape, dtype=torch.float32, device=hidden.device) if need_h else None - grad_w = torch.zeros_like(weight) if need_w else None - for row_start in range(0, n_tokens, _TEMPERATURE_MATERIALIZE_ROW_CHUNK): - row_end = min(row_start + _TEMPERATURE_MATERIALIZE_ROW_CHUNK, n_tokens) - hidden_chunk = hidden[row_start:row_end] - labels_chunk = labels[row_start:row_end] - valid_chunk = valid_mask[row_start:row_end] - temperature_chunk = None if temp_row is None else temp_row[row_start:row_end] - if ctx.use_v2: - exact_logits, _ = head_v2_full_logits_with_lse(hidden_chunk, weight, temperature=None) - else: - exact_logits = bi_lm_head_full_logits(hidden_chunk, weight, vocab_chunk=vocab_chunk) - transformed_logits = ( - exact_logits - if temperature_chunk is None - else exact_temperature_scale_fp32_logits(exact_logits, temperature_chunk) - ) - support = exact_sampling_support( - transformed_logits, - top_ks[row_start:row_end], - top_ps[row_start:row_end], - min_ps[row_start:row_end], - ) - selected_support = support.gather(1, labels_chunk.unsqueeze(1)).squeeze(1) - g = (grad_ce[row_start:row_end] * valid_chunk * selected_support).float() - g_col = g.unsqueeze(1) - lse_col = lse[row_start:row_end].unsqueeze(1) - inv_t = None if temperature_chunk is None else (1.0 / temperature_chunk).unsqueeze(1) - grad_h_chunk = ( - torch.zeros(hidden_chunk.shape, dtype=torch.float32, device=hidden.device) if need_h else None - ) - rows = torch.arange(row_end - row_start, device=hidden.device) - - for col_start in range(0, vocab, vocab_chunk): - col_end = min(col_start + vocab_chunk, vocab) - w_c = weight[col_start:col_end] - logits_c = torch.mm(hidden_chunk, w_c.t(), out_dtype=torch.float32) - if inv_t is not None: - logits_c *= inv_t - grad_z = logits_c.sub_(lse_col).exp_().mul_(g_col) - grad_z *= support[:, col_start:col_end] - in_chunk = selected_support & (labels_chunk >= col_start) & (labels_chunk < col_end) - grad_z[rows[in_chunk], labels_chunk[in_chunk] - col_start] -= g[in_chunk] - if inv_t is not None: - grad_z *= inv_t - grad_z16 = grad_z.to(hidden.dtype) - if need_h: - torch.addmm( - grad_h_chunk, - grad_z16, - w_c, - out_dtype=torch.float32, - out=grad_h_chunk, - ) - if need_w: - grad_w[col_start:col_end].add_( - torch.mm(grad_z16.t(), hidden_chunk, out_dtype=torch.float32).to(weight.dtype) - ) - if need_h: - grad_h[row_start:row_end] = grad_h_chunk - - return ( - grad_h.to(hidden.dtype) if need_h else None, - grad_w if need_w else None, - None, - None, - None, - None, - None, - None, - None, - ) - - selected_support = torch.ones_like(valid_mask) - g = (grad_ce * valid_mask).float() - g_col = g.unsqueeze(1) - lse_col = lse.unsqueeze(1) - inv_t = None if temp_row is None else (1.0 / temp_row).unsqueeze(1) - grad_h = torch.zeros(hidden.shape, dtype=torch.float32, device=hidden.device) if need_h else None - grad_w = torch.empty_like(weight) if need_w else None - rows = torch.arange(n_tokens, device=hidden.device) - - for col_start in range(0, vocab, vocab_chunk): - col_end = min(col_start + vocab_chunk, vocab) - w_c = weight[col_start:col_end] - # bf16 tensor-core GEMM, fp32 accumulate + fp32 out (no fp32 copies) - logits_c = torch.mm(hidden, w_c.t(), out_dtype=torch.float32) - if inv_t is not None: - logits_c *= inv_t - grad_z = logits_c.sub_(lse_col).exp_().mul_(g_col) - in_chunk = selected_support & (labels >= col_start) & (labels < col_end) - grad_z[rows[in_chunk], labels[in_chunk] - col_start] -= g[in_chunk] - if inv_t is not None: - grad_z *= inv_t # dy/dz = 1/T - grad_z16 = grad_z.to(hidden.dtype) - if need_h: - torch.addmm(grad_h, grad_z16, w_c, out_dtype=torch.float32, out=grad_h) - if need_w: - grad_w[col_start:col_end] = torch.mm(grad_z16.t(), hidden, out_dtype=torch.float32).to(weight.dtype) - - return ( - grad_h.to(hidden.dtype) if need_h else None, - grad_w if need_w else None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def bi_fused_per_token_ce( - hidden_states: torch.Tensor, - weight: torch.Tensor, - labels: torch.Tensor, - ignore_index: int = -100, - temperature: float | torch.Tensor = 1.0, - top_ks: torch.Tensor | None = None, - top_ps: torch.Tensor | None = None, - min_ps: torch.Tensor | None = None, - vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, -) -> torch.Tensor: - """Per-token CE (``-log p(labels)``; 0 at ignored positions) through the - batch-invariant lm-head contract. Requires CUDA bf16 hidden/weight; the - fp32-class numerics come from the contract itself, so ``lm_head_fp32`` is - implied rather than materialized. Per-row temperature scores the same - materialized FP32 ``z * (1/T)`` tensor as serving.""" - if not hidden_states.is_cuda: - raise ValueError("ce_mode='bi_fused' requires CUDA tensors") - if hidden_states.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: - raise ValueError("ce_mode='bi_fused' requires bf16 hidden states and lm-head weight") - valid_mask = labels != ignore_index - labels_safe = torch.where(valid_mask, labels, torch.zeros_like(labels)) - if isinstance(temperature, torch.Tensor): - if temperature.dtype is not torch.float32: - raise TypeError("ce_mode='bi_fused' requires per-row FP32 temperature") - if temperature.device != hidden_states.device or tuple(temperature.shape) != (hidden_states.shape[0],): - raise ValueError("ce_mode='bi_fused' requires temperature aligned one-to-one with hidden-state rows") - if not temperature.is_contiguous() or temperature.requires_grad: - raise ValueError("ce_mode='bi_fused' requires contiguous, non-differentiable temperature metadata") - torch._assert_async( - (torch.isfinite(temperature) & (temperature > 0)).all(), - "ce_mode='bi_fused' requires finite temperature > 0", - ) - temp_row = temperature - elif temperature == 1.0: - temp_row = None - else: - temperature = float(temperature) - if not math.isfinite(temperature) or temperature <= 0: - raise ValueError("ce_mode='bi_fused' requires finite temperature > 0") - temp_row = torch.full((hidden_states.shape[0],), temperature, dtype=torch.float32, device=hidden_states.device) - if (top_ks is None, top_ps is None, min_ps is None).count(True) not in (0, 3): - raise ValueError("ce_mode='bi_fused' requires all or none of top-k/top-p/min-p row metadata") - return _BiFusedLmHeadPerTokenCE.apply( - hidden_states, - weight, - labels_safe, - valid_mask, - temp_row, - top_ks, - top_ps, - min_ps, - vocab_chunk, - ) - - -def bi_fused_vocab_parallel_per_token_ce( - hidden_states: torch.Tensor, - local_weight: torch.Tensor, - labels: torch.Tensor, - tp_group: dist.ProcessGroup, - ignore_index: int = -100, - temperature: float | torch.Tensor = 1.0, - top_ks: torch.Tensor | None = None, - top_ps: torch.Tensor | None = None, - min_ps: torch.Tensor | None = None, - vocab_chunk: int = BI_LM_HEAD_VOCAB_CHUNK, -) -> torch.Tensor: - """Run the batch-invariant LM head over rank-local vocab and token shards. - - LM-head-only TP is composed from DP/CP row owners, so ranks can own - different (or zero) token rows while also owning disjoint vocab rows. The - forward broadcasts bounded chunks from each row owner, computes the exact - local FP32 logits, gathers vocabulary shards in group-rank order, and - scores through the same from-logits tail as serving. Only the source - owner's rows are returned. The custom backward reduces hidden gradients - over vocab ranks and leaves each weight-shard gradient local for the - existing replica sync. - """ - - if not hidden_states.is_cuda or not local_weight.is_cuda: - raise ValueError("ce_mode='bi_fused' tensor parallelism requires CUDA tensors") - if hidden_states.ndim != 2 or local_weight.ndim != 2: - raise ValueError("ce_mode='bi_fused' tensor parallelism requires two-dimensional hidden and weight") - if hidden_states.shape[1] != local_weight.shape[1]: - raise ValueError("ce_mode='bi_fused' tensor parallelism received mismatched hidden dimensions") - if hidden_states.dtype is not torch.bfloat16 or local_weight.dtype is not torch.bfloat16: - raise ValueError("ce_mode='bi_fused' tensor parallelism requires bf16 hidden states and lm-head weight") - if labels.ndim != 1 or labels.shape[0] != hidden_states.shape[0]: - raise ValueError("ce_mode='bi_fused' tensor parallelism requires one label per local hidden row") - hidden_states = hidden_states.contiguous() - local_weight = local_weight.contiguous() - labels = labels.contiguous() - valid_mask = labels != ignore_index - labels_safe = torch.where(valid_mask, labels, torch.zeros_like(labels)) - - if isinstance(temperature, torch.Tensor): - if temperature.dtype is not torch.float32: - raise TypeError("ce_mode='bi_fused' requires per-row FP32 temperature") - if temperature.device != hidden_states.device or tuple(temperature.shape) != (hidden_states.shape[0],): - raise ValueError("ce_mode='bi_fused' requires temperature aligned one-to-one with local hidden rows") - if not temperature.is_contiguous() or temperature.requires_grad: - raise ValueError("ce_mode='bi_fused' requires contiguous, non-differentiable temperature metadata") - torch._assert_async( - (torch.isfinite(temperature) & (temperature > 0)).all(), - "ce_mode='bi_fused' requires finite temperature > 0", - ) - temp_row = temperature - elif temperature == 1.0: - temp_row = None - else: - temperature = float(temperature) - if not math.isfinite(temperature) or temperature <= 0: - raise ValueError("ce_mode='bi_fused' requires finite temperature > 0") - temp_row = torch.full( - (hidden_states.shape[0],), - temperature, - dtype=torch.float32, - device=hidden_states.device, - ) - if (top_ks is None, top_ps is None, min_ps is None).count(True) not in (0, 3): - raise ValueError("ce_mode='bi_fused' requires all or none of top-k/top-p/min-p row metadata") - for name, value in (("top_ks", top_ks), ("top_ps", top_ps), ("min_ps", min_ps)): - if value is not None and ( - value.device != hidden_states.device - or tuple(value.shape) != (hidden_states.shape[0],) - or not value.is_contiguous() - ): - raise ValueError(f"ce_mode='bi_fused' requires contiguous local-row {name} metadata") - - return _BiFusedVocabParallelPerTokenCE.apply( - hidden_states, - local_weight, - labels_safe, - valid_mask, - temp_row, - top_ks, - top_ps, - min_ps, - tp_group, - vocab_chunk, - ) +_sys.modules[__name__] = _importlib.import_module("xorl.ops.loss.batch_invariant_lm_head") diff --git a/src/xorl/ops/loss/per_token_ce.py b/src/xorl/ops/loss/per_token_ce.py index 84d4af4b..a2a1c3f1 100644 --- a/src/xorl/ops/loss/per_token_ce.py +++ b/src/xorl/ops/loss/per_token_ce.py @@ -29,7 +29,7 @@ LogprobProbability = float | torch.Tensor -def resolve_bi_fused_lm_head_tp_groups( +def resolve_batch_invariant_lm_head_tp_groups( ce_mode: str, tp_group: Optional[dist.ProcessGroup], lm_head: Optional[torch.nn.Module], @@ -41,7 +41,7 @@ def resolve_bi_fused_lm_head_tp_groups( being passed as ``tp_group``. """ - if ce_mode != "bi_fused" or lm_head is None: + if ce_mode != "batch_invariant" or lm_head is None: return None if getattr(lm_head, "_glm52_exact_tp16_lm_head", False) or getattr(lm_head, "_dsv4_exact_tp8_lm_head", False): return None @@ -57,7 +57,7 @@ def resolve_bi_fused_lm_head_tp_groups( or tp_group is not dedicated_group ): raise NotImplementedError( - "ce_mode='bi_fused' requires the marked vocabulary-sharded lm_head to use its dedicated LM-head TP group" + "ce_mode='batch_invariant' requires the marked vocabulary-sharded lm_head to use its dedicated LM-head TP group" ) return dedicated_group, getattr(ps, "lm_head_tp_replica_group", None) @@ -288,26 +288,26 @@ def compute_per_token_ce( ignore_index=ignore_index, ) - # ``bi_fused`` is the K3 lm-head contract (vendored identically in SGLang). + # ``batch_invariant`` is the K3 lm-head contract (vendored identically in SGLang). # Hidden states stay bf16. Per-row temperature materializes the same FP32 # transformed logits that serving samples and scores, unlike the # scale-hidden-pre-GEMM convention used by the other modes. - if ce_mode == "bi_fused": - from xorl.ops.loss.bi_fused_lm_head import ( # noqa: PLC0415 - bi_fused_per_token_ce, - bi_fused_vocab_parallel_per_token_ce, + if ce_mode == "batch_invariant": + from xorl.ops.loss.batch_invariant_lm_head import ( # noqa: PLC0415 + batch_invariant_per_token_ce, + batch_invariant_vocab_parallel_per_token_ce, ) - bi_fused_tp_groups = resolve_bi_fused_lm_head_tp_groups(ce_mode, tp_group, lm_head) + batch_invariant_tp_groups = resolve_batch_invariant_lm_head_tp_groups(ce_mode, tp_group, lm_head) if use_lm_head_module: - raise NotImplementedError("ce_mode='bi_fused' does not support FP8 lm_head modules") + raise NotImplementedError("ce_mode='batch_invariant' does not support FP8 lm_head modules") if not lm_head_fp32: raise NotImplementedError( - "ce_mode='bi_fused' implements the fp32-class lm-head contract; set lm_head_fp32: true" + "ce_mode='batch_invariant' implements the fp32-class lm-head contract; set lm_head_fp32: true" ) local_weight = weight.to_local() if hasattr(weight, "to_local") else weight - if bi_fused_tp_groups is not None: - return bi_fused_vocab_parallel_per_token_ce( + if batch_invariant_tp_groups is not None: + return batch_invariant_vocab_parallel_per_token_ce( hidden_states_flat, local_weight, labels_flat, @@ -320,9 +320,9 @@ def compute_per_token_ce( ) if tp_group is not None: raise NotImplementedError( - "ce_mode='bi_fused' supports TP only through the dedicated vocabulary-sharded LM-head TP path" + "ce_mode='batch_invariant' supports TP only through the dedicated vocabulary-sharded LM-head TP path" ) - return bi_fused_per_token_ce( + return batch_invariant_per_token_ce( hidden_states_flat, local_weight, labels_flat, diff --git a/src/xorl/ops/sglang/batch_invariant_ops.py b/src/xorl/ops/sglang/batch_invariant_ops.py index 79f4d56b..d93d0f0c 100644 --- a/src/xorl/ops/sglang/batch_invariant_ops.py +++ b/src/xorl/ops/sglang/batch_invariant_ops.py @@ -23,7 +23,7 @@ import triton.language as tl from triton.runtime.errors import OutOfResources -from xorl.ops.exact.bi_gemm_configs import baseline_mm_config, lookup_mm_config +from xorl.ops.exact.batch_invariant_gemm_configs import baseline_mm_config, lookup_mm_config # --- Stubs for SGLang-internal imports --------------------------------------- diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index d6cbcd8b..a2d1063d 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -2114,7 +2114,7 @@ def _get_loss_lm_head_module(self, lm_head): ): return lm_head if ( - self.ce_mode == "bi_fused" + self.ce_mode == "batch_invariant" and lm_head is not None and getattr(lm_head, "_xorl_fsdp_sharded_lm_head_loss", False) ): diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 649caec0..77ab1cde 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -678,7 +678,7 @@ class ServerArguments: default=None, metadata={ "help": "Cross-entropy implementation. Omitted means compiled for ordinary models and " - "bi_fused for canonical GLM-5.2. 'bi_fused' is the batch-invariant " + "batch_invariant for canonical GLM-5.2. 'batch_invariant' is the batch-invariant " "K3 lm-head contract, fp32-class; needs tp=1, no z-loss, bf16 hidden/weight, lm_head_fp32), " "'compiled' (torch.compile), 'quack_linear' (Quack scalar loss; return_per_token uses fused " "selected-logprob CE), 'fused_quack', or 'eager' (baseline, may OOM at 32K)" @@ -1229,6 +1229,16 @@ def __post_init__(self): """Validate and set defaults.""" from xorl.fp8_training.config_compat import normalize_fp8_training_config # noqa: PLC0415 + if self.ce_mode == "bi_fused": + import warnings + + warnings.warn( + "ce_mode='bi_fused' is deprecated; use ce_mode='batch_invariant'", + DeprecationWarning, + stacklevel=2, + ) + self.ce_mode = "batch_invariant" + if self.rope_class_b is not None: import warnings diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 09fb9ce2..597612cb 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -1325,7 +1325,7 @@ def _get_pp_schedule(self, seq_len: int, example_input_ids: "Optional[torch.Tens or getattr(pp_lm_head, "_dsv4_exact_tp8_lm_head", False) ) ) - lm_head_in_loss = ce_mode in {"quack_linear", "bi_fused"} or exact_head + lm_head_in_loss = ce_mode in {"quack_linear", "batch_invariant"} or exact_head stages = [] for model_part, init_stage in zip(self.model_parts, self.pp_stages): stage_index = init_stage.stage_index diff --git a/src/xorl/trainers/training_utils.py b/src/xorl/trainers/training_utils.py index 1eb9adea..a3058ca6 100644 --- a/src/xorl/trainers/training_utils.py +++ b/src/xorl/trainers/training_utils.py @@ -512,7 +512,7 @@ def make_pp_loss_fn( lm_head is not None and (getattr(lm_head, "_glm52_exact_tp16_lm_head", False) or getattr(lm_head, "_dsv4_exact_tp8_lm_head", False)) ) - if ce_mode == "bi_fused" or exact_head: + if ce_mode == "batch_invariant" or exact_head: # Every rank constructs the schedule, but only the terminal stage calls # the loss. Defer the missing-head error so headless PP stages remain # independent of the output projection. @@ -580,7 +580,9 @@ def _quack_loss(hidden, labels): return _pp_quack_linear_ce_sum(hidden, labels, lm_head=lm_head) return _quack_loss - raise ValueError(f"Unknown ce_mode: {ce_mode!r} (expected 'eager', 'compiled', 'quack_linear', or 'bi_fused')") + raise ValueError( + f"Unknown ce_mode: {ce_mode!r} (expected 'eager', 'compiled', 'quack_linear', or 'batch_invariant')" + ) def pad_micro_batches_for_pp( diff --git a/tests/distributed/test_bi_fused_lm_head_tp.py b/tests/distributed/test_batch_invariant_lm_head_tp.py similarity index 92% rename from tests/distributed/test_bi_fused_lm_head_tp.py rename to tests/distributed/test_batch_invariant_lm_head_tp.py index 7ae0e239..60258cf1 100644 --- a/tests/distributed/test_bi_fused_lm_head_tp.py +++ b/tests/distributed/test_batch_invariant_lm_head_tp.py @@ -2,7 +2,7 @@ Run through pytest (which self-launches torchrun) or directly with two ranks:: - torchrun --nproc_per_node=2 tests/distributed/test_bi_fused_lm_head_tp.py + torchrun --nproc_per_node=2 tests/distributed/test_batch_invariant_lm_head_tp.py """ import os @@ -15,9 +15,9 @@ from xorl.objectives.causallm_loss import causallm_loss_function from xorl.objectives.reducers import TokenPartial from xorl.ops import bi_families_v2 -from xorl.ops.loss.bi_fused_lm_head import ( - bi_fused_per_token_ce, - bi_fused_vocab_parallel_per_token_ce, +from xorl.ops.loss.batch_invariant_lm_head import ( + batch_invariant_per_token_ce, + batch_invariant_vocab_parallel_per_token_ce, ) @@ -69,7 +69,7 @@ def _assert_forward_bytes_and_backward( top_ps = torch.ones(hidden.shape[0], dtype=torch.float32, device="cuda") min_ps = torch.zeros(hidden.shape[0], dtype=torch.float32, device="cuda") - actual = bi_fused_vocab_parallel_per_token_ce( + actual = batch_invariant_vocab_parallel_per_token_ce( hidden, local_weight, labels, @@ -81,7 +81,7 @@ def _assert_forward_bytes_and_backward( ) reference_weight = full_weight.detach().clone().requires_grad_(trainable_weight) reference_hidden = hidden.detach().clone().requires_grad_(True) - reference = bi_fused_per_token_ce( + reference = batch_invariant_per_token_ce( reference_hidden, reference_weight, labels, @@ -99,7 +99,7 @@ def _assert_forward_bytes_and_backward( local_weight.detach(), labels, return_per_token=True, - ce_mode="bi_fused", + ce_mode="batch_invariant", tp_group=dist.group.WORLD, lm_head_fp32=True, lm_head=SimpleNamespace(_xorl_fsdp_sharded_lm_head_loss=True), @@ -160,7 +160,7 @@ def _run_cases(rank: int, world_size: int) -> None: local_weight = full_weight[rank * 256 : (rank + 1) * 256].contiguous() hidden = torch.empty((0, _HIDDEN), dtype=torch.bfloat16, device="cuda", requires_grad=True) labels = torch.empty((0,), dtype=torch.int64, device="cuda") - empty_ce = bi_fused_vocab_parallel_per_token_ce( + empty_ce = batch_invariant_vocab_parallel_per_token_ce( hidden, local_weight, labels, @@ -186,7 +186,7 @@ def main() -> None: _run_cases(rank, world_size) dist.barrier() if rank == 0: - print("bi_fused LM-head TP ragged/empty forward+backward passed") + print("batch_invariant LM-head TP ragged/empty forward+backward passed") finally: parallel_state_impl._PARALLEL_STATE = previous_parallel_state dist.destroy_process_group() @@ -200,7 +200,7 @@ def main() -> None: @pytest.mark.gpu @pytest.mark.distributed @skip_if_gpu_count_less_than(2) - def test_bi_fused_lm_head_tp_2gpu() -> None: + def test_batch_invariant_lm_head_tp_2gpu() -> None: result = run_distributed_script(os.path.abspath(__file__), num_gpus=2, timeout=240) result.assert_success() diff --git a/tests/distributed/test_canonical_moe_contract.py b/tests/distributed/test_canonical_moe_contract.py index bbbcd072..d5e206b2 100644 --- a/tests/distributed/test_canonical_moe_contract.py +++ b/tests/distributed/test_canonical_moe_contract.py @@ -28,11 +28,11 @@ _RuntimePlan, _transport_and_fold, canonical_moe_fold_fp64_v3, - moe_fixed_order_leaf_fp32_v1, canonical_moe_reduce_cp_sharded_v3, canonical_moe_reduce_fp64_v3, canonical_moe_reduce_packed_ep16_v2, canonical_moe_reduce_reference, + moe_fixed_order_leaf_fp32_v1, resolve_canonical_moe_transport, ) from xorl.distributed.parallel_state import init_ep_mesh_matrix diff --git a/tests/distributed/test_pp_byte_alignment.py b/tests/distributed/test_pp_byte_alignment.py index e448d769..b6138c37 100644 --- a/tests/distributed/test_pp_byte_alignment.py +++ b/tests/distributed/test_pp_byte_alignment.py @@ -11,7 +11,7 @@ - ``generic``: native RMSNorm + eager attention + eager CE. Pins the PP plumbing itself (send/recv, pruning, metadata routing, stage-local RoPE). - ``exact``: ``_qwen35_exact_contract`` v1 program — ``sglang_fused`` BI - RMSNorm kernels, FA4 varlen attention, exact-contract GDN, ``bi_fused`` + RMSNorm kernels, FA4 varlen attention, exact-contract GDN, ``batch_invariant`` head with ``lm_head_fp32``. Pins the contract lane, and additionally gates microbatch composition invariance, PP-mandated padding, and fail-closed metadata handling. @@ -122,7 +122,7 @@ def _build_config(mode: str) -> Qwen3_5Config: tie_word_embeddings=False, ) if mode in _EXACT_MODES: - # The exact v1 value program: BI RMSNorm families + FA4 + bi_fused head. + # The exact v1 value program: BI RMSNorm families + FA4 + batch_invariant head. set_rmsnorm_mode("sglang_fused") config._attn_implementation = "flash_attention_4" config._qwen35_exact_contract = True @@ -137,7 +137,7 @@ def _build_config(mode: str) -> Qwen3_5Config: def _ce_kwargs(mode: str) -> dict: if mode in _EXACT_MODES: - return {"ce_mode": "bi_fused", "lm_head_fp32": True} + return {"ce_mode": "batch_invariant", "lm_head_fp32": True} return {"ce_mode": "eager"} diff --git a/tests/distributed/test_pp_byte_contract.py b/tests/distributed/test_pp_byte_contract.py index 419c8e6d..c27c1e5c 100644 --- a/tests/distributed/test_pp_byte_contract.py +++ b/tests/distributed/test_pp_byte_contract.py @@ -447,8 +447,8 @@ def test_generic_part_keeps_silent_fallback(): assert hidden.shape == (1, 8, model.config.hidden_size) -def test_bi_fused_pp_loss_defers_terminal_head_lookup(): - loss_fn = make_pp_loss_fn("bi_fused") +def test_batch_invariant_pp_loss_defers_terminal_head_lookup(): + loss_fn = make_pp_loss_fn("batch_invariant") with pytest.raises(ValueError, match="terminal-stage lm_head"): loss_fn(torch.zeros(1, 2, 4), torch.zeros(1, 2, dtype=torch.long)) diff --git a/tests/distributed/test_ulysses_byte_alignment.py b/tests/distributed/test_ulysses_byte_alignment.py index 7848bddd..5fc831a5 100644 --- a/tests/distributed/test_ulysses_byte_alignment.py +++ b/tests/distributed/test_ulysses_byte_alignment.py @@ -6,7 +6,7 @@ - phase "ref" (1 GPU, no torch.distributed): the exact-contract program on a FULL-ATTENTION-ONLY tiny model with the production head geometry (8 Q-heads / 2 KV-heads / head_dim 256, bf16, FA4, sglang_fused BI RMSNorm, - bi_fused head) over packed varlen inputs; writes last-hidden and per-token + batch_invariant head) over packed varlen inputs; writes last-hidden and per-token logprob bytes to an npz. - phase "shard" (torchrun, ULYSSES_GATE_DEGREE ranks): the same model and tokens through the production Ulysses path (sequence-sharded input_ids, @@ -152,7 +152,7 @@ def _logprobs(model, hidden, labels): weight=model.lm_head.weight, labels=labels, return_per_token=True, - ce_mode="bi_fused", + ce_mode="batch_invariant", lm_head_fp32=True, ) return result.per_token_logprobs diff --git a/tests/models/test_bi_fused_lm_head_tp_wiring.py b/tests/models/test_batch_invariant_lm_head_tp_wiring.py similarity index 89% rename from tests/models/test_bi_fused_lm_head_tp_wiring.py rename to tests/models/test_batch_invariant_lm_head_tp_wiring.py index 4e618422..72ec4e08 100644 --- a/tests/models/test_bi_fused_lm_head_tp_wiring.py +++ b/tests/models/test_batch_invariant_lm_head_tp_wiring.py @@ -12,7 +12,7 @@ from xorl.objectives.reducers import TokenPartial -def test_causallm_routes_bi_fused_tp_before_ordinary_vocab_ce(monkeypatch): +def test_causallm_routes_batch_invariant_tp_before_ordinary_vocab_ce(monkeypatch): tp_group = object() replica_group = object() ps = SimpleNamespace( @@ -54,7 +54,7 @@ def fake_all_reduce(value, *, group, **_kwargs): hidden, weight, labels, - ce_mode="bi_fused", + ce_mode="batch_invariant", tp_group=tp_group, lm_head_fp32=True, lm_head=lm_head, @@ -63,14 +63,14 @@ def fake_all_reduce(value, *, group, **_kwargs): assert captured["tp_group"] is tp_group assert captured["lm_head"] is lm_head - assert captured["ce_mode"] == "bi_fused" + assert captured["ce_mode"] == "batch_invariant" assert captured["hidden"].shape == (3, 4) assert reduced_groups == [tp_group, replica_group, tp_group, replica_group] torch.testing.assert_close(result.loss, torch.tensor(1.0)) assert torch.equal(result.per_token_loss, torch.tensor([[0.0, 1.0, 2.0]])) -def test_causallm_explicit_reducer_returns_local_bi_fused_tp_partial(monkeypatch): +def test_causallm_explicit_reducer_returns_local_batch_invariant_tp_partial(monkeypatch): tp_group = object() replica_group = object() ps = SimpleNamespace( @@ -98,7 +98,7 @@ def test_causallm_explicit_reducer_returns_local_bi_fused_tp_partial(monkeypatch torch.randn(1, 3, 4, dtype=torch.bfloat16), lm_head.weight, torch.tensor([[1, -100, 3]]), - ce_mode="bi_fused", + ce_mode="batch_invariant", tp_group=tp_group, lm_head_fp32=True, lm_head=lm_head, @@ -119,7 +119,7 @@ def test_causallm_explicit_reducer_returns_local_bi_fused_tp_partial(monkeypatch (True, 2, True, False, True), ], ) -def test_causallm_rejects_malformed_bi_fused_tp_topology( +def test_causallm_rejects_malformed_batch_invariant_tp_topology( monkeypatch, marked, lm_head_tp_size, @@ -147,14 +147,14 @@ def test_causallm_rejects_malformed_bi_fused_tp_topology( torch.randn(1, 3, 4, dtype=torch.bfloat16), torch.randn(6, 4, dtype=torch.bfloat16), torch.tensor([[1, -100, 3]]), - ce_mode="bi_fused", + ce_mode="batch_invariant", tp_group=tp_group, lm_head_fp32=True, lm_head=lm_head, ) -def test_compute_loss_keeps_sharded_bi_fused_tp_and_global_token_scale(monkeypatch): +def test_compute_loss_keeps_sharded_batch_invariant_tp_and_global_token_scale(monkeypatch): tp_group = object() replica_group = object() ps = SimpleNamespace( @@ -190,7 +190,7 @@ def fake_loss(**kwargs): loss_fn_name="causallm_loss", loss_fn_inputs={"labels": labels}, loss_fn_params={ - "ce_mode": "bi_fused", + "ce_mode": "batch_invariant", "lm_head_fp32": True, "fsdp_sharded_lm_head_loss_num_chunks": 4, "fsdp_sharded_lm_head_loss_global_valid_tokens": torch.tensor(8), @@ -201,15 +201,15 @@ def fake_loss(**kwargs): kwargs = captured["loss_kwargs"] assert kwargs["tp_group"] is tp_group assert kwargs["lm_head"] is lm_head - assert "bi_fused_vocab_parallel" not in kwargs - assert "bi_fused_loss_reduce_group" not in kwargs + assert "batch_invariant_vocab_parallel" not in kwargs + assert "batch_invariant_loss_reduce_group" not in kwargs assert "fsdp_sharded_lm_head_loss_num_chunks" not in kwargs assert "fsdp_sharded_lm_head_loss_global_valid_tokens" not in kwargs torch.testing.assert_close(kwargs["loss_reducer"].scale, torch.tensor(8.0)) torch.testing.assert_close(result.loss, torch.tensor(0.25)) -def test_compute_loss_rejects_unsharded_bi_fused_lm_head_tp(monkeypatch): +def test_compute_loss_rejects_unsharded_batch_invariant_lm_head_tp(monkeypatch): ps = SimpleNamespace( lm_head_tp_size=2, lm_head_tp_group=object(), @@ -224,11 +224,11 @@ def test_compute_loss_rejects_unsharded_bi_fused_lm_head_tp(monkeypatch): torch.randn(1, 3, 4, dtype=torch.bfloat16), loss_fn_name="causallm_loss", loss_fn_inputs={"labels": torch.tensor([[1, -100, 3]])}, - loss_fn_params={"ce_mode": "bi_fused", "lm_head_fp32": True}, + loss_fn_params={"ce_mode": "batch_invariant", "lm_head_fp32": True}, ) -def test_pp_routes_only_sharded_dedicated_bi_fused_lm_head_tp(monkeypatch): +def test_pp_routes_only_sharded_dedicated_batch_invariant_lm_head_tp(monkeypatch): lm_head = nn.Linear(4, 6, bias=False, dtype=torch.bfloat16) hidden = torch.randn(1, 3, 4, dtype=torch.bfloat16) labels = torch.tensor([[1, -100, 3]]) @@ -239,7 +239,7 @@ def test_pp_routes_only_sharded_dedicated_bi_fused_lm_head_tp(monkeypatch): hidden, labels, lm_head=lm_head, - ce_mode="bi_fused", + ce_mode="batch_invariant", tp_group=tp_group, lm_head_fp32=True, ) @@ -261,7 +261,7 @@ def fake_per_token_ce(*_args, **kwargs): hidden, labels, lm_head=lm_head, - ce_mode="bi_fused", + ce_mode="batch_invariant", tp_group=tp_group, lm_head_fp32=True, ) @@ -269,5 +269,5 @@ def fake_per_token_ce(*_args, **kwargs): assert captured["fsdp_sharded_loss"] is True assert captured["tp_group"] is tp_group assert captured["lm_head"] is lm_head - assert "bi_fused_vocab_parallel" not in captured + assert "batch_invariant_vocab_parallel" not in captured torch.testing.assert_close(loss_sum, torch.tensor(3.0)) diff --git a/tests/models/test_dsv4_exact_contract.py b/tests/models/test_dsv4_exact_contract.py index 735e9f48..a0252104 100644 --- a/tests/models/test_dsv4_exact_contract.py +++ b/tests/models/test_dsv4_exact_contract.py @@ -156,14 +156,14 @@ def test_adapter_program_rejects_partial_targets_and_non_rank_one() -> None: ) -def test_exact_loss_mode_rejects_non_tp_aware_bi_fused_path() -> None: +def test_exact_loss_mode_rejects_non_tp_aware_batch_invariant_path() -> None: config = _official_config() config._dsv4_flash_exact_mode = True assert resolve_cross_entropy_mode(config, None) == "compiled" assert resolve_cross_entropy_mode(config, "compiled") == "compiled" with pytest.raises(ValueError, match="requires ce_mode='compiled'"): - resolve_cross_entropy_mode(config, "bi_fused") + resolve_cross_entropy_mode(config, "batch_invariant") def test_inventory_derives_exact_345_non_routed_43_banks_and_948_factors() -> None: diff --git a/tests/models/test_glm52_exact_lm_head_loss_integration.py b/tests/models/test_glm52_exact_lm_head_loss_integration.py index 58be2420..b7f894f2 100644 --- a/tests/models/test_glm52_exact_lm_head_loss_integration.py +++ b/tests/models/test_glm52_exact_lm_head_loss_integration.py @@ -45,7 +45,7 @@ def _fake_exact(hidden, weight, labels, **kwargs): weight, labels, -100, - "bi_fused", + "batch_invariant", tp_group=tp_group, lm_head_fp32=True, lm_head=lm_head, @@ -55,7 +55,7 @@ def _fake_exact(hidden, weight, labels, **kwargs): assert captures["lm_head"] is lm_head assert captures["tp_group"] is tp_group assert captures["ignore_index"] == -100 - assert captures["ce_mode"] == "bi_fused" + assert captures["ce_mode"] == "batch_invariant" assert captures["lm_head_fp32"] is True assert captures["logprob_temperature"] == 1.0 assert captures["hidden"] is hidden @@ -80,7 +80,7 @@ def _fake_exact(_hidden, _weight, _labels, **kwargs): torch.zeros((6, 4), dtype=torch.bfloat16), torch.tensor([1, 2], dtype=torch.int64), -100, - "bi_fused", + "batch_invariant", tp_group=object(), lm_head_fp32=True, lm_head=lm_head, @@ -108,7 +108,7 @@ def _fake_exact(_hidden, _weight, _labels, **kwargs): torch.zeros((6, 4), dtype=torch.bfloat16), torch.tensor([1, 2], dtype=torch.int64), -100, - "bi_fused", + "batch_invariant", tp_group=object(), lm_head_fp32=True, lm_head=lm_head, @@ -141,7 +141,7 @@ def test_causallm_exact_head_admits_its_tp_group_and_rejects_z_loss(monkeypatch: hidden, weight, labels, - ce_mode="bi_fused", + ce_mode="batch_invariant", tp_group=object(), lm_head_fp32=True, lm_head=lm_head, @@ -156,7 +156,7 @@ def test_causallm_exact_head_admits_its_tp_group_and_rejects_z_loss(monkeypatch: hidden.detach(), weight, labels, - ce_mode="bi_fused", + ce_mode="batch_invariant", tp_group=object(), lm_head_fp32=True, lm_head=lm_head, @@ -209,13 +209,13 @@ def _differentiable_exact(hidden, weight, labels, *, lm_head, ignore_index, **_k get_lm_head_weight(direct_head), labels.reshape(-1), ignore_index=-100, - ce_mode="bi_fused", + ce_mode="batch_invariant", tp_group=tp_group, lm_head_fp32=True, lm_head=direct_head, ).sum() pp = make_pp_loss_fn( - "bi_fused", + "batch_invariant", lm_head=pp_head, tp_group=tp_group, lm_head_fp32=True, @@ -252,7 +252,7 @@ def _temperature_aware_exact(hidden, weight, labels, *, logprob_temperature, ign owner = SimpleNamespace(_pp_loss_temperatures=deque([temperatures])) actual = make_pp_loss_fn( - "bi_fused", + "batch_invariant", lm_head=lm_head, tp_group=object(), lm_head_fp32=True, @@ -297,7 +297,7 @@ def _capture_ce(hidden, weight, labels, *, logprob_temperature, **_kwargs): first_temperature = torch.tensor([[0.5, 1.0]], dtype=torch.float32) second_temperature = torch.tensor([[1.25, 2.0]], dtype=torch.float32) owner = SimpleNamespace(_pp_loss_temperatures=deque([first_temperature, second_temperature])) - loss_fn = make_pp_loss_fn("bi_fused", lm_head=lm_head, loss_owner=owner) + loss_fn = make_pp_loss_fn("batch_invariant", lm_head=lm_head, loss_owner=owner) labels = torch.tensor([[1, 2]], dtype=torch.int64) loss_fn(torch.randn(1, 2, 4, dtype=torch.bfloat16), labels) @@ -319,7 +319,7 @@ def _capture_ce(hidden, _weight, _labels, *, logprob_temperature, **_kwargs): return hidden.float().sum(dim=-1) * 0 monkeypatch.setattr(per_token_ce_impl, "compute_per_token_ce", _capture_ce) - cached_loss_fn = make_pp_loss_fn("bi_fused", lm_head=lm_head, loss_owner=owner) + cached_loss_fn = make_pp_loss_fn("batch_invariant", lm_head=lm_head, loss_owner=owner) hidden = torch.randn(1, 2, 4, dtype=torch.bfloat16) labels = torch.tensor([[1, 2]], dtype=torch.int64) diff --git a/tests/ops/test_bi_gemm_config_table.py b/tests/ops/test_batch_invariant_gemm_config_table.py similarity index 97% rename from tests/ops/test_bi_gemm_config_table.py rename to tests/ops/test_batch_invariant_gemm_config_table.py index cb8f8a45..225a3532 100644 --- a/tests/ops/test_bi_gemm_config_table.py +++ b/tests/ops/test_batch_invariant_gemm_config_table.py @@ -11,7 +11,7 @@ import torch import triton -from xorl.ops.exact.bi_gemm_configs import BASELINE_CONFIG, PINNED_BLOCK_K, lookup_mm_config +from xorl.ops.exact.batch_invariant_gemm_configs import BASELINE_CONFIG, PINNED_BLOCK_K, lookup_mm_config from xorl.ops.sglang.batch_invariant_ops import ( _deepgemm_ready, _matmul_persistent_deepgemm, diff --git a/tests/ops/test_bi_fused_lm_head.py b/tests/ops/test_batch_invariant_lm_head.py similarity index 84% rename from tests/ops/test_bi_fused_lm_head.py rename to tests/ops/test_batch_invariant_lm_head.py index 6224a703..33a658e4 100644 --- a/tests/ops/test_bi_fused_lm_head.py +++ b/tests/ops/test_batch_invariant_lm_head.py @@ -2,7 +2,7 @@ import torch from xorl.objectives.causallm_loss import causallm_loss_function -from xorl.ops.loss.bi_fused_lm_head import bi_fused_per_token_ce +from xorl.ops.loss.batch_invariant_lm_head import batch_invariant_per_token_ce requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -22,9 +22,11 @@ def _inputs(seed=0): @requires_cuda @pytest.mark.gpu -def test_bi_fused_matches_eager_fp32_reference(): +def test_batch_invariant_matches_eager_fp32_reference(): hidden, weight, labels = _inputs() - out = causallm_loss_function(hidden, weight, labels, ce_mode="bi_fused", lm_head_fp32=True, return_per_token=True) + out = causallm_loss_function( + hidden, weight, labels, ce_mode="batch_invariant", lm_head_fp32=True, return_per_token=True + ) ref = causallm_loss_function(hidden, weight, labels, ce_mode="eager", lm_head_fp32=True, return_per_token=True) assert torch.allclose(out.per_token_logprobs, ref.per_token_logprobs, rtol=1e-4, atol=1e-5) assert torch.allclose(out.loss.float(), ref.loss.float(), rtol=1e-4, atol=1e-5) @@ -34,12 +36,12 @@ def test_bi_fused_matches_eager_fp32_reference(): @requires_cuda @pytest.mark.gpu -def test_bi_fused_backward_matches_eager_autograd(): +def test_batch_invariant_backward_matches_eager_autograd(): hidden, weight, labels = _inputs(1) h1 = hidden.clone().requires_grad_(True) w1 = weight.clone().requires_grad_(True) - causallm_loss_function(h1, w1, labels, ce_mode="bi_fused", lm_head_fp32=True).loss.backward() + causallm_loss_function(h1, w1, labels, ce_mode="batch_invariant", lm_head_fp32=True).loss.backward() h2 = hidden.clone().requires_grad_(True) w2 = weight.clone().requires_grad_(True) @@ -51,9 +53,9 @@ def test_bi_fused_backward_matches_eager_autograd(): @requires_cuda @pytest.mark.gpu -def test_bi_fused_deterministic_and_batch_invariant(): +def test_batch_invariant_deterministic_and_batch_invariant(): hidden, weight, labels = _inputs(2) - kw = dict(ce_mode="bi_fused", lm_head_fp32=True, return_per_token=True) + kw = dict(ce_mode="batch_invariant", lm_head_fp32=True, return_per_token=True) a = causallm_loss_function(hidden, weight, labels, **kw).per_token_logprobs b = causallm_loss_function(hidden, weight, labels, **kw).per_token_logprobs assert torch.equal(a, b) @@ -63,20 +65,20 @@ def test_bi_fused_deterministic_and_batch_invariant(): @requires_cuda @pytest.mark.gpu -def test_bi_fused_guards(): +def test_batch_invariant_guards(): hidden, weight, labels = _inputs(3) with pytest.raises(NotImplementedError, match="lm_head_fp32"): - causallm_loss_function(hidden, weight, labels, ce_mode="bi_fused", lm_head_fp32=False) + causallm_loss_function(hidden, weight, labels, ce_mode="batch_invariant", lm_head_fp32=False) with pytest.raises(NotImplementedError, match="softmax_auxiliary_loss"): - causallm_loss_function(hidden, weight, labels, ce_mode="bi_fused", lm_head_fp32=True, z_loss_coef=0.1) + causallm_loss_function(hidden, weight, labels, ce_mode="batch_invariant", lm_head_fp32=True, z_loss_coef=0.1) @requires_cuda @pytest.mark.gpu -def test_bi_fused_temperature_matches_eager_reference(): +def test_batch_invariant_temperature_matches_eager_reference(): hidden, weight, labels = _inputs(4) kw = dict(lm_head_fp32=True, logprob_temperature=0.7, return_per_token=True) - out = causallm_loss_function(hidden, weight, labels, ce_mode="bi_fused", **kw) + out = causallm_loss_function(hidden, weight, labels, ce_mode="batch_invariant", **kw) ref = causallm_loss_function(hidden, weight, labels, ce_mode="eager", **kw) assert torch.allclose(out.per_token_logprobs, ref.per_token_logprobs, rtol=1e-4, atol=1e-5) assert (out.per_token_loss.view(-1)[:7] == 0).all() @@ -84,13 +86,13 @@ def test_bi_fused_temperature_matches_eager_reference(): @requires_cuda @pytest.mark.gpu -def test_bi_fused_temperature_backward_matches_eager_autograd(): +def test_batch_invariant_temperature_backward_matches_eager_autograd(): hidden, weight, labels = _inputs(5) h1 = hidden.clone().requires_grad_(True) w1 = weight.clone().requires_grad_(True) causallm_loss_function( - h1, w1, labels, ce_mode="bi_fused", lm_head_fp32=True, logprob_temperature=0.7 + h1, w1, labels, ce_mode="batch_invariant", lm_head_fp32=True, logprob_temperature=0.7 ).loss.backward() h2 = hidden.clone().requires_grad_(True) @@ -103,7 +105,7 @@ def test_bi_fused_temperature_backward_matches_eager_autograd(): @requires_cuda @pytest.mark.gpu -def test_bi_fused_mixed_row_temperature_matches_reference_and_backward(): +def test_batch_invariant_mixed_row_temperature_matches_reference_and_backward(): hidden, weight, labels = _inputs(51) temperature = torch.linspace(0.7, 1.3, N, dtype=torch.float32, device="cuda").reshape(1, N) @@ -113,7 +115,7 @@ def test_bi_fused_mixed_row_temperature_matches_reference_and_backward(): exact_hidden, exact_weight, labels, - ce_mode="bi_fused", + ce_mode="batch_invariant", lm_head_fp32=True, logprob_temperature=temperature, return_per_token=True, @@ -139,13 +141,13 @@ def test_bi_fused_mixed_row_temperature_matches_reference_and_backward(): @requires_cuda @pytest.mark.gpu -def test_bi_fused_per_row_unit_temperature_preserves_forward_bytes(): +def test_batch_invariant_per_row_unit_temperature_preserves_forward_bytes(): hidden, weight, labels = _inputs(52) scalar = causallm_loss_function( hidden, weight, labels, - ce_mode="bi_fused", + ce_mode="batch_invariant", lm_head_fp32=True, return_per_token=True, ) @@ -153,7 +155,7 @@ def test_bi_fused_per_row_unit_temperature_preserves_forward_bytes(): hidden, weight, labels, - ce_mode="bi_fused", + ce_mode="batch_invariant", lm_head_fp32=True, logprob_temperature=torch.ones((1, N), dtype=torch.float32, device="cuda"), return_per_token=True, @@ -164,7 +166,7 @@ def test_bi_fused_per_row_unit_temperature_preserves_forward_bytes(): @requires_cuda @pytest.mark.gpu @pytest.mark.parametrize("family", ["v1", "v2"]) -def test_bi_fused_temperature_matches_serving_materialize_then_score(family): +def test_batch_invariant_temperature_matches_serving_materialize_then_score(family): pytest.importorskip("sglang") from sglang.srt.batch_invariant_ops import ( bi_lm_head_selected_logprob_from_logits as serving_v1_score, @@ -195,13 +197,13 @@ def test_bi_fused_temperature_matches_serving_materialize_then_score(family): logits, _ = head_v2_full_logits_with_lse(hidden, weight) score = serving_v2_score - actual = bi_fused_per_token_ce(hidden, weight, labels, temperature=temperature) + actual = batch_invariant_per_token_ce(hidden, weight, labels, temperature=temperature) transformed = serving_scale(logits, temperature) expected_logprob, _, _ = score(transformed, labels, temperature=None) assert torch.equal(actual.view(torch.uint8), (-expected_logprob).view(torch.uint8)) - scalar_unit = bi_fused_per_token_ce(hidden, weight, labels, temperature=1.0) - row_unit = bi_fused_per_token_ce(hidden, weight, labels, temperature=torch.ones_like(temperature)) + scalar_unit = batch_invariant_per_token_ce(hidden, weight, labels, temperature=1.0) + row_unit = batch_invariant_per_token_ce(hidden, weight, labels, temperature=torch.ones_like(temperature)) assert torch.equal(scalar_unit.view(torch.uint8), row_unit.view(torch.uint8)) finally: bi_families_v2._select_nonexact_families() diff --git a/tests/ops/test_exact_sampling_transforms.py b/tests/ops/test_exact_sampling_transforms.py index 2d5833ba..eb007517 100644 --- a/tests/ops/test_exact_sampling_transforms.py +++ b/tests/ops/test_exact_sampling_transforms.py @@ -15,7 +15,7 @@ exact_support_workspace_bytes, normalize_exact_sampling_transforms, ) -from xorl.ops.loss.bi_fused_lm_head import _score_exact_sampling_rows +from xorl.ops.loss.batch_invariant_lm_head import _score_exact_sampling_rows def _rows(values): @@ -196,7 +196,7 @@ def test_chunked_selected_score_and_vjp_equal_direct_program(): def test_filtered_exact_heads_do_not_save_dense_support_on_autograd_contexts(): modules = [ - importlib.import_module("xorl.ops.loss.bi_fused_lm_head"), + importlib.import_module("xorl.ops.loss.batch_invariant_lm_head"), importlib.import_module("xorl.models.transformers.glm5.exact_lm_head_qlora"), importlib.import_module("xorl.models.transformers.deepseek_v4.exact_lm_head"), ] diff --git a/tests/ops/test_fused_silu_and_mul_scoping.py b/tests/ops/test_fused_silu_and_mul_scoping.py index f25d561a..1af7d75f 100644 --- a/tests/ops/test_fused_silu_and_mul_scoping.py +++ b/tests/ops/test_fused_silu_and_mul_scoping.py @@ -20,7 +20,7 @@ import torch.nn.functional as F from xorl.models.transformers.qwen3_5.modeling_qwen3_5 import Qwen3_5MLP -from xorl.ops.exact.fused_silu_and_mul import one_round_swiglu, fused_silu_and_mul +from xorl.ops.exact.fused_silu_and_mul import fused_silu_and_mul, one_round_swiglu def _two_round_reference(x: torch.Tensor) -> torch.Tensor: diff --git a/tests/server/runner/test_model_runner_drgrpo.py b/tests/server/runner/test_model_runner_drgrpo.py index 389b8f54..0d83b46d 100644 --- a/tests/server/runner/test_model_runner_drgrpo.py +++ b/tests/server/runner/test_model_runner_drgrpo.py @@ -5,7 +5,7 @@ import torch.nn.functional as F import xorl.objectives.causallm_loss as causallm_loss_impl -import xorl.ops.loss.bi_fused_lm_head as bi_fused_lm_head_impl +import xorl.ops.loss.batch_invariant_lm_head as batch_invariant_lm_head_impl import xorl.ops.loss.per_token_ce as per_token_ce_impl import xorl.server.runner.model_runner as model_runner_module from xorl.server.runner.model_runner import ModelRunner @@ -53,7 +53,7 @@ def forward(self, input_ids, **kwargs): return SimpleNamespace(last_hidden_state=self.embed(input_ids)) -def _make_bi_fused_lm_head_tp_runner(monkeypatch): +def _make_batch_invariant_lm_head_tp_runner(monkeypatch): tp_group = object() ps = SimpleNamespace( lm_head_tp_size=2, @@ -85,19 +85,19 @@ def fake_vocab_parallel_ce( return F.cross_entropy(logits, labels, reduction="none", ignore_index=ignore_index) monkeypatch.setattr( - bi_fused_lm_head_impl, - "bi_fused_vocab_parallel_per_token_ce", + batch_invariant_lm_head_impl, + "batch_invariant_vocab_parallel_per_token_ce", fake_vocab_parallel_ce, ) runner = object.__new__(ModelRunner) runner.model = _TinyMarkedBiFusedModel() - runner.ce_mode = "bi_fused" + runner.ce_mode = "batch_invariant" runner.lm_head_fp32 = True return runner, tp_group, routed_groups -def test_compute_micro_batch_loss_routes_marked_bi_fused_causallm(monkeypatch): - runner, tp_group, routed_groups = _make_bi_fused_lm_head_tp_runner(monkeypatch) +def test_compute_micro_batch_loss_routes_marked_batch_invariant_causallm(monkeypatch): + runner, tp_group, routed_groups = _make_batch_invariant_lm_head_tp_runner(monkeypatch) loss, per_token_outputs, _metrics, _metric_ops, _outputs = runner._compute_micro_batch_loss( { @@ -114,18 +114,18 @@ def test_compute_micro_batch_loss_routes_marked_bi_fused_causallm(monkeypatch): assert runner.model.lm_head.weight.grad is not None -def test_marked_ordinary_head_is_exposed_only_for_bi_fused_loss_metadata(): +def test_marked_ordinary_head_is_exposed_only_for_batch_invariant_loss_metadata(): runner = object.__new__(ModelRunner) runner.model = _TinyMarkedBiFusedModel() runner.ce_mode = "eager" assert runner._get_loss_lm_head_module(runner.model.lm_head) is None - runner.ce_mode = "bi_fused" + runner.ce_mode = "batch_invariant" assert runner._get_loss_lm_head_module(runner.model.lm_head) is runner.model.lm_head -def test_compute_micro_batch_loss_routes_marked_bi_fused_drgrpo_backward(monkeypatch): - runner, tp_group, routed_groups = _make_bi_fused_lm_head_tp_runner(monkeypatch) +def test_compute_micro_batch_loss_routes_marked_batch_invariant_drgrpo_backward(monkeypatch): + runner, tp_group, routed_groups = _make_batch_invariant_lm_head_tp_runner(monkeypatch) loss, per_token_outputs, metrics, _metric_ops, _outputs = runner._compute_micro_batch_loss( { diff --git a/tests/test_arguments.py b/tests/test_arguments.py index 14ab9f55..e3fa09db 100644 --- a/tests/test_arguments.py +++ b/tests/test_arguments.py @@ -829,3 +829,27 @@ def _assert_parse_args_load_optimizer_flag(tmp_path, monkeypatch): args = parse_args(Arguments) assert args.train.load_optimizer is expected + + +def test_deprecated_rope_class_b_alias_maps_to_new_field(): + import warnings + + from xorl.arguments import ModelArguments + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + args = ModelArguments(config_path="dummy", rope_class_b=True) + assert args.rope_fp32_single_round is True + assert any("rope_class_b is deprecated" in str(w.message) for w in caught) + + +def test_deprecated_bi_fused_ce_mode_normalizes_to_batch_invariant(): + import warnings + + from xorl.arguments import TrainingArguments + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + args = TrainingArguments(output_dir="/tmp/xorl-test", ce_mode="bi_fused") + assert args.ce_mode == "batch_invariant" + assert any("ce_mode='bi_fused' is deprecated" in str(w.message) for w in caught) diff --git a/tests/trainers/test_qwen3_dense_exact_config.py b/tests/trainers/test_qwen3_dense_exact_config.py index 41e0b73b..e8bfc82a 100644 --- a/tests/trainers/test_qwen3_dense_exact_config.py +++ b/tests/trainers/test_qwen3_dense_exact_config.py @@ -69,7 +69,7 @@ def test_dense_qwen3_resolves_shared_exact_program(): assert not program.activation_native assert program.rope_native assert program.rope_fp32_single_round - assert resolve_cross_entropy_mode(config, None) == "bi_fused" + assert resolve_cross_entropy_mode(config, None) == "batch_invariant" assert _resolve_rope_modes( config, rope_native=None, diff --git a/tests/trainers/test_rope_fp32_single_round_config.py b/tests/trainers/test_rope_fp32_single_round_config.py index d51c5059..6cae83bf 100644 --- a/tests/trainers/test_rope_fp32_single_round_config.py +++ b/tests/trainers/test_rope_fp32_single_round_config.py @@ -155,7 +155,7 @@ def test_canonical_glm_resolves_complete_exact_program(monkeypatch): sparse_mla_enabled=True, sparse_mla_backend="flashmla", ) - assert resolve_cross_entropy_mode(config, None) == "bi_fused" + assert resolve_cross_entropy_mode(config, None) == "batch_invariant" @pytest.mark.parametrize( @@ -193,7 +193,7 @@ def test_canonical_glm_rejects_incompatible_numerical_override(override, value): def test_canonical_glm_rejects_incompatible_ce_override(): config = _exact_glm52_config() - with pytest.raises(ValueError, match="requires ce_mode='bi_fused'"): + with pytest.raises(ValueError, match="requires ce_mode='batch_invariant'"): resolve_cross_entropy_mode(config, "compiled") @@ -262,7 +262,7 @@ def test_exact_qwen35_resolves_the_certified_numerical_program(config_factory): sparse_mla_enabled=False, sparse_mla_backend="auto", ) - assert resolve_cross_entropy_mode(config, None) == "bi_fused" + assert resolve_cross_entropy_mode(config, None) == "batch_invariant" @pytest.mark.parametrize("config_factory", [_exact_qwen35_dense_config, _exact_qwen35_moe_config]) @@ -372,7 +372,7 @@ def test_exact_qwen35_rejects_incompatible_numerical_override(override, value): def test_exact_qwen35_rejects_incompatible_ce_override(): config = _exact_qwen35_moe_config() config._qwen35_exact_contract = True - with pytest.raises(ValueError, match="requires ce_mode='bi_fused'"): + with pytest.raises(ValueError, match="requires ce_mode='batch_invariant'"): resolve_cross_entropy_mode(config, "compiled") From c28f8cb62ce9cf5d95ab2a887683ab80390ea901 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 16:44:29 +0000 Subject: [PATCH 13/14] Fix the ce_mode alias regression test construction (init_device=meta) --- tests/test_arguments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_arguments.py b/tests/test_arguments.py index e3fa09db..f747bc42 100644 --- a/tests/test_arguments.py +++ b/tests/test_arguments.py @@ -850,6 +850,6 @@ def test_deprecated_bi_fused_ce_mode_normalizes_to_batch_invariant(): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") - args = TrainingArguments(output_dir="/tmp/xorl-test", ce_mode="bi_fused") + args = TrainingArguments(output_dir="/tmp/xorl-test", init_device="meta", ce_mode="bi_fused") assert args.ce_mode == "batch_invariant" assert any("ce_mode='bi_fused' is deprecated" in str(w.message) for w in caught) From 5b567ada7142e77528a5ee29164241856f2176f9 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Fri, 21 Aug 2026 17:57:58 +0000 Subject: [PATCH 14/14] Harden _rmtree_owned against a relocated ownership sentinel (#81 security-bot finding) The sentinel alone is not sufficient authorization: whoever controls XORL_EXACT_KERNEL_CONFIG_DIR could pre-place it anywhere. Deletion now also requires the resolved target to live strictly inside the (resolved) pin directory and refuses symlinked targets, bounding any rmtree to the pin tree. Pre-existing behavior surfaced by the phase-3 file move. --- src/xorl/ops/exact/kernel_config_pin.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/src/xorl/ops/exact/kernel_config_pin.py b/src/xorl/ops/exact/kernel_config_pin.py index 4734a935..d8987633 100644 --- a/src/xorl/ops/exact/kernel_config_pin.py +++ b/src/xorl/ops/exact/kernel_config_pin.py @@ -51,10 +51,25 @@ def _mark_owned(path: str) -> None: f.write("created by xorl.ops.exact.kernel_config_pin; safe for it to replace\n") -def _rmtree_owned(path: str) -> None: - """Delete `path` only if this module created it (sentinel present).""" +def _rmtree_owned(path: str, *, containing_dir: str) -> None: + """Delete ``path`` only if this module created it and it stays inside + ``containing_dir`` after symlink resolution. + + The sentinel alone is not sufficient authorization: whoever controls the + pin-directory env var could pre-place it anywhere. Requiring the resolved + target to live under the (also resolved) pin directory, and refusing + symlinked targets, bounds any deletion to the pin tree itself. + """ + if os.path.islink(path): + raise KernelConfigPinError(f"refusing to delete {path!r}: it is a symlink") if not os.path.isdir(path): return + resolved = os.path.realpath(path) + container = os.path.realpath(containing_dir) + if os.path.commonpath([resolved, container]) != container or resolved == container: + raise KernelConfigPinError( + f"refusing to delete {resolved!r}: it escapes the pin directory {container!r}", + ) if not os.path.isfile(os.path.join(path, OWNED_SENTINEL)): raise KernelConfigPinError( f"refusing to delete {path!r}: it lacks the ownership sentinel " @@ -117,7 +132,7 @@ def seed_exact_kernel_config_pin(pin_dir: str, *, source_cache: str | None = Non f"seed source cache {cache_src!r} contains the pin destination {cache_dst!r}; " "copying would recurse into itself", ) - _rmtree_owned(cache_dst) + _rmtree_owned(cache_dst, containing_dir=pin_dir) if os.path.isdir(cache_src): shutil.copytree(cache_src, cache_dst) else: @@ -168,7 +183,7 @@ def pin_exact_kernel_configs(*, rank: int | None = None) -> str: if rank is None: rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) clone = os.path.join(pin_dir, "clones", f"rank{rank}") - _rmtree_owned(clone) + _rmtree_owned(clone, containing_dir=pin_dir) shutil.copytree(os.path.join(pin_dir, CACHE_SUBDIR), clone) _mark_owned(clone) os.environ["TRITON_CACHE_DIR"] = clone