Revised distillation modifier

[natuan]

I made the following changes to the distillation modifier:

• Removed the use of forward hook to capture the latest inputs/outputs because it ignores the kwargs dict passed in to model forward (which is the type of model inputs, e.g. in QA model)
• Removed the computation of teacher outputs inside the modifier's loss_update. Passing in the entire inputs dict into the teacher forward causes cross entropy loss to be calculated, which we don't need; it also creates an issue downstream from autograd when doing loss.backward.

[markurtz]

@natuan our goal is to remove as much complexity for the user with modifiers as possible. This implementation is adding more complexity back to the user, specifically needing to manage the complete lifecycle of the teacher model themselves now. Additionally, the old code, as set up, would work for other pipelines and reduce simplicity for users.

[natuan]

@natuan our goal is to remove as much complexity for the user with modifiers as possible. This implementation is adding more complexity back to the user, specifically needing to manage the complete lifecycle of the teacher model themselves now. Additionally, the old code, as set up, would work for other pipelines and reduce simplicity for users.

Per our discussion, I moved the teacher's logic back into the distillation modifier

[bfineran]

LGTM