diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2dec26a98..596998f1e 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -235,6 +235,13 @@ def __init__(
             **kwargs,
         )
 
+        if (
+            (mixed_precision != "bf16")
+            and getattr(self.state, "downcast_bfloat", False)
+            and (self.state.distributedType != DistributedType.TPU)
+        ):
+            raise ValueError("Can only use `downcast_bf16` when using `mixed_precision='bf16'` and on a TPU")
+
         if gradient_accumulation_steps > 1:
             if self.state.distributed_type == DistributedType.TPU:
                 raise NotImplementedError(
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index 0051d899a..8a4f38f9e 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -260,11 +260,18 @@ def get_cluster_input():
     else:
         mixed_precision = "no"
 
+    downcast_bf16 = "no"
+    if distributed_type == DistributedType.TPU and mixed_precision == "bf16":
+        downcast_bf16 = _ask_field(
+            "Should `torch.float` be cast as `bfloat16` and `torch.double` remain `float32` on TPUs?", default="no"
+        )
+
     return ClusterConfig(
         compute_environment=ComputeEnvironment.LOCAL_MACHINE,
         distributed_type=distributed_type,
         num_processes=num_processes,
         mixed_precision=mixed_precision,
+        downcast_bf16=downcast_bf16,
         machine_rank=machine_rank,
         num_machines=num_machines,
         main_process_ip=main_process_ip,
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 1ef68200b..11ca0ca90 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -143,6 +143,8 @@ class ClusterConfig(BaseConfig):
     deepspeed_config: dict = None
     # args for fsdp
     fsdp_config: dict = None
+    # args for TPU
+    downcast_bf16: bool = False
 
     def __post_init__(self):
         if self.deepspeed_config is None:
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index cffe9990b..75c6c7b1c 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -34,6 +34,7 @@
     get_launch_prefix,
     is_deepspeed_available,
     is_sagemaker_available,
+    patch_environment,
 )
 from accelerate.utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
 from accelerate.utils.dataclasses import SageMakerDistributedType
@@ -198,6 +199,11 @@ def launch_command_parser(subparsers=None):
         default=None,
         help="The name of the main function to be executed in your script (only for TPU training).",
     )
+    parser.add_argument(
+        "--downcast_bf16",
+        action="store_true",
+        help="Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32",
+    )
     parser.add_argument(
         "-m",
         "--module",
@@ -425,9 +431,19 @@ def deepspeed_launcher(args):
 def tpu_launcher(args):
     import torch_xla.distributed.xla_multiprocessing as xmp
 
+    current_env = {}
+
     if args.no_python:
         raise ValueError("--no_python cannot be used with TPU launcher")
 
+    if args.mixed_precision == "bf16":
+        if args.downcast_bf16:
+            current_env["XLA_USE_BF16"] = "0"
+            current_env["XLA_DOWNCAST_BF16"] = "1"
+        else:
+            current_env["XLA_USE_BF16"] = "1"
+            current_env["XLA_DOWNCAST_BF16"] = "0"
+
     if args.module:
         mod_name = args.training_script
     else:
@@ -447,7 +463,8 @@ def tpu_launcher(args):
     sys.argv = [mod.__file__] + args.training_script_args
 
     main_function = getattr(mod, args.main_training_function)
-    xmp.spawn(PrepareForLaunch(main_function), args=(), nprocs=args.num_processes)
+    with patch_environment(**current_env):
+        xmp.spawn(PrepareForLaunch(main_function), args=(), nprocs=args.num_processes)
 
 
 def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 4482341f3..3fcf7d2e2 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -115,7 +115,14 @@ def __init__(
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
                 if mixed_precision == "bf16":
-                    os.environ["XLA_USE_BF16"] = str(1)
+                    if os.environ.get("DOWNCAST_BF16"):
+                        os.environ["XLA_USE_BF16"] = str(0)
+                        os.environ["XLA_DOWNCAST_BF16"] = str(1)
+                        self.downcast_bfloat = True
+                    else:
+                        os.environ["XLA_USE_BF16"] = str(1)
+                        os.environ["XLA_DOWNCAST_BF16"] = str(0)
+                        self.downcast_bfloat = False
                 self.mixed_precision = mixed_precision
             elif os.environ.get("USE_DEEPSPEED", "false") == "true" and not cpu:
                 assert (


Review comments: body: Very much against adding a new Accelerator parameter for such a niche feature. It should be in the config only.
diff_hunk: @@ -126,6 +126,9 @@ class Accelerator:
             accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
             A path to a directory for storing logs of locally-compatible loggers.
+        downcast_bf16 (`bool`, *optional*):
+            If set to `True` and using `mixed_precision="bf16"` then `torch.float` will become `bfloat16` and
+            `torch.double` will remain `float32` on TPUs.
body: Doesn't seem to be able to be None, so I'm confused about this change in typing.
diff_hunk: @@ -156,8 +159,9 @@ def __init__(
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = None,
+        downcast_bf16: Optional[bool] = False,
         dispatch_batches: Optional[bool] = None,
-        step_scheduler_with_optimizer: bool = True,
+        step_scheduler_with_optimizer: Optional[bool] = True,
body: Okay, will adjust on Monday!

(This should be in the release after this next one anyways)
diff_hunk: @@ -126,6 +126,9 @@ class Accelerator:
             accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):
             A path to a directory for storing logs of locally-compatible loggers.
+        downcast_bf16 (`bool`, *optional*):
+            If set to `True` and using `mixed_precision="bf16"` then `torch.float` will become `bfloat16` and
+            `torch.double` will remain `float32` on TPUs.

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._