[BUG]Failed to call forward() multiple times using deepspeed zero3

Describe the bug
I want to call forward() multiple times to calculate different loss using deepspeed zero3 ,but I get this error.
"assert len(set(p.ds_id for p in self.params_in_ipg_bucket)) == len(self.params_in_ipg_bucket)
AssertionError"
It seems like caused by call model twice in forward,but l don't know how to fix it.

Expected behavior

def run_model():
     ......
     loss1 = self.model(input1)
     loss2 = self.model(input2)
     loss = loss1 + loss2
     return {"loss":loss}

def run():
    out = self.run_model(batch)
    loss = out["loss"]
    self.model.backward(loss)
    self.model.step()

Screenshots

System info (please complete the following information):
deepspeed version 0.11.1

Start by reproducing the reported DeepSpeed 0.11.1 failure with the shown run_model pattern: call self.model(input1) and self.model(input2) before backward and step under ZeRO-3. Trace the assertion involving params_in_ipg_bucket and compare it with a single forward call. Done means repeated forward calls complete without the assertion and the combined loss can be backpropagated and stepped.