drop change in get_vocab_size

Signed-off-by: Alexandros Koumparoulis <[email protected]>
NVIDIA · Dec 9, 2024 · ea51639 · ea51639
1 parent 2f7824d
commit ea51639
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/nemo/lightning/base.py b/nemo/lightning/base.py
@@ -39,12 +39,7 @@ def get_vocab_size(
     from nemo.utils import logging
 
     after = vocab_size
-    multiple = make_vocab_size_divisible_by
-
-    tp_size = config.tensor_model_parallel_size
-    if make_vocab_size_divisible_by % tp_size != 0:
-        multiple *= tp_size
-
+    multiple = make_vocab_size_divisible_by * config.tensor_model_parallel_size
     after = ((after + multiple - 1) // multiple) * multiple
     logging.info(
         f"Padded vocab_size: {after}, original vocab_size: {vocab_size}, dummy tokens:" f" {after - vocab_size}."