Why accumulate_grad_batches cannot be used with manual optimization?
I've stumbled upon the problem of not being able to use accumulate_grad_batches argument in the Trainer as I was doing manual optimization in my LightningModule to use adversarial loss functions.
However, I think it would be possible to implement something that would "store" calls to the step method for the module's optimizers and actually apply them once every accumulate_grad_batches iterations. I've seen several related issues about similar behavours when overriding optimizer_step or close to my use case (#5054, #5108). The proposed fixes always leave some manual get-arounds in the final code.
My question: is there a reason for such incompatibility of accumulate_grad_batches with manual optimization ?
One reason might be the need to step different optimizers at different paces (one every batch, another every n batches ...) but this seems to be an extreme use case.
Hey @NathanGodey,
manual optimization was built to provide full control optimization control to the user while abstracting distributed training and precision.
There is no way Lightning can automate properly accumulate grad batches for all the possible use cases and therefore isn't supported.
However, you can easily implement it by not calling zero_grad, step every n batches.