mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
0400d31357
## 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]>