Enabling Access to Currently Training Model in Callback Handler

[jp1924]

Feature request

Update self.callback_handler.model to reference the actual training model instance (model) instead of the pre-prepared model (self.model) in Trainer._inner_training_loop.

Motivation

• When debugging model training, developers often need to monitor model weights and other parameters at each step, similar to wandb.watch
• Currently, callbacks receive the pre-prepared model instance (self.model) rather than the actual training model (model)
• This is evident from different memory IDs:

  id(self.callback_handler.model): 140049287962320
  id(model): 140048547892912  # Model memory id in training_step

• The code flow suggests this might be unintended:

  model = self._wrap_model(self.model_wrapped)
  self.callback_handler.model = self.model  # Current implementation

trainer_code

Your contribution

The proposed change is to update the model reference in the callback handler to point to the actual training model:

# Current implementation
self.callback_handler.model = self.model

# Proposed change
self.callback_handler.model = model

This would allow callbacks to access and monitor the actual model being trained, enabling proper debugging and monitoring capabilities like examining weights during training steps.

This change would ensure that when users implement custom callbacks for debugging or monitoring purposes, they have access to the correct model instance that's being actively trained, rather than the pre-prepared version.