From b6d224bc82827426ff6c905f99de48dd4e046638 Mon Sep 17 00:00:00 2001 From: cm2435 Date: Wed, 24 Jan 2024 18:03:00 +0000 Subject: [PATCH] fixed incorrect bwd pass for GeLU, wrapped functions in pytorch convience handlers --- tests/kernels/test_gelu.py | 8 ++-- unsloth/kernels/gelu.py | 82 +++++++++++++++++++++++++++++--------- 2 files changed, 68 insertions(+), 22 deletions(-) diff --git a/tests/kernels/test_gelu.py b/tests/kernels/test_gelu.py index 820486a778..b2b2d01ed7 100644 --- a/tests/kernels/test_gelu.py +++ b/tests/kernels/test_gelu.py @@ -2,7 +2,7 @@ import pytest import torch import triton -from unsloth.kernels.gelu import gelu_forward_kenel, gelu_backward_kenel +from unsloth.kernels.gelu import gelu_backward_triton, gelu_forward_triton from tests.conftest import set_seed, assert_all_close @set_seed @@ -15,7 +15,7 @@ def test_matrix(request): # Test function def test_relu_kernel_fwd(test_matrix): # Apply your Triton-based ReLU kernel - triton_output = gelu_forward_kenel(test_matrix) + triton_output = gelu_forward_triton(test_matrix) # Apply PyTorch's ReLU for comparison torch_gelu = torch.nn.GELU() @@ -31,11 +31,11 @@ def test_gelu_backward_kernel(test_matrix): grad_input = torch.randn_like(test_matrix) # Apply your Triton-based GeLU backward kernel - triton_output = gelu_backward_kenel(test_matrix, grad_input) + triton_output = gelu_backward_triton(test_matrix, grad_input) # Compute PyTorch's GeLU gradient for comparison torch_gelu = torch.nn.GELU() - torch_output = torch.autograd.grad(torch_output.sum(), test_matrix, grad_outputs=grad_input)[0] + torch_output = torch.autograd.grad(torch_gelu.sum(), test_matrix, grad_outputs=grad_input)[0] # Check if the outputs are close enough using assert_all_close assert_all_close(triton_output, torch_output, rtol=1e-05, atol=1e-08) \ No newline at end of file diff --git a/unsloth/kernels/gelu.py b/unsloth/kernels/gelu.py index 007d86c53f..05ad0e64e2 100644 --- a/unsloth/kernels/gelu.py +++ b/unsloth/kernels/gelu.py @@ -16,8 +16,48 @@ import triton import triton.language as tl import torch +def gelu_forward_triton(x: torch.Tensor): + n_rows, n_cols = x.shape + y = torch.empty_like(x) + + # Define the grid of blocks + # Here, we divide the number of rows by the block size to determine the number of blocks needed + BLOCK_SIZE = 1024 + num_blocks = triton.cdiv(n_rows, BLOCK_SIZE) + + # Launch the kernel with the grid configuration + _gelu_forward_kenel[(num_blocks,)]( + output_ptr=y.data_ptr(), + input_ptr=x.data_ptr(), + n_elements=x.stride(0), + n_cols=n_cols, + BLOCK_SIZE=BLOCK_SIZE, + ) + return y + + +def gelu_backward_triton(x: torch.Tensor, grad_output: torch.Tensor): + n_rows, n_cols = x.shape + grad_input = torch.empty_like(x) + + # Define the grid of blocks + BLOCK_SIZE = 1024 + num_blocks = triton.cdiv(n_rows, BLOCK_SIZE) + + # Launch the kernel with the grid configuration + _gelu_backward_kernel[(num_blocks,)]( + grad_input_ptr=grad_input.data_ptr(), + input_ptr=x.data_ptr(), + grad_output_ptr=grad_output.data_ptr(), + n_elements=x.stride(0), + n_cols=n_cols, + BLOCK_SIZE=BLOCK_SIZE, + ) + return grad_input + + @triton.jit -def gelu_forward_kenel(output_ptr: tl.pointer, input_ptr: tl.pointer, n_elements: tl.int32, n_cols: tl.int32, BLOCK_SIZE : tl.constexpr,): +def _gelu_forward_kenel(output_ptr: tl.pointer, input_ptr: tl.pointer, n_elements: tl.int32, n_cols: tl.int32, BLOCK_SIZE : tl.constexpr,): ''' Triton kernel for the forward pass of a GeLU function based off the equation https://pytorch.org/docs/stable/generated/torch.nn.GELU.html @@ -42,32 +82,38 @@ def gelu_forward_kenel(output_ptr: tl.pointer, input_ptr: tl.pointer, n_elements pass -@triton.jit -def gelu_backward_kernel(output_ptr: tl.pointer, input_ptr: tl.pointer, n_elements: tl.int32, n_cols: tl.int32, BLOCK_SIZE : tl.constexpr,): - ''' - Triton kernel for the backward pass of a GeLU function based off eq [13] - https://arxiv.org/pdf/2305.12073.pdf - - output_ptr : the pointer for the first memory adress of the output tensor - input_ptr : the pointer for the input of the first element of the first row of the input tensor. - n_elements : - ''' +@triton.jit +def _gelu_backward_kernel(grad_input_ptr: tl.pointer, input_ptr: tl.pointer, grad_output_ptr: tl.pointer, n_elements: tl.int32, n_cols: tl.int32, BLOCK_SIZE: tl.constexpr): row_idx = tl.program_id(0) - row_start_ptr = input_ptr + row_idx * n_elements + # Compute pointers to the start of the row for input, gradient output, and gradient input + row_start_input_ptr = input_ptr + row_idx * n_elements + row_start_grad_output_ptr = grad_output_ptr + row_idx * n_elements + row_start_grad_input_ptr = grad_input_ptr + row_idx * n_elements + + # Iterate over the columns of the row col_offsets = tl.arange(0, BLOCK_SIZE) - input_ptrs = row_start_ptr + col_offsets + input_ptrs = row_start_input_ptr + col_offsets + grad_output_ptrs = row_start_grad_output_ptr + col_offsets + grad_input_ptrs = row_start_grad_input_ptr + col_offsets + + # Mask to avoid out-of-bounds memory access mask = col_offsets < n_cols - row = tl.load(input_ptrs, mask=mask, other=0) - output_values = gelu_bwd_pass(x=row) + # Load input and gradient output values + x = tl.load(input_ptrs, mask=mask, other=0) + grad_output = tl.load(grad_output_ptrs, mask=mask, other=0) - output_row_start_ptr = output_ptr + row_idx * n_elements - output_ptrs = output_row_start_ptr + col_offsets - tl.store(output_ptrs, output_values, mask=mask) + # Compute the GELU backward operation using gelu_bwd_pass + dgelu_dx = gelu_bwd_pass(x) + grad_input = dgelu_dx * grad_output + + # Store the computed gradient input + tl.store(grad_input_ptrs, grad_input, mask=mask) pass + @triton.jit def gelu_operation(x): PI = 3.141592653589793 # Define Pi