Right ways to serialize and load DDP model checkpoints

username: Sayak_Paul
cooked: <p>I have trained a model using <code>DistributedDataParallel</code>. After training, I serialized the model like so where the model is wrapped using <code>DistributedDataParallel</code>:</p>
<pre><code class="lang-python">torch.save(model.state_dict(), 'model.pt')
</code></pre>
<p><strong>Note</strong> that this serialization was performed in the launcher function which is typically passed to <code>spawn()</code> of <code>torch.multiprocessing</code>. My training setup consists of 4 GPUs.</p>
<p>Now when I am trying to load the checkpoint in my local inference setup (single GPU) the keys are not matching. The model, in this case, is not wrapped using <code>DistributedDataParallel</code>. Any pointers would be useful.</p>

username: sio277
cooked: <p>Save your DDP model after unwrapping DataParallel, such as,</p>
<pre><code class="lang-auto">torch.save(model.module.state_dict(), 'model.pt')
</code></pre>
<p>Here, <code>model.module</code> is where your original model (before DDP wrapping) is placed.</p>

username: Sayak_Paul
cooked: <p>Thanks <a class="mention" href="/u/sio277">@sio277</a>.</p>
<p>Also, a slightly unrelated question on how to best log/print the loss and other metrics in these settings. Currently, I am only logging loss and other metrics from the master i.e. when <code>rank == 0</code>.</p>

username: sio277
cooked: <p>To get mean metrics all across the ranks, I use all-reduce function something like this:</p>
<pre><code class="lang-auto">import torch.distributed as dist

def global_meters_all_avg(args, *meters):
    """meters: scalar values of loss/accuracy calculated in each rank"""
    tensors = [torch.tensor(meter, device=args.gpu, dtype=torch.float32) for meter in meters]
    for tensor in tensors:
        # each item of `tensors` is all-reduced starting from index 0 (in-place)
        dist.all_reduce(tensor)

    return [(tensor / args.world_size).item() for tensor in tensors]
</code></pre>
<p>See <a href="https://pytorch.org/docs/stable/distributed.html#distributed-communication-package-torch-distributed" class="inline-onebox" rel="noopener nofollow ugc">Distributed communication package - torch.distributed — PyTorch 1.8.1 documentation</a> for more details about the collective communications.</p>

username: Sayak_Paul
cooked: <p>Thank you very much. I suppose <code>meter</code> is the metric quantity?</p>

username: sio277
cooked: <p>Yes, such as loss and accuracy.</p>

username: Sayak_Paul
cooked: <p>I see. Could you also share a minimal example as to where <code>global_meters_all_avg()</code> should be placed inside the training loop? Let’s say the following is my loop (part of the launcher <code>train()</code> function called by <code>mp.spawn()</code>:</p>
<pre><code class="lang-python">for batch in pbar:
    # load image and mask into device memory
    image = batch['image'].cuda(rank, non_blocking=True)
    mask = batch['mask'].cuda(rank, non_blocking=True)

    # pass images into model
    pred = model(image)

    # get loss
    loss = criteria(pred, mask)

    # update the model
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
</code></pre>

username: sio277
cooked: <p>I usually use that function after each epoch (not during the batch iteration), to avoid worsening the training speed. During the batch iterations, I accumulate the <code>loss</code> values, and after 1 epoch, let <code>global_meters_all_avg</code> be called with an input of the accumulated loss value.</p>

username: sio277
cooked: <p>Note that for accumulating the <code>loss</code> in each rank, I use this class:</p>
<pre><code class="lang-auto">class AvgMeter:
    def __init__(self):
        self.reset()

    def reset(self):
        self.sum = 0
        self.count = 0
        self.avg = 0

    def update(self, val, n=1):
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
</code></pre>
<p>In your code, you can use this class something like this:</p>
<pre><code class="lang-auto">losses = AvgMeter()
for batch in pbar:
    # load image and mask into device memory
    image = batch['image'].cuda(rank, non_blocking=True)
    mask = batch['mask'].cuda(rank, non_blocking=True)

    # pass images into model
    pred = model(image)

    # get loss
    loss = criteria(pred, mask)
    losses.update(loss.item(), image.size(0))  # accumulate the loss

    # update the model
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# after each epoch
loss = losses.avg
global_loss = global_meters_all_avg(args, loss) 
</code></pre>
<p>The <code>global_loss</code> is the one all-reduced (averaged) across the ranks.</p>

username: Sayak_Paul
cooked: <p>Thanks so much. Really appreciate your help here.</p>

username: Sayak_Paul
cooked: <p>One more doubt is during logging or printing <code>global_loss</code> I think it needs to be print only from one rank to prevent duplicate entries. Something like <code>if rank == 0: print(global_loss)</code>.</p>

username: sio277
cooked: <p>It depends on your preference <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=9" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:">.</p>