add kl divergence for forward_pass_logit_checker

AI-Hypercomputer / maxtext

A simple, performant and scalable Jax LLM!

Apache License 2.0

1.47k stars 275 forks source link

Closed ZhaoyueCheng closed 1 month ago

ZhaoyueCheng commented 1 month ago

add kl divergence option for forward_pass_logit_checker so that it compares the kl divergence of the output probability distribution with the golden output probability distribution
by default it's still using the original option to check numerical differences so it don't break other tests running

http://shortn/_XlXElFUgnv added the logs of 4 tests (kl passed, kl failed, numerical passed, numerical failed) with updated logging message here

rdyro commented 1 month ago

Great, thanks for implementing this! This looks good to me