diff --git a/docs/High-Level-APIs.md b/docs/High-Level-APIs.md index 6bbe008a..d200419d 100644 --- a/docs/High-Level-APIs.md +++ b/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) ``` diff --git a/src/liger_kernel/megatron/monkey_patch.py b/src/liger_kernel/megatron/monkey_patch.py index e97567ca..5e900ed8 100644 --- a/src/liger_kernel/megatron/monkey_patch.py +++ b/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." diff --git a/test/megatron/test_monkey_patch.py b/test/megatron/test_monkey_patch.py index 634a4abd..f2c830a1 100644 --- a/test/megatron/test_monkey_patch.py +++ b/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 # ===========================================================================