File size: 4,645 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 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 |
import pytest
import torch
from transformers.configuration_utils import PretrainedConfig
from fla.utils import assert_close, device, is_intel_alchemist, is_nvidia_hopper
from .test_modeling_utils import (
GENERATION_UNSUPPORTED,
HOPPER_EXCLUSIVE,
MODELING_UNSUPPORTED_VARLEN,
NOT_READY_FOR_TESTING,
create_model_and_config,
)
# ===================================================================================
# BASE TEST FOR MODELING (FORWARD/BACKWARD PASS)
# ===================================================================================
@pytest.mark.skipif(
is_intel_alchemist,
reason="Skipping test on Intel Alchemist due to known issues with SRAM.",
)
def run_test_model_forward_backward(
L: int,
B: int,
T: int,
H: int,
D: int,
config_class: type,
use_l2warp: bool,
dtype: torch.dtype,
):
"""
A foundational test for the forward and backward passes of a model.
"""
if not is_nvidia_hopper and D == 128:
pytest.skip("D=128 is only tested on Hopper GPUs to save CI time.")
if not is_nvidia_hopper and config_class.__name__ in HOPPER_EXCLUSIVE:
pytest.skip(f"{config_class.__name__} requires Hopper-specific features.")
if config_class.__name__ in NOT_READY_FOR_TESTING:
pytest.skip(f"{config_class.__name__} is not yet ready for testing.")
model, config = create_model_and_config(config_class, L, H, D, use_l2warp=use_l2warp, dtype=dtype)
input_ids = torch.randint(low=0, high=config.vocab_size, size=(B, T), device=device)
output_fixed = model(input_ids, output_hidden_states=True).hidden_states[-1]
assert output_fixed.shape == (B, T, config.hidden_size)
if config_class.__name__ in MODELING_UNSUPPORTED_VARLEN:
pytest.skip(f"Variable length not supported for {config_class.__name__}.")
cu_seqlens = torch.arange(0, B * T + 1, T, dtype=torch.int32, device=device)
output_var = model(
input_ids.view(1, B * T), output_hidden_states=True, cu_seqlens=cu_seqlens,
).hidden_states[-1]
assert output_var.shape == (1, B * T, config.hidden_size)
assert_close("output", output_fixed.view(1, B * T, -1), output_var, 1e-3)
output_var.backward(torch.randn_like(output_var))
# ===================================================================================
# BASE TEST FOR GENERATION (K/V CACHE)
# ===================================================================================
def run_test_generation(
L: int,
B: int,
T: int,
H: int,
D: int,
config_class: type,
dtype: torch.dtype,
use_l2warp: bool = False,
model: torch.nn.Module | None = None,
config: PretrainedConfig | None = None,
tol: float = 2e-3,
):
"""
A foundational test for K/V cache-based generation.
"""
torch.manual_seed(42)
if config_class.__name__ in GENERATION_UNSUPPORTED:
pytest.skip(f"Generation test not supported for {config_class.__name__}.")
if config_class.__name__ in NOT_READY_FOR_TESTING:
pytest.skip(f"{config_class.__name__} is not yet ready for testing.")
if model is None:
model, config = create_model_and_config(config_class, L, H, D, use_l2warp=use_l2warp, dtype=dtype)
model.eval()
model = model.to(dtype).to(device)
num_chunks = 4
chunk_size = T // num_chunks
input_ids = torch.randint(low=0, high=config.vocab_size, size=(B, T)).to(device)
attention_mask = torch.ones((B, T), dtype=torch.bool).to(device)
seq_start = torch.randint(low=1, high=chunk_size - 1, size=(B,))
attention_mask[torch.arange(T) < seq_start[:, None]] = False
ref = torch.cat([
model(input_ids=input_ids[i:i+1, start:], use_cache=False).logits
for i, start in enumerate(seq_start)
], dim=1)
logits = []
out = model(
input_ids=input_ids[:, :chunk_size],
attention_mask=attention_mask[:, :chunk_size],
use_cache=True,
past_key_values=None,
)
logits, past_key_values = [out.logits], out.past_key_values
for i in range(1, num_chunks):
start, end = i * chunk_size, (i + 1) * chunk_size
for j in range(start, end):
out = model(
input_ids=input_ids[:, j:j+1],
attention_mask=attention_mask[:, :j+1],
use_cache=True,
past_key_values=past_key_values,
)
logits.append(out.logits)
past_key_values = out.past_key_values
gen = torch.cat(logits, 1)
gen = torch.cat([gen[i:i+1, start:] for i, start in enumerate(seq_start)], 1)
assert_close('logits', ref, gen, tol)
|