From dfaa00dccc1038e1ad8aa00d5868d80d1598fa4a Mon Sep 17 00:00:00 2001 From: Sotirios Anagnostidis Date: Thu, 5 Jan 2023 00:33:16 +0100 Subject: [PATCH 1/6] gptj 8bit --- .../supervised_finetuning/configs/config.yaml | 5 + .../supervised_finetuning/models/__init__.py | 8 + model/supervised_finetuning/models/gptj.py | 185 ++++++++++++++++++ model/supervised_finetuning/trainer.py | 17 +- model/supervised_finetuning/utils.py | 5 +- 5 files changed, 215 insertions(+), 5 deletions(-) create mode 100644 model/supervised_finetuning/models/__init__.py create mode 100644 model/supervised_finetuning/models/gptj.py diff --git a/model/supervised_finetuning/configs/config.yaml b/model/supervised_finetuning/configs/config.yaml index 29086395..ded4363e 100644 --- a/model/supervised_finetuning/configs/config.yaml +++ b/model/supervised_finetuning/configs/config.yaml @@ -21,6 +21,7 @@ defaults: loss_fn: CrossEntropyLoss eval_size: log_dir: "base" + quantization: galactica-125: learning_rate: 5e-5 @@ -46,3 +47,7 @@ gpt-jt: debug: eval_steps: 20 eval_size: 100 + gradient_accumulation_steps: 2 + per_device_train_batch_size: 1 + per_device_eval_batch_size: 1 + quantization: 8bit \ No newline at end of file diff --git a/model/supervised_finetuning/models/__init__.py b/model/supervised_finetuning/models/__init__.py new file mode 100644 index 00000000..673f4085 --- /dev/null +++ b/model/supervised_finetuning/models/__init__.py @@ -0,0 +1,8 @@ +from transformers import AutoModelForCausalLM +from .gptj import get_model as get_gptj_model + +def get_specific_model(model_name, cache_dir, quantization): + if "gpt-j" in model_name.lower(): + return get_gptj_model(model_name, cache_dir, quantization) + else: + return AutoModelForCausalLM.from_pretrained(conf.model_name, cache_dir=conf.cache_dir) \ No newline at end of file diff --git a/model/supervised_finetuning/models/gptj.py b/model/supervised_finetuning/models/gptj.py new file mode 100644 index 00000000..11234abd --- /dev/null +++ b/model/supervised_finetuning/models/gptj.py @@ -0,0 +1,185 @@ +# Taken from https://github.com/sleekmike/Finetune_GPT-J_6B_8-bit/blob/master/gpt-j-6b-8-bit.py + +import transformers +from transformers import AutoModelForCausalLM +import torch +import torch.nn.functional as F +from torch import nn +from torch.cuda.amp import custom_fwd, custom_bwd +from bitsandbytes.functional import quantize_blockwise, dequantize_blockwise + + +class FrozenBNBLinear(nn.Module): + def __init__(self, weight, absmax, code, bias=None): + assert isinstance(bias, nn.Parameter) or bias is None + super().__init__() + self.out_features, self.in_features = weight.shape + self.register_buffer("weight", weight.requires_grad_(False)) + self.register_buffer("absmax", absmax.requires_grad_(False)) + self.register_buffer("code", code.requires_grad_(False)) + self.adapter = None + self.bias = bias + + def forward(self, input): + output = DequantizeAndLinear.apply(input, self.weight, self.absmax, self.code, self.bias) + if self.adapter: + output += self.adapter(input) + return output + + @classmethod + def from_linear(cls, linear: nn.Linear) -> "FrozenBNBLinear": + weights_int8, state = quantize_blockise_lowmemory(linear.weight) + return cls(weights_int8, *state, linear.bias) + + def __repr__(self): + return f"{self.__class__.__name__}({self.in_features}, {self.out_features})" + + +class DequantizeAndLinear(torch.autograd.Function): + @staticmethod + @custom_fwd + def forward( + ctx, + input: torch.Tensor, + weights_quantized: torch.ByteTensor, + absmax: torch.FloatTensor, + code: torch.FloatTensor, + bias: torch.FloatTensor, + ): + weights_deq = dequantize_blockwise(weights_quantized, absmax=absmax, code=code) + ctx.save_for_backward(input, weights_quantized, absmax, code) + ctx._has_bias = bias is not None + return F.linear(input, weights_deq, bias) + + @staticmethod + @custom_bwd + def backward(ctx, grad_output: torch.Tensor): + assert not ctx.needs_input_grad[1] and not ctx.needs_input_grad[2] and not ctx.needs_input_grad[3] + input, weights_quantized, absmax, code = ctx.saved_tensors + # grad_output: [*batch, out_features] + weights_deq = dequantize_blockwise(weights_quantized, absmax=absmax, code=code) + grad_input = grad_output @ weights_deq + grad_bias = grad_output.flatten(0, -2).sum(dim=0) if ctx._has_bias else None + return grad_input, None, None, None, grad_bias + + +class FrozenBNBEmbedding(nn.Module): + def __init__(self, weight, absmax, code): + super().__init__() + self.num_embeddings, self.embedding_dim = weight.shape + self.register_buffer("weight", weight.requires_grad_(False)) + self.register_buffer("absmax", absmax.requires_grad_(False)) + self.register_buffer("code", code.requires_grad_(False)) + self.adapter = None + + def forward(self, input, **kwargs): + with torch.no_grad(): + # note: both quantuized weights and input indices are *not* differentiable + weight_deq = dequantize_blockwise(self.weight, absmax=self.absmax, code=self.code) + output = F.embedding(input, weight_deq, **kwargs) + if self.adapter: + output += self.adapter(input) + return output + + @classmethod + def from_embedding(cls, embedding: nn.Embedding) -> "FrozenBNBEmbedding": + weights_int8, state = quantize_blockise_lowmemory(embedding.weight) + return cls(weights_int8, *state) + + def __repr__(self): + return f"{self.__class__.__name__}({self.num_embeddings}, {self.embedding_dim})" + + +def quantize_blockise_lowmemory(matrix: torch.Tensor, chunk_size: int = 2**20): + assert chunk_size % 4096 == 0 + code = None + chunks = [] + absmaxes = [] + flat_tensor = matrix.view(-1) + for i in range((matrix.numel() - 1) // chunk_size + 1): + input_chunk = flat_tensor[i * chunk_size : (i + 1) * chunk_size].clone() + quantized_chunk, (absmax_chunk, code) = quantize_blockwise(input_chunk, code=code) + chunks.append(quantized_chunk) + absmaxes.append(absmax_chunk) + + matrix_i8 = torch.cat(chunks).reshape_as(matrix) + absmax = torch.cat(absmaxes) + return matrix_i8, (absmax, code) + + +def convert_to_int8(model): + """Convert linear and embedding modules to 8-bit with optional adapters""" + for module in list(model.modules()): + for name, child in module.named_children(): + if isinstance(child, nn.Linear): + print(name, child) + setattr( + module, + name, + FrozenBNBLinear( + weight=torch.zeros(child.out_features, child.in_features, dtype=torch.uint8), + absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1), + code=torch.zeros(256), + bias=child.bias, + ), + ) + elif isinstance(child, nn.Embedding): + setattr( + module, + name, + FrozenBNBEmbedding( + weight=torch.zeros(child.num_embeddings, child.embedding_dim, dtype=torch.uint8), + absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1), + code=torch.zeros(256), + ), + ) + + +class GPTJBlock(transformers.models.gptj.modeling_gptj.GPTJBlock): + def __init__(self, config): + super().__init__(config) + + convert_to_int8(self.attn) + convert_to_int8(self.mlp) + + +class GPTJModel(transformers.models.gptj.modeling_gptj.GPTJModel): + def __init__(self, config): + super().__init__(config) + convert_to_int8(self) + + +class GPTJForCausalLM(transformers.models.gptj.modeling_gptj.GPTJForCausalLM): + def __init__(self, config): + super().__init__(config) + convert_to_int8(self) + +def add_adapters(model, adapter_dim=16): + assert adapter_dim > 0 + + for module in model.modules(): + if isinstance(module, FrozenBNBLinear): + module.adapter = nn.Sequential( + nn.Linear(module.in_features, adapter_dim, bias=False), + nn.Linear(adapter_dim, module.out_features, bias=False), + ) + nn.init.zeros_(module.adapter[1].weight) + elif isinstance(module, FrozenBNBEmbedding): + module.adapter = nn.Sequential( + nn.Embedding(module.num_embeddings, adapter_dim), + nn.Linear(adapter_dim, module.embedding_dim, bias=False), + ) + nn.init.zeros_(module.adapter[1].weight) + + +def get_model(model_name, cache_dir, quantization): + if quantization is None: + model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) + elif quantization == "8bit": + transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJBlock + model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) + add_adapters(gpt) + else: + raise ValueError(f"Unknown quantization {quantization}") + + return model diff --git a/model/supervised_finetuning/trainer.py b/model/supervised_finetuning/trainer.py index dc7b5934..0db0f443 100644 --- a/model/supervised_finetuning/trainer.py +++ b/model/supervised_finetuning/trainer.py @@ -17,6 +17,7 @@ from transformers import ( TrainingArguments, get_cosine_schedule_with_warmup, ) +import bitsandbytes as bnb from utils import get_dataset, get_loss, get_model, get_tokenizer, read_yamls os.environ["WANDB_PROJECT"] = "supervised-finetuning" @@ -25,6 +26,7 @@ os.environ["WANDB_PROJECT"] = "supervised-finetuning" @dataclass class CustomTrainingArguments(TrainingArguments): loss_function: str = "CrossEntropyLoss" + quantization: str = None def compute_metrics(eval_pred): @@ -71,8 +73,16 @@ class SFTTrainer(Trainer): # By default CrossEntropyLoss ignores padding_index -100, but just in case use our own loss_fct self.loss_fct = get_loss(args.loss_function) - def fetch_scheduler(self): - return get_cosine_schedule_with_warmup( + def create_optimizer_and_scheduler(self, num_training_steps: int): + if self.args.quantization == "8bit": + self.optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001, betas=(0.9, 0.995)) + else: + self.optimizer = torch.optim.AdamW( + self.model.parameters(), lr=self.args.learning_rate, weight_decay=self.args.weight_decay + ) + + print("lr sheduler") + self.lr_scheduler = get_cosine_schedule_with_warmup( self.optimizer, num_warmup_steps=self.args.warmup_steps, num_training_steps=self.num_train_steps, @@ -165,6 +175,7 @@ if __name__ == "__main__": model = get_model(training_conf, tokenizer) train, evals, collate_fn = get_dataset(training_conf, tokenizer) + assert len(evals) > 0 args = CustomTrainingArguments( output_dir=f"{training_conf.model_name}-{training_conf.log_dir}-finetuned", @@ -186,9 +197,9 @@ if __name__ == "__main__": save_steps=training_conf.save_steps, eval_accumulation_steps=training_conf.eval_accumulation_steps, report_to="wandb", + quantization=training_conf.quantization, ) - assert len(evals) > 0 trainer = SFTTrainer( model, args, diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py index a31f74d3..318aeb5e 100644 --- a/model/supervised_finetuning/utils.py +++ b/model/supervised_finetuning/utils.py @@ -6,7 +6,8 @@ from custom_datasets.dialogue_collator import DialogueDataCollator from losses import CrossEntropyLoss from sklearn.model_selection import train_test_split from torch.utils.data import ConcatDataset, Subset -from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers import AutoTokenizer +from models import get_specific_model SUPPORTED_MODELS = ["galactica", "GPT-JT"] # deprecated .. @@ -36,7 +37,7 @@ def get_model(conf, tokenizer): "To include more make sure the masking is dne correctly... (decoder only supported for now)" ) - model = AutoModelForCausalLM.from_pretrained(conf.model_name, cache_dir=conf.cache_dir) + model = get_specific_model(conf.model_name, conf.cache_dir, conf.quantization) if len(tokenizer) != model.get_input_embeddings().num_embeddings: assert not conf.freeze_layer, "Cannot change the number of embeddings if the model is frozen." From ef02693ac9b4705b4ba026a05436709a479dcec7 Mon Sep 17 00:00:00 2001 From: Sotirios Anagnostidis Date: Fri, 6 Jan 2023 18:24:28 +0100 Subject: [PATCH 2/6] quantization --- .../supervised_finetuning/configs/config.yaml | 2 +- .../supervised_finetuning/models/__init__.py | 25 +++- model/supervised_finetuning/models/gptj.py | 60 +++++----- model/supervised_finetuning/trainer.py | 107 ++++++++++++------ model/supervised_finetuning/utils.py | 36 +----- 5 files changed, 128 insertions(+), 102 deletions(-) diff --git a/model/supervised_finetuning/configs/config.yaml b/model/supervised_finetuning/configs/config.yaml index ded4363e..2f529e85 100644 --- a/model/supervised_finetuning/configs/config.yaml +++ b/model/supervised_finetuning/configs/config.yaml @@ -50,4 +50,4 @@ debug: gradient_accumulation_steps: 2 per_device_train_batch_size: 1 per_device_eval_batch_size: 1 - quantization: 8bit \ No newline at end of file + quantization: \ No newline at end of file diff --git a/model/supervised_finetuning/models/__init__.py b/model/supervised_finetuning/models/__init__.py index 673f4085..b99053e3 100644 --- a/model/supervised_finetuning/models/__init__.py +++ b/model/supervised_finetuning/models/__init__.py @@ -1,8 +1,31 @@ from transformers import AutoModelForCausalLM from .gptj import get_model as get_gptj_model +SUPPORTED_MODELS = ["galactica", "gpt-j"] + + +def freeze_top_n_layers(model, target_layers): + # its possible we can simply detect which module is a ModuleList + # and simply freeze the module without doing string parsing + for name, param in model.named_parameters(): + if "embed" in name: + param.requires_grad = False + elif ".layer" in name or ".h." in name: + tokens = name.split(".") + layer_ = None + for token in tokens: + if token.isdigit(): + layer_ = int(token) + break + + if layer_ is not None and layer_ < target_layers: + # print('freeze ', layer_, name) + param.requires_grad = False + return model + + def get_specific_model(model_name, cache_dir, quantization): if "gpt-j" in model_name.lower(): return get_gptj_model(model_name, cache_dir, quantization) else: - return AutoModelForCausalLM.from_pretrained(conf.model_name, cache_dir=conf.cache_dir) \ No newline at end of file + return AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) diff --git a/model/supervised_finetuning/models/gptj.py b/model/supervised_finetuning/models/gptj.py index 11234abd..3cbec3ce 100644 --- a/model/supervised_finetuning/models/gptj.py +++ b/model/supervised_finetuning/models/gptj.py @@ -19,38 +19,32 @@ class FrozenBNBLinear(nn.Module): self.register_buffer("code", code.requires_grad_(False)) self.adapter = None self.bias = bias - + def forward(self, input): output = DequantizeAndLinear.apply(input, self.weight, self.absmax, self.code, self.bias) if self.adapter: output += self.adapter(input) return output - + @classmethod def from_linear(cls, linear: nn.Linear) -> "FrozenBNBLinear": weights_int8, state = quantize_blockise_lowmemory(linear.weight) return cls(weights_int8, *state, linear.bias) - + def __repr__(self): return f"{self.__class__.__name__}({self.in_features}, {self.out_features})" - - -class DequantizeAndLinear(torch.autograd.Function): + + +class DequantizeAndLinear(torch.autograd.Function): @staticmethod @custom_fwd - def forward( - ctx, - input: torch.Tensor, - weights_quantized: torch.ByteTensor, - absmax: torch.FloatTensor, - code: torch.FloatTensor, - bias: torch.FloatTensor, - ): + def forward(ctx, input: torch.Tensor, weights_quantized: torch.ByteTensor, + absmax: torch.FloatTensor, code: torch.FloatTensor, bias: torch.FloatTensor): weights_deq = dequantize_blockwise(weights_quantized, absmax=absmax, code=code) ctx.save_for_backward(input, weights_quantized, absmax, code) ctx._has_bias = bias is not None return F.linear(input, weights_deq, bias) - + @staticmethod @custom_bwd def backward(ctx, grad_output: torch.Tensor): @@ -61,8 +55,8 @@ class DequantizeAndLinear(torch.autograd.Function): grad_input = grad_output @ weights_deq grad_bias = grad_output.flatten(0, -2).sum(dim=0) if ctx._has_bias else None return grad_input, None, None, None, grad_bias - - + + class FrozenBNBEmbedding(nn.Module): def __init__(self, weight, absmax, code): super().__init__() @@ -71,7 +65,7 @@ class FrozenBNBEmbedding(nn.Module): self.register_buffer("absmax", absmax.requires_grad_(False)) self.register_buffer("code", code.requires_grad_(False)) self.adapter = None - + def forward(self, input, **kwargs): with torch.no_grad(): # note: both quantuized weights and input indices are *not* differentiable @@ -79,41 +73,41 @@ class FrozenBNBEmbedding(nn.Module): output = F.embedding(input, weight_deq, **kwargs) if self.adapter: output += self.adapter(input) - return output - + return output + @classmethod def from_embedding(cls, embedding: nn.Embedding) -> "FrozenBNBEmbedding": weights_int8, state = quantize_blockise_lowmemory(embedding.weight) return cls(weights_int8, *state) - + def __repr__(self): return f"{self.__class__.__name__}({self.num_embeddings}, {self.embedding_dim})" - - -def quantize_blockise_lowmemory(matrix: torch.Tensor, chunk_size: int = 2**20): + + +def quantize_blockise_lowmemory(matrix: torch.Tensor, chunk_size: int = 2 ** 20): assert chunk_size % 4096 == 0 code = None chunks = [] absmaxes = [] flat_tensor = matrix.view(-1) for i in range((matrix.numel() - 1) // chunk_size + 1): - input_chunk = flat_tensor[i * chunk_size : (i + 1) * chunk_size].clone() + input_chunk = flat_tensor[i * chunk_size: (i + 1) * chunk_size].clone() quantized_chunk, (absmax_chunk, code) = quantize_blockwise(input_chunk, code=code) chunks.append(quantized_chunk) absmaxes.append(absmax_chunk) - + matrix_i8 = torch.cat(chunks).reshape_as(matrix) absmax = torch.cat(absmaxes) return matrix_i8, (absmax, code) - - + + def convert_to_int8(model): """Convert linear and embedding modules to 8-bit with optional adapters""" for module in list(model.modules()): for name, child in module.named_children(): if isinstance(child, nn.Linear): print(name, child) - setattr( + setattr( module, name, FrozenBNBLinear( @@ -131,7 +125,7 @@ def convert_to_int8(model): weight=torch.zeros(child.num_embeddings, child.embedding_dim, dtype=torch.uint8), absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1), code=torch.zeros(256), - ), + ) ) @@ -147,13 +141,14 @@ class GPTJModel(transformers.models.gptj.modeling_gptj.GPTJModel): def __init__(self, config): super().__init__(config) convert_to_int8(self) - + class GPTJForCausalLM(transformers.models.gptj.modeling_gptj.GPTJForCausalLM): def __init__(self, config): super().__init__(config) convert_to_int8(self) + def add_adapters(model, adapter_dim=16): assert adapter_dim > 0 @@ -176,9 +171,10 @@ def get_model(model_name, cache_dir, quantization): if quantization is None: model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) elif quantization == "8bit": + print("Loading 8-bit model") transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJBlock model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) - add_adapters(gpt) + add_adapters(model) else: raise ValueError(f"Unknown quantization {quantization}") diff --git a/model/supervised_finetuning/trainer.py b/model/supervised_finetuning/trainer.py index 0db0f443..590f9bbd 100644 --- a/model/supervised_finetuning/trainer.py +++ b/model/supervised_finetuning/trainer.py @@ -74,6 +74,7 @@ class SFTTrainer(Trainer): self.loss_fct = get_loss(args.loss_function) def create_optimizer_and_scheduler(self, num_training_steps: int): + print("Optimizer") if self.args.quantization == "8bit": self.optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001, betas=(0.9, 0.995)) else: @@ -174,40 +175,76 @@ if __name__ == "__main__": tokenizer = get_tokenizer(training_conf) model = get_model(training_conf, tokenizer) - train, evals, collate_fn = get_dataset(training_conf, tokenizer) - assert len(evals) > 0 + ### + from datasets import load_dataset + from bitsandbytes.optim import Adam8bit + from torch.nn import functional as F + from tqdm import tqdm - args = CustomTrainingArguments( - output_dir=f"{training_conf.model_name}-{training_conf.log_dir}-finetuned", - num_train_epochs=training_conf.num_train_epochs, - warmup_steps=training_conf.warmup_steps, - loss_function=training_conf.loss_fn, - learning_rate=float(training_conf.learning_rate), - fp16=True, - gradient_checkpointing=training_conf.gradient_checkpointing, - gradient_accumulation_steps=training_conf.gradient_accumulation_steps, - per_device_train_batch_size=training_conf.per_device_train_batch_size, - per_device_eval_batch_size=training_conf.per_device_eval_batch_size, - weight_decay=training_conf.weight_decay, - max_grad_norm=training_conf.max_grad_norm, - logging_steps=training_conf.logging_steps, - save_total_limit=training_conf.save_total_limit, - evaluation_strategy="steps", - eval_steps=training_conf.eval_steps, - save_steps=training_conf.save_steps, - eval_accumulation_steps=training_conf.eval_accumulation_steps, - report_to="wandb", - quantization=training_conf.quantization, - ) + gpt = model.to("cuda") - trainer = SFTTrainer( - model, - args, - train_dataset=train, - eval_dataset=evals, - data_collator=collate_fn, - tokenizer=tokenizer, - compute_metrics=compute_metrics, - preprocess_logits_for_metrics=preprocess_logits_for_metrics, - ) - trainer.train() + gpt.gradient_checkpointing_enable() + + codeparrot = load_dataset("transformersbook/codeparrot-train", streaming=True, cache_dir=training_conf.cache_dir) + optimizer = Adam8bit(gpt.parameters(), lr=1e-5) + + with torch.cuda.amp.autocast(): + for row in tqdm(codeparrot["train"]): + if len(row["content"]) <= 1: + continue + + batch = tokenizer(row["content"], truncation=True, max_length=128, return_tensors="pt") + batch = {k: v.cuda() for k, v in batch.items()} + + out = gpt.forward( + **batch, + ) + + loss = F.cross_entropy( + out.logits[:, :-1, :].flatten(0, -2), batch["input_ids"][:, 1:].flatten(), reduction="mean" + ) + print(loss) + loss.backward() + + optimizer.step() + optimizer.zero_grad() + ### + + + # train, evals, collate_fn = get_dataset(training_conf, tokenizer) + # assert len(evals) > 0 + + # args = CustomTrainingArguments( + # output_dir=f"{training_conf.model_name}-{training_conf.log_dir}-finetuned", + # num_train_epochs=training_conf.num_train_epochs, + # warmup_steps=training_conf.warmup_steps, + # loss_function=training_conf.loss_fn, + # learning_rate=float(training_conf.learning_rate), + # fp16=True, + # gradient_checkpointing=training_conf.gradient_checkpointing, + # gradient_accumulation_steps=training_conf.gradient_accumulation_steps, + # per_device_train_batch_size=training_conf.per_device_train_batch_size, + # per_device_eval_batch_size=training_conf.per_device_eval_batch_size, + # weight_decay=training_conf.weight_decay, + # max_grad_norm=training_conf.max_grad_norm, + # logging_steps=training_conf.logging_steps, + # save_total_limit=training_conf.save_total_limit, + # evaluation_strategy="steps", + # eval_steps=training_conf.eval_steps, + # save_steps=training_conf.save_steps, + # eval_accumulation_steps=training_conf.eval_accumulation_steps, + # report_to="wandb", + # quantization=training_conf.quantization, + # ) + + # trainer = SFTTrainer( + # model, + # args, + # train_dataset=train, + # eval_dataset=evals, + # data_collator=collate_fn, + # tokenizer=tokenizer, + # compute_metrics=compute_metrics, + # preprocess_logits_for_metrics=preprocess_logits_for_metrics, + # ) + # trainer.train() diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py index 318aeb5e..c3cc512c 100644 --- a/model/supervised_finetuning/utils.py +++ b/model/supervised_finetuning/utils.py @@ -7,9 +7,7 @@ from losses import CrossEntropyLoss from sklearn.model_selection import train_test_split from torch.utils.data import ConcatDataset, Subset from transformers import AutoTokenizer -from models import get_specific_model - -SUPPORTED_MODELS = ["galactica", "GPT-JT"] # deprecated .. +from models import get_specific_model, SUPPORTED_MODELS, freeze_top_n_layers def get_tokenizer(conf): @@ -31,10 +29,10 @@ def get_tokenizer(conf): def get_model(conf, tokenizer): - if not any([x in conf.model_name for x in SUPPORTED_MODELS]): + if not any([x in conf.model_name.lower() for x in SUPPORTED_MODELS]): raise ValueError( f"Model {conf.model_name} not supported. Supported models: {SUPPORTED_MODELS}. " - "To include more make sure the masking is dne correctly... (decoder only supported for now)" + "To include more make sure the masking is done correctly... (decoder only supported for now)" ) model = get_specific_model(conf.model_name, conf.cache_dir, conf.quantization) @@ -96,31 +94,3 @@ def train_val_dataset(dataset, val_split=0.2): list(range(len(dataset))), test_size=val_split, random_state=666, shuffle=True ) return Subset(dataset, train_idx), Subset(dataset, val_idx) - - -def freeze_top_n_layers(model, target_layers): - # its possible we can simply detect which module is a ModuleList - # and simply freeze the module without doing string parsing - for name, param in model.named_parameters(): - if "embed" in name: - param.requires_grad = False - elif ".layer" in name or ".h." in name: - tokens = name.split(".") - layer_ = None - for token in tokens: - if token.isdigit(): - layer_ = int(token) - break - - if layer_ is not None and layer_ < target_layers: - # print('freeze ', layer_, name) - param.requires_grad = False - return model - - -if __name__ == "__main__": - from transformers import AutoModelForSequenceClassification - - model = AutoModelForSequenceClassification.from_pretrained("bigscience/bloomz-560m") - freeze_top_n_layers(model, 10) - print(model.state_dict().keys()) From 91853753a8759208ac9c9c8cdb07ef3ae2fffdba Mon Sep 17 00:00:00 2001 From: Sotirios Anagnostidis Date: Fri, 6 Jan 2023 18:25:20 +0100 Subject: [PATCH 3/6] conf --- model/supervised_finetuning/configs/config.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/model/supervised_finetuning/configs/config.yaml b/model/supervised_finetuning/configs/config.yaml index 2f529e85..64c024c2 100644 --- a/model/supervised_finetuning/configs/config.yaml +++ b/model/supervised_finetuning/configs/config.yaml @@ -47,7 +47,8 @@ gpt-jt: debug: eval_steps: 20 eval_size: 100 + model_name: EleutherAI/gpt-j-6B gradient_accumulation_steps: 2 per_device_train_batch_size: 1 per_device_eval_batch_size: 1 - quantization: \ No newline at end of file + quantization: 8bit \ No newline at end of file From 88ee3b32644ee3f8a249dc281150f86505490b5e Mon Sep 17 00:00:00 2001 From: Sotirios Anagnostidis Date: Fri, 6 Jan 2023 21:28:26 +0100 Subject: [PATCH 4/6] merge deepspeed --- model/supervised_finetuning/README.md | 35 +++- .../supervised_finetuning/configs/config.yaml | 19 +- .../custom_datasets/__init__.py | 10 +- .../custom_datasets/prompt_dialogue.py | 66 +++++++ model/supervised_finetuning/losses.py | 2 +- model/supervised_finetuning/trainer.py | 171 ++++++------------ model/supervised_finetuning/utils.py | 10 +- 7 files changed, 181 insertions(+), 132 deletions(-) create mode 100644 model/supervised_finetuning/custom_datasets/prompt_dialogue.py diff --git a/model/supervised_finetuning/README.md b/model/supervised_finetuning/README.md index 014afa95..bd202397 100644 --- a/model/supervised_finetuning/README.md +++ b/model/supervised_finetuning/README.md @@ -23,7 +23,40 @@ open-asisstant dataset are available it will be added here. ## Model -TBD +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 +``` ## Results diff --git a/model/supervised_finetuning/configs/config.yaml b/model/supervised_finetuning/configs/config.yaml index 64c024c2..fb2bdaa0 100644 --- a/model/supervised_finetuning/configs/config.yaml +++ b/model/supervised_finetuning/configs/config.yaml @@ -6,7 +6,7 @@ defaults: per_device_eval_batch_size: 2 weight_decay: 0.00 warmup_steps: 600 - eval_steps: 200 + eval_steps: 100 save_steps: 500 max_length: 512 num_train_epochs: 3 @@ -17,6 +17,7 @@ defaults: freeze_layer: datasets: - webgpt + - prompt_dialogue cache_dir: ~/.cache loss_fn: CrossEntropyLoss eval_size: @@ -44,11 +45,21 @@ gpt-jt: per_device_train_batch_size: 4 per_device_eval_batch_size: 4 +codegen: + learning_rate: 2e-6 + model_name: Salesforce/codegen-2B-multi + 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 + debug: eval_steps: 20 eval_size: 100 - model_name: EleutherAI/gpt-j-6B - gradient_accumulation_steps: 2 + gradient_accumulation_steps: 1 per_device_train_batch_size: 1 per_device_eval_batch_size: 1 - quantization: 8bit \ No newline at end of file + quantization: \ No newline at end of file diff --git a/model/supervised_finetuning/custom_datasets/__init__.py b/model/supervised_finetuning/custom_datasets/__init__.py index 7e3bdc79..5706bfa7 100644 --- a/model/supervised_finetuning/custom_datasets/__init__.py +++ b/model/supervised_finetuning/custom_datasets/__init__.py @@ -2,6 +2,8 @@ from datasets import load_dataset from sklearn.model_selection import train_test_split from torch.utils.data import Dataset, Subset +from .prompt_dialogue import PromptGeneratedDataset + QA_SPECIAL_TOKENS = {"Question": "", "Answer": ""} @@ -14,8 +16,8 @@ class SquadV2Dataset(Dataset): def __getitem__(self, idx): data = self.dataset[idx] - # dummy return first answer - return "".join([data["title"], ". ", data["context"], " " + data["question"]]), data["answers"]["text"][0] + # return first answer form list of possible answers + return data["title"] + ". " + data["context"] + " " + data["question"], data["answers"]["text"][0] class WebGPT(Dataset): @@ -57,12 +59,14 @@ def get_one_dataset(conf, dataset_name): dataset_name = dataset_name.lower() if dataset_name == "squadv2": - raise ValueError("SquadV2 is not diverse enough for generation .. ") train = SquadV2Dataset(conf.cache_dir, "train") eval = SquadV2Dataset(conf.cache_dir, "validation") elif dataset_name == "webgpt": dataset = WebGPT() train, eval = train_val_dataset(dataset, val_split=0.2) + elif dataset_name == "prompt_dialogue": + dataset = PromptGeneratedDataset() + train, eval = train_val_dataset(dataset, val_split=0.2) else: raise ValueError(f"Unknown dataset {dataset_name}") diff --git a/model/supervised_finetuning/custom_datasets/prompt_dialogue.py b/model/supervised_finetuning/custom_datasets/prompt_dialogue.py new file mode 100644 index 00000000..17911141 --- /dev/null +++ b/model/supervised_finetuning/custom_datasets/prompt_dialogue.py @@ -0,0 +1,66 @@ +import os +from urllib.request import urlopen + +from torch.utils.data import Dataset + + +class PromptGeneratedDataset(Dataset): + """Generates from flan 11B + User: What are the best methods for preventing a slave trade? + + Rosey: The best methods .... + <|endoftext|> + + we are ignoring results with multiple lines for now + """ + + url = "https://github.com/Rallio67/language-model-agents/raw/main/chat_dialogue_v2_c.txt" + + def __init__(self) -> None: + super().__init__() + os.makedirs("datasets", exist_ok=True) + chat_dialogue = os.path.join("datasets", "chat_dialogue_v2_c.txt") + if not os.path.exists(chat_dialogue): + with urlopen(self.url) as file: + content = file.read().decode() + with open(chat_dialogue, "w") as fout: + fout.write(content) + + question = "" + answer = "" + self.pairs = [] + with open(chat_dialogue, "r") as f: + corpus = f.read().split("<|endoftext|>") + for dialogue in corpus: + dialogue = dialogue.strip() + if "Rosey:" in dialogue: + user, bot = dialogue.split("Rosey:", maxsplit=1) + question = user.split(":", maxsplit=1)[1].strip() + answer = bot.strip() + if len(answer) and len(question): + self.pairs.append((question, answer)) + + if len(question) > 0 and len(answer) > 0: + self.pairs.append((question, answer)) + + def __len__(self): + return len(self.pairs) + + def __getitem__(self, index): + question, answer = self.pairs[index] + return question, answer + + +if __name__ == "__main__": + from torch.utils.data import DataLoader + from transformers import AutoTokenizer + + from .dialogue_collator import DialogueDataCollator + + tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-2B-multi") + tokenizer.add_special_tokens({"pad_token": "<|endoftext|>", "sep_token": "<|endoftext|>"}) + dataset = PromptGeneratedDataset() + collate_fn = DialogueDataCollator(tokenizer, padding=True, max_length=128) + dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=5) + for batch in dataloader: + print(batch["input_ids"].shape) diff --git a/model/supervised_finetuning/losses.py b/model/supervised_finetuning/losses.py index 795396b9..0cc639cf 100644 --- a/model/supervised_finetuning/losses.py +++ b/model/supervised_finetuning/losses.py @@ -7,7 +7,7 @@ class CrossEntropyLoss(nn.CrossEntropyLoss): def forward(self, input, target, mask=None): if mask is not None: - mask = mask.view(-1) + mask = mask.view(-1).bool() input = input.view(-1, input.size(-1)) target = target.view(-1) input = input[mask] diff --git a/model/supervised_finetuning/trainer.py b/model/supervised_finetuning/trainer.py index 590f9bbd..bb77b9c3 100644 --- a/model/supervised_finetuning/trainer.py +++ b/model/supervised_finetuning/trainer.py @@ -1,34 +1,19 @@ import argparse import os -from dataclasses import dataclass from distutils.util import strtobool -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union import torch from torch import nn from torch.utils.data import Dataset -from transformers import ( - DataCollator, - EvalPrediction, - PreTrainedModel, - PreTrainedTokenizerBase, - Trainer, - TrainerCallback, - TrainingArguments, - get_cosine_schedule_with_warmup, -) +from transformers import PreTrainedModel, Trainer, TrainingArguments, get_cosine_schedule_with_warmup import bitsandbytes as bnb + from utils import get_dataset, get_loss, get_model, get_tokenizer, read_yamls os.environ["WANDB_PROJECT"] = "supervised-finetuning" -@dataclass -class CustomTrainingArguments(TrainingArguments): - loss_function: str = "CrossEntropyLoss" - quantization: str = None - - def compute_metrics(eval_pred): pred_ids = eval_pred.predictions labels = eval_pred.label_ids @@ -46,35 +31,15 @@ class SFTTrainer(Trainer): self, model: Union[PreTrainedModel, nn.Module] = None, args: TrainingArguments = None, - data_collator: Optional[DataCollator] = None, - train_dataset: Optional[Dataset] = None, - eval_dataset: Optional[Dataset] = None, - tokenizer: Optional[PreTrainedTokenizerBase] = None, - model_init: Callable[[], PreTrainedModel] = None, - compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None, - callbacks: Optional[List[TrainerCallback]] = None, - optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), - preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None, + loss_function: str = "CrossEntropyLoss", + **kwargs, ): - super().__init__( - model, - args, - data_collator, - train_dataset, - eval_dataset, - tokenizer, - model_init, - compute_metrics, - callbacks, - optimizers, - preprocess_logits_for_metrics, - ) + super().__init__(model, args, **kwargs) # By default CrossEntropyLoss ignores padding_index -100, but just in case use our own loss_fct - self.loss_fct = get_loss(args.loss_function) + self.loss_fct = get_loss(loss_function) def create_optimizer_and_scheduler(self, num_training_steps: int): - print("Optimizer") if self.args.quantization == "8bit": self.optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001, betas=(0.9, 0.995)) else: @@ -82,7 +47,6 @@ class SFTTrainer(Trainer): self.model.parameters(), lr=self.args.learning_rate, weight_decay=self.args.weight_decay ) - print("lr sheduler") self.lr_scheduler = get_cosine_schedule_with_warmup( self.optimizer, num_warmup_steps=self.args.warmup_steps, @@ -93,16 +57,17 @@ class SFTTrainer(Trainer): def compute_loss(self, model, inputs, return_outputs=False): labels_mask = inputs.pop("label_masks") + targets = inputs.pop("targets") outputs = model(**inputs) - loss = self.loss_fct(outputs.get("logits"), torch.roll(inputs["input_ids"], -1, -1), mask=labels_mask) + loss = self.loss_fct(outputs.get("logits"), targets, mask=labels_mask) return (loss, outputs) if return_outputs else loss def _compute_loss(self, model, inputs): - labels_mask = inputs.pop("label_masks") + targets = inputs.pop("targets") inputs = self._prepare_inputs(inputs) @@ -110,7 +75,6 @@ class SFTTrainer(Trainer): logits = outputs.get("logits") - targets = torch.roll(inputs["input_ids"], -1, -1) loss = self.loss_fct(outputs.get("logits"), targets, mask=labels_mask) return loss, logits, targets, labels_mask @@ -142,12 +106,18 @@ def _strtobool(x): def argument_parsing(notebook=False, notebook_args=None): parser = argparse.ArgumentParser() parser.add_argument("--configs", nargs="+", required=True) + parser.add_argument("--local_rank", type=int, default=-1) + parser.add_argument("--deepspeed", action="store_true") + parser.add_argument("--no-deepspeed", dest="deepspeed", action="store_false") + parser.set_defaults(deepspeed=False) if notebook: args, remaining = parser.parse_known_args(notebook_args) else: args, remaining = parser.parse_known_args() + print(args) + # Config from YAML conf = {} configs = read_yamls("./configs") @@ -158,6 +128,8 @@ def argument_parsing(notebook=False, notebook_args=None): else: conf.update(configs[name]) + conf["local_rank"] = args.local_rank + conf["deepspeed"] = args.deepspeed # Override config from command-line parser = argparse.ArgumentParser() for key, value in conf.items(): @@ -175,76 +147,41 @@ if __name__ == "__main__": tokenizer = get_tokenizer(training_conf) model = get_model(training_conf, tokenizer) - ### - from datasets import load_dataset - from bitsandbytes.optim import Adam8bit - from torch.nn import functional as F - from tqdm import tqdm + train, evals, collate_fn = get_dataset(training_conf, tokenizer) - gpt = model.to("cuda") + args = TrainingArguments( + output_dir=f"{training_conf.model_name}-{training_conf.log_dir}-finetuned", + num_train_epochs=training_conf.num_train_epochs, + warmup_steps=training_conf.warmup_steps, + learning_rate=float(training_conf.learning_rate), + deepspeed="configs/zero_config.json" if training_conf.deepspeed else None, + fp16=True, + local_rank=training_conf.local_rank, + gradient_checkpointing=training_conf.gradient_checkpointing, + gradient_accumulation_steps=training_conf.gradient_accumulation_steps, + per_device_train_batch_size=training_conf.per_device_train_batch_size, + per_device_eval_batch_size=training_conf.per_device_eval_batch_size, + weight_decay=training_conf.weight_decay, + max_grad_norm=training_conf.max_grad_norm, + logging_steps=training_conf.logging_steps, + save_total_limit=training_conf.save_total_limit, + evaluation_strategy="steps", + eval_steps=training_conf.eval_steps, + save_steps=training_conf.save_steps, + eval_accumulation_steps=training_conf.eval_accumulation_steps, + report_to="wandb", + ) - gpt.gradient_checkpointing_enable() - - codeparrot = load_dataset("transformersbook/codeparrot-train", streaming=True, cache_dir=training_conf.cache_dir) - optimizer = Adam8bit(gpt.parameters(), lr=1e-5) - - with torch.cuda.amp.autocast(): - for row in tqdm(codeparrot["train"]): - if len(row["content"]) <= 1: - continue - - batch = tokenizer(row["content"], truncation=True, max_length=128, return_tensors="pt") - batch = {k: v.cuda() for k, v in batch.items()} - - out = gpt.forward( - **batch, - ) - - loss = F.cross_entropy( - out.logits[:, :-1, :].flatten(0, -2), batch["input_ids"][:, 1:].flatten(), reduction="mean" - ) - print(loss) - loss.backward() - - optimizer.step() - optimizer.zero_grad() - ### - - - # train, evals, collate_fn = get_dataset(training_conf, tokenizer) - # assert len(evals) > 0 - - # args = CustomTrainingArguments( - # output_dir=f"{training_conf.model_name}-{training_conf.log_dir}-finetuned", - # num_train_epochs=training_conf.num_train_epochs, - # warmup_steps=training_conf.warmup_steps, - # loss_function=training_conf.loss_fn, - # learning_rate=float(training_conf.learning_rate), - # fp16=True, - # gradient_checkpointing=training_conf.gradient_checkpointing, - # gradient_accumulation_steps=training_conf.gradient_accumulation_steps, - # per_device_train_batch_size=training_conf.per_device_train_batch_size, - # per_device_eval_batch_size=training_conf.per_device_eval_batch_size, - # weight_decay=training_conf.weight_decay, - # max_grad_norm=training_conf.max_grad_norm, - # logging_steps=training_conf.logging_steps, - # save_total_limit=training_conf.save_total_limit, - # evaluation_strategy="steps", - # eval_steps=training_conf.eval_steps, - # save_steps=training_conf.save_steps, - # eval_accumulation_steps=training_conf.eval_accumulation_steps, - # report_to="wandb", - # quantization=training_conf.quantization, - # ) - - # trainer = SFTTrainer( - # model, - # args, - # train_dataset=train, - # eval_dataset=evals, - # data_collator=collate_fn, - # tokenizer=tokenizer, - # compute_metrics=compute_metrics, - # preprocess_logits_for_metrics=preprocess_logits_for_metrics, - # ) - # trainer.train() + assert len(evals) > 0 + trainer = SFTTrainer( + model, + args, + loss_function=training_conf.loss_fn, + train_dataset=train, + eval_dataset=evals, + data_collator=collate_fn, + tokenizer=tokenizer, + compute_metrics=compute_metrics, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + trainer.train() diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py index c3cc512c..8f9ed5ca 100644 --- a/model/supervised_finetuning/utils.py +++ b/model/supervised_finetuning/utils.py @@ -15,6 +15,10 @@ def get_tokenizer(conf): if "galactica" in conf.model_name: tokenizer.add_special_tokens({"pad_token": "", "eos_token": ""}) + elif "GPT-JT" in conf.model_name: + tokenizer.add_special_tokens({"pad_token": tokenizer.eos_token, "sep_token": "<|extratoken_100|>"}) + elif "codegen" in conf.model_name: + tokenizer.add_special_tokens({"pad_token": "<|endoftext|>", "sep_token": "<|endoftext|>"}) additional_special_tokens = ( [] @@ -29,12 +33,6 @@ def get_tokenizer(conf): def get_model(conf, tokenizer): - if not any([x in conf.model_name.lower() for x in SUPPORTED_MODELS]): - raise ValueError( - f"Model {conf.model_name} not supported. Supported models: {SUPPORTED_MODELS}. " - "To include more make sure the masking is done correctly... (decoder only supported for now)" - ) - model = get_specific_model(conf.model_name, conf.cache_dir, conf.quantization) if len(tokenizer) != model.get_input_embeddings().num_embeddings: From 148244455cd29e77cdd3967f44502cab329396f3 Mon Sep 17 00:00:00 2001 From: Sotirios Anagnostidis Date: Fri, 6 Jan 2023 21:29:38 +0100 Subject: [PATCH 5/6] refactor --- .../custom_datasets/dialogue_collator.py | 36 +++++++++---------- model/supervised_finetuning/models/gptj.py | 2 +- model/supervised_finetuning/requirements.txt | 3 ++ 3 files changed, 21 insertions(+), 20 deletions(-) diff --git a/model/supervised_finetuning/custom_datasets/dialogue_collator.py b/model/supervised_finetuning/custom_datasets/dialogue_collator.py index f9e1bb5e..479931f6 100644 --- a/model/supervised_finetuning/custom_datasets/dialogue_collator.py +++ b/model/supervised_finetuning/custom_datasets/dialogue_collator.py @@ -24,26 +24,24 @@ class DialogueDataCollator: flatten_messages = [] label_masks = [] - for messages in features: - assert len(messages) % 2 == 0, "Number of messages must be even" + for feature_one in features: + assert len(feature_one) % 2 == 0, "Number of messages must be even" messages = [ (QA_SPECIAL_TOKENS["Question"] if i % 2 == 0 else "") + x + (QA_SPECIAL_TOKENS["Answer"] if i % 2 == 0 else "") - for i, x in enumerate(messages) + for i, x in enumerate(feature_one) ] # Add a way for the model to terminate generation # When we predict the start of a new expected question, we want to be able to stop generation messages.append(QA_SPECIAL_TOKENS["Question"]) - flatten_messages.append( - self.tokenizer( - "".join(messages), - truncation=True, - max_length=self.max_length, - return_offsets_mapping=True, - ) + flatten_message = self.tokenizer( + "".join(messages), + truncation=True, + max_length=self.max_length, + return_offsets_mapping=True, ) message_change_indices = np.cumsum([len(x) for x in messages[:-1]]) @@ -57,18 +55,19 @@ class DialogueDataCollator: message_indices = list( map( lambda x: next((i for i, val in enumerate(message_change_indices) if val >= x), -2), - list(map(lambda x: x[1], flatten_messages[-1]["offset_mapping"])), + list(map(lambda x: x[1], flatten_message["offset_mapping"])), ) ) label_mask = np.roll(list(map(lambda x: x % 2 == 1, message_indices)), -1, -1) try: label_mask[[i for i in range(len(message_indices)) if message_indices[i] == -2][0] - 1] = True except IndexError: - # an aftermath of padding - pass + # due to truncation, we might not have the last termination token + label_mask[-1] = False label_masks.append(label_mask) - flatten_messages[-1].pop("offset_mapping") + + flatten_messages.append({k: v for k, v in flatten_message.items() if k != "offset_mapping"}) batch = self.tokenizer.pad( flatten_messages, @@ -79,10 +78,9 @@ class DialogueDataCollator: ) dim = batch["input_ids"].shape[-1] - batch["label_masks"] = torch.stack([F.pad(torch.tensor(x), (0, dim - len(x))) for x in label_masks]) - - for k in list(batch.keys()): - if k not in ["input_ids", "attention_mask", "label_masks"]: - batch.pop(k) + batch["label_masks"] = torch.stack( + [F.pad(torch.tensor(x), (0, dim - len(x)), value=False) for x in label_masks] + ) + batch["targets"] = torch.roll(batch["input_ids"], -1, -1) return batch diff --git a/model/supervised_finetuning/models/gptj.py b/model/supervised_finetuning/models/gptj.py index 3cbec3ce..d954c830 100644 --- a/model/supervised_finetuning/models/gptj.py +++ b/model/supervised_finetuning/models/gptj.py @@ -171,7 +171,7 @@ def get_model(model_name, cache_dir, quantization): if quantization is None: model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) elif quantization == "8bit": - print("Loading 8-bit model") + raise ValueError("Loading 8-bit model. Bitsandbytes does not behave so far...") transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJBlock model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) add_adapters(model) diff --git a/model/supervised_finetuning/requirements.txt b/model/supervised_finetuning/requirements.txt index d579468f..798b5950 100644 --- a/model/supervised_finetuning/requirements.txt +++ b/model/supervised_finetuning/requirements.txt @@ -4,3 +4,6 @@ PyYAML==6.0 scikit_learn==1.2.0 torch==1.13.1 transformers==4.25.1 +deepspeed==0.7.7 +mpi4py==3.1.4 +accelerate==0.15.0 \ No newline at end of file From d3952354e2f97e076b72a57184d07cd04b332af3 Mon Sep 17 00:00:00 2001 From: Sotirios Anagnostidis Date: Fri, 6 Jan 2023 22:09:24 +0100 Subject: [PATCH 6/6] pre commits --- .../supervised_finetuning/configs/config.yaml | 2 +- .../configs/zero_config.json | 52 +++++++++++++++ .../supervised_finetuning/models/__init__.py | 12 ++-- model/supervised_finetuning/models/gptj.py | 66 ++++++++++--------- model/supervised_finetuning/requirements.txt | 6 +- model/supervised_finetuning/trainer.py | 31 ++------- model/supervised_finetuning/utils.py | 2 +- 7 files changed, 106 insertions(+), 65 deletions(-) create mode 100644 model/supervised_finetuning/configs/zero_config.json diff --git a/model/supervised_finetuning/configs/config.yaml b/model/supervised_finetuning/configs/config.yaml index fb2bdaa0..97e37121 100644 --- a/model/supervised_finetuning/configs/config.yaml +++ b/model/supervised_finetuning/configs/config.yaml @@ -62,4 +62,4 @@ debug: gradient_accumulation_steps: 1 per_device_train_batch_size: 1 per_device_eval_batch_size: 1 - quantization: \ No newline at end of file + quantization: diff --git a/model/supervised_finetuning/configs/zero_config.json b/model/supervised_finetuning/configs/zero_config.json new file mode 100644 index 00000000..e196d6a0 --- /dev/null +++ b/model/supervised_finetuning/configs/zero_config.json @@ -0,0 +1,52 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 0.1 + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": "auto", + "weight_decay": "auto" + } + }, + "scheduler": { + "type": "WarmupDecayLR", + "params": { + "warmup_min_lr": "auto", + "warmup_max_lr": "auto", + "warmup_num_steps": "auto", + "total_num_steps": "auto" + } + }, + "zero_optimization": { + "stage": 3, + "offload_optimizer": { + "device": "cpu", + "pin_memory": true + }, + "offload_param": { + "device": "cpu", + "pin_memory": true + }, + "overlap_comm": true, + "contiguous_gradients": true, + "reduce_bucket_size": "auto", + "stage3_prefetch_bucket_size": "auto", + "stage3_param_persistence_threshold": "auto", + "sub_group_size": 1e9, + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_gather_16bit_weights_on_model_save": true + }, + "gradient_accumulation_steps": "auto", + "gradient_clipping": 2.0, + "steps_per_print": 2000, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "wall_clock_breakdown": false +} diff --git a/model/supervised_finetuning/models/__init__.py b/model/supervised_finetuning/models/__init__.py index b99053e3..21ddab9d 100644 --- a/model/supervised_finetuning/models/__init__.py +++ b/model/supervised_finetuning/models/__init__.py @@ -1,5 +1,6 @@ from transformers import AutoModelForCausalLM -from .gptj import get_model as get_gptj_model + +# from .gptj import get_model as get_gptj_model SUPPORTED_MODELS = ["galactica", "gpt-j"] @@ -25,7 +26,8 @@ def freeze_top_n_layers(model, target_layers): def get_specific_model(model_name, cache_dir, quantization): - if "gpt-j" in model_name.lower(): - return get_gptj_model(model_name, cache_dir, quantization) - else: - return AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) + return AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) + # if "gpt-j" in model_name.lower(): + # return get_gptj_model(model_name, cache_dir, quantization) + # else: + # return AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) diff --git a/model/supervised_finetuning/models/gptj.py b/model/supervised_finetuning/models/gptj.py index d954c830..b61686c7 100644 --- a/model/supervised_finetuning/models/gptj.py +++ b/model/supervised_finetuning/models/gptj.py @@ -1,12 +1,12 @@ # Taken from https://github.com/sleekmike/Finetune_GPT-J_6B_8-bit/blob/master/gpt-j-6b-8-bit.py -import transformers -from transformers import AutoModelForCausalLM import torch import torch.nn.functional as F +import transformers +from bitsandbytes.functional import dequantize_blockwise, quantize_blockwise from torch import nn -from torch.cuda.amp import custom_fwd, custom_bwd -from bitsandbytes.functional import quantize_blockwise, dequantize_blockwise +from torch.cuda.amp import custom_bwd, custom_fwd +from transformers import AutoModelForCausalLM class FrozenBNBLinear(nn.Module): @@ -19,32 +19,38 @@ class FrozenBNBLinear(nn.Module): self.register_buffer("code", code.requires_grad_(False)) self.adapter = None self.bias = bias - + def forward(self, input): output = DequantizeAndLinear.apply(input, self.weight, self.absmax, self.code, self.bias) if self.adapter: output += self.adapter(input) return output - + @classmethod def from_linear(cls, linear: nn.Linear) -> "FrozenBNBLinear": weights_int8, state = quantize_blockise_lowmemory(linear.weight) return cls(weights_int8, *state, linear.bias) - + def __repr__(self): return f"{self.__class__.__name__}({self.in_features}, {self.out_features})" - - -class DequantizeAndLinear(torch.autograd.Function): + + +class DequantizeAndLinear(torch.autograd.Function): @staticmethod @custom_fwd - def forward(ctx, input: torch.Tensor, weights_quantized: torch.ByteTensor, - absmax: torch.FloatTensor, code: torch.FloatTensor, bias: torch.FloatTensor): + def forward( + ctx, + input: torch.Tensor, + weights_quantized: torch.ByteTensor, + absmax: torch.FloatTensor, + code: torch.FloatTensor, + bias: torch.FloatTensor, + ): weights_deq = dequantize_blockwise(weights_quantized, absmax=absmax, code=code) ctx.save_for_backward(input, weights_quantized, absmax, code) ctx._has_bias = bias is not None return F.linear(input, weights_deq, bias) - + @staticmethod @custom_bwd def backward(ctx, grad_output: torch.Tensor): @@ -55,8 +61,8 @@ class DequantizeAndLinear(torch.autograd.Function): grad_input = grad_output @ weights_deq grad_bias = grad_output.flatten(0, -2).sum(dim=0) if ctx._has_bias else None return grad_input, None, None, None, grad_bias - - + + class FrozenBNBEmbedding(nn.Module): def __init__(self, weight, absmax, code): super().__init__() @@ -65,7 +71,7 @@ class FrozenBNBEmbedding(nn.Module): self.register_buffer("absmax", absmax.requires_grad_(False)) self.register_buffer("code", code.requires_grad_(False)) self.adapter = None - + def forward(self, input, **kwargs): with torch.no_grad(): # note: both quantuized weights and input indices are *not* differentiable @@ -73,41 +79,41 @@ class FrozenBNBEmbedding(nn.Module): output = F.embedding(input, weight_deq, **kwargs) if self.adapter: output += self.adapter(input) - return output - + return output + @classmethod def from_embedding(cls, embedding: nn.Embedding) -> "FrozenBNBEmbedding": weights_int8, state = quantize_blockise_lowmemory(embedding.weight) return cls(weights_int8, *state) - + def __repr__(self): return f"{self.__class__.__name__}({self.num_embeddings}, {self.embedding_dim})" - - -def quantize_blockise_lowmemory(matrix: torch.Tensor, chunk_size: int = 2 ** 20): + + +def quantize_blockise_lowmemory(matrix: torch.Tensor, chunk_size: int = 2**20): assert chunk_size % 4096 == 0 code = None chunks = [] absmaxes = [] flat_tensor = matrix.view(-1) for i in range((matrix.numel() - 1) // chunk_size + 1): - input_chunk = flat_tensor[i * chunk_size: (i + 1) * chunk_size].clone() + input_chunk = flat_tensor[i * chunk_size : (i + 1) * chunk_size].clone() quantized_chunk, (absmax_chunk, code) = quantize_blockwise(input_chunk, code=code) chunks.append(quantized_chunk) absmaxes.append(absmax_chunk) - + matrix_i8 = torch.cat(chunks).reshape_as(matrix) absmax = torch.cat(absmaxes) return matrix_i8, (absmax, code) - - + + def convert_to_int8(model): """Convert linear and embedding modules to 8-bit with optional adapters""" for module in list(model.modules()): for name, child in module.named_children(): if isinstance(child, nn.Linear): print(name, child) - setattr( + setattr( module, name, FrozenBNBLinear( @@ -125,7 +131,7 @@ def convert_to_int8(model): weight=torch.zeros(child.num_embeddings, child.embedding_dim, dtype=torch.uint8), absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1), code=torch.zeros(256), - ) + ), ) @@ -141,7 +147,7 @@ class GPTJModel(transformers.models.gptj.modeling_gptj.GPTJModel): def __init__(self, config): super().__init__(config) convert_to_int8(self) - + class GPTJForCausalLM(transformers.models.gptj.modeling_gptj.GPTJForCausalLM): def __init__(self, config): @@ -171,7 +177,7 @@ def get_model(model_name, cache_dir, quantization): if quantization is None: model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) elif quantization == "8bit": - raise ValueError("Loading 8-bit model. Bitsandbytes does not behave so far...") + raise ValueError("Loading 8-bit model. Use deepspeed instead.") transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJBlock model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir) add_adapters(model) diff --git a/model/supervised_finetuning/requirements.txt b/model/supervised_finetuning/requirements.txt index 798b5950..7d78f36c 100644 --- a/model/supervised_finetuning/requirements.txt +++ b/model/supervised_finetuning/requirements.txt @@ -1,9 +1,9 @@ +accelerate==0.15.0 datasets==2.8.0 +deepspeed==0.7.7 +mpi4py==3.1.4 numpy==1.23.0 PyYAML==6.0 scikit_learn==1.2.0 torch==1.13.1 transformers==4.25.1 -deepspeed==0.7.7 -mpi4py==3.1.4 -accelerate==0.15.0 \ No newline at end of file diff --git a/model/supervised_finetuning/trainer.py b/model/supervised_finetuning/trainer.py index bb77b9c3..cb55131d 100644 --- a/model/supervised_finetuning/trainer.py +++ b/model/supervised_finetuning/trainer.py @@ -5,10 +5,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch from torch import nn -from torch.utils.data import Dataset -from transformers import PreTrainedModel, Trainer, TrainingArguments, get_cosine_schedule_with_warmup -import bitsandbytes as bnb - +from transformers import PreTrainedModel, Trainer, TrainingArguments from utils import get_dataset, get_loss, get_model, get_tokenizer, read_yamls os.environ["WANDB_PROJECT"] = "supervised-finetuning" @@ -39,39 +36,23 @@ class SFTTrainer(Trainer): # By default CrossEntropyLoss ignores padding_index -100, but just in case use our own loss_fct self.loss_fct = get_loss(loss_function) - def create_optimizer_and_scheduler(self, num_training_steps: int): - if self.args.quantization == "8bit": - self.optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001, betas=(0.9, 0.995)) - else: - self.optimizer = torch.optim.AdamW( - self.model.parameters(), lr=self.args.learning_rate, weight_decay=self.args.weight_decay - ) - - self.lr_scheduler = get_cosine_schedule_with_warmup( - self.optimizer, - num_warmup_steps=self.args.warmup_steps, - num_training_steps=self.num_train_steps, - num_cycles=1, - last_epoch=-1, - ) - def compute_loss(self, model, inputs, return_outputs=False): labels_mask = inputs.pop("label_masks") targets = inputs.pop("targets") - outputs = model(**inputs) + outputs = model(input_ids=inputs["input_ids"], attention_mask=inputs.get("attention_mask", None)) loss = self.loss_fct(outputs.get("logits"), targets, mask=labels_mask) return (loss, outputs) if return_outputs else loss def _compute_loss(self, model, inputs): + inputs = self._prepare_inputs(inputs) + labels_mask = inputs.pop("label_masks") targets = inputs.pop("targets") - inputs = self._prepare_inputs(inputs) - - outputs = model(**inputs) + outputs = model(input_ids=inputs["input_ids"], attention_mask=inputs.get("attention_mask", None)) logits = outputs.get("logits") @@ -89,7 +70,7 @@ class SFTTrainer(Trainer): with torch.no_grad(): loss, logits, labels, labels_mask = self._compute_loss(model, inputs) - labels[~labels_mask] = -100 # padding_index + labels[~labels_mask.bool()] = -100 # padding_index loss = loss.mean().detach() diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py index 8f9ed5ca..d6abcff2 100644 --- a/model/supervised_finetuning/utils.py +++ b/model/supervised_finetuning/utils.py @@ -4,10 +4,10 @@ import yaml from custom_datasets import QA_SPECIAL_TOKENS, get_one_dataset from custom_datasets.dialogue_collator import DialogueDataCollator from losses import CrossEntropyLoss +from models import freeze_top_n_layers, get_specific_model from sklearn.model_selection import train_test_split from torch.utils.data import ConcatDataset, Subset from transformers import AutoTokenizer -from models import get_specific_model, SUPPORTED_MODELS, freeze_top_n_layers def get_tokenizer(conf):