diff --git a/model/supervised_finetuning/README.md b/model/supervised_finetuning/README.md index d5b10e01..387e91e4 100644 --- a/model/supervised_finetuning/README.md +++ b/model/supervised_finetuning/README.md @@ -1,62 +1,18 @@ # Train using supervised examples -Requirements +## Requirements -``` -wandb -evaluate -datasets -transformers -torch -``` +`pip install -r requirements.txt` -Start training reward model +Start training SFT model ```bash -python trainer.py --configs defaults galactica-125 +python trainer.py --configs defaults galactica-125m ``` -## Dataset - -For now we only support webgpt and summary dataset from OpenAI. Once -open-asisstant dataset are available it will be added here. - -## Model - -Normally you should be able to add new models in configs/config.yml - -``` -your-model-name: - learning_rate: 2e-6 - model_name: - weight_decay: 0.01 - max_length: 812 - warmup_steps: 600 - gradient_checkpointing: false - gradient_accumulation_steps: 5 - per_device_train_batch_size: 4 - per_device_eval_batch_size: 4 -``` - -``` -python trainer.py --configs defaults your-model-name -``` - -However, if the model of your choice doesn't have pad_token, eos_token, -sep_token, you have to update utils.py `get_tokenizer` to use the right token. - -## Deepspeed support - -You can edit the configs/zero_config.json and use any stage you wish. The -current config uses zero-stage 3. For more details on how to setup the config -checkout [this page](https://www.deepspeed.ai/tutorials/zero/) - -Once you are satisfy with your deepzero config, you can add --deepspeed flag at -the end to trigger deepspeed - -``` -python trainer.py --configs defaults your-model-name --deepspeed -``` +For `wandb`: update the `entity` argument in `trainer.py`'s call to `wandb.init` +to be your weights and biases username per +[docs](https://docs.wandb.ai/ref/python/init). ## Dataset choices @@ -80,6 +36,74 @@ Currently only these languages are supported via prompt translation: ar,de,fr,en,it,nl,tr,ru,ms,ko,ja,zh ``` +## Dataset sub-sampling + +We can subsample the **training** data by passing either the `fraction` or +`size` argument in the `configs/config.yml` file. Don't forget the additional +colon ":" after the dataset name when doing this. + +Example: + +``` + datasets: + - webgpt: + fraction : 0.05 + - prompt_dialogue: + size : 500 + - adversarial_qa + - trivia_qa_nocontext +``` + +In this example, per epoch we will use: + +- A random 5% of `webgpt`; +- A random 500 examples from `prompt_dialogue`; +- All examples from datasets for which we don't specify the `fraction` or `size` + argument. + +In the above example, per epoch we'll use a different 5% from `webgpt` and a +different 500 examples from `prompt_dialogue`. + +This works with `torch.distributed`. + +## Model + +Normally you should be able to add new models in `configs/config.yml` + +``` +your-model-name: + learning_rate: 2e-6 + model_name: + weight_decay: 0.01 + max_length: 812 + warmup_steps: 600 + gradient_checkpointing: false + gradient_accumulation_steps: 5 + per_device_train_batch_size: 4 + per_device_eval_batch_size: 4 +``` + +``` +python trainer.py --configs defaults your-model-name +``` + +However, if the model of your choice doesn't have `pad_token`, `eos_token`, +`sep_token`, you have to update `get_tokenizer` in `utils.py` to use the right +token. + +## Deepspeed support + +You can edit the configs/zero_config.json and use any stage you wish. The +current config uses zero-stage 3. For more details on how to setup the config +checkout [this page](https://www.deepspeed.ai/tutorials/zero/). + +Once you are satisfy with your deepzero config, you can add --deepspeed flag at +the end to trigger deepspeed + +``` +python trainer.py --configs defaults your-model-name --deepspeed +``` + ## Results Experimental results in wandb @@ -87,7 +111,7 @@ Experimental results in wandb ## TODOS -- decide on a model +- Decide on a model - Merge utils etc with reward model - Casual Modelling for GPT-JT does not leverage the bidirectional mask for the prompt? (https://huggingface.co/togethercomputer/GPT-JT-6B-v1) diff --git a/model/supervised_finetuning/configs/config.yaml b/model/supervised_finetuning/configs/config.yaml index 925b6dda..289e49ad 100644 --- a/model/supervised_finetuning/configs/config.yaml +++ b/model/supervised_finetuning/configs/config.yaml @@ -16,32 +16,29 @@ defaults: eval_accumulation_steps: freeze_layer: datasets: - - webgpt: - fraction : 0.001 - - prompt_dialogue: - size : 10 - - squad_v2: - size : 10 - # - adversarial_qa - # - trivia_qa_nocontext - # - xsum - # - cnn_dailymail + - webgpt # - prompt_dialogue - # - multi_news - # - scitldr - # - soda - # - joke - # - gsm8k - # - dive_mt - # - wmt2019_zh-en - # - wmt2019_ru-en - # - wmt2019_de-en - # - ted_trans_nl-en - # - ted_trans_de-ja - # - instruct_tuning - # - wmt2019_de-en - # - samsum - # - soda_dialogue + - squad_v2 + - adversarial_qa + - trivia_qa_nocontext + - xsum + - cnn_dailymail + - prompt_dialogue + - multi_news + - scitldr + - soda + - joke + - gsm8k + - dive_mt + - wmt2019_zh-en + - wmt2019_ru-en + - wmt2019_de-en + - ted_trans_nl-en + - ted_trans_de-ja + - instruct_tuning + - wmt2019_de-en + - samsum + - soda_dialogue cache_dir: .cache loss_fn: CrossEntropyLoss eval_size: diff --git a/model/supervised_finetuning/custom_datasets/qa_datasets.py b/model/supervised_finetuning/custom_datasets/qa_datasets.py index 2c5c7ee2..206b679c 100644 --- a/model/supervised_finetuning/custom_datasets/qa_datasets.py +++ b/model/supervised_finetuning/custom_datasets/qa_datasets.py @@ -219,7 +219,7 @@ class SODA(Dataset): return pairs - def __init__(self, cache_dir, max_sample_size=10000, input_max_length=1024) -> None: + def __init__(self, cache_dir, input_max_length=1024) -> None: super().__init__() self.pairs = [] @@ -230,9 +230,6 @@ class SODA(Dataset): if len(prompt) < input_max_length: self.pairs.append((prompt, answer)) - if len(self.pairs) > max_sample_size: - break - def __len__(self): return len(self.pairs) diff --git a/model/supervised_finetuning/custom_datasets/translation.py b/model/supervised_finetuning/custom_datasets/translation.py index f9a71a8e..008de751 100644 --- a/model/supervised_finetuning/custom_datasets/translation.py +++ b/model/supervised_finetuning/custom_datasets/translation.py @@ -100,8 +100,6 @@ class WMT2019(TranslationPair): else: # translating in reverse direction source = random.choice(TRANSLATION_PROMPT[src]).format(row[tgt]) self.pairs.append((source, row[src])) - if len(self.pairs) > 100000: - break class DiveMT(TranslationPair): diff --git a/model/supervised_finetuning/requirements.txt b/model/supervised_finetuning/requirements.txt index efe6df89..95e5a472 100644 --- a/model/supervised_finetuning/requirements.txt +++ b/model/supervised_finetuning/requirements.txt @@ -4,7 +4,6 @@ datasets==2.8.0 deepspeed==0.7.7 evaluate==0.4.0 gdown -# mpi4py==3.1.4 nltk==3.8.1 numpy>=1.22.4 py7zr @@ -12,3 +11,4 @@ PyYAML>=6.0 scikit_learn==1.2.0 torch>=1.11.0 transformers==4.25.1 +wandb diff --git a/model/supervised_finetuning/trainer.py b/model/supervised_finetuning/trainer.py index 06ae39db..922be10b 100644 --- a/model/supervised_finetuning/trainer.py +++ b/model/supervised_finetuning/trainer.py @@ -8,8 +8,7 @@ import torch from torch import nn from transformers import PreTrainedModel, Trainer, TrainingArguments from transformers.training_args import OptimizerNames -from utils import (build_train_sampler, get_dataset, get_loss, get_metrics, - get_model, get_tokenizer, read_yamls) +from utils import PerDatasetSampler, get_dataset, get_loss, get_metrics, get_model, get_tokenizer, read_yamls def compute_metrics(eval_pred, preprocess_fns, metrics): @@ -92,20 +91,30 @@ class SFTTrainer(Trainer): return (loss, logits, labels) def get_train_dataloader(self): + """Inject custom data sampling behaviour into training loop""" if self.sampler is None: torch.utils.data.DataLoader( self.train_dataset, batch_size=self.args.per_device_train_batch_size, shuffle=True, - collate_fn=self.data_collator + collate_fn=self.data_collator, ) else: - return torch.utils.data.DataLoader( + dataloader = torch.utils.data.DataLoader( self.train_dataset, batch_size=self.args.per_device_train_batch_size, sampler=self.sampler, - collate_fn=self.data_collator + collate_fn=self.data_collator, ) + if torch.cuda.device_count() <= 1: + return dataloader + else: + # Not strictly necessary to use accelerate, currently just + # ensures batches are padded to be divisible by # devices + from accelerate import Accelerator + + accelerator = Accelerator() + return accelerator.prepare(dataloader) def _strtobool(x): @@ -160,7 +169,7 @@ if __name__ == "__main__": model = get_model(training_conf, tokenizer) train, evals, collate_fn = get_dataset(training_conf, tokenizer) - sampler = build_train_sampler(training_conf, train.datasets) + sampler = PerDatasetSampler.build_sampler_from_config(training_conf, train.datasets) metrics, preprocess_fns = get_metrics(training_conf, tokenizer) optimizer = OptimizerNames.ADAMW_BNB if training_conf.quantization else OptimizerNames.ADAMW_HF @@ -178,7 +187,7 @@ if __name__ == "__main__": learning_rate=float(training_conf.learning_rate), deepspeed="configs/zero_config.json" if training_conf.deepspeed else None, optim=optimizer, - # fp16=True, + fp16=True, local_rank=training_conf.local_rank, gradient_checkpointing=training_conf.gradient_checkpointing, gradient_accumulation_steps=training_conf.gradient_accumulation_steps, @@ -205,7 +214,7 @@ if __name__ == "__main__": entity="maw501", name=f"{training_conf.model_name}-{training_conf.log_dir}-finetuned", ) - + trainer = SFTTrainer( model=model, args=args, diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py index 6d6b271a..c18cb380 100644 --- a/model/supervised_finetuning/utils.py +++ b/model/supervised_finetuning/utils.py @@ -1,11 +1,8 @@ -# from functools import partial +import random from pathlib import Path +from typing import List import evaluate -import random - -# import nltk -import numpy as np import transformers import yaml from custom_datasets import get_one_dataset @@ -18,32 +15,55 @@ from torch.utils.data import ConcatDataset, Subset from torch.utils.data.sampler import Sampler -class ClassSampler(Sampler): - """Sampler which returns a fixed number of samples per class, per epoch""" - def __init__(self, class_labels, class_sizes): - self.class_labels = class_labels - self.class_sizes = class_sizes +class PerDatasetSampler(Sampler): + """Sampler which returns a fixed number of samples per dataset, per epoch. + + Example: + + Dataset 1 has 10,000 examples and we want 200 per epoch + Dataset 2 has 500 examples and we want all 500 per epoch + + Epoch size will be 700 and every epoch we'll sample a different + 200 from dataset 1. + + Parameters + ---------- + dataset_sizes : List[int] + A list with the size of each dataset. + dataset_size_per_epoch : List[int] + How many examples to get from each dataset per epoch. + + Note: dataset_sizes & dataset_size_per_epoch must be in the same order. + Further the examples in the underlying torch.utils.data.Dataset + must per ordered as dataset_1, dataset_2, ..., dataset_n. This is fine + if we concatenate a bunch of datasets together + e.g. using torch.utils.data.ConcatDataset which is current behaviour. + """ + + def __init__(self, dataset_sizes: List[int], dataset_size_per_epoch: List[int]): + self.dataset_sizes = dataset_sizes + self.dataset_size_per_epoch = dataset_size_per_epoch + self.num_datasets = len(dataset_sizes) def __iter__(self): - out = [] - for i, _class in enumerate(np.unique(self.class_labels)): - class_idx = np.argwhere(self.class_labels == _class).flatten() - sampled_idx = random.sample(list(class_idx), int(self.class_sizes[i])) - out.extend(sampled_idx) - random.shuffle(out) - return iter(out) + epoch_idx = [] + n = 0 + for i in range(self.num_datasets): + sampled_idx = random.sample(range(n, self.dataset_sizes[i] + n), self.dataset_size_per_epoch[i]) + n += self.dataset_sizes[i] + epoch_idx.extend(sampled_idx) + random.shuffle(epoch_idx) + return iter(epoch_idx) def __len__(self): - return int(sum(self.class_sizes)) + return int(sum(self.dataset_size_per_epoch)) - -def build_train_sampler(training_conf, datasets): - train_sizes = [len(x) for x in datasets] - fractions = get_dataset_fractions(training_conf.datasets, train_sizes) - dataset_size_per_epoch = [int(size * frac) for size, frac in zip(train_sizes, fractions)] - dataset_labels = [[i] * d for i, d in zip(range(len(dataset_size_per_epoch)), dataset_size_per_epoch)] - dataset_labels = [i for s in dataset_labels for i in s] - return ClassSampler(dataset_labels, dataset_size_per_epoch) + @classmethod + def build_sampler_from_config(cls, training_conf, datasets): + dataset_sizes = [len(x) for x in datasets] + fractions = get_dataset_fractions(training_conf.datasets, dataset_sizes) + dataset_size_per_epoch = [int(size * frac) for size, frac in zip(dataset_sizes, fractions)] + return cls(dataset_sizes, dataset_size_per_epoch) def get_tokenizer(conf): @@ -157,13 +177,15 @@ def get_dataset_fractions(conf, dataset_sizes): dataset_name = get_dataset_name_from_data_config(data_config) if isinstance(data_config, dict): if "fraction" in data_config[dataset_name]: + if data_config[dataset_name]["fraction"] <= 0: + raise ValueError("Please specify fraction as a value between 0 < fraction <= 1") fractions.append(min(1, data_config[dataset_name]["fraction"])) elif "size" in data_config[dataset_name]: if data_config[dataset_name]["size"] > dataset_sizes[i]: - raise ValueError(f"Please specify a size smaller than number of examples ({dataset_sizes[i]})") + raise ValueError(f"Please specify a size smaller than number of examples: {dataset_sizes[i]:,.0f}") fractions.append(data_config[dataset_name]["size"] / dataset_sizes[i]) else: - raise ValueError("Please specify either fraction or size in config.yaml") + raise ValueError("Please specify either fraction or size in config.yaml. See README for instructions.") else: fractions.append(1) return fractions