Error when calling forward() multiple times using deepspeed zero3

[powpos360]

Hi all. I try to do something like the following:

output1 = model.forward(input1)
input2 = some_func(input1, output1)
output2 = model.forward(input2)
loss = some_loss(input1, input2, output1, output2)
loss.backward()

However, this gives me File "/somepath/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 1111, in __reduce_and_partition_ipg_grads assert len(set(p.ds_id for p in self.params_in_ipg_bucket)) == len(self.params_in_ipg_bucket) AssertionError. Does anyone have an idea how to resolve this? I'm using deepspeed zero3 0.10.0 with gradient checkpointing. Thanks!

[zhuole1025]

I met the same problem.. Hope someone can fix this

[zhilif]

A simple workaround is to merge input1 and input2 and you can do one forward pass