seeker
70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
import torch
|
|
from transformers import Qwen2Config
|
|
|
|
from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert, PrefixKV
|
|
|
|
|
|
def tiny_expert(non_causal_attention: bool = False) -> KVCachedDiffusionExpert:
|
|
vlm_config = Qwen2Config(
|
|
vocab_size=32,
|
|
hidden_size=64,
|
|
intermediate_size=128,
|
|
num_hidden_layers=2,
|
|
num_attention_heads=4,
|
|
num_key_value_heads=2,
|
|
max_position_embeddings=128,
|
|
)
|
|
architecture = ExpertArchitecture(
|
|
hidden_size=32,
|
|
intermediate_size=64,
|
|
num_layers=2,
|
|
num_attention_heads=2,
|
|
num_key_value_heads=2,
|
|
head_dim=16,
|
|
action_projection_hidden_size=16,
|
|
action_fourier_features=4,
|
|
non_causal_attention=non_causal_attention,
|
|
)
|
|
return KVCachedDiffusionExpert(vlm_config, architecture)
|
|
|
|
|
|
def test_expert_consumes_compatible_past_kv() -> None:
|
|
expert = tiny_expert()
|
|
prefix = PrefixKV(
|
|
keys=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
|
|
values=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
|
|
)
|
|
output = expert(torch.randn(1, 3, 2), torch.rand(1, 1, 1), prefix)
|
|
assert output.shape == (1, 3, 2)
|
|
|
|
|
|
def test_expert_rejects_incompatible_kv_heads() -> None:
|
|
expert = tiny_expert()
|
|
prefix = PrefixKV(
|
|
keys=tuple(torch.randn(1, 1, 5, 16) for _ in range(2)),
|
|
values=tuple(torch.randn(1, 1, 5, 16) for _ in range(2)),
|
|
)
|
|
try:
|
|
expert.validate_prefix(prefix)
|
|
except ValueError as error:
|
|
assert "KV-head mismatch" in str(error)
|
|
else:
|
|
raise AssertionError("Expected incompatible KV heads to be rejected")
|
|
|
|
|
|
def test_non_causal_expert_allows_future_action_attention() -> None:
|
|
torch.manual_seed(3)
|
|
expert = tiny_expert(non_causal_attention=True).eval()
|
|
prefix = PrefixKV(
|
|
keys=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
|
|
values=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
|
|
)
|
|
first = torch.randn(1, 3, 2)
|
|
changed_future = first.clone()
|
|
changed_future[:, -1] += 0.137
|
|
timestep = torch.rand(1, 1, 1)
|
|
with torch.no_grad():
|
|
first_output = expert(first, timestep, prefix)
|
|
second_output = expert(changed_future, timestep, prefix)
|
|
assert not torch.allclose(first_output[:, 0], second_output[:, 0])
|