Principle:Huggingface Open r1 Training Callbacks
Overview
An event-driven extensibility mechanism that hooks custom logic into the training loop lifecycle to perform actions like per-checkpoint Hub publishing and automated benchmark evaluation.
Description
Modern training frameworks use a callback pattern to inject custom behavior at specific training lifecycle events (on_save, on_evaluate, on_train_end, etc.) without modifying the core training loop. Open-R1 uses this pattern for two key capabilities:
- PushToHubRevisionCallback - on every checkpoint save, pushes model weights to a unique Hub branch (e.g.,
main-step-000001000) excluding optimizer states, enabling fine-grained model selection post-training. - Automated benchmark evaluation - after each Hub push completes, submits Slurm evaluation jobs for all configured benchmarks.
The callback registry pattern allows selecting callbacks by name in YAML configs.
Usage
Use when you want per-checkpoint Hub uploads and/or continuous benchmark evaluation during training. Configure via the callbacks field in SFTConfig or GRPOConfig YAML files.
Theoretical Basis
The callback lifecycle pattern decouples custom training-time actions from the core training loop. At each lifecycle event, the trainer iterates over registered callbacks and invokes the corresponding hook method. This allows composable, modular extensions without subclassing or modifying the trainer itself.
Pseudocode:
class TrainingCallback:
def on_save(args, state, control):
revision = f"{base_revision}-step-{global_step:09d}"
future = push_to_hub(model_dir, revision, exclude=["*.pt"])
if slurm_available:
future.add_done_callback(
lambda _: run_benchmarks(model, revision)
)
CALLBACKS_REGISTRY = {"push_to_hub_revision": PushToHubRevisionCallback}
callbacks = [CALLBACKS_REGISTRY[name](model_config) for name in config.callbacks]