Remove composer.trainer.ddp; replace with composer.utils.ddp (mosaicml#105)

Before mosaicml#65, composer.trainer.ddp ensured that DDP functionality was accessed only after ddp was initialized. Now, DDP is available from process start, so this class is no longer needed. Moved all the functionality from this class to the global composer.utils.ddp.

This change allows callbacks, algorithms, etc... to use DDP (such as barriers and reductions) as needed. mosaicml#97 and mosaicml#101 depend on this functionality.

Also removed DDP from the state, as that is available globally.

Author: ravi-mosaicml

composer/algorithms/curriculum_learning/curriculum_learning.py

@@ -10,7 +10,7 @@
 from composer.algorithms import AlgorithmHparams
 from composer.core.types import Algorithm, Batch, Event, Logger, State, Tensor
 from composer.models.transformer_shared import MosaicTransformer
-from composer.utils import ensure_tuple
+from composer.utils import ddp, ensure_tuple
 
 
 def apply_curriculum(batch: Dict[str, Tensor], curr_seq_len: int, truncate: bool) -> Batch:
@@ -153,8 +153,8 @@ def apply(self, event: Event, state: State, logger: Logger) -> Optional[int]:
             # all of the parameters
             device = next(state.model.parameters()).device
 
-            assert (state.train_batch_size % state.world_size) == 0
-            per_gpu_batch = math.ceil(state.train_batch_size / (state.world_size * state.grad_accum))
+            assert (state.train_batch_size % ddp.get_world_size()) == 0
+            per_gpu_batch = math.ceil(state.train_batch_size / (ddp.get_world_size() * state.grad_accum))
             input_ids = torch.randint(low=0,
                                       high=vocab_size - 1,
                                       size=(per_gpu_batch, self.hparams.max_seq_length),

composer/algorithms/seq_length_warmup/seq_length_warmup.py

@@ -10,7 +10,7 @@
 from composer.algorithms import AlgorithmHparams
 from composer.core.types import Algorithm, Batch, Event, Logger, State, Tensor
 from composer.models.transformer_shared import MosaicTransformer
-from composer.utils import ensure_tuple
+from composer.utils import ddp, ensure_tuple
 
 
 def apply_seq_length_warmup(batch: Dict[str, Tensor], curr_seq_len: int, truncate: bool) -> Batch:
@@ -180,8 +180,8 @@ def apply(self, event: Event, state: State, logger: Logger) -> Optional[int]:
             # all of the parameters
             device = next(state.model.parameters()).device
 
-            assert (state.train_batch_size % state.world_size) == 0
-            per_gpu_batch = math.ceil(state.train_batch_size / (state.world_size * state.grad_accum))
+            assert (state.train_batch_size % ddp.get_world_size()) == 0
+            per_gpu_batch = math.ceil(state.train_batch_size / (ddp.get_world_size() * state.grad_accum))
             input_ids = torch.randint(low=0,
                                       high=vocab_size - 1,
                                       size=(per_gpu_batch, self.hparams.max_seq_length),

composer/callbacks/benchmarker.py

@@ -11,6 +11,7 @@
 from composer.callbacks.callback_hparams import BenchmarkerHparams
 from composer.core.callback import Callback
 from composer.core.types import BreakEpochException
+from composer.utils import ddp
 
 log = logging.getLogger(__name__)
 
@@ -158,7 +159,7 @@ def batch_end(self, state: State, logger: Logger):
             now = time.time()
             elapsed = now - self.current_time
             self.current_time = now
-            self.profile_examples += state.last_batch_size * state.world_size
+            self.profile_examples += state.last_batch_size * ddp.get_world_size()
             self.profile_steps += 1
             self.profile_time += elapsed
 

composer/callbacks/speed_monitor.py

@@ -10,6 +10,7 @@
 from composer.callbacks.callback_hparams import SpeedMonitorHparams
 from composer.core.callback import RankZeroCallback
 from composer.core.types import StateDict
+from composer.utils import ddp
 
 
 class SpeedMonitor(RankZeroCallback):
@@ -83,7 +84,7 @@ def batch_end(self, state: State, logger: Logger):
         # Ideally, callbacks would have a way of reducing tensors.
         # It assumes that each process has equal batch sizing
         # For the speed monitor, we might be able to use the static step converter with num_samples
-        batch_num_samples *= state.world_size
+        batch_num_samples *= ddp.get_world_size()
         self.batch_num_samples.append(batch_num_samples)
         self.train_examples_per_epoch += batch_num_samples
         if len(self.batch_end_times) == self.hparams.window_size + 1:

composer/core/callback.py

@@ -8,7 +8,12 @@
 from typing import TYPE_CHECKING
 
 from composer.core.serializable import Serializable
-from composer.utils.ddp import is_rank_zero
+from composer.utils import ddp
+
+try:
+    from typing import final
+except ImportError:
+    final = lambda x: x  # final is not available in python 3.7
 
 try:
     from typing import final
@@ -299,7 +304,7 @@ def eval_end(self, state: State, logger: Logger) -> None:
 
 
 class RankZeroCallback(Callback, abc.ABC):
-    """Base class for callbacks that only run on the rank zero process.
+    """Base class for callbacks that only run on the local rank zero process.
 
     Callbacks can be implemented in two ways:
 
@@ -314,6 +319,6 @@ class RankZeroCallback(Callback, abc.ABC):
 
     @final
     def run_event(self, event: Event, state: State, logger: Logger) -> None:
-        if not is_rank_zero():
+        if ddp.get_local_rank() != 0:
             return
         return self._run_event(event, state, logger)

composer/core/logging/base_backend.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 from composer.core.callback import Callback, RankZeroCallback
-from composer.utils.ddp import is_rank_zero
+from composer.utils import ddp
 
 if TYPE_CHECKING:
     from composer.core.logging.logger import LogLevel, TLogData
@@ -104,7 +104,7 @@ def _will_log(self, state: State, log_level: LogLevel) -> bool:
 
     @final
     def will_log(self, state: State, log_level: LogLevel) -> bool:
-        if not state.is_rank_zero:
+        if ddp.get_local_rank() != 0:
             return False
         return self._will_log(state, log_level)
 
@@ -126,6 +126,6 @@ def _log_metric(self, epoch: int, step: int, log_level: LogLevel, data: TLogData
 
     @final
     def log_metric(self, epoch: int, step: int, log_level: LogLevel, data: TLogData) -> None:
-        if not is_rank_zero():
+        if ddp.get_local_rank() != 0:
             return
         return self._log_metric(epoch, step, log_level, data)

composer/core/state.py

@@ -16,7 +16,6 @@
 from composer.core.precision import Precision
 from composer.core.serializable import Serializable
 from composer.utils import ensure_tuple
-from composer.utils.ddp import get_global_rank, get_local_rank, get_local_world_size, get_world_size
 from composer.utils.precision import default_precision_factory
 
 if TYPE_CHECKING:
@@ -142,26 +141,6 @@ class State(Serializable):
     algorithms: Sequence[Algorithm] = tuple()
     callbacks: Sequence[Callback] = tuple()
 
-    @property
-    def world_size(self) -> int:
-        return get_world_size()
-
-    @property
-    def global_rank(self) -> int:
-        return get_global_rank()
-
-    @property
-    def local_world_size(self) -> int:
-        return get_local_world_size()
-
-    @property
-    def local_rank(self) -> int:
-        return get_local_rank()
-
-    @property
-    def is_rank_zero(self) -> bool:
-        return self.global_rank == 0
-
     def state_dict(self) -> types.StateDict:
         """Returns the state as a :class:`dict`."""
         state_dict: types.StateDict = {}

composer/datasets/__init__.py

@@ -3,6 +3,7 @@
 from composer.datasets.brats import BratsDatasetHparams as BratsDatasetHparams
 from composer.datasets.cifar10 import CIFAR10DatasetHparams as CIFAR10DatasetHparams
 from composer.datasets.dataloader import DataloaderHparams as DataloaderHparams
+from composer.datasets.dataloader import DDPDataLoader as DDPDataLoader
 from composer.datasets.dataloader import WrappedDataLoader as WrappedDataLoader
 from composer.datasets.hparams import DataloaderSpec as DataloaderSpec
 from composer.datasets.hparams import DatasetHparams as DatasetHparams

composer/datasets/dataloader.py

@@ -2,13 +2,15 @@
 
 from __future__ import annotations
 
+import warnings
 from dataclasses import dataclass
-from typing import Any, Iterator
+from typing import Any, Iterator, Optional
 
 import torch
 import torch.distributed
 import torch.utils.data
 import yahp as hp
+from torch.utils.data.distributed import DistributedSampler
 from torch.utils.data.sampler import Sampler
 
 from composer.core.types import Batch, DataLoader
@@ -44,6 +46,44 @@ def __setattr__(self, name: str, value: Any) -> None:
         return super().__setattr__(name, value)
 
 
+class DDPDataLoader(WrappedDataLoader):
+    """Ensure sampler.set_epoch() is called after each iteration.
+
+    DDPDataLoader wraps a dataloader and a distributed sampler and is
+    called after each iteration (epoch) through the dataset.
+    See: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
+    """
+
+    def __init__(self, dataloader: DataLoader) -> None:
+        super().__init__(dataloader)
+        if not isinstance(self.dataloader.sampler, DistributedSampler):
+            raise ValueError("When using the DDP data loader, the sampler must be a DistributedSampler")
+        self._iterator: Optional[Iterator[Batch]] = None
+
+    def __iter__(self) -> DDPDataLoader:
+        if self._iterator is not None:
+            warnings.warn(
+                "DataloaderMultipleIterationWarning: "
+                "The dataloader detected the start of a new iteration before the previous iteration finished. "
+                "The dataloader is skipping ahead to the start of the next epoch. "
+                "Multiple simultaneous iterations through the DDP dataloader prohibited, since "
+                "it automatically tracks the current epoch.")
+            assert isinstance(self.sampler, DistributedSampler)
+            self.sampler.set_epoch(epoch=self.sampler.epoch + 1)
+        self._iterator = iter(self.dataloader)
+        return self
+
+    def __next__(self) -> Batch:
+        assert self._iterator is not None
+        try:
+            return next(self._iterator)
+        except StopIteration:
+            self._iterator = None
+            assert isinstance(self.sampler, DistributedSampler)
+            self.sampler.set_epoch(epoch=self.sampler.epoch + 1)
+            raise
+
+
 @dataclass
 class DataloaderHparams(hp.Hparams):
     """Hyperparameters to initialize a ``torch.utils.data.Dataloader``."""

composer/trainer/__init__.py

@@ -1,7 +1,6 @@
 # Copyright 2021 MosaicML. All Rights Reserved.
 
 from composer.trainer import devices as devices
-from composer.trainer.ddp import DDPDataLoader as DDPDataLoader
 from composer.trainer.trainer import Trainer as Trainer
 from composer.trainer.trainer_hparams import TrainerHparams as TrainerHparams