Merge branch 'master' into fix/tpu_sync_dist

Author: kaushikb11

.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

@@ -243,6 +243,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - 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/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/utilities/imports.py

@@ -69,6 +69,7 @@ def _compare_version(package: str, op, version) -> bool:
 _TORCH_GREATER_EQUAL_1_6 = _compare_version("torch", operator.ge, "1.6.0")
 _TORCH_GREATER_EQUAL_1_7 = _compare_version("torch", operator.ge, "1.7.0")
 _TORCH_GREATER_EQUAL_1_8 = _compare_version("torch", operator.ge, "1.8.0")
+_TORCH_GREATER_EQUAL_1_8_1 = _compare_version("torch", operator.ge, "1.8.1")
 _TORCH_GREATER_EQUAL_1_9 = _compare_version("torch", operator.ge, "1.9.0")
 
 _APEX_AVAILABLE = _module_available("apex.amp")
@@ -80,7 +81,7 @@ def _compare_version(package: str, op, version) -> bool:
 _HOROVOD_AVAILABLE = _module_available("horovod.torch")
 _HYDRA_AVAILABLE = _module_available("hydra")
 _HYDRA_EXPERIMENTAL_AVAILABLE = _module_available("hydra.experimental")
-_KINETO_AVAILABLE = torch.profiler.kineto_available() if _TORCH_GREATER_EQUAL_1_8 else False
+_KINETO_AVAILABLE = _TORCH_GREATER_EQUAL_1_8_1 and torch.profiler.kineto_available()
 _NATIVE_AMP_AVAILABLE = _module_available("torch.cuda.amp") and hasattr(torch.cuda.amp, "autocast")
 _OMEGACONF_AVAILABLE = _module_available("omegaconf")
 _RPC_AVAILABLE = not _IS_WINDOWS and _module_available('torch.distributed.rpc')

requirements/extra.txt

@@ -10,4 +10,3 @@ hydra-core>=1.0
 # todo: when switch to standard package stream, drop `fairscale` from hard mocked docs libs
 https://github.com/PyTorchLightning/fairscale/archive/pl_1.2.0.zip
 jsonargparse[signatures]>=3.3.1
-deepspeed>=0.3.13

tests/models/test_horovod.py

@@ -376,6 +376,8 @@ def _compute_batch():
     horovod.run(_compute_batch, np=2)
 
 
+# todo: need to be fixed :]
+@pytest.mark.skip(reason="TODO Breaking CI: Aborted (core dumped)")
 @RunIf(skip_windows=True, horovod=True)
 def test_horovod_multi_optimizer_with_scheduling_stepping(tmpdir):
 

tests/overrides/test_data_parallel.py

@@ -2,8 +2,11 @@
 
 import pytest
 import torch
+import torch.nn as nn
 from torch.nn import DataParallel
 
+from pytorch_lightning import LightningModule
+from pytorch_lightning.core.decorators import auto_move_data
 from pytorch_lightning.overrides import LightningDistributedModule
 from pytorch_lightning.overrides.data_parallel import (
     LightningParallelModule,
@@ -123,3 +126,68 @@ def training_step(self, batch, batch_idx):
     wrapped_model = LightningParallelModule(model)
     output = wrapped_model(batch, batch_idx)
     assert output["python scalar"] == torch.tensor([12.3], device=device)
+
+
+@RunIf(min_gpus=2)
+@pytest.mark.parametrize(
+    "nest, unnest", [
+        (lambda x: x, lambda x: x),
+        (lambda x: dict(data=x), lambda x: x["data"]),
+        (lambda x: [x, (x, x)], lambda x: x[1][0]),
+    ]
+)
+def test_lightning_parallel_module_device_access(nest, unnest):
+    """ Test that self.device returns the correct value in replicas of DataParallel. """
+
+    class DeviceAccessModel(LightningModule):
+
+        def __init__(self):
+            super().__init__()
+            self.layer = nn.Linear(2, 3)
+
+        @auto_move_data
+        def training_step(self, batch, batch_idx):
+            batch = unnest(batch)
+            assert batch.shape == torch.Size([1, 1])
+            assert self.device.index == batch.item()
+            assert self.device == self.layer.weight.device
+            return torch.tensor(1, device=self.device)
+
+    pl_module = DeviceAccessModel()
+    # required for redirecting the forward call to training_step
+    pl_module.trainer = Mock()
+    pl_module.trainer._running_stage = RunningStage.TRAINING
+
+    root_device = torch.device("cuda", 0)
+    wrapped_module = LightningParallelModule(pl_module).to(root_device)
+    model = DataParallel(wrapped_module, device_ids=[0, 1])
+
+    data = torch.tensor([0.0, 1.0], device=root_device).view(2, 1)  # one value per gpu
+    data = data.to(root_device)
+    data = nest(data)
+    output = model(data, 0)
+    assert output.device == root_device
+    assert pl_module.device == root_device
+    assert torch.all(output.cpu().eq(torch.tensor([1, 1])))
+
+
+@RunIf(min_gpus=2)
+def test_lightning_parallel_module_device_access_warning():
+    """ Test that we show a warning when the device can't be inferred from the input. """
+
+    class DeviceAccessModel(LightningModule):
+
+        def training_step(self, batch, batch_idx):
+            pass
+
+    pl_module = DeviceAccessModel()
+    # required for redirecting the forward call to training_step
+    pl_module.trainer = Mock()
+    pl_module.trainer._running_stage = RunningStage.TRAINING
+
+    wrapped_module = LightningParallelModule(pl_module).cuda()
+    model = DataParallel(wrapped_module, device_ids=[0, 1])
+
+    data = dict(x=1)  # contains no tensors
+    with pytest.warns(UserWarning, match="Could not determine on which device the inputs are."):
+        _ = model(data, 0)