pad vocab for CUDA (#10)

cuda likes embed layer (vocab) to be a multiple of 64/128 so we do that

---------

Signed-off-by: Peter Szemraj <peterszemraj+dev@gmail.com>
Co-authored-by: Peter Szemraj <peterszemraj+dev@gmail.com>

Author: pszemraj

annotated_mpnet/modeling/mpnet_for_pretraining.py

@@ -52,11 +52,14 @@ class MPNetForPretraining(nn.Module):
     def __init__(self, args, tokenizer) -> None:
         super().__init__()
 
+        # Use padded_vocab_size if available, otherwise use the tokenizer's vocab_size
+        vocab_size = getattr(args, "padded_vocab_size", tokenizer.vocab_size)
+
         # Let's define the encoder here
         self.args = args
         self.sentence_encoder = SentenceEncoder(
             padding_idx=tokenizer.vocab[tokenizer.pad_token],
-            vocab_size=tokenizer.vocab_size,
+            vocab_size=vocab_size,  # Use the padded vocab size
             num_encoder_layers=args.encoder_layers,
             embedding_dim=args.encoder_embed_dim,
             ffn_embedding_dim=args.encoder_ffn_dim,
@@ -71,15 +74,15 @@ def __init__(self, args, tokenizer) -> None:
             normalize_before=args.normalize_before,
         )
 
-        # Add the language modeling head so that we can do pretraining
+        # Add the language modeling head
         self.lm_head = MPNetLMHead(
             embed_dim=args.encoder_embed_dim,
-            output_dim=tokenizer.vocab_size,
+            output_dim=vocab_size,  # Use the padded vocab size
             activation_fn=args.activation_fn,
             weight=self.sentence_encoder.embed_tokens.weight,
         )
 
-        # Finally initialize the weights according to the guidelines in the original BERT paper
+        # Initialize the weights
         self.apply(init_final_params)
 
     def output_layer(

cli_tools/pretrain_mpnet.py

@@ -139,7 +139,25 @@ def main(args) -> None:
     # Now let's instantiate the tokenizer
     tokenizer = AutoTokenizer.from_pretrained("microsoft/mpnet-base")
 
-    # Instantiate the tensorboard writers here as well
+    # Check and adjust vocab_size parameter for better GPU performance
+    original_vocab_size = tokenizer.vocab_size
+    target_vocab_size = (
+        (original_vocab_size + 127) // 128
+    ) * 128  # Round up to nearest multiple of 128
+
+    if target_vocab_size > original_vocab_size:
+        LOGGER.info(
+            f"Padding model's vocab_size from {original_vocab_size} to {target_vocab_size} "
+            "(div. by 128) for GPU performance"
+        )
+        # Store both sizes in args for reference during conversion
+        args.original_vocab_size = original_vocab_size
+        args.padded_vocab_size = target_vocab_size
+    else:
+        args.original_vocab_size = original_vocab_size
+        args.padded_vocab_size = original_vocab_size
+
+    # Instantiate the tensorboard writers here
     if args.tensorboard_log_dir is not None:
         writers = {
             "train": SummaryWriter(os.path.join(args.tensorboard_log_dir, "train")),