[Megatron] Add SwiGLU integration

benchmark/scripts/benchmark_megatron_swiglu.py

@@ -0,0 +1,186 @@
+"""Benchmark Liger's Megatron-LM SwiGLU wrapper.
+
+Compares four providers on the fused gate-up activation call shape
+``[seq, batch, 2 * ffn_local]`` -- exactly the tensor Megatron's ``linear_fc1`` hands to
+``bias_swiglu_impl``:
+
+  - **liger**: ``LigerMegatronSwiGLU`` -- Liger's fused gate-up Triton kernel in the
+    Megatron-shaped wrapper. This is the default configuration.
+  - **liger_in_place**: the same, with ``in_place=True``, which writes the backward
+    gradient into the fc1 output buffer rather than allocating a new one. Same speed,
+    one fewer activation-sized allocation; opt-in because it destroys that buffer.
+  - **megatron**: Megatron's ``bias_swiglu_impl`` (``bias_activation_fusion=True``), a
+    chain of ``@jit_fuser`` TorchScript helpers. This is the symbol Liger displaces.
+  - **torch**: eager ``F.silu(y_1) * y_2`` over ``torch.chunk(y, 2, -1)`` -- the
+    unfused reference, and structurally what Megatron runs when
+    ``bias_activation_fusion=False`` (there it's a closure inside ``MLP.forward``).
+
+Why there is no ``--tp-size`` flag (unlike the Megatron CE benchmark): SwiGLU is
+elementwise and token-local. It issues **no collectives**, and tensor parallelism affects
+it only by shrinking the per-rank column count to ``ffn_hidden_size / tp``. Sweeping
+``ffn_local`` on a single GPU therefore already covers every TP configuration -- TP=8 at
+``ffn=28672`` is the same kernel work as the ``ffn_local=3584`` point on this curve. Peak
+memory is likewise per-rank and scales as 1/TP.
+
+The x-axis spans the Blackwell tiling threshold: ``liger_kernel.ops.swiglu`` switches to a
+column-tiled 2D grid when ``next_pow2(n_cols) >= 16384`` on Blackwell, so the two largest
+points exercise that path on B200 and the one-row path everywhere else.
+
+Requires a Liger-supported accelerator (CUDA / ROCm). With megatron-core not installed the
+``megatron`` provider is silently dropped and the run proceeds with ``liger`` + ``torch``.
+
+Output goes to the shared ``benchmark/data/all_benchmark_data.csv`` -- rows are tagged with
+``kernel_name="megatron_swiglu"`` and the standard visualizer renders them via:
+
+    python benchmark/benchmarks_visualizer.py \\
+        --kernel-name megatron_swiglu --metric-name speed
+    python benchmark/benchmarks_visualizer.py \\
+        --kernel-name megatron_swiglu --metric-name memory
+"""
+
+import torch
+import torch.nn.functional as F
+import triton
+
+from utils import QUANTILES
+from utils import SingleBenchmarkRunInput
+from utils import SingleBenchmarkRunOutput
+from utils import _test_memory
+from utils import parse_benchmark_script_args
+from utils import run_benchmarks
+
+from liger_kernel.megatron import LigerMegatronSwiGLU
+from liger_kernel.utils import infer_device
+
+device = infer_device()
+
+try:
+    from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl
+
+    _MEGATRON_AVAILABLE = True
+except ImportError:
+    bias_swiglu_impl = None
+    _MEGATRON_AVAILABLE = False
+
+
+def _torch_swiglu(y):
+    """Eager reference — identical math to Megatron's ``swiglu``, minus the JIT fusion."""
+    y_1, y_2 = torch.chunk(y, 2, -1)
+    return F.silu(y_1) * y_2
+
+
+def _make_fwd(provider: str):
+    if provider == "liger":
+        module = LigerMegatronSwiGLU()
+        return lambda y: module(y, None, False, False)
+    if provider == "liger_in_place":
+        module = LigerMegatronSwiGLU(in_place=True)
+        return lambda y: module(y, None, False, False)
+    if provider == "torch":
+        return _torch_swiglu
+    if provider == "megatron":
+        if not _MEGATRON_AVAILABLE:
+            raise RuntimeError("megatron-core not installed; cannot benchmark 'megatron' provider")
+        return lambda y: bias_swiglu_impl(y, None, False, False)
+    raise ValueError(f"unknown provider: {provider!r}")
+
+
+def _make_input(s: int, b: int, ffn_local: int, requires_grad: bool = True) -> torch.Tensor:
+    # 2 * ffn_local: Megatron's linear_fc1 emits gate and up concatenated on the last dim.
+    return torch.randn(s, b, 2 * ffn_local, device=device, dtype=torch.bfloat16, requires_grad=requires_grad)
+
+
+def bench_speed_megatron_swiglu(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    ffn_local = input.x
+    provider = input.kernel_provider
+    mode = input.kernel_operation_mode
+    s = input.extra_benchmark_config["S"]
+    b = input.extra_benchmark_config["B"]
+
+    fwd_fn = _make_fwd(provider)
+    x = _make_input(s, b, ffn_local)
+
+    def fwd():
+        return fwd_fn(x)
+
+    if mode == "forward":
+        ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES)
+    elif mode == "backward":
+        # Rerun fwd each iteration: Liger's in-place backward consumes the saved buffers,
+        # so a retained graph would corrupt on the second pass. Subtract the "forward"
+        # row to get backward-only timing.
+        def _fwd_bwd():
+            if x.grad is not None:
+                x.grad = None
+            out = fwd()
+            out.sum().backward()
+
+        ms_50, ms_20, ms_80 = triton.testing.do_bench(_fwd_bwd, rep=100, quantiles=QUANTILES)
+    elif mode == "full":
+
+        def full():
+            if x.grad is not None:
+                x.grad = None
+            y = fwd()
+            y.sum().backward()
+
+        ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES)
+    else:
+        raise ValueError(f"unknown mode: {mode!r}")
+
+    return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80)
+
+
+def bench_memory_megatron_swiglu(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    ffn_local = input.x
+    provider = input.kernel_provider
+    s = input.extra_benchmark_config["S"]
+    b = input.extra_benchmark_config["B"]
+
+    fwd_fn = _make_fwd(provider)
+    x = _make_input(s, b, ffn_local)
+
+    def full():
+        if x.grad is not None:
+            x.grad = None
+        y = fwd_fn(x)
+        y.sum().backward()
+
+    mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES)
+    return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80)
+
+
+if __name__ == "__main__":
+    args = parse_benchmark_script_args()
+
+    providers = ["liger", "liger_in_place", "torch"]
+    if _MEGATRON_AVAILABLE:
+        providers.append("megatron")
+
+    common_configs = {
+        "kernel_name": "megatron_swiglu",
+        "x_name": "ffn_local",
+        "x_label": "per-rank FFN hidden size",
+        # 1024 → 32768. Llama-7B is 11008 and Llama-70B is 28672, so this brackets
+        # production sizes; the top two points cross the Blackwell tiling threshold.
+        "x_values": [2**i for i in range(10, 16)],
+        "kernel_providers": providers,
+        # Megatron's standard training shape, matching the megatron CE benchmark.
+        "extra_benchmark_configs": [{"S": 2048, "B": 4}],
+        "overwrite": args.overwrite,
+    }
+
+    run_benchmarks(
+        bench_test_fn=bench_speed_megatron_swiglu,
+        kernel_operation_modes=["forward", "backward", "full"],
+        metric_name="speed",
+        metric_unit="ms",
+        **common_configs,
+    )
+    run_benchmarks(
+        bench_test_fn=bench_memory_megatron_swiglu,
+        kernel_operation_modes=["full"],
+        metric_name="memory",
+        metric_unit="MB",
+        **common_configs,
+    )

examples/megatron/README.md

@@ -2,21 +2,39 @@
 
 Two self-contained scripts demonstrating the integration modes shipped by
 `liger_kernel.megatron`. Both train a tiny GPT model with mock data on
-2 × GPU (TP=2, PP=1) for 5 iterations and print the resolved norm classes
+2 × GPU (TP=1, PP=1, DP=2) for 5 iterations and print the resolved bindings
 so you can see which slots picked up Liger.
 
+## Op support matrix
+
+| Op | Mode 1 flag | Patched symbol(s) | Mode 2 class | Slot used in Mode 2 |
+|---|---|---|---|---|
+| RMSNorm | `rms_norm=True` (on by default) | `LocalSpecProvider.layer_norm`, `transformer_block.LayerNormImpl` | `LigerMegatronRMSNorm` | every norm slot, incl. block-level `final_layernorm` |
+| Cross-entropy | `cross_entropy=True` (opt-in) | `fused_cross_entropy.fused_vocab_parallel_cross_entropy`, `tensor_parallel.cross_entropy.vocab_parallel_cross_entropy` | `LigerMegatronCrossEntropy` | none — `GPTModel` subclass overriding `compute_language_model_loss` |
+| SwiGLU | `swiglu=True` (opt-in) | `fusions.fused_bias_swiglu.SwiGLUFunction` | `LigerMegatronSwiGLU` | the `mlp` module slot — an `MLP` subclass |
+
+Notes that apply to the table:
+
+- RMSNorm only covers the local (non-TE) backend.
+- SwiGLU is used only when `gated_linear_unit=True`, `activation_func=F.silu`,
+  and `bias_activation_fusion=True`; otherwise the patch is applied but not
+  exercised.
+- Liger replaces only `SwiGLUFunction`, so bias and MoE variants stay on
+  Megatron.
+- Cross-entropy and SwiGLU are wired through subclasses (no dedicated spec
+  slot).
+
 ## Prerequisites
 
 - A working Megatron-Core install (`pip install megatron-core`).
 - `liger-kernel` installed (editable or from PyPI).
-- `psutil` (used by Megatron's async checkpoint worker pool).
+- `psutil`.
 - At least 2 GPUs.
 
 ## Mode 1 — `apply_liger_kernel_to_megatron()` (monkey-patch)
 
-One-line opt-in. Patches `LocalSpecProvider.layer_norm` and
-`transformer_block.LayerNormImpl` so every RMSNorm slot becomes Liger
-without changing the spec the user constructs.
+One-line opt-in. Patches the symbols in the matrix above without changing
+the spec the user constructs.
 
 ```bash
 torchrun --nproc_per_node=2 \
@@ -26,10 +44,9 @@ torchrun --nproc_per_node=2 \
 
 ## Mode 2 — hand-assembled `TransformerBlockSubmodules`
 
-Slot-level control. Explicitly places `LigerMegatronRMSNorm` into each
-norm slot, including the block-level `final_layernorm`. Useful when you
-want to mix Liger with other backends (e.g. TransformerEngine) on a
-per-slot basis.
+Slot-level control — see the "Mode 2 class" column of the matrix above.
+Useful when you want to mix Liger with other backends (e.g.
+TransformerEngine) on a per-slot basis.
 
 ```bash
 torchrun --nproc_per_node=2 \
@@ -41,13 +58,20 @@ torchrun --nproc_per_node=2 \
 
 For both scripts:
 
-- 5 lines of `[modeN] iter <i> loss=<float>` with the loss decreasing.
+- 5 lines of `[modeN] iter <i> loss=<float>`. (Five iterations on a
+  12-hidden-size model with random mock data is far too short to show a
+  trend — the value just hovers.)
 - A printed module tree with `LigerMegatronRMSNorm` in **5 of 5** norm
   slots: four per-layer (`input_layernorm`, `pre_mlp_layernorm` × 2 layers)
   and one block-level (`final_layernorm`).
 - `Successfully loaded the model` after the distributed checkpoint
   round-trip.
 
+Mode 1 additionally prints `=== Resolved SwiGLU symbols ===` with all three
+bindings tagged `[Liger]` — the defining module plus the two consumers that
+import the symbol by name. Mode 2 prints `=== Resolved SwiGLU MLPs ===`
+listing both `_LigerSwiGLUMLP` layers.
+
 If your environment doesn't have Apex or TransformerEngine installed, you
 will see harmless warnings — Megatron falls back to the local backend,
 which is exactly where Liger plugs in.

examples/megatron/run_mode1_monkey_patch.py

@@ -1,24 +1,34 @@
-"""Mode 1 — monkey-patch Megatron-Core to use Liger RMSNorm + cross-entropy.
+"""Mode 1 — monkey-patch Megatron-Core to use Liger RMSNorm + cross-entropy + SwiGLU.
 
 Adapted from Megatron's ``examples/run_simple_mcore_train_loop.py``. The
 relevant additions (vs. that file) are:
 
-  1. ``apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)``
-     called once at the top of ``model_provider()``. This patches:
+  1. ``apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True,
+     swiglu=True)`` called once at the top of ``model_provider()``. This patches:
        - ``LocalSpecProvider.layer_norm`` (per-layer norm slots)
        - ``transformer_block.LayerNormImpl`` (block-level ``final_layernorm``)
        - ``fused_cross_entropy.fused_vocab_parallel_cross_entropy``
          (the fused CE path)
        - ``tensor_parallel.cross_entropy.vocab_parallel_cross_entropy``
          (the unfused CE path)
+       - ``fusions.fused_bias_swiglu.SwiGLUFunction`` (the SwiGLU
+         activation, reached via Megatron's own ``bias_swiglu_impl``)
 
   2. ``normalization="RMSNorm"`` added to ``TransformerConfig`` so the
      model actually has RMSNorm slots to patch (Megatron defaults to
      ``LayerNorm``).
 
-  3. ``_print_norm_classes`` + ``_print_ce_symbols`` after model construction
-     — print the resolved class/function bindings so you can verify Liger
-     took over for every slot.
+  3. The SwiGLU dispatch flags added to ``TransformerConfig``
+     (``gated_linear_unit``, ``activation_func=F.silu``,
+     ``bias_activation_fusion``). ``MLP.forward`` only routes through
+     ``bias_swiglu_impl`` when all three hold, so without them the patch is
+     applied but never reached. ``add_bias_linear=False`` keeps ``bias`` at
+     ``None``, which is the configuration Liger accelerates; with a bias the
+     wrapper transparently falls back to Megatron's own kernel.
+
+  4. ``_print_norm_classes`` + ``_print_ce_symbols`` + ``_print_swiglu_symbols``
+     after model construction — print the resolved class/function bindings so
+     you can verify Liger took over for every slot.
 
 Run with:
     torchrun --nproc_per_node=2 --master_addr=127.0.0.1 --master_port=29500 \\
@@ -34,6 +44,7 @@
 from typing import Iterator
 
 import torch
+import torch.nn.functional as F
 
 from megatron.core import dist_checkpointing
 from megatron.core import parallel_state
@@ -74,7 +85,7 @@ def initialize_distributed(tp: int = 2, pp: int = 1) -> None:
 
 def model_provider() -> GPTModel:
     # ↓↓ Mode 1 — patch once, everything below picks up Liger ↓↓
-    apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)
+    apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True, swiglu=True)
     # ↑↑ ------------------------------------------------------ ↑↑
 
     cfg = TransformerConfig(
@@ -84,6 +95,10 @@ def model_provider() -> GPTModel:
         use_cpu_initialization=True,
         pipeline_dtype=torch.float32,
         normalization="RMSNorm",
+        gated_linear_unit=True,
+        activation_func=F.silu,
+        bias_activation_fusion=True,
+        add_bias_linear=False,
     )
     return GPTModel(
         config=cfg,
@@ -157,6 +172,31 @@ def _print_ce_symbols() -> None:
     print()
 
 
+def _print_swiglu_symbols() -> None:
+    """Show which ``SwiGLUFunction`` the dense MLP and shared experts end up running.
+
+    Liger replaces the class, not ``bias_swiglu_impl``. That function is a plain
+    dispatcher that looks the class up in its own module globals on every call,
+    so the two consumer modules keep their import-time ``bias_swiglu_impl``
+    binding and still route to Liger. Both facts are printed below.
+    """
+    import megatron.core.fusions.fused_bias_swiglu as defining
+    import megatron.core.transformer.mlp as mlp
+    import megatron.core.transformer.moe.shared_experts as shared
+
+    tag = "Liger" if getattr(defining.SwiGLUFunction, "__liger_patched__", False) else "Megatron"
+    print("\n=== Resolved SwiGLU symbols ===")
+    print(f"  {'fusions.fused_bias_swiglu.SwiGLUFunction':52s} \u2192  [{tag}]")
+    for label, mod in (
+        ("transformer.mlp", mlp),
+        ("transformer.moe.shared_experts", shared),
+    ):
+        same = mod.bias_swiglu_impl is defining.bias_swiglu_impl
+        note = "unpatched, resolves the class above" if same else "REBOUND \u2014 unexpected"
+        print(f"  {label + '.bias_swiglu_impl':52s} \u2192  {note}")
+    print()
+
+
 def main() -> None:
     # TP=1, DP=2 — CE patch (TP=1 only). Norms are correct under any TP value, so
     # demonstrating both Liger features in one script means running data-parallel.
@@ -171,6 +211,7 @@ def main() -> None:
         print(gpt_model)
         _print_norm_classes(gpt_model)
         _print_ce_symbols()
+        _print_swiglu_symbols()
 
     ddp_cfg = DistributedDataParallelConfig(
         grad_reduce_in_fp32=False,