Precision Plugins should be part of Training Type Plugins

[ananthsub]

🚀 Feature

Precision Plugin should be a property of the training type plugin, not the accelerator. these can overlap, and interleaving sequences of operations across the two in the accelerator can be challenging to get fully right.

cc @justusschock @awaelchli @kaushikb11 @tchaton

todo: fill out in more detail

Motivation

This would make management of things like this easier:

• Lightning + DDP + apex.amp fails to initialize #7271
• LR scheduler called before optimizer step when using DDP sharded plugin #7217

Pitch

Alternatives

Additional context

[justusschock]

@ananthsub This was discussed quite a few times already now and I think the agreement was to change this, when we maybe also move the accelerator inside the training type plugin

cc @SeanNaren

[SeanNaren]

Thanks @ananthsub for making this. I recall doing an initial POC of this a few months back and realised there is quite some work to make this change when being careful of BW compat.

It would also make these 'shell' classes obsolete, like we see with deepspeed: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/plugins/precision/deepspeed_precision.py

all the above logic should live within the TrainingTypePlugin really. We do have to be careful though, I do like distinctions to remain (precision logic to be separate if possible).

[carmocca]

It might be helpful to compare how the implementation of one training type or plugin would look if this was done. As it stands, it is not clear to me exactly what the benefits of this are