DPO integration test - logits comparision against Pytorch

[igorts-git]

Description

Add a logits comparison test for MaxText DPO vs HuggingFace TRL implementation.

The new test consists of two parts:

1. generate_dpo_golden_data_and_compare_pytorch_logits.py - compares DPO in MaxText vs HF TRL and saves golden data json file.
2. A unit test that only runs MaxText DPO and compares the results against the golden data in the json files.

This separation allows us to not have a dependency on TRL in CI.
The tests are running on CPU to reduce the GPU/TPU floating point differences. We use a tiny 2-layer qwen2-like model to make sure that this unit test runs fast.

A separate PR will be sent out to add a similar test for ORPO. The Tunix implementation would first need to be adjusted a little bit to match the canonical implementation (see google/tunix#1574).

Tests

Ran the test manually.

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!

[github-actions]

🤖 Hi @igorts-git, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @igorts-git, but I was unable to process your request. Please see the logs for more details.

[github-actions]

🤖 Hi @igorts-git, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR introduces a robust integration test verifying JAX DPO correctness against Hugging Face's canonical TRL implementation. By running a structurally identical tiny architecture and saving golden results, the test achieves high verification confidence while maintaining zero PyTorch/TRL dependencies in CI. The implementation quality is extremely high, with clear documentation and strict adherence to project standards.

🔍 General Feedback

• Symmetric Architecture Design: The miniaturized 2-layer model approach running entirely on CPU is excellent for avoiding GPU/TPU non-determinism while remaining extremely fast.
• Well-documented Tolerances: The rigorous calibration of DPO loss and log probability tolerances across multiple random seeds is an exceptional testing practice and is highly commendable.
• Robust Environment Handling: Clean setup and restoration of JAX environment configurations and training hooks prevent state leakage to neighboring tests.
• Suggested Improvements: A few minor optimizations around dynamic temporary directories, offline/hermetic tokenizer resolution, and explicit class-level attribute mutations have been suggested for maximum reliability.

[github-actions]

🤖 Hi @igorts-git, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @igorts-git, but I was unable to process your request. Please see the logs for more details.