mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Buffett LiuandVaibhav Jindal
2120862ba5
## 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%** | <img width="1680" height="1200" alt="megatron_swiglu_panel_H100" src="https://github.com/user-attachments/assets/0c0a2710-42df-4354-a0ce-bd7a25837058" /> ### 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 ```python 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 ```python 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 ```python 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 1. `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`. 2. MoE routed experts: `weighted_bias_swiglu_impl` requires a `per_token_scale`-aware kernel (left as a follow-up). 3. 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`. 4. 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) - [x] `pytest test/megatron/` **with** `megatron-core` installed — **175 passed, 49 skipped** - [x] `pytest test/megatron/` **without** `megatron-core` (default CI shape) — **168 passed, 56 skipped** - [x] `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) - [x] run `make checkstyle` to ensure code style — clean; `prek run -a` also clean - [x] H100 + B200 benchmark sweeps; H100 rows committed to `all_benchmark_data.csv` - [ ] run `make test-convergence` to ensure convergence ## Benchmark reproduction ```bash python benchmark/scripts/benchmark_megatron_swiglu.py --overwrite python benchmark/benchmarks_visualizer.py \ --kernel-name megatron_swiglu --metric-name speed --overwrite ``` --------- Co-authored-by: Vaibhav Jindal <[email protected]>