diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index a15d9e3b2..7c81352bd 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/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( diff --git a/test/transformers/test_fused_linear_cross_entropy.py b/test/transformers/test_fused_linear_cross_entropy.py index e1d5bc1a3..eb1f48f96 100644 --- a/test/transformers/test_fused_linear_cross_entropy.py +++ b/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 #############################################################################