Liger-Kernel

linkedin/Liger-Kernel

Fork 0

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

Files

T

History

Eric(Qun)andShao Tang 9743e507d1 Add PolyNorm operator (#901 )

Resolve #865

- Add LigerPolyNormFunction
- Add Accuracy Test and Perfomance Benchmark

Performance improvement vs PyTorch autograd:
- Speed: 12-40x faster
- Memory: 6.4x reduction (10x with in-place)

Reference: 
https://github.com/BryceZhuo/PolyCom/
https://arxiv.org/abs/2508.09148
https://arxiv.org/pdf/2411.03884v1

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Introduce **PolyNorm** (Type-II PolyCom, fix r=3) as a training-side
fused operator for Liger-Kernel.



<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
<img width="1227" height="255" alt="image"
src="https://github.com/user-attachments/assets/998f125f-ec72-448e-84ee-19536634800a"
/>

**Kernels:** `LigerPolyNormFunction` (Triton) with **fwd+bwd fused**.
- **Forward:** per row compute `mean(x^2)`, `mean(x^4)`, `mean(x^6)` →
cache `rstd_p` (p=1,2,3; 3 scalars/row); write `y` once (no intermediate
tensors).
- **Backward (closed-form):** per row compute `S_p = ⟨grad, x^p⟩`, then
dX = Σ_{p=1..3} w_p * [ p * x^{p-1} * rstd_p * grad − (p/H) * x^{2p-1} *
(rstd_p)^3 * S_p ]
    dW_p = rstd_p * S_p,  
    dB = sum(grad).  

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

Performance improvement vs PyTorch autograd:
- Speed: 12-40x faster
- Memory: 6.4x reduction (10x with in-place)

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

---------

Co-authored-by: Shao Tang <[email protected]>

2025-10-14 13:36:57 +08:00

data

Add GSPO (#845 )

2025-09-03 09:44:13 -07:00

scripts

Add PolyNorm operator (#901 )

2025-10-14 13:36:57 +08:00

__init__.py

first commit

2024-08-06 18:44:05 +00:00

benchmarks_visualizer.py

Multi-token Attention (#689 )

2025-05-26 18:56:52 +08:00

README.md

Visualizer Update (#717 )

2025-05-16 21:56:49 -07:00

README.md

Benchmarking Liger Kernels

Follow these steps to benchmark and visualize kernel performance:

Create a benchmark script
- Add your script under benchmark/scripts/
- Name it according to the kernel (e.g., benchmark_<kernel_name>.py)
Run the benchmark
- Results will be saved to benchmark/data/all_benchmark_data.csv
Example: Benchmarking KTO Loss
```
cd benchmark
python scripts/benchmark_kto_loss.py
```
Visualize results
- Use the visualization script with optional modes:
  - To target specific mode(s), pass --kernel-operation-mode one or more values.
  - If you omit --kernel-operation-mode, the script will:
    - For speed metrics: generate plots for all available modes (forward/backward/full).
    - For memory metrics: generate only the full plot.
Examples:
1. Specific modes (speed):
```
python benchmarks_visualizer.py \
    --kernel-name kto_loss \
    --metric-name speed \
    --kernel-operation-mode forward backward
```
1. All modes (speed):
```
python benchmarks_visualizer.py \
    --kernel-name kto_loss \
    --metric-name speed
```
1. Memory (always full):
```
python benchmarks_visualizer.py \
    --kernel-name kto_loss \
    --metric-name memory
```
View results
- Generated plots will be saved in benchmark/visualizations/