ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: rename optimizer step callback hooks to before and after

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-09-01 14:22:54 +08:00
parent 08721f6d31
commit 0e7dafad8e
5 changed files with 24 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/developer/architecture.md
+6 -5
View File
@@ -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
+5 -3
View File
@@ -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` |
docs/guides/training.md
+5 -3
View File
@@ -70,12 +70,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
```
@@ -87,8 +88,9 @@ on_train_end
| `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` |