Strategies and Accelerator teardown logic clean up

[four4fish]

Proposed refactor

After #11022 Move accelerator to strategies, the teardown logic can be simplified and cleaned up a bit.

Motivation

Reduce redundant code and make code more clean

Pitch

In ddp, ddp_spawning and hovorod plugins, teardown() all has device specific logic like:
https://github.com/four4fish/pytorch-lightning/blob/move-accelerator/pytorch_lightning/plugins/training_type/ddp.py#L507-L511

with accelerator moved into strategy, device related teardown should move to X_accelerator

class GPUAccelerator()
        def teardown(self):
               # self.lightning_module.cpu()  suggested by @ananthsub module movement should be handled by strategies
               torch.cuda.empty_cache()

Some teardown() logic are shared between different strategy and same device type, like single_tpu and tpu_spawn has the same teardown() logic:
https://github.com/four4fish/pytorch-lightning/blob/move-accelerator/pytorch_lightning/plugins/training_type/single_tpu.py#L90-L92

we can move it to TPUAccelerator

class TPUAccelerator()
          def teardown(self) -> None:
                 os.environ.pop("PT_XLA_DEBUG", None)

Additional context

---

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 @kaushikb11 @ananthsub

[awaelchli]

As discussed in other places before, I agree with this. The strategy calls into accelerator.teardown() which handles device specific teardown. I am also working on teardown related things, e.g., #10990, happy to help with this one or otherwise ping me for review <3

[four4fish]

As discussed in other places before, I agree with this. The strategy calls into accelerator.teardown() which handles device specific teardown. I am also working on teardown related things, e.g., #10990, happy to help with this one or otherwise ping me for review <3

@awaelchli Oh Awesome, you already work on the similar things. Will assign to you then :) Thank you for helps and happy to review when it's done ~~

[ananthsub]

@four4fish I don't fully understand the proposal for the GPU accelerator

1. The accelerator does not have a reference to the LightningModule.
2. We explicitly made teardown control the device movement logic to accommodate parallelization strategies like FSDP: refactor accelerator teardown -> training type plugin teardown #7579

If the proposal is for the GPU accelerator to call torch.cuda.empty_cache() in teardown that's perfectly reasonable. But the accelerator should not control the device movement of the model - that should be determined by the strategy. This way, users can specify a manual parallel strategy that works with the GPU accelerator without again needing to provide both custom strategies and custom accelerators.

[four4fish]

@ananthsub yeah make sense, the self.lightning_module.cpu() should definitely in strategy not accelerator. I saw it was user self.on_gpu() check, so suggested move to accelerator. But you are right, module/model movement logic should be in strategy. (will update the proposal)
Current fsdp has lightning_module.cpu() logic too, but we should revisit these. I think YIfu find out that ligtning_module.cpu() cause OOM for both DDP and FSDP. I started #10415 but it was won't fix.