Removing NullTokenizer from Finetuning scripts

Signed-off-by: Raghav Hrishikeshan Mukundan <rmukundan@nvidia.com>

Author: rhmukundan

scripts/performance/llm/finetune_llama31_405b.py

@@ -108,16 +108,7 @@ def override_recipe_configs(
     )
 
     # data module configs
-    if args.use_hf_tokenizer:
-        recipe.data.tokenizer = hf_tokenizer(HF_MODEL_URI)
-    else:
-        recipe.data.tokenizer = run.Config(
-            get_nmt_tokenizer, library="null", model_name="NullTokenizer", vocab_size=128256
-        )
-        recipe.model.tokenizer = recipe.data.tokenizer
-    if recipe.data.__fn_or_cls__ == SquadDataModule and not isfile_train_pack_metadata(HF_MODEL_URI, recipe.data):
-        # flag is valid only for SquadDataModule
-        recipe.data.force_redownload = True
+    recipe.data.tokenizer = hf_tokenizer(HF_MODEL_URI)
 
     comm_overlap_callback_idx = get_comm_overlap_callback_idx(recipe.trainer.callbacks)
     assert comm_overlap_callback_idx is not None, "MegatronCommOverlapCallback missing. Required for performance."
@@ -227,11 +218,7 @@ def override_recipe_configs(
             assert args.hf_token is not None, "HF token is required for importing checkpoint from HuggingFace"
             exp.add(*import_ckpt_experiment(executor, model(), source=f"hf://{HF_MODEL_URI}"))
         if not SKIP_DATASET_DOWNLOAD:
-            exp.add(
-                *prepare_squad_dataset_experiment(
-                    executor, HF_MODEL_URI, seq_length=4096, nemo_home=args.nemo_home, vocab_size=128256
-                )
-            )
+            exp.add(*prepare_squad_dataset_experiment(executor, HF_MODEL_URI, seq_length=4096, nemo_home=args.nemo_home))
         exp.add(
             recipe,
             executor=executor,

scripts/performance/llm/finetune_llama3_70b.py

@@ -115,16 +115,7 @@ def override_recipe_configs(
     )
 
     # data module configs
-    if args.use_hf_tokenizer:
-        recipe.data.tokenizer = hf_tokenizer(HF_MODEL_URI)
-    else:
-        recipe.data.tokenizer = run.Config(
-            get_nmt_tokenizer, library="null", model_name="NullTokenizer", vocab_size=128256
-        )
-        recipe.model.tokenizer = recipe.data.tokenizer
-    # if recipe.data.__fn_or_cls__ == SquadDataModule and not isfile_train_pack_metadata(HF_MODEL_URI, recipe.data):
-    #     # flag is valid only for SquadDataModule
-    #     recipe.data.force_redownload = True
+    recipe.data.tokenizer = hf_tokenizer(HF_MODEL_URI)
 
     comm_overlap_callback_idx = get_comm_overlap_callback_idx(recipe.trainer.callbacks)
     assert comm_overlap_callback_idx is not None, "MegatronCommOverlapCallback missing. Required for performance."
@@ -234,11 +225,7 @@ def override_recipe_configs(
             assert args.hf_token is not None, "HF token is required for importing checkpoint from HuggingFace"
             exp.add(*import_ckpt_experiment(executor, model(), source=f"hf://{HF_MODEL_URI}"))
         if not SKIP_DATASET_DOWNLOAD:
-            exp.add(
-                *prepare_squad_dataset_experiment(
-                    executor, HF_MODEL_URI, seq_length=4096, nemo_home=args.nemo_home, vocab_size=128256
-                )
-            )
+            exp.add(*prepare_squad_dataset_experiment(executor, HF_MODEL_URI, seq_length=4096, nemo_home=args.nemo_home))
 
         exp.add(
             recipe,

scripts/performance/llm/finetune_llama3_8b.py

@@ -102,16 +102,7 @@ def override_recipe_configs(
     )
 
     # data module configs
-    if args.use_hf_tokenizer:
-        recipe.data.tokenizer = hf_tokenizer(HF_MODEL_URI)
-    else:
-        recipe.data.tokenizer = run.Config(
-            get_nmt_tokenizer, library="null", model_name="NullTokenizer", vocab_size=128256
-        )
-        recipe.model.tokenizer = recipe.data.tokenizer
-    if recipe.data.__fn_or_cls__ == SquadDataModule and not isfile_train_pack_metadata(HF_MODEL_URI, recipe.data):
-        # flag is valid only for SquadDataModule
-        recipe.data.force_redownload = True
+    recipe.data.tokenizer = hf_tokenizer(HF_MODEL_URI)
 
     recipe.optim.config.use_distributed_optimizer = True
     recipe.model.config.disable_parameter_transpose_cache = True
@@ -168,11 +159,7 @@ def override_recipe_configs(
             assert args.hf_token is not None, "HF token is required for importing checkpoint from HuggingFace"
             exp.add(*import_ckpt_experiment(executor, model(), source=f"hf://{HF_MODEL_URI}"))
         if not SKIP_DATASET_DOWNLOAD:
-            exp.add(
-                *prepare_squad_dataset_experiment(
-                    executor, HF_MODEL_URI, seq_length=4096, nemo_home=args.nemo_home, vocab_size=128256
-                )
-            )
+            exp.add(*prepare_squad_dataset_experiment(executor, HF_MODEL_URI, seq_length=4096, nemo_home=args.nemo_home))
         exp.add(
             recipe,
             executor=executor,

scripts/performance/llm/finetune_llama4_e128.py

@@ -108,16 +108,7 @@ def override_recipe_configs(
     )
 
     # data module configs
-    if args.use_hf_tokenizer:
-        recipe.data.tokenizer = hf_tokenizer(HF_MODEL_URI)
-    else:
-        recipe.data.tokenizer = run.Config(
-            get_nmt_tokenizer, library="null", model_name="NullTokenizer", vocab_size=202048
-        )
-        recipe.model.tokenizer = recipe.data.tokenizer
-    # #If you want to force redownload for SquadDataModule, uncomment and adjust the following:
-    # if recipe.data.__fn_or_cls__ == SquadDataModule and not isfile_train_pack_metadata(HF_MODEL_URI, recipe.data):
-    #     SKIP_DATASET_DOWNLOAD = True
+    recipe.data.tokenizer = hf_tokenizer(HF_MODEL_URI)
 
     # Compute dtype configs
     if args.compute_dtype.lower() == "fp8":
@@ -201,16 +192,7 @@ def override_recipe_configs(
             assert args.hf_token is not None, "HF token is required for importing checkpoint from HuggingFace"
             exp.add(*import_ckpt_experiment(executor, model(), source=f"hf://{HF_MODEL_URI}"))
         if not SKIP_DATASET_DOWNLOAD:
-            exp.add(
-                *prepare_squad_dataset_experiment(
-                    executor,
-                    HF_MODEL_URI,
-                    seq_length=4096,
-                    nemo_home=args.nemo_home,
-                    use_hf_tokenizer=args.use_hf_tokenizer,
-                    vocab_size=202048,
-                )
-            )
+            exp.add(*prepare_squad_dataset_experiment(executor, HF_MODEL_URI, seq_length=4096, nemo_home=args.nemo_home))
         exp.add(
             recipe,
             executor=executor,

scripts/performance/utils.py

@@ -91,39 +91,28 @@ def get_nemo_home(nemo_home=None):
     if env_nemo_set:
         return os.environ["NEMO_HOME"]
 
-    raise ValueError("Neither nemo_home argument nor NEMO_HOME environment variable is set")
+    raise ValueError("Neither -nh/--nemo_home argument nor NEMO_HOME environment variable is set")
 
 
-def prepare_squad_dataset(
-    model_name: str, seq_length: int = 2048, nemo_home=None, use_hf_tokenizer=True, vocab_size=None
-):
+def prepare_squad_dataset(model_name: str, seq_length: int = 2048, nemo_home=None):
     """Prepare the SQuAD dataset for fine-tuning.
 
     Args:
         model_name (str): The name of the model
         seq_length (int): The sequence length to use for packing. Defaults to 2048.
         nemo_home: Optional path to NEMO home directory set via args.nemo_home
-        use_hf_tokenizer: Whether to use HuggingFace tokenizer or NullTokenizer
-        vocab_size: Vocabulary size to use when use_hf_tokenizer is False. Required when use_hf_tokenizer is False.
     """
     from pathlib import Path
 
     from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
-    from nemo.collections.common.tokenizers.null_tokenizer import NullTokenizer
     from nemo.collections.llm.gpt.data.packed_sequence import PackedSequenceSpecs
     from nemo.collections.llm.gpt.data.squad import SquadDataModule
 
-    if not use_hf_tokenizer and vocab_size is None:
-        raise ValueError("vocab_size must be provided when use_hf_tokenizer is False")
-
     nemo_home_path = Path(get_nemo_home(nemo_home))
     dataset_root = nemo_home_path / "datasets" / "squad"
     dataset_root.mkdir(parents=True, exist_ok=True)
 
-    if use_hf_tokenizer:
-        tokenizer = AutoTokenizer(pretrained_model_name=model_name)
-    else:
-        tokenizer = NullTokenizer(vocab_size=vocab_size)
+    tokenizer = AutoTokenizer(pretrained_model_name=model_name)
 
     # Configure SquadDataModule with packing specs
     datamodule = SquadDataModule(
@@ -150,14 +139,7 @@ def prepare_squad_dataset(
         raise FileNotFoundError(f"Packed dataset dir not found at {packed_dir}. Dataset download failed")
 
 
-def prepare_squad_dataset_experiment(
-    executor: run.SlurmExecutor,
-    model_name: str,
-    seq_length: int = 2048,
-    nemo_home=None,
-    use_hf_tokenizer=True,
-    vocab_size=None,
-):
+def prepare_squad_dataset_experiment(executor: run.SlurmExecutor, model_name: str, seq_length: int = 2048, nemo_home=None):
     """
     Downloads and prepares the SQuAD dataset for fine-tuning.
     """
@@ -173,7 +155,6 @@ def prepare_squad_dataset_experiment(
             model_name=model_name,
             seq_length=seq_length,
             nemo_home=nemo_home,
-            use_hf_tokenizer=use_hf_tokenizer,
         ),
         dataset_executor,
         "prepare_squad_dataset_exp",