Rollback to 98eaf8f

Author: zhenyan-zhang-meta

torchchat/cli/builder.py

@@ -16,14 +16,14 @@
 import torch._dynamo.config
 import torch._inductor.config
 import torch.distributed as dist
-from torchchat.distributed.logging_utils import SingletonLogger
 
-from torchchat.distributed.utils import (
+from torchchat.distributed.utils import(
     Color as color,
     CUDATrackTime,
-    GPUMemoryMonitor,
     init_distributed,
+    GPUMemoryMonitor,
 )
+from torchchat.distributed.logging_utils import SingletonLogger
 
 from torchchat.model import Model, ModelArgs, ModelType, Transformer, TransformerArgs
 from torchchat.model_config.model_config import resolve_model_config
@@ -36,6 +36,7 @@
 from torchchat.utils.measure_time import measure_time
 from torchchat.utils.quantize import quantize_model
 
+
 from torchtune.models.convert_weights import meta_to_tune
 
 from torchtune.models.llama3_1._position_embeddings import Llama3ScaledRoPE
@@ -187,19 +188,15 @@ def from_args(cls, args: argparse.Namespace) -> "BuilderArgs":
         tp = getattr(args, "tp", 1)
         chpt_from = getattr(args, "chpt_from", "hf")
         sdp_backend_dict = {
-            "math": torch.nn.attention.SDPBackend.MATH,
-            "flash_attention": torch.nn.attention.SDPBackend.FLASH_ATTENTION,
-            "efficient_attention": torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
-            "cudnn_attention": torch.nn.attention.SDPBackend.CUDNN_ATTENTION,
+            'math': torch.nn.attention.SDPBackend.MATH,
+            'flash_attention': torch.nn.attention.SDPBackend.FLASH_ATTENTION,
+            'efficient_attention': torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
+            'cudnn_attention': torch.nn.attention.SDPBackend.CUDNN_ATTENTION,
         }
         attention_backend = sdp_backend_dict[args.attention_backend]
-        if args.device == "cpu" and (
-            args.attention_backend == "efficient_attention"
-            or args.attention_backend == "cudnn_attention"
-        ):
-            print(
-                f"Warning: {args.attention_backend} is not supported on CPU. Using math instead."
-            )
+        if args.device == "cpu" and (args.attention_backend == "efficient_attention"
+                                     or args.attention_backend == "cudnn_attention"):
+            print(f"Warning: {args.attention_backend} is not supported on CPU. Using math instead.")
             attention_backend = torch.nn.attention.SDPBackend.MATH
         return cls(
             checkpoint_dir=checkpoint_dir,
@@ -281,6 +278,8 @@ def __post_init__(self):
         except:
             pass
 
+        self.tokenizer_type = TokenizerType.NONE
+        self.t = None
         return
 
     def is_tiktoken(self) -> bool:
@@ -292,13 +291,6 @@ def is_sentencepiece(self) -> bool:
     def is_hf_tokenizer(self) -> bool:
         return self.tokenizer_type == TokenizerType.HF_TOKENIZER
 
-    def is_tokenizer_none(self) -> bool:
-        if self.tokenizer_type != TokenizerType.NONE:
-            return False
-
-        assert self.t is None, "tokenizer_type is NONE but t is not None"
-        return True
-
     def validate_model(
         self,
         model: Optional[Model],
@@ -307,21 +299,22 @@ def validate_model(
         if model is None:
             return
 
-        if self.is_tokenizer_none():
-            raise RuntimeError(f"no tokenizer was found at {self.tokenizer_path}")
 
         is_tiktoken = self.is_tiktoken()
         is_sentencepiece = self.is_sentencepiece()
         is_hf_tokenizer = self.is_hf_tokenizer()
 
+        if sum([is_tiktoken, is_hf_tokenizer, is_sentencepiece]) != 1:
+            raise RuntimeError(f"no tokenizer was found at {self.tokenizer_path}")
+
         use_tiktoken = model.config.use_tiktoken
         use_hf_tokenizer = model.config.use_hf_tokenizer
         use_sentencepiece = not (use_tiktoken or use_hf_tokenizer)
 
         if (
-            (is_tiktoken and not use_tiktoken)
-            or (is_hf_tokenizer and not use_hf_tokenizer)
-            or (is_sentencepiece and not use_sentencepiece)
+            (is_tiktoken and not use_tiktoken) or
+            (is_hf_tokenizer and not use_hf_tokenizer) or
+            (is_sentencepiece and not use_sentencepiece)
         ):
             raise RuntimeError(
                 "model-specified tokenizer ({}) does not match provided tokenizer ({}) for {}".format(
@@ -519,7 +512,6 @@ def _load_model(builder_args: BuilderArgs) -> Model:
         # AOTI-compoiled model will load its own weights.
         # Release weights here to avoid OOM
         import gc
-
         if hasattr(model, "model"):
             model.model = None
         gc.collect()
@@ -577,7 +569,6 @@ def _initialize_model(
 
             def do_nothing(max_batch_size, max_seq_length):
                 pass
-
             model.setup_caches = do_nothing
 
             model.forward = torch._export.aot_load(
@@ -615,7 +606,6 @@ def do_nothing(max_batch_size, max_seq_length):
 
             def do_nothing(max_batch_size, max_seq_length):
                 pass
-
             model.setup_caches = do_nothing
 
             model.forward = aoti_compiled_model
@@ -667,15 +657,12 @@ def do_nothing(max_batch_size, max_seq_length):
         try:
             model = torch.load(builder_args.snapshot_path, weights_only=False)
         except Exception:
-            raise RuntimeError(
-                f"Failed to load torchchat snapshot {builder_args.snapshot_path}"
-            )
+            raise RuntimeError(f"Failed to load torchchat snapshot {builder_args.snapshot_path}")
         # _active_backend() does not allow DSO & AOTI to be true.
         # Choose either.
         from torchchat.utils.build_utils import set_backend
-
-        set_backend(dso=True, pte=False, aoti_package=False)
-        if model.config != config:
+        set_backend (dso=True, pte=False, aoti_package=False)
+        if (model.config != config):
             raise RuntimeError("loaded model architecture mismatch")
         ##
         ## import all libraries with custom kernels ans custom operators
@@ -693,9 +680,7 @@ def do_nothing(max_batch_size, max_seq_length):
         logger = SingletonLogger.get_logger()
 
         gpu_memory_monitor = GPUMemoryMonitor("cuda")
-        logger.info(
-            f"{color.yellow} {gpu_memory_monitor.get_device_info()}{color.reset}"
-        )
+        logger.info(f"{color.yellow} {gpu_memory_monitor.get_device_info()}{color.reset}")
 
         # Model-level config
         if builder_args.params_table:
@@ -706,16 +691,20 @@ def do_nothing(max_batch_size, max_seq_length):
         config = TransformerArgs.from_params(model_config.transformer_args["text"])
         logger.info(f"Transformer Config: {config}")
 
-        # TODO: Move into head of file after solving circular import
-        from torchchat.distributed.checkpoint_utils import load_model_weights
+        #TODO: Move into head of file after solving circular import
+        from torchchat.distributed.checkpoint_utils import (
+            load_model_weights,
+            )
 
         # Validate pipeline degree
         assert config.n_layers % pp_degree == 0
 
         # Create device mesh
         device_mesh = dist.init_device_mesh(
-            "cuda", (pp_degree, tp_degree), mesh_dim_names=("pp", "tp")
-        )
+            "cuda",
+            (pp_degree, tp_degree),
+            mesh_dim_names=("pp", "tp")
+            )
         tp_mesh = device_mesh["tp"]
         pp_mesh = device_mesh["pp"]
         logger.info(f"Created device mesh: {device_mesh}\n{tp_mesh=}, {pp_mesh=}")
@@ -744,13 +733,7 @@ def do_nothing(max_batch_size, max_seq_length):
         # Load weights
         logger.info(f"Loading weights for {pp_rank=} on {device=}")
         with CUDATrackTime() as timer:
-            load_model_weights(
-                model,
-                builder_args.distribution_path,
-                device,
-                config,
-                builder_args.chpt_from,
-            )
+            load_model_weights(model, builder_args.distribution_path, device, config, builder_args.chpt_from)
 
         logger.info(
             f"{color.green}Total weight loading time: {timer.get_time()} {timer.unit} for rank {rank}{color.reset}"
@@ -764,7 +747,7 @@ def do_nothing(max_batch_size, max_seq_length):
         # lanes.
         # TODO: bump up the lane count
         pipeline_lanes = 1
-        seqlen_prefill = 1024
+        seqlen_prefill=1024
         with device:
             model.setup_caches(1, seqlen_prefill, cache_lanes=pipeline_lanes)