XLA cluster environment plugin

[awaelchli]

🚀 Feature

Add an environment for the TPU.
The TPUSpawnStrategy will use this cluster environment by default.

Motivation

Decoupling the environment lookup from the strategy will reduce the code divergence of the generic ddpspawn strategy and tpu spawn strategy.

Pitch

class TPUEnvironment(ClusterEnvironment):
    # here access the env variables set by XLA directly

List of environment variables: https://github.com/pytorch/xla/blob/master/torch_xla/core/xla_env_vars.py

Name? XLAEnvironment? TPUEnvironment?

Alternatives

Additional context

Discussion sparked in #11163
Proposed by @tchaton

---

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 @Borda @kaushikb11 @rohitgr7 @awaelchli @ananthsub

[kaushikb11]

Sounds good to me!

XLAEnvironment would be the ideal name. As it could be further extended to other hardwares and we have a checkpointing plugin called XLACheckpointIO