[Main Issue] Accelerator and Plugin refactor

[four4fish]

Proposed refactoring or deprecation

Motivation

Accelerator is not stable API yet, we can improve the Accelerator related logic and move towards stable Accelerator version for 1.6

Pitch

Steps

1. Collective refactor Consolidate collective functions #7534
   • deprecate accelerator collective directly call from TTP 1/n Call training_type_plugin collective functions directly instead of going through the Accelerator  #9677
   • collective refactor 2/n Consolidate collective functions - collective base and subclasses #9414
2. Move Precision Plugin into TTP Precision Plugins should be part of Training Type Plugins #7324
3. Move Accelerator into Strategy [Accelerator refactor] Move Accelerator into Strategy #10648
4. Simplify the Spawning logic Simplify multiprocessing logic in DDPSpawn plugins #10059
5. [RFC] Simplifying the Accelerator Connector logic and flags (can be done in parallel with aboves) Rewrite Accelerator_connector and follow up tasks #11449 11449
6. [RFC] Revisit the inheritance of TTP Flatten the Strategy inheritance #11863

More details in: Accelerator Refactor Proposal
[updating]

FAQ

1. Will this be a lot of breaking changes?
   Not much user facing API changes from 1,2,3,4.(Unless we found out other existing bugs during refactor) The only breaking change will be for custom plugins
   5 and 6 is still RFC stage, may have breaking changes which impact user facing APIs

2. How does this impact lightningLite?
   Should be helpful for lightningLite too, there maybe function refactor/simplification could happen for lightningLite. (@awaelchli any suggestion about this part?)

Follow up TODOs:

• check if trainer.FITTING before setting up optimizers 2/n Move Precision Plugin into strategy - move optimizer related logics #10596 (comment)

• Send a PR to https://github.com/ray-project/ray_lightning to update their plugins

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

cc @justusschock @awaelchli @akihironitta @tchaton @Borda @kaushikb11 @ananthsub

[awaelchli]

Great summary! About

How does this impact lightningLite?

As best as I can see right now:
Step 1, 2 do not impact Lite. If there are changes in required in the Trainer, they will be mirrored directly by Lite. The Trainer as well as Lite do not make any strong assumptions on the internal composition of plugins.
Step 3: not sure
Step 4: Will impact Trainer more than Lite. Lite is sort of already setting an example here how it could/should look like in the Trainer.
Step 5, 6: should not impact Lite at all (accelerator connector is shared)

[tchaton]

Yes, definitely really excited about this. I am quite eager to finally see the accelerators marked as a stable API targetting v1.6.

[awaelchli]

@ananthsub

#10570 (review)

It'd be very good to list what APIs this lets us simplify or remove from the Accelerator/strategy/precision plugin. Those could be listed as immediate followups to this PR

Here are some that I found and also some are mentioned in the doc

After 2)

Methods:

• Accelerator.optimizer_step (move)
• Accelerator.backward (move)
• Accelerator.*_dispatch (move + reduce)
• Accelerator.*_step (move)
• some miscellaneous ones

Property:

• Accelerator.amp_backend (move)
• Accelerator.precision (move)
• Accelerator.scaler (move)

After 3)

• TrainingTypePlugin.on_gpu/on_tpu
• TrainingTypePlugin.model_to_device

These will either make use of the accelerator or be moved to it.

[awaelchli]

Now that the precision plugin has moved, I will take a look at this TODO here:
https://github.com/PyTorchLightning/pytorch-lightning/blob/af0bb96f0ff645102680a7adc99dc131cfeb9c0b/pytorch_lightning/lite/lite.py#L433-L439

[awaelchli]

Here is another follow up we need to do IMO: #10686. This will unblock also #10657

[carmocca]

Closing this issue in favor of the smaller linked issues for the pending tasks.