File size: 2,044 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 69 70 71 72 |
import os
import pytest
import torch
from fla.models import Mamba2Config, Mamba2ForCausalLM
from fla.utils import device
# ===================================================================================
# Test for Modeling (Forward/Backward Pass)
# ===================================================================================
@pytest.mark.parametrize(
['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype', 'conv_backend'],
[
pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}-conv-{}".format(*test))
for test in [
(4, 4, 1024, 4, 64, True, torch.bfloat16, 'cuda'),
(4, 4, 1024, 4, 64, False, torch.bfloat16, 'cuda'),
(4, 4, 1024, 4, 128, False, torch.bfloat16, 'cuda'),
]
],
)
def test_modeling(
L: int,
B: int,
T: int,
H: int,
D: int,
use_l2warp: bool,
dtype: torch.dtype,
conv_backend: str,
):
"""
Test the forward and backward pass of the Mamba2 model by manually
instantiating the configuration and the model.
"""
os.environ['FLA_CONV_BACKEND'] = conv_backend
# Manually create a consistent configuration
# The key relationship is: num_heads = expand * hidden_size / head_dim
# To ensure consistency, we derive hidden_size from other parameters.
expand = 2
hidden_size = H * D // expand
config = Mamba2Config(
num_hidden_layers=L,
hidden_size=hidden_size,
expand=expand,
num_heads=H,
head_dim=D,
use_l2warp=use_l2warp,
vocab_size=1000, # dummy vocab size
)
model = Mamba2ForCausalLM(config).to(device=device, dtype=dtype)
model.eval()
# Create random input tensor
x = torch.randint(0, config.vocab_size, (B, T), device=device)
# Forward pass
y = model(x)
# Assert output shape is correct
assert y.logits.shape == (B, T, config.vocab_size)
# Backward pass
y.logits.sum().backward()
print(f"Test test_modeling passed with H={H}, D={D}, backend={conv_backend}.")
|