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

Merge pull request #27 from 0z5a/codex/fix-checkpoint-after-step

Browse Source
This commit is contained in:
ViperEkura
2026-09-01 14:24:18 +08:00
parent 432dfec3c2 0e7dafad8e
commit 63f23a4454
6 changed files with 63 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/trainer/train_callback.py
+9 -6
View File
@@ -54,8 +54,11 @@ class TrainCallback(Protocol):
def on_batch_end(self, context: TrainContext):
"""Called at the end of each batch."""
def on_optimizer_step(self, context: TrainContext):
"""Called on every optimizer step (sync step only)."""
def before_optimizer_step(self, context: TrainContext):
"""Called immediately before every optimizer step (sync step only)."""
def after_optimizer_step(self, context: TrainContext):
"""Called after the optimizer and scheduler step (sync step only)."""
def on_error(self, context: TrainContext):
"""Called when an error occurs during training."""
@@ -82,7 +85,7 @@ class GradientClippingCallback(TrainCallback):
def __init__(self, max_grad_norm: float):
self.max_grad_norm = max_grad_norm
def on_optimizer_step(self, context: TrainContext):
def before_optimizer_step(self, context: TrainContext):
context.grad_norm = context.executor.clip_grad_norm(
context.model, self.max_grad_norm
)
@@ -170,7 +173,7 @@ class CheckpointCallback(TrainCallback):
)
context.checkpoint.save(save_path)
def on_batch_end(self, context: TrainContext):
def after_optimizer_step(self, context: TrainContext):
if context.optimizer_step - self.last_ckpt_step >= self.interval:
self._save_checkpoint(context)
@@ -216,7 +219,7 @@ class ProgressBarCallback(TrainCallback):
)
@only_on_rank(0)
def on_optimizer_step(self, context: TrainContext):
def before_optimizer_step(self, context: TrainContext):
postfix = {
"step": f"{context.optimizer_step:d}",
"loss": f"{context.loss:.4f}",
@@ -343,7 +346,7 @@ class MetricCallback(TrainCallback):
for log in self.log_cache:
f.write(json.dumps(log) + "\n")
def on_optimizer_step(self, context):
def before_optimizer_step(self, context):
context.grad_snr_tracker.update(context.model)
if (
astrai/trainer/trainer.py
+3 -1
View File
@@ -93,7 +93,7 @@ class Trainer:
self._call_callbacks("on_batch_end", context)
if executor.sync_gradients:
self._call_callbacks("on_optimizer_step", context)
self._call_callbacks("before_optimizer_step", context)
context.optimizer.step()
context.strategy.on_optimizer_step()
context.optimizer.zero_grad()
@@ -101,6 +101,8 @@ class Trainer:
if context.scheduler:
context.scheduler.step()
self._call_callbacks("after_optimizer_step", context)
self._call_callbacks("on_epoch_end", context)
if context.stop_requested:
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` |
tests/trainer/test_callbacks.py
+35
View File
@@ -1,8 +1,12 @@
from pathlib import Path
import torch
from astrai.model.components.decoder_block import DecoderBlock
from astrai.serialization import Checkpoint
from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
from astrai.trainer.trainer import Trainer
from tests.helpers import RandomTokenDataset
def test_gradient_checkpointing_enable_disable(test_model):
@@ -135,3 +139,34 @@ def test_callback_integration(
assert "on_train_begin" in callback_calls
assert "on_batch_end" in callback_calls
assert "on_epoch_end" in callback_calls
def test_checkpoint_captures_completed_optimizer_step(
base_test_env, train_config_factory, device
):
"""Checkpoint state must include the update represented by its step number."""
model = base_test_env["model"]
initial_state = {
name: tensor.detach().cpu().clone()
for name, tensor in model.state_dict().items()
}
train_config = train_config_factory(
model_fn=lambda: model,
dataset=RandomTokenDataset(length=2),
test_dir=base_test_env["test_dir"],
device=device,
batch_per_device=2,
ckpt_interval=1,
)
Trainer(train_config).train()
checkpoint = Checkpoint.load(
str(Path(base_test_env["test_dir"]) / "epoch_0_step_1")
)
assert any(
not torch.equal(checkpoint.state_dict[name].cpu(), initial_tensor)
for name, initial_tensor in initial_state.items()
)
assert checkpoint.extra["optimizer"]["state"]
assert checkpoint.extra["scheduler"]["last_epoch"] == 1