Pass sharing args from CLI

Author: borzunov

dalle_pytorch/dalle_pytorch.py

@@ -326,7 +326,9 @@ def __init__(
         stable = False,
         sandwich_norm = False,
         shift_tokens = True,
-        rotary_emb = True
+        rotary_emb = True,
+        shared_attn_ids = None,
+        shared_ff_ids = None,
     ):
         super().__init__()
         assert isinstance(vae, (DiscreteVAE, OpenAIDiscreteVAE, VQGanVAE)), 'vae must be an instance of DiscreteVAE'
@@ -374,7 +376,9 @@ def __init__(
             stable = stable,
             sandwich_norm = sandwich_norm,
             shift_tokens = shift_tokens,
-            rotary_emb = rotary_emb
+            rotary_emb = rotary_emb,
+            shared_attn_ids = shared_attn_ids,
+            shared_ff_ids = shared_ff_ids,
         )
 
         self.stable = stable
@@ -417,7 +421,7 @@ def generate_texts(
             text_tokens = torch.tensor([[0]]).cuda()
         else:
             text_tokens = torch.tensor(tokenizer.tokenizer.encode(text)).cuda().unsqueeze(0)
-   
+
         for _ in range(text_tokens.shape[1], text_seq_len):
             device = text_tokens.device
 
@@ -443,9 +447,9 @@ def generate_texts(
             filtered_logits = top_k(logits, thres = filter_thres)
             probs = F.softmax(filtered_logits / temperature, dim = -1)
             sample = torch.multinomial(probs, 1)
- 
+
             text_tokens = torch.cat((text_tokens, sample), dim=-1)
-    
+
         padding_tokens = set(np.arange(self.text_seq_len) + (self.num_text_tokens - self.text_seq_len))
         texts = [tokenizer.tokenizer.decode(text_token, pad_tokens=padding_tokens) for text_token in text_tokens]
         return text_tokens, texts

train_dalle.py

@@ -46,9 +46,9 @@
                     help='path to your folder of images and text for learning the DALL-E')
 
 parser.add_argument(
-    '--wds', 
-    type = str, 
-    default='', 
+    '--wds',
+    type = str,
+    default='',
     help = 'Comma separated list of WebDataset (1) image and (2) text column names. Must contain 2 values, e.g. img,cap.'
 )
 
@@ -134,6 +134,10 @@
 
 model_group.add_argument('--rotary_emb', help = 'Use rotary embeddings', action = 'store_true')
 
+model_group.add_argument('--shared_attn_ids', default = None, type = str, help = 'Comma separated list of shared attention layer ids. Default: sharing is disabled')
+
+model_group.add_argument('--shared_ff_ids', default = None, type = str, help = 'Comma separated list of shared feed forward layer ids. Default: sharing is disabled')
+
 args = parser.parse_args()
 
 # helpers
@@ -191,6 +195,8 @@ def cp_path_to_dir(cp_path, tag):
 ROTARY_EMB = args.rotary_emb
 
 ATTN_TYPES = tuple(args.attn_types.split(','))
+SHARED_ATTN_IDS = tuple(args.shared_attn_ids.split(',')) if exists(args.shared_attn_ids) else None
+SHARED_FF_IDS = tuple(args.shared_ff_ids.split(',')) if exists(args.shared_ff_ids) else None
 
 DEEPSPEED_CP_AUX_FILENAME = 'auxiliary.pt'
 
@@ -303,6 +309,8 @@ def cp_path_to_dir(cp_path, tag):
         stable=STABLE,
         shift_tokens=SHIFT_TOKENS,
         rotary_emb=ROTARY_EMB,
+        shared_attn_ids=SHARED_ATTN_IDS,
+        shared_ff_ids=SHARED_FF_IDS,
     )
     resume_epoch = 0
 
@@ -368,7 +376,7 @@ def filter_dataset(item): # For e.g. C@H which (rarely) has no caption available
         if myimg not in item:
             return False
         return True
-	
+
     w_dataset = wds.WebDataset(DATASET, handler=wds.warn_and_continue)
     filtered_dataset = w_dataset.select(filter_dataset)
     ds = filtered_dataset.map_dict(**image_text_mapping).map_dict(**image_mapping).to_tuple(mycap, myimg).batched(BATCH_SIZE, partial=True)
@@ -600,7 +608,7 @@ def save_model(path, epoch=0):
 
         if i % SAVE_EVERY_N_STEPS == 0:
             save_model(DALLE_OUTPUT_FILE_NAME, epoch=epoch)
-	
+
         if i % 100 == 0:
             if distr_backend.is_root_worker():
                 sample_text = text[:1]
@@ -633,7 +641,7 @@ def save_model(path, epoch=0):
         distr_scheduler.step(avg_loss)
 
     save_model(DALLE_OUTPUT_FILE_NAME, epoch=epoch)
-    
+
     if distr_backend.is_root_worker():
         # save trained model to wandb as an artifact every epoch's end