Skip to content

Commit 933ee07

Browse files
igorts-gitigorts-git
committed
Add logits correctness integration tests for DPO against HF TRL
This commit introduces a JAX DPO correctness integration test that validates JAX DPO training step metrics (loss, margin, chosen/rejected logprobs) against stored golden outputs. It also includes the CPU-only JAX/PyTorch parallel validation script used to verify parity against Hugging Face TRL and generate the remote-canonical golden assets.
1 parent 7ac8f66 commit 933ee07

8 files changed

Lines changed: 1011 additions & 0 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,3 +155,4 @@ gha-creds-*.json
155155

156156
# vscode workspace
157157
maxtext.code-workspace
158+
maxtext_output/

‎tests/assets/golden_logits/golden_dpo_correctness.json‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
{
2+
"explicit_prompt_len_3_column": {
3+
"loss_step_1": 0.6931471824645996,
4+
"margin_step_1": 0.0,
5+
"loss": 0.6573728919029236,
6+
"margin": 0.0728759765625,
7+
"chosen_logps": -854.4208984375,
8+
"rejected_logps": -530.0099487304688
9+
},
10+
"default_prompt_len_2_column": {
11+
"loss_step_1": 0.6931471824645996,
12+
"margin_step_1": 0.0,
13+
"loss": 0.6954630613327026,
14+
"margin": -0.004626465495675802,
15+
"chosen_logps": -875.7848510742188,
16+
"rejected_logps": -502.052978515625
17+
}
18+
}

‎tests/assets/local_datasets/dpo/dpo_2_column_dataset.json‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
[
2+
{
3+
"chosen": "\n\nHuman: How does gradient descent work?\n\nAssistant: Updating parameters in the opposite direction of the gradient is how gradient descent works.",
4+
"rejected": "\n\nHuman: How does gradient descent work?\n\nAssistant: Climbing up hills to find local maxima is how gradient descent works."
5+
},
6+
{
7+
"chosen": "\n\nHuman: What is a neural network?\n\nAssistant: Interconnected nodes that learn representations from data form a neural network.",
8+
"rejected": "\n\nHuman: What is a neural network?\n\nAssistant: Storing tabular data in a designed schema is the purpose of a neural network."
9+
},
10+
{
11+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
12+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
13+
},
14+
{
15+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
16+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
17+
},
18+
{
19+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
20+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
21+
},
22+
{
23+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
24+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
25+
},
26+
{
27+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
28+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
29+
},
30+
{
31+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
32+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
33+
},
34+
{
35+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
36+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
37+
},
38+
{
39+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
40+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
41+
},
42+
{
43+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
44+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
45+
},
46+
{
47+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
48+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
49+
},
50+
{
51+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
52+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
53+
},
54+
{
55+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
56+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
57+
},
58+
{
59+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
60+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
61+
},
62+
{
63+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
64+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
65+
},
66+
{
67+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
68+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
69+
},
70+
{
71+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
72+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
73+
},
74+
{
75+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
76+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
77+
},
78+
{
79+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
80+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
81+
},
82+
{
83+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
84+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
85+
},
86+
{
87+
"chosen": "\n\nHuman: What is preference optimization?\n\nAssistant: Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
88+
"rejected": "\n\nHuman: What is preference optimization?\n\nAssistant: Database operations to choose preferred options are called preference optimization."
89+
}
90+
]

‎tests/assets/local_datasets/dpo/dpo_3_column_dataset.json‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
[
2+
{
3+
"prompt": "How does gradient descent work?",
4+
"chosen": "Updating parameters in the opposite direction of the gradient is how gradient descent works.",
5+
"rejected": "Climbing up hills to find local maxima is how gradient descent works."
6+
},
7+
{
8+
"prompt": "What is a neural network?",
9+
"chosen": "Interconnected nodes that learn representations from data form a neural network.",
10+
"rejected": "Storing tabular data in a designed schema is the purpose of a neural network."
11+
},
12+
{
13+
"prompt": "What is preference optimization?",
14+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
15+
"rejected": "Database operations to choose preferred options are called preference optimization."
16+
},
17+
{
18+
"prompt": "What is preference optimization?",
19+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
20+
"rejected": "Database operations to choose preferred options are called preference optimization."
21+
},
22+
{
23+
"prompt": "What is preference optimization?",
24+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
25+
"rejected": "Database operations to choose preferred options are called preference optimization."
26+
},
27+
{
28+
"prompt": "What is preference optimization?",
29+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
30+
"rejected": "Database operations to choose preferred options are called preference optimization."
31+
},
32+
{
33+
"prompt": "What is preference optimization?",
34+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
35+
"rejected": "Database operations to choose preferred options are called preference optimization."
36+
},
37+
{
38+
"prompt": "What is preference optimization?",
39+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
40+
"rejected": "Database operations to choose preferred options are called preference optimization."
41+
},
42+
{
43+
"prompt": "What is preference optimization?",
44+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
45+
"rejected": "Database operations to choose preferred options are called preference optimization."
46+
},
47+
{
48+
"prompt": "What is preference optimization?",
49+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
50+
"rejected": "Database operations to choose preferred options are called preference optimization."
51+
},
52+
{
53+
"prompt": "What is preference optimization?",
54+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
55+
"rejected": "Database operations to choose preferred options are called preference optimization."
56+
},
57+
{
58+
"prompt": "What is preference optimization?",
59+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
60+
"rejected": "Database operations to choose preferred options are called preference optimization."
61+
},
62+
{
63+
"prompt": "What is preference optimization?",
64+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
65+
"rejected": "Database operations to choose preferred options are called preference optimization."
66+
},
67+
{
68+
"prompt": "What is preference optimization?",
69+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
70+
"rejected": "Database operations to choose preferred options are called preference optimization."
71+
},
72+
{
73+
"prompt": "What is preference optimization?",
74+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
75+
"rejected": "Database operations to choose preferred options are called preference optimization."
76+
},
77+
{
78+
"prompt": "What is preference optimization?",
79+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
80+
"rejected": "Database operations to choose preferred options are called preference optimization."
81+
},
82+
{
83+
"prompt": "What is preference optimization?",
84+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
85+
"rejected": "Database operations to choose preferred options are called preference optimization."
86+
},
87+
{
88+
"prompt": "What is preference optimization?",
89+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
90+
"rejected": "Database operations to choose preferred options are called preference optimization."
91+
},
92+
{
93+
"prompt": "What is preference optimization?",
94+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
95+
"rejected": "Database operations to choose preferred options are called preference optimization."
96+
},
97+
{
98+
"prompt": "What is preference optimization?",
99+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
100+
"rejected": "Database operations to choose preferred options are called preference optimization."
101+
},
102+
{
103+
"prompt": "What is preference optimization?",
104+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
105+
"rejected": "Database operations to choose preferred options are called preference optimization."
106+
},
107+
{
108+
"prompt": "What is preference optimization?",
109+
"chosen": "Aligning LLMs using pairs of chosen and rejected responses is called preference optimization.",
110+
"rejected": "Database operations to choose preferred options are called preference optimization."
111+
}
112+
]

0 commit comments

Comments
 (0)