Refactor LLMEdgeManager's to_dtype

Author: Martin Yuan

examples/models/llama/export_llama_lib.py

@@ -588,25 +588,9 @@ def _prepare_for_llama_export(args) -> LLMEdgeManager:
     )
 
     # At this point, the model is loaded in the default fp32.
-
-    # Checkpoint dtype should be lower or equal precision to the dtype override.
+    # override dtype
     checkpoint_dtype = edge_manager.model.checkpoint_dtype
-    if not (
-        checkpoint_dtype == dtype_override.to_torch_dtype()
-        or (
-            checkpoint_dtype == torch.float16
-            and dtype_override.to_torch_dtype() == torch.float32
-        )
-        or (
-            checkpoint_dtype == torch.bfloat16
-            and dtype_override.to_torch_dtype() == torch.float32
-        )
-    ):
-        logging.warning(
-            f"Checkpoint dtype {checkpoint_dtype} precision is higher than dtype override {dtype_override.to_torch_dtype()}."
-        )
-
-    edge_manager.model = edge_manager.model.to(dtype=dtype_override.to_torch_dtype())
+    edge_manager.to_dtype(dtype_override)
 
     # We want to quantize (in the source transforms) the weights of the model
     # in the checkpoint dtype.

extension/llm/export/builder.py

@@ -147,7 +147,25 @@ def to_dtype(self, dtype_override: Optional[DType]) -> "LLMEdgeManager":
         assert not dtype_override or isinstance(
             dtype_override, DType
         ), "Override dtype needs to be of type <DType>"
-        if dtype_override is not None and dtype_override != self.dtype:
+
+        # Checkpoint dtype should be lower or equal precision to the dtype override.
+        checkpoint_dtype = self.model.checkpoint_dtype
+        if not (
+                checkpoint_dtype == dtype_override.to_torch_dtype()
+                or (
+                        checkpoint_dtype == torch.float16
+                        and dtype_override.to_torch_dtype() == torch.float32
+                )
+                or (
+                        checkpoint_dtype == torch.bfloat16
+                        and dtype_override.to_torch_dtype() == torch.float32
+                )
+        ):
+            logging.warning(
+                f"Checkpoint dtype {checkpoint_dtype} precision is higher than dtype override {dtype_override.to_torch_dtype()}."
+            )
+
+        if dtype_override != self.dtype:
             torch_dtype = dtype_override.to_torch_dtype()
             logging.info(f"model.to {torch_dtype}")
             self.model = self.model.to(dtype=torch_dtype)