[megatron] Fused LM-head log-prob + entropy (avoid full [*, seq, vocab] logit materialization)

[dyurk-lila]

Summary

Adds an opt-in fused LM-head path for the Megatron backend that folds the vocab projection into the chunked log-prob (and entropy), so the full [*, seq, vocab // TP] logits tensor is never materialized. That tensor is the dominant activation transient on the log-prob path at large vocab + long context (e.g. ~4 GiB fp32 at micro-bs 2 / 32k ctx / 131k vocab / TP4); eliminating it unlocks larger micro-batches / no-recompute without OOM.

Serves SFT and all RL losses (cross_entropy, ppo, grpo, …): the fused kernel returns (log_probs, entropy) — verl's established dual-output — so RL's entropy term is satisfied from the fused output instead of needing the full logits. Numerically equivalent to the stock logits path; default OFF → byte-identical to today.

Design ported from verl's fused linear-cross-entropy (volcengine/verl, Apache-2.0). Two backends behind one flag:

• torch (default): pure-PyTorch, vocab-parallel FusedLinearLogprob. Chunks over the seq dim, recomputes per-chunk logits in backward, fp32 softmax/entropy math, masks out-of-shard targets. Runs anywhere (CPU/GPU), no extra deps.
• triton: vendored flash-style Triton kernel — tiles over the vocab dim so per-chunk logits never materialize (lower memory floor + faster). GPU + triton required; falls back to torch with a warning if unavailable.

Config

trainer.fused_linear_logprob: false          # default off → stock behaviour, byte-identical
trainer.fused_linear_logprob_backend: torch  # "torch" | "triton"

Megatron-only. Uses the existing logprobs_chunk_size for the chunk width.

How it integrates (no megatron-core edits)

MegatronModelWrapper wraps output_layer.forward at model-build to return the pre-projection hidden state (skipping the vocab GEMM) and capture the LM-head weight shard. Correctness points:

• Logprob and entropy come from one fused call. SFT (cross_entropy) requests logprob-only (skips the entropy all-reduce); RL requests return_entropy=True. The fused per-token entropy = logsumexp(logits) − Σ(softmax·logits), algebraically identical to stock _VocabParallelEntropy and verl; backward adds the matching dentropy term. RL entropy is sliced to the action window (non-packed) or fed through vocab_parallel_entropy_packed_sequences via a new precomputed_entropy_tokens arg (packed action-weighting + CP all-reduce unchanged — verified a bit-identical substitution). KL keeps using logprobs.
• dtype: the fused path casts hidden → LM-head weight.dtype before the GEMM (the common hybrid case is fp32 pre-head hidden + bf16 head), mirroring ColumnParallelLinear; grad cast back on the way out.
• Tied embeddings: weight taken from the weight= kwarg the model passes (shared_embedding_or_output_weight()) → no output_layer.weight is None crash (cf. verl#3730).
• Sequence parallelism: SP hidden gathered with tensor_parallel_output_grad=True; its reduce-scatter backward performs the TP reduction of grad_hidden exactly as stock ColumnParallelLinear. Vocab range derived from the captured weight shard, not hidden.shape[-1].
• Both training loss_func and forward-only collection_func route through the fused path. Auto-fallback to stock under MuP output scaling; default fused-off is byte-identical.

Verification

• CPU (tests/backends/skyrl_train/cpu/megatron/test_fused_linear_logprob.py): gloo TP1/TP2, chunked + single, with/without OOV. Logprob fwd + grad-hidden bit-exact, grad-weight to fp32 rounding vs the stock logits path; entropy fwd + backward parity vs stock vocab_parallel_entropy (~5e-7), entropy target-independence, packed precomputed_entropy_tokens bit-identical substitution, slice-alignment guard. Runs on CPU CI (no GPU / no megatron-core needed — megatron.core.parallel_state is stubbed when absent).
• GPU (tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py), 2×H100: 24/24 pass. Logprob 16/16 (incl. fp32+bf16 dtype-mix) + entropy 8/8 (fp32 ~5e-7 / ~1e-8 grads under IEEE; bf16 within tolerance; OOV-invariant). Marked @pytest.mark.megatron so the megatron GPU CI job (-m megatron) selects it. The fp32 cases pin IEEE (FORCE_FP32_IEEE_PRECISION) as a math-is-exact sanity check (default-precision fwd error 1.27e-3 TF32 vs 9.5e-7 IEEE); production uses the fast default (TF32) precision — the model trains in bf16 where it matches stock within ~3e-3.

Licensing / attribution

verl is Apache-2.0 (same as SkyRL); SkyRL already vendors verl in ~8 files with the # This code is adapted from VERL … original copyright reproduced below: convention. The vendored files (verl/utils/kernel/{kernels.py,linear_cross_entropy.py}) carry a dual copyright — Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES and Copyright 2024 Bytedance Ltd. and/or its affiliates (verl is a Bytedance/volcengine project; the NVIDIA line is an NVIDIA-contributed portion). Both are reproduced verbatim in the vendored file header with a Changes: note (§4(b)). Adds a top-level NOTICE (SkyRL had none) reproducing verl's Notice.txt. Preserves verl's out-of-bounds-vocab -inf masking fix (verl#2656).

Test plan

• uv run --isolated --extra dev -- pytest -s tests/backends/skyrl_train/cpu/megatron/test_fused_linear_logprob.py (CPU, no megatron-core)
• pytest -s tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py -m megatron on 2×H100 — 24/24
• ruff + black clean (pre-commit)
• Default fused_linear_logprob: false path byte-identical to current behaviour

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces a fused LM-head log-prob optimization for the Megatron backend, featuring both a pure-PyTorch chunked implementation and a high-performance Triton-based backend vendored from verl. By fusing the LM-head projection into the chunked log-prob and entropy computation, the PR avoids materializing the massive full-vocab logits tensor, significantly reducing peak memory usage. The feedback recommends enhancing the robustness of the model wrapper by using *args and **kwargs to guard against future Megatron-core signature changes, broadening the exception handling when importing the Triton backend to catch driver or library mismatches, and checking CUDA availability alongside Triton to prevent cryptic errors in CPU-only environments.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Using explicit keyword arguments like weight=None and runtime_gather_output=None can lead to TypeError if Megatron-core updates its ParallelLinear.forward signature or if other wrappers pass unexpected arguments. Using *args and **kwargs makes the wrapper future-proof and robust against signature changes.

            def fused_output_layer_forward(input_, *args, **kwargs):
                w = kwargs.get("weight", None)
                if w is None and len(args) > 0:
                    w = args[0]
                w = w if w is not None else getattr(layer, "weight", None)
                if w is None:
                    # Tied-embedding stage with no allocated head weight AND none passed in:
                    # we cannot fuse safely — defer to the original projection.
                    return original(input_, *args, **kwargs)
                hidden = input_
                if getattr(layer, "sequence_parallel", False):
                    # Gather the sequence-parallel-scattered hidden across TP. Its backward is a
                    # reduce-scatter (tensor_parallel_output_grad=True), which reduces the fused
                    # function's per-shard grad_hidden exactly as ColumnParallelLinear would.
                    hidden = gather_from_sequence_parallel_region(
                        hidden, tensor_parallel_output_grad=True, group=tpg
                    )
                holder["weight"] = w
                # Return hidden in the logits position (+ no bias). The model's _scale_logits is
                # identity here (MuP excluded above) and its transpose yields [b, s, h].
                return hidden, None

[dyurk-lila]

Claude suggested that this is a proper copyright notice but this should definitely be checked by someone more familiar with the repo's conventions around this