tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

add manual seed for flaky test_geglu_backward test (#2763) [skip ci]

Browse Source
This commit is contained in:
Wing Lian
2025-06-05 09:23:17 -07:00
committed by GitHub
parent cb03c765a1
commit 7909bfb076
1 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/e2e/kernels/test_geglu.py
View File

@@ -1,7 +1,6 @@
"""Tests for GEGLU activation function Triton kernels.""" """Tests for GEGLU activation function Triton kernels."""
# pylint: disable=duplicate-code import pytest
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
@@ -34,8 +33,14 @@ def test_geglu_forward_values():
assert torch.allclose(triton_out, torch_out, rtol=1e-3) assert torch.allclose(triton_out, torch_out, rtol=1e-3)
def test_geglu_backward(): @pytest.mark.parametrize(
"torch_seed",
[0, 42],
)
def test_geglu_backward(torch_seed):
"""Test GEGLU backward pass matches PyTorch autograd.""" """Test GEGLU backward pass matches PyTorch autograd."""
torch.manual_seed(torch_seed)
gate = torch.randn(2, 3, 64, device="cuda", requires_grad=True) gate = torch.randn(2, 3, 64, device="cuda", requires_grad=True)
up = torch.randn(2, 3, 64, device="cuda", requires_grad=True) up = torch.randn(2, 3, 64, device="cuda", requires_grad=True)
grad_output = torch.randn(2, 3, 64, device="cuda") grad_output = torch.randn(2, 3, 64, device="cuda")