Stage1

adaptdl/adaptdl/torch/context.py

@@ -0,0 +1,98 @@
+# Copyright 2020 Petuum, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import collections
+import math
+import adaptdl.checkpoint
+import adaptdl.collective
+import adaptdl.env
+from adaptdl.torch._metrics import get_goodput_fn
+import adaptdl.torch.data
+from adaptdl.torch.scaling_rules import ScalingRuleBase
+
+class AdaptiveDLContext(object):
+    """
+    This class provides context tool to get AdaptDL-suggest parameters,
+    such as batch_size, accum_steps and lr_scale.
+    """
+
+    def __init__(self, batch_size):
+        self._elastic = adaptdl.torch.data.AdaptiveDataLoaderHelper(batch_size)
+        # Autoscale batch size fields.
+        self._speedup_threshold = 1.05
+        self.adapt_batch_size = None
+        self.adapt_accum_steps = None
+        self.adapt_lr_scale = None
+
+    def autoscale_batch_size(self, max_batch_size, local_bsz_bounds=None,
+                             gradient_accumulation=False):
+        self._elastic.autoscale_batch_size(max_batch_size, local_bsz_bounds,
+                                           gradient_accumulation)
+
+    def get_batch_size(self):
+        _, self.adapt_batch_size, _ = self._sync_local_bsz()
+        return self.adapt_batch_size
+
+    def get_accum_steps(self):
+        _, _, self.adapt_accum_steps = self._sync_local_bsz()
+        return self.adapt_accum_steps
+
+    def get_lr_scale(self):
+        self.adapt_lr_scale = ScalingRuleBase._get_adapt_lr_scale()
+        return float(self.adapt_lr_scale)
+
+    def _sync_local_bsz(self):
+        goodput_fn = get_goodput_fn()
+        if self._elastic.max_batch_size is None or goodput_fn is None:
+            # No autoscale batch size, just divide batch size evenly.
+            self._elastic._state.current_local_bsz = math.ceil(
+                self._elastic.batch_size / adaptdl.env.num_replicas())
+            self._elastic._state.accumulation_steps = 0
+        elif not self._elastic._state.current_local_bsz:
+            # if init, use the batch size suggested
+            _, atomic_bsz, accum_steps = goodput_fn.optimize(
+                adaptdl.env.num_nodes(), adaptdl.env.num_replicas(),
+                max_batch_size=self._elastic._max_batch_size,
+                atomic_bsz_range=self._elastic._local_bsz_bounds,
+                accumulation=self._elastic._gradient_accumulation)
+            self._elastic._state.current_local_bsz = atomic_bsz
+            self._elastic._state.accumulation_steps = accum_steps
+        else:
+            # if not first time, we check against the relative speedup
+            suggest_goodput, atomic_bsz, accum_steps = goodput_fn.optimize(
+                adaptdl.env.num_nodes(), adaptdl.env.num_replicas(),
+                max_batch_size=self._elastic._max_batch_size,
+                atomic_bsz_range=self._elastic._local_bsz_bounds,
+                accumulation=self._elastic._gradient_accumulation)
+            # get current goodput
+            current_goodput = goodput_fn(
+                adaptdl.env.num_nodes(), adaptdl.env.num_replicas(),
+                self._elastic.current_local_bsz, self._elastic.accumulation_steps)
+            # use only if speedup is significant
+            speedup = suggest_goodput / max(current_goodput, 1e-8)
+            if speedup > self._speedup_threshold:
+                self._elastic._state.current_local_bsz = atomic_bsz
+                self._elastic._state.accumulation_steps = accum_steps
+        self._elastic._state.current_local_bsz, self._elastic._state.accumulation_steps = \
+            adaptdl.collective.broadcast((self._elastic._state.current_local_bsz,
+                                          self._elastic._state.accumulation_steps))
+        return self._elastic.current_local_bsz, self._elastic._state.current_local_bsz, self._elastic._state.accumulation_steps
+
+    @property
+    def training(self):
+        return self._elastic.training
+
+    def to_tensorboard(self, writer, global_step, tag_prefix=""):
+        self._elastic.to_tensorboard(writer, global_step, tag_prefix)
+    # to_tensorboard.__doc__ = adaptdl.torch.data.AdaptiveDataLoaderHelper.to_tensorboard.__doc__

adaptdl/adaptdl/torch/data.py

@@ -32,6 +32,7 @@
     profile_step_start, profile_step_commit,
     set_batch_size, get_goodput_fn, get_progress)
 from adaptdl._signal import get_exit_flag
+from adaptdl.torch.context import AdaptiveDLContext
 
 logging.basicConfig(level=logging.INFO)
 LOG = logging.getLogger(__name__)
@@ -267,43 +268,6 @@ def autoscale_batch_size(self, max_batch_size, local_bsz_bounds=None,
         self._gradient_accumulation = gradient_accumulation
         self.train()
 
-    def _sync_local_bsz(self):
-        goodput_fn = get_goodput_fn()
-        if self.max_batch_size is None or goodput_fn is None:
-            # No autoscale batch size, just divide batch size evenly.
-            self._state.current_local_bsz = math.ceil(
-                self.batch_size / adaptdl.env.num_replicas())
-            self._state.accumulation_steps = 0
-        elif not self._state.current_local_bsz:
-            # if init, use the batch size suggested
-            _, atomic_bsz, accum_steps = goodput_fn.optimize(
-                adaptdl.env.num_nodes(), adaptdl.env.num_replicas(),
-                max_batch_size=self._max_batch_size,
-                atomic_bsz_range=self._local_bsz_bounds,
-                accumulation=self._gradient_accumulation)
-            self._state.current_local_bsz = atomic_bsz
-            self._state.accumulation_steps = accum_steps
-        else:
-            # if not first time, we check against the relative speedup
-            suggest_goodput, atomic_bsz, accum_steps = goodput_fn.optimize(
-                adaptdl.env.num_nodes(), adaptdl.env.num_replicas(),
-                max_batch_size=self._max_batch_size,
-                atomic_bsz_range=self._local_bsz_bounds,
-                accumulation=self._gradient_accumulation)
-            # get current goodput
-            current_goodput = goodput_fn(
-                adaptdl.env.num_nodes(), adaptdl.env.num_replicas(),
-                self.current_local_bsz, self.accumulation_steps)
-            # use only if speedup is significant
-            speedup = suggest_goodput / max(current_goodput, 1e-8)
-            if speedup > self._speedup_threshold:
-                self._state.current_local_bsz = atomic_bsz
-                self._state.accumulation_steps = accum_steps
-        self._state.current_local_bsz, self._state.accumulation_steps = \
-            adaptdl.collective.broadcast((self._state.current_local_bsz,
-                                          self._state.accumulation_steps))
-        return self.current_local_bsz
-
     @property
     def training(self):
         return self is AdaptiveDataLoaderHelper._training
@@ -398,53 +362,6 @@ def to_tensorboard(self, writer, global_step, tag_prefix=""):
                           self.accumulation_steps, global_step)
 
 
-class AdaptiveDataLoaderMixin(object):
-    """
-    This class provides elastic functionality to any custom DataLoader which
-    inherits it. It defines a member _elastic of type
-    :class:`AdaptiveDataLoaderHelper` which has useful methods and members to
-    implement restart-safe, elastic DataLoaders. It also exposes public methods
-    which can be used inside training loops directly from
-    :class:`AdaptiveDataLoader`.
-    """
-
-    def __init__(self, batch_size):
-        self._elastic = AdaptiveDataLoaderHelper(batch_size)
-
-    def autoscale_batch_size(self, max_batch_size, local_bsz_bounds=None,
-                             gradient_accumulation=False):
-        self._elastic.autoscale_batch_size(max_batch_size, local_bsz_bounds,
-                                           gradient_accumulation)
-
-    @property
-    def current_local_bsz(self):
-        if AdaptiveDataLoaderHelper._current is not self._elastic:
-            return None
-        return self._elastic.current_local_bsz
-
-    @property
-    def accumulation_steps(self):
-        """
-        The number of batches returned by the dataloader before a
-        step is taken.
-        """
-        return self._elastic.accumulation_steps
-
-    @property
-    def training(self):
-        return self._elastic.training
-
-    @property
-    def current_batch_size(self):
-        if AdaptiveDataLoaderHelper._current is not self._elastic:
-            return None
-        return self._elastic.current_batch_size
-
-    def to_tensorboard(self, writer, global_step, tag_prefix=""):
-        self._elastic.to_tensorboard(writer, global_step, tag_prefix)
-    to_tensorboard.__doc__ = AdaptiveDataLoaderHelper.to_tensorboard.__doc__
-
-
 def _worker_init_wrapper(worker_init_fn, num_workers):
     # Set globally-unique python and numpy seeds for each worker.
 
@@ -462,7 +379,7 @@ def wrapper(worker_id):
     return wrapper
 
 
-class AdaptiveDataLoader(DataLoader, AdaptiveDataLoaderMixin):
+class AdaptiveDataLoader(DataLoader, AdaptiveDLContext):
     """
     This class is a PyTorch DataLoader that also supports adaptive batch sizes
     and checkpoint-restart elasticity. Applications can typically use objects
@@ -501,7 +418,7 @@ def __init__(self, dataset, batch_size=1, shuffle=False, **kwargs):
         kwargs["worker_init_fn"] = _worker_init_wrapper(
             kwargs.get("worker_init_fn"), kwargs.get("num_workers"))
         super().__init__(dataset, batch_size, shuffle=False, **kwargs)
-        AdaptiveDataLoaderMixin.__init__(self, batch_size)
+        AdaptiveDLContext.__init__(self, batch_size)
 
     def __iter__(self):
         """
@@ -526,7 +443,7 @@ def __iter__(self):
             while not done:
                 self.sampler.set_epoch(
                     epoch, index=self._elastic.current_index)
-                self.batch_sampler.batch_size = self._elastic._sync_local_bsz()
+                self.batch_sampler.batch_size = self.get_batch_size()
                 for idx, batch in enumerate(super().__iter__()):
                     with self._elastic.profile(self.training and idx >= 1):
                         yield batch

adaptdl/adaptdl/torch/scaling_rules.py

@@ -19,7 +19,7 @@
 
 from types import MethodType
 
-from adaptdl.torch.data import current_dataloader
+# from adaptdl.torch.data import current_dataloader
 
 
 __all__ = ["ScalingRuleBase", "AdaScale", "LinearScale", "SqrtScale",
@@ -45,6 +45,9 @@ class ScalingRuleBase(object):
                 loss.backward()
                 adascale.step()
     """
+
+    _adaptlr = None
+
     def __init__(self):
         # instance of AdaptiveDataParallel, needs to be set before any of the
         # methods can be used
@@ -77,6 +80,7 @@ def step(self, *args, **kwargs):
         scale = self.adp.gns.accum_scale * self.adp.gns.accum_count
         initial_lr = [pg["lr"] for pg in self._optimizer.param_groups]
         scaled_lr = np.multiply(self.scale_lr(scale), initial_lr)
+        ScalingRuleBase._adaptlr = scaled_lr
         for lr, pg in zip(scaled_lr, self._optimizer.param_groups):
             pg["lr"] = lr
         self._orig_optimizer_step(*args, **kwargs)
@@ -107,6 +111,10 @@ def initialize(self, adp, optimizer, patch_optimizer=False):
         if patch_optimizer:
             self._patch_optimizer()
 
+    @staticmethod
+    def _get_adapt_lr_scale():
+        return ScalingRuleBase._adaptlr
+
 
 class AdaScale(ScalingRuleBase):
     """