mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00

Files

T

0400d31357 [Megatron] Add cross-entropy integration (#1207 )

## Summary

This PR adds Liger's Triton cross-entropy kernel to Megatron-LM.

Megatron has two cross-entropy code paths. This PR replaces both with
Liger:

- The **fused** path: `fused_vocab_parallel_cross_entropy` (used when
`cross_entropy_loss_fusion=True`).
- The **unfused** path: `vocab_parallel_cross_entropy` (used when
`cross_entropy_loss_fusion=False`).

You get Liger's kernel no matter which path your Megatron config picks.

Liger's kernel uses online softmax. It writes gradients in place. It
never builds the full softmax tensor. The net effect is faster training
and lower memory use.

This PR also keeps the RMSNorm patch that was added in PR #1254. So one
call patches both kernels at once.

## How to use

There are two ways. Pick whichever fits your code better.

### Mode 1: One-line monkey patch

Best for users who don't want to touch their model code. Call this once
before you build the model:

```python
from liger_kernel.megatron import apply_liger_kernel_to_megatron

apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)
```

That's the whole API. No CE-specific knobs. The defaults match
Megatron's native behavior.

### Mode 2: Use the classes directly

Best for users who build their model spec by hand or need custom config.
Import the classes and wire them in yourself:

```python
from liger_kernel.megatron import LigerMegatronRMSNorm, LigerMegatronCrossEntropy

# In your TransformerBlockSubmodules:
input_layernorm = LigerMegatronRMSNorm
pre_mlp_layernorm = LigerMegatronRMSNorm

# In a GPTModel subclass:
class MyModel(GPTModel):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.liger_ce = LigerMegatronCrossEntropy(
            ignore_index=-100,
            label_smoothing=0.1,
        )

    def compute_language_model_loss(self, labels, logits):
        labels_sb = labels.transpose(0, 1).contiguous()
        loss_sb = self.liger_ce(logits, labels_sb, self.pg_collection.tp)
        return loss_sb.transpose(0, 1).contiguous()
```

The class accepts the same knobs Liger's standalone CE accepts
(`ignore_index`, `label_smoothing`, etc.). Use Mode 2 when you need
them.

## Scope

This PR supports `tensor_model_parallel_size=1` only.

Why? At TP>1, each GPU holds a slice of the vocabulary, not the full
vocab. Cross-entropy then needs cross-rank all-reduces to get a correct
loss. Liger's kernel does not do those all-reduces. So the math would be
wrong at TP>1.

To stop you from training with a broken loss, the patch raises
`RuntimeError` if it sees TP>1. The check runs at patch time and again
at call time, so misconfigurations fail loudly.

**[TODO]**: TP>1 support is planned as follow-up work (a separate
"vocab-parallel" Triton kernel).

You can still use this PR with pipeline parallelism, data parallelism,
or single-GPU runs. Those are all TP=1 cases.


## Testing

All unit tests live in `test/megatron/`.

### Run all megatron tests

```bash
python -m pytest test/megatron/ --no-cov
```

That runs three files:

- `test/megatron/test_rms_norm.py` — RMSNorm class correctness (forward,
backward, dtypes).
- `test/megatron/test_cross_entropy.py` — `LigerMegatronCrossEntropy`
class correctness. Covers shape sweeps, `ignore_index`,
`label_smoothing`, combined config, out-of-bounds targets,
`torch.no_grad()` path, and gradient parity vs. PyTorch's
`F.cross_entropy`.
- `test/megatron/test_monkey_patch.py` —
`apply_liger_kernel_to_megatron()` plumbing. Covers symbol replacement
for both CE paths and both RMSNorm targets, idempotency,
missing-megatron errors, TP>1 guard, dispatch behavior, and end-to-end
correctness through the patched symbols.

You should see all tests pass (currently 105 tests in the megatron
suite).

### Run a single test file

```bash
python -m pytest test/megatron/test_cross_entropy.py -v --no-cov
```

### What the tests do NOT need

Megatron-LM is **not** a test dependency. The monkey-patch tests use
stub modules in `sys.modules`, so they run on CPU without a real
Megatron install. You only need PyTorch + Triton + Liger itself.

The correctness tests (RMSNorm, CE) need a GPU because they run the
actual Triton kernel.

## Benchmarking

### Step 1: Run the benchmark

```bash
python benchmark/scripts/benchmark_megatron_cross_entropy.py --overwrite
```

This compares four providers across vocab sizes 4K → 128K (with S=2048,
B=4, BF16):

- **liger** — `LigerMegatronCrossEntropy`
- **torch** — vanilla `F.cross_entropy`
- **megatron** — Megatron's fused CE (`@jit_fuser` via TorchScript)
- **megatron-unfused** — Megatron's unfused CE (eager Python)

The script needs a CUDA/ROCm GPU. If Megatron-LM is not installed, the
two megatron providers are silently skipped and you only get `liger` vs
`torch`.

Results land in `benchmark/data/all_benchmark_data.csv`. Each row is
tagged with `kernel_name="megatron_cross_entropy"`.

### Step 2: Render plots

```bash
python benchmark/benchmarks_visualizer.py \
    --kernel-name megatron_cross_entropy --metric-name speed

python benchmark/benchmarks_visualizer.py \
    --kernel-name megatron_cross_entropy --metric-name memory
```

PNGs land in `benchmark/visualizations/`.

### Memory savings
<img width="1000" height="600"
alt="megatron_cross_entropy_memory_full_token_length"
src="https://github.com/user-attachments/assets/d08d7ee7-f6de-4e8e-a548-98b50da066ea"
/>

### Backward speedup
<img width="1000" height="600"
alt="megatron_cross_entropy_speed_backward_token_length"
src="https://github.com/user-attachments/assets/f8b524e2-658f-4206-b5a5-9d4d5e52e816"
/>

### Forward speedup
<img width="1000" height="600"
alt="megatron_cross_entropy_speed_forward_token_length"
src="https://github.com/user-attachments/assets/771115ff-143a-4c63-b112-c50bc0dddd84"
/>

### Backward+Forward speedup
<img width="1000" height="600"
alt="megatron_cross_entropy_speed_full_token_length"
src="https://github.com/user-attachments/assets/2152656c-c13e-4a6f-8b3d-577d6365a26e"
/>


## Example training scripts

Two end-to-end examples in `examples/megatron/`. Both train a tiny GPT
model with mock data for 5 iterations. They print which classes filled
which norm slots and which function backs each CE entry point, so you
can see Liger took over.

### Prerequisites

- A working Megatron-Core install (`pip install megatron-core`).
- `liger-kernel` installed.
- `psutil` (Megatron's async checkpoint worker pool uses it).
- At least 2 GPUs.

### Example 1: Mode 1 (monkey patch)

```bash
torchrun --nproc_per_node=2 \
    --master_addr=127.0.0.1 --master_port=29500 \
    examples/megatron/run_mode1_monkey_patch.py
```

One `apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)`
call at the top. Everything else is stock Megatron code.

### Example 2: Mode 2 (hand-built spec)

```bash
torchrun --nproc_per_node=2 \
    --master_addr=127.0.0.1 --master_port=29500 \
    examples/megatron/run_mode2_hand_spec.py
```

No monkey patch. The script builds a `TransformerBlockSubmodules` with
`LigerMegatronRMSNorm` slots and subclasses `GPTModel` to use
`LigerMegatronCrossEntropy` in `compute_language_model_loss`.

### What you should see

For both:

- 5 lines of `[modeN] iter <i> loss=<float>` with the loss going down.
- A printed module tree with `LigerMegatronRMSNorm` in 5 of 5 norm slots
(4 per-layer + 1 block-level final).
- The CE symbols print as `liger_fused_vocab_parallel_cross_entropy` and
`liger_vocab_parallel_cross_entropy` (Mode 1) or
`LigerMegatronCrossEntropy` (Mode 2).
- `Successfully loaded the model` after a distributed checkpoint
round-trip.

If Apex or TransformerEngine isn't installed, you'll see harmless
warnings. Megatron falls back to the local backend, which is where Liger
plugs in.

## Testing Done

- Hardware Type: H100
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Co-authored-by: Vaibhav Jindal <[email protected]>
Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>

2026-06-10 20:51:13 +00:00

5.5 KiB

Raw Permalink Blame History

High-Level APIs

AutoModel

AutoModel Variant	API
AutoModelForCausalLM	`liger_kernel.transformers.AutoLigerKernelForCausalLM`

This API extends the implementation of the AutoModelForCausalLM within the transformers library from Hugging Face.

::: liger_kernel.transformers.AutoLigerKernelForCausalLM options: extra: show_docstring: true show_signature: true show_source: true

!!! Example "Try it Out" You can experiment as shown in this example here.

Patching

You can also use the Patching APIs to use the kernels for a specific model architecture.

Model	API	Supported Operations
LLaMA 2 & 3	`liger_kernel.transformers.apply_liger_kernel_to_llama`	RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy
LLaMA 3.2-Vision	`liger_kernel.transformers.apply_liger_kernel_to_mllama`	RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy
Mistral	`liger_kernel.transformers.apply_liger_kernel_to_mistral`	RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy
Mixtral	`liger_kernel.transformers.apply_liger_kernel_to_mixtral`	RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy
Gemma1	`liger_kernel.transformers.apply_liger_kernel_to_gemma`	RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy
Gemma2	`liger_kernel.transformers.apply_liger_kernel_to_gemma2`	RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy
Qwen2, Qwen2.5, & QwQ	`liger_kernel.transformers.apply_liger_kernel_to_qwen2`	RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy
Qwen2-VL	`liger_kernel.transformers.apply_liger_kernel_to_qwen2_vl`	RMSNorm, LayerNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy
Phi3 & Phi3.5	`liger_kernel.transformers.apply_liger_kernel_to_phi3`	RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy

Function Signatures

::: liger_kernel.transformers.apply_liger_kernel_to_llama options: extra: show_docstring: true show_signature: true

::: liger_kernel.transformers.apply_liger_kernel_to_mllama options: extra: show_docstring: true show_signature: true

::: liger_kernel.transformers.apply_liger_kernel_to_mistral options: extra: show_docstring: true show_signature: true

::: liger_kernel.transformers.apply_liger_kernel_to_mixtral options: extra: show_docstring: true show_signature: true

::: liger_kernel.transformers.apply_liger_kernel_to_gemma options: extra: show_docstring: true show_signature: true

::: liger_kernel.transformers.apply_liger_kernel_to_gemma2 options: extra: show_docstring: true show_signature: true

::: liger_kernel.transformers.apply_liger_kernel_to_qwen2 options: extra: show_docstring: true show_signature: true

::: liger_kernel.transformers.apply_liger_kernel_to_qwen2_vl options: extra: show_docstring: true show_signature: true

::: liger_kernel.transformers.apply_liger_kernel_to_phi3 options: extra: show_docstring: true show_signature: true

Megatron-LM

Liger also exposes a patch for the Megatron-LM training framework, replacing Megatron's native RMSNorm and both vocab-parallel cross-entropy paths (fused and unfused) with Liger's Triton kernels.

Framework	API	Supported Operations
Megatron-LM	`liger_kernel.megatron.apply_liger_kernel_to_megatron`	RMSNorm, CrossEntropyLoss

Scope: Initial release supports tensor_model_parallel_size=1 only for cross-entropy. Vocab-parallel cross-entropy (TP>1) is follow-up work — with TP>1, each rank holds a sharded [N, V/tp] logits slice and cross-entropy requires cross-rank all-reduces that Liger's kernel does not perform. The patch raises a RuntimeError at patch time or call time if TP>1 is detected.

Usage:

from liger_kernel.megatron import apply_liger_kernel_to_megatron

# Call before Megatron's forward pass reaches compute_language_model_loss.
# Defaults match Megatron's native CE behavior; no CE-specific config needed.
apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)

Both the fused (config.cross_entropy_loss_fusion=True, cross_entropy_fusion_impl='native') and unfused (config.cross_entropy_loss_fusion=False) CE paths are patched in a single call, so Megatron picks up Liger regardless of which path your config selects.

For training setups that need explicit kernel configuration (custom ignore_index, label_smoothing, etc.), instantiate LigerMegatronCrossEntropy directly and wire it into your model — see examples/megatron/run_mode2_hand_spec.py.

::: liger_kernel.megatron.apply_liger_kernel_to_megatron options: extra: show_docstring: true show_signature: true

5.5 KiB Raw Permalink Blame History