diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index d3b253bce..403587f10 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -137,7 +137,7 @@ def training_function(config, args):
     lr_scheduler = get_linear_schedule_with_warmup(
         optimizer=optimizer,
         num_warmup_steps=100,
-        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+        num_training_steps=(len(train_dataloader) * num_epochs),
     )
 
     # Prepare everything
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index c5ed66cd4..4fad12c72 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -39,6 +39,9 @@ class AcceleratedOptimizer(torch.optim.Optimizer):
     """
     Internal wrapper around a torch optimizer.
 
+    Conditionally will perform `step` and `zero_grad` if gradients should be synchronized when performing gradient
+    accumulation.
+
     Args:
         optimizer (`torch.optim.optimizer.Optimizer`):
             The optimizer to wrap.
diff --git a/src/accelerate/scheduler.py b/src/accelerate/scheduler.py
index cfd6d719f..75695091e 100644
--- a/src/accelerate/scheduler.py
+++ b/src/accelerate/scheduler.py
@@ -16,7 +16,7 @@
 
 import warnings
 
-from .state import AcceleratorState, GradientState
+from .state import AcceleratorState
 
 
 warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler")
@@ -25,10 +25,11 @@
 class AcceleratedScheduler:
     """
     A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful
-    to avoid making a scheduler step too fast when:
+    to avoid making a scheduler step too fast when gradients went overflow and there was no training step (in mixed
+    precision training)
 
-    - gradients went overflow and there was no training step (in mixed precision training)
-    - step was skipped because of gradient accumulation
+    When performing gradient accumulation scheduler lengths should not be changed accordingly, accelerate will always
+    step the scheduler to account for it.
 
     Args:
         scheduler (`torch.optim.lr_scheduler._LRScheduler`):
@@ -47,7 +48,6 @@ def __init__(self, scheduler, optimizers, step_with_optimizer: bool = True, spli
         self.scheduler = scheduler
         self.optimizers = optimizers if isinstance(optimizers, (list, tuple)) else [optimizers]
         self.split_batches = split_batches
-        self.gradient_state = GradientState()
         self.step_with_optimizer = step_with_optimizer
 
     def step(self, *args, **kwargs):


Review comments: body: ```suggestion
    When performing gradient accumulation scheduler lengths should not be changed accordingly, accelerate will
    always step the scheduler to account for it.
```
diff_hunk: @@ -25,10 +25,11 @@
 class AcceleratedScheduler:
     """
     A wrapper around a learning rate scheduler that will only step when the optimizer(s) have a training step. Useful
-    to avoid making a scheduler step too fast when:
+    to avoid making a scheduler step too fast when gradients went overflow and there was no training step (in mixed
+    precision training)
 
-    - gradients went overflow and there was no training step (in mixed precision training)
-    - step was skipped because of gradient accumulation
+    When performing gradient accumulation scheduler lengths should not be changed accordingly, accelerate will handle
+    the right times to step.

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