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>
124 lines
4.0 KiB
Python
124 lines
4.0 KiB
Python
"""
|
|
e2e tests for kd trainer support in Axolotl
|
|
"""
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from e2e.utils import check_tensorboard, require_torch_2_5_1
|
|
|
|
from axolotl.cli.args import TrainerCliArgs
|
|
from axolotl.common.datasets import load_datasets
|
|
from axolotl.train import train
|
|
from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
|
|
from axolotl.utils.dict import DictDefault
|
|
|
|
|
|
@pytest.fixture(name="kd_min_cfg")
|
|
def min_cfg(temp_dir):
|
|
return {
|
|
"base_model": "osllmai-community/Llama-3.2-1B",
|
|
"tokenizer_config": "axolotl-ai-co/Llama-3.3-70B-Instruct-tokenizer",
|
|
"plugins": [
|
|
"axolotl.integrations.kd.KDPlugin",
|
|
"axolotl.integrations.liger.LigerPlugin",
|
|
],
|
|
"liger_rms_norm": True,
|
|
"liger_glu_activation": True,
|
|
"torch_compile": True,
|
|
"chat_template": "llama3",
|
|
"kd_trainer": True,
|
|
"kd_ce_alpha": 0.1,
|
|
"kd_alpha": 0.9,
|
|
"kd_temperature": 1.0,
|
|
"dataloader_prefetch_factor": 8,
|
|
"dataloader_num_workers": 4,
|
|
"dataloader_pin_memory": True,
|
|
"datasets": [
|
|
{
|
|
"path": "axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample",
|
|
"type": "axolotl.integrations.kd.chat_template",
|
|
"field_messages": "messages_combined",
|
|
"split": "train",
|
|
"logprobs_field": "llm_text_generation_vllm_logprobs",
|
|
"temperature": 1.0,
|
|
"preprocess_shards": 2,
|
|
},
|
|
],
|
|
"val_set_size": 0.0,
|
|
"sequence_len": 2048,
|
|
"sample_packing": True,
|
|
"pad_to_sequence_len": True,
|
|
"gradient_accumulation_steps": 2,
|
|
"micro_batch_size": 1,
|
|
"num_epochs": 1,
|
|
"optimizer": "adamw_8bit",
|
|
"lr_scheduler": "cosine",
|
|
"learning_rate": 0.00001,
|
|
"bf16": "auto",
|
|
"gradient_checkpointing": True,
|
|
"flash_attention": True,
|
|
"special_tokens": {
|
|
"pad_token": "<|end_of_text|>",
|
|
"eos_token": "<|eot_id|>",
|
|
},
|
|
"max_steps": 5,
|
|
"output_dir": temp_dir,
|
|
"save_safetensors": True,
|
|
"use_tensorboard": True,
|
|
}
|
|
|
|
|
|
class TestKnowledgeDistillation:
|
|
"""
|
|
Test case for Knowledge Distillation
|
|
"""
|
|
|
|
# While this will run on torch 2.4.x without torch_compile enabled
|
|
# the VRAM requirement is higher than what is available in CI
|
|
@require_torch_2_5_1
|
|
def test_llama_kd(self, temp_dir, kd_min_cfg):
|
|
cfg = DictDefault(kd_min_cfg)
|
|
# pylint: disable=duplicate-code
|
|
cfg = validate_config(cfg)
|
|
prepare_plugins(cfg)
|
|
normalize_config(cfg)
|
|
cli_args = TrainerCliArgs()
|
|
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
|
|
|
|
train(cfg=cfg, dataset_meta=dataset_meta)
|
|
assert (Path(temp_dir) / "model.safetensors").exists()
|
|
check_tensorboard(
|
|
temp_dir + "/runs", "train/loss", 1.0, "Train Loss is too high"
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"load_in_8bit",
|
|
[True, False],
|
|
)
|
|
def test_llama_lora_kd(self, temp_dir, kd_min_cfg, load_in_8bit):
|
|
cfg = DictDefault(
|
|
{
|
|
"load_in_8bit": load_in_8bit,
|
|
"torch_compile": False,
|
|
"adapter": "lora",
|
|
"peft_use_dora": True,
|
|
"lora_target_linear": True,
|
|
"lora_r": 16,
|
|
"lora_alpha": 32,
|
|
"lora_dropout": 0.0,
|
|
}
|
|
| kd_min_cfg
|
|
)
|
|
# pylint: disable=duplicate-code
|
|
cfg = validate_config(cfg)
|
|
prepare_plugins(cfg)
|
|
normalize_config(cfg)
|
|
cli_args = TrainerCliArgs()
|
|
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
|
|
|
|
train(cfg=cfg, dataset_meta=dataset_meta)
|
|
assert (Path(temp_dir) / "adapter_model.safetensors").exists()
|
|
check_tensorboard(
|
|
temp_dir + "/runs", "train/loss", 1.0, "Train Loss is too high"
|
|
)
|