wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-07-13 00:50:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

Get working on multi-gpu

Browse Source
This commit is contained in:
Mark Worrall
2023-02-08 20:49:25 +00:00
parent e2caf53654
commit 283df8ec84
7 changed files with 167 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/supervised_finetuning/README.md
+76 -52
View File
@@ -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: <huggingface 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: <huggingface 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)
model/supervised_finetuning/configs/config.yaml
+22 -25
View File
@@ -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:
model/supervised_finetuning/custom_datasets/qa_datasets.py
+1 -4
View File
@@ -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)
model/supervised_finetuning/custom_datasets/translation.py
-2
View File
@@ -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):
model/supervised_finetuning/requirements.txt
+1 -1
View File
@@ -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
model/supervised_finetuning/trainer.py
+17 -8
View File
@@ -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,
model/supervised_finetuning/utils.py
+50 -28
View File
@@ -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