[Auto Sync] Update test_deterministic.py (20260131) (#18034)
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Jiayi Yuan <34369239+jy-yuan@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
github-actions[bot]
Jiayi Yuan
gemini-code-assist[bot]
parent
fb609669ca
commit
1acae30806
@@ -320,6 +320,9 @@ class TokenIdsAndLogprobs:
|
|||||||
token_ids: List[int]
|
token_ids: List[int]
|
||||||
logprobs: List[float]
|
logprobs: List[float]
|
||||||
|
|
||||||
|
# Logprob differences smaller than this are treated as non-divergent.
|
||||||
|
DIVERGENCE_EPS = 0.0
|
||||||
|
|
||||||
def __add__(self, other):
|
def __add__(self, other):
|
||||||
return TokenIdsAndLogprobs(
|
return TokenIdsAndLogprobs(
|
||||||
token_ids=self.token_ids + other.token_ids,
|
token_ids=self.token_ids + other.token_ids,
|
||||||
@@ -381,13 +384,24 @@ class TokenIdsAndLogprobs:
|
|||||||
|
|
||||||
# K3 approximation: KL(A||B) ≈ (exp(logr) - 1) - logr, where logr = log_a - log_b
|
# K3 approximation: KL(A||B) ≈ (exp(logr) - 1) - logr, where logr = log_a - log_b
|
||||||
logr = logprobs_a - logprobs_b
|
logr = logprobs_a - logprobs_b
|
||||||
kl_per_token = (np.exp(logr) - 1) - logr
|
diverge_mask = np.abs(logr) > cls.DIVERGENCE_EPS
|
||||||
kl_mean = np.mean(kl_per_token)
|
diverge_count = int(np.count_nonzero(diverge_mask))
|
||||||
kl_max = np.max(kl_per_token)
|
total_count = int(logr.shape[0])
|
||||||
|
|
||||||
print(f" KL(A||B) mean: {kl_mean:.10e}")
|
if diverge_count > 0:
|
||||||
print(f" KL(A||B) max : {kl_max:.10e}")
|
kl_per_token = (np.exp(logr) - 1) - logr
|
||||||
print(f" Mean absolute logprob diff: {np.mean(np.abs(logr)):.10e}")
|
kl_divergent = kl_per_token[diverge_mask]
|
||||||
|
kl_mean = float(np.mean(kl_divergent))
|
||||||
|
kl_max = float(np.max(kl_divergent))
|
||||||
|
mean_abs_logr = float(np.mean(np.abs(logr[diverge_mask])))
|
||||||
|
print(f" Divergent tokens: {diverge_count}/{total_count}")
|
||||||
|
print(f" KL(A||B) mean (divergent): {kl_mean:.10e}")
|
||||||
|
print(f" KL(A||B) max (divergent): {kl_max:.10e}")
|
||||||
|
print(
|
||||||
|
f" Mean absolute logprob diff (divergent): {mean_abs_logr:.10e}"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
print(f" Divergent tokens: 0/{total_count}")
|
||||||
|
|
||||||
return token_match and logprobs_match
|
return token_match and logprobs_match
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user