[HELP] Perform backpropagation separately for the two losses and extract the corresponding gradients.

Hello, I currently have a requirement: during a single-time forward pass, I compute two losses, and I would like to backpropagate each loss separately to obtain their respective gradients. Since I made the modifications within a complex project codebase and cannot paste the original code, the following pseudocode illustrates the core logic:

#step1:  forward, and calculate two losses
output = model(input)
loss1 = loss1_func(output, target_1)
loss2 = loss2_func(output, target_2)

#step2:  loss1 backward and extract its grad
model.zero_grad()
accelerator.backward(loss1,  retain_graph=True)
grad_loss1 = safe_get_full_grad(model)

#step3:  loss2 backward and extract its grad
model.zero_grad()
accelerator.backward(loss2)
grad_loss2 = safe_get_full_grad(model)

However, when I run the code with ZeRO-1 and gradient_accumulation_steps=1, I find that the second call to model.zero_grad() (the first line of step3 ) does not take effect — that is, even though this line is executed, the gradients from loss1 still remain in the model. Moreover, the gradients from loss2 cannot overwrite or accumulate onto the existing gradients. In other words, once the backward pass for loss1 is completed, its gradients are forcibly retained in the model and cannot be cleared or overwritten.

I further found that when I set gradient_accumulation_steps=2, the gradients for the second loss can be backpropagated correctly. However, in my project, gradient_accumulation_steps must be strictly set to 1. Therefore, is there any way to successfully extract the gradients corresponding to both losses when gradient_accumulation_steps=1?

Thanks a lot!

Reproduce the pseudocode with ZeRO-1 and gradient_accumulation_steps=1, then compare it with accumulation_steps=2. Start by tracing model.zero_grad(), accelerator.backward(), and safe_get_full_grad() through the ZeRO-1 gradient-management path. Done means both losses produce independently extractable gradients without loss1 gradients persisting into the loss2 pass.