Why accumulate_grad_batches cannot be used with manual optimization?

[NathanGodey]

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.

[rohitgr7]

gradient accumulation revolves around loss.backward and optimizer.step calls. Since both of them are not controlled by PL during manual optimization and is up to the user to call it, it can't really be integrated natively for this case.

[NathanGodey]

Yes, that is right. But what I end up doing is storing my own accumulate_grad_batches as a custom attribute of my LightningModule, and using it in the manual optimization process, which looks like an ugly fix. My question is: is there room for improvement in the library here or are there deeper reasons for not changing this ?

[NathanGodey]

For instance, there could be some kind of buffer regarding step calls to the module's optimizers when both manual optimization and accumulate_grad_batches are required.

[tchaton]

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.

[NathanGodey]

Thanks for your answer! That's how I implemented it indeed, the question was more on the reasoning behind this "incompatibility", but as you mentioned, it would be almost impossible to cover all use cases.