Refactor LLMEdgeManager's to_dtype (#9542)

Summary:
Small refactor to aggregate dtype override in to_dtype method.


Test Plan: CI

Reviewed By: billmguo

Differential Revision: D71736608

Pulled By: iseeyuan

Author: Martin Yuan

examples/models/llama/export_llama_lib.py

@@ -618,25 +618,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,26 @@ 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.
+        if hasattr(self.model, "checkpoint_dtype"):
+            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)