331 lines
12 KiB
Python
331 lines
12 KiB
Python
import math
|
|
import unittest
|
|
|
|
import torch
|
|
|
|
from backbones import (
|
|
SharedEventTrajectoryCore,
|
|
SharedTrajectoryMixer,
|
|
TrajectoryCrossAttention,
|
|
)
|
|
from models import (
|
|
EVENT_TRAJECTORY_ARCHITECTURE,
|
|
MODEL_SIZE_PRESETS,
|
|
DeepHealth,
|
|
resolve_model_size,
|
|
validate_event_trajectory_config,
|
|
validate_event_trajectory_state_dict,
|
|
)
|
|
from train_util import get_model_parameter_counts
|
|
|
|
|
|
def build_test_model(
|
|
*,
|
|
target_mode: str = "next_token",
|
|
time_mode: str = "absolute",
|
|
n_reasoning_rounds: int = 3,
|
|
) -> DeepHealth:
|
|
return DeepHealth(
|
|
vocab_size=32,
|
|
model_size="nano",
|
|
n_reasoning_rounds=n_reasoning_rounds,
|
|
n_types=2,
|
|
n_cont_types=0,
|
|
n_categories=2,
|
|
cont_type_ids=[],
|
|
target_mode=target_mode,
|
|
time_mode=time_mode,
|
|
)
|
|
|
|
|
|
def model_inputs() -> dict[str, torch.Tensor]:
|
|
return {
|
|
"event_seq": torch.tensor([[1, 2, 3, 4], [5, 6, 0, 0]]),
|
|
"time_seq": torch.tensor(
|
|
[[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 0.0, 0.0]]
|
|
),
|
|
"sex": torch.tensor([0, 1]),
|
|
"padding_mask": torch.tensor(
|
|
[[True, True, True, True], [True, True, False, False]]
|
|
),
|
|
"other_type": torch.zeros(2, 1, dtype=torch.long),
|
|
"other_value": torch.zeros(2, 1),
|
|
"other_value_kind": torch.zeros(2, 1, dtype=torch.long),
|
|
"other_time": torch.zeros(2, 1),
|
|
}
|
|
|
|
|
|
class EventTrajectoryBackboneTest(unittest.TestCase):
|
|
def test_model_size_presets(self) -> None:
|
|
expected = {
|
|
"nano": (120, 6, 20, 24),
|
|
"tiny": (256, 8, 32, 32),
|
|
"small": (512, 8, 64, 32),
|
|
"medium": (768, 12, 64, 48),
|
|
"huge": (1024, 16, 64, 64),
|
|
}
|
|
self.assertEqual(set(MODEL_SIZE_PRESETS), set(expected))
|
|
for name, values in expected.items():
|
|
preset = resolve_model_size(name)
|
|
self.assertEqual(
|
|
(
|
|
preset.d_model,
|
|
preset.n_trajectory,
|
|
preset.trajectory_dim,
|
|
preset.traj_hidden,
|
|
),
|
|
values,
|
|
)
|
|
|
|
def test_default_mixer_shapes_and_parameter_count(self) -> None:
|
|
mixer = SharedTrajectoryMixer(
|
|
n_trajectory=8,
|
|
trajectory_dim=32,
|
|
)
|
|
state = torch.randn(2, 5, 8, 32)
|
|
self.assertEqual(mixer(state).shape, state.shape)
|
|
self.assertEqual(mixer.traj_hidden, 32)
|
|
self.assertEqual(tuple(mixer.gate_proj.shape), (32, 8, 32))
|
|
self.assertEqual(tuple(mixer.value_proj.shape), (32, 8, 32))
|
|
self.assertEqual(tuple(mixer.output_proj.shape), (32, 32, 8))
|
|
self.assertEqual(
|
|
get_model_parameter_counts(mixer),
|
|
{
|
|
"model_parameter_count": 24_576,
|
|
"trainable_parameter_count": 24_576,
|
|
},
|
|
)
|
|
|
|
def test_reasoning_rounds_share_one_core_parameter_set(self) -> None:
|
|
core_one = SharedEventTrajectoryCore(
|
|
d_model=256,
|
|
n_trajectory=8,
|
|
n_reasoning_rounds=1,
|
|
)
|
|
core_twelve = SharedEventTrajectoryCore(
|
|
d_model=256,
|
|
n_trajectory=8,
|
|
n_reasoning_rounds=12,
|
|
)
|
|
self.assertEqual(
|
|
sum(p.numel() for p in core_one.parameters()),
|
|
sum(p.numel() for p in core_twelve.parameters()),
|
|
)
|
|
self.assertAlmostEqual(core_one.attn_scale.item(), 1.0)
|
|
self.assertAlmostEqual(
|
|
core_twelve.attn_scale.item(),
|
|
1.0 / math.sqrt(12),
|
|
places=6,
|
|
)
|
|
self.assertAlmostEqual(
|
|
core_twelve.mixer_scale.item(),
|
|
1.0 / math.sqrt(12),
|
|
places=6,
|
|
)
|
|
|
|
def test_all_masked_attention_is_finite_and_zero(self) -> None:
|
|
attention = TrajectoryCrossAttention(
|
|
d_model=32,
|
|
n_trajectory=4,
|
|
)
|
|
memory = torch.randn(2, 3, 32)
|
|
key_value = attention.project_event_memory(memory)
|
|
state = torch.randn(2, 2, 4, 8)
|
|
invalid_mask = torch.ones(2, 2, 3, dtype=torch.bool)
|
|
output = attention(
|
|
trajectory_state=state,
|
|
event_key_value=key_value,
|
|
event_invalid_mask=invalid_mask,
|
|
)
|
|
self.assertTrue(torch.isfinite(output).all())
|
|
torch.testing.assert_close(output, torch.zeros_like(output))
|
|
|
|
def test_next_token_future_events_do_not_change_earlier_query(self) -> None:
|
|
torch.manual_seed(0)
|
|
model = build_test_model(n_reasoning_rounds=2)
|
|
model.eval()
|
|
inputs = model_inputs()
|
|
original = model(**inputs)
|
|
changed_inputs = dict(inputs)
|
|
changed_inputs["event_seq"] = inputs["event_seq"].clone()
|
|
changed_inputs["event_seq"][0, 3] = 9
|
|
changed = model(**changed_inputs)
|
|
torch.testing.assert_close(original[0, 1], changed[0, 1])
|
|
|
|
def test_next_token_later_equal_time_event_is_not_visible(self) -> None:
|
|
torch.manual_seed(0)
|
|
model = build_test_model(n_reasoning_rounds=2)
|
|
model.eval()
|
|
inputs = model_inputs()
|
|
inputs["time_seq"] = inputs["time_seq"].clone()
|
|
inputs["time_seq"][0] = torch.tensor([1.0, 1.0, 2.0, 3.0])
|
|
original = model(**inputs)
|
|
changed_inputs = dict(inputs)
|
|
changed_inputs["event_seq"] = inputs["event_seq"].clone()
|
|
changed_inputs["event_seq"][0, 1] = 9
|
|
changed = model(**changed_inputs)
|
|
torch.testing.assert_close(original[0, 0], changed[0, 0])
|
|
|
|
def test_padding_content_does_not_change_valid_queries(self) -> None:
|
|
torch.manual_seed(0)
|
|
model = build_test_model(
|
|
time_mode="relative",
|
|
n_reasoning_rounds=2,
|
|
)
|
|
model.eval()
|
|
inputs = model_inputs()
|
|
original = model(**inputs)
|
|
changed_inputs = dict(inputs)
|
|
changed_inputs["event_seq"] = inputs["event_seq"].clone()
|
|
changed_inputs["time_seq"] = inputs["time_seq"].clone()
|
|
changed_inputs["event_seq"][1, 2:] = torch.tensor([9, 10])
|
|
changed_inputs["time_seq"][1, 2:] = torch.tensor([30.0, 40.0])
|
|
changed = model(**changed_inputs)
|
|
torch.testing.assert_close(original[1, :2], changed[1, :2])
|
|
|
|
def test_next_token_and_all_future_output_contracts(self) -> None:
|
|
inputs = model_inputs()
|
|
next_model = build_test_model(target_mode="next_token")
|
|
next_hidden = next_model(**inputs)
|
|
self.assertEqual(tuple(next_hidden.shape), (2, 4, 120))
|
|
next_output = next_model(**inputs, return_output=True)
|
|
self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 120))
|
|
self.assertEqual(tuple(next_output.padding_mask.shape), (2, 4))
|
|
|
|
future_model = build_test_model(target_mode="all_future")
|
|
future_hidden = future_model(
|
|
**inputs,
|
|
t_query=torch.tensor([5.0, 3.0]),
|
|
)
|
|
self.assertEqual(tuple(future_hidden.shape), (2, 120))
|
|
|
|
def test_model_contains_one_shared_core_and_no_block_stack(self) -> None:
|
|
model = build_test_model(n_reasoning_rounds=12)
|
|
self.assertFalse(hasattr(model, "blocks"))
|
|
reasoning_keys = [
|
|
key
|
|
for key in model.state_dict()
|
|
if key.startswith("reasoning_core.")
|
|
]
|
|
self.assertTrue(reasoning_keys)
|
|
self.assertFalse(any("blocks." in key for key in model.state_dict()))
|
|
self.assertFalse(any("out_proj" in key for key in reasoning_keys))
|
|
self.assertFalse(any("group_align" in key for key in reasoning_keys))
|
|
|
|
def test_event_key_and_value_are_projected_once_per_forward(self) -> None:
|
|
model = build_test_model(n_reasoning_rounds=12)
|
|
call_counts = {"key": 0, "value": 0}
|
|
|
|
def count_key(*_args) -> None:
|
|
call_counts["key"] += 1
|
|
|
|
def count_value(*_args) -> None:
|
|
call_counts["value"] += 1
|
|
|
|
key_handle = (
|
|
model.reasoning_core.cross_attention.k_proj
|
|
.register_forward_hook(count_key)
|
|
)
|
|
value_handle = (
|
|
model.reasoning_core.cross_attention.v_proj
|
|
.register_forward_hook(count_value)
|
|
)
|
|
try:
|
|
model(**model_inputs())
|
|
finally:
|
|
key_handle.remove()
|
|
value_handle.remove()
|
|
self.assertEqual(call_counts, {"key": 1, "value": 1})
|
|
|
|
def test_relative_time_forward_and_backward_are_finite(self) -> None:
|
|
torch.manual_seed(0)
|
|
model = build_test_model(
|
|
target_mode="all_future",
|
|
time_mode="relative",
|
|
n_reasoning_rounds=2,
|
|
)
|
|
hidden = model(
|
|
**model_inputs(),
|
|
t_query=torch.tensor([5.0, 3.0]),
|
|
)
|
|
(hidden * torch.randn_like(hidden)).sum().backward()
|
|
self.assertTrue(torch.isfinite(hidden).all())
|
|
self.assertIsNotNone(model.event_projection.weight.grad)
|
|
self.assertTrue(torch.isfinite(model.event_projection.weight.grad).all())
|
|
time_scale = model.reasoning_core.cross_attention.time_bias_scale
|
|
self.assertIsNotNone(time_scale)
|
|
self.assertIsNotNone(time_scale.grad)
|
|
self.assertGreater(abs(float(time_scale.grad)), 0.0)
|
|
|
|
def test_architecture_marker_and_checkpoint_are_required(self) -> None:
|
|
validate_event_trajectory_config(
|
|
{
|
|
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
|
|
"model_size": "nano",
|
|
"d_model": 120,
|
|
"n_trajectory": 6,
|
|
"trajectory_dim": 20,
|
|
"traj_hidden": 24,
|
|
"n_reasoning_rounds": 3,
|
|
}
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
|
|
validate_event_trajectory_config(
|
|
{"model_architecture": "traj_mixer_v2"}
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
|
|
validate_event_trajectory_config(
|
|
{"model_architecture": "event_trajectory_shared_v1"}
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "trajectory_dim"):
|
|
validate_event_trajectory_config(
|
|
{
|
|
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
|
|
"model_size": "nano",
|
|
"d_model": 120,
|
|
"n_trajectory": 6,
|
|
"trajectory_dim": 10,
|
|
"traj_hidden": 24,
|
|
"n_reasoning_rounds": 3,
|
|
}
|
|
)
|
|
|
|
model = build_test_model()
|
|
state_dict = model.state_dict()
|
|
validate_event_trajectory_state_dict(
|
|
state_dict,
|
|
expected_d_model=120,
|
|
expected_n_trajectory=6,
|
|
expected_n_reasoning_rounds=3,
|
|
)
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
"Checkpoint architecture does not match",
|
|
):
|
|
validate_event_trajectory_state_dict(
|
|
state_dict,
|
|
expected_n_reasoning_rounds=12,
|
|
)
|
|
state_dict.pop("reasoning_core.attn_scale")
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
"not a shared event-trajectory checkpoint",
|
|
):
|
|
validate_event_trajectory_state_dict(state_dict)
|
|
|
|
def test_unknown_model_size_is_rejected(self) -> None:
|
|
with self.assertRaisesRegex(ValueError, "Unknown model_size"):
|
|
DeepHealth(
|
|
vocab_size=32,
|
|
model_size="giant",
|
|
n_reasoning_rounds=2,
|
|
n_types=2,
|
|
n_cont_types=0,
|
|
n_categories=2,
|
|
cont_type_ids=[],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|