Fine-grained control of Parallel Trainer: GPU Assignment in DDP
Dear all,
really enjoy using Lightning - amazing framework! I have recently begun to run into limits with its distributed trainer implementation,when seeking to assign different models which feed into the same gradient flow to specific GPUs, but have been unable to find any documentation for it and kind of presume for it to not exist in Lightning yet.
The specific use-case is that at GAN training time I seek to include a pre-trained model to nudge my generator into a certain direction. The way I am currently doing this is by "nudging" with the pre-trained model every third epoch. But this causes some data movement, which I would like to minimize. E.g. for
def training_step(self):
    (..)
    if optimizer_idx == 0:
        if self.trainer.current_epoch % 3 == 0:
            model_1
        else:
            model_2

def configure_optimizers(self):
    (..)
    return [opt_g, opt_d], _


In the example of a 4-GPU cluster, I would now like to be to assign/launch 3 instances of model_2 on GPU <0, 1, 2> and launch one instance of  model_1 on GPU3 with all of them using the same optimizer.
This functionality should also be highly advantageous for (sparse) mixture of experts and sharding-reliant models.
Thank you for reading my issue,
Ludger