diff --git a/model/supervised_finetuning/configs/config.yaml b/model/supervised_finetuning/configs/config.yaml
index 2ceb1388..97e37121 100644
--- a/model/supervised_finetuning/configs/config.yaml
+++ b/model/supervised_finetuning/configs/config.yaml
@@ -22,6 +22,7 @@ defaults:
   loss_fn: CrossEntropyLoss
   eval_size:
   log_dir: "base"
+  quantization:
 
 galactica-125:
   learning_rate: 5e-5
@@ -58,3 +59,7 @@ codegen:
 debug:
   eval_steps: 20
   eval_size: 100
+  gradient_accumulation_steps: 1
+  per_device_train_batch_size: 1
+  per_device_eval_batch_size: 1
+  quantization:
diff --git a/model/supervised_finetuning/custom_datasets/__init__.py b/model/supervised_finetuning/custom_datasets/__init__.py
index 60462186..5706bfa7 100644
--- a/model/supervised_finetuning/custom_datasets/__init__.py
+++ b/model/supervised_finetuning/custom_datasets/__init__.py
@@ -16,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):
@@ -59,7 +59,6 @@ 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":
diff --git a/model/supervised_finetuning/custom_datasets/dialogue_collator.py b/model/supervised_finetuning/custom_datasets/dialogue_collator.py
index 535b9007..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,11 +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])
-
-        # why the fuck?
-        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/__init__.py b/model/supervised_finetuning/models/__init__.py
new file mode 100644
index 00000000..21ddab9d
--- /dev/null
+++ b/model/supervised_finetuning/models/__init__.py
@@ -0,0 +1,33 @@
+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):
+    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
new file mode 100644
index 00000000..b61686c7
--- /dev/null
+++ b/model/supervised_finetuning/models/gptj.py
@@ -0,0 +1,187 @@
+# Taken from https://github.com/sleekmike/Finetune_GPT-J_6B_8-bit/blob/master/gpt-j-6b-8-bit.py
+
+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_bwd, custom_fwd
+from transformers import AutoModelForCausalLM
+
+
+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":
+        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)
+    else:
+        raise ValueError(f"Unknown quantization {quantization}")
+
+    return model
diff --git a/model/supervised_finetuning/requirements.txt b/model/supervised_finetuning/requirements.txt
index d579468f..7d78f36c 100644
--- a/model/supervised_finetuning/requirements.txt
+++ b/model/supervised_finetuning/requirements.txt
@@ -1,4 +1,7 @@
+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
diff --git a/model/supervised_finetuning/trainer.py b/model/supervised_finetuning/trainer.py
index a66f2d11..cb55131d 100644
--- a/model/supervised_finetuning/trainer.py
+++ b/model/supervised_finetuning/trainer.py
@@ -5,7 +5,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import nn
-from transformers import PreTrainedModel, Trainer, TrainingArguments, get_cosine_schedule_with_warmup
+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"
@@ -36,35 +36,26 @@ 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 fetch_scheduler(self):
-        return 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"), 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")
-
         inputs = self._prepare_inputs(inputs)
 
-        outputs = model(**inputs)
+        labels_mask = inputs.pop("label_masks")
+        targets = inputs.pop("targets")
+
+        outputs = model(input_ids=inputs["input_ids"], attention_mask=inputs.get("attention_mask", None))
 
         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
@@ -79,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()
 
@@ -106,6 +97,8 @@ def argument_parsing(notebook=False, notebook_args=None):
     else:
         args, remaining = parser.parse_known_args()
 
+    print(args)
+
     # Config from YAML
     conf = {}
     configs = read_yamls("./configs")
diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py
index 7147dcb9..d6abcff2 100644
--- a/model/supervised_finetuning/utils.py
+++ b/model/supervised_finetuning/utils.py
@@ -4,9 +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 AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer
 
 
 def get_tokenizer(conf):
@@ -32,8 +33,7 @@ def get_tokenizer(conf):
 
 
 def get_model(conf, tokenizer):
-
-    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."
@@ -92,31 +92,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())