254 lines
9.1 KiB
Python
254 lines
9.1 KiB
Python
import unittest
|
|
|
|
import torch
|
|
|
|
from backbones import GPTBlock, TemporalAttention, TrajMixer
|
|
from models import (
|
|
TRAJ_MIXER_ARCHITECTURE,
|
|
validate_traj_mixer_config,
|
|
validate_traj_mixer_state_dict,
|
|
)
|
|
from train_util import get_model_parameter_counts
|
|
|
|
|
|
class TrajMixerTest(unittest.TestCase):
|
|
def test_zero_rbf_bias_has_live_projection_gradient(self) -> None:
|
|
attention = TemporalAttention(
|
|
n_embd=12,
|
|
n_head=3,
|
|
use_time_rope=False,
|
|
use_rbf_bias=True,
|
|
)
|
|
features = torch.randn(2, 4, 4, 16)
|
|
target = torch.randn(2, 4, 4, 3)
|
|
|
|
initial_bias = (
|
|
attention.time_bias_scale.tanh()
|
|
* attention.rbf_proj(features)
|
|
)
|
|
torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias))
|
|
|
|
loss = (initial_bias * target).sum()
|
|
loss.backward()
|
|
|
|
projection_grad = attention.rbf_proj.weight.grad
|
|
self.assertIsNotNone(projection_grad)
|
|
self.assertGreater(projection_grad.abs().sum().item(), 0.0)
|
|
|
|
with torch.no_grad():
|
|
attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3)
|
|
attention.zero_grad(set_to_none=True)
|
|
updated_bias = (
|
|
attention.time_bias_scale.tanh()
|
|
* attention.rbf_proj(features)
|
|
)
|
|
(updated_bias * target).sum().backward()
|
|
|
|
scale_grad = attention.time_bias_scale.grad
|
|
self.assertIsNotNone(scale_grad)
|
|
self.assertGreater(scale_grad.abs().item(), 0.0)
|
|
|
|
def test_default_shape_parameters_and_initialization(self) -> None:
|
|
mixer = TrajMixer(
|
|
n_embd=120,
|
|
n_head=10,
|
|
dropout=0.0,
|
|
)
|
|
|
|
x = torch.randn(2, 7, 120)
|
|
self.assertEqual(mixer(x).shape, x.shape)
|
|
self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040)
|
|
self.assertFalse(hasattr(mixer, "group_align"))
|
|
self.assertFalse(hasattr(mixer, "intra_norm"))
|
|
self.assertFalse(hasattr(mixer, "cross_norm"))
|
|
self.assertEqual(tuple(mixer.norm.normalized_shape), (120,))
|
|
self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12))
|
|
torch.testing.assert_close(
|
|
torch.sigmoid(mixer.intra_gate_logits.detach()),
|
|
torch.full((10, 12), 0.1),
|
|
)
|
|
self.assertEqual(mixer.intra_hidden, 48)
|
|
self.assertEqual(
|
|
tuple(mixer.intra_gate_proj.shape),
|
|
(10, 12, 48),
|
|
)
|
|
self.assertEqual(
|
|
tuple(mixer.intra_value_proj.shape),
|
|
(10, 12, 48),
|
|
)
|
|
self.assertEqual(
|
|
tuple(mixer.intra_output_proj.shape),
|
|
(10, 48, 12),
|
|
)
|
|
self.assertEqual(mixer.hidden_group, 40)
|
|
self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40))
|
|
self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40))
|
|
self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10))
|
|
|
|
def test_zero_final_output_projection_makes_mixer_identity(self) -> None:
|
|
torch.manual_seed(0)
|
|
mixer = TrajMixer(120, n_head=10, dropout=0.0)
|
|
with torch.no_grad():
|
|
mixer.output_proj.zero_()
|
|
x = torch.randn(2, 5, 120)
|
|
torch.testing.assert_close(mixer(x), x)
|
|
|
|
def test_forward_matches_single_outer_residual_formula(self) -> None:
|
|
torch.manual_seed(0)
|
|
mixer = TrajMixer(120, n_head=10, dropout=0.0)
|
|
mixer.eval()
|
|
x = torch.randn(2, 5, 120)
|
|
|
|
grouped = mixer.norm(x).reshape(2, 5, 10, 12)
|
|
intra_output = mixer._intra_mix(grouped)
|
|
static_gate = torch.sigmoid(mixer.intra_gate_logits).view(
|
|
1, 1, 10, 12
|
|
)
|
|
mixed_input = grouped + static_gate * intra_output
|
|
update = mixer._cross_mix(mixed_input).reshape(2, 5, 120)
|
|
|
|
torch.testing.assert_close(mixer(x), x + update)
|
|
|
|
def test_intra_stage_is_independent_across_groups(self) -> None:
|
|
torch.manual_seed(0)
|
|
mixer = TrajMixer(120, n_head=10, dropout=0.0)
|
|
mixer.eval()
|
|
|
|
grouped = torch.randn(2, 4, 10, 12)
|
|
changed = grouped.clone()
|
|
changed[:, :, 3, :] += torch.randn_like(changed[:, :, 3, :])
|
|
|
|
original_out = mixer._intra_mix(grouped)
|
|
changed_out = mixer._intra_mix(changed)
|
|
unchanged_groups = torch.tensor([0, 1, 2, 4, 5, 6, 7, 8, 9])
|
|
torch.testing.assert_close(
|
|
original_out.index_select(2, unchanged_groups),
|
|
changed_out.index_select(2, unchanged_groups),
|
|
)
|
|
|
|
def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None:
|
|
mixer = TrajMixer(6, n_head=3, dropout=0.0)
|
|
mixer.eval()
|
|
with torch.no_grad():
|
|
mixer.gate_proj.zero_()
|
|
mixer.value_proj.zero_()
|
|
mixer.output_proj.zero_()
|
|
|
|
# For coordinate 0 only, read group 0 through hidden unit 0 and
|
|
# write the resulting gated value into group 1.
|
|
mixer.gate_proj[0, 0, 0] = 1.0
|
|
mixer.value_proj[0, 0, 0] = 1.0
|
|
mixer.output_proj[0, 0, 1] = 1.0
|
|
|
|
grouped = torch.tensor(
|
|
[[[
|
|
[-1.0, 4.0],
|
|
[0.0, 5.0],
|
|
[1.0, 6.0],
|
|
]]]
|
|
)
|
|
changed = grouped.clone()
|
|
changed[0, 0, 0, 0] = 2.0
|
|
|
|
original_out = mixer._cross_mix(grouped)
|
|
changed_out = mixer._cross_mix(changed)
|
|
|
|
self.assertNotEqual(
|
|
original_out[0, 0, 1, 0].item(),
|
|
changed_out[0, 0, 1, 0].item(),
|
|
)
|
|
torch.testing.assert_close(
|
|
original_out[..., 1],
|
|
changed_out[..., 1],
|
|
)
|
|
|
|
def test_mixer_does_not_mix_sequence_positions(self) -> None:
|
|
torch.manual_seed(0)
|
|
mixer = TrajMixer(120, n_head=10, dropout=0.0)
|
|
mixer.eval()
|
|
x = torch.randn(2, 5, 120)
|
|
changed = x.clone()
|
|
changed[:, 3, :] += torch.randn_like(changed[:, 3, :])
|
|
|
|
original_out = mixer(x)
|
|
changed_out = mixer(changed)
|
|
unchanged_positions = torch.tensor([0, 1, 2, 4])
|
|
torch.testing.assert_close(
|
|
original_out.index_select(1, unchanged_positions),
|
|
changed_out.index_select(1, unchanged_positions),
|
|
)
|
|
|
|
def test_gradients_reach_all_projection_families(self) -> None:
|
|
torch.manual_seed(1)
|
|
mixer = TrajMixer(120, n_head=10, dropout=0.0)
|
|
x = torch.randn(2, 4, 120, requires_grad=True)
|
|
|
|
mixer(x).square().mean().backward()
|
|
|
|
self.assertIsNotNone(x.grad)
|
|
for name, parameter in mixer.named_parameters():
|
|
self.assertIsNotNone(parameter.grad, name)
|
|
self.assertTrue(torch.isfinite(parameter.grad).all(), name)
|
|
|
|
def test_gpt_block_delegates_single_mixer_residual_to_traj_mixer(self) -> None:
|
|
block = GPTBlock(n_embd=120, n_head=10)
|
|
self.assertIsInstance(block.mlp, TrajMixer)
|
|
self.assertFalse(hasattr(block, "ln2"))
|
|
self.assertIsInstance(block.mlp.norm, torch.nn.LayerNorm)
|
|
self.assertFalse(hasattr(block.mlp, "intra_norm"))
|
|
self.assertFalse(hasattr(block.mlp, "cross_norm"))
|
|
|
|
x = torch.randn(2, 6, 120)
|
|
self.assertEqual(block(x).shape, x.shape)
|
|
|
|
def test_architecture_marker_is_required(self) -> None:
|
|
validate_traj_mixer_config(
|
|
{"model_architecture": TRAJ_MIXER_ARCHITECTURE}
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
|
|
validate_traj_mixer_config({})
|
|
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
|
|
validate_traj_mixer_config({"model_architecture": "delphi_swiglu"})
|
|
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
|
|
validate_traj_mixer_config(
|
|
{"model_architecture": "traj_mixer_v2"}
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
|
|
validate_traj_mixer_config(
|
|
{"model_architecture": "traj_mixer_v3"}
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
|
|
validate_traj_mixer_config(
|
|
{"model_architecture": "traj_mixer_v4"}
|
|
)
|
|
|
|
def test_checkpoint_must_contain_traj_mixer_parameters(self) -> None:
|
|
block = GPTBlock(n_embd=120, n_head=10)
|
|
state_dict = {
|
|
f"blocks.0.{key}": value
|
|
for key, value in block.state_dict().items()
|
|
}
|
|
validate_traj_mixer_state_dict(state_dict)
|
|
|
|
state_dict.pop("blocks.0.mlp.intra_gate_proj")
|
|
with self.assertRaisesRegex(ValueError, "not a TrajMixer checkpoint"):
|
|
validate_traj_mixer_state_dict(state_dict)
|
|
|
|
def test_invalid_group_partition_is_rejected(self) -> None:
|
|
with self.assertRaisesRegex(ValueError, "divisible"):
|
|
TrajMixer(n_embd=121, n_head=10)
|
|
|
|
def test_parameter_counts_match_traj_mixer_parameters(self) -> None:
|
|
mixer = TrajMixer(n_embd=120, n_head=10)
|
|
self.assertEqual(
|
|
get_model_parameter_counts(mixer),
|
|
{
|
|
"model_parameter_count": 32_040,
|
|
"trainable_parameter_count": 32_040,
|
|
},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|