Repository: pytorch/ignite

Issue:
Checkpoint and ModelCheckpoint automatically saves base model for DDP wrapper
Idea is to avoid saving [DDP wrapped](https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel) model with `Checkpoint`/`ModelCheckpoint` classes and save instead its base model. Such that state dictionary wont contain prefix `module.` in the keys.

Discussion:

Patch:
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
index dd62d70f..3e2d6d01 100644
--- a/ignite/handlers/checkpoint.py
+++ b/ignite/handlers/checkpoint.py
@@ -8,6 +8,7 @@ from collections import namedtuple
 from typing import Callable, Mapping, Optional, Union
 
 import torch
+import torch.nn as nn
 
 from ignite.engine import Engine, Events
 
@@ -58,7 +59,9 @@ class Checkpoint:
 
     Args:
         to_save (Mapping): Dictionary with the objects to save. Objects should have implemented `state_dict` and `
-            load_state_dict` methods.
+            load_state_dict` methods. If contains objects of type torch `DistributedDataParallel`_ or
+            `DataParallel`_, their internal wrapped model is automatically saved (to avoid additional key ``module.`` in
+            the state dictionary).
         save_handler (callable or :class:`~ignite.handlers.checkpoint.BaseSaveHandler`): Method or callable class to
             use to save engine and other provided objects. Function receives two objects: checkpoint as a dictionary
             and filename. If `save_handler` is callable class, it can
@@ -79,6 +82,9 @@ class Checkpoint:
             To setup global step from another engine, please use :meth:`~ignite.handlers.global_step_from_engine`.
         archived (bool, optional): Deprecated argument as models saved by `torch.save` are already compressed.
 
+    .. _DistributedDataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel
+    .. _DataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel
+
     Note:
         This class stores a single file as a dictionary of provided objects to save.
         The filename has the following structure: `{filename_prefix}_{name}_{suffix}.{ext}` where
@@ -284,6 +290,8 @@ class Checkpoint:
     def _setup_checkpoint(self) -> dict:
         checkpoint = {}
         for k, obj in self.to_save.items():
+            if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
+                obj = obj.module
             checkpoint[k] = obj.state_dict()
         return checkpoint
 


Test patch:
diff --git a/tests/ignite/handlers/test_checkpoint.py b/tests/ignite/handlers/test_checkpoint.py
index 495d9272..4d74e54c 100644
--- a/tests/ignite/handlers/test_checkpoint.py
+++ b/tests/ignite/handlers/test_checkpoint.py
@@ -110,6 +110,24 @@ def test_checkpoint_default():
     _test(to_save, {"model": model.state_dict(), "optimizer": optimizer.state_dict()}, "checkpoint")
 
 
+def test_checkpoint_with_dp():
+
+    model = DummyModel()
+    dp_model = nn.DataParallel(model)
+    to_save = {"model": dp_model}
+
+    save_handler = MagicMock(spec=BaseSaveHandler)
+    checkpointer = Checkpoint(to_save, save_handler=save_handler)
+
+    trainer = Engine(lambda e, b: None)
+    trainer.state = State(epoch=0, iteration=0)
+
+    checkpointer(trainer)
+    assert save_handler.call_count == 1
+    metadata = {"basename": "model", "score_name": None, "priority": 0}
+    save_handler.assert_called_with(model.state_dict(), "model_0.pt", metadata)
+
+
 def test_checkpoint_with_global_step_transform():
     def _test(filename_prefix, to_save, obj, name):
         save_handler = MagicMock(spec=BaseSaveHandler)
@@ -867,3 +885,33 @@ def test_disksaver_wrong_input(dirname):
             DiskSaver(dirname, require_empty=True)
 
     _test(".pt")
+
+
+def _test_checkpoint_with_ddp(device):
+    model = DummyModel().to(device)
+    ddp_model = nn.parallel.DistributedDataParallel(model)
+    to_save = {"model": ddp_model}
+
+    save_handler = MagicMock(spec=BaseSaveHandler)
+    checkpointer = Checkpoint(to_save, save_handler=save_handler)
+
+    trainer = Engine(lambda e, b: None)
+    trainer.state = State(epoch=0, iteration=0)
+
+    checkpointer(trainer)
+    assert save_handler.call_count == 1
+    metadata = {"basename": "model", "score_name": None, "priority": 0}
+    save_handler.assert_called_with(model.state_dict(), "model_0.pt", metadata)
+
+
+@pytest.mark.distributed
+def test_checkpoint_with_ddp_gloo(distributed_context_single_node_gloo):
+    device = "cpu"
+    _test_checkpoint_with_ddp(device=device)
+
+
+@pytest.mark.distributed
+@pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
+def test_checkpoint_with_ddp_nccl(local_rank, distributed_context_single_node_nccl):
+    device = "cuda:{}".format(local_rank)
+    _test_checkpoint_with_ddp(device=device)