Distinguish between dist and DDP (#201)

* dist, not ddp

* simplify ClosureGradScaler

* formatting

* formatting and more fixes

* that did not save

* small fixes

* dont need to worry about circular dependencies any longer

* dumb pyright fix

* woops

Author: jbloxham

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 ddp, ensure_tuple
+from composer.utils import dist, 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 % ddp.get_world_size()) == 0
-            per_gpu_batch = math.ceil(state.train_batch_size / (ddp.get_world_size() * state.grad_accum))
+            assert (state.train_batch_size % dist.get_world_size()) == 0
+            per_gpu_batch = math.ceil(state.train_batch_size / (dist.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,7 +11,7 @@
 from composer.core import Logger, State
 from composer.core.callback import Callback
 from composer.core.types import BreakEpochException
-from composer.utils import ddp
+from composer.utils import dist
 
 log = logging.getLogger(__name__)
 
@@ -159,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 * ddp.get_world_size()
+            self.profile_examples += state.last_batch_size * dist.get_world_size()
             self.profile_steps += 1
             self.profile_time += elapsed
 

composer/callbacks/run_directory_uploader.py

@@ -20,7 +20,7 @@
 from composer.core.logging import Logger
 from composer.core.logging.logger import LogLevel
 from composer.core.state import State
-from composer.utils import ddp
+from composer.utils import dist
 from composer.utils.run_directory import get_modified_files, get_run_directory
 
 log = logging.getLogger(__name__)
@@ -148,13 +148,13 @@ def __init__(
         self._finished: Union[None, multiprocessing._EventType, threading.Event] = None
         self._workers = []
 
-        if ddp.get_local_rank() == 0:
+        if dist.get_local_rank() == 0:
             _validate_credentials(provider, container, self._object_name_prefix, provider_init_kwargs)
 
     def init(self, state: State, logger: Logger) -> None:
         if get_run_directory() is None:
             return
-        if not ddp.get_local_rank() == 0:
+        if not dist.get_local_rank() == 0:
             return
         del state, logger  # unused
         self._finished = self._finished_cls()
@@ -194,7 +194,7 @@ def training_end(self, state: State, logger: Logger) -> None:
     def post_close(self):
         # Cleaning up on post_close to ensure that all artifacts are uploaded
         self._trigger_upload(logger=None, log_level=None)
-        if not ddp.get_local_rank() == 0:
+        if not dist.get_local_rank() == 0:
             return
         if self._finished is not None:
             self._finished.set()
@@ -207,8 +207,8 @@ def _trigger_upload(self, logger: Optional[Logger], log_level: Optional[LogLevel
         # Ensure that every rank is at this point
         # Assuming only the main thread on each rank writes to the run directory, then the barrier here will ensure
         # that the run directory is not being modified after we pass this barrier
-        ddp.barrier()
-        if ddp.get_local_rank() == 0:
+        dist.barrier()
+        if dist.get_local_rank() == 0:
             run_directory = get_run_directory()
             assert run_directory is not None, "invariant error"
             # the disk time can differ from system time, so going to touch a file and then read the timestamp from it to get the real time
@@ -244,7 +244,7 @@ def _trigger_upload(self, logger: Optional[Logger], log_level: Optional[LogLevel
                 logger.metric(log_level, {"run_directory/uploaded_files": files_to_be_uploaded})
 
         # Ensure that other callbacks do not start writing to the run directory until we copying everything
-        ddp.barrier()
+        dist.barrier()
 
 
 def _validate_credentials(

composer/callbacks/speed_monitor.py

@@ -10,7 +10,7 @@
 from composer.core import Logger, State
 from composer.core.callback import RankZeroCallback
 from composer.core.types import StateDict
-from composer.utils import ddp
+from composer.utils import dist
 
 
 class SpeedMonitor(RankZeroCallback):
@@ -84,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 *= ddp.get_world_size()
+        batch_num_samples *= dist.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/callbacks/torch_profiler.py

@@ -12,7 +12,7 @@
 from composer.callbacks.callback_hparams import TorchProfilerHparams
 from composer.core import Callback, Logger, State
 from composer.core.types import StateDict
-from composer.utils.ddp import get_global_rank
+from composer.utils.dist import get_global_rank
 from composer.utils.run_directory import get_relative_to_run_directory
 
 _PROFILE_MISSING_ERROR = "The profiler has not been setup. Please call profiler.training_start() before training starts."

composer/core/callback.py

@@ -8,7 +8,7 @@
 from typing import TYPE_CHECKING
 
 from composer.core.serializable import Serializable
-from composer.utils import ddp
+from composer.utils import dist
 
 try:
     from typing import final
@@ -333,6 +333,6 @@ class RankZeroCallback(Callback, abc.ABC):
 
     @final
     def run_event(self, event: Event, state: State, logger: Logger) -> None:
-        if ddp.get_local_rank() != 0:
+        if dist.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 import ddp
+from composer.utils import dist
 
 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 ddp.get_local_rank() != 0:
+        if dist.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 ddp.get_local_rank() != 0:
+        if dist.get_local_rank() != 0:
             return
         return self._log_metric(epoch, step, log_level, data)

composer/core/state.py

@@ -13,7 +13,7 @@
 import composer.core.types as types
 from composer.core.precision import Precision
 from composer.core.serializable import Serializable
-from composer.utils import ddp, ensure_tuple
+from composer.utils import dist, ensure_tuple
 from composer.utils.precision import default_precision_factory
 
 if TYPE_CHECKING:
@@ -185,14 +185,14 @@ def train_batch_size(self):
         """The global batch size used for training."""
         if self.train_dataloader.batch_size is None:
             raise RuntimeError("train dataloader batch size is undefined")
-        return self.train_dataloader.batch_size * ddp.get_world_size()
+        return self.train_dataloader.batch_size * dist.get_world_size()
 
     @property
     def eval_batch_size(self):
         """The batch size used for evaluation."""
         if self.eval_dataloader.batch_size is None:
             raise RuntimeError("eval dataloader batch size is undefined")
-        return self.eval_dataloader.batch_size * ddp.get_world_size()
+        return self.eval_dataloader.batch_size * dist.get_world_size()
 
     def state_dict(self) -> types.StateDict:
         """Returns the state as a :class:`dict`."""

composer/datasets/brats.py

@@ -14,7 +14,7 @@
 from composer.core.types import DataLoader, Dataset
 from composer.datasets.dataloader import DataloaderHparams
 from composer.datasets.hparams import DatasetHparams
-from composer.utils import ddp
+from composer.utils import dist
 
 PATCH_SIZE = [1, 192, 160]
 
@@ -48,7 +48,7 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
         x_train, y_train, x_val, y_val = get_data_split(self.datadir)
         dataset = PytTrain(x_train, y_train, oversampling) if self.is_train else PytVal(x_val, y_val)
         collate_fn = None if self.is_train else _my_collate
-        sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
+        sampler = dist.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
 
         return dataloader_hparams.initialize_object(
             dataset=dataset,

composer/datasets/cifar10.py

@@ -10,7 +10,7 @@
 from composer.datasets.dataloader import DataloaderHparams
 from composer.datasets.hparams import DatasetHparams, SyntheticHparamsMixin
 from composer.datasets.synthetic import SyntheticBatchPairDataset
-from composer.utils import ddp
+from composer.utils import dist
 
 
 @dataclass
@@ -59,7 +59,7 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
                 download=self.download,
                 transform=transformation,
             )
-        sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
+        sampler = dist.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
 
         return dataloader_hparams.initialize_object(dataset,
                                                     batch_size=batch_size,