Repository: pytorch/ignite

Issue:
Do not wrap existing DistributedSampler in DistributedProxySampler in auto_dataloader()
## 🚀 Feature

I've looked at the implementation of ``ignite.distributed.auto_dataloader()`` and found that when the data loader has a sampler set, it will always wrap this sampler in a `DistributedProxySampler`. This works fine when the data loader sampler is of type `Sampler` and does itself not handle distribution of samples across processes. However, in my case I want to distribute samples differently across processes to minimize cache misses with my custom map-style dataset which buffers chunks of data. This places some constraints on how indices should be grouped based on rank, i.e., which indices should be assigned to the same process. Currently, the only work-around is to use an iterable dataset instead.

However, I would think that the lines:
https://github.com/pytorch/ignite/blob/0f7819a5d2f5c17697eaf4f01848a54ed43e6e53/ignite/distributed/auto.py#L86-L89

could be changed to keep the user provided instance (of a subclass) of `DistributedSampler` if set.



Discussion:
@aschuh-hf Thank you for this request !!

Maybe you could try providing the `batch_sampler` rather than the `sampler` ? Anyway, we will have a look ! 
@aschuh-hf thanks for FR and giving the context, it helps to understand !

> However, I would think that the lines: ... could be changed to keep the user provided instance (of a subclass) of DistributedSampler if set.

Yes, I agree that we can think about adding this possiblity to skip wraping `sampler` as `DistributedSampler` with `DistributedProxySampler`.

@sdesrozis I see, providing a `BatchSampler` that wraps my distributed sampler is a better option indeed than using a proxy `IterableDataset` that wraps the map-style dataset and custom sampler.

Still, having a distributed sampler wrapped by another distributed sampler doesn't seem quite right. I think it should at least trigger a warning or error if not supported by Ignite? But passing on the users distributed sampler to the data loader should be an easy enough change, possibly checking only that the `rank` and `num_replicas` attributes of the existing sampler matches with `ignite.distributed.get_rank()` and `ignite.distributed.get_world_size()`.

I can provide a PR for this.
> Still, having a distributed sampler wrapped by another distributed sampler doesn't seem quite right. I think it should at least trigger a warning or error if not supported by Ignite? But passing on the users distributed sampler to the data loader should be an easy enough change, possibly checking only that the rank and num_replicas attributes of the existing sampler matches with ignite.distributed.get_rank() and ignite.distributed.get_world_size().

Yes, I agree. I was thinking we were raising an issue when wraping the sampler with `DistributedProxySampler`, but actually it only checks the type as Sampler: https://github.com/pytorch/ignite/blob/0f7819a5d2f5c17697eaf4f01848a54ed43e6e53/ignite/distributed/auto.py#L296-L297

It would be better to raise an issue in the constructor to avoid wraping `DistributedSampler` with `DistributedProxySampler`.

>  I can provide a PR for this.

Sounds great ! Let me assign it to you.

Resolving patch:
diff --git a/ignite/distributed/auto.py b/ignite/distributed/auto.py
--- a/ignite/distributed/auto.py
+++ b/ignite/distributed/auto.py
@@ -25,8 +25,9 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> Union[DataLoader, "_MpDe
 
     - batch size is scaled by world size: ``batch_size / world_size`` if larger or equal world size.
     - number of workers is scaled by number of local processes: ``num_workers / nprocs`` if larger or equal world size.
-    - if no sampler provided by user, `torch DistributedSampler`_ is setup.
-    - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.
+    - if no sampler provided by user, a `torch DistributedSampler`_ is setup.
+    - if a `torch DistributedSampler`_ is provided by user, it is used without wrapping it.
+    - if another sampler is provided, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.
     - if the default device is 'cuda', `pin_memory` is automatically set to `True`.
 
     .. warning::
@@ -83,18 +84,22 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> Union[DataLoader, "_MpDe
                     "Please, make sure that the dataset itself produces different data on different ranks."
                 )
             else:
-                if kwargs.get("sampler", None) is not None:
-                    sampler = DistributedProxySampler(
-                        kwargs["sampler"], num_replicas=world_size, rank=rank
-                    )  # type: Union[DistributedProxySampler, DistributedSampler, Sampler]
+                sampler: Optional[Union[DistributedProxySampler, DistributedSampler, Sampler]]
+                sampler = kwargs.get("sampler", None)
+                if isinstance(sampler, DistributedSampler):
+                    if sampler.rank != rank:
+                        warnings.warn(f"Found distributed sampler with rank={sampler.rank}, but process rank is {rank}")
+                    if sampler.num_replicas != world_size:
+                        warnings.warn(
+                            f"Found distributed sampler with num_replicas={sampler.num_replicas}, "
+                            f"but world size is {world_size}"
+                        )
+                elif sampler is None:
+                    # removes "shuffle" from kwargs if sampler is used
+                    shuffle = kwargs.pop("shuffle", True)
+                    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=shuffle)
                 else:
-                    sampler = DistributedSampler(
-                        dataset, num_replicas=world_size, rank=rank, shuffle=kwargs.get("shuffle", True)
-                    )
-                    # we need to remove "shuffle" from kwargs if sampler is used
-                    if "shuffle" in kwargs:
-                        del kwargs["shuffle"]
-
+                    sampler = DistributedProxySampler(sampler, num_replicas=world_size, rank=rank)
                 kwargs["sampler"] = sampler
         else:
             warnings.warn(


Test patch:
diff --git a/tests/ignite/distributed/test_auto.py b/tests/ignite/distributed/test_auto.py
--- a/tests/ignite/distributed/test_auto.py
+++ b/tests/ignite/distributed/test_auto.py
@@ -16,6 +16,12 @@
 
 
 class DummyDS(Dataset):
+    def __init__(self, length=10):
+        self.length = length
+
+    def __len__(self):
+        return self.length
+
     def __getitem__(self, index):
         return index
 
@@ -43,6 +49,24 @@ def test_auto_dataloader_warning(distributed_context_single_node_gloo):
         )
 
 
+@pytest.mark.distributed
+@pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
+@pytest.mark.skipif("WORLD_SIZE" not in os.environ, reason="Skip if WORLD_SIZE not in env vars")
+def test_auto_dataloader_warning_distributed_sampler(distributed_context_single_node_gloo):
+    dataset = DummyDS()
+    rank = idist.get_rank()
+    world_size = idist.get_world_size()
+    auto_dataloader(dataset, sampler=DistributedSampler(dataset, num_replicas=world_size, rank=rank))
+    if world_size > 1:
+        wrong_rank = (rank + 1) % world_size
+        expected_warning = f"Found distributed sampler with rank={wrong_rank}, but process rank is {rank}"
+        with pytest.warns(UserWarning, match=expected_warning):
+            auto_dataloader(dataset, sampler=DistributedSampler(dataset, num_replicas=world_size, rank=wrong_rank))
+    expected_warning = f"Found distributed sampler with num_replicas={world_size + 1}, but world size is {world_size}"
+    with pytest.warns(UserWarning, match=expected_warning):
+        auto_dataloader(dataset, sampler=DistributedSampler(dataset, num_replicas=world_size + 1, rank=rank))
+
+
 @pytest.mark.tpu
 @pytest.mark.skipif("NUM_TPU_WORKERS" in os.environ, reason="Skip if NUM_TPU_WORKERS in env vars")
 @pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
@@ -57,6 +81,8 @@ def _test(data):
             sampler = None
         elif sampler_name == "WeightedRandomSampler":
             sampler = WeightedRandomSampler(weights=torch.ones(100), num_samples=100)
+        elif sampler_name == "DistributedSampler":
+            sampler = DistributedSampler(data, num_replicas=ws, rank=idist.get_rank())
         else:
             raise RuntimeError(f"Unknown sampler name: {sampler_name}")
 
@@ -83,7 +109,10 @@ def _test(data):
         if isinstance(data, IterableDataset):
             sampler_type = _InfiniteConstantSampler
         elif ws > 1:
-            sampler_type = DistributedSampler if sampler is None else DistributedProxySampler
+            if sampler is None or isinstance(sampler, DistributedSampler):
+                sampler_type = DistributedSampler
+            else:
+                sampler_type = DistributedProxySampler
         else:
             sampler_type = RandomSampler if sampler is None else type(sampler)
 
@@ -149,6 +178,7 @@ def test_auto_methods_no_dist():
     _test_auto_dataloader(1, 1, batch_size=1)
     _test_auto_dataloader(1, 1, batch_size=10, num_workers=2)
     _test_auto_dataloader(1, 1, batch_size=10, sampler_name="WeightedRandomSampler")
+    _test_auto_dataloader(1, 1, batch_size=10, sampler_name="DistributedSampler")
 
     _test_auto_model_optimizer(1, "cuda" if torch.cuda.is_available() else "cpu")
 
@@ -161,6 +191,7 @@ def test_auto_methods_gloo(distributed_context_single_node_gloo):
     _test_auto_dataloader(ws=ws, nproc=ws, batch_size=1)
     _test_auto_dataloader(ws=ws, nproc=ws, batch_size=10, num_workers=2)
     _test_auto_dataloader(ws=ws, nproc=ws, batch_size=10, sampler_name="WeightedRandomSampler")
+    _test_auto_dataloader(ws=ws, nproc=ws, batch_size=10, sampler_name="DistributedSampler")
 
     device = idist.device()
     _test_auto_model_optimizer(ws, device)
@@ -181,6 +212,7 @@ def test_auto_methods_nccl(distributed_context_single_node_nccl):
     _test_auto_dataloader(ws=ws, nproc=ws, batch_size=1)
     _test_auto_dataloader(ws=ws, nproc=ws, batch_size=10, num_workers=10)
     _test_auto_dataloader(ws=ws, nproc=ws, batch_size=1, sampler_name="WeightedRandomSampler")
+    _test_auto_dataloader(ws=ws, nproc=ws, batch_size=1, sampler_name="DistributedSampler")
 
     device = idist.device()
     _test_auto_model_optimizer(ws, device)
@@ -201,6 +233,7 @@ def test_auto_methods_hvd(gloo_hvd_executor):
     gloo_hvd_executor(_test_auto_dataloader, args=(np, np, 1), np=np, do_init=True)
     gloo_hvd_executor(_test_auto_dataloader, args=(np, np, 10, 10), np=np, do_init=True)
     gloo_hvd_executor(_test_auto_dataloader, args=(np, np, 1, 1, "WeightedRandomSampler"), np=np, do_init=True)
+    gloo_hvd_executor(_test_auto_dataloader, args=(np, np, 1, 1, "DistributedSampler"), np=np, do_init=True)
 
     gloo_hvd_executor(_test_auto_model_optimizer, args=(np, device), np=np, do_init=True)
 
@@ -223,6 +256,7 @@ def _test_auto_methods_xla(index, ws):
     _test_auto_dataloader(ws=ws, nproc=ws, batch_size=1, dl_type=dl_type)
     _test_auto_dataloader(ws=ws, nproc=ws, batch_size=10, num_workers=2, dl_type=dl_type)
     _test_auto_dataloader(ws=ws, nproc=ws, batch_size=1, sampler_name="WeightedRandomSampler", dl_type=dl_type)
+    _test_auto_dataloader(ws=ws, nproc=ws, batch_size=1, sampler_name="DistributedSampler", dl_type=dl_type)
 
     device = "xla"
     _test_auto_model_optimizer(ws, device)