perf(flce): drop host syncs that do not feed the computation

src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -127,20 +127,25 @@ def fused_linear_cross_entropy_forward(
     token_accuracy_1d = torch.zeros(BT, dtype=torch.float32, device=device) if return_token_accuracy else None
     predicted_tokens_1d = torch.full((BT,), -1, dtype=torch.int64, device=device) if return_predicted_tokens else None
 
-    # TODO: evaluate how CUDA synchronization caused by .item() affects the speed
+    # Each .item() below synchronizes the device, and a sync waits for everything already queued on it (e.g.
+    # activation-offload copies), so read a value on the host only when the kernel uses it: the non-ignored count
+    # and target-weight sum only scale reduction="mean", and the class-weight sum only enters label smoothing.
     target_mask = target != ignore_index
-    total_n_non_ignore = target_mask.sum().item()
+    normalize = reduction == "mean"
+    total_n_non_ignore = target_mask.sum().item() if normalize else 1
     total_sum_non_ignore_ce_weight = total_n_non_ignore
     ce_weight_sum = 0.0
     if ce_weight is not None:
         assert ce_weight.shape[0] == V, f"If given, weight has to be a Tensor of size V. Got: {ce_weight.shape}"
         assert torch.is_floating_point(ce_weight), (
             f"If given, weight has to be a Tensor of floating point dtype. Got: {ce_weight.dtype}"
         )
-        total_sum_non_ignore_ce_weight = (
-            torch.gather(ce_weight, dim=0, index=target.masked_select(target_mask)).sum().item()
-        )
-        ce_weight_sum = ce_weight.sum().item()
+        if normalize:
+            total_sum_non_ignore_ce_weight = (
+                torch.gather(ce_weight, dim=0, index=target.masked_select(target_mask)).sum().item()
+            )
+        if label_smoothing > 0:
+            ce_weight_sum = ce_weight.sum().item()
         if ce_weight.stride(-1) != 1:
             ce_weight = ce_weight.contiguous()
 
@@ -404,7 +409,7 @@ def fused_linear_cross_entropy_forward(
         loss = torch.sum(loss_1d)
         z_loss = torch.sum(z_loss_1d) if return_z_loss else None
         # For accuracy, we compute the mean across all non-ignored tokens
-        token_accuracy = torch.sum(token_accuracy_1d) / total_n_non_ignore if return_token_accuracy else None
+        token_accuracy = torch.sum(token_accuracy_1d) / target_mask.sum() if return_token_accuracy else None
 
     predicted_tokens = predicted_tokens_1d if return_predicted_tokens else None
 
@@ -416,51 +421,55 @@ def fused_linear_cross_entropy_forward(
 
 
 def fused_linear_cross_entropy_backward(grad_output, grad_input, grad_weight, grad_bias):
-    # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time
-    if not torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
-        # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
-        # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
-        BT, H = grad_input.shape
-        n_rows = BT
-        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
+    # If cross entropy is the last layer, grad_output is 1.0 and the kernels skip the mul (SKIP_IF_ONE). Checking
+    # that on the host with torch.equal would synchronize the device every backward pass; when the stream also has
+    # queued device-to-host copies (e.g. activation offload), the host then stalls until all of them finish.
+    # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
+    # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
+    BT, H = grad_input.shape
+    n_rows = BT
+    BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
+
+    element_mul_kernel[(n_rows,)](
+        grad_input,
+        grad_input.stride(-2),
+        grad_output,
+        H,
+        BLOCK_SIZE=BLOCK_SIZE,
+        X_col_stride=grad_input.stride(-1),
+        SKIP_IF_ONE=True,
+        num_warps=32 if not is_hip() else 16,
+    )
+
+    # handle grad_weight
+    if grad_weight is not None:
+        V, H = grad_weight.shape
+        n_rows = V
 
         element_mul_kernel[(n_rows,)](
-            grad_input,
-            grad_input.stride(-2),
+            grad_weight,
+            grad_weight.stride(-2),
             grad_output,
             H,
             BLOCK_SIZE=BLOCK_SIZE,
-            X_col_stride=grad_input.stride(-1),
+            X_col_stride=grad_weight.stride(-1),
+            SKIP_IF_ONE=True,
             num_warps=32 if not is_hip() else 16,
         )
 
-        # handle grad_weight
-        if grad_weight is not None:
-            V, H = grad_weight.shape
-            n_rows = V
-
-            element_mul_kernel[(n_rows,)](
-                grad_weight,
-                grad_weight.stride(-2),
-                grad_output,
-                H,
-                BLOCK_SIZE=BLOCK_SIZE,
-                X_col_stride=grad_weight.stride(-1),
-                num_warps=32 if not is_hip() else 16,
-            )
-
-        if grad_bias is not None:
-            V = grad_bias.shape[0]
-            n_rows = V
+    if grad_bias is not None:
+        V = grad_bias.shape[0]
+        n_rows = V
 
-            element_mul_kernel[(n_rows,)](
-                grad_bias,
-                grad_bias.stride(-1),
-                grad_output,
-                1,
-                BLOCK_SIZE=BLOCK_SIZE,
-                num_warps=32 if not is_hip() else 16,
-            )
+        element_mul_kernel[(n_rows,)](
+            grad_bias,
+            grad_bias.stride(-1),
+            grad_output,
+            1,
+            BLOCK_SIZE=BLOCK_SIZE,
+            SKIP_IF_ONE=True,
+            num_warps=32 if not is_hip() else 16,
+        )
     return grad_input, grad_weight, grad_bias
 
 

src/liger_kernel/ops/utils.py

@@ -113,6 +113,7 @@ def element_mul_kernel(
     n_cols,
     BLOCK_SIZE: tl.constexpr,
     X_col_stride: tl.constexpr = 1,
+    SKIP_IF_ONE: tl.constexpr = False,
 ):
     """
     This function multiplies each element of the tensor pointed by X_ptr with the value pointed by grad_output_ptr.
@@ -125,6 +126,8 @@ def element_mul_kernel(
     n_cols (int): The number of columns in the input tensor.
     BLOCK_SIZE (int): The block size for Triton operations.
     X_col_stride (int): The stride between columns of the input tensor.
+    SKIP_IF_ONE (bool): Leave the tensor untouched when the gradient output is 1.0. Callers use this instead of
+        comparing the gradient output with 1.0 on the host, which would synchronize the device.
     """
 
     # Get the program ID and convert it to int64 to avoid overflow
@@ -135,6 +138,9 @@ def element_mul_kernel(
 
     # Load the gradient output value
     grad_output = tl.load(grad_output_ptr)
+    if SKIP_IF_ONE:
+        if grad_output == 1.0:
+            return
 
     # Perform the element-wise multiplication
     for i in range(0, n_cols, BLOCK_SIZE):

test/transformers/test_fused_linear_cross_entropy.py

@@ -1108,3 +1108,63 @@ def liger():
     assert_verbose_allclose(lig_gw, ref_gw, atol=atol, rtol=rtol)
     if bias:
         assert_verbose_allclose(lig_gb, ref_gb, atol=atol, rtol=rtol)
+
+
+@pytest.mark.skipif(device != "cuda", reason="torch.cuda.set_sync_debug_mode is CUDA-only")
+@pytest.mark.parametrize("reduction", ["sum", "none"])
+@pytest.mark.parametrize("grad_scale", [1.0, 0.37])
+@pytest.mark.parametrize("bias", [True, False])
+@pytest.mark.parametrize(
+    "class_weight, label_smoothing",
+    [(False, 0.0), (False, 0.1), (True, 0.0)],
+    ids=["plain", "label_smoothing", "class_weight"],
+)
+def test_no_host_sync_without_mean_reduction(reduction, grad_scale, bias, class_weight, label_smoothing):
+    """Forward and backward must not synchronize the host when no value has to be read back.
+
+    A host sync waits for everything queued on the device, e.g. activation-offload copies, so it stalls training.
+    Only reduction="mean" needs the non-ignored count, and only label smoothing needs the class-weight sum.
+    grad_scale=1.0 covers the backward path that skips the multiply, 0.37 the path that scales the saved gradients.
+    """
+    set_seed(42)
+    BT, H, V = 64, 128, 256
+    _input = torch.randn(BT, H, device=device)
+    weight = torch.randn(V, H, device=device)
+    bias_t = torch.randn(V, device=device) if bias else None
+    ce_weight = torch.rand(V, device=device) + 0.5 if class_weight else None
+    target = torch.randint(0, V, (BT,), device=device)
+    target[::5] = -100
+
+    def leaves():
+        x = _input.clone().requires_grad_(True)
+        w = weight.clone().requires_grad_(True)
+        b = bias_t.clone().requires_grad_(True) if bias else None
+        return x, w, b
+
+    x, w, b = leaves()
+    torch.cuda.synchronize()
+    torch.cuda.set_sync_debug_mode("error")
+    try:
+        loss, _, accuracy, _ = LigerFusedLinearCrossEntropyFunction.apply(
+            x, w, target, b, ce_weight, -100, 0.0, label_smoothing, reduction, None, False, None, False, True
+        )
+        (loss * grad_scale).sum().backward()
+    finally:
+        torch.cuda.set_sync_debug_mode("default")
+
+    xr, wr, br = leaves()
+    logits = torch.nn.functional.linear(xr, wr, br)
+    ref = torch.nn.functional.cross_entropy(
+        logits, target, weight=ce_weight, ignore_index=-100, label_smoothing=label_smoothing, reduction=reduction
+    )
+    (ref * grad_scale).sum().backward()
+    mask = target != -100
+    ref_accuracy = (logits.argmax(-1) == target)[mask].float().mean()
+
+    assert_verbose_allclose(loss, ref, atol=1e-4, rtol=1e-4)
+    if reduction == "sum":
+        assert_verbose_allclose(accuracy, ref_accuracy, atol=1e-6, rtol=1e-6)
+    assert_verbose_allclose(x.grad, xr.grad, atol=1e-4, rtol=1e-4)
+    assert_verbose_allclose(w.grad, wr.grad, atol=1e-4, rtol=1e-4)
+    if bias:
+        assert_verbose_allclose(b.grad, br.grad, atol=1e-4, rtol=1e-4)