File size: 2,219 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 58 59 60 61 62 63 64 65 66 67 68 |
import pytest
import torch
from transformers import AutoModelForCausalLM
from fla.models import GatedDeltaProductConfig
from fla.utils import device
from .test_modeling_base import run_test_generation, run_test_model_forward_backward
from .test_modeling_utils import init_weights_recursively
# ===================================================================================
# 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, GatedDeltaProductConfig, use_l2warp=use_l2warp, dtype=dtype)
# ===================================================================================
# Test for Generation
# ===================================================================================
@pytest.mark.parametrize(
['L', 'B', 'T', 'use_forget_gate', 'num_householders', 'dtype'],
[
pytest.param(*test, id="L{}-B{}-T{}-use_forget_gate{}-num_householders{}".format(*test))
for test in [
(1, 3, 2000, False, 2, torch.float16),
(2, 4, 4000, True, 3, torch.float16),
]
],
)
def test_generation(
L: int,
B: int,
T: int,
use_forget_gate: bool,
num_householders: int,
dtype: torch.dtype,
):
config = GatedDeltaProductConfig()
config.num_hidden_layers = L
config.use_forget_gate = use_forget_gate
config.num_householders = num_householders
model = AutoModelForCausalLM.from_config(config)
model.apply(init_weights_recursively)
model = model.to(dtype).to(device)
run_test_generation(L, B, T, None, None, GatedDeltaProductConfig, dtype, model=model, config=config, tol=3e-3)
|