From 6438fdbe2c7b387173b2c604177c5faee3894003 Mon Sep 17 00:00:00 2001 From: Sotirios Anagnostidis Date: Wed, 11 Jan 2023 22:44:20 +0100 Subject: [PATCH] quantization from #582 --- model/supervised_finetuning/trainer.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/model/supervised_finetuning/trainer.py b/model/supervised_finetuning/trainer.py index 88a725b3..9a6bf148 100644 --- a/model/supervised_finetuning/trainer.py +++ b/model/supervised_finetuning/trainer.py @@ -3,9 +3,11 @@ import os from distutils.util import strtobool from typing import Any, Dict, List, Optional, Tuple, Union +import bitsandbytes import torch from torch import nn from transformers import PreTrainedModel, Trainer, TrainingArguments +from transformers.training_args import OptimizerNames from utils import get_dataset, get_loss, get_metrics, get_model, get_tokenizer, read_yamls from functools import partial @@ -141,12 +143,22 @@ if __name__ == "__main__": train, evals, collate_fn = get_dataset(training_conf, tokenizer) metrics, preprocess_fns = get_metrics(training_conf, tokenizer) + optimizer = OptimizerNames.ADAMW_BNB if training_conf.quantization else OptimizerNames.ADAMW_HF + + if training_conf.quantization: + for module in model.modules(): + if isinstance(module, torch.nn.Embedding): + bitsandbytes.optim.GlobalOptimManager.get_instance().register_module_override( + module, "weight", {"optim_bits": 32} + ) + 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, + optim=optimizer, fp16=True, local_rank=training_conf.local_rank, gradient_checkpointing=training_conf.gradient_checkpointing,