linkedin / linkedin/Liger-Kernel

Reduce FLCE weight-gradient allocation overhead across FP16, BF16, and FP32

Open
#1,323 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
6.6k
Forks
603
Avg merge
1d 20h
Merged PRs (30d)
47

Description

# 🐛 Describe the bug field

## Background and motivation

Fused linear cross entropy reduces the cost of language-model training by avoiding a single full `[tokens, vocabulary]` logits tensor. Liger FLCE does this by processing the linear projection and cross entropy in token chunks.

PyTorch 2.13 now provides a native chunked `linear_cross_entropy` implementation for FP16, BF16, and FP32. That gives Liger a useful framework-level comparison point: Liger FLCE should remain at least competitive with native PyTorch LCE in both execution time and memory use across the same training dtypes.

While comparing the two implementations, I found an avoidable allocation in Liger's weight-gradient update. Liger chunks the logits, but its fallback weight-gradient expression first creates a complete `[vocabulary, hidden_size]` matrix for each chunk before adding that contribution to the persistent weight-gradient buffer:

```python
grad_weight += torch.mm(grad_logits_chunk.t(), input_chunk).float()
```

The size of this intermediate is determined by the model dimensions `V` and `H`, not by the number of tokens in the chunk. Therefore, logits chunking does not reduce this part of the peak memory or the associated allocation and memory-traffic cost.

For FP16 and BF16, `torch.mm` creates a low-precision `[V, H]` result and `.float()` creates an additional FP32 copy. For FP32, the cast is a no-op, but the separate `[V, H]` matrix-multiplication result is still created.

At `V=32000` and `H=4096`, the avoidable intermediates are:

| Dtype | Matrix-multiplication result | Additional FP32 conversion | Total avoidable allocation |
|---|---:|---:|---:|
| BF16 | 250 MiB | 500 MiB | 750 MiB |
| FP16 | 250 MiB | 500 MiB | 750 MiB |
| FP32 | 500 MiB | none | 500 MiB |

This is a general FLCE implementation cost. Larger GPUs can accommodate larger allocations, but the allocation and memory traffic are still unnecessary and scale with `V * H`.

## Proposed direction

When `grad_weight`, `grad_logits`, and the input chunk have the same native training dtype, Liger can accumulate the matrix product directly into the existing destination:

```python
grad_weight.addmm_(grad_logits_chunk.t(), input_chunk)
```

The proposed support matrix is:

| Configuration | Accumulation path |
|---|---|
| Same-dtype FP16 on NVIDIA CUDA | Direct `addmm_` |
| Same-dtype BF16 on NVIDIA CUDA, compute capability 8.0+ | Direct `addmm_` |
| Same-dtype FP32 on NVIDIA CUDA | Direct `addmm_` |
| FP16/BF16 operands with an FP32 destination | Existing `addmm(..., out_dtype=torch.float32, out=...)` path |
| Other devices, dtypes, or unsupported combinations | Existing fallback |

This changes only how each chunk's weight-gradient contribution is accumulated. It does not change the FLCE API, token chunking, loss calculation, or returned gradient dtype.

## Comparison with PyTorch 2.13 native LCE

The following end-to-end benchmark measures the complete forward and backward training operation at `N=16384`, `H=4096`, and `V=32000`. Providers use identical tensors, run in an interleaved seeded order, and are timed with CUDA events after three warm-up rounds. Each latency result contains 10 samples; memory contains two isolated samples.

`torch-lce-auto` is PyTorch 2.13 native LCE with automatic chunking. `torch-lce-512` is the same native implementation with a fixed 512-token chunk. `liger` is Liger FLCE with the proposed direct accumulation.

| Dtype | Provider | Tokens per chunk | p20 / median / p80 | Throughput | Peak allocated | Peak increase during operation |
|---|---|---:|---:|---:|---:|---:|
| BF16 | PyTorch LCE auto | 1,024 | 1,020.31 / 1,026.39 / 1,029.58 ms | 15,962.8 tokens/s | 1,150.38 MiB | 756.00 MiB |
| BF16 | PyTorch LCE 512 | 512 | 1,050.54 / 1,057.01 / 1,058.44 ms | 15,500.3 tokens/s | 1,150.38 MiB | 756.00 MiB |
| BF16 | Liger FLCE | 2,048 | 892.08 / 897.07 / 900.65 ms | 18,263.8 tokens/s | 1,024.45 MiB | 630.08 MiB |
| FP16 | PyTorch LCE auto | 1,024 | 1,011.01 / 1,012.98 / 1,014.89 ms | 16,174.1 tokens/s | 1,150.38 MiB | 756.00 MiB |
| FP16 | PyTorch LCE 512 | 512 | 1,065.28 / 1,066.65 / 1,068.67 ms | 15,360.3 tokens/s | 1,150.38 MiB | 756.00 MiB |
| FP16 | Liger FLCE | 2,048 | 897.90 / 899.94 / 901.24 ms | 18,205.7 tokens/s | 1,024.45 MiB | 630.08 MiB |
| FP32 | PyTorch LCE auto | 1,024 | 3,313.70 / 3,317.26 / 3,324.13 ms | 4,939.0 tokens/s | 2,284.38 MiB | 1,512.00 MiB |
| FP32 | PyTorch LCE 512 | 512 | 3,329.99 / 3,332.67 / 3,338.46 ms | 4,916.2 tokens/s | 2,284.38 MiB | 1,512.00 MiB |
| FP32 | Liger FLCE | 2,048 | 3,451.62 / 3,454.83 / 3,462.58 ms | 4,742.3 tokens/s | 2,028.45 MiB | 1,256.08 MiB |

Compared with PyTorch LCE auto at this shape:

- BF16 Liger has 12.6% lower median latency and 10.9% lower peak allocated memory.
- FP16 Liger has 11.2% lower median latency and 10.9% lower peak allocated memory.
- FP32 Liger has 11.2% lower peak allocated memory, while its median latency is 4.1% higher. The remaining FP32 latency difference is separate follow-up optimization work; this change improves the accumulation path without claiming that every dtype is already faster.

The correctness gate compared loss, input gradients, and weight gradients for ordinary PyTorch, both PyTorch LCE chunking modes, and Liger in BF16, FP16, and FP32. All comparisons passed their dtype-specific tolerances. The complete Liger FLCE test file also passes (`140 passed`), including direct-path regression tests for all three dtypes.

## Relationship to earlier work

This follows issue #1232 and PR #1239, which optimized FP16/BF16 operands accumulating into an FP32 destination. The proposed change covers the separate same-dtype FP16, BF16, and FP32 paths.

Issue #533 discusses the lifetime of the required persistent weight-gradient buffer. This proposal instead removes intermediate `[V, H]` results that do not need to exist.

# Reproduce field

The allocation pattern can be inspected independently of total GPU capacity. This example measures the extra peak allocation of the existing same-dtype update at `V=32000`, `H=4096`, and a 512-token chunk:

```python
import torch

N, H, V = 512, 4096, 32000
dtype = torch.float16 # also test torch.bfloat16 and torch.float32

grad_weight = torch.zeros(V, H, device="cuda", dtype=dtype)
grad_logits = torch.randn(N, V, device="cuda", dtype=dtype)
input_chunk = torch.randn(N, H, device="cuda", dtype=dtype)

torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
baseline = torch.cuda.memory_allocated()

grad_weight += torch.mm(grad_logits.t(), input_chunk).float()

torch.cuda.synchronize()
extra_mib = (torch.cuda.max_memory_allocated() - baseline) / 1024**2
print(f"extra peak allocation: {extra_mib:.2f} MiB")
```

The direct form to compare is:

```python
grad_weight.addmm_(grad_logits.t(), input_chunk)
```

# Versions field

- Liger Kernel baseline: `main` at `ed08f6ea75e388a63b7493b5bc97042b1f142712`
- Python: `3.14.4`
- PyTorch: `2.13.0+cu130`
- Triton: `3.7.1`
- Transformers: `5.14.1`
- CUDA: `13.0`
- Benchmark hardware: NVIDIA GeForce RTX 3050 Laptop GPU, compute capability 8.6

Contributor guide

Open the contributing guide

Research direction

Start with the FLCE weight-gradient update described in the issue and reproduce the allocation pattern using the provided CUDA example at the stated dimensions and dtypes. Compare the direct accumulation form with the existing fallback, then run the complete Liger FLCE test file and verify the direct-path regression tests, dtype-specific correctness checks, and benchmark memory results.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.