[ExecuTorch] Allow setting dtype to bf16 in export_llama

Pull Request resolved: #4985

Support creating bf16 PTEs.
ghstack-source-id: 240511594
@exported-using-ghexport

Differential Revision: [D61981363](https://our.internmc.facebook.com/intern/diff/D61981363/)

Author: swolchok

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(

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}")