[DeepSpeed] ZeRO Stage 3

docs/source/main_classes/trainer.rst

@@ -736,7 +736,7 @@ which ZeRO stages you want to enable and how to configure them.
        }
     }
 
-Notes:
+Performance tuning:
 
 - enabling ``cpu_offload`` should reduce GPU RAM usage (it requires ``"stage": 2``)
 - ``"overlap_comm": true`` trades off increased GPU RAM usage to lower all-reduce latency. ``overlap_comm`` uses 4.5x
@@ -752,6 +752,31 @@ This section has to be configured exclusively via DeepSpeed configuration - the
 no equivalent command line arguments.
 
 
+ZeRO-3
+"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+
+Performance tuning:
+
+- ``stage3_max_live_parameters``: ``1e9``
+- ``stage3_max_reuse_distance``:  ``1e9``
+- ``stage3_prefetch_bucket_size``:  ``0.9 * hidden_size*hidden_size``
+- ``stage3_param_persitance_threshold``:  ``10*hidden_size``
+- ``reduce_bucket_size``: ``hidden_size*hidden_size``
+- ``sub_group_size`` ~``1e15``
+
+If hitting OOM reduce ``stage3_max_live_parameters`` and ``stage3_max_reuse_distance``. They should have minimal impact on perf unless you are doing activation checkpointing. ``1e9`` would consume ~2GB.
+The memory is shared by ``stage3_max_live_parameters`` and ``stage3_max_reuse_distance``, so its not additive, its just 2GB total.
+
+``stage3_max_live_parameters`` is the upper limit on how many full parameters you want to keep on the GPU at any given time.  ``reuse_distance`` is a metric we are using to figure out when will a parameter be used again in the future, and we use the ``stage3_max_reuse_distance`` to decide whether to throw away the parameter or to keep it. If a parameter is going to be used again in near future (less than  ``stage3_max_reuse_distance``) then we keep it to reduce communication overhead. This is super helpful when you have activation check-pointing enabled, where we do a forward recompute and backward passes a a single layer granularity and want to keep the parameter in the forward recompute till the backward
+
+
+If you set ``reduce_bucket_size``, ``stage3_prefetch_bucket_size`` and ``stage3_param_persitance_threshold``: as recommended above, they will already be fairly small so you wont have to tune those much.
+
+Other:
+
+``allgather_partitions``, ``allgather_bucket_size`` and ``reduce_scatter`` are not used in ZeRO-3.
+
 
 Optimizer and Scheduler
 =======================================================================================================================

examples/seq2seq/run_summarization.py

@@ -592,4 +592,9 @@ def _mp_fn(index):
 
 
 if __name__ == "__main__":
+
+    import faulthandler, signal
+
+    #faulthandler.dump_traceback_later(20, repeat=True)
+
     main()

examples/tests/deepspeed/ds_config.json → examples/tests/deepspeed/ds_config_zero2.json

@@ -3,7 +3,7 @@
         "enabled": true,
         "loss_scale": 0,
         "loss_scale_window": 1000,
-        "initial_scale_power": 32,
+        "initial_scale_power": 16,
         "hysteresis": 2,
         "min_loss_scale": 1
     },
@@ -19,6 +19,8 @@
         "cpu_offload": true
     },
 
+    "zero_allow_untested_optimizer": true,
+
     "optimizer": {
         "type": "AdamW",
         "params": {

examples/tests/deepspeed/ds_config_zero3.json

@@ -0,0 +1,48 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "cpu_offload": true,
+        "cpu_offload_params": true,
+        "cpu_offload_use_pin_memory" : true,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e8,
+        "stage3_prefetch_bucket_size": 2e5,
+        "stage3_param_persitance_threshold": 1e5,
+        "reduce_bucket_size": 3e6,
+        "prefetch_bucket_size": 3e6,
+        "sub_group_size": 1e6
+    },
+
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": 3e-5,
+            "betas": [0.8, 0.999],
+            "eps": 1e-8,
+            "weight_decay": 3e-7
+        }
+    },
+
+    "scheduler": {
+        "type": "WarmupLR",
+        "params": {
+            "warmup_min_lr": 0,
+            "warmup_max_lr": 3e-5,
+            "warmup_num_steps": 500
+        }
+    },
+
+    "steps_per_print": 2000,
+    "wall_clock_breakdown": false
+}