[Perf] Use direct FLCE weight-gradient accumulation for FP16, BF16, and FP32

[Anubhav-2003]

Summary

This PR reduces allocation and memory-traffic overhead in Liger fused linear cross entropy by accumulating same-dtype FP16, BF16, and FP32 weight-gradient contributions directly into the existing grad_weight buffer on supported NVIDIA CUDA devices.

The broader goal is to keep Liger FLCE competitive with PyTorch 2.13 native linear_cross_entropy across the dtypes used for training. Both implementations avoid retaining a full [tokens, vocabulary] logits tensor, but Liger's fallback weight-gradient update still materialized a complete [vocabulary, hidden_size] matrix before adding each chunk's contribution to the persistent gradient. Direct accumulation removes that intermediate without changing the API, chunking policy, loss calculation, or returned gradient dtype.

The existing mixed-precision path for FP16/BF16 operands with an FP32 destination is preserved. Unsupported devices and dtype combinations continue to use the existing fallback.

Fixes #1323 .

Details

The previous same-dtype update used:

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

This creates a separate [V, H] matrix-multiplication result. FP16 and BF16 then create an additional FP32 conversion of that result. At V=32000 and H=4096, those intermediates account for 750 MiB in FP16/BF16 and 500 MiB in FP32, independently of the token chunk size.

This PR introduces a shared accumulation helper and selects direct accumulation when all three operands use the same supported dtype:

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

Dispatch behavior

Configuration	Behavior after this PR
Same-dtype FP16 on NVIDIA CUDA	Direct accumulation into grad_weight
Same-dtype BF16 on NVIDIA CUDA, compute capability 8.0+	Direct accumulation into grad_weight
Same-dtype FP32 on NVIDIA CUDA	Direct accumulation into grad_weight
FP16/BF16 operands with an FP32 destination	Preserve existing out_dtype=torch.float32 path
ROCm, older BF16 hardware, other dtypes, or mixed unsupported operands	Preserve existing fallback

The BF16 compute-capability check remains BF16-specific. FP16 and FP32 use the same CUDA operand types as the previously supported matrix-multiplication fallback and do not inherit the BF16 hardware restriction.

The new regression test parametrizes FP16, BF16, and FP32 and deliberately fails if a supported same-dtype case reaches the allocating torch.mm fallback.

Benchmark results

Methodology

• Complete forward and backward training operation; forward-only timings are not compared.
• Shape: N=16384, H=4096, V=32000.
• Identical inputs, weights, and targets for every provider in a dtype.
• Providers interleaved in a seeded random order.
• Three warm-up rounds and 10 CUDA-event timing samples per provider.
• Two isolated peak-memory samples per provider.
• Correctness gate for loss, input gradient, and weight gradient before timing.
• PyTorch 2.13.0+cu130; Liger candidate applied to baseline ed08f6e.

torch-lce-auto is PyTorch 2.13 native LCE with automatic chunking. torch-lce-512 uses the same native implementation with a fixed 512-token chunk.

Full all-dtype comparison

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

Comparison with PyTorch LCE auto

Dtype	Liger median-latency difference	Liger peak-memory difference	Interpretation
BF16	12.6% lower	10.9% lower	Liger is faster and uses less peak memory
FP16	11.2% lower	10.9% lower	Liger is faster and uses less peak memory
FP32	4.1% higher	11.2% lower	Liger uses less memory; FP32 latency remains follow-up work

These results do not treat one GPU's capacity as the motivation for the change. The optimization removes an allocation proportional to V * H; the benchmark hardware is reported only so that the measurements can be reproduced and compared fairly.

Focused accumulation result

At V=32000, H=4096, and a 512-token chunk, extending direct accumulation beyond the BF16-only branch produced:

Dtype	Before	After	Before temporary	After temporary
BF16 control	12.726 ms	12.738 ms	0 MiB	0 MiB
FP16	20.774 ms	10.262 ms	750 MiB	0 MiB
FP32	35.277 ms	28.643 ms	500 MiB	0 MiB

BF16 remains within measurement noise, while FP16 and FP32 remove the full-size temporary and reduce accumulation latency.

Correctness

The standalone gate compared ordinary PyTorch, PyTorch LCE auto, PyTorch LCE with 512-token chunks, and Liger. All loss, input-gradient, and weight-gradient comparisons passed in BF16, FP16, and FP32.

Dtype	Absolute tolerance	Relative tolerance	Maximum loss error	Maximum input-gradient error	Maximum weight-gradient error
BF16	0.005	0.05	0.25	3.74e-4	3.05e-4
FP16	0.005	0.05	1.80e-3	7.63e-6	1.53e-5
FP32	1e-5	5e-4	3.81e-6	6.05e-8	3.73e-8

The BF16 loss comparison passes through the relative-tolerance term; its gradient errors remain below 4e-4.

Testing Done

• Hardware Type: NVIDIA GeForce RTX 3050 Laptop GPU, compute capability 8.6
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

Additional validation completed:

• Full FLCE test file: python -m pytest test/transformers/test_fused_linear_cross_entropy.py -x -q (140 passed)
• Direct-accumulation branch tests for FP16, BF16, and FP32 (3 passed)
• Repository-wide ruff check .
• Repository-wide ruff format --check .
• Standalone all-dtype correctness gate across four providers
• Interleaved PyTorch 2.13 LCE versus Liger latency and memory benchmark

[Tcc0403]

wouldn't it be more straightforward to set out_dtype=grad_weight.dtype? with a little bit change on branching condition.

[Anubhav-2003]

Thank you very much for the review. I see how that would simplify the path for PyTorch ≥2.8. Since Liger supports PyTorch ≥2.1.2, using out_dtype for all these cases would either require retaining a separate path for older versions or leave them on the allocating fallback. That’s why I kept addmm_ for same-dtype accumulation. Would you prefer limiting this optimization to PyTorch ≥2.8 instead? But then older versions will keep using the unoptimized path.

[Tcc0403]

yes, modify output_dtype in 2.8+ path.

add_ supports mismatching dtype but I don't recall which version starts supporting it. We can change the other path to grad_weight.add_(torch.mm(grad_logits_chunk.t(), _input_chunk)) if 2.1.2 has it covered

[Anubhav-2003]

Thank you for the clarification. I checked PyTorch 2.1.2, and add_ supports the floating-point dtype conversions needed here. I’ll use out_dtype=grad_weight.dtype for the eligible 2.8+ path and grad_weight.add_(torch.mm(...)) as the fallback, with coverage for both paths.

[Anubhav-2003]

@Tcc0403
Respectd Sir,

I have made the changes as per your advice. You may take a look whenever you are free.

Thank You.