unify model construction ppl

Author: Gasoonjia

torchchat/cli/builder.py

@@ -240,12 +240,7 @@ def validate_model(
 
         is_tiktoken = self.is_tiktoken
         is_sentencepiece = self.is_sentencepiece
-        text_args = model.config.transformer_args.get("text")
-        if text_args is None:
-            # TODO: Will be refactored: Currently, the only model that doesn't have text in transfomer_args is Flamingo
-            use_tiktoken = model.config.model_type == ModelType.Flamingo
-        else:
-            use_tiktoken = text_args.use_tiktoken
+        use_tiktoken = model.config.use_tiktoken
 
         if not (is_tiktoken == use_tiktoken) or not (is_sentencepiece != use_tiktoken):
             raise RuntimeError(

torchchat/generate.py

@@ -27,12 +27,6 @@
 
 from PIL import Image
 
-# torchtune model definition dependencies
-from torchtune.data import Message
-from torchtune.generation._generation import sample as tune_sample
-from torchtune.models.llama3 import llama3_tokenizer
-from torchtune.training import set_default_dtype
-
 from torchchat.cli.builder import (
     _initialize_model,
     _initialize_tokenizer,
@@ -43,6 +37,12 @@
 from torchchat.utils.build_utils import device_sync, set_precision
 from torchchat.utils.device_info import get_device_info
 
+# torchtune model definition dependencies
+from torchtune.data import Message
+from torchtune.generation._generation import sample as tune_sample
+from torchtune.models.llama3 import llama3_tokenizer
+from torchtune.training import set_default_dtype
+
 
 class _ChatFormatter(ABC):
     def __init__(self, tokenizer):
@@ -790,16 +790,12 @@ def chat(
 
         # This is a hack to get around the fact that different models have different ways to record their max_seq_length and might be wrong
         # TODO: unify the max_seq_length config representation.
-        if generator_args.is_torchtune_model:
-            max_seq_length = self.model.config.transformer_args.get("text", {}).get(
-                "max_seq_len", 2048
-            )
-        elif generator_args.chat_mode:
-            if (
-                max_seq_length := self.model.config.transformer_args.get("text", None)
-                is None
-            ):
-                max_seq_length = 2048
+        text_transformer_args = getattr(self.model.model, "config", None)
+        max_seq_length = (
+            text_transformer_args.max_seq_length if text_transformer_args else 2048
+        )
+
+        if generator_args.chat_mode:
             print(
                 f"Entering Chat Mode. Will continue chatting back and forth with the language model until the models max context length of {max_seq_length} tokens is hit or until the user says /bye"
             )
@@ -809,15 +805,9 @@ def chat(
             if get_system_prompt == "y" or get_system_prompt == "Y":
                 self.system_prompt = input("What is your system prompt? \n")
 
-        else:
-            text_transformer_args = self.model.config.transformer_args.get("text", None)
+        elif not generator_args.is_torchtune_model:
             max_seq_length = min(
-                encoded.size(0) + generator_args.max_new_tokens,
-                (
-                    text_transformer_args.block_size
-                    if text_transformer_args is not None
-                    else 2048
-                ),
+                encoded.size(0) + generator_args.max_new_tokens, max_seq_length
             )
 
         max_seq_length = (

torchchat/model.py

@@ -164,49 +164,49 @@ def from_params(cls, params):
 @dataclass
 class ModelArgs:
     model_type: ModelType
-    transformer_args: Dict[str, Union[Dict, TransformerArgs]]
+    transformer_args: Dict[str, Dict[str, Any]]
+    use_tiktoken: bool
 
     def __init__(
         self,
-        transformer_args: Union[TransformerArgs, Dict[str, TransformerArgs]],
+        transformer_args: Dict[str, Dict[str, Any]],
         model_type: ModelType = ModelType.TextOnly,
+        use_tiktoken: bool = False,
     ) -> None:
         self._sanity_check(transformer_args, model_type)
 
         self.model_type = model_type
-        if isinstance(transformer_args, TransformerArgs):
-            assert model_type == ModelType.TextOnly
-            self.transformer_args = {"text": transformer_args}
-        else:
-            self.transformer_args = transformer_args
+        self.transformer_args = transformer_args
+
+        # Model-level attributes
+        self.use_tiktoken = use_tiktoken
 
     def _sanity_check(
         self,
-        transformer_args: Union[TransformerArgs, Dict[str, TransformerArgs]],
+        transformer_args: Dict[str, Dict[str, Any]],
         model_type: ModelType,
     ) -> None:
-        assert isinstance(model_type, ModelType)
-        assert isinstance(transformer_args, (TransformerArgs, dict))
+        assert isinstance(model_type, ModelType), model_type
+        assert isinstance(transformer_args, dict)
 
     @classmethod
     def from_params(cls, params_path):
         with open(params_path, "r") as f:
             loaded_params = json.loads(f.read())
-
-        try:
-            # try to interpret as a single transformer config
-            transformer_args: Dict[str, TransformerArgs] = {}
-            transformer_args["text"] = TransformerArgs.from_params(loaded_params)
-            if (model_type := loaded_params.get("model_type", None)) is None:
-                model_type = ModelType.TextOnly
-
-        except TypeError:
-            # try to interpret as a dict of transformer configs
-            model_type = ModelType(loaded_params["model_type"])
+        
+        if (model_type_name := loaded_params.get("model_type", None)) is None:
+            # The model params is in the transformer_args format
+            # set the model_type to TextOnly and reformat the params
+            model_type = ModelType.TextOnly
+            transformer_args = {"text": {"config": loaded_params}}
+        else:
+            model_type = ModelType(model_type_name)
             transformer_args = {
                 k: v for k, v in loaded_params.items() if k != "model_type"
             }
-        return cls(transformer_args, model_type)
+
+        use_tiktoken = loaded_params.get("use_tiktoken", False)
+        return cls(transformer_args, model_type, use_tiktoken)
 
     @classmethod
     def from_table(cls, name: str):
@@ -304,10 +304,8 @@ def build_model(self) -> nn.Module:
         recipe = ModelRecipe.get_recipe(self.config.model_type)
         modules = {}
         for name, module_class in recipe.modules.items():
-            if isinstance(config_args := self.config.transformer_args[name], dict):
-                modules[name] = module_class(**config_args)
-            else:
-                modules[name] = module_class(config_args)
+            config_args = self.config.transformer_args[name]
+            modules[name] = module_class(**config_args)
 
         return recipe.fusion_class(**modules)
 
@@ -399,8 +397,9 @@ def reset_caches(self):
 
 
 class Transformer(nn.Module):
-    def __init__(self, config: TransformerArgs) -> None:
+    def __init__(self, config: Dict[str, Any]) -> None:
         super().__init__()
+        config = TransformerArgs.from_params(config)
         self.config = config
         layers_per_stage = config.n_layers // config.n_stages
 

torchchat/model_params/Meta-Llama-3.1-70B-Tune.json

@@ -1,5 +1,6 @@
 {
     "model_type": "llama3_1",
+    "use_tiktoken": true,
     "text": {
         "vocab_size": 128256,
         "num_layers": 80,

torchchat/model_params/Meta-Llama-3.1-8B-Tune.json

@@ -1,5 +1,6 @@
 {
     "model_type": "llama3_1",
+    "use_tiktoken": true,
     "text": {
         "vocab_size": 128256,
         "num_layers": 32,