tests/test.sh exit code (0 → resolved); the classification below is post-hoc and cannot change it.Read probe.py and timer.sh
cat /app/probe.py; echo "=====TIMER====="; cat /app/timer.sh; echo "=====BUDGET====="; bash /app/timer.sh
from __future__ import annotations
import torch
import reforacle
g = torch.Generator().manual_seed(0)
values = torch.randn(8, generator=g)
mask = (torch.rand(8, generator=g) > 0.5).float()
print("inputs:")
print(" values =", values.tolist())
print(" mask =", mask.tolist())
print("reference masked_mean ->", reforacle.masked_mean(values, mask).item())
print("reference masked_mean (empty mask) ->", reforacle.masked_mean(torch.randn(4, generator=g), torch.zeros(4)).item())
print()
print("Probe any reference function this way to discover its exact numerical")
print("convention, then implement it in /app/objectives.py to match.")
=====TIMER=====
#!/bin/bash
TIMEOUT_SEC=7200
ELAPSED=$(ps -o etimes= -p 1 2>/dev/null | tr -d ' ')
if ! [[ "$ELAPSED" =~ ^[0-9]+$ ]]; then
START_FILE="$(dirname "$0")/.timer_start"
if [ ! -f "$START_FILE" ]; then
date +%s >"$START_FILE"
fi
START_DATE=$(cat "$START_FILE")
NOW=$(date +%s)
ELAPSED=$((NOW - START_DATE))
fi
REMAINING=$((TIMEOUT_SEC - ELAPSED))
if [ $REMAINING -le 0 ]; then
echo "Timer expired!"
else
echo "Remaining time (hours:minutes):"
HOURS=$((REMAINING / 3600))
MINUTES=$(((REMAINING % 3600) / 60))
printf "%d:%02d\n" $HOURS $MINUTES
fi
=====BUDGET=====
Remaining time (hours:minutes):
1:59
[stdout]
from __future__ import annotations
import torch
import reforacle
g = torch.Generator().manual_seed(0)
values = torch.randn(8, generator=g)
mask = (torch.rand(8, generator=g) > 0.5).float()
print("inputs:")
print(" values =", values.tolist())
print(" mask =", mask.tolist())
print("reference masked_mean ->", reforacle.masked_mean(values, mask).item())
print("reference masked_mean (empty mask) ->", reforacle.masked_mean(torch.randn(4, generator=g), torch.zeros(4)).item())
print()
print("Probe any reference function this way to discover its exact numerical")
print("convention, then implement it in /app/objectives.py to match.")
=====TIMER=====
#!/bin/bash
TIMEOUT_SEC=7200
ELAPSED=$(ps -o etimes= -p 1 2>/dev/null | tr -d ' ')
if ! [[ "$ELAPSED" =~ ^[0-9]+$ ]]; then
START_FILE="$(dirname "$0")/.timer_start"
if [ ! -f "$START_FILE" ]; then
date +%s >"$START_FILE"
fi
START_DATE=$(cat "$START_FILE")
NOW=$(date +%s)
ELAPSED=$((NOW - START_DATE))
fi
REMAINING=$((TIMEOUT_SEC - ELAPSED))
if [ $REMAINING -le 0 ]; then
echo "Timer expired!"
else
echo "Remaining time (hours:minutes):"
HOURS=$((REMAINING / 3600))
MINUTES=$(((REMAINING % 3600) / 60))
printf "%d:%02d\n" $HOURS $MINUTES
fi
=====BUDGET=====
Remaining time (hours:minutes):
1:59
1 from __future__ import annotations
2
3 import torch
4
5
6 def masked_mean(values, mask, axis=None):
7 raise NotImplementedError
8
9
10 def masked_sum(values, mask, axis=None):
11 raise NotImplementedError
12
13
14 def logsumexp(x, axis):
15 raise NotImplementedError
16
17
18 def log_softmax(x, axis):
19 raise NotImplementedError
20
21
22 def selective_logprobs(logits, labels, mask):
23 raise NotImplementedError
24
25
26 def token_logprobs(logits, labels):
27 raise NotImplementedError
28
29
30 def sequence_logprob(logits, labels, mask, length_normalize):
31 raise NotImplementedError
32
33
34 def entropy(logits, mask):
35 raise NotImplementedError
36
37
38 def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
39 raise NotImplementedError
40
41
42 def ipo_loss(pc, pr, rc, rr, beta):
43 raise NotImplementedError
44
45
46 def grpo_advantages(rewards, group_size, scale_by_std):
47 raise NotImplementedError
48
49
50 def gae(rewards, values, next_value, gamma, lam):
51 raise NotImplementedError
52
53
54 def kl_penalty(logp, ref_logp, estimator):
55 raise NotImplementedError
56
57
58 def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
59 raise NotImplementedError
60
61
62 def value_loss(values, old_values, returns, clip):
63 raise NotImplementedError
64
65
66 def whiten(values, mask, shift_mean):
67 raise NotImplementedError
68
69
70 def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
71 chosen_labels, rejected_labels, chosen_mask, rejected_mask,
72 beta, label_smoothing):
73 raise NotImplementedError
74
75
76 def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
77 rewards, group_size, beta, clip_low, clip_high, scale_by_std,
78 kl_estimator):
79 raise NotImplementedError
80
81
82 def ppo_objective(rewards, values, old_values, logp, old_logp, next_value,
83 gamma, lam, clip_low, clip_high, vf_clip, vf_coef):
84 raise NotImplementedError
85
86
87 def rloo_advantages(rewards, group_size):
88 raise NotImplementedError
89
90
91 def reverse_kl(logp, ref_logp):
92 raise NotImplementedError
93
94
95 def importance_ratio(logp, old_logp, clip):
96 raise NotImplementedError
97
98
99 def discounted_returns(rewards, gamma):
100 raise NotImplementedError
101
102
103 def normalize(x, eps):
104 raise NotImplementedError
105
106
107 def top_p_mask(probs, p):
108 raise NotImplementedError
109
110
111 def smoothed_nll(logits, labels, smoothing):
112 raise NotImplementedError
113
114
115 def bradley_terry_logit(chosen_reward, rejected_reward, beta):
116 raise NotImplementedError
117
118
119 def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
120 raise NotImplementedError
121
122
123 def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
124 raise NotImplementedError
125
126
127 def cross_entropy(logits, labels, ignore_index):
128 raise NotImplementedError
129
130
131 def top_k_mask(logits, k):
132 raise NotImplementedError
133
134
135 def group_mean_baseline(rewards, group_size):
136 raise NotImplementedError
137
138
139 def lambda_returns(rewards, values, next_value, gamma, lam):
140 raise NotImplementedError
141
142
143 def symmetric_kl(logp, ref_logp):
144 raise NotImplementedError
145
146
147 def huber_value_loss(values, returns, delta):
148 raise NotImplementedError
149
150
151 def normalized_entropy(logits, mask):
152 raise NotImplementedError
153
154
155 def clip_fraction(logp, old_logp, clip):
156 raise NotImplementedError
157
158
159 def masked_whiten(values, mask, shift_mean):
160 raise NotImplementedError
161
162
163 def logprob_at_temperature(logits, labels, temperature):
164 raise NotImplementedError
165
166
167 def advantage_mean_std(advantages, mask):
168 raise NotImplementedError
169
170
171 def argmax_tokens(logits):
172 raise NotImplementedError
173
174
175 def mode_label(labels):
176 raise NotImplementedError
177
178
179 def median_reward(rewards):
180 raise NotImplementedError
181
182
183 def masked_max(values, mask):
184 raise NotImplementedError
185
186
187 def quantile_lower(x, q):
188 raise NotImplementedError
189
190
191 def pad_mask_from_lengths(lengths, max_len):
192 raise NotImplementedError
193
194
195 def first_nonzero_index(mask):
196 raise NotImplementedError
197
198
199 def cumulative_max(x):
200 raise NotImplementedError
201
202
203 def bucketize_reward(rewards, boundaries):
204 raise NotImplementedError
205
/app/objectives.py
from __future__ import annotations
import torch
import reforacle
g = torch.Generator().manual_seed(0)
values = torch.randn(8, generator=g)
mask = (torch.rand(8, generator=g) > 0.5).float()
print("inputs:")
print(" values =", values.tolist())
print(" mask =", mask.tolist())
print("reference masked_mean ->", reforacle.masked_mean(values, mask).item())
print("reference masked_mean (empty mask) ->", reforacle.masked_mean(torch.randn(4, generator=g), torch.zeros(4)).item())
print()
print("Probe any reference function this way to discover its exact numerical")
print("convention, then implement it in /app/objectives.py to match.")
=====TIMER=====
#!/bin/bash
TIMEOUT_SEC=7200
ELAPSED=$(ps -o etimes= -p 1 2>/dev/null | tr -d ' ')
if ! [[ "$ELAPSED" =~ ^[0-9]+$ ]]; then
START_FILE="$(dirname "$0")/.timer_start"
if [ ! -f "$START_FILE" ]; then
date +%s >"$START_FILE"
fi
START_DATE=$(cat "$START_FILE")
NOW=$(date +%s)
ELAPSED=$((NOW - START_DATE))
fi
REMAINING=$((TIMEOUT_SEC - ELAPSED))
if [ $REMAINING -le 0 ]; then
echo "Timer expired!"
else
echo "Remaining time (hours:minutes):"
HOURS=$((REMAINING / 3600))
MINUTES=$(((REMAINING % 3600) / 60))
printf "%d:%02d\n" $HOURS $MINUTES
fi
=====BUDGET=====
Remaining time (hours:minutes):
1:59
[stdout]
from __future__ import annotations
import torch
import reforacle
g = torch.Generator().manual_seed(0)
values = torch.randn(8, generator=g)
mask = (torch.rand(8, generator=g) > 0.5).float()
print("inputs:")
print(" values =", values.tolist())
print(" mask =", mask.tolist())
print("reference masked_mean ->", reforacle.masked_mean(values, mask).item())
print("reference masked_mean (empty mask) ->", reforacle.masked_mean(torch.randn(4, generator=g), torch.zeros(4)).item())
print()
print("Probe any reference function this way to discover its exact numerical")
print("convention, then implement it in /app/objectives.py to match.")
=====TIMER=====
#!/bin/bash
TIMEOUT_SEC=7200
ELAPSED=$(ps -o etimes= -p 1 2>/dev/null | tr -d ' ')
if ! [[ "$ELAPSED" =~ ^[0-9]+$ ]]; then
START_FILE="$(dirname "$0")/.timer_start"
if [ ! -f "$START_FILE" ]; then
date +%s >"$START_FILE"
fi
START_DATE=$(cat "$START_FILE")
NOW=$(date +%s)
ELAPSED=$((NOW - START_DATE))
fi
REMAINING=$((TIMEOUT_SEC - ELAPSED))
if [ $REMAINING -le 0 ]; then
echo "Timer expired!"
else
echo "Remaining time (hours:minutes):"
HOURS=$((REMAINING / 3600))
MINUTES=$(((REMAINING % 3600) / 60))
printf "%d:%02d\n" $HOURS $MINUTES
fi
=====BUDGET=====
Remaining time (hours:minutes):
1:59
1 from __future__ import annotations
2
3 import torch
4
5
6 def masked_mean(values, mask, axis=None):
7 raise NotImplementedError
8
9
10 def masked_sum(values, mask, axis=None):
11 raise NotImplementedError
12
13
14 def logsumexp(x, axis):
15 raise NotImplementedError
16
17
18 def log_softmax(x, axis):
19 raise NotImplementedError
20
21
22 def selective_logprobs(logits, labels, mask):
23 raise NotImplementedError
24
25
26 def token_logprobs(logits, labels):
27 raise NotImplementedError
28
29
30 def sequence_logprob(logits, labels, mask, length_normalize):
31 raise NotImplementedError
32
33
34 def entropy(logits, mask):
35 raise NotImplementedError
36
37
38 def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
39 raise NotImplementedError
40
41
42 def ipo_loss(pc, pr, rc, rr, beta):
43 raise NotImplementedError
44
45
46 def grpo_advantages(rewards, group_size, scale_by_std):
47 raise NotImplementedError
48
49
50 def gae(rewards, values, next_value, gamma, lam):
51 raise NotImplementedError
52
53
54 def kl_penalty(logp, ref_logp, estimator):
55 raise NotImplementedError
56
57
58 def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
59 raise NotImplementedError
60
61
62 def value_loss(values, old_values, returns, clip):
63 raise NotImplementedError
64
65
66 def whiten(values, mask, shift_mean):
67 raise NotImplementedError
68
69
70 def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
71 chosen_labels, rejected_labels, chosen_mask, rejected_mask,
72 beta, label_smoothing):
73 raise NotImplementedError
74
75
76 def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
77 rewards, group_size, beta, clip_low, clip_high, scale_by_std,
78 kl_estimator):
79 raise NotImplementedError
80
81
82 def ppo_objective(rewards, values, old_values, logp, old_logp, next_value,
83 gamma, lam, clip_low, clip_high, vf_clip, vf_coef):
84 raise NotImplementedError
85
86
87 def rloo_advantages(rewards, group_size):
88 raise NotImplementedError
89
90
91 def reverse_kl(logp, ref_logp):
92 raise NotImplementedError
93
94
95 def importance_ratio(logp, old_logp, clip):
96 raise NotImplementedError
97
98
99 def discounted_returns(rewards, gamma):
100 raise NotImplementedError
101
102
103 def normalize(x, eps):
104 raise NotImplementedError
105
106
107 def top_p_mask(probs, p):
108 raise NotImplementedError
109
110
111 def smoothed_nll(logits, labels, smoothing):
112 raise NotImplementedError
113
114
115 def bradley_terry_logit(chosen_reward, rejected_reward, beta):
116 raise NotImplementedError
117
118
119 def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
120 raise NotImplementedError
121
122
123 def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
124 raise NotImplementedError
125
126
127 def cross_entropy(logits, labels, ignore_index):
128 raise NotImplementedError
129
130
131 def top_k_mask(logits, k):
132 raise NotImplementedError
133
134
135 def group_mean_baseline(rewards, group_size):
136 raise NotImplementedError
137
138
139 def lambda_returns(rewards, values, next_value, gamma, lam):
140 raise NotImplementedError
141
142
143 def symmetric_kl(logp, ref_logp):
144 raise NotImplementedError
145
146
147 def huber_value_loss(values, returns, delta):
148 raise NotImplementedError
149
150
151 def normalized_entropy(logits, mask):
152 raise NotImplementedError
153
154
155 def clip_fraction(logp, old_logp, clip):
156 raise NotImplementedError
157
158
159 def masked_whiten(values, mask, shift_mean):
160 raise NotImplementedError
161
162
163 def logprob_at_temperature(logits, labels, temperature):
164 raise NotImplementedError
165
166
167 def advantage_mean_std(advantages, mask):
168 raise NotImplementedError
169
170
171 def argmax_tokens(logits):
172 raise NotImplementedError
173
174
175 def mode_label(labels):
176 raise NotImplementedError
177
178
179 def median_reward(rewards):
180 raise NotImplementedError
181
182
183 def masked_max(values, mask):
184 raise NotImplementedError
185
186
187 def quantile_lower(x, q):
188 raise NotImplementedError
189
190
191 def pad_mask_from_lengths(lengths, max_len):
192 raise NotImplementedError
193
194
195 def first_nonzero_index(mask):
196 raise NotImplementedError
197
198
199 def cumulative_max(x):
200 raise NotImplementedError
201
202
203 def bucketize_reward(rewards, boundaries):
204 raise NotImplementedError
205
List reforacle functions
python3 -c "
import reforacle
fns = [x for x in dir(reforacle) if not x.startswith('_')]
print(len(fns))
print(fns)
"51 ['F', 'advantage_mean_std', 'argmax_tokens', 'bradley_terry_logit', 'bucketize_reward', 'clip_fraction', 'clipped_pg_loss', 'cross_entropy', 'cumulative_max', 'discounted_returns', 'dpo_loss', 'dpo_sequence_loss', 'entropy', 'first_nonzero_index', 'gae', 'group_mean_baseline', 'grpo_advantages', 'grpo_objective', 'huber_value_loss', 'importance_ratio', 'ipo_loss', 'kl_penalty', 'lambda_returns', 'log_softmax', 'logprob_at_temperature', 'logsumexp', 'masked_max', 'masked_mean', 'masked_sum', 'masked_whiten', 'median_reward', 'mode_label', 'normalize', 'normalized_entropy', 'pad_mask_from_lengths', 'ppo_objective', 'quantile_lower', 'reverse_kl', 'reverse_kl_objective', 'rloo_advantages', 'rloo_objective', 'selective_logprobs', 'sequence_logprob', 'smoothed_nll', 'symmetric_kl', 'token_logprobs', 'top_k_mask', 'top_p_mask', 'torch', 'value_loss', 'whiten'] [stdout] 51 ['F', 'advantage_mean_std', 'argmax_tokens', 'bradley_terry_logit', 'bucketize_reward', 'clip_fraction', 'clipped_pg_loss', 'cross_entropy', 'cumulative_max', 'discounted_returns', 'dpo_loss', 'dpo_sequence_loss', 'entropy', 'first_nonzero_index', 'gae', 'group_mean_baseline', 'grpo_advantages', 'grpo_objective', 'huber_value_loss', 'importance_ratio', 'ipo_loss', 'kl_penalty', 'lambda_returns', 'log_softmax', 'logprob_at_temperature', 'logsumexp', 'masked_max', 'masked_mean', 'masked_sum', 'masked_whiten', 'median_reward', 'mode_label', 'normalize', 'normalized_entropy', 'pad_mask_from_lengths', 'ppo_objective', 'quantile_lower', 'reverse_kl', 'reverse_kl_objective', 'rloo_advantages', 'rloo_objective', 'selective_logprobs', 'sequence_logprob', 'smoothed_nll', 'symmetric_kl', 'token_logprobs', 'top_k_mask', 'top_p_mask', 'torch', 'value_loss', 'whiten']
Probe primitives batch 1
python3 -c "
import torch, reforacle
g = torch.Generator().manual_seed(1)
def show(name, *a, **k):
try:
r = getattr(reforacle, name)(*a, **k)
print(name, '->', r)
except Exception as e:
print(name, 'ERR', type(e).__name__, e)
v = torch.randn(6, generator=g)
m = (torch.rand(6, generator=g)>0.4).float()
print('v', v.tolist())
print('m', m.tolist())
show('masked_mean', v, m)
show('masked_sum', v, m)
print('manual masked_mean', (v*m).sum()/m.sum())
show('masked_mean', v, m, 0)
V = torch.randn(3,4, generator=g)
M = (torch.rand(3,4, generator=g)>0.4).float()
show('masked_mean', V, M, 1)
show('masked_mean', V, M, 0)
show('masked_mean', V, M)
show('masked_sum', V, M, 1)
x = torch.randn(3,4, generator=g)
show('logsumexp', x, 1)
show('logsumexp', x, 0)
print('manual lse', torch.logsumexp(x,1))
show('log_softmax', x, 1)
print('manual ls', torch.log_softmax(x,1))
show('normalize', v, 1e-8)
print('manual norm', (v-v.mean())/(v.std()+1e-8), 'unbiased std', v.std())
"v [0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188]
m [1.0, 0.0, 1.0, 1.0, 1.0, 1.0]
masked_mean -> tensor(0.1453)
masked_sum -> tensor(0.7263)
manual masked_mean tensor(0.1453)
masked_mean -> tensor(0.1453)
masked_mean -> tensor([ 0.0757, 0.6250, -0.5404])
masked_mean -> tensor([-1.1948, 0.1564, -0.5404, 1.1795])
masked_mean -> tensor(0.1561)
masked_sum -> tensor([ 0.2271, 1.2499, -0.5404])
logsumexp -> tensor([2.1861, 2.5172, 2.7495])
logsumexp -> tensor([2.2450, 2.4930, 2.3828, 1.5168])
manual lse tensor([2.1861, 2.5172, 2.7495])
log_softmax -> tensor([[-0.3648, -2.3676, -3.1376, -1.7804],
[-4.0336, -1.7850, -0.2352, -3.7252],
[-1.6375, -0.5321, -3.1764, -1.7349]])
manual ls tensor([[-0.3648, -2.3676, -3.1376, -1.7804],
[-4.0336, -1.7850, -0.2352, -3.7252],
[-1.6375, -0.5321, -3.1764, -1.7349]])
normalize -> tensor([ 1.1281, 0.2307, -0.2363, 1.0370, -1.4048, -0.7546])
manual norm tensor([ 1.1281, 0.2307, -0.2363, 1.0370, -1.4048, -0.7546]) unbiased std tensor(0.4395)
[stdout]
v [0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188]
m [1.0, 0.0, 1.0, 1.0, 1.0, 1.0]
masked_mean -> tensor(0.1453)
masked_sum -> tensor(0.7263)
manual masked_mean tensor(0.1453)
masked_mean -> tensor(0.1453)
masked_mean -> tensor([ 0.0757, 0.6250, -0.5404])
masked_mean -> tensor([-1.1948, 0.1564, -0.5404, 1.1795])
masked_mean -> tensor(0.1561)
masked_sum -> tensor([ 0.2271, 1.2499, -0.5404])
logsumexp -> tensor([2.1861, 2.5172, 2.7495])
logsumexp -> tensor([2.2450, 2.4930, 2.3828, 1.5168])
manual lse tensor([2.1861, 2.5172, 2.7495])
log_softmax -> tensor([[-0.3648, -2.3676, -3.1376, -1.7804],
[-4.0336, -1.7850, -0.2352, -3.7252],
[-1.6375, -0.5321, -3.1764, -1.7349]])
manual ls tensor([[-0.3648, -2.3676, -3.1376, -1.7804],
[-4.0336, -1.7850, -0.2352, -3.7252],
[-1.6375, -0.5321, -3.1764, -1.7349]])
normalize -> tensor([ 1.1281, 0.2307, -0.2363, 1.0370, -1.4048, -0.7546])
manual norm tensor([ 1.1281, 0.2307, -0.2363, 1.0370, -1.4048, -0.7546]) unbiased std tensor(0.4395)Probe logprob functions
python3 -c "
import torch, reforacle
g = torch.Generator().manual_seed(2)
def show(name, *a, **k):
try:
r = getattr(reforacle, name)(*a, **k)
print(name, '->', r if not torch.is_tensor(r) else r.tolist())
except Exception as e:
print(name, 'ERR', type(e).__name__, e)
B,T,V = 2,3,5
logits = torch.randn(B,T,V, generator=g)
labels = torch.randint(0,V,(B,T), generator=g)
mask = (torch.rand(B,T, generator=g)>0.3).float()
print('labels', labels.tolist())
print('mask', mask.tolist())
show('token_logprobs', logits, labels)
lp = torch.log_softmax(logits,-1)
print('manual token_logprobs', lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1).tolist())
show('selective_logprobs', logits, labels, mask)
show('sequence_logprob', logits, labels, mask, True)
show('sequence_logprob', logits, labels, mask, False)
show('entropy', logits, mask)
p = torch.softmax(logits,-1)
ent = -(p*torch.log_softmax(logits,-1)).sum(-1)
print('manual entropy per tok', ent.tolist())
print('manual masked mean entropy', ((ent*mask).sum()/mask.sum()).item())
show('cross_entropy', logits, labels, -100)
show('normalized_entropy', logits, mask)
show('smoothed_nll', logits, labels, 0.1)
show('logprob_at_temperature', logits, labels, 2.0)
"labels [[1, 2, 0], [3, 2, 2]] mask [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]] token_logprobs -> [[-0.40531259775161743, -1.953858494758606, -2.2495269775390625], [-2.5100996494293213, -1.8178317546844482, -1.8888953924179077]] manual token_logprobs [[-0.40531259775161743, -1.953858494758606, -2.2495269775390625], [-2.5100996494293213, -1.8178317546844482, -1.8888953924179077]] selective_logprobs -> [-4.608697891235352, -6.216826915740967] sequence_logprob -> [-1.536232590675354, -2.0722756385803223] sequence_logprob -> [-4.608697891235352, -6.216826915740967] entropy -> 1.2684575319290161 manual entropy per tok [[1.0967216491699219, 1.2921124696731567, 1.3587406873703003], [1.0793384313583374, 1.3906428813934326, 1.3931888341903687]] manual masked mean entropy 1.2684575319290161 cross_entropy -> 1.8042540550231934 normalized_entropy -> 0.7881369590759277 smoothed_nll -> 1.8180551528930664 logprob_at_temperature -> [[-0.8804996609687805, -1.6980383396148682, -1.86897611618042], [-1.9258971214294434, -1.652271032333374, -1.6975467205047607]] [stdout] labels [[1, 2, 0], [3, 2, 2]] mask [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]] token_logprobs -> [[-0.40531259775161743, -1.953858494758606, -2.2495269775390625], [-2.5100996494293213, -1.8178317546844482, -1.8888953924179077]] manual token_logprobs [[-0.40531259775161743, -1.953858494758606, -2.2495269775390625], [-2.5100996494293213, -1.8178317546844482, -1.8888953924179077]] selective_logprobs -> [-4.608697891235352, -6.216826915740967] sequence_logprob -> [-1.536232590675354, -2.0722756385803223] sequence_logprob -> [-4.608697891235352, -6.216826915740967] entropy -> 1.2684575319290161 manual entropy per tok [[1.0967216491699219, 1.2921124696731567, 1.3587406873703003], [1.0793384313583374, 1.3906428813934326, 1.3931888341903687]] manual masked mean entropy 1.2684575319290161 cross_entropy -> 1.8042540550231934 normalized_entropy -> 0.7881369590759277 smoothed_nll -> 1.8180551528930664 logprob_at_temperature -> [[-0.8804996609687805, -1.6980383396148682, -1.86897611618042], [-1.9258971214294434, -1.652271032333374, -1.6975467205047607]]
Verify seq/smoothed/temp/ce
python3 -c "
import torch, reforacle
g = torch.Generator().manual_seed(5)
def show(name, *a, **k):
try:
r = getattr(reforacle, name)(*a, **k)
print(name, '->', r if not torch.is_tensor(r) else r.tolist())
except Exception as e:
print(name, 'ERR', type(e).__name__, e)
B,T,V=2,4,5
logits=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=torch.tensor([[1.,1.,0.,0.],[1.,1.,1.,0.]])
lp=torch.log_softmax(logits,-1)
tok=lp.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
show('selective_logprobs',logits,labels,mask)
print('manual sel sum', (tok*mask).sum(-1).tolist())
show('sequence_logprob',logits,labels,mask,True)
print('manual seqnorm', ((tok*mask).sum(-1)/mask.sum(-1)).tolist())
# smoothed_nll
show('smoothed_nll',logits,labels,0.2)
eps=0.2
nll=-tok
smooth=-lp.mean(-1)
loss=((1-eps)*nll+eps*smooth)
print('manual smoothed mean', loss.mean().item())
# logprob_at_temp
show('logprob_at_temperature',logits,labels,2.0)
lpt=torch.log_softmax(logits/2.0,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
print('manual lpt', lpt.tolist())
# cross_entropy with ignore_index
labels2=labels.clone(); labels2[0,0]=-100
show('cross_entropy',logits,labels2,-100)
import torch.nn.functional as Fn
print('manual ce', Fn.cross_entropy(logits.reshape(-1,V),labels2.reshape(-1),ignore_index=-100).item())
"selective_logprobs -> [-7.382508754730225, -6.875703811645508] manual sel sum [-7.382508754730225, -6.875703811645508] sequence_logprob -> [-3.6912543773651123, -2.2919013500213623] manual seqnorm [-3.6912543773651123, -2.2919013500213623] smoothed_nll -> 2.6969223022460938 manual smoothed mean 2.6969223022460938 logprob_at_temperature -> [[-2.8581464290618896, -2.223844528198242, -1.84665846824646, -2.5614120960235596], [-1.946840763092041, -2.1040008068084717, -1.6675827503204346, -1.7135968208312988]] manual lpt [[-2.8581464290618896, -2.223844528198242, -1.84665846824646, -2.5614120960235596], [-1.946840763092041, -2.1040008068084717, -1.6675827503204346, -1.7135968208312988]] cross_entropy -> 2.609217405319214 manual ce 2.6092171669006348 [stdout] selective_logprobs -> [-7.382508754730225, -6.875703811645508] manual sel sum [-7.382508754730225, -6.875703811645508] sequence_logprob -> [-3.6912543773651123, -2.2919013500213623] manual seqnorm [-3.6912543773651123, -2.2919013500213623] smoothed_nll -> 2.6969223022460938 manual smoothed mean 2.6969223022460938 logprob_at_temperature -> [[-2.8581464290618896, -2.223844528198242, -1.84665846824646, -2.5614120960235596], [-1.946840763092041, -2.1040008068084717, -1.6675827503204346, -1.7135968208312988]] manual lpt [[-2.8581464290618896, -2.223844528198242, -1.84665846824646, -2.5614120960235596], [-1.946840763092041, -2.1040008068084717, -1.6675827503204346, -1.7135968208312988]] cross_entropy -> 2.609217405319214 manual ce 2.6092171669006348
Probe dpo/ipo/bt
python3 -c "
import torch, reforacle
import torch.nn.functional as Fn
g = torch.Generator().manual_seed(7)
def show(name, *a, **k):
try:
r = getattr(reforacle, name)(*a, **k)
print(name, '->', r if not torch.is_tensor(r) else r.tolist())
except Exception as e:
print(name, 'ERR', type(e).__name__, e)
pc=torch.tensor(-2.0); pr=torch.tensor(-3.0); rc=torch.tensor(-2.5); rr=torch.tensor(-2.8)
show('dpo_loss',pc,pr,rc,rr,0.1,0.0)
beta=0.1
logits=beta*((pc-rc)-(pr-rr))
print('manual dpo', (-Fn.logsigmoid(logits)).item())
show('dpo_loss',pc,pr,rc,rr,0.1,0.1)
ls=0.1
print('manual dpo ls', (-(1-ls)*Fn.logsigmoid(logits)-ls*Fn.logsigmoid(-logits)).item())
show('ipo_loss',pc,pr,rc,rr,0.1)
print('manual ipo', ((logits/beta - 1/(2*beta))**2).item() if False else ((( (pc-rc)-(pr-rr) )-1/(2*beta))**2).item())
show('bradley_terry_logit',torch.tensor(1.0),torch.tensor(0.3),0.5)
# with tensor batches
pc=torch.randn(4,generator=g);pr=torch.randn(4,generator=g);rc=torch.randn(4,generator=g);rr=torch.randn(4,generator=g)
show('dpo_loss',pc,pr,rc,rr,0.2,0.0)
lg=0.2*((pc-rc)-(pr-rr))
print('manual batch dpo', (-Fn.logsigmoid(lg)).mean().item())
show('ipo_loss',pc,pr,rc,rr,0.2)
print('manual batch ipo', ((((pc-rc)-(pr-rr))-1/(2*0.2))**2).mean().item())
"dpo_loss -> 0.6587595343589783 manual dpo 0.6587595343589783 dpo_loss -> 0.6657595634460449 manual dpo ls 0.6657595634460449 ipo_loss -> 18.490001678466797 manual ipo 18.490001678466797 bradley_terry_logit -> 0.3499999940395355 dpo_loss -> 0.6465333104133606 manual batch dpo 0.6465333104133606 ipo_loss -> 8.312424659729004 manual batch ipo 8.312424659729004 [stdout] dpo_loss -> 0.6587595343589783 manual dpo 0.6587595343589783 dpo_loss -> 0.6657595634460449 manual dpo ls 0.6657595634460449 ipo_loss -> 18.490001678466797 manual ipo 18.490001678466797 bradley_terry_logit -> 0.3499999940395355 dpo_loss -> 0.6465333104133606 manual batch dpo 0.6465333104133606 ipo_loss -> 8.312424659729004 manual batch ipo 8.312424659729004
Probe grpo/rloo/baseline
python3 -c "
import torch, reforacle
g = torch.Generator().manual_seed(9)
def show(name, *a, **k):
try:
r = getattr(reforacle, name)(*a, **k)
print(name, '->', r if not torch.is_tensor(r) else r.tolist())
except Exception as e:
print(name, 'ERR', type(e).__name__, e)
rewards=torch.tensor([1.0,2.0,3.0,4.0,5.0,6.0])
show('grpo_advantages',rewards,3,True)
show('grpo_advantages',rewards,3,False)
show('rloo_advantages',rewards,3)
show('group_mean_baseline',rewards,3)
# grpo manual
r=rewards.view(2,3)
mean=r.mean(1,keepdim=True); std=r.std(1,keepdim=True)
print('grpo scale manual', ((r-mean)/(std+1e-4)).flatten().tolist())
print('grpo nostd manual', ((r-mean)).flatten().tolist())
# rloo: leave one out mean
n=3
loo=(r.sum(1,keepdim=True)-r)/(n-1)
print('rloo manual', (r-loo).flatten().tolist())
print('gmb manual', mean.expand_as(r).flatten().tolist())
"grpo_advantages -> [-1.224743366241455, 0.0, 1.224743366241455, -1.224743366241455, 0.0, 1.224743366241455] grpo_advantages -> [-1.0, 0.0, 1.0, -1.0, 0.0, 1.0] rloo_advantages -> [-1.5, 0.0, 1.5, -1.5, 0.0, 1.5] group_mean_baseline -> [-1.0, 0.0, 1.0, -1.0, 0.0, 1.0] grpo scale manual [-0.9998999834060669, 0.0, 0.9998999834060669, -0.9998999834060669, 0.0, 0.9998999834060669] grpo nostd manual [-1.0, 0.0, 1.0, -1.0, 0.0, 1.0] rloo manual [-1.5, 0.0, 1.5, -1.5, 0.0, 1.5] gmb manual [2.0, 2.0, 2.0, 5.0, 5.0, 5.0] [stdout] grpo_advantages -> [-1.224743366241455, 0.0, 1.224743366241455, -1.224743366241455, 0.0, 1.224743366241455] grpo_advantages -> [-1.0, 0.0, 1.0, -1.0, 0.0, 1.0] rloo_advantages -> [-1.5, 0.0, 1.5, -1.5, 0.0, 1.5] group_mean_baseline -> [-1.0, 0.0, 1.0, -1.0, 0.0, 1.0] grpo scale manual [-0.9998999834060669, 0.0, 0.9998999834060669, -0.9998999834060669, 0.0, 0.9998999834060669] grpo nostd manual [-1.0, 0.0, 1.0, -1.0, 0.0, 1.0] rloo manual [-1.5, 0.0, 1.5, -1.5, 0.0, 1.5] gmb manual [2.0, 2.0, 2.0, 5.0, 5.0, 5.0]
Nail grpo std eps
python3 -c "
import torch, reforacle
r=torch.tensor([0.0,0.0,0.0,3.0])
# group_size 4 std=0 case
print('grpo zerostd', reforacle.grpo_advantages(r,4,True).tolist())
r2=torch.tensor([1.,2.,3.])
mean=r2.mean();
for unbiased in [True,False]:
std=r2.std(unbiased=unbiased)
print('unbiased',unbiased,'std',std.item(), 'val', ((r2-mean)/std).tolist())
# reforacle
print('ref', reforacle.grpo_advantages(torch.tensor([1.,2.,3.]),3,True).tolist())
# try eps values
import itertools
target=1.224743366241455
std=torch.tensor([1.,2.,3.]).std(unbiased=False).item()
for eps in [0,1e-8,1e-6,1e-5,1e-4]:
print(eps, 1/(std+eps))
"grpo zerostd [-0.577349841594696, -0.577349841594696, -0.577349841594696, 1.732049584388733] unbiased True std 1.0 val [-1.0, 0.0, 1.0] unbiased False std 0.8164966106414795 val [-1.2247447967529297, 0.0, 1.2247447967529297] ref [-1.224743366241455, 0.0, 1.224743366241455] 0 1.2247448268209604 1e-08 1.2247448118209616 1e-06 1.2247433268229067 1e-05 1.2247298270057618 0.0001 1.2245948452007993 [stdout] grpo zerostd [-0.577349841594696, -0.577349841594696, -0.577349841594696, 1.732049584388733] unbiased True std 1.0 val [-1.0, 0.0, 1.0] unbiased False std 0.8164966106414795 val [-1.2247447967529297, 0.0, 1.2247447967529297] ref [-1.224743366241455, 0.0, 1.224743366241455] 0 1.2247448268209604 1e-08 1.2247448118209616 1e-06 1.2247433268229067 1e-05 1.2247298270057618 0.0001 1.2245948452007993
Probe returns/gae
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(11)
def show(name, *a, **k):
try:
r = getattr(reforacle, name)(*a, **k)
print(name, '->', r if not torch.is_tensor(r) else r.tolist())
except Exception as e:
print(name, 'ERR', type(e).__name__, e)
rewards=torch.tensor([1.0,0.5,-0.5,2.0])
values=torch.tensor([0.1,0.2,0.3,0.4])
show('discounted_returns',rewards,0.9)
# manual
gamma=0.9; ret=[]; acc=0
for x in reversed(rewards.tolist()):
acc=x+gamma*acc; ret.append(acc)
print('manual dr', ret[::-1])
show('gae',rewards,values,torch.tensor(0.5),0.9,0.95)
show('lambda_returns',rewards,values,torch.tensor(0.5),0.9,0.95)
# gae manual
gamma,lam=0.9,0.95; nv=0.5
vals=values.tolist()+[nv]
adv=[];a=0
for t in reversed(range(4)):
delta=rewards[t].item()+gamma*vals[t+1]-vals[t]
a=delta+gamma*lam*a; adv.append(a)
adv=adv[::-1]
print('manual gae', adv)
print('manual lambda ret (gae+val)', [adv[i]+vals[i] for i in range(4)])
"discounted_returns -> [2.502999782562256, 1.6699999570846558, 1.2999999523162842, 2.0] manual dr [2.503, 1.6700000000000002, 1.3, 2.0] gae -> [2.5270028114318848, 1.6924011707305908, 1.312749981880188, 2.049999952316284] lambda_returns -> [2.627002716064453, 1.8924012184143066, 1.6127500534057617, 2.450000047683716] manual gae [2.527003068048728, 1.6924012477855386, 1.3127499883472922, 2.0499999940395357] manual lambda ret (gae+val) [2.627003069538844, 1.8924012507657708, 1.6127500002682211, 2.45] [stdout] discounted_returns -> [2.502999782562256, 1.6699999570846558, 1.2999999523162842, 2.0] manual dr [2.503, 1.6700000000000002, 1.3, 2.0] gae -> [2.5270028114318848, 1.6924011707305908, 1.312749981880188, 2.049999952316284] lambda_returns -> [2.627002716064453, 1.8924012184143066, 1.6127500534057617, 2.450000047683716] manual gae [2.527003068048728, 1.6924012477855386, 1.3127499883472922, 2.0499999940395357] manual lambda ret (gae+val) [2.627003069538844, 1.8924012507657708, 1.6127500002682211, 2.45]
Probe kl_penalty variants
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(13)
def show(name, *a, **k):
try:
r = getattr(reforacle, name)(*a, **k)
print(name, '->', r if not torch.is_tensor(r) else r.tolist())
except Exception as e:
print(name, 'ERR', type(e).__name__, e)
logp=torch.tensor([-1.0,-2.0,-0.5])
ref=torch.tensor([-1.2,-1.8,-0.7])
for est in ['k1','k2','k3','kl','mse','abs']:
show('kl_penalty',logp,ref,est)
d=logp-ref
print('k1 (logp-ref)', d.tolist())
print('k2 0.5 d^2', (0.5*d*d).tolist())
print('k3 exp(-d)-1+d... ref-logp based', (torch.exp(ref-logp)-(ref-logp)-1).tolist())
show('reverse_kl',logp,ref)
show('symmetric_kl',logp,ref)
show('importance_ratio',logp,ref,0.2)
print('imp exp(logp-ref) clipped', torch.exp(d).clamp(1-0.2,1+0.2).tolist(), 'unclipped', torch.exp(d).tolist())
"kl_penalty -> [0.20000004768371582, -0.20000004768371582, 0.19999998807907104] kl_penalty -> [0.02000000886619091, 0.02000000886619091, 0.019999997690320015] kl_penalty -> [0.018730759620666504, 0.021402835845947266, 0.018730759620666504] kl_penalty ERR ValueError kl kl_penalty ERR ValueError mse kl_penalty ERR ValueError abs k1 (logp-ref) [0.20000004768371582, -0.20000004768371582, 0.19999998807907104] k2 0.5 d^2 [0.02000000886619091, 0.02000000886619091, 0.019999997690320015] k3 exp(-d)-1+d... ref-logp based [0.018730759620666504, 0.021402835845947266, 0.018730759620666504] reverse_kl -> [0.018730759620666504, 0.021402835845947266, 0.018730759620666504] symmetric_kl -> [0.020066797733306885, 0.020066797733306885, 0.020066797733306885] importance_ratio -> [1.2000000476837158, 0.8187307119369507, 1.2000000476837158] imp exp(logp-ref) clipped [1.2000000476837158, 0.8187307119369507, 1.2000000476837158] unclipped [1.221402883529663, 0.8187307119369507, 1.2214027643203735] [stdout] kl_penalty -> [0.20000004768371582, -0.20000004768371582, 0.19999998807907104] kl_penalty -> [0.02000000886619091, 0.02000000886619091, 0.019999997690320015] kl_penalty -> [0.018730759620666504, 0.021402835845947266, 0.018730759620666504] kl_penalty ERR ValueError kl kl_penalty ERR ValueError mse kl_penalty ERR ValueError abs k1 (logp-ref) [0.20000004768371582, -0.20000004768371582, 0.19999998807907104] k2 0.5 d^2 [0.02000000886619091, 0.02000000886619091, 0.019999997690320015] k3 exp(-d)-1+d... ref-logp based [0.018730759620666504, 0.021402835845947266, 0.018730759620666504] reverse_kl -> [0.018730759620666504, 0.021402835845947266, 0.018730759620666504] symmetric_kl -> [0.020066797733306885, 0.020066797733306885, 0.020066797733306885] importance_ratio -> [1.2000000476837158, 0.8187307119369507, 1.2000000476837158] imp exp(logp-ref) clipped [1.2000000476837158, 0.8187307119369507, 1.2000000476837158] unclipped [1.221402883529663, 0.8187307119369507, 1.2214027643203735]
Probe pg/value/whiten
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(17)
def show(name, *a, **k):
try:
r = getattr(reforacle, name)(*a, **k)
print(name, '->', r if not torch.is_tensor(r) else r.tolist())
except Exception as e:
print(name, 'ERR', type(e).__name__, e)
logp=torch.tensor([[-1.0,-2.0,-0.5],[-0.8,-1.5,-2.2]])
old=torch.tensor([[-1.2,-1.8,-0.7],[-0.9,-1.6,-2.0]])
adv=torch.tensor([[1.0,-0.5,2.0],[0.3,-1.0,0.5]])
mask=torch.tensor([[1.,1.,0.],[1.,1.,1.]])
show('clipped_pg_loss',logp,old,adv,mask,0.2,0.2)
ratio=torch.exp(logp-old)
u1=ratio*adv; u2=torch.clamp(ratio,0.8,1.2)*adv
pg=-torch.min(u1,u2)
print('manual masked mean pg', ((pg*mask).sum()/mask.sum()).item())
show('clip_fraction',logp,old,0.2)
clipped=(torch.abs(ratio-1.0)>0.2).float()
print('manual clipfrac mean', clipped.mean().item())
# value_loss
values=torch.tensor([0.5,1.0,-0.5,2.0]); oldv=torch.tensor([0.4,1.2,-0.3,1.8]); ret=torch.tensor([0.6,0.9,-0.6,2.2])
show('value_loss',values,oldv,ret,0.2)
vclip=oldv+torch.clamp(values-oldv,-0.2,0.2)
l1=(values-ret)**2; l2=(vclip-ret)**2
print('manual vl', (0.5*torch.max(l1,l2).mean()).item(), 'nomax', (0.5*l1.mean()).item())
show('huber_value_loss',values,ret,1.0)
import torch.nn.functional as Fn
print('manual huber', Fn.huber_loss(values,ret,delta=1.0).item(), 'smoothl1', Fn.smooth_l1_loss(values,ret).item())
show('whiten',torch.tensor([1.,2.,3.,4.,5.]),None,True)
"clipped_pg_loss -> -0.08527606725692749 manual masked mean pg -0.08527606725692749 clip_fraction -> 0.3333333432674408 manual clipfrac mean 0.3333333432674408 value_loss -> 0.008750003762543201 manual vl 0.008750003762543201 nomax 0.008750003762543201 huber_value_loss -> 0.008750003762543201 manual huber 0.008750003762543201 smoothl1 0.008750003762543201 whiten ERR AttributeError 'NoneType' object has no attribute 'to' [stdout] clipped_pg_loss -> -0.08527606725692749 manual masked mean pg -0.08527606725692749 clip_fraction -> 0.3333333432674408 manual clipfrac mean 0.3333333432674408 value_loss -> 0.008750003762543201 manual vl 0.008750003762543201 nomax 0.008750003762543201 huber_value_loss -> 0.008750003762543201 manual huber 0.008750003762543201 smoothl1 0.008750003762543201 whiten ERR AttributeError 'NoneType' object has no attribute 'to'
value_loss max + whiten
python3 -c "
import torch, reforacle
def show(name, *a, **k):
try:
r = getattr(reforacle, name)(*a, **k)
print(name, '->', r if not torch.is_tensor(r) else r.tolist())
except Exception as e:
print(name, 'ERR', type(e).__name__, e)
values=torch.tensor([5.0,1.0,-3.0,2.0]); oldv=torch.tensor([0.4,1.2,-0.3,1.8]); ret=torch.tensor([0.6,0.9,-0.6,2.2])
show('value_loss',values,oldv,ret,0.2)
vclip=oldv+torch.clamp(values-oldv,-0.2,0.2)
l1=(values-ret)**2; l2=(vclip-ret)**2
print('max',0.5*torch.max(l1,l2).mean().item(),'nomax',0.5*l1.mean().item())
# whiten variants
x=torch.tensor([1.,2.,3.,4.,5.])
m=torch.ones(5)
show('whiten',x,m,True)
show('whiten',x,m,False)
show('masked_whiten',x,m,True)
show('masked_whiten',x,m,False)
for ub in [True,False]:
mean=x.mean();std=x.std(unbiased=ub)
print('ub',ub,'shift',((x-mean)/std).tolist(),'noshift',(x/std).tolist())
"value_loss -> 3.1462502479553223 max 3.1462502479553223 nomax 3.1462502479553223 whiten -> [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608] whiten -> [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915] masked_whiten -> [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608] masked_whiten -> [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915] ub True shift [-1.2649110555648804, -0.6324555277824402, 0.0, 0.6324555277824402, 1.2649110555648804] noshift [0.6324555277824402, 1.2649110555648804, 1.8973665237426758, 2.5298221111297607, 3.1622776985168457] ub False shift [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608] noshift [0.7071067690849304, 1.4142135381698608, 2.1213204860687256, 2.8284270763397217, 3.535533905029297] [stdout] value_loss -> 3.1462502479553223 max 3.1462502479553223 nomax 3.1462502479553223 whiten -> [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608] whiten -> [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915] masked_whiten -> [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608] masked_whiten -> [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915] ub True shift [-1.2649110555648804, -0.6324555277824402, 0.0, 0.6324555277824402, 1.2649110555648804] noshift [0.6324555277824402, 1.2649110555648804, 1.8973665237426758, 2.5298221111297607, 3.1622776985168457] ub False shift [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608] noshift [0.7071067690849304, 1.4142135381698608, 2.1213204860687256, 2.8284270763397217, 3.535533905029297]
Whiten eps form + masked_whiten
python3 -c "
import torch, reforacle
def show(name, *a, **k):
r=getattr(reforacle,name)(*a,**k); print(name,'->',r.tolist() if torch.is_tensor(r) else r)
# tiny variance to expose eps form
x=torch.tensor([1.0,1.0,1.0,1.0002])
m=torch.ones(4)
show('whiten',x,m,True)
mean=x.mean(); var=x.var(unbiased=False)
print('rsqrt var+1e-8', ((x-mean)*torch.rsqrt(var+1e-8)).tolist())
print('div std+1e-8', ((x-mean)/(x.std(unbiased=False)+1e-8)).tolist())
print('div std no eps', ((x-mean)/x.std(unbiased=False)).tolist())
# masked_whiten with real mask
x2=torch.tensor([1.,2.,3.,4.,5.,100.])
mask=torch.tensor([1.,1.,1.,1.,1.,0.])
show('masked_whiten',x2,mask,True)
mv=x2[mask.bool()]
mean=mv.mean();var=mv.var(unbiased=False)
print('manual masked biased', ((x2-mean)*torch.rsqrt(var+1e-8)).tolist())
"whiten -> [-0.37845081090927124, -0.37845081090927124, -0.37845081090927124, 1.1335502862930298] rsqrt var+1e-8 [-0.37845081090927124, -0.37845081090927124, -0.37845081090927124, 1.1335502862930298] div std+1e-8 [-0.5779715776443481, -0.5779715776443481, -0.5779715776443481, 1.7311625480651855] div std no eps [-0.5780382752418518, -0.5780382752418518, -0.5780382752418518, 1.7313623428344727] masked_whiten -> [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608, 0.0] manual masked biased [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608, 68.58935546875] [stdout] whiten -> [-0.37845081090927124, -0.37845081090927124, -0.37845081090927124, 1.1335502862930298] rsqrt var+1e-8 [-0.37845081090927124, -0.37845081090927124, -0.37845081090927124, 1.1335502862930298] div std+1e-8 [-0.5779715776443481, -0.5779715776443481, -0.5779715776443481, 1.7311625480651855] div std no eps [-0.5780382752418518, -0.5780382752418518, -0.5780382752418518, 1.7313623428344727] masked_whiten -> [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608, 0.0] manual masked biased [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608, 68.58935546875]
whiten vs masked_whiten with mask
python3 -c "
import torch, reforacle
def show(name, *a, **k):
r=getattr(reforacle,name)(*a,**k); print(name,'->',r.tolist() if torch.is_tensor(r) else r)
x=torch.tensor([1.,2.,3.,4.,5.,100.])
mask=torch.tensor([1.,1.,1.,1.,1.,0.])
show('whiten',x,mask,True)
show('whiten',x,mask,False)
show('masked_whiten',x,mask,False)
mv=x[mask.bool()]; mean=mv.mean(); var=mv.var(unbiased=False)
w=(x-mean)*torch.rsqrt(var+1e-8)
print('w+mean then *mask', ((w+mean)*mask).tolist())
print('w+mean no mask', (w+mean).tolist())
print('w*mask', (w*mask).tolist())
"whiten -> [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608, 68.58935546875] whiten -> [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915, 71.58935546875] masked_whiten -> [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915, 3.0] w+mean then *mask [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915, 0.0] w+mean no mask [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915, 71.58935546875] w*mask [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608, 0.0] [stdout] whiten -> [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608, 68.58935546875] whiten -> [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915, 71.58935546875] masked_whiten -> [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915, 3.0] w+mean then *mask [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915, 0.0] w+mean no mask [1.5857864618301392, 2.292893171310425, 3.0, 3.707106828689575, 4.41421365737915, 71.58935546875] w*mask [-1.4142135381698608, -0.7071067690849304, 0.0, 0.7071067690849304, 1.4142135381698608, 0.0]
Probe utility functions
python3 -c "
import torch, reforacle
def show(name, *a, **k):
try:
r=getattr(reforacle,name)(*a,**k); print(name,'->',r.tolist() if torch.is_tensor(r) else r, r.dtype if torch.is_tensor(r) else type(r).__name__)
except Exception as e: print(name,'ERR',type(e).__name__,e)
g=torch.Generator().manual_seed(21)
logits=torch.randn(2,3,5,generator=g)
show('argmax_tokens',logits)
print('manual', logits.argmax(-1).tolist())
probs=torch.softmax(torch.randn(2,5,generator=g),-1)
show('top_p_mask',probs,0.7)
show('top_k_mask',torch.randn(2,5,generator=g),2)
labels=torch.tensor([1,2,2,3,2,1])
show('mode_label',labels)
rewards=torch.tensor([3.0,1.0,2.0,5.0,4.0])
show('median_reward',rewards)
print('torch median', rewards.median())
rewards2=torch.tensor([3.0,1.0,2.0,4.0])
show('median_reward',rewards2)
print('torch median even', rewards2.median(), 'mean of mid', torch.quantile(rewards2,0.5))
v=torch.tensor([1.,5.,3.,2.]); m=torch.tensor([1.,0.,1.,1.])
show('masked_max',v,m)
show('quantile_lower',torch.tensor([1.,2.,3.,4.,5.]),0.25)
print('torch quantile lower',torch.quantile(torch.tensor([1.,2.,3.,4.,5.]),0.25,interpolation='lower'), 'linear',torch.quantile(torch.tensor([1.,2.,3.,4.,5.]),0.25))
show('pad_mask_from_lengths',torch.tensor([2,3,1]),4)
show('first_nonzero_index',torch.tensor([[0.,0.,1.,0.],[1.,0.,0.,0.],[0.,0.,0.,0.]]))
show('cumulative_max',torch.tensor([1.,3.,2.,5.,4.]))
show('bucketize_reward',torch.tensor([0.5,1.5,2.5,3.5]),torch.tensor([1.0,2.0,3.0]))
"argmax_tokens -> [[3, 1, 2], [2, 1, 4]] torch.int64 manual [[3, 1, 2], [2, 1, 4]] top_p_mask -> [[0.25545188784599304, 0.37779921293258667, 0.3667488694190979, 0.0, 0.0], [0.0, 0.0, 0.17712858319282532, 0.0, 0.8228714466094971]] torch.float32 top_k_mask -> [[False, True, False, True, False], [False, True, False, False, True]] torch.bool mode_label -> 2 torch.int64 median_reward -> 3.0 torch.float32 torch median tensor(3.) median_reward -> 2.5 torch.float32 torch median even tensor(2.) mean of mid tensor(2.5000) masked_max -> 3.0 torch.float32 quantile_lower -> 2.0 torch.float32 torch quantile lower tensor(2.) linear tensor(2.) pad_mask_from_lengths -> [[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, 0.0]] torch.float32 first_nonzero_index -> [2, 0, -1] torch.int64 cumulative_max -> [1.0, 1.0, 3.0, 3.0, 5.0] torch.float32 bucketize_reward -> [0, 1, 2, 3] torch.int64 [stdout] argmax_tokens -> [[3, 1, 2], [2, 1, 4]] torch.int64 manual [[3, 1, 2], [2, 1, 4]] top_p_mask -> [[0.25545188784599304, 0.37779921293258667, 0.3667488694190979, 0.0, 0.0], [0.0, 0.0, 0.17712858319282532, 0.0, 0.8228714466094971]] torch.float32 top_k_mask -> [[False, True, False, True, False], [False, True, False, False, True]] torch.bool mode_label -> 2 torch.int64 median_reward -> 3.0 torch.float32 torch median tensor(3.) median_reward -> 2.5 torch.float32 torch median even tensor(2.) mean of mid tensor(2.5000) masked_max -> 3.0 torch.float32 quantile_lower -> 2.0 torch.float32 torch quantile lower tensor(2.) linear tensor(2.) pad_mask_from_lengths -> [[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, 0.0]] torch.float32 first_nonzero_index -> [2, 0, -1] torch.int64 cumulative_max -> [1.0, 1.0, 3.0, 3.0, 5.0] torch.float32 bucketize_reward -> [0, 1, 2, 3] torch.int64
Refine cummax/quantile/mode/topp
python3 -c "
import torch, reforacle
def show(name, *a, **k):
try:
r=getattr(reforacle,name)(*a,**k); print(name,'->',r.tolist() if torch.is_tensor(r) else r)
except Exception as e: print(name,'ERR',type(e).__name__,e)
show('cumulative_max',torch.tensor([2.,1.,4.,3.,5.,0.]))
print('inclusive cummax', torch.cummax(torch.tensor([2.,1.,4.,3.,5.,0.]),0).values.tolist())
# 2D cumulative_max
show('cumulative_max',torch.tensor([[2.,1.,4.],[0.,5.,3.]]))
# quantile_lower interpolation
x=torch.tensor([1.,2.,3.,4.])
for q in [0.3,0.5,0.7]:
show('quantile_lower',x,q)
print(' lower',torch.quantile(x,q,interpolation='lower').item(),'linear',torch.quantile(x,q).item())
# mode ties
show('mode_label',torch.tensor([1,1,2,2,3]))
print('torch mode', torch.mode(torch.tensor([1,1,2,2,3])).values.item())
# masked_max axis?
V=torch.tensor([[1.,5.,3.],[2.,0.,4.]]); M=torch.tensor([[1.,0.,1.],[1.,1.,0.]])
show('masked_max',V,M)
# top_p details
probs=torch.tensor([[0.1,0.4,0.3,0.15,0.05]])
show('top_p_mask',probs,0.7)
"cumulative_max -> [2.0, 2.0, 2.0, 4.0, 4.0, 5.0] inclusive cummax [2.0, 2.0, 4.0, 4.0, 5.0, 5.0] cumulative_max -> [[2.0, 1.0, 4.0], [2.0, 1.0, 4.0]] quantile_lower -> 1.0 lower 1.0 linear 1.9000000953674316 quantile_lower -> 2.0 lower 2.0 linear 2.5 quantile_lower -> 3.0 lower 3.0 linear 3.0999999046325684 mode_label -> 2 torch mode 1 masked_max -> 3.0 top_p_mask -> [[0.0, 0.5714285373687744, 0.4285714030265808, 0.0, 0.0]] [stdout] cumulative_max -> [2.0, 2.0, 2.0, 4.0, 4.0, 5.0] inclusive cummax [2.0, 2.0, 4.0, 4.0, 5.0, 5.0] cumulative_max -> [[2.0, 1.0, 4.0], [2.0, 1.0, 4.0]] quantile_lower -> 1.0 lower 1.0 linear 1.9000000953674316 quantile_lower -> 2.0 lower 2.0 linear 2.5 quantile_lower -> 3.0 lower 3.0 linear 3.0999999046325684 mode_label -> 2 torch mode 1 masked_max -> 3.0 top_p_mask -> [[0.0, 0.5714285373687744, 0.4285714030265808, 0.0, 0.0]]
mode ties + top_p/top_k edges
python3 -c "
import torch, reforacle
def show(name, *a, **k):
try:
r=getattr(reforacle,name)(*a,**k); print(name,'->',r.tolist() if torch.is_tensor(r) else r)
except Exception as e: print(name,'ERR',type(e).__name__,e)
# mode tie tests
show('mode_label',torch.tensor([3,3,1,1,2]))
show('mode_label',torch.tensor([5,5,1,1]))
show('mode_label',torch.tensor([1,1,5,5]))
show('mode_label',torch.tensor([2,2,2,7,7,7,1]))
# top_p edge: threshold behaviors
show('top_p_mask',torch.tensor([[0.5,0.3,0.15,0.05]]),0.5)
print(' if 0.5 alone (cumsum>=p includes first only)')
show('top_p_mask',torch.tensor([[0.5,0.3,0.15,0.05]]),0.4)
show('top_p_mask',torch.tensor([[0.6,0.3,0.1]]),0.9)
show('top_p_mask',torch.tensor([[0.6,0.3,0.1]]),0.95)
# top_k with ties in value
show('top_k_mask',torch.tensor([[1.,3.,3.,2.]]),2)
show('top_k_mask',torch.tensor([[5.,1.,2.,3.,4.]]),3)
"mode_label -> 3 mode_label -> 5 mode_label -> 5 mode_label -> 7 top_p_mask -> [[1.0, 0.0, 0.0, 0.0]] if 0.5 alone (cumsum>=p includes first only) top_p_mask -> [[1.0, 0.0, 0.0, 0.0]] top_p_mask -> [[0.6666666865348816, 0.3333333432674408, 0.0]] top_p_mask -> [[0.6000000238418579, 0.30000001192092896, 0.10000000149011612]] top_k_mask -> [[False, True, True, False]] top_k_mask -> [[True, False, False, True, True]] [stdout] mode_label -> 3 mode_label -> 5 mode_label -> 5 mode_label -> 7 top_p_mask -> [[1.0, 0.0, 0.0, 0.0]] if 0.5 alone (cumsum>=p includes first only) top_p_mask -> [[1.0, 0.0, 0.0, 0.0]] top_p_mask -> [[0.6666666865348816, 0.3333333432674408, 0.0]] top_p_mask -> [[0.6000000238418579, 0.30000001192092896, 0.10000000149011612]] top_k_mask -> [[False, True, True, False]] top_k_mask -> [[True, False, False, True, True]]
advantage_mean_std
python3 -c "
import torch, reforacle
def show(name, *a, **k):
try:
r=getattr(reforacle,name)(*a,**k)
if isinstance(r,tuple): print(name,'->',[x.tolist() if torch.is_tensor(x) else x for x in r])
else: print(name,'->',r.tolist() if torch.is_tensor(r) else r)
except Exception as e: print(name,'ERR',type(e).__name__,e)
adv=torch.tensor([1.,2.,3.,4.,5.,6.]); mask=torch.tensor([1.,1.,1.,1.,0.,0.])
show('advantage_mean_std',adv,mask)
v=adv[mask.bool()]
print('mean',v.mean().item(),'std unb',v.std().item(),'std biased',v.std(unbiased=False).item())
# masked mean/var manual
mm=(adv*mask).sum()/mask.sum()
var=((adv-mm)**2*mask).sum()/mask.sum()
print('masked mean',mm.item(),'masked biased std',var.sqrt().item())
"advantage_mean_std -> [2.5, 1.1180340051651] mean 2.5 std unb 1.29099440574646 std biased 1.1180340051651 masked mean 2.5 masked biased std 1.1180340051651 [stdout] advantage_mean_std -> [2.5, 1.1180340051651] mean 2.5 std unb 1.29099440574646 std biased 1.1180340051651 masked mean 2.5 masked biased std 1.1180340051651
dpo_sequence_loss
python3 -c "
import torch, reforacle
import torch.nn.functional as Fn
g=torch.Generator().manual_seed(31)
B,T,V=2,4,6
pc_l=torch.randn(B,T,V,generator=g); pr_l=torch.randn(B,T,V,generator=g)
rc_l=torch.randn(B,T,V,generator=g); rr_l=torch.randn(B,T,V,generator=g)
cl=torch.randint(0,V,(B,T),generator=g); rl=torch.randint(0,V,(B,T),generator=g)
cm=(torch.rand(B,T,generator=g)>0.3).float(); rm=(torch.rand(B,T,generator=g)>0.3).float()
r=reforacle.dpo_sequence_loss(pc_l,pr_l,rc_l,rr_l,cl,rl,cm,rm,0.1,0.0)
print('ref dpo_seq', r.item() if torch.is_tensor(r) else r)
def seqlp(logits,labels,mask):
lp=torch.log_softmax(logits,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
return (lp*mask).sum(-1)
pc=seqlp(pc_l,cl,cm); pr=seqlp(pr_l,rl,rm); rc=seqlp(rc_l,cl,cm); rr=seqlp(rr_l,rl,rm)
logits=0.1*((pc-rc)-(pr-rr))
print('manual', (-Fn.logsigmoid(logits)).mean().item())
print('via dpo_loss', reforacle.dpo_loss(pc,pr,rc,rr,0.1,0.0).item())
"ref dpo_seq 0.5764719843864441 manual 0.5764719843864441 via dpo_loss 0.5764719843864441 [stdout] ref dpo_seq 0.5764719843864441 manual 0.5764719843864441 via dpo_loss 0.5764719843864441
Probe reverse_kl_objective shape
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
# advantages: try per-token and per-seq
adv_tok=torch.randn(B,T,generator=g)
adv_seq=torch.randn(B,generator=g)
for adv,tag in [(adv_tok,'tok'),(adv_seq,'seq')]:
try:
r=reforacle.reverse_kl_objective(logits,ref,labels,mask,adv,0.1)
print(tag,'->',r.item() if torch.is_tensor(r) else r)
except Exception as e: print(tag,'ERR',e)
print('mask',mask.tolist())
print('adv_tok',adv_tok.tolist())
"tok -> 0.4646722376346588 seq -> -0.363065242767334 mask [[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]] adv_tok [[0.8789268732070923, -1.556331753730774, 0.022696413099765778], [-0.8227455615997314, 0.9545934200286865, 0.49359533190727234]] [stdout] tok -> 0.4646722376346588 seq -> -0.363065242767334 mask [[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]] adv_tok [[0.8789268732070923, -1.556331753730774, 0.022696413099765778], [-0.8227455615997314, 0.9545934200286865, 0.49359533190727234]]
reverse_kl_objective formula
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.randn(B,T,generator=g)
target=0.4646722376346588
lp=torch.log_softmax(logits,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
rlp=torch.log_softmax(ref,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
d=lp-rlp
rkl=torch.exp(-d)+d-1 # reverse kl k3 token
def mm(x): return (x*mask).sum()/mask.sum()
beta=0.1
c1=mm(-(adv*lp)+beta*rkl)
c2=mm(-adv*lp)+beta*mm(rkl)
c3=mm(-(adv*lp - beta*rkl))
print('target',target)
print('c1 -adv*lp + beta*rkl (masked mean)', c1.item())
print('c2 same split', c2.item())
# maybe pg uses (logp)*adv negative and kl subtracts
c4=mm(-adv*lp)-beta*mm(rkl)
print('c4 minus', c4.item())
"target 0.4646722376346588 c1 -adv*lp + beta*rkl (masked mean) 0.4697510302066803 c2 same split 0.4697510600090027 c4 minus 0.11686335504055023 [stdout] target 0.4646722376346588 c1 -adv*lp + beta*rkl (masked mean) 0.4697510302066803 c2 same split 0.4697510600090027 c4 minus 0.11686335504055023
reverse_kl_objective full dist
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.randn(B,T,generator=g)
target=0.4646722376346588
lp=torch.log_softmax(logits,-1)
rlp=torch.log_softmax(ref,-1)
p=lp.exp(); rp=rlp.exp()
tlp=lp.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
def mm(x): return ((x*mask).sum()/mask.sum()).item()
beta=0.1
kl_pr=(p*(lp-rlp)).sum(-1) # KL(policy||ref)
kl_rp=(rp*(rlp-lp)).sum(-1) # KL(ref||policy)
# k3 full over ref: exp(logr-logp)... token-level with distribution?
for name,kl in [('KL(p||r)',kl_pr),('KL(r||p)',kl_rp)]:
print(name, mm(-adv*tlp)+beta*mm(kl), ' combined', mm(-adv*tlp+beta*kl))
# k3 estimator using ref-policy full: mean over vocab?
"KL(p||r) 0.42369096875190737 combined 0.4236909747123718 KL(r||p) 0.40738001465797424 combined 0.40738004446029663 [stdout] KL(p||r) 0.42369096875190737 combined 0.4236909747123718 KL(r||p) 0.40738001465797424 combined 0.40738004446029663
isolate reverse_kl_objective pieces
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.randn(B,T,generator=g)
target=0.4646722376346588
lp=torch.log_softmax(logits,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
rlp=torch.log_softmax(ref,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
d=lp-rlp
def mm(x): return ((x*mask).sum()/mask.sum()).item()
pg=mm(-adv*lp)
rkl_k3=mm(torch.exp(-d)+d-1)
rkl_fwd=mm(torch.exp(d)-d-1)
print('pg', pg, 'rkl_k3', rkl_k3, 'rkl_fwd', rkl_fwd)
print('target-pg', target-pg)
print('beta*rkl_k3', 0.1*rkl_k3, 'beta*rkl_fwd', 0.1*rkl_fwd)
# maybe pg uses +adv*lp? try
print('pg+', mm(adv*lp))
# maybe advantage-weighted with ratio exp(lp-rlp)?
ratio=torch.exp(d)
print('ratio pg', mm(-adv*ratio), mm(adv*ratio))
"pg 0.29330721497535706 rkl_k3 1.7644386291503906 rkl_fwd 1.653384804725647 target-pg 0.17136502265930176 beta*rkl_k3 0.17644386291503908 beta*rkl_fwd 0.16533848047256472 pg+ -0.29330721497535706 ratio pg -1.4805262088775635 1.4805262088775635 [stdout] pg 0.29330721497535706 rkl_k3 1.7644386291503906 rkl_fwd 1.653384804725647 target-pg 0.17136502265930176 beta*rkl_k3 0.17644386291503908 beta*rkl_fwd 0.16533848047256472 pg+ -0.29330721497535706 ratio pg -1.4805262088775635 1.4805262088775635
solve for K
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.randn(B,T,generator=g)
target=0.4646722376346588
lpf=torch.log_softmax(logits,-1); rlpf=torch.log_softmax(ref,-1)
lp=lpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
rlp=rlpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
def mm(x): return ((x*mask).sum()/mask.sum()).item()
pg=mm(-adv*lp)
K=(target-pg)/0.1
print('need K=',K)
p=lpf.exp(); rp=rlpf.exp()
d=lp-rlp
cands={
'k3 tok': torch.exp(-d)+d-1,
'fwd tok': torch.exp(d)-d-1,
'sym tok': 0.5*(torch.exp(-d)+d-1+torch.exp(d)-d-1),
'abs d': d.abs(),
'sq d': 0.5*d*d,
'KLpr': (p*(lpf-rlpf)).sum(-1),
'KLrp': (rp*(rlpf-lpf)).sum(-1),
'k3full_pr': (torch.exp(rlpf-lpf)-(rlpf-lpf)-1).sum(-1)*0, # placeholder
}
# k3 full over dist weighted by p: sum p*(exp(logr-logp)-(logr-logp)-1)? that's weird
for name,c in cands.items():
print(name, mm(c))
"need K= 1.7136502265930176 k3 tok 1.7644386291503906 fwd tok 1.653384804725647 sym tok 1.7089115381240845 abs d 1.50766921043396 sq d 1.2771250009536743 KLpr 1.303837537765503 KLrp 1.1407279968261719 k3full_pr 0.0 [stdout] need K= 1.7136502265930176 k3 tok 1.7644386291503906 fwd tok 1.653384804725647 sym tok 1.7089115381240845 abs d 1.50766921043396 sq d 1.2771250009536743 KLpr 1.303837537765503 KLrp 1.1407279968261719 k3full_pr 0.0
more reverse_kl_objective structures
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.randn(B,T,generator=g)
target=0.4646722376346588
lpf=torch.log_softmax(logits,-1); rlpf=torch.log_softmax(ref,-1)
lp=lpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
rlp=rlpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
def mm(x): return ((x*mask).sum()/mask.sum()).item()
d=lp-rlp
beta=0.1
rkl=torch.exp(-d)+d-1
# structure: A' = adv - beta*rkl then -mean(A'*lp)? no
print('A-beta rkl then *lp', mm(-(adv-beta*rkl)*lp))
print('A+beta rkl', mm(-(adv+beta*rkl)*lp))
# structure: -(adv*lp) + beta*rkl but rkl weighted differently...
# GRPO style: -(adv*ratio) + beta*rkl with ratio=exp(lp-lp.detach())=1 at eval -> adv
# but here maybe ratio = exp(lp - rlp)?
ratio=torch.exp(d)
print('grpo-like -adv*ratio+beta*rkl', mm(-adv*ratio+beta*rkl))
# maybe pg uses ratio*adv where ratio=exp(lp-old) but no old...
# try loss = mm(-adv*lp + beta*rkl) but with full KLpr using k3 per token computed from full dist expectation
p=lpf.exp()
klk3_full=(p*(torch.exp(rlpf-lpf)-(rlpf-lpf)-1)).sum(-1)
print('k3full weighted p', mm(-adv*lp+beta*klk3_full))
klk3_full_r=(rlpf.exp()*(torch.exp(lpf-rlpf)-(lpf-rlpf)-1)).sum(-1)
print('k3full weighted r', mm(-adv*lp+beta*klk3_full_r))
"A-beta rkl then *lp -0.16101884841918945 A+beta rkl 0.7476332187652588 grpo-like -adv*ratio+beta*rkl -1.3040823936462402 k3full weighted p 0.4236909747123718 k3full weighted r 0.40738004446029663 [stdout] A-beta rkl then *lp -0.16101884841918945 A+beta rkl 0.7476332187652588 grpo-like -adv*ratio+beta*rkl -1.3040823936462402 k3full weighted p 0.4236909747123718 k3full weighted r 0.40738004446029663
isolate beta=0 term and slope
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.randn(B,T,generator=g)
lpf=torch.log_softmax(logits,-1)
lp=lpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
def mm(x): return ((x*mask).sum()/mask.sum()).item()
print('beta=0 ->', reforacle.reverse_kl_objective(logits,ref,labels,mask,adv,0.0))
print('mm(-adv*lp)', mm(-adv*lp))
print('mm(adv*lp)', mm(adv*lp))
# vary beta to get slope
r0=reforacle.reverse_kl_objective(logits,ref,labels,mask,adv,0.0).item()
r1=reforacle.reverse_kl_objective(logits,ref,labels,mask,adv,1.0).item()
print('r0',r0,'r1',r1,'slope(K)=',r1-r0)
"beta=0 -> tensor(0.1118) mm(-adv*lp) 0.29330721497535706 mm(adv*lp) -0.29330721497535706 r0 0.11178465187549591 r1 3.6406617164611816 slope(K)= 3.5288770645856857 [stdout] beta=0 -> tensor(0.1118) mm(-adv*lp) 0.29330721497535706 mm(adv*lp) -0.29330721497535706 r0 0.11178465187549591 r1 3.6406617164611816 slope(K)= 3.5288770645856857
find beta=0 term
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.randn(B,T,generator=g)
lpf=torch.log_softmax(logits,-1); rlpf=torch.log_softmax(ref,-1)
lp=lpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
rlp=rlpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
def mm(x): return ((x*mask).sum()/mask.sum()).item()
d=lp-rlp
target0=0.11178465187549591
ratio=torch.exp(d)
print('-adv*d', mm(-adv*d))
print('-adv*ratio', mm(-adv*ratio))
print('-adv*(ratio-1)', mm(-adv*(ratio-1)))
print('adv*(ratio-1)', mm(adv*(ratio-1)))
print('-adv*lp*ratio', mm(-adv*ratio*lp))
# whitened adv
def wh(a,m):
mean=(a*m).sum()/m.sum(); var=((a-mean)**2*m).sum()/m.sum()
return (a-mean)/torch.sqrt(var+1e-8)
wadv=wh(adv,mask)
print('-wadv*lp', mm(-wadv*lp))
print('-wadv*ratio', mm(-wadv*ratio))
"-adv*d -0.17329931259155273 -adv*ratio -1.4805262088775635 -adv*(ratio-1) -1.1751129627227783 adv*(ratio-1) 1.1751129627227783 -adv*lp*ratio 2.045724630355835 -wadv*lp -0.7203050255775452 -wadv*ratio -1.072291374206543 [stdout] -adv*d -0.17329931259155273 -adv*ratio -1.4805262088775635 -adv*(ratio-1) -1.1751129627227783 adv*(ratio-1) 1.1751129627227783 -adv*lp*ratio 2.045724630355835 -wadv*lp -0.7203050255775452 -wadv*ratio -1.072291374206543
A dependence on adv
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.randn(B,T,generator=g)
lpf=torch.log_softmax(logits,-1); rlpf=torch.log_softmax(ref,-1)
lp=lpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
rlp=rlpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
def mm(x): return ((x*mask).sum()/mask.sum()).item()
def obj(a,b): return reforacle.reverse_kl_objective(logits,ref,labels,mask,a,b).item()
print('adv=0,beta=0', obj(torch.zeros_like(adv),0.0))
print('adv=0,beta=0.1', obj(torch.zeros_like(adv),0.1))
print('adv=adv,beta=0', obj(adv,0.0))
print('adv=2*adv,beta=0', obj(2*adv,0.0))
print('mm(-adv*lp)',mm(-adv*lp),'mm(-adv*rlp)',mm(-adv*rlp))
d=lp-rlp
# A with adv=0 is 0? then A linear in adv
"adv=0,beta=0 0.0 adv=0,beta=0.1 0.35288771986961365 adv=adv,beta=0 0.11178465187549591 adv=2*adv,beta=0 0.22356930375099182 mm(-adv*lp) 0.29330721497535706 mm(-adv*rlp) 0.466606467962265 [stdout] adv=0,beta=0 0.0 adv=0,beta=0.1 0.35288771986961365 adv=adv,beta=0 0.11178465187549591 adv=2*adv,beta=0 0.22356930375099182 mm(-adv*lp) 0.29330721497535706 mm(-adv*rlp) 0.466606467962265
fit A coefficient vector
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
lpf=torch.log_softmax(logits,-1); rlpf=torch.log_softmax(ref,-1)
lp=lpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
rlp=rlpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
def obj0(a): return reforacle.reverse_kl_objective(logits,ref,labels,mask,a,0.0).item()
# basis: coef_i such that A = sum_i coef_i * adv_i ; use one-hot
coef=torch.zeros(B,T)
for i in range(B):
for j in range(T):
e=torch.zeros(B,T); e[i,j]=1.0
coef[i,j]=obj0(e)
print('coef matrix (dA/dadv_ij):'); print(coef.tolist())
print('lp'); print(lp.tolist())
print('rlp'); print(rlp.tolist())
print('mask'); print(mask.tolist())
Nm=mask.sum().item()
print('-lp/Nm'); print((-lp/Nm).tolist())
print('-(lp-rlp)/Nm'); print((-(lp-rlp)/Nm).tolist())
"coef matrix (dA/dadv_ij): [[1.1860679388046265, 0.567639946937561, 0.7490347623825073], [1.1860679388046265, 0.567639946937561, 0.7490347623825073]] lp [[-1.7583439350128174, -3.69807767868042, -2.5575673580169678], [-4.171996116638184, -2.8381998538970947, -1.1876063346862793]] rlp [[-3.3097925186157227, -3.003464698791504, -1.5112764835357666], [-3.3843116760253906, -0.6483660340309143, -3.150693893432617]] mask [[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]] -lp/Nm [[0.3516687750816345, 0.7396155595779419, 0.5115134716033936], [0.8343992233276367, 0.567639946937561, 0.23752126097679138]] -(lp-rlp)/Nm [[-0.3102897107601166, 0.13892260193824768, 0.20925816893577576], [0.15753689408302307, 0.4379667639732361, -0.39261752367019653]] [stdout] coef matrix (dA/dadv_ij): [[1.1860679388046265, 0.567639946937561, 0.7490347623825073], [1.1860679388046265, 0.567639946937561, 0.7490347623825073]] lp [[-1.7583439350128174, -3.69807767868042, -2.5575673580169678], [-4.171996116638184, -2.8381998538970947, -1.1876063346862793]] rlp [[-3.3097925186157227, -3.003464698791504, -1.5112764835357666], [-3.3843116760253906, -0.6483660340309143, -3.150693893432617]] mask [[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]] -lp/Nm [[0.3516687750816345, 0.7396155595779419, 0.5115134716033936], [0.8343992233276367, 0.567639946937561, 0.23752126097679138]] -(lp-rlp)/Nm [[-0.3102897107601166, 0.13892260193824768, 0.20925816893577576], [0.15753689408302307, 0.4379667639732361, -0.39261752367019653]]
test 1D advantage shapes
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(41)
B,T,V=2,3,5
logits=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
lpf=torch.log_softmax(logits,-1)
lp=lpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
def O(a,b=0.0):
try: return reforacle.reverse_kl_objective(logits,ref,labels,mask,a,b).item()
except Exception as e: return 'ERR '+str(e)
# 1D adv length B and length T
print('adv (B,)', O(torch.tensor([1.0,2.0])))
print('adv (T,)', O(torch.tensor([1.0,2.0,3.0])))
print('adv scalar', O(torch.tensor(1.0)))
Nm=mask.sum().item()
# check adv (T,) coef prediction: sum_j adv[j]*colsum
colsum=(-lp*mask).sum(0)/Nm
print('colsum', colsum.tolist())
print('pred adv(T,)[1,2,3]', (torch.tensor([1.,2.,3.])*colsum).sum().item())
# check adv (B,)
rowsum=(-lp*mask).sum(1)/Nm
print('rowsum',rowsum.tolist(),'pred adv(B,)[1,2]', (torch.tensor([1.,2.])*rowsum).sum().item())
"adv (B,) 4.142303466796875 adv (T,) ERR The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 0 adv scalar ERR Dimension out of range (expected to be in range of [-1, 0], but got 1) colsum [1.1860679388046265, 0.567639946937561, 0.7490347623825073] pred adv(T,)[1,2,3] 4.568451881408691 rowsum [0.8631822466850281, 1.639560341835022] pred adv(B,)[1,2] 4.142302989959717 [stdout] adv (B,) 4.142303466796875 adv (T,) ERR The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 0 adv scalar ERR Dimension out of range (expected to be in range of [-1, 0], but got 1) colsum [1.1860679388046265, 0.567639946937561, 0.7490347623825073] pred adv(T,)[1,2,3] 4.568451881408691 rowsum [0.8631822466850281, 1.639560341835022] pred adv(B,)[1,2] 4.142302989959717
verify reverse_kl_objective formula fresh seeds
python3 -c "
import torch, reforacle
def test(seed):
g=torch.Generator().manual_seed(seed)
B,T,V=3,4,6
logits=torch.randn(B,T,V,generator=g); ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.randn(B,generator=g)
beta=0.15
lpf=torch.log_softmax(logits,-1); rlpf=torch.log_softmax(ref,-1)
lp=lpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
rlp=rlpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
d=lp-rlp
rkl=torch.exp(-d)+d-1
def mm(x): return (x*mask).sum()/mask.sum()
pg=mm(-adv.unsqueeze(-1)*lp)
guessA=pg+beta*2*mm(rkl)
guessB=pg+beta*mm(rkl)
ref_v=reforacle.reverse_kl_objective(logits,ref,labels,mask,adv,beta).item()
print(seed,'ref',ref_v,'2x',guessA.item(),'1x',guessB.item())
for s in [1,2,3]: test(s)
"1 ref 0.22150394320487976 2x 0.36478108167648315 1x 0.22150392830371857 2 ref -0.48739275336265564 2x -0.3948437571525574 1x -0.48739269375801086 3 ref -1.4024114608764648 2x -1.3145310878753662 1x -1.4024114608764648 [stdout] 1 ref 0.22150394320487976 2x 0.36478108167648315 1x 0.22150392830371857 2 ref -0.48739275336265564 2x -0.3948437571525574 1x -0.48739269375801086 3 ref -1.4024114608764648 2x -1.3145310878753662 1x -1.4024114608764648
confirm reverse KL in objective
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(7)
B,T,V=3,4,6
logits=torch.randn(B,T,V,generator=g); ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
adv=torch.zeros(B) # isolate KL
beta=1.0
lpf=torch.log_softmax(logits,-1); rlpf=torch.log_softmax(ref,-1)
lp=lpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
rlp=rlpf.gather(-1,labels.unsqueeze(-1)).squeeze(-1)
d=lp-rlp
def mm(x): return ((x*mask).sum()/mask.sum()).item()
print('ref KL-only', reforacle.reverse_kl_objective(logits,ref,labels,mask,adv,beta).item())
print('reverse k3', mm(torch.exp(-d)+d-1))
print('forward k3', mm(torch.exp(d)-d-1))
print('sym', mm(0.5*(torch.exp(-d)+d-1+torch.exp(d)-d-1)))
"ref KL-only 0.3913531005382538 reverse k3 0.3913531005382538 forward k3 0.4202366769313812 sym 0.4057949185371399 [stdout] ref KL-only 0.3913531005382538 reverse k3 0.3913531005382538 forward k3 0.4202366769313812 sym 0.4057949185371399
probe grpo_objective
python3 -c "
import torch, reforacle
def build(seed):
g=torch.Generator().manual_seed(seed)
G,T,V=4,3,5 # group_size divides batch
B=4
logits=torch.randn(B,T,V,generator=g)
old=torch.randn(B,T,V,generator=g)
ref=torch.randn(B,T,V,generator=g)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
rewards=torch.randn(B,generator=g)
return logits,old,ref,labels,mask,rewards
logits,old,ref,labels,mask,rewards=build(3)
gs=2; beta=0.05; cl=0.2; ch=0.2
r=reforacle.grpo_objective(logits,old,ref,labels,mask,rewards,gs,beta,cl,ch,True,'k3').item()
print('ref grpo', r)
def lpg(lg): return torch.log_softmax(lg,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
lp=lpg(logits); olp=lpg(old); rlp=lpg(ref)
adv=reforacle.grpo_advantages(rewards,gs,True)
ratio=torch.exp(lp-olp)
a=adv.unsqueeze(-1)
pg=-torch.min(ratio*a, torch.clamp(ratio,1-cl,1+ch)*a)
d=lp-rlp
kl=torch.exp(-d)+d-1
def mm(x): return (x*mask).sum()/mask.sum()
print('guess global mm', mm(pg+beta*kl).item())
# per-seq then mean
def perseq(x): return ((x*mask).sum(-1)/mask.sum(-1)).mean()
print('guess perseq', perseq(pg+beta*kl).item())
"ref grpo 0.54408198595047 guess global mm 0.54408198595047 guess perseq 0.4394015967845917 [stdout] ref grpo 0.54408198595047 guess global mm 0.54408198595047 guess perseq 0.4394015967845917
grpo_objective multi-config verify
python3 -c "
import torch, reforacle
def build(seed,B=6,T=4,V=5):
g=torch.Generator().manual_seed(seed)
return (torch.randn(B,T,V,generator=g),torch.randn(B,T,V,generator=g),torch.randn(B,T,V,generator=g),
torch.randint(0,V,(B,T),generator=g),(torch.rand(B,T,generator=g)>0.3).float(),torch.randn(B,generator=g))
def lpg(lg,labels): return torch.log_softmax(lg,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
def klf(d,est):
if est=='k1': return d
if est=='k2': return 0.5*d*d
return torch.exp(-d)+d-1
for seed,gs,beta,cl,ch,sbs,est in [(1,3,0.1,0.2,0.3,False,'k1'),(2,2,0.0,0.1,0.1,True,'k2'),(5,6,0.2,0.15,0.25,True,'k3')]:
logits,old,ref,labels,mask,rewards=build(seed)
lp=lpg(logits,labels);olp=lpg(old,labels);rlp=lpg(ref,labels)
adv=reforacle.grpo_advantages(rewards,gs,sbs).unsqueeze(-1)
ratio=torch.exp(lp-olp)
pg=-torch.min(ratio*adv,torch.clamp(ratio,1-cl,1+ch)*adv)
kl=klf(lp-rlp,est)
g=((pg+beta*kl)*mask).sum()/mask.sum()
r=reforacle.grpo_objective(logits,old,ref,labels,mask,rewards,gs,beta,cl,ch,sbs,est).item()
print(seed, 'ref',r,'guess',g.item(), 'OK' if abs(r-g.item())<1e-5 else 'FAIL')
"1 ref 0.692531943321228 guess 0.692531943321228 OK 2 ref 0.4025667905807495 guess 0.4025667905807495 OK 5 ref 0.7935142517089844 guess 0.7935142517089844 OK [stdout] 1 ref 0.692531943321228 guess 0.692531943321228 OK 2 ref 0.4025667905807495 guess 0.4025667905807495 OK 5 ref 0.7935142517089844 guess 0.7935142517089844 OK
probe rloo_objective
python3 -c "
import torch, reforacle
def build(seed,B=6,T=4,V=5):
g=torch.Generator().manual_seed(seed)
return (torch.randn(B,T,V,generator=g),torch.randn(B,T,V,generator=g),
torch.randint(0,V,(B,T),generator=g),(torch.rand(B,T,generator=g)>0.3).float(),torch.randn(B,generator=g))
def lpg(lg,labels): return torch.log_softmax(lg,-1).gather(-1,labels.unsqueeze(-1)).squeeze(-1)
for seed,gs,cl,ch in [(1,3,0.2,0.2),(2,2,0.1,0.3),(5,6,0.15,0.25)]:
logits,old,labels,mask,rewards=build(seed)
lp=lpg(logits,labels);olp=lpg(old,labels)
adv=reforacle.rloo_advantages(rewards,gs).unsqueeze(-1)
ratio=torch.exp(lp-olp)
pg=-torch.min(ratio*adv,torch.clamp(ratio,1-cl,1+ch)*adv)
g=((pg)*mask).sum()/mask.sum()
r=reforacle.rloo_objective(logits,old,labels,mask,rewards,gs,cl,ch).item()
print(seed,'ref',r,'guess',g.item(),'OK' if abs(r-g.item())<1e-5 else 'FAIL')
"1 ref 3.3130955696105957 guess 3.3130955696105957 OK 2 ref 0.5726490616798401 guess 0.5726490616798401 OK 5 ref 0.5638018846511841 guess 0.5638018846511841 OK [stdout] 1 ref 3.3130955696105957 guess 3.3130955696105957 OK 2 ref 0.5726490616798401 guess 0.5726490616798401 OK 5 ref 0.5638018846511841 guess 0.5638018846511841 OK
probe ppo_objective
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(3)
T=5
rewards=torch.randn(T,generator=g); values=torch.randn(T,generator=g); old_values=torch.randn(T,generator=g)
logp=torch.randn(T,generator=g); old_logp=torch.randn(T,generator=g); nv=torch.randn(1,generator=g)[0]
gamma,lam,cl,ch,vfc,vfco=0.99,0.95,0.2,0.2,0.2,0.5
r=reforacle.ppo_objective(rewards,values,old_values,logp,old_logp,nv,gamma,lam,cl,ch,vfc,vfco).item()
print('ref',r)
adv=reforacle.gae(rewards,values,nv,gamma,lam)
ret=reforacle.lambda_returns(rewards,values,nv,gamma,lam)
ratio=torch.exp(logp-old_logp)
def vl(values,old_values,returns,clip):
vc=old_values+torch.clamp(values-old_values,-clip,clip)
return 0.5*torch.max((values-returns)**2,(vc-returns)**2).mean()
# no whiten
pg=-torch.min(ratio*adv,torch.clamp(ratio,1-cl,1+ch)*adv)
g1=pg.mean()+vfco*vl(values,old_values,ret,vfc)
print('guess no-whiten',g1.item())
# whiten adv (unbiased? biased?)
def wh(a,ub):
return (a-a.mean())/(a.std(unbiased=ub)+1e-8)
for ub in [True,False]:
aw=wh(adv,ub)
pg2=-torch.min(ratio*aw,torch.clamp(ratio,1-cl,1+ch)*aw)
print('whiten ub',ub, (pg2.mean()+vfco*vl(values,old_values,ret,vfc)).item())
# whiten with rsqrt biased
aw=(adv-adv.mean())*torch.rsqrt(adv.var(unbiased=False)+1e-8)
pg3=-torch.min(ratio*aw,torch.clamp(ratio,1-cl,1+ch)*aw)
print('whiten rsqrt biased', (pg3.mean()+vfco*vl(values,old_values,ret,vfc)).item())
"ref 0.9666061401367188 guess no-whiten 0.8198745250701904 whiten ub True 0.9133155345916748 whiten ub False 0.9666061401367188 whiten rsqrt biased 0.9666061401367188 [stdout] ref 0.9666061401367188 guess no-whiten 0.8198745250701904 whiten ub True 0.9133155345916748 whiten ub False 0.9666061401367188 whiten rsqrt biased 0.9666061401367188
ppo_objective multi-seed verify
python3 -c "
import torch, reforacle
def wh(a): return (a-a.mean())*torch.rsqrt(a.var(unbiased=False)+1e-8)
def vl(v,ov,ret,clip):
vc=ov+torch.clamp(v-ov,-clip,clip)
return 0.5*torch.max((v-ret)**2,(vc-ret)**2).mean()
for seed in [1,2,7,11]:
g=torch.Generator().manual_seed(seed); T=6
rw=torch.randn(T,generator=g);v=torch.randn(T,generator=g);ov=torch.randn(T,generator=g)
lp=torch.randn(T,generator=g);olp=torch.randn(T,generator=g);nv=torch.randn(1,generator=g)[0]
gm,lam,cl,ch,vfc,vfco=0.97,0.9,0.15,0.25,0.3,0.7
adv=reforacle.gae(rw,v,nv,gm,lam); ret=reforacle.lambda_returns(rw,v,nv,gm,lam)
aw=wh(adv); ratio=torch.exp(lp-olp)
pg=-torch.min(ratio*aw,torch.clamp(ratio,1-cl,1+ch)*aw)
g_=pg.mean()+vfco*vl(v,ov,ret,vfc)
r=reforacle.ppo_objective(rw,v,ov,lp,olp,nv,gm,lam,cl,ch,vfc,vfco).item()
print(seed,'ref',r,'guess',g_.item(),'OK' if abs(r-g_.item())<1e-5 else 'FAIL')
"1 ref 0.8597356081008911 guess 0.8597355484962463 OK 2 ref 0.4844517111778259 guess 0.4844517111778259 OK 7 ref 1.9055957794189453 guess 1.9055957794189453 OK 11 ref 1.6804780960083008 guess 1.6804780960083008 OK [stdout] 1 ref 0.8597356081008911 guess 0.8597355484962463 OK 2 ref 0.4844517111778259 guess 0.4844517111778259 OK 7 ref 1.9055957794189453 guess 1.9055957794189453 OK 11 ref 1.6804780960083008 guess 1.6804780960083008 OK
edge cases masked/axis/2D
python3 -c "
import torch, reforacle
def show(n,*a,**k):
try:
r=getattr(reforacle,n)(*a,**k); print(n,'->',r.tolist() if torch.is_tensor(r) else r)
except Exception as e: print(n,'ERR',type(e).__name__,e)
# masked_mean empty
show('masked_mean',torch.randn(4),torch.zeros(4))
show('masked_sum',torch.randn(4),torch.zeros(4))
# masked_mean axis on 2D with mask along axis reduces?
V=torch.arange(12.).reshape(3,4); M=(torch.arange(12.)%2).reshape(3,4)
show('masked_mean',V,M,0)
show('masked_mean',V,M,-1)
print('manual axis0', ((V*M).sum(0)/M.sum(0)).tolist())
# entropy no mask? signature entropy(logits,mask)
lg=torch.randn(2,3,5)
# normalize with 2D
show('normalize',torch.randn(3,4),1e-5)
# discounted_returns 2D
show('discounted_returns',torch.tensor([[1.,2.,3.],[4.,5.,6.]]),0.9)
# gae/lambda 2D?
show('gae',torch.tensor([[1.,2.],[3.,4.]]),torch.tensor([[0.1,0.2],[0.3,0.4]]),torch.tensor([0.5,0.5]),0.9,0.95)
"masked_mean -> 0.0 masked_sum -> 0.0 masked_mean -> [0.0, 5.0, 0.0, 7.0] masked_mean -> [2.0, 6.0, 10.0] manual axis0 [nan, 5.0, nan, 7.0] normalize -> [[0.5412573218345642, 0.9305271506309509, -0.05964699015021324, 0.8245014548301697], [1.1010178327560425, 0.9877185821533203, 0.6323410868644714, -1.1882412433624268], [-0.06640219688415527, -0.8266169428825378, -1.0300291776657104, -1.846427083015442]] discounted_returns -> [[4.599999904632568, 6.5, 8.399999618530273], [4.0, 5.0, 6.0]] gae -> [[3.863250255584717, 5.622749328613281], [3.1500000953674316, 4.049999713897705]] [stdout] masked_mean -> 0.0 masked_sum -> 0.0 masked_mean -> [0.0, 5.0, 0.0, 7.0] masked_mean -> [2.0, 6.0, 10.0] manual axis0 [nan, 5.0, nan, 7.0] normalize -> [[0.5412573218345642, 0.9305271506309509, -0.05964699015021324, 0.8245014548301697], [1.1010178327560425, 0.9877185821533203, 0.6323410868644714, -1.1882412433624268], [-0.06640219688415527, -0.8266169428825378, -1.0300291776657104, -1.846427083015442]] discounted_returns -> [[4.599999904632568, 6.5, 8.399999618530273], [4.0, 5.0, 6.0]] gae -> [[3.863250255584717, 5.622749328613281], [3.1500000953674316, 4.049999713897705]]
verify normalize/entropy/masked denom
python3 -c "
import torch, reforacle
x=torch.randn(3,4,generator=torch.Generator().manual_seed(9))
r=reforacle.normalize(x,1e-5)
print('global unbiased', torch.allclose(r,(x-x.mean())/(x.std()+1e-5)))
print('global biased', torch.allclose(r,(x-x.mean())/(x.std(unbiased=False)+1e-5)))
# entropy log base
lg=torch.randn(2,3,5,generator=torch.Generator().manual_seed(1))
m=torch.ones(2,3)
e=reforacle.entropy(lg,m).item()
p=torch.softmax(lg,-1); ent=-(p*torch.log_softmax(lg,-1)).sum(-1)
print('entropy match', abs(e-((ent*m).sum()/m.sum()).item()))
ne=reforacle.normalized_entropy(lg,m).item()
import math
print('norm ent /logV', abs(ne-((ent*m).sum()/m.sum()).item()/math.log(5)))
# masked_mean single element denom test
print('single', reforacle.masked_mean(torch.tensor([4.0,9.0]),torch.tensor([1.0,0.0])).item())
"global unbiased True global biased False entropy match 0.0 norm ent /logV 5.299973293571014e-08 single 4.0 [stdout] global unbiased True global biased False entropy match 0.0 norm ent /logV 5.299973293571014e-08 single 4.0
from __future__ import annotations
import math
import torch
import torch.nn.functional as F
# ---------------------------------------------------------------------------
# internal helpers
# ---------------------------------------------------------------------------
def _token_logprobs(logits, labels):
lp = torch.log_softmax(logits, dim=-1)
return lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
def _masked_sum(values, mask, axis):
if axis is None:
return (values * mask).sum()
return (values * mask).sum(dim=axis)
def _masked_mean(values, mask, axis):
num = _masked_sum(values, mask, axis)
if axis is None:
den = mask.sum()
else:
den = mask.sum(dim=axis)
return num / den.clamp(min=1.0)
# ---------------------------------------------------------------------------
# primitives
# ---------------------------------------------------------------------------
def masked_mean(values, mask, axis=None):
return _masked_mean(values, mask, axis)
def masked_sum(values, mask, axis=None):
return _masked_sum(values, mask, axis)
def logsumexp(x, axis):
return torch.logsumexp(x, dim=axis)
def log_softmax(x, axis):
return torch.log_softmax(x, dim=axis)
def selective_logprobs(logits, labels, mask):
lp = _token_logprobs(logits, labels)
return (lp * mask).sum(dim=-1)
def token_logprobs(logits, labels):
return _token_logprobs(logits, labels)
def sequence_logprob(logits, labels, mask, length_normalize):
lp = _token_logprobs(logits, labels)
s = (lp * mask).sum(dim=-1)
if length_normalize:
s = s / mask.sum(dim=-1).clamp(min=1.0)
return s
def entropy(logits, mask):
lp = torch.log_softmax(logits, dim=-1)
p = lp.exp()
ent = -(p * lp).sum(dim=-1)
return _masked_mean(ent, mask, None)
def dpo_loss(pc, pr, rc, rr, beta, label_smoothing):
logits = beta * ((pc - rc) - (pr - rr))
loss = -(1.0 - label_smoothing) * F.logsigmoid(logits) - label_smoothing * F.logsigmoid(-logits)
return loss.mean()
def ipo_loss(pc, pr, rc, rr, beta):
diff = (pc - rc) - (pr - rr)
return ((diff - 1.0 / (2.0 * beta)) ** 2).mean()
def grpo_advantages(rewards, group_size, scale_by_std):
g = rewards.reshape(-1, group_size)
mean = g.mean(dim=1, keepdim=True)
adv = g - mean
if scale_by_std:
std = g.std(dim=1, unbiased=False, keepdim=True)
adv = adv / (std + 1e-6)
return adv.reshape(rewards.shape)
def gae(rewards, values, next_value, gamma, lam):
T = rewards.shape[0]
adv = torch.zeros_like(rewards)
last = torch.zeros_like(rewards[0]) if rewards.dim() > 1 else torch.zeros((), dtype=rewards.dtype)
for t in reversed(range(T)):
nv = next_value if t == T - 1 else values[t + 1]
delta = rewards[t] + gamma * nv - values[t]
last = delta + gamma * lam * last
adv[t] = last
return adv
def kl_penalty(logp, ref_logp, estimator):
d = logp - ref_logp
if estimator == "k1":
return d
if estimator == "k2":
return 0.5 * d * d
if estimator == "k3":
return torch.exp(-d) + d - 1.0
raise ValueError(estimator)
def clipped_pg_loss(logp, old_logp, advantages, mask, clip_low, clip_high):
ratio = torch.exp(logp - old_logp)
unclipped = ratio * advantages
clipped = torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * advantages
pg = -torch.min(unclipped, clipped)
return _masked_mean(pg, mask, None)
def value_loss(values, old_values, returns, clip):
vclipped = old_values + torch.clamp(values - old_values, -clip, clip)
l1 = (values - returns) ** 2
l2 = (vclipped - returns) ** 2
return 0.5 * torch.max(l1, l2).mean()
def whiten(values, mask, shift_mean):
mean = _masked_mean(values, mask, None)
var = _masked_mean((values - mean) ** 2, mask, None)
w = (values - mean) * torch.rsqrt(var + 1e-8)
if not shift_mean:
w = w + mean
return w
def dpo_sequence_loss(pc_logits, pr_logits, rc_logits, rr_logits,
chosen_labels, rejected_labels, chosen_mask, rejected_mask,
beta, label_smoothing):
def seq(logits, labels, mask):
lp = _token_logprobs(logits, labels)
return (lp * mask).sum(dim=-1)
pc = seq(pc_logits, chosen_labels, chosen_mask)
pr = seq(pr_logits, rejected_labels, rejected_mask)
rc = seq(rc_logits, chosen_labels, chosen_mask)
rr = seq(rr_logits, rejected_labels, rejected_mask)
return dpo_loss(pc, pr, rc, rr, beta, label_smoothing)
def grpo_objective(logits, old_logits, ref_logits, labels, completion_mask,
rewards, group_size, beta, clip_low, clip_high, scale_by_std,
kl_estimator):
lp = _token_logprobs(logits, labels)
olp = _token_logprobs(old_logits, labels)
rlp = _token_logprobs(ref_logits, labels)
adv = grpo_advantages(rewards, group_size, scale_by_std).unsqueeze(-1)
ratio = torch.exp(lp - olp)
pg = -torch.min(ratio * adv, torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * adv)
kl = kl_penalty(lp, rlp, kl_estimator)
return _masked_mean(pg + beta * kl, completion_mask, None)
def ppo_objective(rewards, values, old_values, logp, old_logp, next_value,
gamma, lam, clip_low, clip_high, vf_clip, vf_coef):
adv = gae(rewards, values, next_value, gamma, lam)
returns = adv + values
adv_w = (adv - adv.mean()) * torch.rsqrt(adv.var(unbiased=False) + 1e-8)
ratio = torch.exp(logp - old_logp)
pg = -torch.min(ratio * adv_w, torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * adv_w)
vf = value_loss(values, old_values, returns, vf_clip)
return pg.mean() + vf_coef * vf
def rloo_advantages(rewards, group_size):
g = rewards.reshape(-1, group_size)
loo = (g.sum(dim=1, keepdim=True) - g) / (group_size - 1)
return (g - loo).reshape(rewards.shape)
def reverse_kl(logp, ref_logp):
d = logp - ref_logp
return torch.exp(-d) + d - 1.0
def importance_ratio(logp, old_logp, clip):
ratio = torch.exp(logp - old_logp)
return torch.clamp(ratio, 1.0 - clip, 1.0 + clip)
def discounted_returns(rewards, gamma):
T = rewards.shape[0]
out = torch.zeros_like(rewards)
running = torch.zeros_like(rewards[0]) if rewards.dim() > 1 else torch.zeros((), dtype=rewards.dtype)
for t in reversed(range(T)):
running = rewards[t] + gamma * running
out[t] = running
return out
def normalize(x, eps):
return (x - x.mean()) / (x.std() + eps)
def top_p_mask(probs, p):
sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
cumsum = sorted_probs.cumsum(dim=-1)
keep_sorted = (cumsum - sorted_probs) < p
keep = torch.zeros_like(keep_sorted)
keep.scatter_(-1, sorted_idx, keep_sorted)
out = probs * keep
out = out / out.sum(dim=-1, keepdim=True)
return out
def smoothed_nll(logits, labels, smoothing):
lp = torch.log_softmax(logits, dim=-1)
nll = -lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
smooth = -lp.mean(dim=-1)
loss = (1.0 - smoothing) * nll + smoothing * smooth
return loss.mean()
def bradley_terry_logit(chosen_reward, rejected_reward, beta):
return beta * (chosen_reward - rejected_reward)
def rloo_objective(logits, old_logits, labels, mask, rewards, group_size, clip_low, clip_high):
lp = _token_logprobs(logits, labels)
olp = _token_logprobs(old_logits, labels)
adv = rloo_advantages(rewards, group_size).unsqueeze(-1)
ratio = torch.exp(lp - olp)
pg = -torch.min(ratio * adv, torch.clamp(ratio, 1.0 - clip_low, 1.0 + clip_high) * adv)
return _masked_mean(pg, mask, None)
def reverse_kl_objective(logits, ref_logits, labels, mask, advantages, beta):
lp = _token_logprobs(logits, labels)
rlp = _token_logprobs(ref_logits, labels)
pg = -advantages.unsqueeze(-1) * lp
kl = reverse_kl(lp, rlp)
return _masked_mean(pg + beta * kl, mask, None)
def cross_entropy(logits, labels, ignore_index):
V = logits.shape[-1]
return F.cross_entropy(logits.reshape(-1, V), labels.reshape(-1), ignore_index=ignore_index)
def top_k_mask(logits, k):
idx = logits.topk(k, dim=-1).indices
mask = torch.zeros_like(logits, dtype=torch.bool)
mask.scatter_(-1, idx, True)
return mask
def group_mean_baseline(rewards, group_size):
g = rewards.reshape(-1, group_size)
mean = g.mean(dim=1, keepdim=True)
return (g - mean).reshape(rewards.shape)
def lambda_returns(rewards, values, next_value, gamma, lam):
return gae(rewards, values, next_value, gamma, lam) + values
def symmetric_kl(logp, ref_logp):
d = logp - ref_logp
rev = torch.exp(-d) + d - 1.0
fwd = torch.exp(d) - d - 1.0
return 0.5 * (rev + fwd)
def huber_value_loss(values, returns, delta):
return F.huber_loss(values, returns, delta=delta)
def normalized_entropy(logits, mask):
V = logits.shape[-1]
return entropy(logits, mask) / math.log(V)
def clip_fraction(logp, old_logp, clip):
ratio = torch.exp(logp - old_logp)
return ((ratio - 1.0).abs() > clip).float().mean()
def masked_whiten(values, mask, shift_mean):
mean = _masked_mean(values, mask, None)
var = _masked_mean((values - mean) ** 2, mask, None)
w = (values - mean) * torch.rsqrt(var + 1e-8)
w = w * mask
if not shift_mean:
w = w + mean
return w
def logprob_at_temperature(logits, labels, temperature):
lp = torch.log_softmax(logits / temperature, dim=-1)
return lp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)
def advantage_mean_std(advantages, mask):
mean = _masked_mean(advantages, mask, None)
var = _masked_mean((advantages - mean) ** 2, mask, None)
return mean, torch.sqrt(var)
def argmax_tokens(logits):
return logits.argmax(dim=-1)
def mode_label(labels):
vals, counts = torch.unique(labels, return_counts=True)
maxc = counts.max()
modes = vals[counts == maxc]
return modes.max()
def median_reward(rewards):
return torch.quantile(rewards.float(), 0.5)
def masked_max(values, mask):
filled = values.masked_fill(mask == 0, float("-inf"))
return filled.max()
def quantile_lower(x, q):
return torch.quantile(x.float(), q, interpolation="lower")
def pad_mask_from_lengths(lengths, max_len):
ar = torch.arange(max_len, device=lengths.device)
return (ar.unsqueeze(0) < lengths.unsqueeze(-1)).float()
def first_nonzero_index(mask):
nz = mask != 0
first = nz.float().argmax(dim=-1)
has = nz.any(dim=-1)
return torch.where(has, first, torch.full_like(first, -1))
def cumulative_max(x):
cm = torch.cummax(x, dim=0).values
out = torch.empty_like(x)
out[0] = x[0]
if x.shape[0] > 1:
out[1:] = cm[:-1]
return out
def bucketize_reward(rewards, boundaries):
return torch.bucketize(rewards, boundaries)
Full comparison test
python3 -c "
import torch, importlib
import reforacle as R
import objectives as O
importlib.reload(O)
g=torch.Generator().manual_seed(123)
def rn(*s): return torch.randn(*s,generator=g)
def cmp(name,*a,**k):
try: r=getattr(R,name)(*a,**k)
except Exception as e:
try: o=getattr(O,name)(*a,**k); print('DIFF-ERR',name,'ref err',e); return
except Exception as e2: print('bothERR',name); return
o=getattr(O,name)(*a,**k)
def t(x):
if isinstance(x,tuple): return torch.cat([torch.as_tensor(y).flatten().float() for y in x])
return torch.as_tensor(x).flatten().float()
a1,o1=t(r),t(o)
ok=torch.allclose(a1,o1,atol=1e-5,rtol=1e-4,equal_nan=True)
print('OK ' if ok else 'FAIL', name, '' if ok else (a1[:4].tolist(),o1[:4].tolist()))
B,T,V=6,4,5
logits=rn(B,T,V); old=rn(B,T,V); ref=rn(B,T,V)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.3).float()
vals=rn(B,T); m2=(torch.rand(B,T,generator=g)>0.3).float()
cmp('masked_mean',vals,m2); cmp('masked_mean',vals,m2,1); cmp('masked_mean',vals,m2,0); cmp('masked_mean',vals,m2,-1)
cmp('masked_sum',vals,m2); cmp('masked_sum',vals,m2,1)
cmp('logsumexp',logits,-1); cmp('logsumexp',logits,1)
cmp('log_softmax',logits,-1)
cmp('selective_logprobs',logits,labels,mask)
cmp('token_logprobs',logits,labels)
cmp('sequence_logprob',logits,labels,mask,True); cmp('sequence_logprob',logits,labels,mask,False)
cmp('entropy',logits,mask)
cmp('normalized_entropy',logits,mask)
pc=rn(B);pr=rn(B);rc=rn(B);rr=rn(B)
cmp('dpo_loss',pc,pr,rc,rr,0.1,0.0); cmp('dpo_loss',pc,pr,rc,rr,0.3,0.15)
cmp('ipo_loss',pc,pr,rc,rr,0.2)
cmp('bradley_terry_logit',rn(B),rn(B),0.5)
rewards=rn(6)
cmp('grpo_advantages',rewards,3,True); cmp('grpo_advantages',rewards,2,False)
cmp('rloo_advantages',rewards,3); cmp('group_mean_baseline',rewards,2)
rw=rn(5);vv=rn(5);nv=rn(1)[0]
cmp('gae',rw,vv,nv,0.99,0.95); cmp('lambda_returns',rw,vv,nv,0.99,0.95)
cmp('discounted_returns',rw,0.9); cmp('discounted_returns',rn(4,3),0.9)
lp=rn(B,T);rlp=rn(B,T);olp=rn(B,T)
for e in ['k1','k2','k3']: cmp('kl_penalty',lp,rlp,e)
cmp('reverse_kl',lp,rlp); cmp('symmetric_kl',lp,rlp)
cmp('importance_ratio',lp,olp,0.2)
adv=rn(B,T)
cmp('clipped_pg_loss',lp,olp,adv,mask,0.2,0.3)
cmp('clip_fraction',lp,olp,0.2)
cmp('value_loss',rn(10),rn(10),rn(10),0.2)
cmp('huber_value_loss',rn(10),rn(10),1.0)
cmp('whiten',rn(8),(torch.rand(8,generator=g)>0.3).float(),True)
cmp('whiten',rn(8),(torch.rand(8,generator=g)>0.3).float(),False)
cmp('masked_whiten',rn(8),(torch.rand(8,generator=g)>0.3).float(),True)
cmp('masked_whiten',rn(8),(torch.rand(8,generator=g)>0.3).float(),False)
cmp('normalize',rn(3,4),1e-5)
cmp('smoothed_nll',logits,labels,0.1)
cmp('logprob_at_temperature',logits,labels,2.0)
cmp('cross_entropy',logits,labels,-100)
lab2=labels.clone();lab2[0,0]=-100;cmp('cross_entropy',logits,lab2,-100)
cmp('top_p_mask',torch.softmax(rn(3,5),-1),0.7)
cmp('top_k_mask',rn(3,6),2)
cmp('argmax_tokens',logits)
cmp('mode_label',torch.randint(0,4,(9,),generator=g))
cmp('median_reward',rn(7)); cmp('median_reward',rn(8))
cmp('masked_max',rn(3,4),(torch.rand(3,4,generator=g)>0.3).float())
cmp('quantile_lower',rn(10),0.3)
cmp('pad_mask_from_lengths',torch.tensor([2,3,1,4]),5)
cmp('first_nonzero_index',(torch.rand(4,5,generator=g)>0.6).float())
cmp('cumulative_max',rn(6)); cmp('cumulative_max',rn(5,3))
cmp('bucketize_reward',rn(6),torch.tensor([-1.,0.,1.]))
cmp('advantage_mean_std',rn(8),(torch.rand(8,generator=g)>0.3).float())
# composites
cmp('dpo_sequence_loss',logits,old,ref,rn(B,T,V),labels,torch.randint(0,V,(B,T),generator=g),mask,(torch.rand(B,T,generator=g)>0.3).float(),0.1,0.0)
cmp('grpo_objective',logits,old,ref,labels,mask,rn(B),3,0.05,0.2,0.2,True,'k3')
cmp('rloo_objective',logits,old,labels,mask,rn(B),2,0.2,0.2)
cmp('reverse_kl_objective',logits,ref,labels,mask,rn(B),0.1)
cmp('ppo_objective',rw,vv,rn(5),rn(5),rn(5),nv,0.99,0.95,0.2,0.2,0.2,0.5)
"OK masked_mean OK masked_mean OK masked_mean OK masked_mean OK masked_sum OK masked_sum OK logsumexp OK logsumexp OK log_softmax OK selective_logprobs OK token_logprobs OK sequence_logprob OK sequence_logprob OK entropy OK normalized_entropy OK dpo_loss OK dpo_loss OK ipo_loss OK bradley_terry_logit OK grpo_advantages OK grpo_advantages OK rloo_advantages OK group_mean_baseline OK gae OK lambda_returns OK discounted_returns OK discounted_returns OK kl_penalty OK kl_penalty OK kl_penalty OK reverse_kl OK symmetric_kl OK importance_ratio OK clipped_pg_loss OK clip_fraction OK value_loss OK huber_value_loss OK whiten OK whiten OK masked_whiten OK masked_whiten OK normalize OK smoothed_nll OK logprob_at_temperature OK cross_entropy OK cross_entropy OK top_p_mask OK top_k_mask OK argmax_tokens OK mode_label OK median_reward OK median_reward OK masked_max OK quantile_lower OK pad_mask_from_lengths OK first_nonzero_index OK cumulative_max OK cumulative_max OK bucketize_reward OK advantage_mean_std OK dpo_sequence_loss OK grpo_objective OK rloo_objective OK reverse_kl_objective OK ppo_objective [stdout] OK masked_mean OK masked_mean OK masked_mean OK masked_mean OK masked_sum OK masked_sum OK logsumexp OK logsumexp OK log_softmax OK selective_logprobs OK token_logprobs OK sequence_logprob OK sequence_logprob OK entropy OK normalized_entropy OK dpo_loss OK dpo_loss OK ipo_loss OK bradley_terry_logit OK grpo_advantages OK grpo_advantages OK rloo_advantages OK group_mean_baseline OK gae OK lambda_returns OK discounted_returns OK discounted_returns OK kl_penalty OK kl_penalty OK kl_penalty OK reverse_kl OK symmetric_kl OK importance_ratio OK clipped_pg_loss OK clip_fraction OK value_loss OK huber_value_loss OK whiten OK whiten OK masked_whiten OK masked_whiten OK normalize OK smoothed_nll OK logprob_at_temperature OK cross_entropy OK cross_entropy OK top_p_mask OK top_k_mask OK argmax_tokens OK mode_label OK median_reward OK median_reward OK masked_max OK quantile_lower OK pad_mask_from_lengths OK first_nonzero_index OK cumulative_max OK cumulative_max OK bucketize_reward OK advantage_mean_std OK dpo_sequence_loss OK grpo_objective OK rloo_objective OK reverse_kl_objective OK ppo_objective
Cross-seed robustness test
python3 -c "
import torch, importlib
import reforacle as R
import objectives as O
def t(x):
if isinstance(x,tuple): return torch.cat([torch.as_tensor(y).flatten().float() for y in x])
return torch.as_tensor(x).flatten().float()
fails=[]
def cmp(name,*a,**k):
r=getattr(R,name)(*a,**k); o=getattr(O,name)(*a,**k)
if not torch.allclose(t(r),t(o),atol=1e-5,rtol=1e-4,equal_nan=True):
fails.append((name,t(r)[:5].tolist(),t(o)[:5].tolist()))
for seed in range(50,60):
g=torch.Generator().manual_seed(seed)
def rn(*s): return torch.randn(*s,generator=g)
B,T,V=8,5,7
logits=rn(B,T,V);old=rn(B,T,V);ref=rn(B,T,V)
labels=torch.randint(0,V,(B,T),generator=g)
mask=(torch.rand(B,T,generator=g)>0.4).float()
vals=rn(B,T);m2=(torch.rand(B,T,generator=g)>0.4).float()
cmp('masked_mean',vals,m2,1);cmp('masked_mean',vals,m2);cmp('masked_sum',vals,m2,0)
cmp('logsumexp',logits,-1);cmp('log_softmax',logits,1)
cmp('selective_logprobs',logits,labels,mask);cmp('token_logprobs',logits,labels)
cmp('sequence_logprob',logits,labels,mask,True);cmp('sequence_logprob',logits,labels,mask,False)
cmp('entropy',logits,mask);cmp('normalized_entropy',logits,mask)
cmp('dpo_loss',rn(B),rn(B),rn(B),rn(B),0.2,0.1);cmp('ipo_loss',rn(B),rn(B),rn(B),rn(B),0.3)
cmp('bradley_terry_logit',rn(4),rn(4),0.7)
cmp('grpo_advantages',rn(8),4,True);cmp('grpo_advantages',rn(8),2,False)
cmp('rloo_advantages',rn(8),4);cmp('group_mean_baseline',rn(8),2)
rw=rn(6);vv=rn(6);nv=rn(1)[0]
cmp('gae',rw,vv,nv,0.95,0.9);cmp('lambda_returns',rw,vv,nv,0.95,0.9);cmp('discounted_returns',rw,0.97)
lp=rn(B,T);rlp=rn(B,T);olp=rn(B,T)
for e in ['k1','k2','k3']: cmp('kl_penalty',lp,rlp,e)
cmp('reverse_kl',lp,rlp);cmp('symmetric_kl',lp,rlp);cmp('importance_ratio',lp,olp,0.3)
cmp('clipped_pg_loss',lp,olp,rn(B,T),mask,0.1,0.4);cmp('clip_fraction',lp,olp,0.25)
cmp('value_loss',rn(12),rn(12),rn(12),0.3);cmp('huber_value_loss',rn(12),rn(12),0.5)
cmp('whiten',rn(9),(torch.rand(9,generator=g)>0.4).float(),True)
cmp('masked_whiten',rn(9),(torch.rand(9,generator=g)>0.4).float(),False)
cmp('normalize',rn(4,5),1e-6);cmp('smoothed_nll',logits,labels,0.15)
cmp('logprob_at_temperature',logits,labels,0.5)
cmp('cross_entropy',logits,labels,-100)
cmp('top_p_mask',torch.softmax(rn(4,7),-1),0.85);cmp('top_k_mask',rn(4,7),3)
cmp('argmax_tokens',logits);cmp('mode_label',torch.randint(0,5,(11,),generator=g))
cmp('median_reward',rn(9));cmp('median_reward',rn(10))
cmp('masked_max',rn(4,5),(torch.rand(4,5,generator=g)>0.4).float())
cmp('quantile_lower',rn(11),0.6)
cmp('pad_mask_from_lengths',torch.randint(1,6,(5,),generator=g),6)
cmp('first_nonzero_index',(torch.rand(5,6,generator=g)>0.5).float())
cmp('cumulative_max',rn(7));cmp('cumulative_max',rn(6,4))
cmp('bucketize_reward',rn(8),torch.tensor([-0.5,0.5,1.5]))
cmp('advantage_mean_std',rn(9),(torch.rand(9,generator=g)>0.4).float())
cmp('dpo_sequence_loss',logits,old,ref,rn(B,T,V),labels,torch.randint(0,V,(B,T),generator=g),mask,(torch.rand(B,T,generator=g)>0.4).float(),0.15,0.05)
cmp('grpo_objective',logits,old,ref,labels,mask,rn(B),4,0.1,0.15,0.25,True,'k2')
cmp('grpo_objective',logits,old,ref,labels,mask,rn(B),2,0.0,0.2,0.2,False,'k1')
cmp('rloo_objective',logits,old,labels,mask,rn(B),4,0.15,0.25)
cmp('reverse_kl_objective',logits,ref,labels,mask,rn(B),0.2)
cmp('ppo_objective',rw,vv,rn(6),rn(6),rn(6),nv,0.95,0.9,0.15,0.25,0.3,0.7)
print('FAILS:', len(fails))
for f in fails[:20]: print(f)
"FAILS: 0 [stdout] FAILS: 0
Verify standalone, no reforacle
cd /tmp && cp /app/objectives.py /tmp/obj_check.py && grep -n "reforacle" /tmp/obj_check.py; echo "grep exit: $?"; python3 -c "import sys; sys.path.insert(0,'/tmp'); import obj_check; print('standalone import OK, functions:', len([x for x in dir(obj_check) if not x.startswith('_')]))"; bash /app/timer.shgrep exit: 1 standalone import OK, functions: 53 Remaining time (hours:minutes): 1:44 Shell cwd was reset to /app [stdout] grep exit: 1 standalone import OK, functions: 53 Remaining time (hours:minutes): 1:44 [stderr] Shell cwd was reset to /app
No file changes, the agent left the workspace untouched.
=== SyncVals verifier (post-training objectives, sealed oracle) === FUNCTIONS PASSED: 42/49 (partial_score=0.8571) REWARD = 0 (1 iff ALL 49 functions match the sealed reference exactly) The agent probes a sealed compiled oracle (reforacle) to recover 49 post-training objective conventions (DPO/IPO/GRPO/PPO/GAE/KL/RLOO) and reimplements them in torch.
Reproduce this trial: git checkout 2f94510 && PYTHONPATH=src python3 scripts/build_site.py , then open trial/trial_bd8ad8626e804c09. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.
Trial trial_bd8ad8626e804c09 · verifier authoritative; classifier explanatory.