[Megatron] Add SwiGLU integration by buffett0323 · Pull Request #1326 · linkedin/Liger-Kernel

buffett0323 · 2026-07-27T03:36:01Z

Summary

Adds Liger's Triton SwiGLU to Megatron-Core. Megatron's dense MLP routes its activation through bias_swiglu_impl in megatron.core.fusions.fused_bias_swiglu, a TorchScript @jit_fuser implementation. This PR replaces it with a new fused gate-up Triton kernel written for Megatron's memory layout.

At the Megatron training shape (S=2048, B=4, bf16) on H100, against megatron-core==0.16.1:

	Forward	Full (fwd+bwd)	Memory	Memory (`in_place=True`)
`ffn_local=4096`	+57%	+146%	+20%	−20%
`ffn_local=32768`	+63%	+166%	+20%	−20%

Background

Liger's existing LigerSiLUMulFunction takes gate and up as two separate tensors (a, b), because that is how HuggingFace models arrive: gate_proj(x) and up_proj(x) are distinct GEMMs.

Megatron is the opposite. linear_fc1 is a single fused ColumnParallelLinear producing one [s, b, 2 * ffn_local] tensor with gate in columns [0, n) and up in [n, 2n). Bridging to the two-tensor kernel therefore requires a torch.chunk, and chunk on the last dimension returns non-contiguous views, so @ensure_contiguous materializes 2 full copies on every forward. We measured that bridge before writing anything new.

H100, forward, `ffn_local=32768`	Time	vs Megatron
Megatron `bias_swiglu_impl`	0.853 ms	—
Liger via `chunk` bridge	1.828 ms	−53% (2.1× slower)
Liger fused gate-up (this PR)	0.536 ms	+59%

The chunk bridge is 3.4× slower than the fused kernel and also inflated memory by +40%. The new LigerFusedGateUpSiLUMulFunction reads both halves via a column offset into the single buffer, with no copies, no cat. Input row stride is 2n, output row stride is n.

How to use

Mode 1: One-line monkey patch

from liger_kernel.megatron import apply_liger_kernel_to_megatron

apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True, swiglu=True) # swiglu defaults to False

Patches bias_swiglu_impl across Megatron's core modules (fused_bias_swiglu, mlp, and moe.shared_experts).

Mode 2: Use the class directly

from liger_kernel.megatron import LigerMegatronSwiGLU

swiglu = LigerMegatronSwiGLU()
out = swiglu(intermediate, bias, fp8_input_store, cpu_offload_input)

The signature mirrors bias_swiglu_impl positionally. To use Mode 2, subclass MLP (see examples/megatron/run_mode2_hand_spec.py).

Optional: `in_place=True` for a 20% memory reduction

LigerMegatronSwiGLU(in_place=True) # Opt-in, off by default

Overwrites the linear_fc1 output buffer during backward passes to reduce memory by 20% vs Megatron with zero speed penalty.

Scope

Supported: Dense MLP (bias_swiglu_impl) and shared experts at any TP size with bias=None.
Transparent Fallbacks to Megatron: Auto-delegates to Megatron (and logs once) when:
- bias is not None (Liger lacks bias support)
- fp8_input_store=True (requires FP8 input retention)
- cpu_offload_input=True (requires offloadable inputs)
Pass fallback_impl=None to raise an error instead of falling back.
Out of Scope: MoE routed experts (weighted_bias_swiglu_impl).

What's added

File	Purpose
`src/liger_kernel/ops/swiglu.py`	New `LigerFusedGateUpSiLUMulFunction` + fwd/bwd Triton kernels for the `[tokens, 2n]` layout. Existing two-tensor path untouched.
`src/liger_kernel/megatron/swiglu.py`	`LigerMegatronSwiGLU` nn.Module, fallback policy, shape validation.
`src/liger_kernel/megatron/monkey_patch.py`	`swiglu=` flag, `_patch_bias_swiglu_impl()`, `_rebind_consumer_symbol()`.
`src/liger_kernel/megatron/__init__.py`	Re-exports `LigerMegatronSwiGLU`.
`test/megatron/test_swiglu.py`	36 tests — 28 dependency-free + 8 against real `megatron-core`.
`test/megatron/test_monkey_patch.py`	+14 SwiGLU patch-mechanism tests; stub-harness snapshot/restore rewrite.
`benchmark/scripts/benchmark_megatron_swiglu.py`	4 providers: `liger`, `liger_in_place`, `megatron`, `torch`.
`examples/megatron/*`	Mode 1 + Mode 2 extended to cover SwiGLU.

Megatron feature compatibility

Verified: Tensor Parallelism (TP), fp32 + bf16, 2D and 3D inputs, distributed checkpointing (save/load round-trip in both example scripts), real MLP forward+backward parity against Megatron's TorchScript, and both consumer module paths.

Should work, not yet E2E-verified: multi-rank patch application (every rank must call apply_... before model build — single-process patching is verified, the multi-rank ordering is not), CP, use_distributed_optimizer=True.

Out of scope: activation recompute and CUDA-graph capture with in_place=True — precisely why that flag is opt-in. The default path is recompute-safe.

Known issues / follow-ups

in_place=True and retain_graph=True incompatibility: In-place Triton kernel writes don't bump PyTorch's _version counter, so running backward a second time silently yields garbage gradients without triggering autograd guards. We added a ctx.already_backward flag to explicitly raise an error when reusing the graph with in_place=True.
MoE routed experts: weighted_bias_swiglu_impl requires a per_token_scale-aware kernel (left as a follow-up).
RMSNorm patch fix: Fixed a TypeError in _patch_local_spec_provider_layer_norm caused by forwarding an unsupported has_residual kwarg on megatron-core 0.16.1. Signature matching is now version-agnostic via inspect.signature.
Example script fixes: Fixed three pre-existing megatron-core 0.16.1 compatibility issues in the example scripts (AttributeError on as_mlp_submodule, Makefile build errors in pip wheels, and invalid tokenizer metadata).

Testing Done

Hardware Type: NVIDIA A10G (correctness, 1× and 2×), H100 80GB HBM3 (benchmarks), B200 (benchmarks)
pytest test/megatron/ with megatron-core installed — 175 passed, 49 skipped
pytest test/megatron/ without megatron-core (default CI shape) — 168 passed, 56 skipped
pytest test/transformers/test_swiglu.py — 49 passed, 4 skipped, 4 xfailed (existing two-tensor SwiGLU not regressed by the shared ops/swiglu.py change)
run make checkstyle to ensure code style — clean; prek run -a also clean
H100 + B200 benchmark sweeps; H100 rows committed to all_benchmark_data.csv
run make test-convergence to ensure convergence

Benchmark reproduction

python benchmark/scripts/benchmark_megatron_swiglu.py --overwrite

python benchmark/benchmarks_visualizer.py \
    --kernel-name megatron_swiglu --metric-name speed --overwrite

Tcc0403

Thanks for the contribution, the kernel impl looks good.

However, there are plenty of new added functions to resolve environment issue with megatron-core installed in our test suite. I prefer moving them to another PR, ideally a proposal issue first, to address such issues.

Tcc0403 · 2026-07-28T15:26:23Z

+
+@triton.jit
+def _swiglu_fused_gate_up_backward_kernel(
+    dc_ptr, y_ptr, dy_ptr, in_stride, out_stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr


Suggested change

dc_ptr, y_ptr, dy_ptr, in_stride, out_stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr

dc_ptr, y_ptr, dy_ptr, in_stride, out_stride, ffn_size: tl.constexpr, BLOCK_SIZE: tl.constexpr

nit: name it ffn_size for readability

Tcc0403 · 2026-07-28T15:26:35Z

+    two_n = ori_shape[-1]
+    if two_n % 2 != 0:
+        raise ValueError(f"fused gate-up input must have an even trailing dim; got {two_n}.")
+    n_cols = two_n // 2


Tcc0403 · 2026-07-28T15:32:38Z

-2 × GPU (TP=2, PP=1) for 5 iterations and print the resolved norm classes
-so you can see which slots picked up Liger.
+2 × GPU (TP=1, PP=1, DP=2) for 5 iterations and print the resolved norm
+classes, CE bindings and SwiGLU bindings so you can see which slots picked


it would be better to create a op support matrix and append new op to it instead of inline new changes

Tcc0403 · 2026-07-28T15:36:40Z

    )


+def _ensure_dataset_helpers() -> None:


out of this PR scope? is it a required change to make this script work?

Tcc0403 · 2026-07-28T15:37:14Z

    )


+def _ensure_dataset_helpers() -> None:


Tcc0403 · 2026-07-28T15:47:57Z

+    )
+
+
+def _rebind_consumer_symbol(


I don't fully understand the purpose of this function. In what scenario do we need this rebind function, could you elaborate further?

Should we also apply it to ce and norm?

Tcc0403 · 2026-07-28T15:52:56Z

+# Force-import the submodule so liger_kernel.ops.swiglu can resolve
+# torch.distributed.tensor.DTensor on torch 2.11+, where the subpackage is no longer
+# auto-loaded as an attribute of torch.distributed.
+import torch.distributed.tensor  # noqa: F401  # isort: skip
+
+from liger_kernel.ops import LigerFusedGateUpSiLUMulFunction


would from liger_kernel.ops import LigerFusedGateUpSiLUMulFunction fail without importing torch.distributed.tensor beforehand?

Tcc0403 · 2026-07-28T16:04:04Z

+_SAVED_REAL_MEGATRON: dict = {}
+_STUBS_ACTIVE = False
+
+
+def _is_megatron_module(name: str) -> bool:
+    return name == "megatron" or name.startswith("megatron.")
+
+
+def _snapshot_real_megatron():


is there any conflict with stub megatron when megatron-core is installed? would the test fail without these functions?

Tcc0403 · 2026-07-28T16:19:30Z

+# Unlike CE and RMSNorm, its consumers use by-name imports, so the patch has to rewrite
+# their module attributes too; several tests below exist specifically to enforce that.


Follow CE and RMSNorm tests convention for self-consistent and maintainability. Issues in megatron-installed environment should be landed in another PR

Tcc0403 · 2026-07-28T16:27:24Z

+# ---------------------------------------------------------------------------
+# 6. End-to-end validation against a real megatron-core install
+# ---------------------------------------------------------------------------
+#
+# Everything above verifies the wrapper against our own reproduction of Megatron's
+# formula. That cannot catch a wrong belief about Megatron's API. The tests below import
+# the real package and pin: the true signature and positional order of
+# ``bias_swiglu_impl``, numerical parity with Megatron's own TorchScript implementation,
+# that the consumer module paths the monkey patch rebinds actually exist and are rebound,
+# and that a real ``MLP`` produces identical output and gradients when patched.
+#
+# The import is conditional rather than a module-level ``pytest.importorskip`` so that the
+# dependency-free tests above still run when megatron-core is absent.


Although this section is skipped when megatron-core is not installed, I still suggest removing it to keep this PR concise. Feel free to raise an issue if you plan to handle actual megatron-core environment in the test suite.

buffett0323 · 2026-07-29T07:13:03Z

Hi @Tcc0403 , thanks for reviewing. The comments are pretty thorough, and I updated the code.

A big change is instead of patching bias_swiglu_impl and then rebinding it in a hardcoded list of consumer modules, I now patch fused_bias_swiglu.SwiGLUFunction directly.

bias_swiglu_impl is what consumers import, and they import it by name (transformer/mlp.py, moe/shared_experts.py), so reassigning it in the defining module doesn't reach them. so the old rebind list, which would possibly miss any new consumer in a future release.

SwiGLUFunction doesn't have that problem. Instead, bias_swiglu_impl resolves it from its own module globals on every call, and it's referenced nowhere else in megatron-core. One binding, no stale copies, so the patch reaches every consumer regardless of import order.

Tcc0403 reviewed Jul 28, 2026

View reviewed changes

buffett0323 added 2 commits July 28, 2026 23:27

[Megatron] Add SwiGLU integration

6bf2e5a

Fix and revision

9bfdbff

buffett0323 force-pushed the megatron-swiglu branch from 6f87072 to 9bfdbff Compare July 29, 2026 06:28

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Megatron] Add SwiGLU integration - #1326

[Megatron] Add SwiGLU integration#1326
buffett0323 wants to merge 2 commits into
linkedin:mainfrom
buffett0323:megatron-swiglu

buffett0323 commented Jul 27, 2026

Uh oh!

Tcc0403 left a comment

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

Tcc0403 Jul 28, 2026

Uh oh!

buffett0323 commented Jul 29, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	dc_ptr, y_ptr, dy_ptr, in_stride, out_stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
	dc_ptr, y_ptr, dy_ptr, in_stride, out_stride, ffn_size: tl.constexpr, BLOCK_SIZE: tl.constexpr

		# Unlike CE and RMSNorm, its consumers use by-name imports, so the patch has to rewrite
		# their module attributes too; several tests below exist specifically to enforce that.

		)


		def _rebind_consumer_symbol(

Uh oh!

Conversation

buffett0323 commented Jul 27, 2026

Summary

Background

How to use

Mode 1: One-line monkey patch

Mode 2: Use the class directly

Optional: in_place=True for a 20% memory reduction

Scope

What's added

Megatron feature compatibility

Known issues / follow-ups

Testing Done

Benchmark reproduction

Uh oh!

Tcc0403 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

buffett0323 commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Optional: `in_place=True` for a 20% memory reduction

buffett0323 commented Jul 29, 2026 •

edited

Loading