mlx-examples
mlx-examples copied to clipboard
Extension of lora.py for supervised ML of configurable dataset formats with YAML-based configuration of parameters
Based on lora.py and meant to perform supervised instruction fine-tuning for the various supported models, separating the model-specific logic from the common training and supporting logic (and management of training configuration parameters).
It breaks out argument parameters into a YAML file and allows arbitrary training data (and prompt) formats
A configuration .yaml file (the only command line argument) is expected to be in the following format:
parameters:
model: "..."
num_tokens: 100
write_every: 1
temp: 0.8
train: true
[..]
Each entry under parameters is the default argument name of the command-line arguments originally provided to lora.py. The default values are the same as the original argument parameters they are based on.
An epoch parameter was added, determining the number of iterations if provided (the number needed for a full pass of the data, i.e., an epoch). If the value is -1, it is ignored, and the iters parameter is used as before. If iters is -1 then 1 epoch is performed.
A module for a particular prompt syntax or training dataset format can be defined with a class that specializes TrainingRecordHandler, and provides an instance of it to main, along with a function for getting model and tokenizer instances, and another to generate from a given prompt using the tokenizer. Otherwise, the iterating batch implementation calculates (Q)LoRa finetuning with configurable parameters.
It needs to define its own get_input and get_output methods, which take a Python dictionary and return an instruction input and output as strings, respectively, and with the appropriate prompt formatting.
This allows handling arbitrary dataset JSON formats and prompt formats separate from the training.
Will probably need to reconcile with #213 once that PR is merged to main
See mistral_supervised.py for an example.
Checklist
Put an x in the boxes that apply.
- [x] I have read the CONTRIBUTING document
- [x] I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
Another pass at separating model-specific bits from training logic. Still keeping an eye on #213 to see if there is any synergy
Just pushed (proposed) final version of #213 . Take a look and let me know how I can help utilize our changes together!
Just pushed (proposed) final version of #213 . Take a look and let me know how I can help utilize our changes together!
That would be fantastic! Sorry I only just saw this. However, I see this comment from today and will be consolidating #337 merges into this PR.
I still think a separate, configuration-based Lora example would be handy even just for the purposes of Lora hacking on this framework. Later this week, I'm planning on also adding the writing of loss/validation structured data for plotting purposes (perhaps as a separate PR). Let me know your updated thoughts regarding this and #213
Amazing job @chimezie 🚀