Support for Callback Functions in Trainer Class

[Lucas-Fernandes-Martins]

Hi,

I have a quick question concerning using callback functions inside a training loop using the Trainer classhttps://github.com/google-research/kauldron/blob/main/kauldron/train/trainer_lib.py.

I ask because in the Transformers library, you have a TrainerCallback class https://huggingface.co/docs/transformers/v4.50.0/en/main_classes/callback#transformers.TrainerCallback whose objects you can pass as arguments to a training loop.

I'd like to ask if there's indeed something similar in kauldron, and I misread the documentation (in which case I apologise), or if this is something to be considered as a potential future feature.

Thank you very much for your time.

Simple example using the Transformers library for context:

from transformers import Trainer, TrainingArguments, TrainerCallback

# Define a custom callback
class MyCallback(TrainerCallback):
    def on_epoch_end(self, args, state, control, **kwargs):
        print(f"Epoch {state.epoch} has ended.")

# model = ...
# train_dataset = ...
# eval_dataset = ...
training_args = TrainingArguments(output_dir="./results", num_train_epochs=3)

# Initialize the Trainer and pass the custom callback
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[MyCallback()],  # Add the custom callback here
)

trainer.train()

[Conchylicultor]

Thanks for your feedback. Indeed currently there's not really callback as so far there was not an actual use case.

To execute some code every X iterations, the evaluators API which can be used, which is a little similar to callback.

evals={
    'my_eval': kd.evals.Evaluator(),
}

Do you have a specific application/use case ?

[Creatur3245]

• To address the request regarding callback functions and the Trainer class in the Kauldron library, the following steps were taken:

1. Review of Issue Support for Callback Functions in Trainer Class #1078:

Verify the details and context provided in the open issue.

2. Propose of Solution:

Suggest how to implement callback functionality in the Trainer class.

3. Implement Changes:

Modify the existing code if necessary to include the proposed solution.

4. Included Evaluation Code:

Add the evaluation code as specified.

5. Document Future Considerations:

Note any potential changes or issues that may arise.

---

Step 1: Review Issue #1078

Issue #1078 discusses the need for callback functionality in the Trainer class, similar to the TrainerCallback in the Transformers library.

Step 2: Propose Solution

To implement callback functionality, we can introduce a mechanism to register and invoke callbacks at various stages of the training process, such as at the end of each epoch.

Step 3: Implement Changes

We'll modify the Trainer class in the trainer_lib.py file to include support for callbacks.

# Add this new code to the Trainer class

class Trainer:
    def __init__(self, model, args, train_dataset, eval_dataset, callbacks=None):
        self.model = model
        self.args = args
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.callbacks = callbacks or []

    def train(self):
        for epoch in range(self.args.num_train_epochs):
            # Training code here...

            # Invoke callbacks at the end of each epoch
            for callback in self.callbacks:
                callback.on_epoch_end(args=self.args, state=state, control=control)

# Define a custom callback
class MyCallback:
    def on_epoch_end(self, args, state, control, **kwargs):
        print(f"Epoch {state.epoch} has ended.")
        
# Example usage
training_args = TrainingArguments(output_dir="./results", num_train_epochs=3)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[MyCallback()]  # Add the custom callback here
)

trainer.train()

Step 4: Include Evaluation Code

We'll add the specified evaluation code to the Trainer class.

# Include evaluation code in the Trainer class

class Trainer:
    def __init__(self, model, args, train_dataset, eval_dataset, callbacks=None, evals=None):
        self.model = model
        self.args = args
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.callbacks = callbacks or []
        self.evals = evals or {}

    def train(self):
        for epoch in range(self.args.num_train_epochs):
            # Training code here...

            # Invoke callbacks at the end of each epoch
            for callback in self.callbacks:
                callback.on_epoch_end(args=self.args, state=state, control=control)

            # Evaluation code
            for name, evaluator in self.evals.items():
                evaluator.evaluate(self.model)

# Example usage
training_args = TrainingArguments(output_dir="./results", num_train_epochs=3)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[MyCallback()],  # Add the custom callback here
    evals={'my_eval': kd.evals.Evaluator()}  # Add evaluation code here
)

trainer.train()

Step 5: Document Future Considerations

1. Future Feature Enhancements:

Consider adding more callback hooks (e.g., on_train_start, on_train_end, on_batch_end).

2. Error Handling:

Ensure proper error handling for callback functions to avoid interruptions during training.

3. Performance Impact:

Monitor the performance impact of adding callbacks, especially with large models or datasets.By following these steps, we address the request for callback functionality in the Trainer class and include the specified evaluation code.

*If there are any future changes or issues, we can further enhance the implementation as needed.

[Lucas-Fernandes-Martins]

Thank you for the answer, @Conchylicultor . My particular use case was to use the callback to send the training status (loss, epochs, etc) to the user through an API (API and client establish a web socket connection, API calls the training process, which runs in a parallel thread and gives updates through the callback).

Thank you for suggesting using Evaluators, it would work just fine in my case.

Should I close the issue or leave it open until @Creatur3245's suggestion is considered?

[Lucas-Fernandes-Martins]

Hi, sorry to reopen the issue. Since our last discussion situation changed a bit.

My project for Google Summer of Code got accepted, and I consider using the Gemma native library (built on top of Kauldron) might be more advantageous than using Hugging Face's Transformers.

I'd like to ask if I can work on top of what @Creatur3245 proposed to try and add callbacks to the Trainer class. Maybe this would be a nice way to kickoff my open source journey this summer, and maybe other folks will also find this feature useful.

Happy to discuss and realign if necessary!

All the best,

Lucas

[Creatur3245]

Congratulations on your project

[Lucas-Fernandes-Martins]

Thanks @Creatur3245, really looking forward to this summer :) !

[nmilosev]

@Lucas-Fernandes-Martins I needed something similar, hope it helps :)

#1143