b194e17c28
* feat: add config for optional parameters in a chat message * chore: cleanup * chore: fix nits and add light docs * docs: update docs/dataset-formats/conversation.qmd Co-authored-by: NanoCode012 <kevinvong@rocketmail.com> * feat: configurable message mappings, jinja template analyzer * chore: handle bradley terry * docs: update docs * refactor: change order of mappings, improve message transform * refactor: make chat awware of property mappings * chore: remove .python-version * chore: revert change * chore: add dataset validation to tests where appropriate * chore: add dataset validation to tests where appropriate * chore: clean up handling of ds_cfg * chore: recursively serialize config * make sure to use the return value from validate_config * DefaultDict pickle/unpickle fix * fix super call for override * refactor: message fields * chore: empty commit * tests: validate config before using * chore: add config validation to all e2e tests * chore: add unneeded logging * chore: add missed config validation * chore: pass field_messages to prompter * test: fix borked test * chore: remove uninteded file * chore: add deprecation warning and update chat_datasets script * chore: lint * refactor: message fields * feat: update axolotlinputconfig and test_models - add configdict import in axolotl/utils/config/models/input/v0_4_1/__init__.py - remove unnecessary line breaks in sftdataset, dpodataset, ktodataset, stepwisesuperviseddataset classes - update model_dump method in axolotlinputconfig to exclude none values - correct typo in test_models.py comment * feat: simplify dpodataset and ktodataset classes in config models removed several optional fields from dpodataset and ktodataset classes in axolotl/utils/config/models/input/v0_4_1. this simplifies the configuration subsets for these datasets. * feat: improve readability and structure in dataset configuration models this commit enhances the readability and structure of the dataset configuration models in the `axolotl/utils/config/models/input/v0_4_1` module. it removes unused `configdict` import and adds line breaks to separate class definitions for better clarity. additionally, a minor documentation fix is included to ensure a newline at the end of the `stepwise_supervised.qmd` file. * feat: change log level from info to debug in chattemplatestrategy * feat(prompt_strategies): refactor chattemplateprompter and chattemplatestrategy - Make `chat_template` a required parameter in `ChatTemplatePrompter` constructor - Add default value for `message_property_mappings` in `ChatTemplatePrompter` constructor - Add `messages_array_name` property to `ChatTemplatePrompter` - Change `processor` type to Optional in `ChatTemplatePrompter` - Add TypeError check for `processor` in `ChatTemplatePrompter.build_prompt` - Remove `_messages` property from `ChatTemplateStrategy` - Make `prompter` a required parameter and add type hint in `ChatTemplateStrategy` constructor - Remove `messages` getter and setter from `ChatTemplateStrategy` - Use `prompter.messages_array_name` in `ChatTemplateStrategy.get_conversation_thread` - Remove condition to set `messages` field in `load` function * feat(tests/utils): ignore type check in load_model call in test_models.py * feat: improve type handling and test structure in chat templates - Add return type hint for `get_chat_template` function in `chat_templates.py` - Remove unnecessary assignment of `strategy.messages` in several test cases - Add `messages_array_name` parameter to various test configurations in `test_chat_templates.py` and `test_chat_templates_advanced.py` - Remove redundant `strategy.messages` assignment in `test_chat_templates_advanced.py` * feat(axolotl): enhance chat strategy with datasetconfig support This commit introduces support for DatasetConfig in the ChatTemplateStrategy. It also refines the strategy loader to handle different types of ds_cfg inputs and improves the clarity of the code by formatting and reordering. The key changes include: - Importing Union from typing and BaseModel from pydantic. - Adding DatasetConfig as an optional type for ds_cfg in StrategyLoader. - Adjusting the handling of ds_cfg in StrategyLoader to account for BaseModel instances. - Refactoring the prompter_params and strategy_params for better readability. - Changing the reference from prompt[self.messages] to prompt[self.prompter.messages_array_name] in the is_prompt_batched method. * feat: update message handling in btchattemplatestrategy * Replace `self.messages` with direct string references to "chosen_messages" and "rejected_messages" * Append system, user, and assistant content directly to "chosen_messages" and "rejected_messages" * Add a new attribute "messages_array_name" to the `load` function parameters * Remove the conditional attribute assignment for "field_messages" in the `load` function * feat: add config validation in test_kd.py - Import `validate_config` from `axolotl.utils.config` - Validate the configuration in `test_llama_kd` and another function in `TestKnowledgeDistillation` class * feat: enhance config validation and capabilities handling * Import `EnvCapabilities` and `GPUCapabilities` from `axolotl.utils.config.models.internals` * Update `validate_config` function to create `KTODataset` and `SFTDataset` instances using `dict(ds_cfg)` * Replace `capabilities` and `env_capabilities` with instances of `GPUCapabilities` and `EnvCapabilities` respectively in `AxolotlConfigWCapabilities` model dump * feat: update config validation in axolotl utils - Remove import of `EnvCapabilities` and `GPUCapabilities` from `axolotl.utils.config.models.internals` - Update `validate_config` function to use `capabilities` and `env_capabilities` directly instead of creating new instances of `GPUCapabilities` and `EnvCapabilities` * feat: refactor strategyloader in chat_template.py - Extracted the creation of strategy parameters into a separate function, `_get_strategy_params(cfg, dataset_config)` - Created a new function, `_get_strategy_cls()`, to obtain the strategy class - Replaced `ChatTemplateStrategy` with `strategy_cls` for strategy instantiation * trigger CI * chore: revert dataset config changes for kto/dpo * subject: refactor: rename 'messages_array_name' to 'field_messages' Body: - Renamed 'messages_array_name' to 'field_messages' in 'ChatTemplatePrompter' class and its usages in 'chat_template.py' - Updated 'load' function in 'bradley_terry/chat_template.py' to reflect the change - Adjusted 'get_chat_template_msg_variables' and 'get_message_vars' methods in 'jinja_template_analyzer.py' to use the new variable name - Modified 'StrategyLoader' in 'chat_template.py' to use 'field_messages' - Updated tests in 'test_chat_templates.py' and 'test_chat_templates_advanced.py' to use 'field_messages' instead of 'messages_array_name' * feat: refactor prompt strategies and update config models * Remove redundant 'return None' in `axolotl/prompt_strategies/__init__.py` * Simplify message handling in `axolotl/prompt_strategies/bradley_terry/chat_template.py` by using a single 'messages' list instead of separate 'chosen_messages' and 'rejected_messages' lists * Update default 'message_property_mappings' in `axolotl/prompt_strategies/bradley_terry/chat_template.py` * Add 'field_messages' field to `axolotl/utils/config/models/input/v0_4_1/__init__.py` configuration model * chore: remove unused input * chore: remove redundant type ignore * fix: remove old configs and update examples * fix: type check * fix: remove loading old config in ChatMessage * fix: update faq with potential new undefinederror * fix: add debug if property mapped is not found * chore: improve explanation for unmapped properties * fix: update docs with new config * chore: add note for deprecation config and del old config from dict --------- Co-authored-by: NanoCode012 <kevinvong@rocketmail.com> Co-authored-by: Wing Lian <wing@axolotl.ai> Co-authored-by: NanoCode012 <nano@axolotl.ai>
376 lines
13 KiB
Python
376 lines
13 KiB
Python
"""
|
|
E2E tests for lora llama
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from axolotl.cli.args import TrainerCliArgs
|
|
from axolotl.common.datasets import load_preference_datasets
|
|
from axolotl.train import train
|
|
from axolotl.utils.config import normalize_config, validate_config
|
|
from axolotl.utils.dict import DictDefault
|
|
|
|
from .utils import check_model_output_exists, with_temp_dir
|
|
|
|
LOG = logging.getLogger("axolotl.tests.e2e")
|
|
os.environ["WANDB_DISABLED"] = "true"
|
|
|
|
|
|
class TestDPOLlamaLora(unittest.TestCase):
|
|
"""
|
|
Test case for DPO Llama models using LoRA
|
|
"""
|
|
|
|
@with_temp_dir
|
|
def test_dpo_lora(self, temp_dir):
|
|
# pylint: disable=duplicate-code
|
|
cfg = DictDefault(
|
|
{
|
|
"base_model": "JackFram/llama-68m",
|
|
"tokenizer_type": "LlamaTokenizer",
|
|
"sequence_len": 1024,
|
|
"load_in_8bit": True,
|
|
"adapter": "lora",
|
|
"lora_r": 64,
|
|
"lora_alpha": 32,
|
|
"lora_dropout": 0.1,
|
|
"lora_target_linear": True,
|
|
"special_tokens": {},
|
|
"rl": "dpo",
|
|
"datasets": [
|
|
{
|
|
"path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
|
|
"type": "chatml.ultra",
|
|
"split": "train",
|
|
},
|
|
],
|
|
"num_epochs": 1,
|
|
"micro_batch_size": 4,
|
|
"gradient_accumulation_steps": 1,
|
|
"output_dir": temp_dir,
|
|
"learning_rate": 0.00001,
|
|
"optimizer": "paged_adamw_8bit",
|
|
"lr_scheduler": "cosine",
|
|
"max_steps": 20,
|
|
"save_steps": 10,
|
|
"warmup_steps": 5,
|
|
"gradient_checkpointing": True,
|
|
"gradient_checkpointing_kwargs": {"use_reentrant": True},
|
|
}
|
|
)
|
|
|
|
cfg = validate_config(cfg)
|
|
normalize_config(cfg)
|
|
cli_args = TrainerCliArgs()
|
|
dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
|
|
|
|
train(cfg=cfg, dataset_meta=dataset_meta)
|
|
check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
|
|
|
|
@with_temp_dir
|
|
def test_dpo_nll_lora(self, temp_dir):
|
|
# pylint: disable=duplicate-code
|
|
cfg = DictDefault(
|
|
{
|
|
"base_model": "JackFram/llama-68m",
|
|
"tokenizer_type": "LlamaTokenizer",
|
|
"sequence_len": 1024,
|
|
"load_in_8bit": True,
|
|
"adapter": "lora",
|
|
"lora_r": 64,
|
|
"lora_alpha": 32,
|
|
"lora_dropout": 0.1,
|
|
"lora_target_linear": True,
|
|
"special_tokens": {},
|
|
"rl": "dpo",
|
|
"rpo_alpha": 0.5,
|
|
"datasets": [
|
|
{
|
|
"path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
|
|
"type": "chatml.ultra",
|
|
"split": "train",
|
|
},
|
|
],
|
|
"num_epochs": 1,
|
|
"micro_batch_size": 4,
|
|
"gradient_accumulation_steps": 1,
|
|
"output_dir": temp_dir,
|
|
"learning_rate": 0.00001,
|
|
"optimizer": "paged_adamw_8bit",
|
|
"lr_scheduler": "cosine",
|
|
"max_steps": 20,
|
|
"save_steps": 10,
|
|
"warmup_steps": 5,
|
|
"gradient_checkpointing": True,
|
|
"gradient_checkpointing_kwargs": {"use_reentrant": True},
|
|
}
|
|
)
|
|
|
|
cfg = validate_config(cfg)
|
|
normalize_config(cfg)
|
|
cli_args = TrainerCliArgs()
|
|
dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
|
|
|
|
train(cfg=cfg, dataset_meta=dataset_meta)
|
|
check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
|
|
|
|
@with_temp_dir
|
|
def test_dpo_use_weighting(self, temp_dir):
|
|
# pylint: disable=duplicate-code
|
|
cfg = DictDefault(
|
|
{
|
|
"base_model": "JackFram/llama-68m",
|
|
"tokenizer_type": "LlamaTokenizer",
|
|
"sequence_len": 1024,
|
|
"load_in_8bit": True,
|
|
"adapter": "lora",
|
|
"lora_r": 64,
|
|
"lora_alpha": 32,
|
|
"lora_dropout": 0.1,
|
|
"lora_target_linear": True,
|
|
"special_tokens": {},
|
|
"rl": "dpo",
|
|
"dpo_use_weighting": True,
|
|
"datasets": [
|
|
{
|
|
"path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
|
|
"type": "chatml.ultra",
|
|
"split": "train",
|
|
},
|
|
],
|
|
"num_epochs": 1,
|
|
"micro_batch_size": 4,
|
|
"gradient_accumulation_steps": 1,
|
|
"output_dir": temp_dir,
|
|
"learning_rate": 0.00001,
|
|
"optimizer": "paged_adamw_8bit",
|
|
"lr_scheduler": "cosine",
|
|
"max_steps": 20,
|
|
"save_steps": 10,
|
|
"warmup_steps": 5,
|
|
"gradient_checkpointing": True,
|
|
"gradient_checkpointing_kwargs": {"use_reentrant": True},
|
|
}
|
|
)
|
|
|
|
cfg = validate_config(cfg)
|
|
normalize_config(cfg)
|
|
cli_args = TrainerCliArgs()
|
|
dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
|
|
|
|
train(cfg=cfg, dataset_meta=dataset_meta)
|
|
check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
|
|
|
|
@pytest.mark.skip("kto_pair no longer supported in trl")
|
|
@with_temp_dir
|
|
def test_kto_pair_lora(self, temp_dir):
|
|
# pylint: disable=duplicate-code
|
|
cfg = DictDefault(
|
|
{
|
|
"base_model": "JackFram/llama-68m",
|
|
"tokenizer_type": "LlamaTokenizer",
|
|
"sequence_len": 1024,
|
|
"load_in_8bit": True,
|
|
"adapter": "lora",
|
|
"lora_r": 64,
|
|
"lora_alpha": 32,
|
|
"lora_dropout": 0.1,
|
|
"lora_target_linear": True,
|
|
"special_tokens": {},
|
|
"rl": "kto_pair",
|
|
"datasets": [
|
|
{
|
|
"path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
|
|
"type": "chatml.ultra",
|
|
"split": "train",
|
|
},
|
|
],
|
|
"num_epochs": 1,
|
|
"micro_batch_size": 4,
|
|
"gradient_accumulation_steps": 1,
|
|
"output_dir": temp_dir,
|
|
"learning_rate": 0.00001,
|
|
"optimizer": "paged_adamw_8bit",
|
|
"lr_scheduler": "cosine",
|
|
"max_steps": 20,
|
|
"save_steps": 10,
|
|
"warmup_steps": 5,
|
|
"gradient_checkpointing": True,
|
|
"gradient_checkpointing_kwargs": {"use_reentrant": True},
|
|
}
|
|
)
|
|
|
|
cfg = validate_config(cfg)
|
|
normalize_config(cfg)
|
|
cli_args = TrainerCliArgs()
|
|
dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
|
|
|
|
train(cfg=cfg, dataset_meta=dataset_meta)
|
|
check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
|
|
|
|
@with_temp_dir
|
|
def test_ipo_lora(self, temp_dir):
|
|
# pylint: disable=duplicate-code
|
|
cfg = DictDefault(
|
|
{
|
|
"base_model": "JackFram/llama-68m",
|
|
"tokenizer_type": "LlamaTokenizer",
|
|
"sequence_len": 1024,
|
|
"load_in_8bit": True,
|
|
"adapter": "lora",
|
|
"lora_r": 64,
|
|
"lora_alpha": 32,
|
|
"lora_dropout": 0.1,
|
|
"lora_target_linear": True,
|
|
"special_tokens": {},
|
|
"rl": "ipo",
|
|
"datasets": [
|
|
{
|
|
"path": "arcee-ai/distilabel-intel-orca-dpo-pairs-binarized",
|
|
"type": "chatml.ultra",
|
|
"split": "train",
|
|
},
|
|
],
|
|
"num_epochs": 1,
|
|
"micro_batch_size": 4,
|
|
"gradient_accumulation_steps": 1,
|
|
"output_dir": temp_dir,
|
|
"learning_rate": 0.00001,
|
|
"optimizer": "paged_adamw_8bit",
|
|
"lr_scheduler": "cosine",
|
|
"max_steps": 20,
|
|
"save_steps": 10,
|
|
"warmup_steps": 5,
|
|
"gradient_checkpointing": True,
|
|
"gradient_checkpointing_kwargs": {"use_reentrant": True},
|
|
}
|
|
)
|
|
|
|
cfg = validate_config(cfg)
|
|
normalize_config(cfg)
|
|
cli_args = TrainerCliArgs()
|
|
dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
|
|
|
|
train(cfg=cfg, dataset_meta=dataset_meta)
|
|
check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
|
|
|
|
@with_temp_dir
|
|
def test_orpo_lora(self, temp_dir):
|
|
# pylint: disable=duplicate-code
|
|
cfg = DictDefault(
|
|
{
|
|
"base_model": "JackFram/llama-68m",
|
|
"tokenizer_type": "LlamaTokenizer",
|
|
"sequence_len": 1024,
|
|
"load_in_8bit": True,
|
|
"adapter": "lora",
|
|
"lora_r": 64,
|
|
"lora_alpha": 32,
|
|
"lora_dropout": 0.1,
|
|
"lora_target_linear": True,
|
|
"special_tokens": {},
|
|
"rl": "orpo",
|
|
"orpo_alpha": 0.1,
|
|
"remove_unused_columns": False,
|
|
"chat_template": "chatml",
|
|
"datasets": [
|
|
{
|
|
"path": "argilla/distilabel-capybara-dpo-7k-binarized",
|
|
"type": "chat_template.argilla",
|
|
"split": "train",
|
|
},
|
|
],
|
|
"num_epochs": 1,
|
|
"micro_batch_size": 4,
|
|
"gradient_accumulation_steps": 1,
|
|
"output_dir": temp_dir,
|
|
"learning_rate": 0.00001,
|
|
"optimizer": "paged_adamw_8bit",
|
|
"lr_scheduler": "cosine",
|
|
"max_steps": 20,
|
|
"save_steps": 10,
|
|
"warmup_steps": 5,
|
|
"gradient_checkpointing": True,
|
|
"gradient_checkpointing_kwargs": {"use_reentrant": True},
|
|
}
|
|
)
|
|
|
|
cfg = validate_config(cfg)
|
|
normalize_config(cfg)
|
|
cli_args = TrainerCliArgs()
|
|
dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
|
|
|
|
train(cfg=cfg, dataset_meta=dataset_meta)
|
|
check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
|
|
|
|
@pytest.mark.skip(reason="Fix the implementation")
|
|
@with_temp_dir
|
|
def test_kto_lora(self, temp_dir):
|
|
# pylint: disable=duplicate-code
|
|
cfg = DictDefault(
|
|
{
|
|
"base_model": "JackFram/llama-68m",
|
|
"tokenizer_type": "LlamaTokenizer",
|
|
"sequence_len": 1024,
|
|
"load_in_8bit": True,
|
|
"adapter": "lora",
|
|
"lora_r": 64,
|
|
"lora_alpha": 32,
|
|
"lora_dropout": 0.1,
|
|
"lora_target_linear": True,
|
|
"special_tokens": {},
|
|
"rl": "kto",
|
|
"rl_beta": 0.5,
|
|
"kto_desirable_weight": 1.0,
|
|
"kto_undesirable_weight": 1.0,
|
|
"remove_unused_columns": False,
|
|
"datasets": [
|
|
# {
|
|
# "path": "argilla/kto-mix-15k",
|
|
# "type": "chatml.argilla_chat",
|
|
# "split": "train",
|
|
# },
|
|
{
|
|
"path": "argilla/ultrafeedback-binarized-preferences-cleaned-kto",
|
|
"type": "chatml.ultra",
|
|
"split": "train",
|
|
},
|
|
# {
|
|
# "path": "argilla/kto-mix-15k",
|
|
# "type": "llama3.argilla_chat",
|
|
# "split": "train",
|
|
# },
|
|
{
|
|
"path": "argilla/ultrafeedback-binarized-preferences-cleaned-kto",
|
|
"type": "llama3.ultra",
|
|
"split": "train",
|
|
},
|
|
],
|
|
"num_epochs": 1,
|
|
"micro_batch_size": 4,
|
|
"gradient_accumulation_steps": 1,
|
|
"output_dir": temp_dir,
|
|
"learning_rate": 0.00001,
|
|
"optimizer": "paged_adamw_8bit",
|
|
"lr_scheduler": "cosine",
|
|
"max_steps": 20,
|
|
"save_steps": 10,
|
|
"warmup_steps": 5,
|
|
"gradient_checkpointing": True,
|
|
"gradient_checkpointing_kwargs": {"use_reentrant": True},
|
|
}
|
|
)
|
|
|
|
cfg = validate_config(cfg)
|
|
normalize_config(cfg)
|
|
cli_args = TrainerCliArgs()
|
|
dataset_meta = load_preference_datasets(cfg=cfg, cli_args=cli_args)
|
|
|
|
train(cfg=cfg, dataset_meta=dataset_meta)
|
|
check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
|