Add model checkpointing

Author: gonzalobenegas

README.md

@@ -98,3 +98,13 @@ You can override any parameter from command line like this
 ```bash
 uv run python glm_experiments/train.py trainer.max_epochs=20 data.batch_size=64
 ```
+
+### Loading a Checkpoint
+
+```python
+# replace with CLMLitModule if necessary
+from glm_experiments.models.lm_lit_module import MLMLitModule
+
+# Load plain pytorch model from checkpoint
+model = MLMLitModule.load_from_checkpoint("logs/train/runs/<timestamp>/checkpoints/{step}.ckpt").net
+```

configs/callbacks/default.yaml

@@ -2,20 +2,19 @@ defaults:
   - model_summary
   - step_progress_bar
   - lr_monitor
+  - model_checkpoint
   - _self_
 
 model_checkpoint:
   dirpath: ${paths.output_dir}/checkpoints
-  filename: "epoch_{epoch:03d}"
-  monitor: "val/acc"
-  mode: "max"
-  save_last: True
-  auto_insert_metric_name: False
-
-early_stopping:
-  monitor: "val/acc"
-  patience: 100
-  mode: "max"
+  filename: "{step}"
+  monitor: null
+  save_last: false # prefer the explicit step-based checkpointing
+  save_top_k: -1
+  auto_insert_metric_name: false
+  every_n_epochs: 0
+  verbose: true
+  every_n_train_steps: ${trainer.max_steps}
 
 model_summary:
   max_depth: 1

configs/experiment/transformer_small.yaml …gs/experiment/mlm_transformer_small.yamlconfigs/experiment/transformer_small.yaml renamed to configs/experiment/mlm_transformer_small.yaml configs/experiment/transformer_small.yaml renamed to configs/experiment/mlm_transformer_small.yaml

@@ -66,8 +66,7 @@ def __init__(
     ):
         super().__init__()
 
-        # Save hyperparameters (excluding net for cleaner logs)
-        self.save_hyperparameters(ignore=["net"], logger=False)
+        self.save_hyperparameters(logger=False)
 
         self.net = net