Use utils.checkpoint with DDP report an error

HELLORPG: <p>I use utils.checkpoint with DistributedDataParallel at the same time to save the CUDA memory.</p>
<p>When I use a single GPU to run my code, it works.</p>
<p>But if I use multi GPUs to run my code, it will report an error:</p>
<blockquote>
<p>RuntimeError: RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument <code>find_unused_parameters=True</code> to <code>torch.nn.parallel.DistributedDataParallel</code>, and by<br>
making sure all <code>forward</code> function outputs participate in calculating loss.<br>
If you already have done the above, then the distributed data parallel module wasn’t able to locate the output tensors in the return value of your module’s <code>forward</code> function. Please include the loss function and the structure of the return value of <code>forward</code> of your module when reporting this issue (e.g. list, dict, iterable).<br>
Parameter indices which did not receive grad for rank 2: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 …<br>
In addition, you can set the environment variable TORCH_DISTRIBUTED_DEBUG to either INFO or DETAIL to print out information about which particular parameters did not receive gradient on this rank as part of this error</p>
</blockquote>
<p>But I checked my code on single GPU, all params have their grad. I do not know what cause this error. It looks like all my params have no grad.</p>

HELLORPG: <p>If I use <code>use_reentrant=True</code>, there will be the same error but with different params:</p>
<blockquote>
<p>Parameter indices which did not receive grad for rank 1: 0 1 2 3 4 5 6 7 8 9 10 11 12<br>
13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41<br>
42 346 347 348 349 350 351 352 353 354 355 356 357 370 371 372 373 374 375 376 377 37<br>
8 379 380 381 382 383 384 385</p>
</blockquote>
<p>It makes me really confuse.</p>
<p>BTW, I really don’t understand what dose <code>reentrant</code> mean? If a part of my model’s outputs will be the inputs for next frame, should I use <code>True</code> or <code>False</code> for the argument <code>use_reentrant</code>?</p>

agu: <p>Is it possible to provide a minimal script to reproduce your error?</p>
<p>If your model’s computational graph does not change from iteration to iteration, can you try passing <code>static_graph=True</code> to the <code>DistributedDataParallel</code> constructor?</p>

HELLORPG: <p>My project is so huge, thus I can’t find a way to provide a small script. I will try it, but not easy.</p>
<p>My model is just like <a href="https://github.com/megvii-research/MOTR" rel="noopener nofollow ugc">megvii-research/MOTR: [ECCV2022] MOTR: End-to-End Multiple-Object Tracking with TRansformer (github.com)</a>, in their project, they do not use the official <code>checkpoint</code> but <a href="https://github.com/csrhddlam/pytorch-checkpoint" rel="noopener nofollow ugc">csrhddlam/pytorch-checkpoint (github.com)</a> instead, I don’t know why. But I want to use official <code>checkpoint</code> to make sure my code is more precise and maintainable.</p>
<p>And I can’t use <code>static_graph=True</code> because I have if condition in my code, for first frame some params will not be used, but in the second frame all params will be used. For an iter it will have at least two frame.</p>
<p>BTW, as above mentioned, I really don’t understand what dose <code>reentrant</code> mean? If a part of my model’s outputs will be the inputs for next frame, should I use <code>True</code> or <code>False</code> for the argument <code>use_reentrant</code> ?</p>
<p>Thanks for your reply~</p>

agu: <p>Let me check if we have activation checkpointing support for non-static graph models when using DDP and get back to you.</p>
<p>Regarding <code>reentrant</code>, there are two versions of activation checkpointing implemented in PyTorch today: one is so-called “reentrant” and the other is “non-reentrant”. At a high level, the non-reentrant one (<code>checkpoint(use_reentrant=False)</code>) is supposed to add support for additional use cases that the reentrant cannot handle, but there are some gaps in both versions at the moment. Your case of having a model’s outputs be inputs for the next frame should not affect which version to use.</p>

HELLORPG: <p>Truly appreciate for your reply ~</p>
<p>It makes my mind clear. For a long long long time, I misunderstood the <code>reentrant</code> and <code>non-reentrant</code> means two different models type (like LSTM and CNN). <img src="https://discuss.pytorch.org/images/emoji/apple/face_with_head_bandage.png?v=12" title=":face_with_head_bandage:" class="emoji" alt=":face_with_head_bandage:" loading="lazy" width="20" height="20"></p>

agu: <p>I checked and non-reentrant checkpoint should support DDP and non-static graph models. Could you try to pass <code>use_reentrant=False</code> and see if that works?</p>

HELLORPG: <p>I’ve done this before, I meet the same error mentioned above:</p>
<blockquote>
<p>RuntimeError: RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument <code>find_unused_parameters=True</code> to <code>torch.nn.parallel.DistributedDataParallel</code> , and by<br>
making sure all <code>forward</code> function outputs participate in calculating loss.<br>
If you already have done the above, then the distributed data parallel module wasn’t able to locate the output tensors in the return value of your module’s <code>forward</code> function. Please include the loss function and the structure of the return value of <code>forward</code> of your module when reporting this issue (e.g. list, dict, iterable).<br>
Parameter indices which did not receive grad for rank 2: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 …<br>
In addition, you can set the environment variable TORCH_DISTRIBUTED_DEBUG to either INFO or DETAIL to print out information about which particular parameters did not receive gradient on this rank as part of this error</p>
</blockquote>
<p>The params indices are different from <code>use_reentrant=True</code>. But without DDP, on single GPU this code will work fine. I believe this means these params do have grad on single GPU.</p>
<p>I find a not elegant solution, my model just like the pseudo-code as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">class Model:
    def forward(x1, x2):
        feature = self.backbone(x1)
        [some other code]
        x2_tensor = modified_x2_to_tensor(x2)
        res = self.transformer(feature, x2_tensor)
        [some other code]
        return res
</code></pre>
<p>The x1 is a NestTensor (same as Deformable DETR), x2 is a custom object list. x2 is empty in the first frame, and will be a part of the res from the first frame’s outputs when input to the second frame (same model for different frame in an iteration).</p>
<p>If I use checkpoint like this:</p>
<pre data-code-wrap="python"><code class="lang-python">model = Model()
......
x2 = init_x2()
for x1 in frames:
    res = torch.utils.checkpoint(model, x1, x2, use_reentrant=False)
    x2 = generate_x2(res)
.......
</code></pre>
<p>It will raise the above error.</p>
<p>But if I use it in the forward function like this:</p>
<pre data-code-wrap="python"><code class="lang-python">class Model:
    def forward(x1, x2):
        feature = torch.utils.checkpoint(self.backbone, x1, use_reentrant=False)
        # feature = self.backbone(x1)
        [some other code]
        x2_tensor = modified_x2_to_tensor(x2)
        res = torch.utils.checkpoint(self.transformer, feature, x2_tensor, use_reentrant=False)
        # res = self.transformer(feature, x2)
        [some other code]
        return res

model = Model()
......
x2 = init_x2()
for x1 in frames:
    res = model(x1, x2)
    x2 = generate_x2(res)
.......
</code></pre>
<p>It works! But I don’t know why. Don’t know whether above code helps you understand the details about my case.</p>
<p>I don’t know if the above pseudo-code confused you, if so, I’m truly sorry that I can’t find a better way to  show my code logic. <img src="https://discuss.pytorch.org/images/emoji/apple/disappointed.png?v=12" title=":disappointed:" class="emoji" alt=":disappointed:" loading="lazy" width="20" height="20"></p>
<p>BTW, I will still try to make a demo to reproduce my case, but it may case some time because my code is complex and I have a lot of deadlines lately. <img src="https://discuss.pytorch.org/images/emoji/apple/exploding_head.png?v=12" title=":exploding_head:" class="emoji" alt=":exploding_head:" loading="lazy" width="20" height="20"></p>
<p>Really really thank you for your patiently reply.</p>

rvarm1: <aside class="quote no-group" data-username="HELLORPG" data-post="8" data-topic="163187">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hellorpg/48/53892_2.png" class="avatar"> HELLORPG:</div>
<blockquote>
<p>It works! But I don’t know why. Don’t know whether above code helps you understand the details about my case.</p>
</blockquote>
</aside>
<p><a class="mention" href="/u/hellorpg">@HELLORPG</a> the reason this change works is because you are checkpointing the wrapped DDP module (i.e. the local module) instead of the entire DDP wrapped module, this will change the forward pass that is re-run.</p>

HELLORPG: <aside class="quote no-group quote-modified" data-username="rvarm1" data-post="9" data-topic="163187" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/rvarm1/48/15821_2.png" class="avatar"> rvarm1:</div>
<blockquote>
<aside class="quote no-group" data-username="HELLORPG" data-post="8" data-topic="163187">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hellorpg/48/53892_2.png" class="avatar"> HELLORPG:</div>
<blockquote>
<p>It works! But I don’t know why. Don’t know whether above code helps you understand the details about my case.</p>
</blockquote>
</aside>
<p><a class="mention" href="/u/hellorpg">@HELLORPG</a> the reason this change works is because you are checkpointing the wrapped DDP module (i.e. the local module) instead of the entire DDP wrapped module, this will change the forward pass that is re-run.</p>
</blockquote>
</aside>
<p>Yes, I think so. But I think this change does not change the code’s logic, just split a model into two modules. I don’t understand the difference between this two implementation. <img src="https://discuss.pytorch.org/images/emoji/apple/rofl.png?v=12" title=":rofl:" class="emoji" alt=":rofl:" loading="lazy" width="20" height="20"></p>

HELLORPG: <p>Hi, I do make a demo, not the same error raised, but this error is also what I met.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
from torch.optim import SGD


class Tracks:
    def __init__(self):
        self.embed = torch.zeros((1, 0, 10), dtype=torch.float)


def select_tracks(ts: Tracks):
    ts_len = len(ts.embed)
    idx = int(ts_len * 0.5)
    ts.embed = ts.embed[idx:]
    return ts


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=3, padding=1)
        self.query = nn.Parameter(torch.randn(5, 10))
        self.attn = nn.MultiheadAttention(embed_dim=10, num_heads=1, batch_first=True)

    def forward(self, f, t: Tracks):
        feature = self.conv(f).flatten(2).transpose(1, 2)
        query = torch.cat((t.embed, self.query.repeat(1, 1, 1)), dim=1)
        r, _ = self.attn(query, feature, feature)
        return r


if __name__ == '__main__':
    image1 = torch.randn((1, 3, 10, 10))
    image2 = torch.randn((1, 3, 10, 10))
    images = [image1, image2]
    tracks = Tracks()
    model = Model()
    criterion = nn.L1Loss()
    optimizer = SGD(model.parameters(), lr=0.00001)
    loss = torch.zeros((1,), dtype=float)
    for frame in images:
        from torch.utils.checkpoint import checkpoint
        # res = model(frame, tracks)
        res = checkpoint(model, frame, tracks, use_reentrant=False)
        gt = torch.zeros(res.shape, dtype=torch.float)
        loss += criterion(res, gt)
        tracks.embed = res
        tracks = select_tracks(tracks)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    print("Finish")
</code></pre>
<p>When I run this code, it will raise an error like this:</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “/home/gaoruopeng/Code/MeMOTR/demo.py”, line 59, in <br>
loss.backward()<br>
File “/home/gaoruopeng/anaconda3/envs/MeMOTR/lib/python3.10/site-packages/torch/_tensor.py”, line 396, in backward<br>
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)<br>
File “/home/gaoruopeng/anaconda3/envs/MeMOTR/lib/python3.10/site-packages/torch/autograd/<strong>init</strong>.py”, line 173, in backward<br>
Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass<br>
RuntimeError: mat1 and mat2 shapes cannot be multiplied (10x8 and 9x10)</p>
</blockquote>
<p>It’s not the same error that we talked before, but if I fix this error, I think my code will work fine. Cause the error we talked above happends because I want to bypass this error.</p>