From 3954f169769abb215c62a55a1bcd906d0b1f738e Mon Sep 17 00:00:00 2001
From: cm2435 <cm2435@bath.ac.uk>
Date: Mon, 22 Jan 2024 21:34:21 +0000
Subject: [PATCH] Contributed relu triton kernl, removed a small amount of
 boilerplate from the layernorm, added test suite for kernls.

---
 tests/kernels/test_layernorm.py | 56 +++++++++++++++++++++++++++++++++
 tests/kernels/test_relu.py      | 27 ++++++++++++++++
 tests/kernels/test_rmsprop.py   |  0
 unsloth/__init__.py             |  3 +-
 unsloth/kernels/layernorm.py    | 10 +-----
 unsloth/kernels/relu.py         | 55 ++++++++++++++++++++++++++++++++
 6 files changed, 141 insertions(+), 10 deletions(-)
 create mode 100644 tests/kernels/test_layernorm.py
 create mode 100644 tests/kernels/test_relu.py
 delete mode 100644 tests/kernels/test_rmsprop.py
 create mode 100644 unsloth/kernels/relu.py

diff --git a/tests/kernels/test_layernorm.py b/tests/kernels/test_layernorm.py
new file mode 100644
index 0000000000..4f65a14b2f
--- /dev/null
+++ b/tests/kernels/test_layernorm.py
@@ -0,0 +1,56 @@
+import pytest
+import torch
+import triton
+
+from unsloth.kernels.layernorm import LayerNorm
+from tests.conftest import set_seed 
+
+# Fixture for test matrices and associated parameters
+@set_seed()
+@pytest.fixture(params=[(64, 64), (1024, 512), (2048, 1024)])
+def test_data(request):
+    torch.manual_seed(0)  # For reproducibility
+    batch_size, num_features = request.param
+    x = torch.randn(batch_size, num_features, device='cuda')
+    weight = torch.randn(num_features, device='cuda')
+    bias = torch.randn(num_features, device='cuda')
+    eps = 1e-5
+    return x, weight, bias, eps
+
+# Test forward pass
+def test_layer_norm_forward(test_data):
+    x, weight, bias, eps = test_data
+
+    # Triton layer norm forward
+    triton_output = LayerNorm.apply(x, x.size(), weight, bias, eps)
+
+    # PyTorch layer norm forward
+    pytorch_layer_norm = torch.nn.LayerNorm(x.size()[1:], eps=eps, elementwise_affine=True)
+    pytorch_layer_norm.weight = torch.nn.Parameter(weight)
+    pytorch_layer_norm.bias = torch.nn.Parameter(bias)
+    pytorch_output = pytorch_layer_norm(x)
+
+    # Check if outputs are close
+    assert torch.allclose(triton_output, pytorch_output, rtol=1e-05, atol=1e-08), \
+           "Forward pass outputs differ between Triton and PyTorch."
+
+
+def test_layer_norm_backward(test_data):
+    x, weight, bias, eps = test_data
+    x.requires_grad = True
+
+    # Triton layer norm backward
+    triton_output = LayerNorm.apply(x, x.size(), weight, bias, eps)
+    triton_grad = torch.autograd.grad(triton_output.sum(), x)[0]
+
+    # PyTorch layer norm backward
+    pytorch_layer_norm = torch.nn.LayerNorm(x.size()[1:], eps=eps, elementwise_affine=True)
+    pytorch_layer_norm.weight = torch.nn.Parameter(weight)
+    pytorch_layer_norm.bias = torch.nn.Parameter(bias)
+    pytorch_output = pytorch_layer_norm(x)
+    pytorch_output.sum().backward()
+    pytorch_grad = x.grad
+
+    # Check if gradients are close
+    assert torch.allclose(triton_grad, pytorch_grad, rtol=1e-05, atol=1e-08), \
+           "Backward pass gradients differ between Triton and PyTorch."
diff --git a/tests/kernels/test_relu.py b/tests/kernels/test_relu.py
new file mode 100644
index 0000000000..3c2a09ec34
--- /dev/null
+++ b/tests/kernels/test_relu.py
@@ -0,0 +1,27 @@
+import pytest
+import torch
+import triton
+
+from unsloth.kernels.relu import relu_kernel  # Import your relu_kernel function
+from tests.conftest import set_seed
+
+@set_seed
+@pytest.fixture(params=[(100, 100), (1024, 1024), (5000, 1024), (12345, 5678)])
+def test_matrix(request):
+    shape = request.param
+    x = torch.randn(shape, device='cuda')
+    return x
+
+# Test function
+def test_relu_kernel(test_matrix):
+    # Apply your Triton-based ReLU kernel
+    triton_output = relu_kernel(test_matrix)
+
+    # Apply PyTorch's ReLU for comparison
+    torch_relu = torch.nn.ReLU()
+    torch_output = torch_relu(test_matrix)
+
+    # Check if the outputs are close enough
+    # You can adjust rtol and atol based on the precision you expect
+    assert torch.allclose(triton_output, torch_output, rtol=1e-05, atol=1e-08), \
+           "The outputs are not close enough between Triton and PyTorch implementation."
diff --git a/tests/kernels/test_rmsprop.py b/tests/kernels/test_rmsprop.py
deleted file mode 100644
index e69de29bb2..0000000000
diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index c3baac21d4..ba55e0acfc 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+""""# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -83,3 +83,4 @@ pass
 
 from .models import *
 from .save import *
+"""
\ No newline at end of file
diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py
index 4d83b607ce..8ca79ed40c 100644
--- a/unsloth/kernels/layernorm.py
+++ b/unsloth/kernels/layernorm.py
@@ -11,19 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import torch
 
+import torch
 import triton
 import triton.language as tl
 
-try:
-    # This is https://github.com/NVIDIA/apex, NOT the apex on PyPi, so it
-    # should not be added to extras_require in setup.py.
-    import apex
-    HAS_APEX = True
-except ModuleNotFoundError:
-    HAS_APEX = False
-
 
 @triton.jit
 def _layer_norm_fwd_fused(
diff --git a/unsloth/kernels/relu.py b/unsloth/kernels/relu.py
new file mode 100644
index 0000000000..072c39dd17
--- /dev/null
+++ b/unsloth/kernels/relu.py
@@ -0,0 +1,55 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import triton 
+import triton.language as tl 
+import torch 
+from .utils import calculate_settings
+
+@triton.jit 
+def _relu_kernel(output_ptr, input_ptr, n_elements, n_cols, BLOCK_SIZE : tl.constexpr,): 
+    row_idx = tl.program_id(0)
+    row_start_ptr = input_ptr + row_idx * n_elements
+
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    input_ptrs = row_start_ptr + col_offsets
+    mask = col_offsets < n_cols
+
+    #x = max(0,x)
+    row = tl.load(input_ptrs, mask=mask, other=0)
+    row_relu = row * (row>0)
+
+    output_row_start_ptr = output_ptr + row_idx * n_elements
+    output_ptrs = output_row_start_ptr + col_offsets
+    tl.store(output_ptrs, row_relu, mask=mask)
+pass 
+
+def relu_kernel(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
+    _relu_kernel[(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