chore(megatron): reconcile fused-logprob onto upstream #1841

Rebase re-scoped this PR to layer only its distinguishing additions (the Triton
backend + its GPU tests) on top of the merged #1841 fused-linear-cross-entropy.
The torch chunked-logprob core is now upstream's FusedLinearChunkedDistributedLogprob;
this commit fixes the vendored Triton adapter's docstrings, the benchmark, and the
GPU test to reference the upstream class / the reconciled fused_lm_head_logprob_backend
selector instead of this PR's now-dropped FusedLinearLogprob / _fused_linear_logprob_apply.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: dyurk-lila

skyrl/backends/skyrl_train/distributed/megatron/fused_linear_logprob_triton.py

@@ -25,8 +25,9 @@
 #     fallback; we additionally surface a module-level flag and a hard guard.)
 #   * Added ``FusedLinearLogprobTriton`` autograd Function — a SkyRL-shaped
 #     adapter around verl's ``LinearCrossEntropy``/kernels that matches the
-#     pure-torch ``FusedLinearLogprob`` contract in ``model_utils.py`` exactly
-#     (see that class for the contract). This is NEW code, not from verl.
+#     pure-torch ``FusedLinearChunkedDistributedLogprob`` contract in
+#     ``model_utils.py`` exactly (see that class for the contract). This is NEW
+#     code, not from verl.
 #
 # NOTE on the #2656 out-of-bounds-vocab fix: verl's forward mainloop masks
 # out-of-bounds vocab columns to -inf for the LSE max/exp/denominator while
@@ -68,7 +69,7 @@
 This module vendors verl's online-softmax fused linear-cross-entropy Triton
 kernels (with the #2656 TP all-reduce / OOB-vocab masking scheme) and exposes
 :class:`FusedLinearLogprobTriton`, an autograd Function that matches the exact
-interface of the pure-torch ``FusedLinearLogprob`` in ``model_utils.py``.
+interface of the pure-torch ``FusedLinearChunkedDistributedLogprob`` in ``model_utils.py``.
 
 Triton/CUDA is required to actually *run* the kernel, but the module is import-safe
 without triton (``TRITON_AVAILABLE`` will be ``False`` and the adapter raises a
@@ -1888,8 +1889,9 @@ def backward(ctx, dlogprobs: torch.Tensor, dentropy: torch.Tensor) -> list:
 
 
 # ======================================================================================
-# SkyRL adapter (NEW code — not from verl). Matches the pure-torch FusedLinearLogprob
-# contract in model_utils.py exactly so the two backends are interchangeable.
+# SkyRL adapter (NEW code — not from verl). Matches the pure-torch
+# FusedLinearChunkedDistributedLogprob contract in model_utils.py exactly so the two
+# backends are interchangeable.
 # ======================================================================================
 
 # A label sentinel guaranteed to fall in NO shard's column range after re-keying. verl
@@ -1905,11 +1907,13 @@ def _verl_logprob_kernel_available() -> bool:
 class FusedLinearLogprobTriton(torch.autograd.Function):
     """Triton fused LM-head + vocab-parallel log-prob of the target token.
 
-    Drop-in, interface-identical replacement for the pure-torch ``FusedLinearLogprob``
-    in ``model_utils.py`` (same ``apply`` signature, same return/grad contract), backed by
-    verl's online-softmax fused linear-cross-entropy Triton kernels (vendored above).
+    Drop-in, interface-identical replacement for the pure-torch
+    ``FusedLinearChunkedDistributedLogprob`` in ``model_utils.py`` (same ``apply`` signature, same
+    return/grad contract), backed by verl's online-softmax fused linear-cross-entropy Triton
+    kernels (vendored above). Selected via ``fused_lm_head_logprob_backend="triton"`` and dispatched
+    from ``model_utils._fused_lm_head_logprob_apply``.
 
-    Contract (must match ``FusedLinearLogprob`` exactly):
+    Contract (must match ``FusedLinearChunkedDistributedLogprob`` exactly):
       * ``apply(hidden, weight, target, vocab_start_index, vocab_end_index, chunk_size,
         tp_group, inference_only=False) -> log_probs``.
       * ``hidden``: ``[B, S, H]`` (TP-replicated / SP-already-gathered by the caller).
@@ -1942,7 +1946,7 @@ class FusedLinearLogprobTriton(torch.autograd.Function):
          label-logit all-reduce. verl's epilogue would then return ``max + log(accu) - 0`` = the
          LSE (garbage, not 0). We therefore force ``log_prob = 0`` for fully-OOV positions in
          forward. We do NOT touch the backward gradient at those positions: the references (stock
-         ``DistributedLogprob`` and the pure-torch ``FusedLinearLogprob``) still emit
+         ``DistributedLogprob`` and the pure-torch ``FusedLinearChunkedDistributedLogprob``) still emit
          ``-softmax * grad_output`` there (the chosen-position term is absent because no shard owns
          the column), and verl's kernel already reproduces that exact value for sentinel-keyed
          labels. Zeroing the backward gradient at those positions would drop that term and diverge
@@ -1964,8 +1968,8 @@ def forward(  # pyrefly: ignore[bad-override]
         if not _verl_logprob_kernel_available():
             raise RuntimeError(
                 "FusedLinearLogprobTriton requires Triton (and a CUDA device), but triton is not "
-                "importable. Install the optional 'triton' extra, or use the pure-torch "
-                "FusedLinearLogprob backend instead."
+                "importable. Install the optional 'triton' extra, or use the default pure-torch "
+                "'torch' backend (FusedLinearChunkedDistributedLogprob) instead."
             )
 
         B, S, H = hidden.shape
@@ -1983,7 +1987,7 @@ def forward(  # pyrefly: ignore[bad-override]
 
         # hidden must be contiguous & 2D for the kernel. The kernel's ``tl.dot(hidden, weight.T)``
         # requires BOTH operands to share a dtype, so cast hidden to the LM-head weight dtype here
-        # — exactly what the pure-torch ``FusedLinearLogprob`` does (``hidden.to(weight.dtype) @
+        # — exactly what the pure-torch ``FusedLinearChunkedDistributedLogprob`` does (``hidden.to(weight.dtype) @
         # weight.t()``) and what ``ColumnParallelLinear`` does before its bf16 GEMM. On a real
         # hybrid model the pre-head hidden is fp32 (final norm) while the head is bf16, so without
         # this the kernel asserts "Both operands must be same dtype. Got fp32 and bf16". When the
@@ -2047,7 +2051,7 @@ def backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple:
         # would otherwise return the LSE there (no shard contributes the label logit). But the
         # GRADIENT contract is set by the references this backend must match exactly — stock
         # ``DistributedLogprob`` / ``ChunkedDistributedLogprob`` and the pure-torch
-        # ``FusedLinearLogprob`` in model_utils.py. Those compute, for every position
+        # ``FusedLinearChunkedDistributedLogprob`` in model_utils.py. Those compute, for every position
         # (in-shard, out-of-shard, or fully-OOV alike):
         #     d_logits = grad_output * (onehot(target) - softmax)
         # and only ADD the ``onehot`` (chosen-position) term where the target is owned by this

skyrl/backends/skyrl_train/distributed/megatron/model_utils.py

@@ -16,6 +16,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from typing import Any, Optional
 
 import megatron.core.parallel_state as mpu

skyrl/benchmarks/bench_fused_ce_h100.py

@@ -11,8 +11,8 @@
   nvidia       : same logits, then ``fused_vocab_parallel_cross_entropy``
                  (@jit_fuser fused CE — fuses the softmax/CE stages, still
                  materializes the logits).
-  fused-torch  : FusedLinearLogprob (pure-torch chunked fused linear logprob —
-                 folds the projection into the chunked TP log-prob; logits never
+  fused-torch  : FusedLinearChunkedDistributedLogprob (pure-torch chunked fused linear
+                 logprob — folds the projection into the chunked TP log-prob; logits never
                  materialized).
   fused-triton : FusedLinearLogprobTriton (vendored verl/bytedance flash kernel).
 
@@ -49,14 +49,16 @@ def _loss(mode, hidden, weight, target, vocab_local, chunk_size, tp_group):
     from megatron.core.tensor_parallel.cross_entropy import vocab_parallel_cross_entropy
 
     from skyrl.backends.skyrl_train.distributed.megatron.model_utils import (
-        FusedLinearLogprob,
+        FusedLinearChunkedDistributedLogprob,
     )
     from skyrl.backends.skyrl_train.distributed.megatron.fused_linear_logprob_triton import (
         FusedLinearLogprobTriton,
     )
 
     if mode == "fused-torch":
-        lp = FusedLinearLogprob.apply(hidden, weight, target, 0, vocab_local, chunk_size, tp_group, False)
+        lp = FusedLinearChunkedDistributedLogprob.apply(
+            hidden, weight, target, 0, vocab_local, chunk_size, tp_group, False
+        )
         return (-lp).sum()
     if mode == "fused-triton":
         lp = FusedLinearLogprobTriton.apply(hidden, weight, target, 0, vocab_local, chunk_size, tp_group, False)
@@ -138,7 +140,7 @@ def main():
         print(f"Device {torch.cuda.get_device_name(device)} | TP(world)={world} | hidden={HIDDEN} | chunk={args.chunk_size}")
         print(
             "baseline = vocab_parallel_cross_entropy (eager) | nvidia = fused_vocab_parallel_cross_entropy | "
-            "fused-torch = FusedLinearLogprob | fused-triton = FusedLinearLogprobTriton"
+            "fused-torch = FusedLinearChunkedDistributedLogprob | fused-triton = FusedLinearLogprobTriton"
         )
         print("all: hidden+weight -> per-token CE -> sum -> backward\n")
         print("correctness (TP=%d, vocab=%d):" % (world, vocab_shards[0]))

tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py

@@ -56,7 +56,7 @@ def _direct_fused_logprobs(hidden, weight_shard, target_shifted, vstart, vend, c
 
     The public entry point shifts targets internally; here we pass an ALREADY-shifted target and
     compare against a directly-driven stock reference, so the Triton kernel is exercised in
-    isolation from the ``_fused_linear_logprob_apply`` backend selection.
+    isolation from the ``_fused_lm_head_logprob_apply`` backend selection.
     """
     leaf_h = hidden.detach().clone().requires_grad_(True)
     leaf_w = weight_shard.detach().clone().requires_grad_(True)