Repository: pytorch/ignite

Issue:
Pass kwargs to idist.auto_model
## 🚀 Feature

It would be nice to have a possibility to pass kwargs to `idist.auto_model`.
For example, a use-case if using multi-gpu DDP model training, I'd like to pass `find_unused_parameters=True` to 
`torch.nn.parallel.DistributedDataParallel`. 
```python
    # model = idist.auto_model(model, find_unused_parameters=True)

    import torch
    lrank = idist.get_local_rank()
    model = model.to(idist.device())
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,], find_unused_parameters=True)
```



Discussion:
@ryanwongsa here I'd like to pass a kwargs to `auto_*` method, namely to pass a parameter like `find_unused_parameters` which makes sense only for DDP GPUs... Do you think that this is similarly misleading as for the case with `barrier` for `auto_optim` in the other issue ?
I think to me adding the above type of parameter (`find_unused_parameters`) is fine as in this case I don't think it requires other additional code. In the the barrier case in `auto_optim`, you would need to add `xm.mark_step()` if you set `barrier=False` so it could be misleading. Hopefully that makes sense.
@ryanwongsa thanks for the feedback. Just in case, if you are interested in helping with this issue, do not hesitate to send a PR with the fix :) 

Resolving patch:
diff --git a/ignite/distributed/auto.py b/ignite/distributed/auto.py
--- a/ignite/distributed/auto.py
+++ b/ignite/distributed/auto.py
@@ -125,7 +125,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> Union[DataLoader, "_MpDe
     return dataloader
 
 
-def auto_model(model: nn.Module, sync_bn: bool = False) -> nn.Module:
+def auto_model(model: nn.Module, sync_bn: bool = False, **kwargs: Any) -> nn.Module:
     """Helper method to adapt provided model for non-distributed and distributed configurations (supporting
     all available backends from :meth:`~ignite.distributed.utils.available_backends()`).
 
@@ -158,6 +158,8 @@ def auto_model(model: nn.Module, sync_bn: bool = False) -> nn.Module:
         sync_bn (bool): if True, applies `torch convert_sync_batchnorm`_ to the model for native torch
             distributed only. Default, False. Note, if using Nvidia/Apex, batchnorm conversion should be
             applied before calling ``amp.initialize``.
+        **kwargs: kwargs to model's wrapping class: `torch DistributedDataParallel`_ or `torch DataParallel`_
+            if applicable. Please, make sure to use acceptable kwargs for given backend.
 
     Returns:
         torch.nn.Module
@@ -184,16 +186,19 @@ def auto_model(model: nn.Module, sync_bn: bool = False) -> nn.Module:
                 logger.info("Convert batch norm to sync batch norm")
                 model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
 
+            if "device_ids" in kwargs:
+                raise ValueError(f"Argument kwargs should not contain 'device_ids', but got {kwargs}")
+
             lrank = idist.get_local_rank()
             logger.info(f"Apply torch DistributedDataParallel on model, device id: {lrank}")
-            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,])
+            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,], **kwargs)
         elif idist.has_native_dist_support and bnd == idist_native.GLOO:
             if sync_bn:
                 logger.info("Convert batch norm to sync batch norm")
                 model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
 
             logger.info("Apply torch DistributedDataParallel on model")
-            model = torch.nn.parallel.DistributedDataParallel(model)
+            model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         elif idist.has_hvd_support and bnd == idist_hvd.HOROVOD:
             import horovod.torch as hvd
 
@@ -203,7 +208,7 @@ def auto_model(model: nn.Module, sync_bn: bool = False) -> nn.Module:
     # not distributed but multiple GPUs reachable so data parallel model
     elif torch.cuda.device_count() > 1 and "cuda" in idist.device().type:
         logger.info("Apply torch DataParallel on model")
-        model = torch.nn.parallel.DataParallel(model)
+        model = torch.nn.parallel.DataParallel(model, **kwargs)
 
     return model
 


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
@@ -51,14 +51,16 @@ def _test_auto_dataloader(ws, nproc, batch_size, num_workers=1, sampler_name=Non
         assert dataloader.pin_memory == ("cuda" in idist.device().type)
 
 
-def _test_auto_model(model, ws, device, sync_bn=False):
-    model = auto_model(model, sync_bn=sync_bn)
+def _test_auto_model(model, ws, device, sync_bn=False, **kwargs):
+    model = auto_model(model, sync_bn=sync_bn, **kwargs)
     bnd = idist.backend()
     if ws > 1 and device in ("cuda", "cpu"):
-        if idist.has_native_dist_support and bnd in ("nccl" or "gloo"):
+        if idist.has_native_dist_support and bnd in ("nccl", "gloo"):
             assert isinstance(model, nn.parallel.DistributedDataParallel)
             if sync_bn:
                 assert any([isinstance(m, nn.SyncBatchNorm) for m in model.modules()])
+            if "find_unused_parameters" in kwargs:
+                assert model.find_unused_parameters == kwargs["find_unused_parameters"]
         elif idist.has_hvd_support and bnd in ("horovod",):
             assert isinstance(model, nn.Module)
     elif device != "cpu" and torch.cuda.is_available() and torch.cuda.device_count() > 1:
@@ -78,6 +80,9 @@ def _test_auto_model_optimizer(ws, device):
 
     model = nn.Sequential(nn.Linear(20, 100), nn.BatchNorm1d(100))
     _test_auto_model(model, ws, device, sync_bn="cuda" in device)
+    if ws > 1:
+        _test_auto_model(model, ws, device, find_unused_parameters=True)
+        _test_auto_model(model, ws, device, find_unused_parameters=False)
 
     # Test auto_optim
     bnd = idist.backend()