Pass the model intializer into `Trainer.__init__`

I think it could be helpful to add a model_initializer: Optional[Callable[[torch.nn.Module], None]] argument to the trainer.

If this argument is provided, then trainer would call model.apply(model_initializer) after the random seed is set, after deterministic mode is configured, and after Event.INIT fires (i.e. after surgery occurs), but before checkpoints are loaded. The advantages of initializing the model here (rather than relying on the user to initialize the model before passing it into the trainer) is that:

1. the initializing would be deterministic w.r.t the random seed (so the user no longer needs to call composer.utils.reproducibility.set_random_seed() or composer.utils.reproducibility.configure_deterministic_mode() before creating the model and constructing the trainer)
2. model layers replaced via surgery would be initialized.

This would not affect checkpoints, since checkpoint weights would override any initialized weights.

Thoughts?

cc: @jbloxham @ajaysaini725 @anisehsani @hanlint @A-Jacobson

[A-Jacobson]

Alternatively, we can have the recommended way be to have the user call composer.utils.reproducibility, and then remove these arguments from the trainer’s __init__. However, I think this would make it slightly less intuitive, as now a user would need another import and know where to call these functions.

forgot to address this. I’d like to keep it on the trainer init if possible too (I don’t like the additional import) but only if it actually does what it says it does.

[mvpatel2000]

Closing because we don’t plan on adding this