[DeepSpeed] ZeRO Stage 3

docs/source/main_classes/trainer.rst

@@ -589,6 +589,7 @@ or with bash magic, where you can write a multi-line code for the shell to run:
 
 
 
+
 Configuration
 =======================================================================================================================
 
@@ -736,7 +737,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 +753,41 @@ 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_persistence_threshold``: ``10*hidden_size``
+- ``reduce_bucket_size``: ``hidden_size*hidden_size``
+- ``sub_group_size``: ``1e14``
+
+If hitting OOM reduce ``stage3_max_live_parameters`` and ``stage3_max_reuse_distance``. They should have minimal impact
+on performance 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_persistence_threshold``: as
+recommended above, they will already be fairly small so you won't have to tune those much.
+
+Other:
+
+``allgather_partitions``, ``allgather_bucket_size`` and ``reduce_scatter`` are not used in ZeRO-3.
+
+
+
+
 
 Optimizer and Scheduler
 =======================================================================================================================
@@ -963,6 +999,89 @@ Here is an example of the ``gradient_clipping`` configuration:
 
 
 
+ZeRO 3 Nuances
+=======================================================================================================================
+
+ZeRO 3 is quite different from ZeRO 2 because of its param sharding, so the following might be of use in certain
+circumstances.
+
+
+Registering External Parameters
+"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+If layer A needs to access weights belonging to layer B, currently layer A needs to tell DeepSpeed about it. This is
+done with the help of ``deepspeed.zero.register_external_parameter`` that needs to be called in ``A.__init__`` and can
+be seen in the following example:
+
+.. code-block:: python
+
+   class ModuleZ3(torch.nn.Module):
+       def __init__(self, *args):
+           super().__init__(self, *args)
+           self.layer1 = SomeLayer()
+           self.layer2 = OtherLayer()
+           deepspeed.zero.register_external_parameter(self, self.layer1.weight)
+
+       def forward(self, input):
+           x = self.layer1(input)
+           # self.layer1.weight is needed in ModuleZ3.forward
+           y = self.layer2(x, self.layer1.weight)
+           return y
+
+
+For full details on this method please refer to `Registering External Parameters
+<https://deepspeed.readthedocs.io/en/latest/zero3.html#registering-external-parameters>`__
+
+
+
+Constructing Massive Models
+"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+DeepSpeed/ZeRO-3 can handle models with Trillions of parameters which may not fit onto the existing RAM. In such cases, but also if you want the initialization to happen much faster, initialize the model using `deepspeed.zero.Init()`
+context manager (which is also a function decorator), like so:
+
+.. code-block:: python
+
+    from transformers import T5ForConditionalGeneration, T5Config
+    import deepspeed
+    with deepspeed.zero.Init():
+       config = T5Config.from_pretrained("t5-small")
+       model = T5ForConditionalGeneration(config)
+
+As you can see this gives you a randomly initialized model.
+
+If you want to use a pretrained model, ``model_class.from_pretrained`` will activate this feature as long as ``deepspeed_is_zero3_enabled()`` returns ``True``. Therefore to enable this feature here is the required sequence:
+
+.. code-block:: python
+
+    from transformers.integrations import deepspeed_is_zero3_enabled
+    deepspeed_is_zero3_enabled(True)
+    model = T5ForConditionalGeneration.from_pretrained("t5-small")
+
+If you're using Trainer command line arguments which include ``--deepspeed ds_config.json`` with ZeRO3 config enabled, then you can skip ``deepspeed_is_zero3_enabled(True)`` as it will try to discover whether it'll be run under zero3 and ``from_pretrained`` will automatically activate this feature.
+
+Note: If the fp16 weights of the model can't fit onto the memory of a single GPU this feature must be used.
+
+For full details on this method and other related features please refer to `Constructing Massive Models
+<https://deepspeed.readthedocs.io/en/latest/zero3.html#constructing-massive-models>`__.
+
+
+
+
+External Parameters
+"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+When using ``deepspeed.zero.Init()`` discussed earlier, a layer A may need to access layer B's weights
+
+
+XXX: complete
+
+
+For full details on this method please refer to `Gathering Parameters
+<https://deepspeed.readthedocs.io/en/latest/zero3.html#manual-parameter-coordination>`__
+
+
+
 Notes
 =======================================================================================================================
 

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
     },

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": 1e9,
+        "stage3_prefetch_bucket_size": 0.94e6,
+        "stage3_param_persistence_threshold": 1e4,
+        "reduce_bucket_size": 1e6,
+        "prefetch_bucket_size": 3e6,
+        "sub_group_size": 1e14
+    },
+
+    "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
+}