Skip to content

Commit

Permalink
re-enable DPO for tests in modal ci (#1374)
Browse files Browse the repository at this point in the history
* re-enable DPO for tests in modal ci

* workaround for training args

* don't mixin AxolotlTrainingArguments

* fix mixin order so MRO doesn't result in

 TypeError: non-default argument follows default argument error

* use smaller datasets for dpo tests
  • Loading branch information
winglian authored Jun 3, 2024
1 parent 5cde065 commit 1f151c0
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 13 deletions.
16 changes: 11 additions & 5 deletions src/axolotl/prompt_strategies/orpo/chat_template.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,9 @@ def get_chosen_conversation_thread(self, prompt) -> MessageList:
messages: List[Message] = []
if system := prompt.get("system", None):
messages.append(Message(role="system", content=system, label=False))
messages.append(Message(role="user", content=prompt["prompt"], label=False))
messages.append(
Message(role="user", content=prompt["chosen"][0]["content"], label=False)
)
messages.append(
Message(
role="assistant", content=prompt["chosen"][1]["content"], label=True
Expand All @@ -70,7 +72,9 @@ def get_rejected_conversation_thread(self, prompt) -> MessageList:
messages: List[Message] = []
if system := prompt.get("system", None):
messages.append(Message(role="system", content=system, label=False))
messages.append(Message(role="user", content=prompt["prompt"], label=False))
messages.append(
Message(role="user", content=prompt["rejected"][0]["content"], label=False)
)
messages.append(
Message(
role="assistant", content=prompt["rejected"][1]["content"], label=True
Expand Down Expand Up @@ -152,8 +156,8 @@ def __init__(
def tokenize_prompt(self, prompt):
# pass the rejected prompt/row to the Prompter to get the formatted prompt
prompt_len = 0
rejected_message_list = self.dataset_parser.get_rejected_conversation_thread(
prompt
rejected_message_list: MessageList = (
self.dataset_parser.get_rejected_conversation_thread(prompt)
)
input_ids = []
labels = []
Expand All @@ -174,7 +178,9 @@ def tokenize_prompt(self, prompt):
rejected_input_ids = input_ids
rejected_labels = labels
# pass the chosen prompt/row to the Prompter to get the formatted prompt
chosen_message_list = self.dataset_parser.get_chosen_conversation_thread(prompt)
chosen_message_list: MessageList = (
self.dataset_parser.get_chosen_conversation_thread(prompt)
)
input_ids = []
labels = []
for _, (part, label) in enumerate(
Expand Down
16 changes: 8 additions & 8 deletions tests/e2e/test_dpo.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
os.environ["WANDB_DISABLED"] = "true"


@pytest.mark.skip(reason="doesn't seem to work on modal")
class TestDPOLlamaLora(unittest.TestCase):
"""
Test case for DPO Llama models using LoRA
Expand All @@ -45,8 +44,8 @@ def test_dpo_lora(self, temp_dir):
"rl": "dpo",
"datasets": [
{
"path": "Intel/orca_dpo_pairs",
"type": "chatml.intel",
"path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
"type": "chatml.ultra",
"split": "train",
},
],
Expand Down Expand Up @@ -89,8 +88,8 @@ def test_kto_pair_lora(self, temp_dir):
"rl": "kto_pair",
"datasets": [
{
"path": "Intel/orca_dpo_pairs",
"type": "chatml.intel",
"path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
"type": "chatml.ultra",
"split": "train",
},
],
Expand Down Expand Up @@ -133,8 +132,8 @@ def test_ipo_lora(self, temp_dir):
"rl": "ipo",
"datasets": [
{
"path": "Intel/orca_dpo_pairs",
"type": "chatml.intel",
"path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
"type": "chatml.ultra",
"split": "train",
},
],
Expand Down Expand Up @@ -180,7 +179,7 @@ def test_orpo_lora(self, temp_dir):
"chat_template": "chatml",
"datasets": [
{
"path": "argilla/ultrafeedback-binarized-preferences-cleaned",
"path": "argilla/distilabel-capybara-dpo-7k-binarized",
"type": "chat_template.argilla",
"split": "train",
},
Expand All @@ -206,6 +205,7 @@ def test_orpo_lora(self, temp_dir):
train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()

@pytest.mark.skip(reason="Fix the implementation")
@with_temp_dir
def test_kto_lora(self, temp_dir):
# pylint: disable=duplicate-code
Expand Down

0 comments on commit 1f151c0

Please sign in to comment.