Add DPO correctness integration test and respect dpo.max_prompt_length in train_dpo.py

igorts-git · igorts-git · commit 8de36c3b7cf8 · 2026-06-02T10:24:04.000-07:00
diff --git a/.gitignore b/.gitignore
@@ -155,3 +155,4 @@ gha-creds-*.json
 
 # vscode workspace
 maxtext.code-workspace
+maxtext_output/
diff --git a/src/maxtext/trainers/post_train/dpo/train_dpo.py b/src/maxtext/trainers/post_train/dpo/train_dpo.py
@@ -91,6 +91,7 @@ def get_tunix_config(mt_config: MaxTextConfig) -> DPOTrainingConfig:
         set_profile_options=set_profile_options,
     )
 
+
   return DPOTrainingConfig(
       eval_every_n_steps=mt_config.eval_interval,
       max_steps=mt_config.steps,
diff --git a/tests/post_training/integration/dpo_trainer_correctness_test.py b/tests/post_training/integration/dpo_trainer_correctness_test.py

Original file line number	Diff line number	Diff line change
`@@ -91,6 +91,7 @@ def get_tunix_config(mt_config: MaxTextConfig) -> DPOTrainingConfig:`
`91`	`91`	`set_profile_options=set_profile_options,`
`92`	`92`	`)`
`93`	`93`
	`94`	`+`
`94`	`95`	`return DPOTrainingConfig(`
`95`	`96`	`eval_every_n_steps=mt_config.eval_interval,`
`96`	`97`	`max_steps=mt_config.steps,`