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

igorts-git · igorts-git · commit d8fbfc83e00d · 2026-05-29T23:50:56.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
@@ -90,7 +90,7 @@ def get_tunix_config(mt_config: MaxTextConfig) -> DPOTrainingConfig:
         set_profile_options=set_profile_options,
     )
 
-  max_prompt_length = mt_config.max_target_length // 2
+  max_prompt_length = mt_config.dpo.max_prompt_length or (mt_config.max_target_length // 2)
   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
`@@ -90,7 +90,7 @@ def get_tunix_config(mt_config: MaxTextConfig) -> DPOTrainingConfig:`
`90`	`90`	`set_profile_options=set_profile_options,`
`91`	`91`	`)`
`92`	`92`
`93`		`- max_prompt_length = mt_config.max_target_length // 2`
	`93`	`+ max_prompt_length = mt_config.dpo.max_prompt_length or (mt_config.max_target_length // 2)`
`94`	`94`	`return DPOTrainingConfig(`
`95`	`95`	`eval_every_n_steps=mt_config.eval_interval,`
`96`	`96`	`max_steps=mt_config.steps,`