Have each DDP worker optimizing a specific layer of a common model

Nobody has claimed this yet.

### Description & Motivation

Consider the case where we have one model that works with different input types (e.g. RGB and grayscale). Assume we know we should update only certain layers of the model depending on the input type.

Now I want to train one model with these different input types, so that I train it on worker 1 only with inputs of type t_1 and on worker 2 only with inputs of type t_2.

### Pitch

It would be great if this could be achieved in PyTorch lightning.

### Alternatives

The closest solution I have so far is:

```python
import numpy as np
import pytorch_lightning as pl
import torch
import time
import torch.nn as nn

from torch.utils.data import DataLoader, Dataset

class MyDataset(Dataset):
def __init__(self):
self.X = torch.randn((1000, 1, 128, 128))

def __getitem__(self, idx):
return self.X[idx, :]

def __len__(self):
return self.X.shape[0]

class Model(nn.Module):

def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 2, 1, bias=False)
self.conv2 = nn.Conv2d(2, 1, 1, bias=False)

def forward(self, x):
return self.conv2(self.conv1(x))

class Module(pl.LightningModule):

def __init__(self):
super().__init__()
self.automatic_optimization = False
self.model = Model()

def training_step(self, batch, batch_idx):
optA, optB = self.optimizers()
pred = self.model(batch)

if self.trainer.is_global_zero:
optA.zero_grad()
optB.zero_grad()
loss = torch.nn.functional.mse_loss(pred, batch)
self.manual_backward(loss)
optA.step()
else:
optA.zero_grad()
optB.zero_grad()
loss = torch.nn.functional.mse_loss(pred, batch)
self.manual_backward(loss)
optB.step()

def configure_optimizers(self):
optA = torch.optim.Adam(self.model.conv1.parameters())

optB = torch.optim.Adam(self.model.conv2.parameters())
return optA, optB

def train_dataloader(self):
dataset = MyDataset()
return DataLoader(dataset, batch_size=10)

if __name__ == '__main__':
model = Module()
model.requires_grad_(True)
trainer = pl.Trainer(max_epochs=20, devices=2)
trainer.fit(model)

```
Here I activate `requires_grad` for all parameters of the model although I want per worker only to update a subset of it. In DDP, the set of parameters with `requires_grad` must be the same across workers if I am not mistaken.

Now each worker uses its specific optimizer to update the gradients.
However, the current code is not correct. Lets consider the case of rank 0:
The call `self.manual_backward(loss)` will sum up and sync the gradient for `self.model.conv1.parameters()` computed from rank 0 **and** rank 1. This is not want we wanted. Also, it will divided it by 2, which is again wrong.

So if we have a set a workers `S_1` that shall optimise `self.model.conv1` and a set of workers `S_2` that shall optimise `self.model.conv2` , then
`self.manual_backward(loss)` must compute the gradient within the specified group `S_1`: sum up the gradients of `self.model.conv1.parameters()` from all workers of `S_1` and divide the final gradient by `|S_1|` . Finally, the gradient for the parameters of `self.model.conv1.parameters()` must be sent to the workers of `S_1` and `S_2`.
Likewise, the workers of workers `S_2` compute within their group the gradient of `self.model.conv2.parameters()` and sync it with all workers.
Then all workers will have the same gradient information and the optimizer `optA` updates the parameters `self.model.conv1`, `optB` updates the parameters `self.model.conv2`. Finally, the result must be synced again so that all workers end-up with the same model weights.

### Additional context

_No response_

cc @lantiga @borda @justusschock

Start with the issue's `training_step`, `manual_backward`, `configure_optimizers`, and `Trainer` example to understand the requested DDP behavior. The feature should let worker groups synchronize gradients only for their assigned parameter subsets, then synchronize updated model weights across all workers; no repository files or tests are named.

Written by the indexing model from the issue text.