Gradient checkpointing with DDP in a loop

[shivammehta25]

Since my method is an Autoregressive algorithm It is making a huge gradient tape, I am trying to do something like this

for i in range(len(maxtrix.shape)):
    output = torch.utils.checkpoint.checkpoint(NNModel(matrix[i]))

It works fine on single GPU but on DDP it throws this error

RuntimeError: Expected to mark a variable ready only once. This error is caused by one of the following reasons: 1) Use of a module parameter outside the `forward` function. Please make sure model parameters are not shared across multiple concurrent forward-backward passes. or try to use _set_static_graph() as a workaround if this module graph does not change during training loop.2) Reused parameters in multiple reentrant backward passes. For example, if you use multiple `checkpoint` functions to wrap the same part of your model, it would result in the same set of parameters been used by different reentrant backward passes multiple times, and hence marking a variable ready multiple times. DDP does not support such use cases in default. You can try to use _set_static_graph() as a workaround if your module graph does not change over iterations.
Parameter at index 30 with name module.model.decoder.decoder_network.layers.1.weight has been marked as ready twice. This means that multiple autograd engine  hooks have fired for this particular parameter during this iteration.

I am running it with

pytorch_lightning.plugins.DDPPlugin(find_unused_parameters=False)

Any workaround for this?

[kuixu]

Dear @shivammehta007 , I also got this error, does it have been solved?

[ananthsub]

This does not appear to be a lightning issue, but rather with DistributedDataParallel from torch.distributed not supporting gradient checkpointing

[Bhavay-2001]

Hi @ananthsub , I guess they have added the functionality of _set_static_graph() in the torch DDP but how to work it with pytorchlightning DDP?

[kuixu]

Currently, I solved this problem. Cause of the model has parameters that were not used in producing a loss. Do the following two settings, and you will find the unused parameter names.

• setfind_unused_parameters=False in DDP

• set the environment variable TORCH_DISTRIBUTED_DEBUG to either INFO or DETAIL.

[stuartzong]

can you elaborate on your solution?

[Bhavay-2001]

Hi @kuixu, we can find the parameter names but how to go about the solution? Do we need to remove them or what? How will the issue be solved?
I'm facing a similar problem

[stuartzong]

any update on this. have you solve the problem?