diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index df35068ce..e049629b0 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/fused_linear_cross_entropy.py @@ -127,9 +127,12 @@ 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: @@ -137,10 +140,12 @@ def fused_linear_cross_entropy_forward( 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 diff --git a/src/liger_kernel/ops/utils.py b/src/liger_kernel/ops/utils.py index 45ad47880..72cbafd80 100644 --- a/src/liger_kernel/ops/utils.py +++ b/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): diff --git a/test/transformers/test_fused_linear_cross_entropy.py b/test/transformers/test_fused_linear_cross_entropy.py index 1b8a10e12..2ef1bf877 100644 --- a/test/transformers/test_fused_linear_cross_entropy.py +++ b/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)