From ef02693ac9b4705b4ba026a05436709a479dcec7 Mon Sep 17 00:00:00 2001 From: Sotirios Anagnostidis Date: Fri, 6 Jan 2023 18:24:28 +0100 Subject: [PATCH] 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())