feat: add Megatron fused linear cross entropy backends

docs/High-Level-APIs.md

@@ -98,38 +98,57 @@ You can also use the Patching APIs to use the kernels for a specific model archi
 
 Liger also exposes a patch for the [Megatron-LM](https://github.com/NVIDIA/Megatron-LM)
 training framework, replacing Megatron's native RMSNorm and both vocab-parallel
-cross-entropy paths (fused and unfused) with Liger's Triton kernels.
+cross-entropy paths (fused and unfused) with Liger kernels. It can also fuse
+the GPT output projection with cross-entropy.
 
 | **Framework** | **API**                                                | **Supported Operations** |
 |---------------|--------------------------------------------------------|--------------------------|
-| Megatron-LM   | `liger_kernel.megatron.apply_liger_kernel_to_megatron` | RMSNorm, CrossEntropyLoss |
+| Megatron-LM   | `liger_kernel.megatron.apply_liger_kernel_to_megatron` | RMSNorm, CrossEntropyLoss, fused output projection + CrossEntropyLoss |
+| Megatron-LM   | `liger_kernel.megatron.LigerMegatronFusedLinearCrossEntropy` | Fused output projection + CrossEntropyLoss |
 
-**Scope**: Initial release supports `tensor_model_parallel_size=1` only for
-cross-entropy. Vocab-parallel cross-entropy (TP>1) is follow-up work — with
-TP>1, each rank holds a sharded `[N, V/tp]` logits slice and cross-entropy
-requires cross-rank all-reduces that Liger's kernel does not perform. The
-patch raises a `RuntimeError` at patch time or call time if TP>1 is detected.
+Both cross-entropy patches and FLCE support TP1 and TP>1. Automatic FLCE
+patching requires Megatron-Core 0.18 or newer, BF16 or FP16, a native
+`ColumnParallelLinear` output layer, and no sequence parallelism,
+gradient-accumulation fusion, or gathered logits. Unsupported configurations
+raise an explicit error.
 
 **Usage**:
 
 ```python
 from liger_kernel.megatron import apply_liger_kernel_to_megatron
 
-# Call before Megatron's forward pass reaches compute_language_model_loss.
-# Defaults match Megatron's native CE behavior; no CE-specific config needed.
-apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)
+apply_liger_kernel_to_megatron(
+    rms_norm=True,
+    cross_entropy=True,
+    fused_linear_cross_entropy=True,
+)
 ```
 
 Both the fused (`config.cross_entropy_loss_fusion=True`,
 `cross_entropy_fusion_impl='native'`) and unfused
 (`config.cross_entropy_loss_fusion=False`) CE paths are patched in a single
-call, so Megatron picks up Liger regardless of which path your config selects.
+call. When `fused_linear_cross_entropy=True`, labeled standard GPT forwards
+bypass both paths and use FLCE directly. For custom integrations, use
+`LigerMegatronFusedLinearCrossEntropy` with replicated hidden states, a
+contiguous local vocabulary shard, and global target indices.
 
 For training setups that need explicit kernel configuration (custom
 `ignore_index`, `label_smoothing`, etc.), instantiate
 `LigerMegatronCrossEntropy` directly and wire it into your model — see
 `examples/megatron/run_mode2_hand_spec.py`.
 
+::: liger_kernel.megatron.LigerMegatronFusedLinearCrossEntropy
+    options:
+      extra:
+        show_docstring: true
+        show_signature: true
+
+::: liger_kernel.megatron.liger_megatron_fused_linear_cross_entropy_output_processor
+    options:
+      extra:
+        show_docstring: true
+        show_signature: true
+
 ::: liger_kernel.megatron.apply_liger_kernel_to_megatron
     options:
       extra:

src/liger_kernel/megatron/__init__.py

@@ -10,10 +10,14 @@
         experts. Mode 1 patches ``fused_bias_swiglu.SwiGLUFunction`` instead;
         both fall back to Megatron for FP8 input store and CPU activation
         offload, and neither touches the bias or MoE-routed variants.
+    LigerMegatronFusedLinearCrossEntropy — hidden-state-to-loss fused output
+        projection for tensor-parallel vocabulary shards.
+    liger_megatron_fused_linear_cross_entropy_output_processor — adapts FLCE
+        to Megatron's GPT output-processor hook.
     apply_liger_kernel_to_megatron — patches Megatron-Core so existing training
         scripts pick up Liger kernels with one line. Currently supports
         RMSNorm (via BackendSpecProvider), both the fused and unfused
-        vocab-parallel cross-entropy paths, and SwiGLU.
+        vocab-parallel cross-entropy paths, opt-in GPT FLCE, and SwiGLU.
 
 The general-purpose ``LigerVocabParallelCrossEntropy`` Module lives under
 ``liger_kernel.transformers`` alongside the other nn.Module wrappers; the
@@ -23,13 +27,17 @@
 """
 
 from liger_kernel.megatron.cross_entropy import LigerMegatronCrossEntropy
+from liger_kernel.megatron.fused_linear_cross_entropy import LigerMegatronFusedLinearCrossEntropy
+from liger_kernel.megatron.fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy_output_processor
 from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron
 from liger_kernel.megatron.rms_norm import LigerMegatronRMSNorm
 from liger_kernel.megatron.swiglu import LigerMegatronSwiGLU
 
 __all__ = [
     "LigerMegatronCrossEntropy",
+    "LigerMegatronFusedLinearCrossEntropy",
     "LigerMegatronRMSNorm",
     "LigerMegatronSwiGLU",
     "apply_liger_kernel_to_megatron",
+    "liger_megatron_fused_linear_cross_entropy_output_processor",
 ]

src/liger_kernel/megatron/fused_linear_cross_entropy.py

@@ -0,0 +1,100 @@
+"""Megatron-facing module for tensor-parallel fused linear cross entropy."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops import LigerMegatronFusedLinearCrossEntropyFunction
+
+
+class LigerMegatronFusedLinearCrossEntropy(nn.Module):
+    """Fuse a vocab-sharded output projection with per-token cross entropy.
+
+    ``hidden`` is replicated across TP ranks and ``weight`` contains the local
+    contiguous vocabulary shard. The output shape matches ``target``.
+    """
+
+    def __init__(
+        self,
+        ignore_index: int = -100,
+    ):
+        super().__init__()
+        self.ignore_index = ignore_index
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        target: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        tp_group=None,
+    ) -> torch.Tensor:
+        return LigerMegatronFusedLinearCrossEntropyFunction.apply(
+            hidden,
+            weight,
+            target,
+            bias,
+            tp_group,
+            self.ignore_index,
+        )
+
+    def extra_repr(self) -> str:
+        return f"ignore_index={self.ignore_index}"
+
+
+def liger_megatron_fused_linear_cross_entropy_output_processor(
+    *,
+    hidden_states: torch.Tensor,
+    output_layer,
+    output_weight: torch.Tensor | None,
+    labels: torch.Tensor,
+    runtime_gather_output: bool | None,
+    config,
+    **_,
+) -> torch.Tensor:
+    """Megatron ``GPTModel`` output processor for the native TP output layer."""
+    unsupported = []
+    if type(output_layer).__name__ != "ColumnParallelLinear":
+        unsupported.append("the output layer is not Megatron's native ColumnParallelLinear")
+    if getattr(output_layer, "sequence_parallel", False):
+        unsupported.append("sequence_parallel=True")
+    if getattr(output_layer, "gradient_accumulation_fusion", False):
+        unsupported.append("gradient_accumulation_fusion=True")
+    if getattr(output_layer, "disable_grad_reduce", False):
+        unsupported.append("output-layer dgrad reduction is disabled")
+    if getattr(output_layer, "explicit_expert_comm", False):
+        unsupported.append("the output layer uses explicit expert communication")
+    if getattr(output_layer, "skip_bias_add", False):
+        unsupported.append("the output layer returns bias separately")
+    if getattr(config, "defer_embedding_wgrad_compute", False):
+        unsupported.append("defer_embedding_wgrad_compute=True")
+    if getattr(config, "mtp_num_layers", None):
+        unsupported.append("MTP is enabled")
+    if getattr(config, "use_mup", False):
+        unsupported.append("MuP output scaling is enabled")
+
+    gather_output = (
+        getattr(output_layer, "gather_output", False) if runtime_gather_output is None else runtime_gather_output
+    )
+    if gather_output:
+        unsupported.append("the output layer gathers TP logits")
+    if unsupported:
+        raise RuntimeError(
+            "Liger Megatron FLCE does not support this GPT output configuration: " + "; ".join(unsupported)
+        )
+
+    weight = output_weight if output_weight is not None else getattr(output_layer, "weight", None)
+    if weight is None:
+        raise RuntimeError("Liger Megatron FLCE requires an output weight tensor.")
+
+    labels_sb = labels.transpose(0, 1).contiguous()
+    loss_sb = LigerMegatronFusedLinearCrossEntropyFunction.apply(
+        hidden_states,
+        weight,
+        labels_sb,
+        getattr(output_layer, "bias", None),
+        getattr(output_layer, "tp_group", None),
+        -100,
+    )
+    return loss_sb.transpose(0, 1).contiguous()

src/liger_kernel/megatron/monkey_patch.py

@@ -2,16 +2,27 @@
 
 from __future__ import annotations
 
+import functools
+import inspect
 import logging
+import sys
 
 logger = logging.getLogger(__name__)
 
 _PATCH_MARKER = "__liger_patched__"
 
 
+def _replace_loaded_binding(module_name: str, symbol_name: str, original, replacement) -> None:
+    """Update a known by-name import without clobbering third-party patches."""
+    module = sys.modules.get(module_name)
+    if module is not None and getattr(module, symbol_name, None) is original:
+        setattr(module, symbol_name, replacement)
+
+
 def apply_liger_kernel_to_megatron(
     rms_norm: bool = True,
     cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = False,
     swiglu: bool = False,
 ) -> None:
     """Patch Megatron-Core to use Liger Triton kernels.
@@ -39,6 +50,11 @@ def apply_liger_kernel_to_megatron(
             wrapper additionally honors a runtime ``label_smoothing``
             argument, matching native's
             ``(logits, target, label_smoothing=0.0, tp_group=None)``.
+        fused_linear_cross_entropy: When ``True`` inject Liger's fused local
+            output projection and vocab-parallel cross-entropy through
+            ``GPTModel._postprocess`` for supported native
+            ``ColumnParallelLinear`` training configurations. This is opt-in
+            and independent of ``cross_entropy``.
         swiglu: When ``True`` replace
             ``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` with Liger's
             Triton SiLU-multiply kernel, covering the dense ``MLP`` and the MoE
@@ -68,6 +84,8 @@ def apply_liger_kernel_to_megatron(
     if cross_entropy:
         _patch_fused_vocab_parallel_cross_entropy()
         _patch_vocab_parallel_cross_entropy()
+    if fused_linear_cross_entropy:
+        _patch_gpt_fused_linear_cross_entropy()
     if swiglu:
         _patch_swiglu_function()
 
@@ -183,7 +201,14 @@ def _patch_fused_vocab_parallel_cross_entropy() -> None:
         )
 
     if getattr(fused_ce.fused_vocab_parallel_cross_entropy, _PATCH_MARKER, False):
-        return  # already patched
+        replacement = fused_ce.fused_vocab_parallel_cross_entropy
+        _replace_loaded_binding(
+            "megatron.core.models.common.language_module.language_module",
+            "fused_vocab_parallel_cross_entropy",
+            replacement.__wrapped__,
+            replacement,
+        )
+        return
 
     original = fused_ce.fused_vocab_parallel_cross_entropy
 
@@ -197,11 +222,15 @@ def liger_fused_vocab_parallel_cross_entropy(vocab_parallel_logits, target, tp_g
     setattr(liger_fused_vocab_parallel_cross_entropy, _PATCH_MARKER, True)
     setattr(liger_fused_vocab_parallel_cross_entropy, "__wrapped__", original)
     fused_ce.fused_vocab_parallel_cross_entropy = liger_fused_vocab_parallel_cross_entropy
-
-    logger.info(
-        "Patched megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy with Liger cross-entropy."
+    _replace_loaded_binding(
+        "megatron.core.models.common.language_module.language_module",
+        "fused_vocab_parallel_cross_entropy",
+        original,
+        liger_fused_vocab_parallel_cross_entropy,
     )
 
+    logger.info("Patched Megatron's fused vocab-parallel cross-entropy definition and loaded LanguageModule binding.")
+
 
 def _patch_vocab_parallel_cross_entropy() -> None:
     """Replace ``megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy``.
@@ -229,7 +258,14 @@ def _patch_vocab_parallel_cross_entropy() -> None:
         )
 
     if getattr(unfused_ce.vocab_parallel_cross_entropy, _PATCH_MARKER, False):
-        return  # already patched
+        replacement = unfused_ce.vocab_parallel_cross_entropy
+        _replace_loaded_binding(
+            "megatron.core.tensor_parallel",
+            "vocab_parallel_cross_entropy",
+            replacement.__wrapped__,
+            replacement,
+        )
+        return
 
     original = unfused_ce.vocab_parallel_cross_entropy
 
@@ -260,11 +296,58 @@ def liger_vocab_parallel_cross_entropy(
     setattr(liger_vocab_parallel_cross_entropy, _PATCH_MARKER, True)
     setattr(liger_vocab_parallel_cross_entropy, "__wrapped__", original)
     unfused_ce.vocab_parallel_cross_entropy = liger_vocab_parallel_cross_entropy
+    _replace_loaded_binding(
+        "megatron.core.tensor_parallel",
+        "vocab_parallel_cross_entropy",
+        original,
+        liger_vocab_parallel_cross_entropy,
+    )
 
-    logger.info(
-        "Patched megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy with Liger cross-entropy."
+    logger.info("Patched Megatron's unfused vocab-parallel cross-entropy definition and tensor_parallel export.")
+
+
+def _patch_gpt_fused_linear_cross_entropy() -> None:
+    """Inject Liger FLCE through Megatron's GPT output-processor hook."""
+    try:
+        import megatron.core.models.gpt.gpt_model as gpt_model
+    except ImportError as exc:
+        raise ImportError(
+            "apply_liger_kernel_to_megatron(fused_linear_cross_entropy=True) requires "
+            "megatron-core with megatron.core.models.gpt.gpt_model.GPTModel."
+        ) from exc
+
+    if not hasattr(gpt_model, "GPTModel") or not hasattr(gpt_model.GPTModel, "_postprocess"):
+        raise ImportError(
+            "megatron.core.models.gpt.gpt_model.GPTModel._postprocess was not found. "
+            "The symbol path may have changed in your Megatron-Core version."
+        )
+
+    current = gpt_model.GPTModel._postprocess
+    if getattr(current, _PATCH_MARKER, False):
+        return
+    signature = inspect.signature(current)
+    if not {"labels", "output_processor"} <= signature.parameters.keys():
+        raise ImportError(
+            "Megatron GPTModel._postprocess does not expose the labels and output_processor hooks "
+            "required by Liger FLCE. Upgrade to Megatron-Core 0.18 or newer."
+        )
+
+    from liger_kernel.megatron.fused_linear_cross_entropy import (
+        liger_megatron_fused_linear_cross_entropy_output_processor,
     )
 
+    @functools.wraps(current)
+    def liger_postprocess(self, *args, **kwargs):
+        bound = signature.bind(self, *args, **kwargs)
+        bound.apply_defaults()
+        if bound.arguments["labels"] is not None and bound.arguments["output_processor"] is None:
+            bound.arguments["output_processor"] = liger_megatron_fused_linear_cross_entropy_output_processor
+        return current(*bound.args, **bound.kwargs)
+
+    setattr(liger_postprocess, _PATCH_MARKER, True)
+    gpt_model.GPTModel._postprocess = liger_postprocess
+    logger.info("Patched Megatron GPTModel._postprocess to inject Liger fused linear cross-entropy.")
+
 
 def _patch_swiglu_function() -> None:
     """Replace ``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` with Liger.

src/liger_kernel/ops/__init__.py

@@ -67,6 +67,10 @@
 from liger_kernel.ops.layer_norm import layer_norm_backward  # noqa: F401
 from liger_kernel.ops.layer_norm import layer_norm_forward  # noqa: F401
 from liger_kernel.ops.llama4_rope import LigerLlama4RopeFunction  # noqa: F401
+from liger_kernel.ops.megatron_fused_linear_cross_entropy import (
+    LigerMegatronFusedLinearCrossEntropyFunction,  # noqa: F401
+)
+from liger_kernel.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy  # noqa: F401
 from liger_kernel.ops.mhc import LigerMHCCoeffsFunction  # noqa: F401
 from liger_kernel.ops.mhc import LigerMHCPostResFunction  # noqa: F401
 from liger_kernel.ops.mhc import LigerMHCPreFunction  # noqa: F401

src/liger_kernel/ops/cutedsl/ops/__init__.py

@@ -18,6 +18,10 @@
 from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import LigerFusedScaledCrossEntropySM90Function
 from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import fused_scaled_cross_entropy_backward
 from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import fused_scaled_cross_entropy_forward
+from liger_kernel.ops.cutedsl.ops.megatron_fused_linear_cross_entropy import (
+    LigerMegatronFusedLinearCrossEntropyFunction,
+)
+from liger_kernel.ops.cutedsl.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy
 from liger_kernel.ops.cutedsl.ops.rms_norm import LigerRMSNormFunction
 from liger_kernel.ops.cutedsl.ops.rms_norm import rms_norm_backward
 from liger_kernel.ops.cutedsl.ops.rms_norm import rms_norm_forward
@@ -41,6 +45,8 @@
     "LigerFusedScaledCrossEntropySM90Function",
     "fused_scaled_cross_entropy_backward",
     "fused_scaled_cross_entropy_forward",
+    "LigerMegatronFusedLinearCrossEntropyFunction",
+    "liger_megatron_fused_linear_cross_entropy",
     "LigerRMSNormFunction",
     "rms_norm_backward",
     "rms_norm_forward",