ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

update test to use common torch.allclose function util

Browse source
This commit is contained in:
cm2435 2024-01-24 17:44:40 +00:00
commit 48eb887037
2 changed files with 9 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

12
tests/kernels/test_layernorm.py
View file

@ -3,7 +3,7 @@ import torch
import triton
from unsloth.kernels.layernorm import LayerNorm
from tests.conftest import set_seed
from tests.conftest import set_seed, assert_all_close
# Fixture for test matrices and associated parameters
@set_seed()
@ -30,9 +30,8 @@ def test_layer_norm_forward(test_data):
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."
# Check if outputs are close using assert_all_close
assert_all_close(triton_output, pytorch_output, rtol=1e-05, atol=1e-08)
def test_layer_norm_backward(test_data):
@ -51,6 +50,5 @@ def test_layer_norm_backward(test_data):
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."
# Check if gradients are close using assert_all_close
assert_all_close(triton_grad, pytorch_grad, rtol=1e-05, atol=1e-08)

10
tests/kernels/test_relu.py
View file

@ -2,8 +2,8 @@ import pytest
import torch
import triton
from unsloth.kernels.relu import relu_kernel # Import your relu_kernel function
from tests.conftest import set_seed
from unsloth.kernels.relu import relu_kernel
from tests.conftest import set_seed, assert_all_close
@set_seed
@pytest.fixture(params=[(100, 100), (1024, 1024), (5000, 1024), (12345, 5678)])
@ -21,7 +21,5 @@ def test_relu_kernel(test_matrix):
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."
# Check if the outputs are close enough using assert_all_close
assert_all_close(triton_output, torch_output, rtol=1e-05, atol=1e-08)