mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
b5cdbf7e34
## Description Addresses #1067 (previously explored in #1068) This PR brings the LoRA normalization optimization across the finish line: - Rebased on top of modern `main` (cleanly resolving the #1116 benchmark conflicts that blocked #1068). - Incorporates maintainer feedback from @Tcc0403 on #1068 (passes `None` directly for unused gradient pointers). - Comprehensive unit tests: 24/24 passing on Float32 and BFloat16. - Hardware-validated on NVIDIA A10G and NVIDIA A100 GPUs (up to 1.38x backward speedup). Full co-author credit is preserved for @yukiu00 in the commit. --- ### Motivation In parameter-efficient fine-tuning (LoRA / PEFT), normalization parameters (weights and biases) are frozen (`requires_grad=False`). Previously, Liger-Kernel unconditionally: 1. Allocated intermediate reduction buffers (`_dW`, `_DB` of size `sm_count * n_cols`) in HBM. 2. Allocated registers and ran Triton accumulation loops across tokens. 3. Wrote intermediate gradients to global HBM. 4. Launched PyTorch `.sum(dim=0)` global reductions to produce tensors that PyTorch autograd immediately discarded. --- ### Key Changes - **`rms_norm.py`**: Added `COMPUTE_DW: tl.constexpr` to `_rms_norm_backward_kernel` and `_block_rms_norm_backward_kernel`. Prunes accumulation and memory writes when `compute_dw=False`; skips `_dW` buffer allocation and reduction; queries `ctx.needs_input_grad[1]`. - **`layer_norm.py`**: Added `COMPUTE_DW` and `COMPUTE_DB` flags to prune weight/bias gradients independently based on `ctx.needs_input_grad[1]` and `[2]`. - **`fused_add_rms_norm.py`**: Added `COMPUTE_DW` to skip `_dW` allocation and reduction based on `ctx.needs_input_grad[2]`. - **`test/ops/test_norm_lora_parity.py`**: Added 24 unit tests verifying numerical parity (FP32 & BF16), checking that `w.grad is None` / `b.grad is None` when frozen, and confirming exact `x.grad` equality. - **`benchmark/scripts/benchmark_norm_lora.py`**: Added dedicated standalone LoRA backward microbenchmark. --- ### Benchmark Results #### NVIDIA A100-SXM4-40GB (`triton.testing.do_bench`, bfloat16) | Operation | Shape (B, D) | Full FT Bwd (ms) | LoRA Bwd (ms) | Speedup | | :--- | :--- | :--- | :--- | :--- | | **LayerNorm Bwd** | (2048, 4096) | 0.0806 | 0.0585 | **1.38x (+38%)** | | **RMSNorm Bwd** | (2048, 4096) | 0.0871 | 0.0707 | **1.23x (+23%)** | | **LayerNorm Bwd** | (4096, 4096) | 0.1217 | 0.0991 | **1.23x (+23%)** | | **RMSNorm Bwd** | (4096, 4096) | 0.1305 | 0.1177 | **1.11x (+11%)** | | **LayerNorm Bwd** | (4096, 8192) | 0.1966 | 0.1729 | **1.14x (+14%)** | | **RMSNorm Bwd** | (4096, 8192) | 0.2275 | 0.2135 | **1.07x (+7%)** | | **LayerNorm Bwd** | (8192, 4096) | 0.2070 | 0.1843 | **1.12x (+12%)** | | **RMSNorm Bwd** | (8192, 4096) | 0.2384 | 0.2243 | **1.06x (+6%)** | | **LayerNorm Bwd** | (8192, 8192) | 0.3514 | 0.3255 | **1.08x (+8%)** | | **RMSNorm Bwd** | (8192, 8192) | 0.4235 | 0.4077 | **1.04x (+4%)** | #### NVIDIA A10G (`triton.testing.do_bench`, bfloat16) - **LayerNorm Bwd**: **1.05x – 1.18x speedup** across shapes. - **RMSNorm Bwd**: **1.02x – 1.07x speedup** across shapes. --- ### Test Validation - `test/ops/test_norm_lora_parity.py`: **24 / 24 PASSED (100%)** - Existing ops correctness: **108 / 108 PASSED (100%)** - `make checkstyle` (`ruff check` & `ruff format`): **PASSED** Co-authored-by: Yuki Uehara <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> Co-authored-by: Arup De <[email protected]>