[Megatron] Add RMSNorm benchmark + measured H100 results

[vaibhavjindal]

Summary

Adds benchmark/scripts/benchmark_megatron_rms_norm.py — the RMSNorm parallel to the Megatron cross-entropy benchmark landed in #1207. Provides empirical speed/memory numbers so we can point at concrete data when claiming RMSNorm wins, instead of just citing the kernel sweep.

Output goes to the shared benchmark/data/all_benchmark_data.csv (tagged kernel_name="megatron_rms_norm"), so the standard visualizer renders the plots:

python benchmark/benchmarks_visualizer.py \
    --kernel-name megatron_rms_norm --metric-name speed
python benchmark/benchmarks_visualizer.py \
    --kernel-name megatron_rms_norm --metric-name memory

Providers compared

• liger — LigerMegatronRMSNorm (Liger's Triton RMSNorm via the Megatron-shaped wrapper from [Megatron] Add RMSNorm integration #1254)
• torch — vanilla torch.nn.RMSNorm
• megatron — Megatron's WrappedTorchNorm — its __new__ returns torch.nn.RMSNorm, so timings should match torch exactly. Included for explicit parity confirmation since WrappedTorchNorm is the symbol Liger displaces in the local-backend path.

If megatron-core is not installed, the megatron provider is silently dropped and the run proceeds with liger + torch.

Results on H100 80GB (S=4096, B=1, bf16)

60 rows committed in the CSV (5 hidden sizes × 3 providers × 4 measurements).

Speed — forward

H	liger	torch	speedup
1024	0.012 ms	0.013 ms	~flat
2048	0.018 ms	0.019 ms	~6%
4096	0.029 ms	0.033 ms	~12%
8192	0.051 ms	0.074 ms	~31%
16384	0.095 ms	0.149 ms	~36%

Liger forward wins, and the gap widens with hidden size.

Speed — full (fwd + bwd)

H	liger	torch	notes
1024	0.30 ms	0.075 ms	Liger SLOWER (Triton launch overhead)
4096	0.31 ms	0.20 ms	Liger slower
8192	0.32 ms	0.40 ms	crossover
16384	0.59 ms	0.77 ms	~24% faster

The flat Liger curve at small H is the giveaway: kernel launch overhead dominates. nn.RMSNorm's backward is a single fused C++/CUDA kernel; Liger's backward launches multiple Triton kernels (dx + dw reduction + element_mul). At small H the actual compute is tiny relative to per-launch overhead, so the math wins from Triton get drowned. At H ≥ ~6K the compute dominates and Liger wins.

Memory — full

Approximately neutral across all H — Liger uses 0.5–1% more than torch, within measurement noise.

nn.RMSNorm is already a single fused CUDA kernel with minimal intermediates, so Liger doesn't get the activation-memory win it gets when replacing eager-PyTorch RMSNorm (which materializes variance + rsqrt + scale separately). Speed is the win in this comparison, not memory.

Parity sanity check

torch and megatron rows are bit-identical across the entire sweep — exactly what we'd expect since WrappedTorchNorm is a factory that returns nn.RMSNorm. Confirms Liger is replacing the right baseline.

Honest read

Production LLMs run at H ≥ 4096. At those sizes:

• Forward: Liger wins ~12–36%
• Full: Liger wins from ~H=6K onward
• Memory: neutral

So Liger's RMSNorm is a real speed improvement for typical training shapes. The tiny-H regression on full is launch overhead, not a numerical/math issue.

Plots

Memory

Backward speed

Forward speed

Forward + Backward speed

Testing Done

• Hardware Type: H100 80GB HBM3
• Benchmark runs end-to-end on H100 with all 3 providers active
• Standard visualizer renders all 4 PNGs (3 speed modes + memory) from the shared CSV
• torch / megatron providers produce bit-identical numbers (parity confirmation)
• make checkstyle passes