Validate tokenizer and model alignment before training

torchtitan/models/utils.py

@@ -468,3 +468,34 @@ def get_moe_model_nparams_and_flops(
         nparams = nparams - nparams_embedding
 
     return nparams, num_flops_per_token
+
+
+def validate_tokenizer_model_compatibility(
+    tokenizer: "BaseTokenizer | None",
+    model_args: "BaseModelArgs",
+) -> None:
+    """
+    Validate that tokenizer configuration is compatible with model configuration.
+
+    Args:
+        tokenizer: Tokenizer instance to validate. Can be None.
+        model_args: Model arguments object containing configuration to validate against.
+
+    Raises:
+        ValueError: If tokenizer vocab_size exceeds model vocab_size, which would
+            cause index out of bounds errors during training.
+    """
+    if tokenizer is None:
+        return
+
+    if hasattr(model_args, "vocab_size"):
+        tokenizer_vocab_size = tokenizer.get_vocab_size()
+        model_vocab_size = model_args.vocab_size
+        if model_vocab_size < tokenizer_vocab_size:
+            raise ValueError(
+                f"Model vocab_size ({model_vocab_size}) is smaller than "
+                f"tokenizer vocab_size ({tokenizer_vocab_size}). "
+                f"This will cause index out of bounds errors during training. "
+                f"The model's embedding layer must be at least as large as the "
+                f"tokenizer's vocabulary size."
+            )

torchtitan/train.py

@@ -25,6 +25,7 @@
 )
 from torchtitan.config import ConfigManager, JobConfig, TORCH_DTYPE_MAP
 from torchtitan.distributed import ParallelDims, utils as dist_utils
+from torchtitan.models.utils import validate_tokenizer_model_alignment
 from torchtitan.protocols.model_converter import build_model_converters
 from torchtitan.tools import utils
 from torchtitan.tools.logging import init_logger, logger
@@ -134,6 +135,8 @@ def __init__(self, job_config: JobConfig):
         model_args.update_from_config(job_config)
         self.model_args = model_args
 
+        validate_tokenizer_model_alignment(self.tokenizer, model_args)
+
         logger.info(
             f"Building {job_config.model.name} {job_config.model.flavor} with {model_args}"
         )