[WIP] Add SPIN trainer #1344
[WIP] Add SPIN trainer #1344
Conversation
| preprocessing_num_workers: int = field( | ||
| default=12, metadata={"help": "The number of processes to use for the preprocessing."} | ||
| ) | ||
| do_generate: bool = field(default=False) |
There was a problem hiding this comment.
These args could arguably live in the SPINConfig
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| import torch | ||
| from datasets import Dataset, load_dataset | ||
| from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, TrainingArguments | ||
|
|
||
| from trl import DPOTrainer, ModelConfig, get_kbit_device_map, get_peft_config, get_quantization_config | ||
|
|
||
|
|
||
| datasets.disable_caching() |
There was a problem hiding this comment.
What is this? Maybe we should do load_from_cache=False?
| if script_args.do_generate: | ||
| print(f"Generating completions for {len(prompt_train_ds)} training examples") | ||
| train_completions = spin_trainer.generate(prompt_train_ds, "prompt", generation_config, batch_size=16) | ||
| test_completions = spin_trainer.generate(prompt_test_ds, "prompt", generation_config, batch_size=16) |
There was a problem hiding this comment.
If the generation only occurs in the end, wdyt about doing a subprocess call on a vllm/tgi script to generate instead?
| chosen_tokens = self.tokenizer(chosen, add_special_tokens=False) | ||
| rejected_tokens = self.tokenizer(rejected, add_special_tokens=False) | ||
| prompt_tokens = self.tokenizer(prompt, add_special_tokens=False) |
There was a problem hiding this comment.
This feels quite similar to the logic in DPOTrainer. Maybe unify the logic somehow (or use similar terminology).
trl/trl/trainer/dpo_trainer.py
Lines 721 to 729 in b32656f
Not necessary but we could also benefit from doing a filediff between SPIN and DPO and try to make the lines of code differences minimal.
|
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread. |
Implements the Self-Play fIne-tuNing (SPIN) algorithm from: https://arxiv.org/abs/2401.01335
TODO
RuntimeError: 'weight' must be 2-Derrors)