Repository: pytorch/ignite

Issue:
Distributed model initialisation - CIFAR 10 example
## ❓ Questions/Help/Support

In the cifar10 example when a [model](https://github.com/pytorch/ignite/blob/master/examples/contrib/cifar10/main.py#L51) is defined inside the process, is the model updated and shared between the `nproc_per_node`?

In most of the TPU examples I have seen they define the model before the `Parallel` process.
- [tpu-training-super-fast-xlmroberta](https://www.kaggle.com/shonenkov/tpu-training-super-fast-xlmroberta)
- [i-like-clean-tpu-training-kernels-i-can-not-lie](https://www.kaggle.com/abhishek/i-like-clean-tpu-training-kernels-i-can-not-lie)

Does the way ignite handles the model ensure the same model is used throughout the training?


Discussion:
@ryanwongsa in the examples you cite seems like they define a model on CPU once and on each proc they move it to the corresponding device. If you take a look at [one of "official" pytorch/xla examples](https://github.com/pytorch/xla/blob/master/test/test_train_mp_imagenet.py), you can see that the model is created inside spawned function: https://github.com/pytorch/xla/blob/master/test/test_train_mp_imagenet.py#L179
I think this should be equivalent. 

- Grads synchronization happens inside `xm.optimizer_step(optimizer)`.
- There is also [MpModelWrapper](https://github.com/pytorch/xla/blob/master/torch_xla/distributed/xla_multiprocessing.py#L303) provided by `pytorch/xla` which helps to reduce memory usage (see class docstring for details). In this case, we can indeed define a single model instance and then move it to corresponding device inside spawned function.

HTH


EDIT: maybe we need to make it also clear by setting manual_seed just before creation of the model.
Ah, I see thanks for the links.
 
I have a second question relating to the CIFAR example. What is the best way to update it support mixed-precision with the ignite auto distribution methods.
Should something like:
```python
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
```
be done before or after the below code
```python
optimizer = idist.auto_optim(optimizer)
model = idist.auto_optim(model)
```

Resolving patch:
diff --git a/ignite/distributed/auto.py b/ignite/distributed/auto.py
--- a/ignite/distributed/auto.py
+++ b/ignite/distributed/auto.py
@@ -127,8 +127,8 @@ def auto_model(model: nn.Module) -> nn.Module:
 
     Internally, we perform to following:
 
-    - send model to current :meth:`~ignite.distributed.utils.device()`.
-    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1
+    - send model to current :meth:`~ignite.distributed.utils.device()` if model's parameters are not on the device.
+    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1.
     - wrap the model to `torch DataParallel`_ if no distributed context found and more than one CUDA devices available.
 
     Examples:
@@ -139,6 +139,15 @@ def auto_model(model: nn.Module) -> nn.Module:
 
         model = idist.auto_model(model)
 
+    In addition with NVidia/Apex, it can be used in the following way:
+
+    .. code-block:: python
+
+        import ignite.distribted as idist
+
+        model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
+        model = idist.auto_model(model)
+
     Args:
         model (torch.nn.Module): model to adapt.
 
@@ -150,7 +159,10 @@ def auto_model(model: nn.Module) -> nn.Module:
     """
     logger = setup_logger(__name__ + ".auto_model")
 
-    model.to(idist.device())
+    # Put model's parameters to device if its parameters are not on the device
+    device = idist.device()
+    if not all([p.device == device for p in model.parameters()]):
+        model.to(device)
 
     # distributed data parallel model
     if idist.get_world_size() > 1:


Test patch: