-
Notifications
You must be signed in to change notification settings - Fork 26.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Weights not aligned for pt and jax #22425
Comments
Hey @crystina-z! Thanks for the great code example! I see that you're running the script on TPU, could you try repeating the benchmark using the highest JAX matmul precision (see #15754 (comment) for details)? I think this should close the gap to PyTorch. Some more details about the behaviour of JAX matmul here: jax-ml/jax#10413 (comment) |
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread. Please note that issues that do not follow the contributing guidelines are likely to be ignored. |
System Info
transformers
version: 4.26.1Who can help?
@sanchit-gandhi since it's about the weights in jax, and @ArthurZucker @younesbelkada since the sample below is based on XLM-R
Information
Tasks
examples
folder (such as GLUE/SQuAD, ...)Reproduction
Expected behavior
The above script was supposed to output the same (or very close) values, however, it would produce:
The first two lines are the same (results using jax model, from jax weights or pytorch weights), however, they are different from the third line, the results produced by pytorch model. The difference is around 1~2 decimal points (e.g.,
64.36 vs 64.297
and even26.371 vs 26.475
, which isn't neglectable)The text was updated successfully, but these errors were encountered: