[ExecuTorch] Allow setting dtype to bf16 in export_llama

Differential Revision: D61981363 Pull Request resolved: #4985
pytorch · Sep 6, 2024 · 1511fc1 · 1511fc1
1 parent 1d420c9
commit 1511fc1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/examples/models/llama2/export_llama_lib.py b/examples/models/llama2/export_llama_lib.py
@@ -256,9 +256,9 @@ def build_args_parser() -> argparse.ArgumentParser:
         "--dtype-override",
         default="fp32",
         type=str,
-        choices=["fp32", "fp16"],
+        choices=["fp32", "fp16", "bf16"],
         help="Override the dtype of the model (default is the checkpoint dtype)."
-        "Options: fp32, fp16. Please be aware that only some backends support fp16.",
+        "Options: fp32, fp16, bf16. Please be aware that only some backends support fp16 and bf16.",
     )
 
     parser.add_argument(

diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py
@@ -46,6 +46,7 @@ def to_torch_dtype(self) -> torch.dtype:
         mapping = {
             DType.fp32: torch.float32,
             DType.fp16: torch.float16,
+            DType.bf16: torch.bfloat16,
         }
         if self not in mapping:
             raise ValueError(f"Unsupported dtype {self}")