Question regarding the gradient accumulation

[lehduong]

Hi, thanks for your implementation.

I noticed you accumulated gradients of two models with two different context managers (here). Could you let me know if you verified your implementation with gradient accumulation step different than 1? Apparently, this approach can be erroneous according to this and followed-up comments. I believe the newer version of hf's accelerate has already allowed the context manager to receive multiple models as in here.

[kxhit]

Hi @lehduong thanks for pointing this out! Yes, in Zero123, the gradient accumulation is 1 so this doesn't matter. And I know people are trying to fix it for more than one model. According to your link, the correct way to do it for multiple models is
with accelerator.accumulate(model1, model2)? Would be very happy if you want to PR! Thank you!

[lehduong]

Yes, you only need to change this line to with accelerator.accumulate(unet, cc_projection) and use accelerate >= 0.23 (not sure if it is the earliest accelerate version that supports this feature but I'm using it). I did a quick experiment and I attached the image below to compare the training loss of two accumulation approaches. I used the resolution of 512, (per device) batch size of 24, and the gradient_accumulation_steps is set to 8 (effective batch size is 1536 as in the original implementation).

[kxhit]

Thanks a lot for your contribution! Just pushed a fix, please reopen the issue if anything needs further change!

[cfeng16]

Hi @lehduong , can I ask do you use gradient accumulation for multiple models in distributed training setting (say multiple GPUs)?

[lehduong]

I trained the model on a single machine (8gpus).