refactor: rename optimizer step callback hooks to before and after

- rename on_optimizer_step to before_optimizer_step across the callback protocol, built-in callbacks, and trainer call site
- rename on_after_optimizer_step to after_optimizer_step for the symmetric post-step hook
- document the hook pair and the checkpoint save location in developer and training guides

docs/developer/architecture.md

@@ -746,13 +746,14 @@ classDiagram
             +on_epoch_end(context)
             +on_batch_begin(context)
             +on_batch_end(context)
-            +on_optimizer_step(context)
+            +before_optimizer_step(context)
+            +after_optimizer_step(context)
             +on_error(context)
         }
 
         class GradientClippingCallback {
             +Optional[float] max_grad_norm
-            +on_optimizer_step(context)
+            +before_optimizer_step(context)
         }
 
         class GradientCheckpointingCallback {
@@ -767,7 +768,7 @@ classDiagram
             +bool weight_only
             +Callable save_extra_fn
             -_save_checkpoint(context)
-            +on_batch_end(context)
+            +after_optimizer_step(context)
             +on_train_end(context)
             +on_error(context)
             +save_extra(context) dict
@@ -779,7 +780,7 @@ classDiagram
             +IO file
             +tqdm progress_bar
             +on_epoch_begin(context)
-            +on_optimizer_step(context)
+            +before_optimizer_step(context)
             +on_epoch_end(context)
         }
 
@@ -788,7 +789,7 @@ classDiagram
             +int save_interval
             +List[str] metrics
             +int val_step
-            +on_optimizer_step(context)
+            +before_optimizer_step(context)
             +on_epoch_end(context)
             +on_train_end(context)
             +on_error(context)

docs/developer/internals.md

@@ -118,12 +118,13 @@ on_train_begin
         on_batch_end
 
         if executor.sync_gradients:
-          on_optimizer_step
+          before_optimizer_step
           optimizer.step()
           strategy.on_optimizer_step()
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()
+          after_optimizer_step
     on_epoch_end
 on_train_end
 ```
@@ -139,8 +140,9 @@ Strategy metrics are detached and converted to Python `float` values before the
 | `on_train_begin` | Before training starts | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
-| `on_optimizer_step` | Every accumulation window | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
-| `on_batch_end` | Every batch | `CheckpointCallback` |
+| `before_optimizer_step` | Every accumulation window, before `optimizer.step()` | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
+| `on_batch_end` | Every batch | — |
+| `after_optimizer_step` | Every accumulation window, after `optimizer.step()` and `scheduler.step()` | `CheckpointCallback` |
 | `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
 | `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
 | `on_train_end` | Training exits after `on_train_begin` completes (via `finally`) | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |