-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathpatch_test_memory.py
More file actions
178 lines (139 loc) · 5.54 KB
/
Copy pathpatch_test_memory.py
File metadata and controls
178 lines (139 loc) · 5.54 KB
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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
with open("tests/test_memory.py", "r") as f:
code = f.read()
search = """ # Manual update
# v_t: (B, d) -> (B, d, 1)
# alpha_t: (B, d) -> (B, 1, d)
update = torch.matmul(v_t.unsqueeze(2), alpha_t.unsqueeze(1))
S_expected = S_prev + update"""
replace = """ # Manual update
# v_t: (B, d) -> (B, d, 1)
# alpha_t: (B, d) -> (B, 1, d)
v_t_f32 = v_t / (torch.norm(v_t, dim=-1, keepdim=True) + 1e-6)
update = torch.matmul(v_t_f32.unsqueeze(2), alpha_t.unsqueeze(1))
S_expected = S_prev + update"""
code = code.replace(search, replace)
search2 = """ # Manual sum
S_manual = torch.zeros(B, d, d)
for t in range(T):
update = torch.matmul(v_seq[t].unsqueeze(2), alpha_seq[t].unsqueeze(1))
S_manual += update"""
replace2 = """ # Manual sum
S_manual = torch.zeros(B, d, d)
for t in range(T):
v_t_f32 = v_seq[t] / (torch.norm(v_seq[t], dim=-1, keepdim=True) + 1e-6)
update = torch.matmul(v_t_f32.unsqueeze(2), alpha_seq[t].unsqueeze(1))
S_manual += update"""
code = code.replace(search2, replace2)
search3 = """ def test_stability_renorm():
\"\"\"
Test 4 — Stability test with renormalization
Construct a sequence where v_t and alpha_t have moderately large values.
Verify renormalization triggers and S_t is scaled down.
\"\"\"
d = 4
threshold = 10.0
manager = MemoryMatrixManager(d_model=d, enable_renorm=True, renorm_threshold=threshold)
B = 1
manager.reset(batch_size=B)
# Create inputs that will definitely exceed threshold
# Threshold is 10.
# If we add v @ alpha^T such that norm > 10.
# Let v = [10, 0...], alpha = [1, 0...] -> update is matrix with 10 at (0,0).
# Norm is 10. If existing S was small, new S has norm >= 10.
v = torch.zeros(B, d)
v[0, 0] = 20.0 # Large value
alpha = torch.zeros(B, d)
alpha[0, 0] = 1.0"""
replace3 = """ def test_stability_renorm():
\"\"\"
Test 4 — Stability test with renormalization
Construct a sequence where v_t and alpha_t have moderately large values.
Verify renormalization triggers and S_t is scaled down.
\"\"\"
d = 4
threshold = 10.0
manager = MemoryMatrixManager(d_model=d, enable_renorm=True, renorm_threshold=threshold)
B = 1
manager.reset(batch_size=B)
# Create inputs that will definitely exceed threshold
# Threshold is 10.
# Since v is normalized, the norm of the update is ||v|| * ||alpha|| = 1 * ||alpha|| = ||alpha||.
# So we just need to make alpha very large.
v = torch.zeros(B, d)
v[0, 0] = 20.0 # Will be normalized
alpha = torch.zeros(B, d)
alpha[0, 0] = 20.0"""
code = code.replace(search3, replace3)
search4 = """ def test_stability_no_renorm():
\"\"\"
Test stability with renorm disabled.
Norm should grow.
\"\"\"
d = 4
threshold = 10.0
manager = MemoryMatrixManager(d_model=d, enable_renorm=False, renorm_threshold=threshold)
B = 1
manager.reset(batch_size=B)
v = torch.zeros(B, d)
v[0, 0] = 20.0
alpha = torch.zeros(B, d)
alpha[0, 0] = 1.0
stats = manager.update(v, alpha)
assert stats['renorm_triggered'] == 0.0, "Renormalization should NOT have triggered"
assert stats['norm_max'] == 20.0"""
replace4 = """ def test_stability_no_renorm():
\"\"\"
Test stability with renorm disabled.
Norm should grow.
\"\"\"
d = 4
threshold = 10.0
manager = MemoryMatrixManager(d_model=d, enable_renorm=False, renorm_threshold=threshold)
B = 1
manager.reset(batch_size=B)
v = torch.zeros(B, d)
v[0, 0] = 20.0
alpha = torch.zeros(B, d)
alpha[0, 0] = 20.0
stats = manager.update(v, alpha)
assert stats['renorm_triggered'] == 0.0, "Renormalization should NOT have triggered"
assert abs(stats['norm_max'] - 20.0) < 1e-4"""
code = code.replace(search4, replace4)
search5 = """ def test_mixed_batch_renorm():
\"\"\"
Test batch where one element needs renorm and another doesn't.
\"\"\"
d = 4
threshold = 10.0
manager = MemoryMatrixManager(d_model=d, enable_renorm=True, renorm_threshold=threshold)
B = 2
manager.reset(batch_size=B)
v = torch.zeros(B, d)
alpha = torch.zeros(B, d)
# Batch 0: Large update (20)
v[0, 0] = 20.0
alpha[0, 0] = 1.0
# Batch 1: Small update (1)
v[1, 0] = 1.0
alpha[1, 0] = 1.0"""
replace5 = """ def test_mixed_batch_renorm():
\"\"\"
Test batch where one element needs renorm and another doesn't.
\"\"\"
d = 4
threshold = 10.0
manager = MemoryMatrixManager(d_model=d, enable_renorm=True, renorm_threshold=threshold)
B = 2
manager.reset(batch_size=B)
v = torch.zeros(B, d)
alpha = torch.zeros(B, d)
# Batch 0: Large update (20)
v[0, 0] = 20.0
alpha[0, 0] = 20.0
# Batch 1: Small update (1)
v[1, 0] = 1.0
alpha[1, 0] = 1.0"""
code = code.replace(search5, replace5)
with open("tests/test_memory.py", "w") as f:
f.write(code)
print("Patched test_memory.py")