diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 3b8714911..6880f4626 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -136,7 +136,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/cross_validation.py b/examples/by_feature/cross_validation.py
index f1f54d78d..e380d2d85 100644
--- a/examples/by_feature/cross_validation.py
+++ b/examples/by_feature/cross_validation.py
@@ -143,7 +143,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 4f0311785..2d1cf2f95 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -140,7 +140,7 @@ def tokenize_function(examples):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index d0a044bcd..b69978023 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -111,6 +111,10 @@ def training_function(config, args):
     batch_size = int(config["batch_size"])
     # New Code #
     gradient_accumulation_steps = int(args.gradient_accumulation_steps)
+    if accelerator.distributed_type == DistributedType.TPU and gradient_accumulation_steps > 1:
+        raise NotImplementedError(
+            "Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`"
+        )
 
     metric = evaluate.load("glue", "mrpc")
 
@@ -151,6 +155,7 @@ def training_function(config, args):
             # We use the new `no_sync` context manager to prevent gradient averaging
             # until we want to at the proper step if we happen to be in a distributed setup
             # otherwise it does nothing
+            # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
             if step % gradient_accumulation_steps != 0:
                 # Accumulate gradients locally
                 with accelerator.no_sync(model):
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 7aef8e1ff..65feed51c 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -119,7 +119,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index 788c54f50..c9445ee81 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -121,7 +121,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 3b7b4c8d4..ef2b1388b 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -131,7 +131,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index f13fdfd27..0433ed9a1 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -104,7 +104,7 @@ def tokenize_function(examples):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 00713e886..55dedae5e 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -105,7 +105,7 @@ def training_function(config, args):
 
     # If the batch size is too big we use gradient accumulation
     gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
         gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
         batch_size = MAX_GPU_BATCH_SIZE
 


Review comments: body: No, you should raise an error at the beginning if on TPU with a gradient accumulation steps argument passed > 1.
diff_hunk: @@ -110,7 +110,10 @@ def training_function(config, args):
     seed = int(config["seed"])
     batch_size = int(config["batch_size"])
     # New Code #
-    gradient_accumulation_steps = int(args.gradient_accumulation_steps)
+    if accelerator.distributed_type != DistributedType.TPU:
+        gradient_accumulation_steps = int(args.gradient_accumulation_steps)
+    else:
+        gradient_accumulation_steps = 1
body: Don't need to change here since `gradient_accumulation_steps` will always be 1.
diff_hunk: @@ -196,7 +196,8 @@ def collate_fn(examples):
             if args.with_tracking:
                 total_loss += loss.detach().float()
             accelerator.backward(loss)
-            if step % gradient_accumulation_steps == 0:
+            # We currently do not support nor recommend gradient accumulation on the TPU as bugs were found on the XLA side when running our tests
+            if step % gradient_accumulation_steps == 0 or accelerator.distributed_type == DistributedType.TPU:
body: Same here.
diff_hunk: @@ -146,7 +146,8 @@ def training_function(config, args):
             loss = outputs.loss
             loss = loss / gradient_accumulation_steps
             accelerator.backward(loss)
-            if step % gradient_accumulation_steps == 0:
+            # We currently do not support nor recommend gradient accumulation on the TPU as bugs were found on the XLA side when running our tests
+            if step % gradient_accumulation_steps == 0 or accelerator.distributed_type == DistributedType.TPU:

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