Best way to add custom callbacks?

[clasqui]

I am a performance analyst and to get the most information possible, I need to instrument the code. I see that PyTorch provides callbacks at the train and step level (very coarse grain IMO), and that axolotl makes use of those for different purposes.

The only way I have found now to add custom callbacks for my needs is to modify the trainer_builder specifically and add a line callbacks.append(...) in the get_callbacks function of the different Trainer Builder classes. Is there a more gentle way to do this? Could this be integrated through the configuration file? I am missing something from the documentation?

[NanoCode012]

Hey! This is actually one of the things that PR #1917 brings . It adds an example of adding a callback for the optimizer step level, but if you check the base class, there’s quite a few callback methods to override.

If you take a look, it requires adding a file to the plugins directory then listing it in the yaml.

[clasqui]

Great thank you! But there is one thing that I'm not sure I understand. Does axolotl already provide the integration of plugins to add callbacks? What is your PR proposing exactly? An example of this, or the support to integrate plugins? From the documentation of axolotl looks like this is already implemented, but the changes https://github.com/axolotl-ai-cloud/axolotl/pull/1917/files#diff-5edc13801ecfdd108e81872527bdc78c6d24a73833968147b0a9ecb8452996f4R1211-R1222 in commit 0b674eb seems like it is adding this functionality.

[clasqui]

I tried it and in fact, the plugin with callbacks only works if you apply all the changes in the PR. Maybe this is something that should be in a separate PR, that introduces support for plugins that add callbacks?

[NanoCode012]

Oh hey, sorry, I thought I replied to this already.

Does axolotl already provide the integration of plugins to add callbacks? What is your PR proposing exactly?

It adds additional support for plugins. I will let wing know that this PR is of interest and to potentially get it in soon. It should already be ready for merge as I have tested it too.

[NanoCode012]

Hello @clasqui , just a follow up, we had a separate PR that merged plugin callback system #2006

You could now use it for your custom callbacks.