Deprecate or re-purpose the `get_mp_spawn_kwargs` hook in spawn plugins

[awaelchli]

Proposed refactor

Discuss and refactor TPU/DDPSpawnPlugin.get_mp_spawn_kwargs.

Motivation

As part of our effort to converge to a stable and well organized Strategy+Accelerator API #10416 we are looking for opportunities to simplify a couple of hooks. One is the get_mp_spawn_kwargs in the Spawn-plugins.

    def get_mp_spawn_kwargs(self, trainer: Optional["pl.Trainer"] = None) -> Dict[str, Any]:
        return {
            "nprocs": len(self.parallel_devices),
            "start_method": self.start_method,
        }

Pitch

Proposal 1:

• Remove/deprecate the method
• Directly use start_method as an attribute and/or constructor argument
• nprocs can be set already in the constructor
• A custom implementation that needs to pass additional arguments would instead override the new spawn() method.

Proposal 2:

• Keep it
• Remove the Trainer argument
• Mark it as protected, since the method is only used internally by the Plugin

Proposal 3:

• Keep as is

Additional context

Original discussion started in #10034 by @four4fish

---

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 @justusschock @awaelchli @akihironitta @tchaton

[ananthsub]

would the choice here be guided by the inheritance/interface structure of strategies going forward for spawn & non-spawn strategies? For instance, would spawn variants extend their non-spawn counterparts, and implement an additional spawn interface? Or the interface described here? https://github.com/PyTorchLightning/pytorch-lightning/pull/10896/files#r761435498

I think this would dictate how args are ultimately passed when spawning

[awaelchli]

I would say get_mp_spawn_kwargs (if we keep it) goes wherever the spawn() method lives.
But other than that I see no direct impact to how these plugins might evolve in the future.

I'm recommending Option 1 to remove get_mp_spawn_kwargs.

My high-level view of the current situation with DDP/DDPSpawn:
The fundamental difference between the DDPSpawn and DDPPlugin is the spawning logic. After #10896, it would theoretically be possible to say: Pull out all spawn logic from DDPSpawnPlugin and pull out all subprocess launching logic from DDPPlugin and what remains of the two classes is identical to each other. This statement would be false today, because execution order is impacted by the current design of how spawn is called and where. Whether that also holds true for subclasses of these plugins would have to be investigated. From here on out, I believe one can go in either of directions you mentioned

[carmocca]

@rohitgr7 @awaelchli is this part of the launcher PR or is this going to be done as a follow-up?

[awaelchli]

Yes this will be completely obsolete and closed by #11643
Basically proposal 1 but directly done inside the launcher classes who need that.
It will be deleted and no backward compatibility.