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

src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -213,12 +213,19 @@ def fused_linear_cross_entropy_forward(
 
         if grad_weight is not None and input_requires_grad:
             grad_logits_t = grad_logits_chunk.t()
+            same_dtype = grad_weight.dtype == grad_logits_t.dtype
+            reduced_to_fp32 = grad_weight.dtype == torch.float32 and grad_logits_t.dtype in (
+                torch.float16,
+                torch.bfloat16,
+            )
+            requires_ampere = grad_logits_t.dtype == torch.bfloat16 or reduced_to_fp32
             if (
                 _ADDMM_SUPPORTS_OUT_DTYPE
                 and grad_weight.device.type == "cuda"
-                and torch.cuda.get_device_capability(grad_weight.device)[0] >= 8
-                and grad_weight.dtype == torch.float32
-                and grad_logits_t.dtype in (torch.float16, torch.bfloat16)
+                and not is_hip()
+                and grad_weight.dtype in (torch.float16, torch.bfloat16, torch.float32)
+                and (same_dtype or reduced_to_fp32)
+                and (not requires_ampere or torch.cuda.get_device_capability(grad_weight.device)[0] >= 8)
             ):
                 # Unlike torch.mm, torch.addmm's out_dtype path does not participate in
                 # autocast operand casting, so under AMP (fp32 params, no bias) _input_chunk
@@ -231,11 +238,11 @@ def fused_linear_cross_entropy_forward(
                     grad_weight,
                     grad_logits_t,
                     input_chunk,
-                    out_dtype=torch.float32,
+                    out_dtype=grad_weight.dtype,
                     out=grad_weight,
                 )
             else:
-                grad_weight += torch.mm(grad_logits_chunk.t(), _input_chunk).float()
+                grad_weight.add_(torch.mm(grad_logits_t, _input_chunk))
 
         if bias is not None and input_requires_grad:
             torch.add(

test/transformers/test_fused_linear_cross_entropy.py

@@ -7,6 +7,8 @@
 from test.utils import assert_verbose_allclose
 from test.utils import set_seed
 
+import liger_kernel.ops.fused_linear_cross_entropy as fused_linear_cross_entropy
+
 from liger_kernel.ops import LigerFusedLinearCrossEntropyFunction
 from liger_kernel.transformers.functional import CrossEntropyOutput
 from liger_kernel.transformers.functional import liger_fused_linear_cross_entropy
@@ -100,6 +102,84 @@ def forward(self, x, y):
         return self.ce_loss(self.lin.weight, x, y, self.lin.bias)
 
 
+def _run_grad_weight_accumulation(operand_dtype, accum_dtype):
+    N, H, V = 11, 41, 37
+    _input = torch.randn(N, H, device=device, dtype=operand_dtype, requires_grad=True)
+    weight = torch.randn(V, H, device=device, dtype=operand_dtype, requires_grad=True)
+    target = torch.randint(0, V, (N,), device=device)
+
+    ref_input = _input.detach().clone().requires_grad_(True)
+    ref_weight = weight.detach().clone().requires_grad_(True)
+    ref_loss = torch.nn.functional.cross_entropy(ref_input @ ref_weight.t(), target)
+    ref_loss.backward()
+
+    _, _, _, _, _, grad_weight, _ = fused_linear_cross_entropy.fused_linear_cross_entropy_forward(
+        _input,
+        weight,
+        target,
+        accum_dtype=accum_dtype,
+    )
+    return grad_weight, ref_weight.grad
+
+
+@pytest.mark.skipif(
+    not fused_linear_cross_entropy._ADDMM_SUPPORTS_OUT_DTYPE or device != "cuda" or torch.version.hip is not None,
+    reason="grad_weight addmm out_dtype requires PyTorch 2.8 or newer on NVIDIA CUDA",
+)
+@pytest.mark.parametrize(
+    "grad_weight_dtype, operand_dtype",
+    [
+        (torch.float16, torch.float16),
+        (torch.bfloat16, torch.bfloat16),
+        (torch.float32, torch.float32),
+        (torch.float32, torch.float16),
+        (torch.float32, torch.bfloat16),
+    ],
+)
+def test_accumulate_grad_weight_out_dtype(grad_weight_dtype, operand_dtype, monkeypatch):
+    requires_ampere = operand_dtype == torch.bfloat16 or grad_weight_dtype != operand_dtype
+    if requires_ampere and torch.cuda.get_device_capability()[0] < 8:
+        pytest.skip("this addmm dtype combination requires compute capability 8.0 or newer")
+
+    original_addmm = torch.addmm
+    observed_out_dtypes = []
+
+    def fail_mm(*args, **kwargs):
+        raise AssertionError("eligible accumulation unexpectedly used the allocating mm fallback")
+
+    def record_addmm(*args, **kwargs):
+        observed_out_dtypes.append(kwargs.get("out_dtype"))
+        return original_addmm(*args, **kwargs)
+
+    monkeypatch.setattr(torch, "mm", fail_mm)
+    monkeypatch.setattr(torch, "addmm", record_addmm)
+    grad_weight, expected = _run_grad_weight_accumulation(operand_dtype, grad_weight_dtype)
+
+    assert observed_out_dtypes == [grad_weight_dtype]
+    torch.testing.assert_close(grad_weight, expected, atol=5e-3, rtol=5e-2)
+
+
+@pytest.mark.parametrize("operand_dtype", [torch.float16, torch.bfloat16])
+def test_accumulate_grad_weight_fallback_mixed_dtype(operand_dtype, monkeypatch):
+    original_mm = torch.mm
+    observed_operand_dtypes = []
+
+    def fail_addmm(*args, **kwargs):
+        raise AssertionError("fallback accumulation unexpectedly used addmm")
+
+    def record_mm(mat1, mat2, *args, **kwargs):
+        observed_operand_dtypes.append((mat1.dtype, mat2.dtype))
+        return original_mm(mat1, mat2, *args, **kwargs)
+
+    monkeypatch.setattr(fused_linear_cross_entropy, "_ADDMM_SUPPORTS_OUT_DTYPE", False)
+    monkeypatch.setattr(torch, "addmm", fail_addmm)
+    monkeypatch.setattr(torch, "mm", record_mm)
+    grad_weight, expected = _run_grad_weight_accumulation(operand_dtype, torch.float32)
+
+    assert observed_operand_dtypes == [(operand_dtype, operand_dtype)]
+    torch.testing.assert_close(grad_weight, expected, atol=5e-3, rtol=5e-2)
+
+
 #############################################################################
 # Test the correctness of the fused linear cross entropy loss
 #############################################################################