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

[Anubhav-2003]

🐛 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:

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:

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:

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:

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