quantization

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
+  quantization:

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)
+        return AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)

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}")
 

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()

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())