mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Vaibhav JindalandCopilot
95b01e9402
## Summary Optimize Liger's Triton fused linear cross entropy (FLCE) weight-gradient accumulation for BF16/FP16 parameters on NVIDIA SM80+ GPUs. The previous default-accumulator path computes a parameter-sized matrix, converts it to FP32, then adds it into the BF16/FP16 running gradient buffer for **every token chunk**: ```python grad_weight += torch.mm(grad_logits_chunk.t(), input_chunk).float() ``` This PR instead accumulates directly into the existing dtype-matched gradient buffer: ```python torch.addmm(grad_weight, grad_logits_t, input_chunk, out=grad_weight) ``` This removes the separate parameter-sized product/conversion temporary and extra addition pass. The stored gradient-buffer dtype is unchanged; the old `.float()` temporary was not an FP32 running accumulator. The fused accumulation can change rounding relative to the old standalone product, so bitwise equivalence is not claimed. The existing guarded FP32 accumulation fast path, AMP operand alignment, and legacy dtype/device fallbacks are preserved. The CE kernel, its launch configuration, the default token-chunk heuristic, and the forward/backward scheduling are unchanged. **Scope:** this is the first stack layer only: Triton accumulation plus regressions in the existing FLCE test file. It does **not** add a public `chunk_size` argument, change cuTile, or include the later three-way benchmark implementation. ## Performance improvements attributable to this PR ### Unchanged production defaults BF16 inputs/weights, `accum_dtype=None`, and the existing `_CHUNK_MEM_CONST=16` heuristic on both sides. The number of token chunks is identical before and after. At the Llama-3-sized reference shape, this is **1.26x faster** with **38.9% lower incremental peak allocated memory** without changing any user settings. | Tokens / hidden / vocab | Chunk rows | Before ms | After ms | Speedup | Before / after peak MiB | Memory reduction | |---|---:|---:|---:|---:|---:|---:| | 8,192 / 4,096 / 128,256 | 4,096 | 25.002 | 19.897 | **1.26x** | 5,074.04 / 3,102.04 | **38.9%** | | 4,096 / 2,048 / 32,000 | 4,096 | 1.774 | 1.366 | **1.30x** | 766.02 / 407.02 | **46.9%** | ### Controlled smaller chunks: benefit per accumulation step These are additional controlled experiments, **not the production defaults**. For the reference shape, the existing private chunk-budget constant was set to 4 (1,024-row chunks) or 1 (256-row chunks) on **both** versions. No new chunk-size API or later stack-layer code was needed. | Tokens / hidden / vocab | Chunk rows | Before ms | After ms | Speedup | Before / after peak MiB | Memory reduction | |---|---:|---:|---:|---:|---:|---:| | 8,192 / 4,096 / 128,256 | 1,024 | 43.167 | 21.836 | **1.98x** | 4,322.54 / 1,575.04 | **63.6%** | | 8,192 / 4,096 / 128,256 | 256 | 123.443 | 31.686 | **3.90x** | 4,134.66 / 1,193.29 | **71.1%** | The larger gain at smaller chunks is expected: the old implementation repeatedly allocated/converted/added a full `[vocab, hidden]` matrix, and this PR removes that work on each iteration. At the reference shape, one such FP32 matrix alone is approximately 2,004 MiB. ### FP32-accumulator control The existing `accum_dtype=torch.float32` fast path is intentionally preserved. At 8,192 tokens / hidden 4,096 / vocab 128,256, with the unchanged default 4,096-row chunk: - Before: **20.473 ms**, **4,104.04 MiB** incremental peak. - After: **20.471 ms**, **4,104.04 MiB** incremental peak. This control is effectively unchanged; the improvement targets the default low-precision gradient-buffer path. ### Measurement methodology - **Hardware:** NVIDIA B200, SM100, driver 580.105.08. - **Software:** PyTorch 2.11.0.1+cu130, Triton 3.6.0. - **Before / after source:** base FLCE implementation `96233e3` versus `6a5077f`, loaded against the same installed dependencies. Current head `d5f3593` changes only a test comment relative to `6a5077f`; its kernel implementation is identical. - Identical seeded BF16 inputs, weights, targets, and token-chunk geometry; both inputs and weights require gradients; no bias or optional loss features in the timing cases. - Measure **full forward + backward**, mean reduction, fixed upstream gradient 1.0. Each iteration uses a fresh graph and clears both leaf gradients. No optimizer step or model-training-throughput claim. - Compilation plus five warmup iterations excluded. Each run uses 30 CUDA-event timed iterations; tables report the median of three runs, alternating before/after order in the same process. - Memory is **incremental peak allocated MiB above the live input/weight baseline**, with both gradients cleared before reset and a dedicated memory-probe iteration. It is not total process memory or reserved memory. - These compare the same Triton-backed implementation before/after the portable cuBLAS accumulation improvement. They are not CuTe-versus-Triton or DSL/compiler speedup claims. ## Testing Done - **Hardware:** NVIDIA B200 / SM100. - Latest relevant suites: **470 passed, 1 environment-gated skipped** across `test/ops/test_fused_linear_cross_entropy.py`, `test/transformers/test_fused_linear_cross_entropy.py`, and `test/cutedsl/test_fused_linear_cross_entropy.py`. - All **70 addmm regression cases** now live in the existing ops FLCE test file alongside its existing 23 cases. Their chunk-budget fixture is class-scoped and does not alter the pre-existing tests. - Coverage includes real aten dispatch checks for direct `addmm.out` and no parameter-shaped upcasts, preserved FP32 `dtype_out`, AMP and FP32-bias promotion, BF16/FP16/FP32, mean/sum/nonuniform per-token gradients, ignored chunks, frozen weights, noncontiguous tensors, and old-Torch/SM capability fallbacks. - Review requests addressed: restored the original inline FP32 guard, removed the extra boolean aliases, consolidated tests, and renamed the marker to `requires_sm80_or_higher`. - [x] `make checkstyle` passes locally; GitHub Checkstyle is also green. - [ ] Full-repository `make test` was not run; the targeted suites above were run. - [ ] `make test-convergence` was not run; no convergence claim is made. --------- Co-authored-by: Copilot <[email protected]>