run linter & fix gradient ckpting

open_flamingo/src/factory.py

@@ -181,6 +181,8 @@ def check_embedding_fns(lang_model):
     if not has_fn(lang_model, "get_input_embeddings"):
         if hasattr_recursive(lang_model, "transformer.wte"):  # MPT
             lang_model.get_input_embeddings = lambda: lang_model.transformer.wte
+        elif hasattr_recursive(lang_model, "model.decoder.embed_tokens"):  # OPT
+            lang_model.get_input_embeddings = lambda: lang_model.decoder.embed_tokens
         else:
             raise ValueError(
                 "We require the language encoder to have a get_input_embeddings method but we couldn't determine the name of the input embeddings attribute. Please supply this manually in factory.py."
@@ -191,6 +193,10 @@ def check_embedding_fns(lang_model):
             lang_model.set_input_embeddings = lambda x: setattr_recursive(
                 lang_model, "transformer.wte", x
             )
+        elif hasattr_recursive(lang_model, "model.decoder.embed_tokens"):  # OPT
+            lang_model.set_input_embeddings = lambda x: setattr_recursive(
+                lang_model, "model.decoder.embed_tokens", x
+            )
         else:
             raise ValueError(
                 "We require the language encoder to have a set_input_embeddings method but we couldn't determine the name of the input embeddings attribute. Please supply this manually in factory.py."
@@ -211,13 +217,14 @@ def check_embedding_fns(lang_model):
             )
         else:
             raise ValueError(
-                "We require the language encoder to have a get_output_embeddings method but we couldn't determine the name of the output embeddings attribute. Please supply this manually in factory.py."
+                "We require the language encoder to have a set_output_embeddings method but we couldn't determine the name of the output embeddings attribute. Please supply this manually in factory.py."
             )
 
+
 def has_fn(model, fn_name):
     """Try to call the fn_name function on the model"""
     try:
         getattr(model, fn_name)()
         return True
     except:
-        return False
+        return False

open_flamingo/src/vlm.py

@@ -75,7 +75,6 @@ def __init__(
         self.lang_model.set_output_embeddings(out_embeds)
 
         # gradient checkpointing
-        self._use_gradient_checkpointing = gradient_checkpointing
         self.vision_tokenizer._use_gradient_checkpointing = gradient_checkpointing
 
     def forward(
@@ -513,8 +512,9 @@ def __init__(
             pad_token_id=pad_token_id,
             gradient_checkpointing=gradient_checkpointing,
         )
-        self.lang_model._use_gradient_checkpointing = gradient_checkpointing
         self.decoder_layers_attr_name = decoder_layers_attr_name
+        for block in getattr_recursive(self.lang_model, self.decoder_layers_attr_name):
+            block._use_gradient_checkpointing = gradient_checkpointing
         assert (
             self.vis_embedding_dim == self.lang_embedding_dim
         ), "To place visual tokens direclty in the language stream, the visual and language tokens need to be the same dim."

open_flamingo/train/train_utils.py

@@ -68,8 +68,6 @@ def train_one_epoch(
         losses_to_log = {}
         batch_metadata_to_log = {}
         for dataset_ix, (images, (input_ids, attention_mask)) in enumerate(batches):
-            print(">> Dataset: ", datasets[dataset_ix].name, "Step: ", step_num)
-
             # unpack the batch and move to device
             images = images.to(device_id, dtype=cast_dtype, non_blocking=True)
             input_ids = input_ids.to(device_id, non_blocking=True)