ZeRO-Offload v1 (squash) (#345)

* update DSE to point to ZeRO-Offload staging

* ZeRO-2 enable CPU offload (#313)

* cpu-offload

* update

* deleted:    deepspeed/pt/deepspeed_zero_optimizer_cpuoffload.py
	modified:   deepspeed/pt/fp16_unfused_optimizer.py
	new file:   install_output.txt
	modified:   tests/unit/test_dynamic_loss_scale.py

* modified:   deepspeed/pt/deepspeed_zero_optimizer.py

* update

* modified:   deepspeed/pt/deepspeed_cpu_adam.py
	modified:   deepspeed/pt/deepspeed_zero_optimizer.py
	modified:   tests/unit/test_checkpointing.py
	modified:   tests/unit/test_fp16.py

* deleted:    install_output.txt

* modified:   deepspeed/pt/fp16_unfused_optimizer.py
	modified:   tests/unit/test_dynamic_loss_scale.py

* modified:   deepspeed/pt/deepspeed_cpu_adam.py

* modified:   deepspeed/pt/deepspeed_zero_optimizer.py

* modified:   deepspeed/pt/deepspeed_cpu_adam.py
	modified:   deepspeed/pt/deepspeed_zero_optimizer.py

* deleted:    deepspeed_cpu_adam.py
	modified:   deepspeed_light.py
	modified:   deepspeed_zero_optimizer.py
	../../deepspeed_zero_optimizer_cpu_offload.py

* modified:   deepspeed/pt/deepspeed_light.py

* modified:   deepspeed/pt/deepspeed_light.py
	modified:   deepspeed/pt/deepspeed_zero_optimizer.py
	modified:   deepspeed/pt/deepspeed_zero_utils.py
	modified:   tests/unit/test_fp16.py

* modified:   deepspeed/pt/deepspeed_config.py
	modified:   deepspeed/pt/deepspeed_light.py
	modified:   deepspeed/pt/deepspeed_zero_optimizer.py
	modified:   tests/unit/test_checkpointing.py
	modified:   tests/unit/test_fp16.py

* modified:   deepspeed/pt/deepspeed_checkpointing.py

* update DSE to ZeRO-Offload commit

Co-authored-by: Jeff Rasley <[email protected]>

* Enable ZeRO checkpointing for ZeRO-Offload (#337)

* Enable ZeRO checkpointing for ZeRO-Offload
Fix unit tests
Bump DSE to 33b9fb77c8cecdb49118188890f662526d8e9397

* Fix accidental revert

* Add ZeRO-Offload checkpointing model tests (#344)

* Enable ZeRO checkpointing for ZeRO-Offload
Fix unit tests
Bump DSE to 33b9fb77c8cecdb49118188890f662526d8e9397

* Fix accidental revert

* Fix ZeRO-Offload checkpointing bug when change gpu count
Add checkpointing model tests for ZeRO-Offload
Remove optimizer key from Megatron model tests
Use different deepspeed master port for Megatron model tests

Co-authored-by: Jie <[email protected]>
Co-authored-by: Olatunji Ruwase <[email protected]>

Author: 3 people

DeepSpeedExamples

@@ -0,0 +1,33 @@
+import torch
+from torch.autograd import Variable
+import collections
+
+
+def async_migrate_to(obj, dev, main_stream=None):
+    if torch.is_tensor(obj):
+        obj = Variable(obj)
+    if isinstance(obj, Variable):
+        v = obj.cuda(dev, async=True)
+        if main_stream is not None:
+            v.data.record_stream(main_stream)
+        return v
+    elif isinstance(obj, collections.Mapping):
+        return {k: async_copy_to(o, dev, main_stream) for k, o in obj.items()}
+    elif isinstance(obj, collections.Sequence):
+        return [async_copy_to(o, dev, main_stream) for o in obj]
+    else:
+        return obj
+
+
+def async_copy_to(obj, dev, main_stream=None):
+    if torch.is_tensor(obj):
+        obj = Variable(obj)
+    if isinstance(obj, Variable):
+        target = torch.empty_like(obj, device=dev).copy_(obj)
+        if main_stream is not None:
+            target.data.record_stream(main_stream)
+        return target
+    elif isinstance(obj, collections.Mapping):
+        return {k: async_copy_to(o, dev, main_stream) for k, o in obj.items()}
+    elif isinstance(obj, collections.Sequence):
+        return [async_copy_to(o, dev, main_stream) for o in obj]

deepspeed/pt/deepspeed_zero_utils.py

@@ -630,10 +630,12 @@ def _do_error_check(self):
         if self.zero_enabled:
             assert self.fp16_enabled, "DeepSpeedConfig: ZeRO is only supported if fp16 is enabled"
             assert self.zero_optimization_stage <= MAX_STAGE_ZERO_OPTIMIZATION, "DeepSpeedConfig: Maximum supported ZeRO stage is {}".format(MAX_STAGE_ZERO_OPTIMIZATION)
+            if self.zero_config.cpu_offload is True:
+                assert self.zero_optimization_stage == ZERO_OPTIMIZATION_GRADIENTS, "DeepSpeedConfig: cpu-offload supported ZeRO stage is {}".format(ZERO_OPTIMIZATION_GRADIENTS)
 
         assert self.train_micro_batch_size_per_gpu, "DeepSpeedConfig: {} is not defined".format(TRAIN_MICRO_BATCH_SIZE_PER_GPU)
 
-        assert self.gradient_accumulation_steps, 'DeepSpeedConfig: {} is not defined'.format(
+        assert self.gradient_accumulation_steps, "DeepSpeedConfig: {} is not defined".format(
             GRADIENT_ACCUMULATION_STEPS)
 
     def _do_warning_check(self):

deepspeed/runtime/config.py

@@ -106,7 +106,6 @@ def __init__(self,
                  collate_fn=None,
                  config_params=None):
         super(DeepSpeedEngine, self).__init__()
-
         self.client_optimizer = optimizer
         self.client_model_parameters = model_parameters
         self.client_lr_scheduler = lr_scheduler
@@ -292,6 +291,9 @@ def zero_reduce_scatter(self):
     def zero_overlap_comm(self):
         return self._config.zero_config.overlap_comm
 
+    def zero_cpu_offload(self):
+        return self._config.zero_config.cpu_offload
+
     def zero_optimization_stage(self):
         return self._config.zero_optimization_stage
 
@@ -491,6 +493,7 @@ def _configure_distributed_model(self, model):
 
     # Configure optimizer
     def _configure_optimizer(self, client_optimizer, model_parameters):
+
         if client_optimizer is not None:
             basic_optimizer = client_optimizer
             logger.info('Using client Optimizer as basic optimizer')
@@ -504,13 +507,14 @@ def _configure_optimizer(self, client_optimizer, model_parameters):
 
         if self.zero_optimization():
             assert not self.amp_enabled(), "Amp and ZeRO are not currently compatible, please use (legacy) fp16 mode which performs similar to amp opt_mode=O2"
-            if self.optimizer_name() != ADAM_OPTIMIZER:
+            if self.optimizer_name() not in [ADAM_OPTIMIZER]:
                 assert self.zero_allow_untested_optimizer(), \
                     'You are using an untested ZeRO Optimizer. Please add <"zero_allow_untested_optimizer": true> in the configuration file to use it.'
 
                 logger.warning(
                     "**** You are using ZeRO with an untested optimizer, proceed with caution *****"
                 )
+
             self.optimizer = self._configure_zero_optimizer(basic_optimizer)
         elif self.amp_enabled():
             assert not self.fp16_enabled(), "Cannot enable both amp with (legacy) fp16 mode"
@@ -522,8 +526,8 @@ def _configure_optimizer(self, client_optimizer, model_parameters):
             self.optimizer = self._configure_fp16_optimizer(basic_optimizer)
         else:
             self.optimizer = basic_optimizer
-
-        # logger.info('DeepSpeed Final Optimizer = {}'.format(self.optimizer.state_dict()))
+        logger.info('DeepSpeed Final Optimizer = {}'.format(self.optimizer))
+        logger.info('DeepSpeed Final Optimizer = {}'.format(self.optimizer.state_dict()))
 
     def _configure_basic_optimizer(self, model_parameters):
         optimizer_parameters = self.optimizer_params()
@@ -532,8 +536,11 @@ def _configure_basic_optimizer(self, model_parameters):
                 "'max_grad_norm' is not supported as an optimizer parameter, please switch to using the deepspeed parameter 'gradient_clipping' see: https://www.deepspeed.ai/docs/config-json/#gradient-clipping for more details"
             )
         if self.optimizer_name() == ADAM_OPTIMIZER:
-            from apex.optimizers.fused_adam import FusedAdam
-            optimizer = FusedAdam(model_parameters, **optimizer_parameters)
+            if self.zero_cpu_offload():
+                optimizer = torch.optim.Adam(model_parameters, **optimizer_parameters)
+            else:
+                from apex.optimizers.fused_adam import FusedAdam
+                optimizer = FusedAdam(model_parameters, **optimizer_parameters)
         elif self.optimizer_name() == LAMB_OPTIMIZER:
             optimizer = FusedLamb(model_parameters, **optimizer_parameters)
         else:
@@ -610,6 +617,7 @@ def _configure_zero_optimizer(self, optimizer):
                 dp_process_group=self.data_parallel_group,
                 reduce_scatter=self.zero_reduce_scatter(),
                 overlap_comm=self.zero_overlap_comm(),
+                cpu_offload=self.zero_cpu_offload(),
                 mpu=self.mpu,
                 postscale_gradients=self.postscale_gradients(),
                 gradient_predivide_factor=self.gradient_predivide_factor())
@@ -844,7 +852,6 @@ def step(self):
                     master_params = amp.master_params(self.optimizer)
                     torch.nn.utils.clip_grad_norm_(parameters=master_params,
                                                    max_norm=self.gradient_clipping())
-
             self.optimizer.step()
 
             #zero grad in basic optimizer could be unreliable and may not exhibit
@@ -947,6 +954,9 @@ def _get_optimizer_param(self, param_name):
     def get_lr(self):
         return self._get_optimizer_param('lr')
 
+    def get_type(self):
+        return self._get_optimizer_param('type')
+
     def get_mom(self):
         return self._get_optimizer_param('betas')
 

deepspeed/runtime/engine.py

@@ -23,6 +23,7 @@
     "overlap_comm": [true|false],
     "reduce_bucket_size": 500000000
     "load_from_fp32_weights": [true|false]
+    "cpu_offload": [true|false]
     }
 }
 '''
@@ -62,21 +63,22 @@
 ZERO_OPTIMIZATION_LOAD_FROM_FP32_WEIGHTS = 'load_from_fp32_weights'
 ZERO_OPTIMIZATION_LOAD_FROM_FP32_WEIGHTS_DEFAULT = True
 
+ZERO_OPTIMIZATION_CPU_OFFLOAD = 'cpu_offload'
+ZERO_OPTIMIZATION_CPU_OFFLOAD_DEFAULT = False
+
 ZERO_OPTIMIZATION_DEFAULT = {
-    ZERO_OPTIMIZATION_STAGE:
-    ZERO_OPTIMIZATION_STAGE_DEFAULT,
+    ZERO_OPTIMIZATION_STAGE: ZERO_OPTIMIZATION_STAGE_DEFAULT,
     ZERO_OPTIMIZATION_CONTIGUOUS_GRADIENTS:
     ZERO_OPTIMIZATION_CONTIGUOUS_GRADIENTS_DEFAULT,
-    ZERO_OPTIMIZATION_REDUCE_SCATTER:
-    ZERO_OPTIMIZATION_REDUCE_SCATTER_DEFAULT,
-    ZERO_OPTIMIZATION_REDUCE_BUCKET_SIZE:
-    ZERO_OPTIMIZATION_REDUCE_BUCKET_SIZE_DEFAULT,
+    ZERO_OPTIMIZATION_REDUCE_SCATTER: ZERO_OPTIMIZATION_REDUCE_SCATTER_DEFAULT,
+    ZERO_OPTIMIZATION_REDUCE_BUCKET_SIZE: ZERO_OPTIMIZATION_REDUCE_BUCKET_SIZE_DEFAULT,
     ZERO_OPTIMIZATION_ALLGATHER_PARTITIONS:
     ZERO_OPTIMIZATION_ALLGATHER_PARTITIONS_DEFAULT,
     ZERO_OPTIMIZATION_ALLGATHER_BUCKET_SIZE:
     ZERO_OPTIMIZATION_ALLGATHER_BUCKET_SIZE_DEFAULT,
     ZERO_OPTIMIZATION_LOAD_FROM_FP32_WEIGHTS:
-    ZERO_OPTIMIZATION_LOAD_FROM_FP32_WEIGHTS_DEFAULT
+    ZERO_OPTIMIZATION_LOAD_FROM_FP32_WEIGHTS_DEFAULT,
+    ZERO_OPTIMIZATION_CPU_OFFLOAD: ZERO_OPTIMIZATION_CPU_OFFLOAD_DEFAULT
 }
 
 
@@ -92,6 +94,7 @@ def __init__(self, param_dict):
         self.allgather_bucket_size = None
         self.overlap_comm = None
         self.load_from_fp32_weights = None
+        self.cpu_offload = None
 
         if ZERO_OPTIMIZATION in param_dict.keys():
             zero_config_dict = param_dict[ZERO_OPTIMIZATION]
@@ -156,7 +159,12 @@ def _initialize(self, zero_config_dict):
             zero_config_dict,
             ZERO_OPTIMIZATION_ALLGATHER_BUCKET_SIZE,
             ZERO_OPTIMIZATION_ALLGATHER_BUCKET_SIZE_DEFAULT)
+
         self.load_from_fp32_weights = get_scalar_param(
             zero_config_dict,
             ZERO_OPTIMIZATION_LOAD_FROM_FP32_WEIGHTS,
             ZERO_OPTIMIZATION_LOAD_FROM_FP32_WEIGHTS_DEFAULT)
+
+        self.cpu_offload = get_scalar_param(zero_config_dict,
+                                            ZERO_OPTIMIZATION_CPU_OFFLOAD,
+                                            ZERO_OPTIMIZATION_CPU_OFFLOAD_DEFAULT)