quantization

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