[Megatron] Add cross-entropy integration

[PrathyushaPolepalli]

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:

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:

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

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

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

Benchmarked against megatron-core==0.19.0+96a894ef6

Step 1: Run the benchmark

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

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

Backward speedup

Forward speedup

Backward+Forward speedup

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)

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)

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
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[Mecoli1219]

Overall looks great! Excited to support Megatron with Liger. Left some comments to address.

Made changes on top