File size: 4,367 Bytes
c335050 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
import pytest
import torch
import torch.nn.functional as F
from fla.modules.activations import logsigmoid, sigmoid, swiglu, swiglu_linear, swish
from fla.utils import assert_close, device
@pytest.mark.parametrize(
('B', 'T', 'D', 'compile'),
[
(1, 1, 64, False),
(2, 500, 128, False),
(2, 512, 128, True),
(3, 2048, 1200, True),
],
)
def test_sigmoid(B: int, T: int, D: int, compile: bool):
torch.manual_seed(42)
x = torch.randn(B, T, D, device=device, requires_grad=True)
y_ref = torch.sigmoid(x)
y_tri = sigmoid(x) if not compile else torch.compile(sigmoid)(x)
g = torch.randn_like(y_ref)
dx_ref = torch.autograd.grad(y_ref, x, g)[0]
dx_tri = torch.autograd.grad(y_tri, x, g)[0]
assert_close('sigmoid fwd', y_ref, y_tri, 1e-3)
assert_close('sigmoid bwd', dx_ref, dx_tri, 1e-3)
@pytest.mark.parametrize(
('B', 'T', 'D', 'temperature', 'compile'),
[
(1, 1, 64, 1.0, False),
(2, 500, 128, 0.5, False),
(2, 512, 128, 0.5, True),
(3, 2048, 1200, 2.0, True),
],
)
def test_logsigmoid(B: int, T: int, D: int, temperature: float, compile: bool):
torch.manual_seed(42)
x = torch.randn(B, T, D, device=device, requires_grad=True)
y_ref = F.logsigmoid(x) / temperature
y_tri = logsigmoid(x, temperature) if not compile else torch.compile(logsigmoid)(x, temperature)
g = torch.randn_like(y_ref)
dx_ref = torch.autograd.grad(y_ref, x, g)[0]
dx_tri = torch.autograd.grad(y_tri, x, g)[0]
assert_close('logsigmoid fwd', y_ref, y_tri, 1e-3)
assert_close('logsigmoid bwd', dx_ref, dx_tri, 1e-3)
@pytest.mark.parametrize(
('B', 'T', 'D', 'compile'),
[
(1, 1, 64, True),
(2, 500, 128, True),
(2, 512, 128, False),
(3, 2048, 1200, False),
],
)
def test_swish(B: int, T: int, D: int, compile: bool):
torch.manual_seed(42)
x = torch.randn(B, T, D, device=device, requires_grad=True)
y_ref = F.silu(x)
y_tri = swish(x) if not compile else torch.compile(swish)(x)
g = torch.randn_like(y_ref)
dx_ref = torch.autograd.grad(y_ref, x, g)[0]
dx_tri = torch.autograd.grad(y_tri, x, g)[0]
assert_close('swish fwd', y_ref, y_tri, 1e-3)
assert_close('swish bwd', dx_ref, dx_tri, 1e-3)
@pytest.mark.parametrize(
('B', 'T', 'D', 'compile'),
[
(1, 1, 64, True),
(2, 500, 128, True),
(2, 512, 128, False),
(3, 2048, 1200, False),
],
)
def test_swiglu(B: int, T: int, D: int, compile: bool):
torch.manual_seed(42)
x = torch.randn(B, T, D, device=device, requires_grad=True)
y = torch.randn(B, T, D, device=device, requires_grad=True)
y_ref = F.silu(x) * y
y_tri = swiglu(x, y) if not compile else torch.compile(swiglu)(x, y)
g = torch.randn_like(y_ref)
dx_ref, dy_ref = torch.autograd.grad(y_ref, (x, y), g)
dx_tri, dy_tri = torch.autograd.grad(y_tri, (x, y), g)
assert_close('swiglu fwd', y_ref, y_tri, 1e-3)
assert_close('swiglu dx', dx_ref, dx_tri, 1e-3)
assert_close('swiglu dy', dy_ref, dy_tri, 1e-3)
@pytest.mark.parametrize(
('B', 'T', 'D', 'O', 'compile'),
[
(2, 512, 128, 256, True),
(1, 1, 64, 32, False),
(2, 500, 128, 64, True),
(3, 2048, 1200, 600, False),
],
)
def test_swiglu_linear(B: int, T: int, D: int, O: int, compile: bool): # noqa: E741
torch.manual_seed(42)
x = torch.randn(B, T, D, device=device, requires_grad=True)
y = torch.randn(B, T, D, device=device, requires_grad=True)
w = torch.randn(O, D, device=device, requires_grad=True)
b = torch.randn(O, device=device, requires_grad=True)
z_ref = F.silu(x) * y
out_ref = F.linear(z_ref, w, b)
out_tri = swiglu_linear(x, y, w, b) if not compile else torch.compile(swiglu_linear)(x, y, w, b)
g = torch.randn_like(out_ref)
dx_ref, dy_ref, dw_ref, db_ref = torch.autograd.grad(out_ref, (x, y, w, b), g)
dx_tri, dy_tri, dw_tri, db_tri = torch.autograd.grad(out_tri, (x, y, w, b), g)
assert_close('swiglu_linear out', out_ref, out_tri, 1e-3)
assert_close('swiglu_linear dx', dx_ref, dx_tri, 1e-3)
assert_close('swiglu_linear dy', dy_ref, dy_tri, 1e-3)
assert_close('swiglu_linear dw', dw_ref, dw_tri, 1e-3)
assert_close('swiglu_linear db', db_ref, db_tri, 1e-3)
|