Two slightly different process for Deq about deq

[jerrybai1995]

Hi @SamChen ,

Thanks for the question. The gradcheck fails because gradcheck works by backpropagating through the same computation graph multiple times (e.g., by adding eps to each entry of the vector output, and then backprop); whereas the self.hook.remove() already removed the hook upon the first backward call. Therefore, the code is correct for DEQ-Transformer training (where each iteration has exactly ONE backward pass through the DEQ), but is incorrect for repetitive backward passes (which is what gradcheck does).

The memory leak, I believe, is a pytorch-related issue. I'm not entirely sure about the source of this problem but pytorch 1.6 and 1.7 should both work well (i.e., no memory leak). If you encounter the SystemError and do not want to downgrade pytorch, then you can also use the tutorial implementation--- basically replacing the current L372-380 with:

z1s_copy = z1s.clone().detach().requires_grad_()
new_z1s_copy = self.func(z1s_copy , *func_args)      # Spend one more NFE in training forward
def backward_hook(grad): 
    new_grad = self.b_solver(lambda y: autograd.grad(new_z1s_copy , z1s_copy , y, retain_graph=True)[0] + grad, \ 
                                  torch.zeros_like(grad), threshold=b_thres)['result'] 
    return new_grad 
new_z1s.register_hook(backward_hook) 

Of course, this means you have to spend one more NFE in the forward pass of training, which means slightly more memory and computation (which is what the current implementation hoped to avoid). But this should help avoid the memory leak.

Let me know if this helps!

[SamChen]

Thanks for the clear explanation. :-)
Your words about gradckeck explains why I saw it calls the DNN function over and over. And, of course, it is not related to the infinite loop

deq/DEQ-Sequence/models/deq_transformer.py

Line 374 in c161644

# To avoid infinite loop