[Megatron] Rebind CE call sites so apply_liger_kernel_to_megatron(cross_entropy=True) takes effect (#1486)

## Summary

Fixes #1485.

`apply_liger_kernel_to_megatron(cross_entropy=True)` replaces the CE
functions only on their defining modules, but Megatron calls copies
bound at import time. When Megatron model code is imported before the
patch (the usual order, also used by
`examples/megatron/run_mode1_monkey_patch.py`), the loss silently stays
on Megatron's kernel:

- Released megatron-core (0.18 / 0.19): `language_module` binds
`fused_vocab_parallel_cross_entropy` with `from ... import`
([L17](https://github.com/NVIDIA/Megatron-LM/blob/571370c829ca768fe37244f4e2e7f28d8accc4ab/megatron/core/models/common/language_module/language_module.py#L17)).
The unfused branch calls the `megatron.core.tensor_parallel` re-export
([L183](https://github.com/NVIDIA/Megatron-LM/blob/571370c829ca768fe37244f4e2e7f28d8accc4ab/megatron/core/models/common/language_module/language_module.py#L183)),
which is bound before the `cross_entropy` submodule can be patched.
Neither path reaches Liger.
- Megatron-LM main (`NVIDIA/Megatron-LM#6766`): `backends` binds the
fused function as `_fused_ce`
([L19-L21](https://github.com/NVIDIA/Megatron-LM/blob/ec806b20a0a976924c1668486c3554ffb90c000e/megatron/core/models/backends.py#L19-L21)).
The fused path misses Liger; the unfused path already works.

This PR rebinds those copies too, but only where they still hold the
original function. No version check is needed, and a call site that
another framework has already overridden is left alone. If no fused call
site ends up on Liger, the patch logs a warning. The docs now say to
call the patch before building the model, matching the docstring:
Megatron-LM main picks the CE function when the model is constructed.

Once the minimum supported megatron-core includes the
`BackendSpecProvider.vocab_parallel_cross_entropy` slot, this can switch
to overriding that slot on the providers and drop the rebinding for
released versions.

## Testing Done

- Hardware Type: 1x NVIDIA GPU
- [ ] run `make test` to ensure correctness: ran `test/megatron` (173
passed, 48 skipped multi-GPU tests)
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence: N/A, no kernel
change

The new tests in `test_monkey_patch.py` compute the loss through stub
call sites laid out like each Megatron version; 4 of 5 fail on current
`main`.

A real 1-layer `GPTModel` (TP=1) was run for one forward and backward
pass, with model code imported before the patch. Calls into Liger CE
(fused / unfused):

| | before | after |
|---|---|---|
| megatron-core 0.18.2 | 0 / 0 | 1 / 1 |
| Megatron-LM main | 0 / 1 | 1 / 1 |

When the patch runs first, both cases also give 1 / 1 after this change
(before: 1 / 0 on 0.18.2).

Co-authored-by: Vaibhav Jindal <[email protected]>

docs/High-Level-APIs.md

@@ -115,7 +115,7 @@ patch raises a `RuntimeError` at patch time or call time if TP>1 is detected.
 ```python
 from liger_kernel.megatron import apply_liger_kernel_to_megatron
 
-# Call before Megatron's forward pass reaches compute_language_model_loss.
+# Call before building the model; newer Megatron picks the CE function at construction.
 # Defaults match Megatron's native CE behavior; no CE-specific config needed.
 apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)
 ```

src/liger_kernel/megatron/monkey_patch.py

@@ -2,12 +2,32 @@
 
 from __future__ import annotations
 
+import importlib
 import logging
 
 logger = logging.getLogger(__name__)
 
 _PATCH_MARKER = "__liger_patched__"
 
+# Megatron modules that bind a cross-entropy function into their own namespace with
+# ``from ... import``. Replacing the function on its defining module does not reach a copy
+# that was bound before the patch ran, so these are rebound as well. Each entry exists
+# only on some Megatron versions; missing modules or attributes are skipped.
+_FUSED_CE_CALL_SITES = (
+    # Released megatron-core (0.19 and earlier): LanguageModule.compute_language_model_loss
+    # calls this module global.
+    ("megatron.core.models.common.language_module.language_module", "fused_vocab_parallel_cross_entropy"),
+    # Megatron-LM main (NVIDIA/Megatron-LM#6766): select_cross_entropy returns this module
+    # global, and LanguageModule caches the result when the model is built.
+    ("megatron.core.models.backends", "_fused_ce"),
+)
+_UNFUSED_CE_CALL_SITES = (
+    # Released megatron-core: compute_language_model_loss calls the package-level re-export
+    # ``tensor_parallel.vocab_parallel_cross_entropy``, which is bound when the package is
+    # imported, i.e. always before its ``cross_entropy`` submodule can be patched.
+    ("megatron.core.tensor_parallel", "vocab_parallel_cross_entropy"),
+)
+
 
 def apply_liger_kernel_to_megatron(
     rms_norm: bool = True,
@@ -31,7 +51,8 @@ def apply_liger_kernel_to_megatron(
             ``megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy``
             (fused path) and
             ``megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy``
-            (unfused path) with Liger's Triton cross-entropy. Default
+            (unfused path) with Liger's Triton cross-entropy, along with the
+            copies Megatron's call sites bound at import time. Default
             ``False`` so adopters opt in explicitly. All tensor-parallel sizes
             are supported (TP=1 and TP>1 share the same vocab-parallel
             kernel). The fused wrapper matches native's
@@ -158,6 +179,27 @@ def _patch_transformer_block_layernorm_impl() -> None:
 _LABEL_SMOOTHING_UNSET = object()
 
 
+def _rebind_call_sites(call_sites, original, replacement) -> bool:
+    """Point every call site still bound to ``original`` at ``replacement``.
+
+    Importing a call-site module here is harmless: the defining module is already patched,
+    so a first import binds ``replacement`` directly. A call site bound to some other
+    function (another framework's override) is left alone.
+
+    Returns whether any call site now holds ``replacement``.
+    """
+    reached = False
+    for module_name, attr in call_sites:
+        try:
+            module = importlib.import_module(module_name)
+        except ImportError:
+            continue
+        if getattr(module, attr, None) is original:
+            setattr(module, attr, replacement)
+        reached = reached or getattr(module, attr, None) is replacement
+    return reached
+
+
 def _patch_fused_vocab_parallel_cross_entropy() -> None:
     """Replace ``megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy``.
 
@@ -165,6 +207,11 @@ def _patch_fused_vocab_parallel_cross_entropy() -> None:
     that match Megatron's native fused-CE behavior) in a closure matching Megatron's fused-CE
     signature ``(logits, target, tp_group)``. Idempotent: a sentinel attribute on the
     replacement prevents wrappers from stacking.
+
+    Megatron never calls the function through this module attribute: each call site binds
+    its own copy with ``from ... import`` (see ``_FUSED_CE_CALL_SITES``). Those copies are
+    rebound too, so the patch works regardless of import order. Logs a warning when no
+    known call site ends up on Liger.
     """
     try:
         import megatron.core.fusions.fused_cross_entropy as fused_ce
@@ -197,10 +244,17 @@ def _patch_fused_vocab_parallel_cross_entropy() -> None:
     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
+    reached = _rebind_call_sites(_FUSED_CE_CALL_SITES, original, liger_fused_vocab_parallel_cross_entropy)
 
     logger.info(
         "Patched megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy with Liger cross-entropy."
     )
+    if not reached:
+        logger.warning(
+            "Could not find where this Megatron version calls fused_vocab_parallel_cross_entropy; "
+            "configs with cross_entropy_loss_fusion=True may keep using Megatron's kernel. Please file "
+            "an issue on https://github.com/linkedin/Liger-Kernel with your megatron-core version."
+        )
 
 
 def _patch_vocab_parallel_cross_entropy() -> None:
@@ -211,6 +265,10 @@ def _patch_vocab_parallel_cross_entropy() -> None:
     at call time, so the wrapper honors a runtime value when the caller actually passed
     one. A sentinel disambiguates "caller passed 0.0" (use 0.0) from "caller didn't pass"
     (use class default).
+
+    Released megatron-core calls the package-level re-export instead, which is rebound too
+    (see ``_UNFUSED_CE_CALL_SITES``). Megatron-LM main looks the function up on this module
+    at call time, so the module attribute alone reaches it.
     """
     try:
         import megatron.core.tensor_parallel.cross_entropy as unfused_ce
@@ -260,6 +318,7 @@ def _patch_vocab_parallel_cross_entropy() -> None:
     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
+    _rebind_call_sites(_UNFUSED_CE_CALL_SITES, original, liger_vocab_parallel_cross_entropy)
 
     logger.info(
         "Patched megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy with Liger cross-entropy."

test/megatron/test_monkey_patch.py

@@ -5,6 +5,7 @@ patch helpers can run entirely on CPU without a real megatron-core install. For
 Liger patches into Megatron, this file verifies:
 
 - the patched Megatron symbol(s) are actually replaced
+- call sites that bound their own copy before patching (CE) still reach Liger
 - patching is idempotent (calling apply twice doesn't stack wrappers)
 - the patch is a no-op when the kernel flag is False
 - missing megatron-core / missing symbol path raise helpful ``ImportError``\\s
@@ -23,6 +24,7 @@ Liger learns to patch:
   6. SwiGLU patch tests
 """
 
+import logging
 import sys
 import types
 
@@ -103,6 +105,68 @@ def _install_fake_megatron_ce(
     return fused_ce, unfused_ce
 
 
+def _install_fake_megatron_ce_call_sites(fused_ce, unfused_ce, layout: str):
+    """Install Megatron's CE call sites, bound the way real Megatron binds them.
+
+    Liger replaces the CE functions on their defining modules, but most Megatron call sites
+    take their own copy at import time. These stubs take that copy before any patching,
+    which is the usual order (model code is imported first).
+
+    - ``"released"`` (megatron-core 0.19 and earlier): ``language_module`` does
+      ``from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy``,
+      and the unfused branch calls the ``megatron.core.tensor_parallel`` package re-export.
+    - ``"main"`` (Megatron-LM main, NVIDIA/Megatron-LM#6766): ``backends`` imports the fused
+      function as ``_fused_ce``; ``unfused_cross_entropy`` looks the unfused one up on its
+      defining module at call time.
+
+    Returns ``compute_loss(logits, labels, cross_entropy_loss_fusion)``, a stand-in for
+    ``LanguageModule.compute_language_model_loss`` that resolves the CE function the way
+    that layout does.
+    """
+    _, megatron_core = _ensure_megatron_roots()
+    tensor_parallel = sys.modules["megatron.core.tensor_parallel"]
+    models = sys.modules.get("megatron.core.models") or types.ModuleType("megatron.core.models")
+    sys.modules["megatron.core.models"] = models
+    megatron_core.models = models
+
+    if layout == "released":
+        common = types.ModuleType("megatron.core.models.common")
+        language_module_pkg = types.ModuleType("megatron.core.models.common.language_module")
+        language_module = types.ModuleType("megatron.core.models.common.language_module.language_module")
+        language_module.fused_vocab_parallel_cross_entropy = fused_ce.fused_vocab_parallel_cross_entropy
+        tensor_parallel.vocab_parallel_cross_entropy = unfused_ce.vocab_parallel_cross_entropy
+
+        sys.modules["megatron.core.models.common"] = common
+        sys.modules["megatron.core.models.common.language_module"] = language_module_pkg
+        sys.modules["megatron.core.models.common.language_module.language_module"] = language_module
+        models.common = common
+        common.language_module = language_module_pkg
+        language_module_pkg.language_module = language_module
+
+        def compute_loss(logits, labels, cross_entropy_loss_fusion):
+            if cross_entropy_loss_fusion:
+                return language_module.fused_vocab_parallel_cross_entropy(logits, labels, None)
+            return tensor_parallel.vocab_parallel_cross_entropy(logits, labels)
+
+        return compute_loss
+
+    assert layout == "main", layout
+    backends = sys.modules.get("megatron.core.models.backends") or types.ModuleType("megatron.core.models.backends")
+    backends._fused_ce = fused_ce.fused_vocab_parallel_cross_entropy
+    sys.modules["megatron.core.models.backends"] = backends
+    models.backends = backends
+
+    def unfused_cross_entropy(logits, labels, tp_group=None):
+        return unfused_ce.vocab_parallel_cross_entropy(logits, labels, tp_group=tp_group)
+
+    def compute_loss(logits, labels, cross_entropy_loss_fusion):
+        # LanguageModule.__init__ runs select_cross_entropy once; the model is built after patching.
+        ce = backends._fused_ce if cross_entropy_loss_fusion else unfused_cross_entropy
+        return ce(logits, labels, None)
+
+    return compute_loss
+
+
 def _install_fake_megatron_rms_norm(
     layer_norm_is_wrapped_torch_norm: bool = True,
     with_backends_module: bool = True,
@@ -195,6 +259,9 @@ def _uninstall_fake_megatron():
         # CE side
         "megatron.core.parallel_state",
         "megatron.core.fusions.fused_cross_entropy",
+        "megatron.core.models.common.language_module.language_module",
+        "megatron.core.models.common.language_module",
+        "megatron.core.models.common",
         # RMSNorm side
         "megatron.core.models.backends",
         "megatron.core.models",
@@ -709,6 +776,62 @@ def test_unfused_wrapper_honors_explicit_zero_label_smoothing(fake_megatron_ce):
     )
 
 
+# ---------------------------------------------------------------------------
+# 2.5 CE call sites that bound their own copy before the patch ran.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("layout", ["released", "main"])
+@pytest.mark.parametrize("cross_entropy_loss_fusion", [True, False])
+def test_patch_reaches_call_sites_bound_before_patching(fake_megatron_ce, caplog, layout, cross_entropy_loss_fusion):
+    """Megatron model code is normally imported before ``apply_liger_kernel_to_megatron``
+    runs, so every call site already holds the original function. The loss computation must
+    still land in Liger; the stub originals raise if it doesn't."""
+    import torch
+
+    fused_ce, unfused_ce = fake_megatron_ce
+    compute_loss = _install_fake_megatron_ce_call_sites(fused_ce, unfused_ce, layout)
+    from liger_kernel.megatron import apply_liger_kernel_to_megatron
+    from liger_kernel.megatron import cross_entropy as ce_mod
+
+    calls = []
+
+    class _FakeCE:
+        def __init__(self, ignore_index=-100, label_smoothing=0.0, reduction="none"):
+            pass
+
+        def __call__(self, logits, target, tp_group=None):
+            calls.append(tuple(logits.shape))
+            return torch.zeros(logits.shape[:2])
+
+    with patch.object(ce_mod, "LigerMegatronCrossEntropy", _FakeCE):
+        apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=True)
+        compute_loss(torch.zeros(2, 1, 4), torch.zeros(2, 1, dtype=torch.long), cross_entropy_loss_fusion)
+
+    assert calls == [(2, 1, 4)]
+    assert "Could not find where this Megatron version calls" not in caplog.text
+
+
+def test_patch_leaves_call_site_bound_to_another_function_alone(fake_megatron_ce, caplog):
+    """A call site that another framework already pointed at its own kernel is not Liger's to
+    take over. The patch leaves it in place and warns that the fused path was not reached."""
+    fused_ce, unfused_ce = fake_megatron_ce
+    _install_fake_megatron_ce_call_sites(fused_ce, unfused_ce, "released")
+    language_module = sys.modules["megatron.core.models.common.language_module.language_module"]
+
+    def framework_fused_ce(vocab_parallel_logits, target, tp_group=None):
+        raise AssertionError("not expected to be called")
+
+    language_module.fused_vocab_parallel_cross_entropy = framework_fused_ce
+    from liger_kernel.megatron import apply_liger_kernel_to_megatron
+
+    with caplog.at_level(logging.WARNING, logger="liger_kernel.megatron.monkey_patch"):
+        apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=True)
+
+    assert language_module.fused_vocab_parallel_cross_entropy is framework_fused_ce
+    assert "Could not find where this Megatron version calls fused_vocab_parallel_cross_entropy" in caplog.text
+
+
 # ===========================================================================
 # 3. RMSNorm patch tests
 # ===========================================================================