Repository: pytorch/ignite

Issue:
Consolidate ZeRO state before checkpoint saving
Is your feature request related to a problem? Please describe.
When checkpoint saving occurs there should be a check if the object that state_dict() is called on is a ZeroRedundancyOptimizer instance

Describe the solution you'd like
Prior to the call state_dict() call a consolidate_state_dict() call should be issued. This call needs to be issued on all ranks and point toward the same consolidating rank. 

Here there are two design solutions, you instantiate a Checkpoint handler on all ranks and only on the designated rank does it save the checkpoint or do you create a Handler that needs to run before the Checkpoint handler in order to consolidate.

Special care must be addressed to the PyTorch version as the naming scheme of the ZeRO's method arguments has changed between recent versions.

Describe alternatives you've considered
The only alternative is for the users to write themselves Handlers to do that, segregating the checkpoint-saving logic. And it can be written as:

```
class ConsolidateZeROHandler:
    """
    Handler that consolidated the Zero Redundancy Optimizer prior to the checkpoint Saving.

    Args:
        zero_optimizer (ZeroRedundancyOptimizer): The optimizer to be consolidated.
        recipient_rank (int): The rank on which the consolidation will happen. Defaults to 0.
        epoch_level (bool): Call every N epochs or every N iterations. `True` is epoch level, `False` is iteration
            level. Defaults to True.
    """

    def __init__(
        self,
        zero_optimizer: ZeroRedundancyOptimizer,
        call_every: int,
        recipient_rank: int = 0,
        epoch_level=True,
    ):
        self.zero_optimizer = zero_optimizer
        self.recipient_rank = recipient_rank
        self.epoch_level = epoch_level
        self.call_every = call_every

    def __call__(self, engine: Engine):
        self.zero_optimizer.consolidate_state_dict(to=self.recipient_rank)

    def attach(self, engine: Engine) -> None:
        if self.epoch_level:
            engine.add_event_handler(Events.EPOCH_COMPLETED(every=self.call_every), self)
        else:
            engine.add_event_handler(Events.ITERATION_COMPLETED(every=self.call_every), self)
```


Discussion:
Thanks for the feature request @danieltudosiu !

Yes, from user side, as you suggested one possible alternative solution could be to attach a specific handler for that:
```python
to_save = {"zero": zero_optimizer}
checkpoint = Checkpoint(to_save, ...)

recipient_rank = 0
trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda _: zero_optimizer.consolidate_state_dict(to=recipient_rank))
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint)
```

> Describe the solution you'd like
Prior to the call state_dict() call a consolidate_state_dict() call should be issued. This call needs to be issued on all ranks and point toward the same consolidating rank.

Do you think we should add this specific hard-coded into `Checkpoint` ? For example, we can assume that `consolidate_state_dict` is an API method and if object has it we can call it `obj.consolidate_state_dict(to=self.consolidate_recipient_rank)` before setting up the state dict:

https://github.com/pytorch/ignite/blob/e45846f850b8788351308859ace5b46368ae2033/ignite/handlers/checkpoint.py#L459-L466

Checkpoint class will have new `consolidate_recipient_rank` argument : 
```
class Checkpoint:
    def __init__(self, ..., consolidate_recipient_rank=0 ...):
        self.consolidate_recipient_rank = consolidate_recipient_rank
```


Hi @vfdev-5,

The whole decision boils down to a design decision. Do you want the user to take care of this specific checkpointing logic or do you want it to be seemingly integrated into Ignite?

> Do you think we should add this specific hard-coded into `Checkpoint` ? For example, we can assume that `consolidate_state_dict` is an API method and if object has it we can call it `obj.consolidate_state_dict(to=self.consolidate_recipient_rank)` before setting up the state dict: 
> Checkpoint class will have new `consolidate_recipient_rank` argument :
> 
> ```
> class Checkpoint:
>     def __init__(self, ..., consolidate_recipient_rank=0 ...):
>         self.consolidate_recipient_rank = consolidate_recipient_rank
> ```

Be careful cuz this is not the complete solution, consolidate_state_dict MUST be called on all ranks, which means that all ranks need to go through the checkpoint-saving logic so the checkpoint saver must be aware of the rank it is being called on so we do not save on all rank the checkpoints. 

Besides that, yes I would incline on integrating the ZeRO consolidation logic in the checkpoint class such that the user does not need to fuss around with it.

Hi @danieltudosiu thanks for your pointers !

Today, `Checkpoint` is intended to be called by all ranks (due to xla which requires that) 
https://github.com/pytorch/ignite/blob/e45846f850b8788351308859ace5b46368ae2033/ignite/handlers/checkpoint.py#L176-L185

and internally `DiskSaver` uses `torch.save` only on rank zero.

> Besides that, yes I would incline on integrating the ZeRO consolidation logic in the checkpoint class such that the user does not need to fuss around with it.

Sounds good to me as well. I'll send a PR and will ask to check it. Hope it works for you :)

Hi @vfdev-5,

Thanks for moving so quickly. 

For my personal use, I already coded the Handler so I am ok ;) I just wanted Ignite to be even more user-friendly <3 

Cheers,

Dan

Resolving patch:
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
--- a/ignite/handlers/checkpoint.py
+++ b/ignite/handlers/checkpoint.py
@@ -11,6 +11,14 @@
 
 import torch
 import torch.nn as nn
+from packaging.version import Version
+
+if Version(torch.__version__) >= Version("1.9.0"):
+    from torch.distributed.optim import ZeroRedundancyOptimizer
+
+    HAVE_ZERO = True
+else:
+    HAVE_ZERO = False
 
 import ignite.distributed as idist
 from ignite.base import Serializable
@@ -166,13 +174,14 @@ class Checkpoint(Serializable):
             > checkpoint_12345.pt
 
     Note:
-        This class is distributed configuration-friendly: it is not required to instantiate the class in rank 0
-        process only.
+        This class is distributed configuration-friendly: it is not required to instantiate the class in rank 0 only
+        process. This class supports automatically distributed configuration and if used with
+        :class:`~ignite.handlers.DiskSaver`, checkpoint is stored by rank 0 process.
 
     .. warning::
 
-        When running on XLA devices, it should be run in all processes, otherwise application can get stuck on
-        saving the checkpoint.
+        When running on XLA devices or using :class:`~torch.distributed.optim.ZeroRedundancyOptimizer`, it
+        should be run in all processes, otherwise application can get stuck while saving the checkpoint.
 
         .. code-block:: python
 
@@ -282,7 +291,7 @@ def __init__(
         filename_pattern: Optional[str] = None,
         include_self: bool = False,
         greater_or_equal: bool = False,
-        save_on_rank: Optional[int] = 0,
+        save_on_rank: int = 0,
     ):
 
         if not isinstance(to_save, collections.Mapping):
@@ -466,6 +475,10 @@ def _setup_checkpoint(self) -> Dict[str, Dict[Any, Any]]:
             for k, obj in self.to_save.items():
                 if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
                     obj = obj.module
+                elif HAVE_ZERO and isinstance(obj, ZeroRedundancyOptimizer):
+                    obj.consolidate_state_dict(to=self.save_on_rank)
+                    if self.save_on_rank != idist.get_rank():
+                        continue
                 checkpoint[k] = obj.state_dict()
         return checkpoint
 
@@ -782,7 +795,7 @@ def __init__(
         atomic: bool = True,
         create_dir: bool = True,
         require_empty: bool = True,
-        save_on_rank: Optional[int] = 0,
+        save_on_rank: int = 0,
         **kwargs: Any,
     ):
         self.dirname = Path(dirname).expanduser()
@@ -948,7 +961,7 @@ def __init__(
         filename_pattern: Optional[str] = None,
         include_self: bool = False,
         greater_or_equal: bool = False,
-        save_on_rank: Optional[int] = 0,
+        save_on_rank: int = 0,
         **kwargs: Any,
     ):
 


Test patch:
diff --git a/tests/ignite/handlers/test_checkpoint.py b/tests/ignite/handlers/test_checkpoint.py
--- a/tests/ignite/handlers/test_checkpoint.py
+++ b/tests/ignite/handlers/test_checkpoint.py
@@ -1243,9 +1243,37 @@ def _test_checkpoint_load_objects_ddp(device):
     Checkpoint.load_objects(to_load, checkpoint)
 
 
+def _test_checkpoint_with_ZeRO(device, dirname, local_rank):
+
+    from torch.distributed.optim import ZeroRedundancyOptimizer
+
+    model = DummyModel().to(device)
+    opt = ZeroRedundancyOptimizer(model.parameters(), torch.optim.SGD, lr=0.01)
+    mocked_opt = MagicMock(ZeroRedundancyOptimizer, wraps=opt)
+
+    # A `step` should be called to optimizer state get populated.
+    out = model(torch.Tensor([1.0]))
+    out.backward()
+    mocked_opt.step()
+
+    to_save = {"model": model, "optim": mocked_opt}
+    checkpointer = Checkpoint(to_save, dirname, save_on_rank=1)
+
+    engine = Engine(lambda e, b: None)
+    checkpointer(engine)
+
+    mocked_opt.consolidate_state_dict.assert_called_once_with(to=1)
+
+    if local_rank == 1:
+
+        loaded_state_dict = torch.load(dirname / "checkpoint_0.pt", map_location=device)["optim"]
+        state_dict = opt.state_dict()
+        assert loaded_state_dict == state_dict
+
+
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
-def test_distrib_gloo_cpu_or_gpu(distributed_context_single_node_gloo, get_rank_zero_dirname):
+def test_distrib_gloo_cpu_or_gpu(distributed_context_single_node_gloo, dirname, get_rank_zero_dirname, local_rank):
 
     device = idist.device()
     rank_zero_dirname = get_rank_zero_dirname()
@@ -1254,6 +1282,11 @@ def test_distrib_gloo_cpu_or_gpu(distributed_context_single_node_gloo, get_rank_
     _test_checkpoint_with_ddp(device)
     _test_checkpoint_load_objects_ddp(device)
 
+    from ignite.handlers.checkpoint import HAVE_ZERO
+
+    if HAVE_ZERO:
+        _test_checkpoint_with_ZeRO(device, dirname, local_rank)
+
 
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")