File size: 1,570 Bytes
b66f552 | 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 |
import pytest
import torch
from fla.models import CombaConfig
from .test_modeling_base import run_test_generation, run_test_model_forward_backward
# ===================================================================================
# Test for Modeling (Forward/Backward Pass)
# ===================================================================================
@pytest.mark.parametrize(
['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'],
[
pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test))
for test in [
(4, 4, 1024, 4, 64, True, torch.bfloat16),
(4, 4, 1024, 4, 64, False, torch.bfloat16),
(4, 4, 1024, 4, 128, False, torch.bfloat16),
]
],
)
def test_modeling(
L: int,
B: int,
T: int,
H: int,
D: int,
use_l2warp: bool,
dtype: torch.dtype,
):
run_test_model_forward_backward(L, B, T, H, D, CombaConfig, use_l2warp=use_l2warp, dtype=dtype)
# ===================================================================================
# Test for Generation
# ===================================================================================
@pytest.mark.parametrize(
['L', 'B', 'T', 'H', 'D', 'dtype'],
[
pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test))
for test in [
(2, 4, 2000, 8, 64, torch.float16),
]
],
)
def test_generation(
L: int,
B: int,
T: int,
H: int,
D: int,
dtype: torch.dtype,
):
run_test_generation(L, B, T, H, D, CombaConfig, dtype)
|