fix test for mixtral logit

[shuningjin]

Description

Fix mixtral-8x7b forward logit check for XLML test

• Previous PR changes the criteria from --atol=3 --rtol=1 to --max_kl_div=2e-3, which is too strict.

AssertionError: KL divergence values exceed the specified threshold of 0.002. Max divergence: 0.0020157103426754475

• Relax it a bit as --max_kl_div=3e-3.

Fix: b/474617159, b/473421624

Tests

Run tests.forward_pass_logit_checker on v6e-256

• cmd: https://paste.googleplex.com/6166083503521792
• log: https://cloudlogging.app.goo.gl/W7dc78udFyw4ibNv8

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!