[doc] fix typos and add improvements

docs/code-docs/source/memory.rst

@@ -136,7 +136,7 @@ The big question is how big of a model you can fit on the hardware you have? Or
 
 * ZeRO-3:
 
-largest_layer_memory = 4*largest_layer_params - GPU memory needed to gather the largest layer on a single GPU. 2 bytes fp16 params are gathered and 2 bytes fp16 grads are computed (total 4x). The optimizer states and fp32 parameters are updated in partitioned form and copied to fp16 params in partitioned form. This happens during the optimizer step. After that the fp16 params are sufficient.
+``largest_layer_memory = 4*largest_layer_params`` - GPU memory needed to gather the largest layer on a single GPU. 2 bytes fp16 params are gathered and 2 bytes fp16 grads are computed (total 4x). The optimizer states and fp32 parameters are updated in partitioned form and copied to fp16 params in partitioned form. This happens during the optimizer step. After that the fp16 params are sufficient.
 
    - case 1: ``"cpu_offload": false, "cpu_offload_params": false`` - largest_layer_memory + 18 * params / total number of gpus across all nodes
    - case 2: ``"cpu_offload": true, "cpu_offload_params": true``- largest_layer_memory. The main limit here is general RAM.
@@ -146,39 +146,39 @@ largest_layer_memory = 4*largest_layer_params - GPU memory needed to gather the
 
 .. code-block:: python
 
-from transformers import AutoModel
-model = AutoModel.from_pretrained("t5-large")
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained("t5-large")
 
-# shared params calculated only ones
-total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())
+    # shared params calculated only ones
+    total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())
 
-largest_layer_params = 0
-for m in model.modules():
-    # assuming no shared params within a single layer
-    layer_params = sum(p.numel() for p in m.parameters(recurse=False))
-    largest_layer_params = max(largest_layer_params, layer_params)
+    largest_layer_params = 0
+    for m in model.modules():
+        # assuming no shared params within a single layer
+        layer_params = sum(p.numel() for p in m.parameters(recurse=False))
+        largest_layer_params = max(largest_layer_params, layer_params)
 
-largest_layer_memory = (4*largest_layer_params)
+    largest_layer_memory = (4*largest_layer_params)
 
-total_gpus = 4
+    total_gpus = 4
 
-case1 = largest_layer_memory + int(18*total_params/total_gpus)
-case2 = largest_layer_memory
-case3 = largest_layer_memory + int(2*total_params/total_gpus)
+    case1 = largest_layer_memory + int(18*total_params/total_gpus)
+    case2 = largest_layer_memory
+    case3 = largest_layer_memory + int(2*total_params/total_gpus)
 
-print(f"total params:         {total_params/1e6:6.2f}M")
-print(f"largest layer params: {largest_layer_params/1e6:6.2f}M")
-print(f"largest layer memory: {largest_layer_memory>>20:6}MB")
-print(f"case1 gpu memory: {(case1)>>20:6}MB")
-print(f"case2 gpu memory: {(case2)>>20:6}MB")
-print(f"case3 gpu memory: {(case3)>>20:6}MB")
+    print(f"total params:         {total_params/1e6:6.2f}M")
+    print(f"largest layer params: {largest_layer_params/1e6:6.2f}M")
+    print(f"largest layer memory: {largest_layer_memory>>20:6}MB")
+    print(f"case1 gpu memory: {(case1)>>20:6}MB")
+    print(f"case2 gpu memory: {(case2)>>20:6}MB")
+    print(f"case3 gpu memory: {(case3)>>20:6}MB")
 
-total params:         737.67M
-largest layer params:  32.90M
-largest layer memory:    125MB
-case1 gpu memory:   3291MB
-case2 gpu memory:    125MB
-case3 gpu memory:    477MB
+    total params:         737.67M
+    largest layer params:  32.90M
+    largest layer memory:    125MB
+    case1 gpu memory:   3291MB
+    case2 gpu memory:    125MB
+    case3 gpu memory:    477MB
 
 
 **General RAM**:
@@ -252,12 +252,14 @@ Here is a breakdown for the 16 and 18 multipliers (b = bytes):
 16:
 
 - 16b for fp32: 4b params, 4b grads, 4b momentum and 4b variance per parameter
-
+
 18:
 
 - 16b for fp32: 4b params, 4b grads, 4b momentum and 4b variance per parameter
 - +2b for fp16 params
-
+
+Note about gradients: While gradients are stored in fp16 (2 bytes), during the weight update, all of them are converted into fp32 before doing the weight updates since the weight updates are done at almost the entire model granularity (param_group granularity) in FusedAdam Optimizer in DeepSpeed. So after that conversion we would need the 4 bytes per gradient for nearly the entire set of weights.
+
 
 **Pinned Memory**