diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 407b90184..0ead08c81 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -49,6 +49,7 @@
     get_pretty_name,
     is_deepspeed_available,
     is_torch_version,
+    is_tpu_available,
     is_transformers_available,
     pad_across_processes,
     reduce,
@@ -68,6 +69,9 @@
         DummyScheduler,
     )
 
+if is_tpu_available():
+    import torch_xla.distributed.xla_multiprocessing as xmp
+
 logger = get_logger(__name__)
 
 
@@ -546,6 +550,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU and self.state.fork_launched:
+            model = xmp.MpModelWrapper(model).to(self.device)
         return model
 
     def _prepare_deepspeed(self, *args):
@@ -712,7 +718,7 @@ def prepare_data_loader(self, data_loader):
             num_processes=self.num_processes,
             process_index=self.process_index,
             split_batches=self.split_batches,
-            put_on_device=self.device_placement,
+            put_on_device=self.device_placement if self.distributed_type != DistributedType.TPU else False,
             rng_types=self.rng_types.copy(),
             dispatch_batches=self.dispatch_batches,
         )
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index 79ba6b3f3..eb26c3a60 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -35,7 +35,7 @@
 
 
 if is_tpu_available():
-    import torch_xla.core.xla_model as xm
+    import torch_xla.distributed.parallel_loader as xpl
 
 
 # kwargs of the DataLoader in min version 1.4.0.
@@ -298,10 +298,7 @@ def __init__(self, dataset, device=None, rng_types=None, generator=None, **kwarg
     def __iter__(self):
         if self.rng_types is not None:
             synchronize_rng_states(self.rng_types, self.generator)
-        state = AcceleratorState()
         for batch in super().__iter__():
-            if state.distributed_type == DistributedType.TPU:
-                xm.mark_step()
             yield batch if self.device is None else send_to_device(batch, self.device)
 
 
@@ -408,9 +405,6 @@ def __iter__(self):
                 batch_size += 1
 
             data_slice = slice(state.process_index * batch_size, (state.process_index + 1) * batch_size)
-
-            if state.distributed_type == DistributedType.TPU:
-                xm.mark_step()
             yield slice_tensors(batch, data_slice)
 
     def __len__(self):
@@ -565,15 +559,19 @@ def prepare_data_loader(
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
     if dispatch_batches:
-        return DataLoaderDispatcher(
+        dataloader = DataLoaderDispatcher(
             new_dataset, split_batches=split_batches, batch_sampler=new_batch_sampler, **kwargs
         )
+    else:
+        dataloader = DataLoaderShard(
+            new_dataset,
+            device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
+            batch_sampler=new_batch_sampler,
+            rng_types=rng_types,
+            generator=generator,
+            **kwargs,
+        )
 
-    return DataLoaderShard(
-        new_dataset,
-        device=device if put_on_device else None,
-        batch_sampler=new_batch_sampler,
-        rng_types=rng_types,
-        generator=generator,
-        **kwargs,
-    )
+    if state.distributed_type == DistributedType.TPU:
+        return xpl.MpDeviceLoader(dataloader, device)
+    return dataloader
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 6e09fdae4..891b0762b 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -75,6 +75,7 @@ def __init__(
         self.__dict__ = self._shared_state
         if parse_flag_from_env("USE_CPU"):
             cpu = True
+        self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0)
         if not getattr(self, "initialized", False):
             self.backend = None
             self.deepspeed_plugin = None
@@ -90,7 +91,9 @@ def __init__(
                 self.process_index = xm.get_ordinal()
                 self.local_process_index = xm.get_local_ordinal()
                 self.device = xm.xla_device()
-                self.mixed_precision = "no"
+                self.mixed_precision = (
+                    parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision
+                )
             elif os.environ.get("USE_DEEPSPEED", "false") == "true" and not cpu:
                 assert (
                     is_deepspeed_available()
diff --git a/src/accelerate/test_utils/scripts/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
index a23179bd7..8ed869edc 100644
--- a/src/accelerate/test_utils/scripts/test_script.py
+++ b/src/accelerate/test_utils/scripts/test_script.py
@@ -326,7 +326,8 @@ def main():
     if state.local_process_index == 0:
         print("\n**DataLoader integration test**")
     dl_preparation_check()
-    central_dl_preparation_check()
+    if state.distributed_type != DistributedType.TPU:
+        central_dl_preparation_check()
 
     # Trainings are not exactly the same in DeepSpeed and CPU mode
     if state.distributed_type == DistributedType.DEEPSPEED:
@@ -337,5 +338,10 @@ def main():
     training_check()
 
 
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
 if __name__ == "__main__":
     main()
diff --git a/src/accelerate/test_utils/scripts/test_sync.py b/src/accelerate/test_utils/scripts/test_sync.py
index 2f5cc3722..b58a8c092 100644
--- a/src/accelerate/test_utils/scripts/test_sync.py
+++ b/src/accelerate/test_utils/scripts/test_sync.py
@@ -126,5 +126,10 @@ def main():
         test_distributed_sync(accelerator)
 
 
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
 if __name__ == "__main__":
     main()
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
index d4f0042fb..c74185c5d 100644
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -68,4 +68,5 @@ def __call__(self, index, *args):
             os.environ["LOCAL_RANK"] = str(index)
             os.environ["RANK"] = str(index)
 
+        os.environ["FORK_LAUNCHED"] = str(1)
         self.launcher(*args)
diff --git a/tests/test_tpu.py b/tests/test_tpu.py
index d0b2f0b72..bffa8b8b4 100644
--- a/tests/test_tpu.py
+++ b/tests/test_tpu.py
@@ -24,7 +24,7 @@
 class MultiTPUTester(unittest.TestCase):
     def setUp(self):
         mod_file = inspect.getfile(accelerate.test_utils)
-        self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["test_script.py"])
+        self.test_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_script.py"])
         self.test_dir = os.path.sep.join(inspect.getfile(self.__class__).split(os.path.sep)[:-1])
 
     @require_tpu
diff --git a/tests/xla_spawn.py b/tests/xla_spawn.py
index c65878056..1a07af29c 100644
--- a/tests/xla_spawn.py
+++ b/tests/xla_spawn.py
@@ -78,7 +78,6 @@ def main():
 
     # Patch sys.argv
     sys.argv = [args.training_script] + args.training_script_args + ["--tpu_num_cores", str(args.num_cores)]
-
     xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
 
 


Review comments: body: No, this shouldn't always be done IMO. This is an optimization for Colab/notebooks so we should only do it when we arrive at this code after a `notebook_launcher`.

Or be ready to benchmark the scripts on TPU to show that this does not affect performance in any way, but usually optimizations for memory come at a cost for speed.
diff_hunk: @@ -546,6 +553,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU:
+            model = xmp.MpModelWrapper(model).to(self.device)
body: Same, would like to see some benchmarks on launching a script on a TPU VM to make sure this does not affect speed.
diff_hunk: @@ -565,15 +563,19 @@ def prepare_data_loader(
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
     if dispatch_batches:
-        return DataLoaderDispatcher(
+        dataloader = DataLoaderDispatcher(
             new_dataset, split_batches=split_batches, batch_sampler=new_batch_sampler, **kwargs
         )
+    else:
+        dataloader = DataLoaderShard(
+            new_dataset,
+            device=device if put_on_device else None,
+            batch_sampler=new_batch_sampler,
+            rng_types=rng_types,
+            generator=generator,
+            **kwargs,
+        )
 
-    return DataLoaderShard(
-        new_dataset,
-        device=device if put_on_device else None,
-        batch_sampler=new_batch_sampler,
-        rng_types=rng_types,
-        generator=generator,
-        **kwargs,
-    )
+    if state.distributed_type == DistributedType.TPU:
+        return xpl.MpDeviceLoader(dataloader, device)
body: Here's the docs on [MpModelWrapper](https://pytorch.org/xla/release/1.11/index.html#torch_xla.distributed.xla_multiprocessing.MpModelWrapper), but here's their comment on it:

> This class should be used together with the spawn(…, start_method=’fork’) API to minimize the use of host memory. Instead of creating models on each multiprocessing process, hence replicating the model’s initial host memory, the model is created once at global scope, and then moved into each device inside the spawn() target function. This method has two advantages. First it uses only one copy of the memory pages to host the original model weights, and second it serializes the move of the wrapped model into each device, by lowering the load onto the system memory during the process.

Will perform benchmarks today and tomorrow
diff_hunk: @@ -546,6 +553,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU:
+            model = xmp.MpModelWrapper(model).to(self.device)
body: Yes, and we don't use that spawn method when launching a script on TPU, hence why I don't think we should use it if we don't come from a colab/notebook here :-)
diff_hunk: @@ -546,6 +553,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU:
+            model = xmp.MpModelWrapper(model).to(self.device)
body: Makes sense! Will test and report back if there's a loss at all :) 
diff_hunk: @@ -546,6 +553,8 @@ def prepare_model(self, model):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        if self.distributed_type == DistributedType.TPU:
+            model = xmp.MpModelWrapper(model).to(self.device)
body: ```suggestion
        self.fork_launched = parse_flag_from_env("FORK_LAUNCHED", 0)
```
Might be clearer this way.
diff_hunk: @@ -75,6 +75,7 @@ def __init__(
         self.__dict__ = self._shared_state
         if parse_flag_from_env("USE_CPU"):
             cpu = True
+        self.use_fork = parse_flag_from_env("USE_LAUNCHER", 0)
body: Makes sense!
diff_hunk: @@ -75,6 +75,7 @@ def __init__(
         self.__dict__ = self._shared_state
         if parse_flag_from_env("USE_CPU"):
             cpu = True
+        self.use_fork = parse_flag_from_env("USE_LAUNCHER", 0)

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
body: Thanks for fixing up the TPU support with best practices! Looks great to me! 🔥 
body: @sgugger this PR also modifies the TPU selection in `AcceleratorState` to actually use mixed precision types.

- [FP16 Support](https://github.com/pytorch/xla/issues/1936)
- [BF16 Support](https://pytorch.org/xla/release/1.11/index.html#xla-tensors-and-bfloat16)
body: We're passing 😄 
![image](https://user-images.githubusercontent.com/7831895/173685310-82547d2a-2219-4236-b899-dd5b7599a9ef.png)