Implement shared event-trajectory reasoning backbone
This commit is contained in:
325
test_event_trajectory_backbone.py
Normal file
325
test_event_trajectory_backbone.py
Normal file
@@ -0,0 +1,325 @@
|
||||
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": (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, 256))
|
||||
next_output = next_model(**inputs, return_output=True)
|
||||
self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 256))
|
||||
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, 256))
|
||||
|
||||
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": 256,
|
||||
"n_trajectory": 8,
|
||||
"trajectory_dim": 32,
|
||||
"traj_hidden": 32,
|
||||
"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, "trajectory_dim"):
|
||||
validate_event_trajectory_config(
|
||||
{
|
||||
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
|
||||
"model_size": "nano",
|
||||
"d_model": 256,
|
||||
"n_trajectory": 8,
|
||||
"trajectory_dim": 16,
|
||||
"traj_hidden": 32,
|
||||
"n_reasoning_rounds": 3,
|
||||
}
|
||||
)
|
||||
|
||||
model = build_test_model()
|
||||
state_dict = model.state_dict()
|
||||
validate_event_trajectory_state_dict(
|
||||
state_dict,
|
||||
expected_d_model=256,
|
||||
expected_n_trajectory=8,
|
||||
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()
|
||||
Reference in New Issue
Block a user