Support legacy optimizer fusion as config option (#75)

* Support legacy optimizer fusion as config option

* Configure for legacy optimizer fusion

* Update configuration jsons for new apex

deepspeed/pt/deepspeed_config.py

@@ -157,6 +157,14 @@ def get_optimizer_gradient_clipping(param_dict):
         return None
 
 
+def get_optimizer_legacy_fusion(param_dict):
+    if OPTIMIZER in param_dict.keys() and \
+        LEGACY_FUSION in param_dict[OPTIMIZER].keys():
+        return param_dict[OPTIMIZER][LEGACY_FUSION]
+    else:
+        return LEGACY_FUSION_DEFAULT
+
+
 def get_scheduler_name(param_dict):
     if SCHEDULER in param_dict.keys() and \
             TYPE in param_dict[SCHEDULER].keys():
@@ -261,6 +269,7 @@ class DeepSpeedConfig(object):
             self.optimizer_name = self.optimizer_name.lower()
 
         self.optimizer_params = get_optimizer_params(param_dict)
+        self.optimizer_legacy_fusion = get_optimizer_legacy_fusion(param_dict)
 
         self.scheduler_name = get_scheduler_name(param_dict)
         self.scheduler_params = get_scheduler_params(param_dict)

deepspeed/pt/deepspeed_constants.py

@@ -24,6 +24,8 @@ OPTIMIZER = "optimizer"
 OPTIMIZER_TYPE_DEFAULT = None
 OPTIMIZER_PARAMS = "params"
 TYPE = "type"
+LEGACY_FUSION = "legacy_fusion"
+LEGACY_FUSION_DEFAULT = True
 SCHEDULER = "scheduler"
 SCHEDULER_TYPE_DEFAULT = None
 SCHEDULER_PARAMS = "params"

deepspeed/pt/deepspeed_light.py

@@ -211,6 +211,9 @@ class DeepSpeedLight(Module):
     def optimizer_params(self):
         return self._config.optimizer_params
 
+    def optimizer_legacy_fusion(self):
+        return self._config.optimizer_legacy_fusion
+
     def scheduler_name(self):
         return self._config.scheduler_name
 
@@ -419,21 +422,23 @@ class DeepSpeedLight(Module):
         if self.optimizer_name() == ADAM_OPTIMIZER:
             if self.dynamic_loss_scale():
                 logging.info('Creating fp16 optimizer with dynamic loss scale')
-                optimizer = FP16_Optimizer(optimizer,
+                optimizer = FP16_Optimizer(
+                    optimizer,
                     dynamic_loss_scale=True,
                     initial_dynamic_scale=initial_dynamic_scale,
                     dynamic_loss_args=dynamic_loss_args,
                     mpu=self.mpu,
                     clip_grad=clip_grad,
-                                           fused_adam_legacy=True)
+                    fused_adam_legacy=self.optimizer_legacy_fusion())
             else:
                 logging.info('Creating fp16 optimizer with static loss scale: {}'.format(
                     self.loss_scale()))
-                optimizer = FP16_Optimizer(optimizer,
+                optimizer = FP16_Optimizer(
+                    optimizer,
                     static_loss_scale=self.loss_scale(),
                     mpu=self.mpu,
                     clip_grad=clip_grad,
-                                           fused_adam_legacy=True)
+                    fused_adam_legacy=self.optimizer_legacy_fusion())
         else:
             logging.info('Creating fp16 unfused optimizer with dynamic loss scale')
             optimizer = FP16_UnfusedOptimizer(
@@ -442,7 +447,7 @@ class DeepSpeedLight(Module):
                 dynamic_loss_args=dynamic_loss_args,
                 mpu=self.mpu,
                 clip_grad=clip_grad,
-                fused_lamb_legacy=True
+                fused_lamb_legacy=self.optimizer_legacy_fusion()
                 if self.optimizer_name() == LAMB_OPTIMIZER else False)
 
         return optimizer

tests/model/Megatron_GPT2/ds_config_func_bs4.json

@@ -5,6 +5,7 @@
   "zero_optimization": true,
   "optimizer": {
     "type": "Adam",
+    "legacy_fusion": false,
     "params": {
       "lr": 0.00015,
       "max_grad_norm": 1.0

tests/model/Megatron_GPT2/ds_config_func_bs8.json

@@ -5,6 +5,7 @@
   "zero_optimization": true,
   "optimizer": {
     "type": "Adam",
+    "legacy_fusion": false,
     "params": {
       "lr": 0.00015,
       "max_grad_norm": 1.0

tests/model/Megatron_GPT2/ds_config_func_bs8_no_zero.json

@@ -5,9 +5,9 @@
   "zero_optimization": false,
   "optimizer": {
     "type": "Adam",
+    "legacy_fusion": false,
     "params": {
-      "lr": 0.00015,
-      "max_grad_norm": 1.0
+      "lr": 0.00015
     }
   },
 

tests/model/Megatron_GPT2/ds_config_func_scheduler.json

@@ -5,6 +5,7 @@
   "zero_optimization": true,
   "optimizer": {
     "type": "Adam",
+    "legacy_fusion": false,
     "params": {
       "lr": 0.00015,
       "max_grad_norm": 1.0

tests/model/Megatron_GPT2/ds_config_perf_bs16.json

@@ -6,6 +6,7 @@
   "disable_allgather": true,
   "optimizer": {
     "type": "Adam",
+    "legacy_fusion": false,
     "params": {
       "lr": 0.00015,
       "max_grad_norm": 1.0

tests/model/Megatron_GPT2/ds_config_perf_bs32.json

@@ -6,6 +6,7 @@
   "disable_allgather": true,
   "optimizer": {
     "type": "Adam",
+    "legacy_fusion": false,
     "params": {
       "lr": 0.00015,
       "max_grad_norm": 1.0

tests/model/Megatron_GPT2/ds_config_perf_bs8.json

@@ -6,6 +6,7 @@
   "disable_allgather": true,
   "optimizer": {
     "type": "Adam",
+    "legacy_fusion": false,
     "params": {
       "lr": 0.00015,
       "max_grad_norm": 1.0