[Auto Sync] Update test_deterministic.py (20251112) (#13128)

Co-authored-by: Stefan He <hebiaobuaa@gmail.com>
This commit is contained in:
Lianmin Zheng
2025-11-12 11:12:23 -08:00
committed by GitHub
parent 4edb240112
commit d646cf6347

View File

@@ -325,32 +325,66 @@ class TokenIdsAndLogprobs:
@classmethod
def compare(cls, a: "TokenIdsAndLogprobs", b: "TokenIdsAndLogprobs"):
import numpy as np
assert len(a.token_ids) == len(b.token_ids)
token_match = a.token_ids == b.token_ids
logprobs_match = a.logprobs == b.logprobs
if token_match:
print(f"Token match: {a.token_ids}")
print(f"Token match")
else:
print(f"Token mismatch: {a.token_ids=} {b.token_ids=}")
print(f"Token mismatch: {a.token_ids=} {b.token_ids=}")
if logprobs_match:
print(f"Logprobs match:", a.logprobs)
print(f"Logprobs match:", a.logprobs[:5])
else:
print(f"Logprobs mismatch")
print(f"Logprobs mismatch")
# Only print first 5 elements for readability
n_show = 5
a_show = a.logprobs[:n_show]
b_show = b.logprobs[:n_show]
print(
" A: ",
[f"{x:.10f}" if x is not None else "None" for x in a.logprobs],
[f"{x:.10f}" if x is not None else "None" for x in a_show],
f"... ({len(a.logprobs)} total)" if len(a.logprobs) > n_show else "",
)
print(
" B: ",
[f"{x:.10f}" if x is not None else "None" for x in b.logprobs],
[f"{x:.10f}" if x is not None else "None" for x in b_show],
f"... ({len(b.logprobs)} total)" if len(b.logprobs) > n_show else "",
)
diff = [
abs(x - y) if x is not None else float("nan")
for x, y in zip(a.logprobs, b.logprobs)
]
print(" Diff:", [f"{x:.10e}" for x in diff])
print(
" Diff:",
[f"{x:.10e}" for x in diff[:n_show]],
f"... ({len(diff)} total)" if len(diff) > n_show else "",
)
# Compute KL-divergence using K3 approximation
# KL(P||Q) ≈ (exp(log(P) - log(Q)) - 1) - (log(P) - log(Q))
# This is based on selected token logprobs only
valid_pairs = [
(lp_a, lp_b)
for lp_a, lp_b in zip(a.logprobs, b.logprobs)
if lp_a is not None and lp_b is not None
]
if valid_pairs and token_match:
logprobs_a = np.array([lp for lp, _ in valid_pairs])
logprobs_b = np.array([lp for _, lp in valid_pairs])
# K3 approximation: KL(A||B) ≈ (exp(logr) - 1) - logr, where logr = log_a - log_b
logr = logprobs_a - logprobs_b
kl_per_token = (np.exp(logr) - 1) - logr
kl_mean = np.mean(kl_per_token)
kl_max = np.max(kl_per_token)
print(f" KL(A||B) mean: {kl_mean:.10e}")
print(f" KL(A||B) max : {kl_max:.10e}")
print(f" Mean absolute logprob diff: {np.mean(np.abs(logr)):.10e}")
return token_match and logprobs_match