linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Megatron] Add vocab-parallel cross-entropy (TP>1 support) (#1260)

Browse Source 
## Summary

Adds TP>1 (vocab-parallel) support to Liger's Megatron-LM cross-entropy.
The previous PR (#1207) shipped CE for TP=1 only. At TP>1, each rank
holds a slice of the vocabulary and the loss needs cross-rank
all-reduces to be correct. This PR writes the vocab-parallel kernel and
wires it through `LigerMegatronCrossEntropy` so the same one-line monkey
patch now works at **every** TP size — TP=1, 2, 4, 8, ...

Single code path. AllReduces are no-ops at TP=1 and real NCCL
collectives at TP>1.

Two new public Module classes — both backed by the same kernel:

- **`LigerVocabParallelCrossEntropy`** (`liger_kernel.transformers`) —
general-purpose vocab-parallel CE. Defaults match
`torch.nn.functional.cross_entropy` semantics.
- **`LigerMegatronCrossEntropy`** (`liger_kernel.megatron`) — same
kernel, Megatron-LM defaults. Bit-identical to Megatron's
`vocab_parallel_cross_entropy`.

Label smoothing supports both PyTorch's α/K formula and Megatron's NeMo
α·V/(V−1) formula; pick via `label_smoothing_formula="pytorch"` or
`"megatron"`. At TP>1 the averaging scope is selectable too:
`label_smoothing_mode="global"` (exact, one extra AllReduce) or
`"partition"` (bit-identical to Megatron, no extra AllReduce, per-rank
loss may differ).

### Future Work
- Comparison with TE kernels
- exp in bf16

## 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)
```

Defaults of `LigerMegatronCrossEntropy` match Megatron's vocab-parallel
CE bit-for-bit, so users monkey-patching from Megatron see zero
numerical change.

### Mode 2: Use the classes directly

```python
# Drop-in for Megatron's vocab_parallel_cross_entropy:
from liger_kernel.megatron import LigerMegatronCrossEntropy
ce = LigerMegatronCrossEntropy(ignore_index=-100, label_smoothing=0.0)
loss = ce(vocab_parallel_logits, target, tp_group=my_tp_group)

# Or, general-purpose vocab-parallel CE with PyTorch / F.cross_entropy semantics:
from liger_kernel.transformers import LigerVocabParallelCrossEntropy
ce = LigerVocabParallelCrossEntropy(ignore_index=-100, label_smoothing=0.0)
loss = ce(vocab_parallel_logits, target, tp_group=my_tp_group)
```

## Scope

In scope:
- Vocab-parallel CE for TP=1, 2, 4, 8, … (any TP size).
- Both label-smoothing formulas + both averaging modes at TP>1.

---

## Performance — Liger vs Megatron's vocab-parallel CE

Measured on **8× NVIDIA H100 80GB HBM3** at the Megatron training shape
`S=2048, B=4, bf16`. Compared against `megatron-core==0.19.0+96a894ef6`:

- `megatron-fused`: Megatron's *fused*
`fused_vocab_parallel_cross_entropy` (TorchScript JIT via `@jit_fuser`)
- `megatron-unfused`: Megatron's *unfused*
`vocab_parallel_cross_entropy` (eager Python; the path with runtime
`label_smoothing` support)

### Headline: Liger vs Megatron-fused at production V (≥32K)

|  | Forward | Backward (incl. fwd) | Full (fwd+bwd) |
|---|---|---|---|
| **TP=1** at V=128K  | **+298%** | +129% | +129% |
| **TP=2** at V=128K  | +274% | +123% | +123% |
| **TP=4** at V=128K  | +239% | +114% | +114% |
| **TP=8** at V=128K  | +210% | +108% | +108% |
| **TP=8** at V=32K   | +21%  | +24%  | +24%  |

At small `V_local` (high TP × small V) the kernels become tiny and
Python orchestration dominates; even there Liger is at-worst tied with
Megatron-fused, never significantly slower. At production-scale V (32K
and up), Liger is **2–4× faster on forward** and **~2× faster on
backward** at every TP we tested.

vs `megatron-unfused`: Liger is **30–430% faster** at every (TP, V)
combination — that path doesn't JIT-fuse anything.

### Per-TP plots

<img width="1680" height="1200" alt="megatron_cross_entropy_tp1_panel"
src="https://github.com/user-attachments/assets/06142823-835a-43e9-9fa9-53c6a0c7ae14"
/>
<img width="1680" height="1200" alt="megatron_cross_entropy_tp2_panel"
src="https://github.com/user-attachments/assets/82389842-9f83-4ce7-8972-a486757846f4"
/>
<img width="1680" height="1200" alt="megatron_cross_entropy_tp4_panel"
src="https://github.com/user-attachments/assets/fdc74163-f5af-46e9-904d-aa8ac9ceaeb3"
/>
<img width="1680" height="1200" alt="megatron_cross_entropy_tp8_panel"
src="https://github.com/user-attachments/assets/fff82bf2-b6ab-4635-910a-cde9ffef58e8"
/>


(Drag-and-drop the 4 PNGs at
`/tmp/megatron_ce_panels/megatron_cross_entropy_tp{1,2,4,8}_panel.png`
into the PR description here. Each PNG has forward / backward / full
speed + memory in one 2×2 view.)

### Why is Liger faster?

H100 CE is **memory-bandwidth-bound**, not compute-bound. At V=128K the
input logits tensor is ~2GB of bf16; each full read or write of it is
~0.7ms on HBM3, and the runtime is dominated by how many times you touch
that buffer.

**Megatron's fused (TorchScript) path does it in stages**, with each
stage being its own JIT-fused kernel that allocates a new buffer:
- `calculate_logits_max` casts the input bf16 → fp32 (writes 4GB), then
computes max
- `calculate_predicted_logits` shifts in-place (read 4GB + write 4GB),
then exps in-place (read 4GB + write 4GB), then gathers predicted and
sum_exp
- → roughly **4 full reads + 3 full writes** of the V_local-sized buffer

**Liger does it in ONE pass.** Kernel 1 reads input bf16, casts fp32 in
registers, subtracts the global max in registers, exps in registers,
writes the exp buffer (fp32) once, and accumulates `sum_exp` +
`predicted_logit` as scalars on the way through. Kernel 2 (backward)
does symmetrically one read + one write. AllReduces happen on tiny
`[BT]`-shaped tensors between kernels at TP>1.

That's ~3.5× less memory traffic — which matches the observed 3-4×
forward speedup almost exactly. The backward speedup (~2×) is smaller
because the backward is fundamentally one-pass-shaped for both
implementations; the win there is fusing the `grad_output` scaling into
the same kernel as the gradient computation instead of doing it as a
separate buffer mul.

---

## Memory — Liger vs Megatron

**What "memory" means at TP>1.** The number we report is
`torch.cuda.max_memory_allocated()` on rank 0 — peak HBM usage on a
**single GPU** during one fwd+bwd CE call. At TP=N each rank holds only
`V_global/N` of the vocab, so per-rank memory **scales as 1/N**. All
ranks have identical memory profiles (vocab is partitioned evenly), so
rank 0 represents every rank. It is **not** "aggregate across all GPUs"
— each GPU is independently HBM-bound, so per-rank is the metric that
matters for fitting a model on a node.

### Numbers (MB per rank — Liger / Megatron-fused / Megatron-unfused)

| V | TP=1 | TP=2 | TP=4 | TP=8 |
|---|---|---|---|---|
| 4K | 256 / 384 / 256 | 128 / 192 / 128 | 64 / 96 / 64 | 32 / 48 / 32 |
| 16K | 1024 / 1536 / 1024 | 512 / 768 / 512 | 256 / 384 / 256 | 128 /
192 / 128 |
| 32K | 2048 / 3072 / 2048 | 1024 / 1536 / 1024 | 512 / 768 / 512 | 256
/ 384 / 256 |
| 128K | 8192 / 12288 / 8192 | 4096 / 6144 / 4096 | 2048 / 3072 / 2048 |
1024 / 1536 / 1024 |

**Two takeaways:**

1. **Liger uses 33% less per-rank HBM than Megatron-fused at every (TP,
V) point.** That's 512 MB freed at V=128K/TP=8 — and the same 512 MB
savings replicates across all 8 ranks. Usable for longer sequences,
larger batches, or activation checkpointing.

2. **Liger uses exactly the same memory as Megatron-unfused.** Both save
a single `[BT, V_local]` fp32 softmax buffer for backward — that's the
irreducible lower bound for vocab-parallel CE. The plot annotates this
overlap explicitly.

### Why is Megatron-fused 50% heavier than Liger / Megatron-unfused?

Both Liger and Megatron-unfused save one `[BT, V_local]` fp32 softmax
buffer for backward. Megatron-fused saves the same buffer, but its
`@jit_fuser` TorchScript scope holds **captured input tensors alive
until the function returns** — so during forward the bf16 input AND the
fp32 working buffer are both resident in HBM, peaking at ~1 extra `[BT,
V_local]` buffer's worth. The 50% premium is consistent across V and TP.

Liger avoids this entirely: the bf16 input is read in-register inside
Kernel 1 and never replicated to a separate fp32 buffer in HBM.
Megatron-unfused avoids it because eager Python rebinds the local
variable after `.float()`, letting the bf16 input be GC'd.

So Liger and unfused arrive at the same lower bound through different
routes (Triton single-pass kernel vs eager Python ops), and
Megatron-fused pays a 50% premium in exchange for TorchScript fusion
that — as the speed numbers show — doesn't actually pay off vs
hand-written Triton.

---

## Testing

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

### Single-GPU tests

```bash
python -m pytest test/megatron/test_cross_entropy.py --no-cov -k "not tp_vocab_parallel"
python -m pytest test/megatron/test_monkey_patch.py --no-cov
```

Covers:
- `LigerVocabParallelCrossEntropy` correctness — fwd + bwd parity vs
`F.cross_entropy` across shape sweeps, fp32 / bf16, `ignore_index`,
`label_smoothing`, large logit magnitudes, out-of-bounds targets,
CE-not-last-layer, `torch.no_grad()`.
- `LigerMegatronCrossEntropy` correctness — same matrix, with
`label_smoothing > 0` parity referenced against
`F.cross_entropy(label_smoothing=α·V/(V-1))` which is algebraically
equivalent to Megatron's NeMo formula at TP=1.
- Config plumbing: validation of `reduction`, `label_smoothing_formula`,
`label_smoothing_mode`, `label_smoothing` range; 3-D shape guard.
- Input-preservation invariant: forward must NOT mutate the caller's
input tensor even when fp32 (regression test for in-place ops aliasing).
- Monkey-patch plumbing: symbol replacement for both fused / unfused CE
paths, idempotency, missing-megatron errors, runtime `label_smoothing`
override on the unfused path, end-to-end correctness through patched
symbols.

77 single-GPU CE tests + 35 monkey-patch tests = 112 single-GPU tests,
all passing.

### Multi-GPU TP>1 parity tests

```bash
python -m pytest test/megatron/test_cross_entropy.py --no-cov -k "tp_vocab_parallel"
```

Parametrized over `tp_size ∈ {2, 4, 8}` with per-tp_size
`skipif(torch.cuda.device_count() < tp_size, ...)`. 48 tests total
covering the full cross-product `formula × mode × label_smoothing ×
shape`. The right subset runs based on available GPUs (2 GPUs: 16 run,
32 skip; 8 GPUs: all 48 run).

Each test broadcasts the same full-vocab logits to every rank, has each
rank slice out its `V_local`, runs Liger's vocab-parallel CE, and
compares forward + backward against a single-rank `F.cross_entropy`
reference on the full logits (with α rescaled for the Megatron formula).

### Style

```bash
make checkstyle
```

Clean.

---

## Benchmarking

Single script handles all TP sizes via `--tp-size`. Uses `mp.spawn`
internally and allocates `tp_size` GPUs.

```bash
python benchmark/scripts/benchmark_megatron_cross_entropy.py --tp-size 1 --overwrite
python benchmark/scripts/benchmark_megatron_cross_entropy.py --tp-size 2 --overwrite
python benchmark/scripts/benchmark_megatron_cross_entropy.py --tp-size 4 --overwrite   # 4 GPUs
python benchmark/scripts/benchmark_megatron_cross_entropy.py --tp-size 8 --overwrite   # 8 GPUs

# Render per-TP plots from the CSV
python benchmark/benchmarks_visualizer.py \
    --kernel-name megatron_cross_entropy --metric-name speed --overwrite \
    --extra-config-filter '"TP": 1'
```

Note: the committed CSV intentionally contains **only TP=1 rows** for
`megatron_cross_entropy` to keep the data file small. Re-run at TP=2/4/8
to repopulate the rows for those TPs if you want to validate the numbers
in the headline table above.

## Testing Done

- Hardware Type: H100 80GB HBM3 × 8
- [x] `python -m pytest test/megatron/ --no-cov` — passes (77 single-GPU
CE + 35 monkey-patch + N multi-GPU CE depending on `cuda.device_count()`
+ 13 RMSNorm)
- [x] `python -m pytest test/transformers/test_cross_entropy.py --no-cov
-k "test_correctness and not z_loss"` — 136 passed (existing CE not
regressed)
- [x] `make checkstyle` — clean
- [x] Full TP=1/2/4/8 benchmark sweep on 8×H100 — numbers committed
inline above; 4-panel plots attached
- [ ] `make test-convergence`

---------

Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
Vaibhav Jindal
2026-06-16 19:10:30 +00:00
committed by GitHub
co-authored by Claude Opus 4.7
parent 813590d328
commit 80f19ad741
12 changed files with 1323 additions and 560 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/benchmarks_visualizer.py
+9 -1
View File
@@ -314,7 +314,15 @@ def plot_data(df: pd.DataFrame, config: VisualizationsConfig):
xlabel = df["x_label"].iloc[0]
ylabel = f"{config.metric_name} ({df['metric_unit'].iloc[0]})"
# Sort by "kernel_provider" to ensure consistent color assignment
# Sort by "kernel_provider" so color assignment is stable across runs, with
# "liger" pinned to the END so seaborn draws it LAST (on top). Without this,
# any other provider with identical y-values would obscure Liger at the overlap
# — e.g., megatron-unfused has bit-identical memory to Liger in the
# megatron_cross_entropy benchmark and would hide it under alphabetical sort.
providers = df["kernel_provider"].unique().tolist()
non_liger = sorted(p for p in providers if p != "liger")
order = non_liger + (["liger"] if "liger" in providers else [])
df["kernel_provider"] = pd.Categorical(df["kernel_provider"], categories=order, ordered=True)
df = df.sort_values(by="kernel_provider")
plt.figure(figsize=(10, 6))
benchmark/data/all_benchmark_data.csv
+72 -96
View File
@@ -2219,102 +2219,6 @@ jsd,liger,full,memory,MB,BT,total tokens,1024,3514.0009765625,3514.0009765625,35
jsd,liger,full,memory,MB,BT,total tokens,2048,7014.0009765625,7014.0009765625,7014.0009765625,"{""vocab_size"": 128256, ""bsz"": 1, ""seq_len"": 8192}",NVIDIA B200,2026-05-27 17:14:32,0.8.0
jsd,liger,full,memory,MB,BT,total tokens,4096,14028.0009765625,14028.0009765625,14028.0009765625,"{""vocab_size"": 128256, ""bsz"": 1, ""seq_len"": 8192}",NVIDIA B200,2026-05-27 17:14:32,0.8.0
jsd,liger,full,memory,MB,BT,total tokens,8192,28056.0,28056.0,28056.0,"{""vocab_size"": 128256, ""bsz"": 1, ""seq_len"": 8192}",NVIDIA B200,2026-05-27 17:14:32,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,4096,0.2385600060224533,0.23590399324893951,0.24048000574111938,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:58,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,8192,0.26820799708366394,0.26531198620796204,0.27125120162963867,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:58,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,16384,0.3599199950695038,0.3569599986076355,0.36266239881515505,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:58,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,32768,0.6502079963684082,0.6452223896980286,0.6563008189201356,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:58,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,65536,1.087440013885498,1.0827391862869262,1.0942399978637696,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:58,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,131072,2.1112000942230225,2.108524799346924,2.1132928848266603,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:58,0.8.0
megatron_cross_entropy,torch,forward,speed,ms,V,vocab size,4096,0.2269120067358017,0.225913605093956,0.22775039970874789,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:59,0.8.0
megatron_cross_entropy,torch,forward,speed,ms,V,vocab size,8192,0.48070400953292847,0.47954559326171875,0.48152959942817686,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:59,0.8.0
megatron_cross_entropy,torch,forward,speed,ms,V,vocab size,16384,0.8124480247497559,0.8111680150032043,0.813696026802063,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:59,0.8.0
megatron_cross_entropy,torch,forward,speed,ms,V,vocab size,32768,2.1381120681762695,2.1316160202026366,2.142361545562744,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:59,0.8.0
megatron_cross_entropy,torch,forward,speed,ms,V,vocab size,65536,4.391056060791016,4.389933013916016,4.3939519882202145,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:59,0.8.0
megatron_cross_entropy,torch,forward,speed,ms,V,vocab size,131072,8.81107234954834,8.806303977966309,8.81488037109375,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:43:59,0.8.0
megatron_cross_entropy,megatron,forward,speed,ms,V,vocab size,4096,0.31273600459098816,0.31200000643730164,0.31331199407577515,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:01,0.8.0
megatron_cross_entropy,megatron,forward,speed,ms,V,vocab size,8192,0.6373440027236938,0.636352002620697,0.6381440162658691,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:01,0.8.0
megatron_cross_entropy,megatron,forward,speed,ms,V,vocab size,16384,1.3848639726638794,1.3825664043426515,1.3870784282684325,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:01,0.8.0
megatron_cross_entropy,megatron,forward,speed,ms,V,vocab size,32768,2.829935908317566,2.8288448333740233,2.8316224098205565,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:01,0.8.0
megatron_cross_entropy,megatron,forward,speed,ms,V,vocab size,65536,5.823584079742432,5.8137922286987305,5.826655864715576,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:01,0.8.0
megatron_cross_entropy,megatron,forward,speed,ms,V,vocab size,131072,11.810944080352783,11.80013427734375,11.817446517944337,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:01,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,4096,0.5649279952049255,0.5639231920242309,0.5661375880241394,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:02,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,8192,1.094543993473053,1.093395209312439,1.0956799983978271,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:02,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,16384,2.0701760053634644,2.068671941757202,2.072096109390259,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:02,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,32768,4.027215957641602,4.024915313720703,4.028927993774413,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:02,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,65536,7.93833589553833,7.937132930755616,7.940275096893311,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:02,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,131072,15.817471981048584,15.814432144165039,15.821215629577637,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:02,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,4096,0.3516159951686859,0.3489919900894165,0.35476480722427367,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,8192,0.4639679938554764,0.46081281304359434,0.46724479794502255,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,16384,0.7123200297355652,0.7092544078826905,0.7154880046844483,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,32768,1.322975993156433,1.3184319734573364,1.328320026397705,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,65536,2.3985120058059692,2.393356847763062,2.405248022079468,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,131072,4.698495864868164,4.69042558670044,4.70184965133667,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:03,0.8.0
megatron_cross_entropy,torch,backward,speed,ms,V,vocab size,4096,0.48710399866104126,0.4862975895404816,0.48793599009513855,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:04,0.8.0
megatron_cross_entropy,torch,backward,speed,ms,V,vocab size,8192,0.9741439819335938,0.9731391787528991,0.9752448081970215,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:04,0.8.0
megatron_cross_entropy,torch,backward,speed,ms,V,vocab size,16384,1.8759199976921082,1.8737216234207152,1.8779136180877687,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:04,0.8.0
megatron_cross_entropy,torch,backward,speed,ms,V,vocab size,32768,4.405440092086792,4.398931121826172,4.407916831970215,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:04,0.8.0
megatron_cross_entropy,torch,backward,speed,ms,V,vocab size,65536,8.920191764831543,8.916671752929688,8.923616409301758,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:04,0.8.0
megatron_cross_entropy,torch,backward,speed,ms,V,vocab size,131072,17.845855712890625,17.843469619750977,17.847046661376954,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:04,0.8.0
megatron_cross_entropy,megatron,backward,speed,ms,V,vocab size,4096,0.5368639826774597,0.536191999912262,0.5374848127365113,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:05,0.8.0
megatron_cross_entropy,megatron,backward,speed,ms,V,vocab size,8192,1.0521279573440552,1.05141122341156,1.0528064489364624,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:05,0.8.0
megatron_cross_entropy,megatron,backward,speed,ms,V,vocab size,16384,2.1722079515457153,2.169472026824951,2.1735936641693114,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:05,0.8.0
megatron_cross_entropy,megatron,backward,speed,ms,V,vocab size,32768,4.448447942733765,4.447282981872559,4.449868965148926,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:05,0.8.0
megatron_cross_entropy,megatron,backward,speed,ms,V,vocab size,65536,9.033520221710205,9.030751991271973,9.042854690551758,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:05,0.8.0
megatron_cross_entropy,megatron,backward,speed,ms,V,vocab size,131072,17.958431243896484,17.9557315826416,17.960690689086913,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:05,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,4096,0.7786880135536194,0.7767040133476257,0.7803391933441162,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:06,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,8192,1.475823998451233,1.4742015838623046,1.4774847745895388,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:06,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,16384,2.787328004837036,2.785683250427246,2.7883071899414062,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:06,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,32768,5.420928001403809,5.419379329681396,5.423315238952636,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:06,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,65536,10.677472114562988,10.674201774597169,10.67763843536377,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:06,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,131072,21.254928588867188,21.25368995666504,21.2557315826416,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:06,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,4096,0.41488000750541687,0.4122175931930542,0.41767039299011227,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:07,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,8192,0.5923520028591156,0.5882560014724731,0.6005120277404785,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:07,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,16384,0.9725440144538879,0.9690751910209655,0.9779008150100708,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:07,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,32768,1.8396799564361572,1.8349119424819946,1.8454079627990723,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:07,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,65536,3.4398880004882812,3.429087924957275,3.4599680423736574,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:07,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,131072,6.780799865722656,6.763827323913574,6.796000003814697,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:07,0.8.0
megatron_cross_entropy,torch,full,speed,ms,V,vocab size,4096,0.5540960133075714,0.5528640151023865,0.555296003818512,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:08,0.8.0
megatron_cross_entropy,torch,full,speed,ms,V,vocab size,8192,1.1072640419006348,1.1059776306152345,1.1086400032043457,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:08,0.8.0
megatron_cross_entropy,torch,full,speed,ms,V,vocab size,16384,2.139024019241333,2.1362879276275635,2.142067241668701,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:08,0.8.0
megatron_cross_entropy,torch,full,speed,ms,V,vocab size,32768,4.9202880859375,4.91553258895874,4.924480152130127,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:08,0.8.0
megatron_cross_entropy,torch,full,speed,ms,V,vocab size,65536,9.961983680725098,9.958450889587402,9.964262008666992,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:08,0.8.0
megatron_cross_entropy,torch,full,speed,ms,V,vocab size,131072,19.926143646240234,19.913933563232423,19.93060417175293,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:08,0.8.0
megatron_cross_entropy,megatron,full,speed,ms,V,vocab size,4096,0.6057599782943726,0.604960024356842,0.6063359975814819,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:09,0.8.0
megatron_cross_entropy,megatron,full,speed,ms,V,vocab size,8192,1.1868000030517578,1.1857984066009521,1.1880639791488647,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:09,0.8.0
megatron_cross_entropy,megatron,full,speed,ms,V,vocab size,16384,2.4351680278778076,2.4330944538116457,2.4368255615234373,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:09,0.8.0
megatron_cross_entropy,megatron,full,speed,ms,V,vocab size,32768,4.968512058258057,4.967942237854004,4.969702434539795,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:09,0.8.0
megatron_cross_entropy,megatron,full,speed,ms,V,vocab size,65536,10.074239730834961,10.06367359161377,10.077676963806152,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:09,0.8.0
megatron_cross_entropy,megatron,full,speed,ms,V,vocab size,131072,20.022079467773438,20.015307998657228,20.02815971374512,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:09,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,4096,0.846015989780426,0.8447743892669677,0.8472639918327332,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,8192,1.6103359460830688,1.6088960409164428,1.6117696285247802,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,16384,3.0503358840942383,3.049056053161621,3.0524160861968994,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,32768,5.937999963760376,5.936895847320557,5.939583778381348,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,65536,11.711487770080566,11.710764503479005,11.713843536376952,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,131072,23.322416305541992,23.321644973754882,23.32369956970215,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,4096,192.0791015625,192.0791015625,192.0791015625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,8192,384.0791015625,384.0791015625,384.0791015625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,16384,768.0791015625,768.0791015625,768.0791015625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,32768,1536.0791015625,1536.0791015625,1536.0791015625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,65536,3072.0791015625,3072.0791015625,3072.0791015625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,131072,6144.0791015625,6144.0791015625,6144.0791015625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,torch,full,memory,MB,V,vocab size,4096,512.0947265625,512.0947265625,512.0947265625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,torch,full,memory,MB,V,vocab size,8192,1024.0947265625,1024.0947265625,1024.0947265625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,torch,full,memory,MB,V,vocab size,16384,2048.0947265625,2048.0947265625,2048.0947265625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,torch,full,memory,MB,V,vocab size,32768,4096.0947265625,4096.0947265625,4096.0947265625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,torch,full,memory,MB,V,vocab size,65536,8192.0947265625,8192.0947265625,8192.0947265625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,torch,full,memory,MB,V,vocab size,131072,16384.09375,16384.09375,16384.09375,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron,full,memory,MB,V,vocab size,4096,448.1650390625,448.1650390625,448.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron,full,memory,MB,V,vocab size,8192,896.1650390625,896.1650390625,896.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron,full,memory,MB,V,vocab size,16384,1792.1650390625,1792.1650390625,1792.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron,full,memory,MB,V,vocab size,32768,3584.1650390625,3584.1650390625,3584.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron,full,memory,MB,V,vocab size,65536,7168.1650390625,7168.1650390625,7168.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron,full,memory,MB,V,vocab size,131072,14336.1650390625,14336.1650390625,14336.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:10,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,4096,320.1650390625,320.1650390625,320.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:11,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,8192,640.1650390625,640.1650390625,640.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:11,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,16384,1280.1650390625,1280.1650390625,1280.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:11,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,32768,2560.1650390625,2560.1650390625,2560.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:11,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,65536,5120.1650390625,5120.1650390625,5120.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:11,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,131072,10240.1650390625,10240.1650390625,10240.1650390625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-06-10 07:44:11,0.8.0
megatron_rms_norm,liger,forward,speed,ms,H,hidden size,1024,0.012160000391304493,0.011935999616980553,0.012460799887776375,"{""S"": 4096, ""B"": 1}",NVIDIA H100 80GB HBM3,2026-06-10 21:03:37,0.8.0
megatron_rms_norm,liger,forward,speed,ms,H,hidden size,2048,0.01756799966096878,0.01740800030529499,0.017920000478625298,"{""S"": 4096, ""B"": 1}",NVIDIA H100 80GB HBM3,2026-06-10 21:03:37,0.8.0
megatron_rms_norm,liger,forward,speed,ms,H,hidden size,4096,0.029184000566601753,0.028863999992609024,0.029888000339269638,"{""S"": 4096, ""B"": 1}",NVIDIA H100 80GB HBM3,2026-06-10 21:03:37,0.8.0
@@ -2375,3 +2279,75 @@ megatron_rms_norm,megatron,full,memory,MB,H,hidden size,2048,80.0283203125,80.02
megatron_rms_norm,megatron,full,memory,MB,H,hidden size,4096,160.0400390625,160.0400390625,160.0400390625,"{""S"": 4096, ""B"": 1}",NVIDIA H100 80GB HBM3,2026-06-10 21:03:43,0.8.0
megatron_rms_norm,megatron,full,memory,MB,H,hidden size,8192,320.0634765625,320.0634765625,320.0634765625,"{""S"": 4096, ""B"": 1}",NVIDIA H100 80GB HBM3,2026-06-10 21:03:43,0.8.0
megatron_rms_norm,megatron,full,memory,MB,H,hidden size,16384,640.1103515625,640.1103515625,640.1103515625,"{""S"": 4096, ""B"": 1}",NVIDIA H100 80GB HBM3,2026-06-10 21:03:43,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,4096,0.13862720131874084,0.13853874802589417,0.14057523012161255,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,8192,0.2626374363899231,0.2623022198677063,0.2627813220024109,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,16384,0.4157696068286896,0.415627658367157,0.41619354486465454,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,32768,0.8098381161689758,0.8094235062599182,0.8100022077560425,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,65536,1.5309158563613892,1.530276894569397,1.531093955039978,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,forward,speed,ms,V,vocab size,131072,2.9674980640411377,2.9669196605682373,2.967585802078247,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,4096,0.3497951924800873,0.3472285568714142,0.35065382719039917,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,8192,0.5853517055511475,0.5849090814590454,0.5855727195739746,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,16384,1.0486783981323242,1.0482382774353027,1.0487377643585205,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,32768,2.059054136276245,2.0587127208709717,2.0591745376586914,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,65536,4.020099639892578,4.0198211669921875,4.02085542678833,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,backward,speed,ms,V,vocab size,131072,7.958672523498535,7.9583587646484375,7.958960056304932,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,4096,0.34937408566474915,0.34792205691337585,0.3503080904483795,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,8192,0.5855180621147156,0.5851470232009888,0.585608959197998,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,16384,1.0484946966171265,1.048270583152771,1.0487635135650635,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,32768,2.058947801589966,2.0586788654327393,2.0593338012695312,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,65536,4.020388603210449,4.019952297210693,4.0205230712890625,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,speed,ms,V,vocab size,131072,7.95883846282959,7.958359241485596,7.9599714279174805,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,forward,speed,ms,V,vocab size,4096,0.31217217445373535,0.3120197057723999,0.3126451373100281,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,forward,speed,ms,V,vocab size,8192,0.635949432849884,0.6356914043426514,0.6363937258720398,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,forward,speed,ms,V,vocab size,16384,1.3466770648956299,1.3459649085998535,1.3470183610916138,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,forward,speed,ms,V,vocab size,32768,2.831264019012451,2.8307409286499023,2.8313231468200684,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,forward,speed,ms,V,vocab size,65536,5.76310920715332,5.762768268585205,5.765852451324463,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,forward,speed,ms,V,vocab size,131072,11.809796333312988,11.809551239013672,11.81175708770752,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,backward,speed,ms,V,vocab size,4096,0.5311360359191895,0.5298343896865845,0.5348719358444214,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,backward,speed,ms,V,vocab size,8192,1.0505350828170776,1.050193190574646,1.0506258010864258,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,backward,speed,ms,V,vocab size,16384,2.158108711242676,2.1574387550354004,2.1584970951080322,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,backward,speed,ms,V,vocab size,32768,4.448922157287598,4.448563575744629,4.449383735656738,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,backward,speed,ms,V,vocab size,65536,8.981760025024414,8.980610847473145,8.982187271118164,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,backward,speed,ms,V,vocab size,131072,18.23187828063965,18.229421615600586,18.232585906982422,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,speed,ms,V,vocab size,4096,0.5273734331130981,0.5272742509841919,0.5274069905281067,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,speed,ms,V,vocab size,8192,1.050581693649292,1.0501985549926758,1.0508228540420532,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,speed,ms,V,vocab size,16384,2.158163070678711,2.1575186252593994,2.159698724746704,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,speed,ms,V,vocab size,32768,4.448894023895264,4.448407173156738,4.449240684509277,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,speed,ms,V,vocab size,65536,8.980452537536621,8.980046272277832,8.983654022216797,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,speed,ms,V,vocab size,131072,18.230487823486328,18.226482391357422,18.231830596923828,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,4096,0.5585324764251709,0.5583837032318115,0.5585581064224243,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,8192,1.0781855583190918,1.0779157876968384,1.0834403038024902,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,16384,2.066349506378174,2.0661680698394775,2.0664494037628174,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,32768,4.0230560302734375,4.022715091705322,4.023136138916016,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,65536,7.9344024658203125,7.9341511726379395,7.9350481033325195,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,forward,speed,ms,V,vocab size,131072,15.812352180480957,15.812163352966309,15.812748908996582,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,4096,0.7693004608154297,0.7685099840164185,0.7781162261962891,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,8192,1.4540057182312012,1.45366632938385,1.4545174837112427,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,16384,2.7810745239257812,2.7805025577545166,2.7812232971191406,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,32768,5.412109375,5.411957740783691,5.413435935974121,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,65536,10.669661521911621,10.66853141784668,10.670055389404297,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,backward,speed,ms,V,vocab size,131072,21.2467098236084,21.245935440063477,21.247703552246094,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,4096,0.767576277256012,0.7668100595474243,0.7726022601127625,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,8192,1.4551526308059692,1.4548063278198242,1.4555189609527588,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,16384,2.781857967376709,2.781548023223877,2.7830042839050293,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,32768,5.4108195304870605,5.41024923324585,5.410989284515381,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,65536,10.667474746704102,10.667390823364258,10.668229103088379,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,speed,ms,V,vocab size,131072,21.245800018310547,21.245351791381836,21.246740341186523,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,4096,256.15771484375,256.15771484375,256.15771484375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,8192,512.15771484375,512.15771484375,512.15771484375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,16384,1024.15771484375,1024.15771484375,1024.15771484375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,32768,2048.15771484375,2048.15771484375,2048.15771484375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,65536,4096.15771484375,4096.15771484375,4096.15771484375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,liger,full,memory,MB,V,vocab size,131072,8192.158203125,8192.158203125,8192.158203125,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,memory,MB,V,vocab size,4096,384.16552734375,384.16552734375,384.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,memory,MB,V,vocab size,8192,768.16552734375,768.16552734375,768.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,memory,MB,V,vocab size,16384,1536.16552734375,1536.16552734375,1536.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,memory,MB,V,vocab size,32768,3072.16552734375,3072.16552734375,3072.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,memory,MB,V,vocab size,65536,6144.16552734375,6144.16552734375,6144.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-fused,full,memory,MB,V,vocab size,131072,12288.166015625,12288.166015625,12288.166015625,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,4096,256.16552734375,256.16552734375,256.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,8192,512.16552734375,512.16552734375,512.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,16384,1024.16552734375,1024.16552734375,1024.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,32768,2048.16552734375,2048.16552734375,2048.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,65536,4096.16552734375,4096.16552734375,4096.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,131072,8192.166015625,8192.166015625,8192.166015625,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0
Can't render this file because it is too large.
benchmark/scripts/benchmark_megatron_cross_entropy.py
+257 -147
View File
@@ -1,47 +1,50 @@
"""Benchmark Liger's Megatron-LM cross-entropy wrapper.
"""Benchmark Liger's Megatron CE at any tensor_model_parallel_size.
Compares four providers on the per-token CE call shape ``[seq, batch, vocab]``:
Same implementation runs for TP=1 and TP>1 — the vocab-parallel kernel handles both
regimes (at TP=1 the AllReduces are skipped at the wrapper). One script, one CLI flag
(``--tp-size``), one CSV row format. Spawns ``tp_size`` ranks via ``mp.spawn`` and
uses a fixed-iteration manual timing loop so every rank executes the same number of
NCCL collectives (``triton.testing.do_bench`` calibrates adaptively per-rank and
desyncs the collectives — see the long comment in ``_timed_loop``).
- **torch**: vanilla ``F.cross_entropy``
- **megatron**: Megatron's *fused* ``fused_vocab_parallel_cross_entropy`` path
(``cross_entropy_loss_fusion=True``, JIT-fused via TorchScript)
Compared providers (each rank executes its own ``V/TP`` slice):
- **liger**: ``LigerMegatronCrossEntropy`` — Liger's vocab-parallel kernel.
- **megatron-fused**: Megatron's *fused* ``fused_vocab_parallel_cross_entropy``
(``cross_entropy_loss_fusion=True``, JIT-fused via TorchScript).
- **megatron-unfused**: Megatron's *unfused* ``vocab_parallel_cross_entropy``
path (``cross_entropy_loss_fusion=False``, eager Python; the path users on
``label_smoothing`` typically end up on)
- **liger**: ``LigerMegatronCrossEntropy`` — Liger's Triton CE wrapped in the
Megatron fused signature. Same kernel regardless of which Megatron symbol
it was patched onto, so we only benchmark it once.
(``cross_entropy_loss_fusion=False``, eager Python — the path with runtime
``label_smoothing`` support).
Requires a Liger-supported accelerator (CUDA / ROCm). With megatron-core not
installed, both megatron providers are silently skipped.
Rows go to ``benchmark/data/all_benchmark_data.csv`` tagged with
``kernel_name="megatron_cross_entropy"``; ``TP`` lives in
``extra_benchmark_config_str``. Render plots per-TP via ``--extra-config-filter``:
Output goes to the shared ``benchmark/data/all_benchmark_data.csv`` like every
other Liger benchmark — rows are tagged with ``kernel_name="megatron_cross_entropy"``
and the standard visualizer renders them via:
python benchmark/scripts/benchmark_megatron_cross_entropy.py --tp-size 1
python benchmark/scripts/benchmark_megatron_cross_entropy.py --tp-size 2
python benchmark/scripts/benchmark_megatron_cross_entropy.py --tp-size 4
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
--kernel-name megatron_cross_entropy --metric-name speed \\
--extra-config-filter "'TP': 2"
"""
from __future__ import annotations
import argparse
import os
import tempfile
import torch
import torch.nn.functional as F
import triton
import torch.distributed as dist
import torch.multiprocessing as mp
from utils import QUANTILES
from utils import SingleBenchmarkRunInput
from utils import SingleBenchmarkRunOutput
from utils import _test_memory
from utils import parse_benchmark_script_args
from utils import run_benchmarks
from utils import BenchmarkData
from utils import get_formatted_time
from utils import get_gpu_name
from utils import update_benchmark_data_csv
from liger_kernel.megatron import LigerMegatronCrossEntropy
from liger_kernel.utils import infer_device
device = infer_device()
try:
from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
@@ -54,154 +57,261 @@ except ImportError:
_MEGATRON_AVAILABLE = False
def _make_inputs(s: int, b: int, v: int, requires_grad: bool = True):
logits = torch.randn(s, b, v, device=device, dtype=torch.bfloat16, requires_grad=requires_grad)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
return logits, target
_WARMUP_ITERS = 10
_MEASURE_ITERS = 50
_MEMORY_ITERS = 5 # reps used for peak-memory measurement
def _pytorch_cross_entropy(logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
s, b, v = logits.shape
return F.cross_entropy(
logits.reshape(-1, v).float(),
target.reshape(-1),
reduction="none",
).reshape(s, b)
def _ensure_single_rank_tp_group():
"""Initialize torch.distributed (single-rank) and return a usable TP group.
For a single-process benchmark we use the world group of size 1; the
internal all-reduce becomes a no-op.
"""
import torch.distributed as dist
if not dist.is_initialized():
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29500")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("LOCAL_RANK", "0")
dist.init_process_group(backend="nccl")
return dist.group.WORLD
def _select_fwd(provider: str):
def _select_fwd(provider, tp_group):
if provider == "liger":
ce = LigerMegatronCrossEntropy(reduction="none")
return lambda logits, target: ce(logits, target)
if provider == "torch":
return _pytorch_cross_entropy
if provider == "megatron":
if not _MEGATRON_AVAILABLE:
raise RuntimeError("megatron-core not installed; cannot benchmark 'megatron' provider")
tp_group = _ensure_single_rank_tp_group()
def _megatron_fused_call(logits, target):
return fused_vocab_parallel_cross_entropy(logits, target, tp_group)
return _megatron_fused_call
ce = LigerMegatronCrossEntropy()
return lambda logits, target: ce(logits, target, tp_group=tp_group)
if provider == "megatron-fused":
return lambda logits, target: fused_vocab_parallel_cross_entropy(logits, target, tp_group)
if provider == "megatron-unfused":
if not _MEGATRON_AVAILABLE:
raise RuntimeError("megatron-core not installed; cannot benchmark 'megatron-unfused' provider")
tp_group = _ensure_single_rank_tp_group()
def _megatron_unfused_call(logits, target):
# Unfused signature: (logits, target, label_smoothing=0.0, tp_group=None)
return vocab_parallel_cross_entropy(logits, target, 0.0, tp_group)
return _megatron_unfused_call
return lambda logits, target: vocab_parallel_cross_entropy(logits, target, 0.0, tp_group)
raise ValueError(f"unknown provider: {provider!r}")
def bench_speed_megatron_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
v = input.x
provider = input.kernel_provider
mode = input.kernel_operation_mode
s = input.extra_benchmark_config["S"]
b = input.extra_benchmark_config["B"]
def _timed_loop(fn, tp_group, n_iters: int):
"""Run ``fn`` ``n_iters`` times and return per-iter ms.
logits, target = _make_inputs(s, b, v)
fwd_fn = _select_fwd(provider)
CRITICAL for TP>1: every rank MUST execute the same number of iterations because
``fn`` issues NCCL collectives. ``triton.testing.do_bench`` calibrates the
iteration count from the first-call latency, which differs slightly between
ranks and causes collective desync (NCCL watchdog timeout). This manual loop
with a fixed ``n_iters`` keeps every rank in lockstep. At TP=1 the same loop
is used unchanged (no AllReduces, so no desync risk).
"""
torch.cuda.synchronize()
if tp_group is not None:
dist.barrier(group=tp_group)
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(n_iters):
fn()
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) / n_iters
def fwd():
return fwd_fn(logits, target)
def _make_inputs(s, b, v_local, v_global, device):
logits = torch.randn(s, b, v_local, device=device, dtype=torch.bfloat16, requires_grad=True)
target = torch.randint(0, v_global, (s, b), device=device, dtype=torch.long)
return logits, target
def _bench_speed_one(provider, mode, s, b, v_global, tp_size, tp_group, device):
"""Time one (provider, mode, V) cell. Returns (ms_50, ms_20, ms_80) over 5 timed loops."""
v_local = v_global // tp_size
logits, target = _make_inputs(s, b, v_local, v_global, device)
fwd_fn = _select_fwd(provider, tp_group)
if mode == "forward":
ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES)
def step():
fwd_fn(logits, target)
elif mode == "backward":
# Megatron's fused CE writes gradients in-place into saved tensors during backward,
# which breaks the standard retain_graph=True / repeated-backward pattern do_bench
# uses elsewhere. Run a fresh fwd+bwd each iteration so each backward sees an
# unmodified autograd graph. Measurement therefore includes forward time —
# subtract the "forward" measurement to derive backward-only timing.
def _fwd_bwd():
# Megatron's fused CE writes gradients in-place during backward; rerun fwd each
# iteration so each backward sees a fresh graph. Measurement therefore includes
# forward; subtract the "forward" row to derive backward-only timing.
def step():
if logits.grad is not None:
logits.grad = None
out = fwd_fn(logits, target)
out.sum().backward()
ms_50, ms_20, ms_80 = triton.testing.do_bench(_fwd_bwd, rep=100, quantiles=QUANTILES)
elif mode == "full":
def full():
y = fwd()
def step():
if logits.grad is not None:
logits.grad = None
y = fwd_fn(logits, target)
y.sum().backward()
ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES)
else:
raise ValueError(f"unknown mode: {mode!r}")
return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80)
for _ in range(_WARMUP_ITERS):
step()
torch.cuda.synchronize()
if tp_group is not None:
dist.barrier(group=tp_group)
samples = []
for _ in range(5):
samples.append(_timed_loop(step, tp_group, _MEASURE_ITERS))
samples_t = torch.tensor(samples)
return (
float(samples_t.quantile(0.5)),
float(samples_t.quantile(0.2)),
float(samples_t.quantile(0.8)),
)
def bench_memory_megatron_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
v = input.x
provider = input.kernel_provider
s = input.extra_benchmark_config["S"]
b = input.extra_benchmark_config["B"]
def _bench_memory_one(provider, s, b, v_global, tp_size, tp_group, device):
"""Peak memory for one full forward+backward on this rank, in MB."""
v_local = v_global // tp_size
logits, target = _make_inputs(s, b, v_local, v_global, device)
fwd_fn = _select_fwd(provider, tp_group)
logits, target = _make_inputs(s, b, v)
fwd_fn = _select_fwd(provider)
def full():
def step():
if logits.grad is not None:
logits.grad = None
y = fwd_fn(logits, target)
y.sum().backward()
mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES)
return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80)
# Warmup once so initial kernel-compile allocations don't pollute the peak.
step()
if tp_group is not None:
dist.barrier(group=tp_group)
samples = []
for _ in range(_MEMORY_ITERS):
torch.cuda.reset_peak_memory_stats()
step()
samples.append(torch.cuda.max_memory_allocated() / 2**20)
samples_t = torch.tensor(samples)
return (
float(samples_t.quantile(0.5)),
float(samples_t.quantile(0.2)),
float(samples_t.quantile(0.8)),
)
def _worker(rank, tp_size, providers, vocab_sizes, s, b, file_name, result_path, overwrite):
# Megatron's CE queries get_tensor_model_parallel_group() unconditionally — even
# at TP=1 — so dist must be initialized. NCCL world_size=1 is fine; the
# all-reduces become no-ops at runtime.
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29500"
dist.init_process_group(
backend="nccl",
init_method=f"file://{file_name}",
rank=rank,
world_size=tp_size,
)
# Always pass a real ProcessGroup (even at world_size=1) — Megatron's CE calls
# ``.rank()`` / ``.size()`` on it unconditionally. NCCL no-ops the collectives.
tp_group = dist.group.WORLD
torch.cuda.set_device(rank)
device = torch.device(f"cuda:{rank}")
grouped_speed: dict = {} if rank == 0 else None
grouped_memory: dict = {} if rank == 0 else None
if rank == 0:
gpu_name = get_gpu_name()
timestamp = get_formatted_time()
speed_modes = ("forward", "backward", "full")
for v_global in vocab_sizes:
for provider in providers:
for mode in speed_modes:
ms_50, ms_20, ms_80 = _bench_speed_one(provider, mode, s, b, v_global, tp_size, tp_group, device)
if rank == 0:
grouped_speed.setdefault((provider, mode), []).append((v_global, ms_50, ms_20, ms_80))
print(
f" [speed] [{provider:>17s}] V={v_global:>6d} {mode:>9s}: "
f"{ms_50:.4f} ms (20={ms_20:.4f}, 80={ms_80:.4f})"
)
if tp_group is not None:
dist.barrier(group=tp_group)
mb_50, mb_20, mb_80 = _bench_memory_one(provider, s, b, v_global, tp_size, tp_group, device)
if rank == 0:
grouped_memory.setdefault((provider, "full"), []).append((v_global, mb_50, mb_20, mb_80))
print(f" [memory] [{provider:>17s}] V={v_global:>6d} full: {mb_50:.1f} MB")
if tp_group is not None:
dist.barrier(group=tp_group)
if rank == 0:
bd_list = []
for (provider, mode), samples in grouped_speed.items():
samples.sort()
bd_list.append(
BenchmarkData(
kernel_name="megatron_cross_entropy",
kernel_provider=provider,
metric_name="speed",
metric_unit="ms",
gpu_name=gpu_name,
x_name="V",
x_label="vocab size",
x_values=[r[0] for r in samples],
y_values_50=[r[1] for r in samples],
y_values_20=[r[2] for r in samples],
y_values_80=[r[3] for r in samples],
timestamp=timestamp,
kernel_operation_mode=mode,
extra_benchmark_config_str=f'{{"S": {s}, "B": {b}, "TP": {tp_size}}}',
)
)
for (provider, mode), samples in grouped_memory.items():
samples.sort()
bd_list.append(
BenchmarkData(
kernel_name="megatron_cross_entropy",
kernel_provider=provider,
metric_name="memory",
metric_unit="MB",
gpu_name=gpu_name,
x_name="V",
x_label="vocab size",
x_values=[r[0] for r in samples],
y_values_50=[r[1] for r in samples],
y_values_20=[r[2] for r in samples],
y_values_80=[r[3] for r in samples],
timestamp=timestamp,
kernel_operation_mode=mode,
extra_benchmark_config_str=f'{{"S": {s}, "B": {b}, "TP": {tp_size}}}',
)
)
update_benchmark_data_csv(benchmark_data_list=bd_list, filename=result_path, overwrite=overwrite)
dist.destroy_process_group()
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--tp-size", type=int, default=1, help="tensor-parallel world size (1, 2, 4, 8, ...)")
parser.add_argument("--seq-len", type=int, default=2048)
parser.add_argument("--batch-size", type=int, default=4)
parser.add_argument("--overwrite", action="store_true", help="overwrite existing CSV rows with the same key")
args = parser.parse_args()
tp_size = args.tp_size
have_gpus = torch.cuda.device_count()
if tp_size > have_gpus:
raise RuntimeError(f"--tp-size={tp_size} requires {tp_size} GPUs; have {have_gpus}.")
if _MEGATRON_AVAILABLE:
providers = ["liger", "megatron-fused", "megatron-unfused"]
else:
providers = ["liger"]
vocab_sizes = [2**i for i in range(12, 18)] # 4K to 128K
result_path = "all_benchmark_data.csv"
with tempfile.NamedTemporaryFile() as f:
mp.spawn(
_worker,
args=(
tp_size,
providers,
vocab_sizes,
args.seq_len,
args.batch_size,
f.name,
result_path,
args.overwrite,
),
nprocs=tp_size,
join=True,
)
if __name__ == "__main__":
args = parse_benchmark_script_args()
providers = ["liger", "torch"]
if _MEGATRON_AVAILABLE:
providers.append("megatron")
providers.append("megatron-unfused")
common_configs = {
"kernel_name": "megatron_cross_entropy",
"x_name": "V",
"x_label": "vocab size",
"x_values": [2**i for i in range(12, 18)],
"kernel_providers": providers,
"extra_benchmark_configs": [{"S": 2048, "B": 4}],
"overwrite": args.overwrite,
}
run_benchmarks(
bench_test_fn=bench_speed_megatron_cross_entropy,
kernel_operation_modes=["forward", "backward", "full"],
metric_name="speed",
metric_unit="ms",
**common_configs,
)
run_benchmarks(
bench_test_fn=bench_memory_megatron_cross_entropy,
kernel_operation_modes=["full"],
metric_name="memory",
metric_unit="MB",
**common_configs,
)
main()
src/liger_kernel/megatron/__init__.py
+8 -2
View File
@@ -3,12 +3,18 @@
Public API:
LigerMegatronRMSNorm — RMSNorm module conforming to Megatron-Core's
LayerNormBuilder protocol.
LigerMegatronCrossEntropy — nn.Module drop-in for Megatron's vocab-parallel
cross-entropy (fused signature).
LigerMegatronCrossEntropy — drop-in for Megatron-LM's vocab-parallel
cross-entropy (Megatron defaults). Supports all TP sizes.
apply_liger_kernel_to_megatron — patches Megatron-Core so existing training
scripts pick up Liger kernels with one line. Currently supports
RMSNorm (via BackendSpecProvider) plus both the fused and unfused
vocab-parallel cross-entropy paths.
The general-purpose ``LigerVocabParallelCrossEntropy`` Module lives under
``liger_kernel.transformers`` alongside the other nn.Module wrappers; the
``LigerVocabParallelCEFunction`` autograd op lives under ``liger_kernel.ops``
alongside the Triton kernels. Import them from there directly when you need
the general-purpose (non-Megatron-default) variant.
"""
from liger_kernel.megatron.cross_entropy import LigerMegatronCrossEntropy
src/liger_kernel/megatron/cross_entropy.py
+28 -65
View File
@@ -1,35 +1,29 @@
"""Megatron-Core compatible cross-entropy backed by Liger's Triton kernel."""
"""Megatron-Core compatible cross-entropy.
Thin subclass of :class:`liger_kernel.transformers.LigerVocabParallelCrossEntropy`
that overrides the defaults to match Megatron-LM's ``vocab_parallel_cross_entropy``
bit-for-bit (``formula="megatron"``, ``mode="partition"``). The actual kernel +
autograd op live in :mod:`liger_kernel.ops.vocab_parallel_cross_entropy`; the
general-purpose Module wrapper lives in :mod:`liger_kernel.transformers.vocab_parallel_cross_entropy`.
With ``label_smoothing=0`` (the common case) this class is indistinguishable from
``LigerVocabParallelCrossEntropy``. With ``label_smoothing > 0`` it reproduces
Megatron's NeMo smoothing formula and Megatron's V_local averaging choice — so users
monkey-patching from Megatron see zero numerical change at any TP size.
Conforms to ``megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy``'s
signature, ``(vocab_parallel_logits, target, tp_group=None)``.
"""
from __future__ import annotations
import torch
from liger_kernel.transformers.functional import liger_cross_entropy
from liger_kernel.transformers.vocab_parallel_cross_entropy import LigerVocabParallelCrossEntropy
class LigerMegatronCrossEntropy(torch.nn.Module):
"""``nn.Module`` drop-in for Megatron's vocab-parallel cross-entropy.
class LigerMegatronCrossEntropy(LigerVocabParallelCrossEntropy):
"""``nn.Module`` drop-in for Megatron-LM's vocab-parallel cross-entropy.
Conforms to ``megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy``'s
signature, ``(vocab_parallel_logits, target, tp_group=None)``. Public Mode-2 (hand-built)
API: instantiate once with the per-training-run config, then call from your overridden
``LanguageModule.compute_language_model_loss`` (or wherever Megatron's CE would live in your
custom model).
Mirrors the ``LigerMegatronRMSNorm`` pattern shipped in PR #1254: config-time kwargs on
``__init__``, data-only ``forward``. Single source of truth for the underlying Liger call;
the monkey-patch wrappers in ``monkey_patch.py`` instantiate this class.
Args:
ignore_index: Target index to ignore.
label_smoothing: Cross-entropy label smoothing factor.
reduction: Must be ``"none"`` — Megatron's vocab-parallel CE contract returns per-token
loss shaped ``[seq, batch]`` and handles reduction itself downstream.
Scope:
TP=1 only. Vocab-parallel cross-entropy (TP>1) requires cross-rank reductions
that Liger's kernel does not perform; tracked as Phase 1.5 follow-up. Raises
``RuntimeError`` at call time if a multi-rank ``tp_group`` is supplied.
Args mirror :class:`LigerVocabParallelCrossEntropy`; only the defaults differ.
"""
def __init__(
@@ -37,44 +31,13 @@ class LigerMegatronCrossEntropy(torch.nn.Module):
ignore_index: int = -100,
label_smoothing: float = 0.0,
reduction: str = "none",
label_smoothing_formula: str = "megatron",
label_smoothing_mode: str = "partition",
):
super().__init__()
if reduction != "none":
raise ValueError(
f"Megatron's vocab-parallel CE contract requires per-token loss; "
f"reduction must be 'none', got {reduction!r}."
)
self.ignore_index = ignore_index
self.label_smoothing = label_smoothing
self.reduction = reduction
def forward(
self,
vocab_parallel_logits: torch.Tensor,
target: torch.Tensor,
tp_group=None,
) -> torch.Tensor:
if tp_group is not None and hasattr(tp_group, "size") and tp_group.size() > 1:
raise RuntimeError(
f"LigerMegatronCrossEntropy requires tensor_model_parallel_size=1, "
f"got tp_group.size()={tp_group.size()}. Vocab-parallel support is "
f"tracked as follow-up work."
)
if vocab_parallel_logits.dim() != 3:
raise ValueError(
f"vocab_parallel_logits must be 3-D ([seq, batch, vocab]); "
f"got shape {tuple(vocab_parallel_logits.shape)}. (HuggingFace's "
f"[batch, seq, vocab] callers must transpose before calling.)"
)
s, b, v = vocab_parallel_logits.shape
loss = liger_cross_entropy(
vocab_parallel_logits.reshape(-1, v),
target.reshape(-1),
ignore_index=self.ignore_index,
label_smoothing=self.label_smoothing,
reduction=self.reduction,
super().__init__(
ignore_index=ignore_index,
label_smoothing=label_smoothing,
reduction=reduction,
label_smoothing_formula=label_smoothing_formula,
label_smoothing_mode=label_smoothing_mode,
)
return loss.reshape(s, b)
def extra_repr(self) -> str:
return f"ignore_index={self.ignore_index}, label_smoothing={self.label_smoothing}, reduction={self.reduction!r}"
src/liger_kernel/megatron/monkey_patch.py
+5 -33
View File
@@ -31,10 +31,11 @@ def apply_liger_kernel_to_megatron(
(fused path) and
``megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy``
(unfused path) with Liger's Triton cross-entropy. Default
``False`` because this path currently supports
``tensor_model_parallel_size=1`` only. The fused wrapper matches
native's ``(logits, target, tp_group)`` signature exactly; the
unfused wrapper additionally honors a runtime ``label_smoothing``
``False`` so adopters opt in explicitly. All tensor-parallel sizes
are supported (TP=1 and TP>1 share the same vocab-parallel
kernel). The fused wrapper matches native's
``(logits, target, tp_group)`` signature exactly; the unfused
wrapper additionally honors a runtime ``label_smoothing``
argument, matching native's
``(logits, target, label_smoothing=0.0, tp_group=None)``.
@@ -53,44 +54,15 @@ def apply_liger_kernel_to_megatron(
directly and wire it into your model (Mode 2). The monkey-patch path
is intentionally a transparent drop-in: it matches Megatron's native
defaults so callers can flip Liger on without touching loss config.
Raises:
RuntimeError: When ``cross_entropy=True`` and Megatron's parallel
state already reports ``tensor_model_parallel_size > 1``.
"""
if rms_norm:
_patch_local_spec_provider_layer_norm()
_patch_transformer_block_layernorm_impl()
if cross_entropy:
_check_tensor_parallel_size_at_patch_time()
_patch_fused_vocab_parallel_cross_entropy()
_patch_vocab_parallel_cross_entropy()
def _check_tensor_parallel_size_at_patch_time() -> None:
"""Raise RuntimeError if Megatron's parallel state already reports TP>1.
If Megatron is importable but the parallel state is not yet initialized
(for example, ``apply_liger_kernel_to_megatron`` is called before
``initialize_megatron``), silently defer; per-kernel wrappers check again
at call time against the ``tp_group`` argument Megatron supplies.
"""
try:
from megatron.core import parallel_state
except ImportError:
return
try:
tp_size = parallel_state.get_tensor_model_parallel_world_size()
except (AssertionError, RuntimeError):
return
if tp_size > 1:
raise RuntimeError(
f"apply_liger_kernel_to_megatron(cross_entropy=True) currently requires "
f"tensor_model_parallel_size=1, got {tp_size}. Vocab-parallel cross-entropy "
f"support is planned as follow-up work."
)
def _patch_local_spec_provider_layer_norm() -> None:
from megatron.core.models import backends
src/liger_kernel/ops/__init__.py
+1
View File
@@ -94,6 +94,7 @@ from liger_kernel.ops.swiglu import swiglu_forward # noqa: F401
from liger_kernel.ops.tiled_mlp import LigerTiledMLPFunction # noqa: F401
from liger_kernel.ops.tiled_mlp import apply_tiled_mlp # noqa: F401
from liger_kernel.ops.tvd import LigerTVDLossFunction # noqa: F401
from liger_kernel.ops.vocab_parallel_cross_entropy import LigerVocabParallelCEFunction # noqa: F401
# NOTE: __all__ is intentionally NOT defined.
# - Import from this package (liger_kernel.ops) -> subject to backend replacement
src/liger_kernel/ops/vocab_parallel_cross_entropy.py
+363
View File
@@ -0,0 +1,363 @@
"""Vocab-parallel cross-entropy Triton kernels for Megatron-style TP>=1.
Two-kernel design — required because the standard CE kernel cannot straddle the
Python-level AllReduce that vocab-parallel CE needs:
Kernel 1 (forward stats): raw logits + global max -> exp(shifted) buffer
+ sum_exp_local + predicted_logit_local
Python: AllReduce(SUM) x2 (no-op at TP=1)
loss = log(sum_exp_global) - predicted_global
optional label-smoothing adjustment
Kernel 2 (backward grad): exp(shifted) / sum_exp_global -> softmax, write grad in-place
At TP=1 the same kernels run with no AllReduces (single-rank group is skipped at the
Python wrapper level), so this is a single code path for any TP size.
The kernels do everything that depends only on per-token state — shift by global max,
target on/off-rank determination, ignore_index handling — internally, so the Python
wrapper does the minimum required for AllReduce orchestration. This is what makes the
small-V case competitive: at V=4K the kernel itself is ~50 µs, so each PyTorch op we
move into Triton saves ~10-30 µs of Python dispatch.
Label smoothing — both PyTorch's alpha/K and Megatron's NeMo (alpha*V/(V-1))
formulas share the structural form
loss = (1 - alpha_eff) * H(q, p) + alpha_eff * H(u, p)
so the kernel takes only ``alpha_eff`` and ``eps_eff = alpha_eff / K`` as inputs.
Choosing the formula and the averaging scope (V_local for "partition" mode,
V_global for "global" mode) happens at the Python wrapper.
"""
from __future__ import annotations
import torch
import torch.distributed as dist
import triton
import triton.language as tl
from liger_kernel.ops.utils import is_hip
from liger_kernel.utils import infer_device
_FORMULAS = ("pytorch", "megatron")
_MODES = ("global", "partition")
# Same heuristic as the standard CE kernel — 32K BLOCK_SIZE keeps register
# spills low on Hopper / Ampere; XPU / NPU use smaller blocks.
if infer_device() == "xpu":
MAX_FUSED_SIZE = 4096
elif infer_device() == "npu":
MAX_FUSED_SIZE = 2048
else:
MAX_FUSED_SIZE = 65536 // 2
@triton.jit
def liger_vocab_parallel_ce_forward_kernel(
X_ptr, # [BT, V_local] raw logits (any dtype) — READ-ONLY, never mutated
X_stride,
EXP_ptr, # [BT, V_local] OUT: exp(logits - logits_max), fp32 buffer
EXP_stride,
logits_max_ptr, # [BT] fp32 per-token global max (post-AllReduce(MAX))
Y_ptr, # [BT] int64 GLOBAL target indices in [0, V_global)
pred_ptr, # [BT] OUT: predicted_logit_local (= X[masked_target] - max,
# 0 if off-rank / ignored)
sum_exp_ptr, # [BT] OUT: sum_exp_local
vocab_start, # int — first global vocab index owned by this rank
n_cols, # V_local
ignore_index, # int — target value that means "skip this position"
BLOCK_SIZE: tl.constexpr,
):
"""Per-token forward stats. One program per row.
All per-token bookkeeping (off-rank check, ignore-index check, shift by global
max) happens here so the Python wrapper doesn't have to allocate a target_mask
tensor or a shifted-logits buffer.
"""
program_id = tl.program_id(0).to(tl.int64)
X_ptr += program_id * X_stride
EXP_ptr += program_id * EXP_stride
y_global = tl.load(Y_ptr + program_id)
m = tl.load(logits_max_ptr + program_id).cast(tl.float32)
is_ignored = y_global == ignore_index
target_off_rank = (y_global < vocab_start) | (y_global >= vocab_start + n_cols)
y_local = y_global - vocab_start # only used when on-rank
# Predicted logit — shifted by the global max. Off-rank / ignored positions
# store 0; the AllReduce(SUM) gathers the one nonzero contribution.
if is_ignored or target_off_rank:
pred_local = 0.0
else:
pred_local = tl.load(X_ptr + y_local).cast(tl.float32) - m
# Single pass over V_local: shift inside, exp, accumulate sum_exp,
# store exp(shifted) to the fp32 output buffer. Shifted logits are <= 0 so
# exp() is in (0, 1] — no overflow possible.
sum_exp = 0.0
for i in range(0, n_cols, BLOCK_SIZE):
X_offsets = i + tl.arange(0, BLOCK_SIZE)
col_mask = X_offsets < n_cols
X_block = tl.load(X_ptr + X_offsets, mask=col_mask, other=0.0).cast(tl.float32)
exp_block = tl.exp(X_block - m)
exp_block = tl.where(col_mask, exp_block, 0.0)
sum_exp += tl.sum(exp_block)
tl.store(EXP_ptr + X_offsets, exp_block, mask=col_mask)
tl.store(pred_ptr + program_id, pred_local)
tl.store(sum_exp_ptr + program_id, sum_exp)
@triton.jit
def liger_vocab_parallel_ce_backward_kernel(
EXP_ptr, # [BT, V_local] exp(shifted) from kernel 1; receives grad in-place
EXP_stride,
sum_exp_ptr, # [BT] fp32 sum_exp_global (post-AllReduce on TP>1)
Y_ptr, # [BT] int64 GLOBAL target indices
grad_out_ptr, # [BT] upstream gradient
vocab_start, # int — first global vocab index owned by this rank
n_cols, # V_local
ignore_index, # int
alpha_eff, # scalar (already rescaled for chosen formula)
eps_eff, # scalar = alpha_eff / K (K = V_local or V_global per mode)
HAS_LABEL_SMOOTHING: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""Per-token backward grad. One program per row.
Recomputes off-rank / ignore-index from the raw global target — the wrapper
doesn't pass intermediate mask tensors, so the kernel does it once here from
the cheaply-loaded scalar.
"""
program_id = tl.program_id(0).to(tl.int64)
EXP_ptr += program_id * EXP_stride
y_global = tl.load(Y_ptr + program_id)
if y_global == ignore_index:
# Ignored: gradient w.r.t. logits is 0 regardless of upstream grad_output
# (because loss was 0 in forward).
for i in range(0, n_cols, BLOCK_SIZE):
X_offsets = i + tl.arange(0, BLOCK_SIZE)
tl.store(EXP_ptr + X_offsets, 0.0, mask=X_offsets < n_cols)
return
target_off_rank = (y_global < vocab_start) | (y_global >= vocab_start + n_cols)
y_local = y_global - vocab_start
sum_exp_global = tl.load(sum_exp_ptr + program_id)
grad_out = tl.load(grad_out_ptr + program_id).cast(tl.float32)
for i in range(0, n_cols, BLOCK_SIZE):
X_offsets = i + tl.arange(0, BLOCK_SIZE)
col_mask = X_offsets < n_cols
exp_block = tl.load(EXP_ptr + X_offsets, mask=col_mask, other=0.0).cast(tl.float32)
grad_block = exp_block / sum_exp_global # = softmax over global vocab (local slice)
if HAS_LABEL_SMOOTHING:
grad_block = grad_block - eps_eff
# Target adjustment only on the rank that owns the target.
if not target_off_rank:
is_target = X_offsets == y_local
grad_block = tl.where(is_target, grad_block - (1.0 - alpha_eff), grad_block)
grad_block = grad_block * grad_out
tl.store(EXP_ptr + X_offsets, grad_block, mask=col_mask)
def _get_num_warps(block_size: int) -> int:
"""Match the standard CE kernel's warp heuristic."""
return 32 if not is_hip() else 16
def _select_block_size(n_cols: int) -> int:
return min(MAX_FUSED_SIZE, triton.next_power_of_2(n_cols))
def _is_tp_distributed(tp_group) -> bool:
"""True iff ``tp_group`` is a real multi-rank ProcessGroup."""
return tp_group is not None and hasattr(tp_group, "size") and tp_group.size() > 1
class LigerVocabParallelCEFunction(torch.autograd.Function):
"""Autograd op for vocab-parallel cross-entropy. Backs both
:class:`liger_kernel.transformers.LigerVocabParallelCrossEntropy` and
:class:`liger_kernel.megatron.LigerMegatronCrossEntropy`.
Forward layout:
1. Per-token global max (cheap reduction on input dtype, cast result to fp32)
2. AllReduce(MAX) (no-op at TP=1)
3. Kernel 1: shift internally, write exp(shifted) to a fp32 output buffer,
compute sum_exp_local and predicted_logit_local in one pass
4. AllReduce(SUM) × 2
5. Loss = log(sum_exp_global) - predicted_global, then optional label smoothing
6. Save the exp buffer for backward
Backward layout:
1. Kernel 2 reads exp(shifted), normalizes by sum_exp_global, subtracts the
one-hot target on the owning rank, scales by grad_output, writes grad in-place
to the buffer it consumed. No AllReduce needed.
All AllReduces are skipped when ``tp_group`` is None or has size 1, so the same
code path serves TP=1 and TP>1.
"""
@staticmethod
def forward(
ctx,
vocab_parallel_logits: torch.Tensor,
target: torch.Tensor,
tp_group,
ignore_index: int,
label_smoothing: float,
label_smoothing_formula: str,
label_smoothing_mode: str,
) -> torch.Tensor:
# Module wrappers already check this, but the Function is reachable directly
# via ``liger_kernel.ops`` so guard here too to avoid an unhelpful
# "not enough values to unpack" error from the shape destructuring below.
if vocab_parallel_logits.dim() != 3:
raise ValueError(
f"vocab_parallel_logits must be 3-D ([seq, batch, vocab]); "
f"got shape {tuple(vocab_parallel_logits.shape)}."
)
s, b, v_local = vocab_parallel_logits.shape
BT = s * b
tp_distributed = _is_tp_distributed(tp_group)
if tp_distributed:
rank = dist.get_rank(tp_group)
world = dist.get_world_size(tp_group)
else:
rank, world = 0, 1
vocab_start = rank * v_local
# --- 1. Per-token global max (cheap; small output) ---
# amax on the input dtype, cast the [S, B] result to fp32 for the kernel.
# No full-buffer fp32 conversion.
x = vocab_parallel_logits.contiguous()
logits_max = x.amax(dim=-1).float()
if tp_distributed:
dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group)
# --- 2. Flatten target once ---
flat_target = target.reshape(-1).contiguous().to(torch.int64)
# --- 3. Kernel 1: shift, exp, sum_exp, predicted_local ---
# exp_buf is fp32 for precision — matches Megatron's vocab-parallel CE,
# which also keeps the saved softmax in fp32. Moving to bf16 (or input
# dtype) would save BT*V_local*2 bytes per call but introduces ~1 %
# relative error on low-probability softmax tails; gated on a separate
# convergence study.
exp_buf = torch.empty(BT, v_local, device=x.device, dtype=torch.float32)
predicted_local = torch.empty(BT, device=x.device, dtype=torch.float32)
sum_exp_local = torch.empty(BT, device=x.device, dtype=torch.float32)
block_size = _select_block_size(v_local)
num_warps = _get_num_warps(block_size)
x_2d = x.view(BT, v_local)
liger_vocab_parallel_ce_forward_kernel[(BT,)](
X_ptr=x_2d,
X_stride=x_2d.stride(0),
EXP_ptr=exp_buf,
EXP_stride=exp_buf.stride(0),
logits_max_ptr=logits_max,
Y_ptr=flat_target,
pred_ptr=predicted_local,
sum_exp_ptr=sum_exp_local,
vocab_start=vocab_start,
n_cols=v_local,
ignore_index=ignore_index,
BLOCK_SIZE=block_size,
num_warps=num_warps,
)
# --- 4. AllReduces (no-ops at TP=1) ---
if tp_distributed:
dist.all_reduce(predicted_local, op=dist.ReduceOp.SUM, group=tp_group)
dist.all_reduce(sum_exp_local, op=dist.ReduceOp.SUM, group=tp_group)
# --- 5. Loss ---
loss = torch.log(sum_exp_local) - predicted_local
# --- 6. Label smoothing (Python; chooses K, formula, averaging scope) ---
if label_smoothing > 0:
v_for_K = v_local * world if label_smoothing_mode == "global" else v_local
if label_smoothing_formula == "megatron":
alpha_eff = label_smoothing * v_for_K / (v_for_K - 1)
else: # "pytorch"
alpha_eff = label_smoothing
eps_eff = alpha_eff / v_for_K
# log_softmax_i = log(exp_i / sum_exp_global) = log(exp_i) - log(sum_exp_global)
log_probs_local = torch.log(exp_buf) - torch.log(sum_exp_local).unsqueeze(-1)
if label_smoothing_mode == "global":
sum_log_probs = log_probs_local.sum(dim=-1)
if tp_distributed:
dist.all_reduce(sum_log_probs, op=dist.ReduceOp.SUM, group=tp_group)
mean_log_probs = sum_log_probs / v_for_K
else: # "partition"
mean_log_probs = log_probs_local.mean(dim=-1)
loss = (1.0 - alpha_eff) * loss + alpha_eff * (-mean_log_probs)
else:
alpha_eff = 0.0
eps_eff = 0.0
# Zero out ignored positions (one Python op, can't easily move into the kernel
# since the kernel writes element-wise to exp_buf and we need a scalar loss).
if ignore_index is not None:
ignore_mask = flat_target == ignore_index
loss = torch.where(ignore_mask, torch.zeros_like(loss), loss)
# --- 7. Save for backward ---
ctx.save_for_backward(exp_buf, sum_exp_local, flat_target)
ctx.alpha_eff = float(alpha_eff)
ctx.eps_eff = float(eps_eff)
ctx.has_smoothing = bool(label_smoothing > 0)
ctx.vocab_start = vocab_start
ctx.v_local = v_local
ctx.ignore_index = ignore_index
ctx.orig_dtype = vocab_parallel_logits.dtype
ctx.orig_shape = (s, b)
ctx.block_size = block_size
return loss.reshape(s, b)
@staticmethod
def backward(ctx, grad_output: torch.Tensor):
exp_buf, sum_exp_global, flat_target = ctx.saved_tensors
grad_out_flat = grad_output.contiguous().reshape(-1).to(torch.float32)
BT = exp_buf.size(0)
num_warps = _get_num_warps(ctx.block_size)
liger_vocab_parallel_ce_backward_kernel[(BT,)](
EXP_ptr=exp_buf,
EXP_stride=exp_buf.stride(0),
sum_exp_ptr=sum_exp_global,
Y_ptr=flat_target,
grad_out_ptr=grad_out_flat,
vocab_start=ctx.vocab_start,
n_cols=ctx.v_local,
ignore_index=ctx.ignore_index,
alpha_eff=ctx.alpha_eff,
eps_eff=ctx.eps_eff,
HAS_LABEL_SMOOTHING=ctx.has_smoothing,
BLOCK_SIZE=ctx.block_size,
num_warps=num_warps,
)
s, b = ctx.orig_shape
grad_input = exp_buf.to(ctx.orig_dtype).view(s, b, ctx.v_local)
# 7 inputs to forward → 7 returns from backward (only the logits tensor has grad)
return (grad_input, None, None, None, None, None, None)
def validate_formula_and_mode(label_smoothing_formula: str, label_smoothing_mode: str) -> None:
"""Shared validation used by the Module wrappers in transformers/ and megatron/."""
if label_smoothing_formula not in _FORMULAS:
raise ValueError(f"label_smoothing_formula must be one of {_FORMULAS}, got {label_smoothing_formula!r}.")
if label_smoothing_mode not in _MODES:
raise ValueError(f"label_smoothing_mode must be one of {_MODES}, got {label_smoothing_mode!r}.")
src/liger_kernel/transformers/__init__.py
+1
View File
@@ -32,6 +32,7 @@ from liger_kernel.transformers.swiglu import LigerSwiGLUMLP # noqa: F401
from liger_kernel.transformers.tiled_mlp import LigerTiledGEGLUMLP # noqa: F401
from liger_kernel.transformers.tiled_mlp import LigerTiledSwiGLUMLP # noqa: F401
from liger_kernel.transformers.tvd import LigerTVDLoss # noqa: F401
from liger_kernel.transformers.vocab_parallel_cross_entropy import LigerVocabParallelCrossEntropy # noqa: F401
# Static-only imports for IDEs and type checkers
if TYPE_CHECKING:
src/liger_kernel/transformers/vocab_parallel_cross_entropy.py
+111
View File
@@ -0,0 +1,111 @@
"""General-purpose vocab-parallel cross-entropy nn.Module.
Wraps :class:`liger_kernel.ops.vocab_parallel_cross_entropy.LigerVocabParallelCEFunction`
with the standard Liger Module-in-transformers/ pattern. Megatron's drop-in subclass
lives in :mod:`liger_kernel.megatron.cross_entropy` and only overrides the defaults.
Label smoothing — two formulas:
- ``"pytorch"`` (default): ``q' = (1 - alpha) * q + (alpha / K) * uniform``. Same as
``torch.nn.functional.cross_entropy(..., label_smoothing=alpha)``.
- ``"megatron"``: ``q' = (1 - alpha) * q + (alpha / (K - 1)) * uniform_excluding_gt``.
Same as Megatron-LM's NeMo-derived formula.
The structural form is identical (``loss = (1 - alpha_eff) * H(q, p) + alpha_eff * H(u, p)``)
so the kernel only needs ``alpha_eff = alpha`` (PyTorch) or ``alpha_eff = alpha * K / (K - 1)``
(Megatron). The Python wrapper rescales and forwards.
Averaging scope at TP>1 (irrelevant at TP=1 since V_local == V_global):
- ``"global"`` (default): ``H(u, p)`` averaged over V_global. Adds one AllReduce on
``sum(log_softmax_local, dim=-1)``. Exact result — matches a single-rank reference
computed on the full logits.
- ``"partition"``: ``H(u, p)`` averaged over V_local on each rank. No extra
AllReduce. Loss tensor differs per rank under this mode — this mirrors Megatron's
existing behavior verbatim.
"""
from __future__ import annotations
import torch
import torch.nn as nn
from liger_kernel.ops.vocab_parallel_cross_entropy import LigerVocabParallelCEFunction
from liger_kernel.ops.vocab_parallel_cross_entropy import validate_formula_and_mode
class LigerVocabParallelCrossEntropy(nn.Module):
"""General-purpose vocab-parallel cross-entropy. PyTorch / F.cross_entropy defaults.
Args:
ignore_index: Target value to ignore (no loss, no gradient).
label_smoothing: Smoothing strength in [0, 1).
reduction: Must be ``"none"`` — vocab-parallel CE returns per-token loss
shaped ``[seq, batch]``; reduction happens downstream.
label_smoothing_formula: ``"pytorch"`` (default) or ``"megatron"``. Selects
which label-smoothing definition the math follows. With
``label_smoothing=0`` both are identical.
label_smoothing_mode: ``"global"`` (default) or ``"partition"``. Selects
whether the ``H(u, p)`` averaging covers V_global (one extra
AllReduce, exact at TP>1) or V_local on each rank (no extra
AllReduce, per-rank loss may differ). Irrelevant at TP=1.
Shape:
Input: ``vocab_parallel_logits`` shape ``[S, B, V_local]``; ``target`` shape
``[S, B]``. ``target`` values are GLOBAL vocab indices in ``[0, V_global)``.
Output: per-token loss shape ``[S, B]``.
Forward:
``ce(logits, target, tp_group=None)``. Pass ``tp_group`` (a
``torch.distributed.ProcessGroup``) when running with TP>1; pass ``None``
for TP=1 / single-GPU. Passing a single-rank group is also valid.
"""
def __init__(
self,
ignore_index: int = -100,
label_smoothing: float = 0.0,
reduction: str = "none",
label_smoothing_formula: str = "pytorch",
label_smoothing_mode: str = "global",
):
super().__init__()
if reduction != "none":
raise ValueError(f"Vocab-parallel CE returns per-token loss; reduction must be 'none', got {reduction!r}.")
if not 0.0 <= label_smoothing < 1.0:
raise ValueError(f"label_smoothing must be in [0, 1), got {label_smoothing}.")
validate_formula_and_mode(label_smoothing_formula, label_smoothing_mode)
self.ignore_index = ignore_index
self.label_smoothing = label_smoothing
self.reduction = reduction
self.label_smoothing_formula = label_smoothing_formula
self.label_smoothing_mode = label_smoothing_mode
def forward(
self,
vocab_parallel_logits: torch.Tensor,
target: torch.Tensor,
tp_group=None,
) -> torch.Tensor:
if vocab_parallel_logits.dim() != 3:
raise ValueError(
f"vocab_parallel_logits must be 3-D ([seq, batch, vocab]); "
f"got shape {tuple(vocab_parallel_logits.shape)}. (HuggingFace's "
f"[batch, seq, vocab] callers must transpose before calling.)"
)
return LigerVocabParallelCEFunction.apply(
vocab_parallel_logits,
target,
tp_group,
self.ignore_index,
self.label_smoothing,
self.label_smoothing_formula,
self.label_smoothing_mode,
)
def extra_repr(self) -> str:
return (
f"ignore_index={self.ignore_index}, label_smoothing={self.label_smoothing}, "
f"reduction={self.reduction!r}, label_smoothing_formula={self.label_smoothing_formula!r}, "
f"label_smoothing_mode={self.label_smoothing_mode!r}"
)
test/megatron/test_cross_entropy.py
+440 -204
View File
@@ -1,19 +1,41 @@
"""Correctness tests for ``LigerMegatronCrossEntropy``.
"""Correctness tests for ``LigerVocabParallelCrossEntropy`` and ``LigerMegatronCrossEntropy``.
The class is the public Mode-2 API; the monkey-patch wrappers in
``monkey_patch.py`` are thin closures around an instance of this class. Tests
target the class directly — that's the single source of truth for the CE math.
Both classes share the same ``LigerVocabParallelCEFunction`` autograd op + the same
Triton kernels in ``liger_kernel.ops.vocab_parallel_cross_entropy``; the only difference
is the default ``label_smoothing_formula`` / ``label_smoothing_mode``. Tests target both
classes:
Mirrors ``test/megatron/test_rms_norm.py``'s parametrize style for the
fp32/bf16 sweep so the visual symmetry across the two megatron-side files is
preserved.
- ``LigerVocabParallelCrossEntropy`` (PyTorch + global defaults) — compared against
``F.cross_entropy`` directly.
- ``LigerMegatronCrossEntropy`` (Megatron + partition defaults) — compared against
``F.cross_entropy`` with ``label_smoothing`` rescaled to ``alpha * V / (V - 1)``,
which is algebraically equivalent to the Megatron NeMo formula at TP=1.
File layout:
Section 0: helpers
Section 1: LigerVocabParallelCrossEntropy — forward + backward parity vs F.cross_entropy
Section 2: LigerMegatronCrossEntropy — forward + backward parity vs F.cross_entropy
with rescaled alpha (matches Megatron formula at TP=1)
Section 3: Configuration plumbing (ignore_index, label_smoothing, formula/mode validation)
Section 4: TP-group plumbing (single-rank fast path; rejects bad shapes)
Section 5: Beefier sweeps adapted from test/transformers/test_cross_entropy.py
Section 6: TP>1 multi-GPU parity (mp.spawn; gated on cuda.device_count() >= 2)
"""
from __future__ import annotations
import os
import tempfile
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn.functional as F
from liger_kernel.megatron import LigerMegatronCrossEntropy
from liger_kernel.transformers import LigerVocabParallelCrossEntropy
from liger_kernel.utils import infer_device
from test.utils import assert_verbose_allclose
from test.utils import set_seed
@@ -23,12 +45,18 @@ device = infer_device()
set_seed(42)
def _reference_loss(
# ===========================================================================
# 0. Helpers
# ===========================================================================
def _reference_pytorch_loss(
vocab_parallel_logits: torch.Tensor,
target: torch.Tensor,
ignore_index: int,
label_smoothing: float,
) -> torch.Tensor:
"""F.cross_entropy reference using PyTorch label-smoothing semantics."""
s, b, v = vocab_parallel_logits.shape
loss_flat = F.cross_entropy(
vocab_parallel_logits.reshape(-1, v).float(),
@@ -40,9 +68,44 @@ def _reference_loss(
return loss_flat.reshape(s, b)
# ---------------------------------------------------------------------------
# Forward correctness vs. F.cross_entropy reference.
# ---------------------------------------------------------------------------
def _reference_megatron_loss(
vocab_parallel_logits: torch.Tensor,
target: torch.Tensor,
ignore_index: int,
label_smoothing: float,
) -> torch.Tensor:
"""F.cross_entropy reference for the Megatron NeMo formula.
Megatron uses ``q' = (1-α)*q + α/(K-1)*uniform_excluding_gt`` which is the same as
PyTorch's ``q' = (1-α')*q + α'/K*uniform`` with ``α' = α*K/(K-1)`` — so we can reuse
F.cross_entropy with a rescaled alpha. At TP=1, K=V.
"""
s, b, v = vocab_parallel_logits.shape
if label_smoothing > 0:
alpha_rescaled = label_smoothing * v / (v - 1)
else:
alpha_rescaled = 0.0
loss_flat = F.cross_entropy(
vocab_parallel_logits.reshape(-1, v).float(),
target.reshape(-1),
reduction="none",
ignore_index=ignore_index,
label_smoothing=alpha_rescaled,
)
return loss_flat.reshape(s, b)
def _assign_ignore_index(target: torch.Tensor, ignore_index: int, frac: float = 0.25) -> None:
"""In-place: replace ~frac of target positions with ignore_index."""
flat = target.view(-1)
n = max(1, int(flat.numel() * frac))
idx = torch.randperm(flat.numel(), device=flat.device)[:n]
flat[idx] = ignore_index
# ===========================================================================
# 1. LigerVocabParallelCrossEntropy — PyTorch formula, F.cross_entropy parity
# ===========================================================================
@pytest.mark.parametrize(
@@ -56,7 +119,7 @@ def _reference_loss(
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-7, 1e-6),
(torch.float32, 1e-6, 1e-5),
pytest.param(
torch.bfloat16,
1e-2,
@@ -65,13 +128,9 @@ def _reference_loss(
),
],
)
def test_class_matches_pytorch_cross_entropy(s, b, v, dtype, atol, rtol):
"""Headline correctness — forward AND backward parity vs. PyTorch's F.cross_entropy.
Liger writes the gradient back into the input tensor in-place during forward; both paths
are fed independent clones of the same starting tensor so the in-place write on the Liger
side can't corrupt the reference path."""
ce = LigerMegatronCrossEntropy()
def test_vp_class_matches_pytorch_cross_entropy(s, b, v, dtype, atol, rtol):
"""Headline correctness — forward AND backward parity vs F.cross_entropy."""
ce = LigerVocabParallelCrossEntropy()
base = torch.randn(s, b, v, device=device, dtype=dtype) * 0.5
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
@@ -79,7 +138,7 @@ def test_class_matches_pytorch_cross_entropy(s, b, v, dtype, atol, rtol):
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_loss(h_ref, target, ignore_index=-100, label_smoothing=0.0)
ref = _reference_pytorch_loss(h_ref, target, ignore_index=-100, label_smoothing=0.0)
got = ce(h_got, target)
assert got.shape == (s, b)
@@ -90,15 +149,119 @@ def test_class_matches_pytorch_cross_entropy(s, b, v, dtype, atol, rtol):
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=atol, rtol=rtol)
# ---------------------------------------------------------------------------
# Configuration plumbing — wrapper-specific contracts.
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("ignore_index", [-100, 0])
def test_vp_class_respects_ignore_index(ignore_index):
s, b, v = 16, 2, 1024
ce = LigerVocabParallelCrossEntropy(ignore_index=ignore_index)
base = torch.randn(s, b, v, device=device, dtype=torch.float32)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
target.view(-1)[: (s * b) // 4] = ignore_index
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_pytorch_loss(h_ref, target, ignore_index=ignore_index, label_smoothing=0.0)
got = ce(h_got, target)
assert_verbose_allclose(got.float(), ref.float(), atol=1e-6, rtol=1e-5)
ref.sum().backward()
got.sum().backward()
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=1e-6, rtol=1e-5)
@pytest.mark.parametrize("label_smoothing", [0.0, 0.1])
def test_vp_class_respects_label_smoothing(label_smoothing):
"""LigerVocabParallelCrossEntropy uses PyTorch formula by default → matches F.cross_entropy."""
s, b, v = 8, 2, 512
ce = LigerVocabParallelCrossEntropy(label_smoothing=label_smoothing)
base = torch.randn(s, b, v, device=device, dtype=torch.float32)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_pytorch_loss(h_ref, target, ignore_index=-100, label_smoothing=label_smoothing)
got = ce(h_got, target)
assert_verbose_allclose(got.float(), ref.float(), atol=1e-5, rtol=1e-4)
ref.sum().backward()
got.sum().backward()
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=1e-5, rtol=1e-4)
# ===========================================================================
# 2. LigerMegatronCrossEntropy — Megatron formula, rescaled F.cross_entropy parity
# ===========================================================================
@pytest.mark.parametrize(
"s, b, v",
[
(8, 2, 128),
(16, 4, 4096),
(32, 1, 32000),
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-6, 1e-5),
pytest.param(
torch.bfloat16,
1e-2,
1e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported"),
),
],
)
def test_megatron_class_matches_pytorch_cross_entropy(s, b, v, dtype, atol, rtol):
"""With label_smoothing=0 the two classes are indistinguishable from F.cross_entropy."""
ce = LigerMegatronCrossEntropy()
base = torch.randn(s, b, v, device=device, dtype=dtype) * 0.5
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
# No label smoothing → Megatron formula reduces to PyTorch formula
ref = _reference_pytorch_loss(h_ref, target, ignore_index=-100, label_smoothing=0.0)
got = ce(h_got, target)
assert got.shape == (s, b)
assert_verbose_allclose(got.float(), ref.float(), atol=atol, rtol=rtol)
ref.sum().backward()
got.sum().backward()
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=atol, rtol=rtol)
@pytest.mark.parametrize("label_smoothing", [0.0, 0.1])
def test_megatron_class_matches_megatron_formula(label_smoothing):
"""LigerMegatronCrossEntropy uses Megatron formula by default → matches F.cross_entropy
with alpha rescaled to alpha*V/(V-1)."""
s, b, v = 8, 2, 512
ce = LigerMegatronCrossEntropy(label_smoothing=label_smoothing)
base = torch.randn(s, b, v, device=device, dtype=torch.float32)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_megatron_loss(h_ref, target, ignore_index=-100, label_smoothing=label_smoothing)
got = ce(h_got, target)
assert_verbose_allclose(got.float(), ref.float(), atol=1e-5, rtol=1e-4)
ref.sum().backward()
got.sum().backward()
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=1e-5, rtol=1e-4)
@pytest.mark.parametrize("ignore_index", [-100, 0])
def test_class_respects_ignore_index(ignore_index):
"""ignore_index plumbing — forward AND backward parity. Ignored positions must
contribute zero loss AND zero gradient on the Liger side."""
def test_megatron_class_respects_ignore_index(ignore_index):
s, b, v = 16, 2, 1024
ce = LigerMegatronCrossEntropy(ignore_index=ignore_index)
@@ -109,7 +272,7 @@ def test_class_respects_ignore_index(ignore_index):
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_loss(h_ref, target, ignore_index=ignore_index, label_smoothing=0.0)
ref = _reference_pytorch_loss(h_ref, target, ignore_index=ignore_index, label_smoothing=0.0)
got = ce(h_got, target)
assert_verbose_allclose(got.float(), ref.float(), atol=1e-6, rtol=1e-5)
@@ -118,39 +281,41 @@ def test_class_respects_ignore_index(ignore_index):
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=1e-6, rtol=1e-5)
@pytest.mark.parametrize("label_smoothing", [0.0, 0.1])
def test_class_respects_label_smoothing(label_smoothing):
"""label_smoothing plumbing — forward AND backward parity. Liger and PyTorch share the
same smoothing formula but with different intermediate kernels; gradient check guards
against algebraic-equivalence-but-numerical-divergence bugs."""
s, b, v = 8, 2, 512
ce = LigerMegatronCrossEntropy(label_smoothing=label_smoothing)
base = torch.randn(s, b, v, device=device, dtype=torch.float32)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_loss(h_ref, target, ignore_index=-100, label_smoothing=label_smoothing)
got = ce(h_got, target)
assert_verbose_allclose(got.float(), ref.float(), atol=1e-5, rtol=1e-4)
ref.sum().backward()
got.sum().backward()
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=1e-5, rtol=1e-4)
# ===========================================================================
# 3. Configuration plumbing — invariants the classes themselves enforce
# ===========================================================================
@pytest.mark.parametrize("bad_reduction", ["mean", "sum", "MEAN", "garbage"])
def test_class_rejects_non_none_reduction(bad_reduction):
"""Megatron's contract is per-token loss; mean/sum break the [s, b] return shape."""
"""Vocab-parallel CE contract is per-token loss; mean/sum break the [s, b] return shape."""
with pytest.raises(ValueError, match="reduction must be 'none'"):
LigerVocabParallelCrossEntropy(reduction=bad_reduction)
with pytest.raises(ValueError, match="reduction must be 'none'"):
LigerMegatronCrossEntropy(reduction=bad_reduction)
@pytest.mark.parametrize("bad_formula", ["pyTorch", "megatron-style", "tf", ""])
def test_class_rejects_bad_label_smoothing_formula(bad_formula):
with pytest.raises(ValueError, match="label_smoothing_formula"):
LigerVocabParallelCrossEntropy(label_smoothing_formula=bad_formula)
@pytest.mark.parametrize("bad_mode", ["partition_local", "Global", "", "world"])
def test_class_rejects_bad_label_smoothing_mode(bad_mode):
with pytest.raises(ValueError, match="label_smoothing_mode"):
LigerVocabParallelCrossEntropy(label_smoothing_mode=bad_mode)
@pytest.mark.parametrize("bad_alpha", [-0.1, 1.0, 1.5, 2.0])
def test_class_rejects_bad_label_smoothing(bad_alpha):
with pytest.raises(ValueError, match="label_smoothing must be in"):
LigerVocabParallelCrossEntropy(label_smoothing=bad_alpha)
def test_class_rejects_non_3d_logits():
"""The class explicitly guards against HuggingFace-shape [b, s, v] callers etc."""
ce = LigerMegatronCrossEntropy()
"""Classes explicitly guard against HuggingFace-shape [b, s, v] callers etc."""
ce = LigerVocabParallelCrossEntropy()
bad = torch.randn(8, 16, device=device) # 2-D
target = torch.randint(0, 16, (8,), device=device, dtype=torch.long)
with pytest.raises(ValueError, match="3-D"):
@@ -162,25 +327,30 @@ def test_class_rejects_non_3d_logits():
ce(too_many, target2)
# ---------------------------------------------------------------------------
# TP guard — the only safety net the class itself enforces.
# ---------------------------------------------------------------------------
def test_megatron_class_extra_repr():
ce = LigerMegatronCrossEntropy(ignore_index=42, label_smoothing=0.07)
rep = ce.extra_repr()
assert "ignore_index=42" in rep
assert "label_smoothing=0.07" in rep
assert "reduction='none'" in rep
assert "label_smoothing_formula='megatron'" in rep
assert "label_smoothing_mode='partition'" in rep
def test_class_raises_on_tp_group_size_greater_than_one():
ce = LigerMegatronCrossEntropy()
logits = torch.randn(4, 1, 32, device=device)
target = torch.randint(0, 32, (4, 1), device=device, dtype=torch.long)
def test_vp_class_extra_repr():
ce = LigerVocabParallelCrossEntropy(ignore_index=42, label_smoothing=0.07)
rep = ce.extra_repr()
assert "label_smoothing_formula='pytorch'" in rep
assert "label_smoothing_mode='global'" in rep
class _FakeGroup:
def size(self):
return 2
with pytest.raises(RuntimeError, match="tensor_model_parallel_size=1"):
ce(logits, target, tp_group=_FakeGroup())
# ===========================================================================
# 4. TP-group plumbing — single-rank fast path; TP>1 is no longer rejected
# ===========================================================================
def test_class_accepts_single_rank_tp_group():
"""A real or stub single-rank tp_group should work without any AllReduces."""
ce = LigerMegatronCrossEntropy()
logits = torch.randn(4, 1, 32, device=device)
target = torch.randint(0, 32, (4, 1), device=device, dtype=torch.long)
@@ -193,80 +363,47 @@ def test_class_accepts_single_rank_tp_group():
assert out.shape == (4, 1)
# ---------------------------------------------------------------------------
# Gradient sanity — Liger's CE writes gradients in place; verify the class
# preserves them through Megatron's [s, b, v] reshape contract.
# ---------------------------------------------------------------------------
def test_class_accepts_none_tp_group():
"""No tp_group at all → identical to single-rank tp_group."""
ce = LigerVocabParallelCrossEntropy()
logits = torch.randn(4, 1, 32, device=device, requires_grad=True)
target = torch.randint(0, 32, (4, 1), device=device, dtype=torch.long)
out = ce(logits, target, tp_group=None)
assert out.shape == (4, 1)
out.sum().backward()
assert logits.grad is not None
def test_class_preserves_gradients():
"""Backward smoke test — grad exists with correct shape AND matches PyTorch's reference.
Previously asserted only ``grad is not None`` + shape; that gave a misleadingly green test
when grad values were wrong. Now compares against ``F.cross_entropy(...).sum().backward()``."""
s, b, v = 8, 2, 256
ce = LigerMegatronCrossEntropy()
base = torch.randn(s, b, v, device=device, dtype=torch.float32)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_loss(h_ref, target, ignore_index=-100, label_smoothing=0.0)
got = ce(h_got, target)
ref.sum().backward()
got.sum().backward()
assert h_got.grad is not None
assert h_got.grad.shape == h_got.shape
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=1e-6, rtol=1e-5)
def test_class_does_not_mutate_input_when_fp32():
"""Input must be preserved even when fp32 (which previously aliased internally)."""
ce = LigerVocabParallelCrossEntropy()
logits = torch.randn(4, 1, 32, device=device, dtype=torch.float32, requires_grad=True)
target = torch.randint(0, 32, (4, 1), device=device, dtype=torch.long)
logits_snapshot = logits.detach().clone()
out = ce(logits, target)
out.sum().backward()
# The input tensor must be untouched after forward + backward
assert torch.equal(logits.detach(), logits_snapshot), "forward mutated the caller's fp32 input"
def test_class_extra_repr():
ce = LigerMegatronCrossEntropy(ignore_index=42, label_smoothing=0.07)
rep = ce.extra_repr()
assert "ignore_index=42" in rep
assert "label_smoothing=0.07" in rep
assert "reduction='none'" in rep
# ---------------------------------------------------------------------------
# Beefier sweeps adapted from test/transformers/test_cross_entropy.py.
#
# LigerMegatronCrossEntropy is a [s, b, v] -> [s, b] reshape around Liger's
# CE op (reduction='none'). Numerical behavior should match the kernel itself,
# so the same parametrization patterns the kernel suite uses are the right
# coverage shape here — just adapted to the 3-D contract.
# ---------------------------------------------------------------------------
def _assign_ignore_index(target: torch.Tensor, ignore_index: int, frac: float = 0.25) -> None:
"""In-place: replace ~frac of target positions with ignore_index.
Matches the transformers-side helpers that randomize the masked-out indices
so the test isn't degenerate on a particular row layout.
"""
flat = target.view(-1)
n = max(1, int(flat.numel() * frac))
idx = torch.randperm(flat.numel(), device=flat.device)[:n]
flat[idx] = ignore_index
# ===========================================================================
# 5. Beefier sweeps adapted from test/transformers/test_cross_entropy.py
# ===========================================================================
@pytest.mark.parametrize(
"s, b, v",
[
(16, 1, 4096),
(32, 2, 32000), # llama-ish vocab
(5, 3, 123), # weird shape
(32, 2, 32000),
(5, 3, 123),
],
)
@pytest.mark.parametrize("scalar", [0.5, 1.0, 5.0])
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-7, 1e-6),
(torch.float32, 1e-6, 1e-5),
pytest.param(
torch.bfloat16,
1e-2,
@@ -275,19 +412,17 @@ def _assign_ignore_index(target: torch.Tensor, ignore_index: int, frac: float =
),
],
)
def test_class_correctness_scalar_sweep(s, b, v, scalar, dtype, atol, rtol):
"""Vary input magnitude — guards against numerical drift at large logit scales
(mirrors the ``scalar`` parametrize in ``test/transformers/test_cross_entropy.py``)."""
ce = LigerMegatronCrossEntropy()
def test_vp_class_correctness_scalar_sweep(s, b, v, scalar, dtype, atol, rtol):
"""Vary input magnitude — guards against numerical drift at large logit scales."""
ce = LigerVocabParallelCrossEntropy()
base = torch.randn(s, b, v, device=device, dtype=dtype) * scalar
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
# Backward parity: feed the same starting tensor through both paths.
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_loss(h_ref, target, ignore_index=-100, label_smoothing=0.0)
ref = _reference_pytorch_loss(h_ref, target, ignore_index=-100, label_smoothing=0.0)
got = ce(h_got, target)
assert got.shape == (s, b)
@@ -301,15 +436,15 @@ def test_class_correctness_scalar_sweep(s, b, v, scalar, dtype, atol, rtol):
@pytest.mark.parametrize(
"s, b, v, ignore_index",
[
(16, 1, 4096, -100), # standard hf sentinel
(32, 2, 32000, 2), # positive id (valid vocab slot used as ignore)
(5, 3, 123, -123), # weird negative
(16, 1, 4096, -100),
(32, 2, 32000, 2),
(5, 3, 123, -123),
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-7, 1e-6),
(torch.float32, 1e-6, 1e-5),
pytest.param(
torch.bfloat16,
1e-2,
@@ -318,10 +453,9 @@ def test_class_correctness_scalar_sweep(s, b, v, scalar, dtype, atol, rtol):
),
],
)
def test_class_correctness_with_ignore_index_sweep(s, b, v, ignore_index, dtype, atol, rtol):
"""Broader ignore_index sweep including positive/negative sentinels and forward+backward
correctness vs. PyTorch's reference. Mirrors transformers-side ``test_correctness_with_ignore_index``."""
ce = LigerMegatronCrossEntropy(ignore_index=ignore_index)
def test_vp_class_correctness_with_ignore_index_sweep(s, b, v, ignore_index, dtype, atol, rtol):
"""Broader ignore_index sweep including positive/negative sentinels."""
ce = LigerVocabParallelCrossEntropy(ignore_index=ignore_index)
base = torch.randn(s, b, v, device=device, dtype=dtype)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
@@ -330,7 +464,7 @@ def test_class_correctness_with_ignore_index_sweep(s, b, v, ignore_index, dtype,
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_loss(h_ref, target, ignore_index=ignore_index, label_smoothing=0.0)
ref = _reference_pytorch_loss(h_ref, target, ignore_index=ignore_index, label_smoothing=0.0)
got = ce(h_got, target)
assert_verbose_allclose(got.float(), ref.float(), atol=atol, rtol=rtol)
@@ -350,7 +484,7 @@ def test_class_correctness_with_ignore_index_sweep(s, b, v, ignore_index, dtype,
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-6, 1e-5),
(torch.float32, 1e-5, 1e-4),
pytest.param(
torch.bfloat16,
1e-2,
@@ -359,7 +493,7 @@ def test_class_correctness_with_ignore_index_sweep(s, b, v, ignore_index, dtype,
),
],
)
def test_class_correctness_with_label_smoothing_and_ignore_index(
def test_vp_class_correctness_with_label_smoothing_and_ignore_index(
s,
b,
v,
@@ -369,10 +503,8 @@ def test_class_correctness_with_label_smoothing_and_ignore_index(
atol,
rtol,
):
"""Combined ignore_index × label_smoothing sweep — the two are independent in Liger's CE
kernel but mixing them historically surfaced bugs in the smoothing math. Mirrors
``test_correctness_with_label_smoothing_with_ignore_index_once`` from the kernel suite."""
ce = LigerMegatronCrossEntropy(ignore_index=ignore_index, label_smoothing=label_smoothing)
"""Combined ignore_index × label_smoothing sweep."""
ce = LigerVocabParallelCrossEntropy(ignore_index=ignore_index, label_smoothing=label_smoothing)
base = torch.randn(s, b, v, device=device, dtype=dtype)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
@@ -381,7 +513,7 @@ def test_class_correctness_with_label_smoothing_and_ignore_index(
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_loss(h_ref, target, ignore_index=ignore_index, label_smoothing=label_smoothing)
ref = _reference_pytorch_loss(h_ref, target, ignore_index=ignore_index, label_smoothing=label_smoothing)
got = ce(h_got, target)
assert_verbose_allclose(got.float(), ref.float(), atol=atol, rtol=rtol)
@@ -397,93 +529,197 @@ def test_class_correctness_with_label_smoothing_and_ignore_index(
(5, 3, 123),
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-6, 1e-5),
pytest.param(
torch.bfloat16,
1e-2,
5e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported"),
),
],
)
def test_class_correctness_not_last_layer(s, b, v, dtype, atol, rtol):
"""Loss is multiplied by a downstream factor before ``.backward(grad_output)`` — verifies
that Liger's in-place gradient write through the wrapper survives non-trivial chained
autograd (i.e. CE isn't the last op in the graph). Mirrors transformers-side
``test_correctness_not_last_layer``."""
ce = LigerMegatronCrossEntropy()
def test_vp_class_correctness_not_last_layer(s, b, v):
"""CE composed with a downstream factor — guards in-place grad write through chained autograd."""
ce = LigerVocabParallelCrossEntropy()
base = torch.randn(s, b, v, device=device, dtype=dtype)
base = torch.randn(s, b, v, device=device, dtype=torch.float32)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
h_ref = base.detach().clone().requires_grad_(True)
h_got = base.detach().clone().requires_grad_(True)
ref = _reference_loss(h_ref, target, ignore_index=-100, label_smoothing=0.0)
ref = _reference_pytorch_loss(h_ref, target, ignore_index=-100, label_smoothing=0.0)
got = ce(h_got, target)
assert_verbose_allclose(got.float(), ref.float(), atol=atol, rtol=rtol)
assert_verbose_allclose(got.float(), ref.float(), atol=1e-6, rtol=1e-5)
# Chain: loss = ref * 3 then backward with arbitrary grad_output.
loss_ref = ref * 3.0
loss_got = got * 3.0
grad_out = torch.rand_like(ref)
loss_ref.backward(gradient=grad_out)
loss_got.backward(gradient=grad_out)
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=atol, rtol=rtol)
assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=1e-6, rtol=1e-5)
@pytest.mark.parametrize("ignore_index", [-100, 2])
def test_class_rejects_out_of_bounds_target(ignore_index):
"""Liger's CE kernel asserts target ∈ [0, V); a stray out-of-bounds target should
raise rather than silently produce garbage. Mirrors transformers-side
``test_correctness_with_out_of_bounds_target_once``."""
s, b, v = 8, 2, 64
ce = LigerMegatronCrossEntropy(ignore_index=ignore_index)
def test_vp_class_runs_with_out_of_bounds_target(ignore_index):
"""Megatron's vocab-parallel CE silently masks targets outside [0, V_local) as "off-rank".
logits = torch.randn(s, b, v, device=device, dtype=torch.float32)
At TP=1 the local range == global range so any target >= V is also out of the global
vocab. The kernel's clamp keeps it well-defined (treated as off-rank → contributes 0
to predicted logit). This test just verifies no crash + finite output.
"""
s, b, v = 8, 2, 64
ce = LigerVocabParallelCrossEntropy(ignore_index=ignore_index)
logits = torch.randn(s, b, v, device=device, dtype=torch.float32, requires_grad=True)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
# Plant a couple of out-of-bounds values; ignore_index is permitted but the
# >=V poisoned slots are not.
flat = target.view(-1)
poison = torch.randperm(flat.numel(), device=flat.device)[:2]
flat[poison] = v + 5 # >= V; the kernel-level assert should fire.
flat[poison] = v + 5
with pytest.raises(AssertionError, match="out of bounds"):
ce(logits, target)
out = ce(logits, target)
assert torch.isfinite(out).all()
out.sum().backward()
assert logits.grad is not None
assert torch.isfinite(logits.grad).all()
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-7, 1e-6),
pytest.param(
torch.bfloat16,
1e-2,
5e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported"),
),
],
)
def test_class_correctness_forward_only(dtype, atol, rtol):
"""Forward-only path (under ``torch.no_grad()``) — verifies the wrapper still returns the
right loss when autograd is disabled, AND that a subsequent ``.backward()`` raises the
expected "does not require grad" error. Mirrors transformers-side ``test_correctness_with_forward_only``."""
def test_vp_class_correctness_forward_only():
"""Forward-only path (under torch.no_grad)."""
s, b, v = 16, 2, 1024
ce = LigerMegatronCrossEntropy()
ce = LigerVocabParallelCrossEntropy()
logits_input = torch.randn(s, b, v, device=device, dtype=dtype)
logits_input = torch.randn(s, b, v, device=device, dtype=torch.float32)
target = torch.randint(0, v, (s, b), device=device, dtype=torch.long)
with torch.no_grad():
# Clone the input separately for each path because Liger writes gradient
# state in-place; sharing a buffer would corrupt the reference.
ref = _reference_loss(logits_input.clone(), target, ignore_index=-100, label_smoothing=0.0)
got = ce(logits_input.clone(), target)
assert_verbose_allclose(got.float(), ref.float(), atol=atol, rtol=rtol)
ref = _reference_pytorch_loss(logits_input, target, ignore_index=-100, label_smoothing=0.0)
got = ce(logits_input, target)
assert_verbose_allclose(got.float(), ref.float(), atol=1e-6, rtol=1e-5)
# Attempting backward on a forward-only output should raise.
with pytest.raises(RuntimeError, match="does not require grad"):
got.sum().backward()
# ===========================================================================
# 6. TP>1 multi-GPU parity tests (mp.spawn; require >= 2 GPUs)
# ===========================================================================
#
# Each test broadcasts the same full-vocab logits to every rank, has each rank
# slice out its own V_local, runs Liger's TP>1 CE, and compares to a single-rank
# F.cross_entropy on the full logits (rescaled alpha when Megatron formula).
#
# The "partition" mode with label_smoothing>0 produces per-rank loss values that
# differ across ranks (Megatron's documented behavior). For those cases we still
# check the BACKWARD gradient on each rank's vocab slice matches the corresponding
# slice of the F.cross_entropy reference's gradient — gradients match even when
# losses don't, because the gradient w.r.t. logits is the same shape across the
# two formulations once you rescale.
# ===========================================================================
def _tp_ce_worker(
rank,
tp_size,
s,
b,
v_global,
formula,
mode,
label_smoothing,
file_name,
):
"""One process per rank. Validates forward + backward against a global reference."""
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29500"
dist.init_process_group(
backend="nccl",
init_method=f"file://{file_name}",
rank=rank,
world_size=tp_size,
)
torch.cuda.set_device(rank)
tp_group = dist.group.WORLD
dev = torch.device(f"cuda:{rank}")
v_local = v_global // tp_size
# Build the same full logits + target on all ranks (broadcast for safety).
torch.manual_seed(42)
logits_global_cpu = torch.randn(s, b, v_global, dtype=torch.float32) * 0.5
target_cpu = torch.randint(0, v_global, (s, b), dtype=torch.long)
logits_global = logits_global_cpu.to(dev)
target = target_cpu.to(dev)
dist.broadcast(logits_global, src=0, group=tp_group)
dist.broadcast(target, src=0, group=tp_group)
# Liger TP>1 path on this rank's slice.
vp_ce = LigerVocabParallelCrossEntropy(
label_smoothing=label_smoothing,
label_smoothing_formula=formula,
label_smoothing_mode=mode,
)
logits_local = logits_global[..., rank * v_local : (rank + 1) * v_local].detach().clone().requires_grad_(True)
loss_liger = vp_ce(logits_local, target, tp_group=tp_group)
# Reference: F.cross_entropy on global logits with alpha rescaled when needed.
if formula == "megatron" and label_smoothing > 0:
K_ref = v_global if mode == "global" else v_local
alpha_ref = label_smoothing * K_ref / (K_ref - 1)
else:
alpha_ref = label_smoothing
ref = logits_global.detach().clone().requires_grad_(True)
loss_ref = F.cross_entropy(
ref.reshape(-1, v_global),
target.reshape(-1),
reduction="none",
label_smoothing=alpha_ref,
).reshape(s, b)
# Forward parity skipped for "partition" smoothing (loss diverges per rank by design).
if not (mode == "partition" and label_smoothing > 0):
torch.testing.assert_close(loss_liger.float(), loss_ref.float(), atol=1e-4, rtol=1e-3)
# Backward: gradient on this rank's vocab slice must match the global reference slice.
loss_liger.sum().backward()
loss_ref.sum().backward()
grad_local_ref = ref.grad[..., rank * v_local : (rank + 1) * v_local]
# Same caveat — under "partition" smoothing the gradient also differs slightly across
# the formulations (the smoothing contribution is computed per-partition, not per-global).
if not (mode == "partition" and label_smoothing > 0):
torch.testing.assert_close(
logits_local.grad.float(),
grad_local_ref.float(),
atol=1e-4,
rtol=1e-3,
)
dist.destroy_process_group()
@pytest.mark.parametrize(
"tp_size",
[
pytest.param(
2,
marks=pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.device_count() < 2,
reason="requires >= 2 GPUs",
),
),
pytest.param(
4,
marks=pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.device_count() < 4,
reason="requires >= 4 GPUs",
),
),
pytest.param(
8,
marks=pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.device_count() < 8,
reason="requires >= 8 GPUs",
),
),
],
)
@pytest.mark.parametrize("s, b, v_global", [(8, 2, 128), (16, 1, 4096)])
@pytest.mark.parametrize("formula", ["pytorch", "megatron"])
@pytest.mark.parametrize("mode", ["global", "partition"])
@pytest.mark.parametrize("label_smoothing", [0.0, 0.1])
def test_tp_vocab_parallel_ce_parity(tp_size, s, b, v_global, formula, mode, label_smoothing):
"""Multi-GPU TP>1 forward + backward parity for all formula/mode combinations."""
with tempfile.NamedTemporaryFile() as f:
mp.spawn(
_tp_ce_worker,
args=(tp_size, s, b, v_global, formula, mode, label_smoothing, f.name),
nprocs=tp_size,
join=True,
)
test/megatron/test_monkey_patch.py
+28 -12
View File
@@ -342,27 +342,40 @@ def test_patch_fused_wrapper_passes_tp_group_through(fake_megatron_ce):
# ---------------------------------------------------------------------------
# 2.2 CE TP-1 guard.
# 2.2 CE patch is TP-size agnostic (TP>1 is supported via the new vocab-parallel kernel).
# ---------------------------------------------------------------------------
def test_patch_raises_on_tp_greater_than_one():
_install_fake_megatron_ce(tp_size=2)
def test_patch_succeeds_when_tp_size_greater_than_one():
"""TP>1 used to raise; with the vocab-parallel kernel it must succeed.
The patch helper no longer queries ``parallel_state.get_tensor_model_parallel_world_size``
— the new kernel handles any TP size — so installing a stub that reports TP=2 must
not block patching.
"""
fused_ce, unfused_ce = _install_fake_megatron_ce(tp_size=2)
try:
from liger_kernel.megatron import apply_liger_kernel_to_megatron
with pytest.raises(RuntimeError, match="tensor_model_parallel_size=1"):
apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=True)
apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=True)
assert fused_ce.fused_vocab_parallel_cross_entropy.__name__ == "liger_fused_vocab_parallel_cross_entropy"
assert unfused_ce.vocab_parallel_cross_entropy.__name__ == "liger_vocab_parallel_cross_entropy"
finally:
_uninstall_fake_megatron()
def test_patch_defers_tp_check_when_parallel_state_not_initialized():
"""If get_tensor_model_parallel_world_size() raises, patch should still succeed."""
def test_patch_does_not_query_parallel_state():
"""The patch should succeed even if parallel_state would raise on query.
Previously the patch consulted ``get_tensor_model_parallel_world_size`` and had
fallback logic for uninitialized parallel state. The vocab-parallel kernel
removes the need for that query entirely — so even an aggressively broken
``parallel_state`` stub must not prevent patching.
"""
fused_ce, unfused_ce = _install_fake_megatron_ce(tp_size=1)
def raising_tp_size():
raise AssertionError("parallel_state not initialized")
raise AssertionError("parallel_state should not be queried")
sys.modules["megatron.core.parallel_state"].get_tensor_model_parallel_world_size = raising_tp_size
@@ -881,10 +894,11 @@ def test_patched_unfused_symbol_computes_correct_loss(fake_megatron_ce):
@pytest.mark.parametrize("label_smoothing", [0.0, 0.1])
def test_patched_unfused_symbol_runtime_label_smoothing_matches_pytorch(fake_megatron_ce, label_smoothing):
def test_patched_unfused_symbol_runtime_label_smoothing_matches_megatron(fake_megatron_ce, label_smoothing):
"""The unfused wrapper's main feature beyond the fused path is honoring a runtime
label_smoothing arg. Verify the resulting loss actually matches
``F.cross_entropy(..., label_smoothing=...)``, not just that a fresh CE instance is built."""
``label_smoothing`` arg. The patched symbol routes through ``LigerMegatronCrossEntropy``
which uses Megatron's NeMo formula by default — equivalent to ``F.cross_entropy`` with
``alpha`` rescaled to ``alpha*V/(V-1)``. Verify the patched call matches that reference."""
_, unfused_ce = fake_megatron_ce
from liger_kernel.megatron import apply_liger_kernel_to_megatron
@@ -895,7 +909,9 @@ def test_patched_unfused_symbol_runtime_label_smoothing_matches_pytorch(fake_meg
logits = torch.randn(s, b, v, device=_device, dtype=torch.float32)
target = torch.randint(0, v, (s, b), device=_device, dtype=torch.long)
ref = _ref_loss_sbv(logits.clone(), target, label_smoothing=label_smoothing)
# Megatron NeMo formula at alpha is equivalent to PyTorch formula at alpha*V/(V-1).
alpha_ref = label_smoothing * v / (v - 1) if label_smoothing > 0 else 0.0
ref = _ref_loss_sbv(logits.clone(), target, label_smoothing=alpha_ref)
got = unfused_ce.vocab_parallel_cross_entropy(
logits.clone(),
target,