Merge branch 'master' into bugfix-#6947-Checkpoint-issue-when-using-Horovod-distributed-backend

Author: awaelchli

.github/workflows/ci_test-base.yml

@@ -31,6 +31,7 @@ jobs:
     - name: Setup macOS
       if: runner.os == 'macOS'
       run: |
+        brew update  # todo: find a better fix (libomp error)
         brew install libomp  # https://github.com/pytorch/pytorch/issues/20030
 
     # Note: This uses an internal pip API and may not always work

.github/workflows/ci_test-full.yml

@@ -41,6 +41,7 @@ jobs:
     - name: Setup macOS
       if: runner.os == 'macOS'
       run: |
+        brew update  # todo: find a better fix (libomp error)
         brew install libomp  # https://github.com/pytorch/pytorch/issues/20030
         brew install openmpi libuv  # Horovod on macOS requires OpenMPI, Gloo not currently supported
 

CHANGELOG.md

@@ -240,6 +240,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed `AttributeError for `require_backward_grad_sync` when running manual optimization with sharded plugin ([#6915](https://github.com/PyTorchLightning/pytorch-lightning/pull/6915))
 
 
+- Fixed `sync_dist` for tpus ([#6950](https://github.com/PyTorchLightning/pytorch-lightning/pull/6950))
+
+
+- Fixed `self.device` not returning the correct device in replicas of data-parallel ([#6414](https://github.com/PyTorchLightning/pytorch-lightning/pull/6414))
+
+
 ## [1.2.7] - 2021-04-06
 
 ### Fixed

dockers/base-cuda/Dockerfile

@@ -113,6 +113,10 @@ RUN \
     pip install --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./apex && \
     rm -rf apex
 
+RUN \
+    # install DeepSpeed
+    pip install deepspeed>=0.3.14
+
 RUN \
     # Show what we have
     pip --version && \

pytorch_lightning/accelerators/accelerator.py

@@ -106,7 +106,7 @@ def pre_dispatch(self, trainer: 'pl.Trainer') -> None:
         self.precision_plugin.pre_dispatch()
 
     def post_dispatch(self, trainer: 'pl.Trainer') -> None:
-        """Hook to do something before the training/evaluation/prediction starts."""
+        """Hook to do something after the training/evaluation/prediction starts."""
         self.training_type_plugin.post_dispatch()
         self.precision_plugin.post_dispatch()
 

pytorch_lightning/core/step_result.py

@@ -21,7 +21,7 @@
 from torch import Tensor
 from torchmetrics import Metric
 
-from pytorch_lightning.utilities.distributed import sync_ddp_if_available
+from pytorch_lightning.utilities.distributed import sync_ddp_if_available, tpu_distributed
 
 
 class Result(Dict):
@@ -105,10 +105,11 @@ def log(
 
         # sync across workers when using distributed training
         sync_fn = sync_fn or sync_ddp_if_available
+
         if sync_dist and isinstance(value, (torch.Tensor, numbers.Number)):
             is_dist_initialized = torch.distributed.is_available() and torch.distributed.is_initialized()
             # TODO: Find a way to make the reduction only once, so we don't need to clone.
-            if is_dist_initialized and isinstance(value, torch.Tensor):
+            if (is_dist_initialized or tpu_distributed) and isinstance(value, torch.Tensor):
                 value = value.clone()
             else:
                 value = torch.tensor(value, device=device, dtype=torch.float)

pytorch_lightning/overrides/data_parallel.py

@@ -22,6 +22,7 @@
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.overrides.base import _LightningModuleWrapperBase
 from pytorch_lightning.overrides.distributed import LightningDistributedModule
+from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.apply_func import apply_to_collection
 
 
@@ -71,6 +72,8 @@ def __init__(self, pl_module: LightningModule):
         super().__init__(pl_module)
 
     def forward(self, *inputs, **kwargs):
+        self.update_replica_device_attributes(inputs)
+        # forward call will redirect to training_step, validation_step, etc.
         output = super().forward(*inputs, **kwargs)
 
         def output_transform(data: Any):
@@ -85,6 +88,37 @@ def output_transform(data: Any):
         )
         return output
 
+    def update_replica_device_attributes(self, inputs: Any) -> None:
+        """
+        Updates the device information of LightningModule by reading the device from the inputs.
+        In :class:`~torch.nn.data_parallel.DataParallel` changes to the state during the `forward` pass
+        are lost when the replicas get discarded. The only way to know the current device is from the
+        inputs passed into the model.
+
+        Args:
+            inputs: A collection of inputs (typically a tuple). If the inputs don't contain tensors,
+                a warning is shown that accessing ``self.device`` will not return the correct device.
+        """
+        replica_device = None
+
+        def find_tensor_with_device(tensor: torch.Tensor) -> torch.Tensor:
+            nonlocal replica_device
+            if replica_device is None and tensor.device != torch.device("cpu"):
+                replica_device = tensor.device
+            return tensor
+
+        apply_to_collection(inputs, dtype=torch.Tensor, function=find_tensor_with_device)
+
+        if replica_device is not None:
+            # by calling .to() we force the update to the self.device property
+            self.module.to(device=replica_device)
+        else:
+            rank_zero_warn(
+                "Could not determine on which device the inputs are."
+                " When using DataParallel (accelerator='dp'), be aware that in case you are using self.device"
+                " in your code, it will reference only the root device."
+            )
+
 
 def python_scalar_to_tensor(data: Any, device: torch.device = torch.device("cpu")) -> Any:
     """ Converts a Python scalar number to a torch tensor and places it on the given device. """

pytorch_lightning/plugins/training_type/deepspeed.py

@@ -102,6 +102,7 @@ def __init__(
         cpu_checkpointing: bool = False,
         contiguous_memory_optimization: bool = False,
         synchronize_checkpoint_boundary: bool = False,
+        save_full_weights: bool = True,
     ) -> None:
         """
 
@@ -177,11 +178,16 @@ def __init__(
                 Not supported by all models
 
             synchronize_checkpoint_boundary: Insert :func:`torch.cuda.synchronize` at each checkpoint boundary.
+
+            save_full_weights: Gathers weights across all processes before saving to disk
+                when using ZeRO Stage 3. This allows a single weight file to contain the entire model,
+                rather than individual sharded weight files.
+                Disable to save sharded states individually. (Default: True)
         """
         if not _DEEPSPEED_AVAILABLE:
             raise MisconfigurationException(
                 "To use the DeepSpeed plugin, you must have DeepSpeed installed."
-                " pip install deepspeed mpi4py"
+                " pip install deepspeed"
             )
         super().__init__(
             parallel_devices=parallel_devices, num_nodes=num_nodes, cluster_environment=cluster_environment
@@ -205,11 +211,13 @@ def __init__(
                 allgather_partitions=allgather_partitions,
                 reduce_scatter=reduce_scatter,
                 allgather_bucket_size=allgather_bucket_size,
-                reduce_bucket_size=reduce_bucket_size
+                reduce_bucket_size=reduce_bucket_size,
             )
         self._config_initialized = False
         deepspeed.utils.logging.logger.setLevel(logging_level)
 
+        self.save_full_weights = save_full_weights
+
         # default FP16 parameters.
         self.loss_scale = loss_scale
         self.initial_scale_power = initial_scale_power
@@ -472,17 +480,27 @@ def save_checkpoint(self, checkpoint: Dict, filepath: str) -> None:
         """Save model/training states as a checkpoint file through state-dump and file-write.
 
         Args:
+            checkpoint: The checkpoint state dictionary
             filepath: write-target file's path
-            weights_only: saving model weights only
         """
         if self.world_size > 1 and self.zero_stage_3:
+            if self.save_full_weights:
+                # todo: expose this as general function in deepspeed
+                state_dict = self.deepspeed_engine._zero3_consolidated_fp16_state_dict()
+                if self.is_global_zero:
+                    # State dict keys will include reference to wrapper LightningDeepSpeedModule
+                    # Delete `module` prefix before saving.
+                    state_dict = {k.partition('module.')[2]: state_dict[k] for k in state_dict.keys()}
+                    checkpoint['state_dict'] = state_dict
+                    return super().save_checkpoint(checkpoint, filepath)
+                return
+
             # Use deepspeed's internal checkpointing function to handle partitioned weights across processes
             # dump states as a checkpoint dictionary object
             save_dir = self._filepath_to_dir(filepath)
             _exclude_keys = ['state_dict', 'optimizer_states', 'lr_schedulers']
             checkpoint = {k: v for k, v in checkpoint.items() if k not in _exclude_keys}
             self.deepspeed_engine.save_checkpoint(save_dir, client_state=checkpoint)
-
         else:
             super().save_checkpoint(checkpoint, filepath)
 
@@ -491,7 +509,8 @@ def restore_model_state_from_ckpt_path(
         ckpt_path: str,
         map_location: Callable = lambda storage, loc: storage,
     ) -> Tuple[Dict, bool]:
-        if self.world_size > 1:
+        if not self.save_full_weights and self.world_size > 1:
+            # Rely on deepspeed to load the checkpoint and necessary information
             from pytorch_lightning.trainer.states import TrainerState
             stage_is_fit = self.lightning_module.trainer.state == TrainerState.FITTING
             save_dir = self._filepath_to_dir(ckpt_path)
@@ -511,6 +530,10 @@ def restore_model_state_from_ckpt_path(
             # hook: give user access to checkpoint if needed.
             self.lightning_module.on_load_checkpoint(client_state)
             return client_state, False
+
+        # Broadcast to ensure we load from the rank 0 checkpoint
+        # This doesn't have to be the case when using deepspeed sharded checkpointing
+        ckpt_path = self.broadcast(ckpt_path)
         return super().restore_model_state_from_ckpt_path(ckpt_path, map_location=map_location)
 
     def update_global_step(self, total_batch_idx: int, current_global_step: int) -> int:

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -15,7 +15,7 @@
 import os
 import re
 import time
-from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING
+from typing import Any, Dict, List, Optional, TYPE_CHECKING, Union
 
 import torch
 import torch.multiprocessing as mp
@@ -41,7 +41,6 @@
 if _OMEGACONF_AVAILABLE:
     from omegaconf import DictConfig, ListConfig, OmegaConf
 
-
 if TYPE_CHECKING:
     from torch.nn import Module
     from torch.utils.data import DataLoader
@@ -278,4 +277,6 @@ def all_gather(self, tensor: torch.Tensor, group: Optional[Any] = None, sync_gra
         Return:
             A tensor of shape (world_size, batch, ...)
         """
-        return xm.all_gather(tensor.unsqueeze(0))
+        if isinstance(tensor, torch.Tensor) and tensor.dim() == 0:
+            tensor = tensor.unsqueeze(0)
+        return xm.all_gather(tensor)

pytorch_lightning/utilities/__init__.py

@@ -53,12 +53,10 @@
     _TORCH_QUANTIZE_AVAILABLE,
     _TORCHTEXT_AVAILABLE,
     _TORCHVISION_AVAILABLE,
+    _TPU_AVAILABLE,
     _XLA_AVAILABLE,
 )
 from pytorch_lightning.utilities.parsing import AttributeDict, flatten_dict, is_picklable  # noqa: F401
-from pytorch_lightning.utilities.xla_device import XLADeviceUtils  # noqa: F401
-
-_TPU_AVAILABLE = XLADeviceUtils.tpu_device_exists()
 
 FLOAT16_EPSILON = numpy.finfo(numpy.float16).eps
 FLOAT32_EPSILON = numpy.finfo(numpy.float32).eps