Delete legacy DPO implementation

[igorts-git]

Description

This pull request deletes the legacy "native" Direct Preference Optimization (DPO) implementation from the pre-training code paths, leaving only the recently added, Tunix-based post-training DPO and ORPO pipelines in place.

Along the way also renamed the metric eval/dpo_reward_accuracy to eval/avg_dpo_reward_accuracy to reflect the fact that Tunix aggregates the eval metrics and provides us with the average value. The legacy metric aggregation code was removed.

As follow up to this PR I plan to send out PRs with the following:

• Add a notebook that runs DPO and ORPO.
• Add support of Grain and TFDS in the new DPO implementation
• (Still haven't started) Enable LoRA in the new DPO.

BUGS: b/485626968

Tests

CI tests
Manual runs of DPO, ORPO, and SFT on a local VM.

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

❌ Patch coverage is 61.53846% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/train_utils.py	0.00%	2 Missing and 1 partial ⚠️
src/maxtext/common/metric_logger.py	0.00%	1 Missing and 1 partial ⚠️
src/maxtext/input_pipeline/tfds_data_processing.py	66.66%	0 Missing and 2 partials ⚠️
src/maxtext/trainers/pre_train/train.py	80.00%	2 Missing ⚠️
src/maxtext/utils/gradient_accumulation.py	50.00%	1 Missing ⚠️

📢 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]

## 📋 Review Summary

This Pull Request successfully deletes the legacy "native" Direct Preference Optimization (DPO) implementation from the pre-training code paths. This significant cleanup simplifies the codebase by removing redundant training logic, specialized state handling for reference models, and legacy input pipelines, directing users toward the newer Tunix-based post-training DPO pipeline.

🔍 General Feedback

• Comprehensive Cleanup: The removal is thorough across training (train.py), state management (train_state_nnx.py), input pipelines (tfds_data_processing.py, grain_data_processing.py), and logging (metric_logger.py).
• Excellent UX: The addition of a ValueError in validate_train_config to explicitly guide users to train_dpo.py is a high-quality detail that improves discoverability and prevents confusion during the transition.
• Maintainable Metrics: Transitioning metric_logger.py to use key-based checks for DPO reward accuracy instead of the use_dpo flag is a more robust approach that supports both current and future post-training workflows.

[ecnal-cienet]

Generally lgtm.
on maxtext_utils.py:
two config.use_dpo branches left here that are now unreachable since validate_train_config raises early:

• line 169 in get_shaped_batch (adds chosen/rejected ShapeDtypeStructs)
• line 1163 in the TFLOPS calc (adds reference-model TFLOPS)

I think both can dropped.