Is there a simple way to use old (PyTorch Lightning < 1.6) global_step update behavior?

[lecacosa]

I have 1 generator and 2 discriminators (d_a and d_b) in my LightningModule for GAN training. d_a will be updated every 2 steps and discriminator d_b will be updated every 3 steps. The generator will be pruned at step 10000.

The following code work in PyTorch Lightning < 1.6,

class PrunedModule(pl.LightningModule):
    """Apply pruning when self.global_step == 10000"""

    def training_step_end(self, _):
        if self.global_step == 10000:
            # apply pruning


class PrunedGAN(PruningLightningModule):

    def training_step(batch, batch_idx):
        # ...
        optimizer_g.step()

        if self.global_step % 2 == 0:
            # ...
            # update d_a
            optimizer_d_a.step()

        if self.global_step % 3 == 0:
            # ...
            # update d_b
            optimizer_d_b.step()

In PyTorch Lightning 1.6, the self.global_step of PrunedGAN will be increased by 1, 2 or 3 in a single training_step.

If I use self.register_buffer("my_global_step", torch.tensor(0)) in PrunedModule and increase self.my_global_step by 1 in training_step_end, I still need to re-implement all classes inherited from PrunedModule and ModelCheckpoint to make sure that they use self.my_global_step instead of self.global_step or trainer.global_step. Besides, my_global_step will be moved to cuda in GPU training mode. It seems to be a bad idea to re-implement global step tracking by myself.

Is there a simple way to use old (PyTorch Lightning < 1.6) global_step update behavior when using multiple optimizers?

[akihironitta]

#13752 (comment) might be relevant to your question.