Add DeepSpeed ZeRO calculations

src/app.py

@@ -9,15 +9,27 @@ from huggingface_hub.utils import HfHubHTTPError
 MODEL = None
 
 
-def get_results(model_name: str, library: str, options: list, access_token: str):
+def get_results(model_name: str, library: str, precision: list, training: list,  access_token: str, zero_stage: int, num_nodes: int, num_gpus: int, offloading: list, zero_init: list):
     global MODEL
     MODEL = get_model(model_name, library, access_token)
     try:
         has_discussion = check_for_discussion(model_name)
     except HfHubHTTPError:
         has_discussion = True
-    title = f"## Memory usage for '{model_name}'"
+
+    options = {
+        "precision": precision,
+        "zero_stage": zero_stage,
+        "cpu_offload": True if "Optimizer" in offloading else False,
+        "cpu_offload_params": True if "Parameters" in offloading else False,
+        "zero_init": True if "zero.Init" in zero_init else False,
+        "num_nodes": num_nodes,
+        "num_gpus_per_node": num_gpus,
+        "training_regime": training,
+    }
     data = calculate_memory(MODEL, options)
+    
+    title = f"## Memory usage for '{model_name}'"
     return [title, gr.update(visible=True, value=pd.DataFrame(data)), gr.update(visible=not has_discussion)]
 
 
@@ -51,21 +63,79 @@ with gr.Blocks() as demo:
             inp = gr.Textbox(label="Model Name or URL", value="bert-base-cased")
         with gr.Row():
             library = gr.Radio(["auto", "transformers", "timm"], label="Library", value="auto")
-            options = gr.CheckboxGroup(
-                ["float32", "float16/bfloat16", "int8", "int4"],
+            precision = gr.CheckboxGroup(
+                ["float32", "float16/bfloat16"],
                 value="float32",
                 label="Model Precision",
             )
+            training = gr.Radio(
+                ["Mixed precision", "Single precision"],
+                value="Mixed precision",
+                label="Training Paradigm",
+            )
             access_token = gr.Textbox(label="API Token", placeholder="Optional (for gated models)")
+        with gr.Row():
+            with gr.Column():
+                zero_stage = gr.Radio(["Stage 0", "Stage 1", "Stage 2", "Stage 3"], label="ZeRO Stage", value="Stage 3", type="index")
+                zero_description = gr.CheckboxGroup(["Optimizer state",  "Gradients", "Parameters"], label="Partitioning", value=["Optimizer state",  "Gradients", "Parameters"], interactive=False)
+                with gr.Row():
+                    offloading = gr.CheckboxGroup(["Optimizer",  "Parameters"], label="ZeRO-Offload", info="Offloading data and compute to CPU", value=["Optimizer",  "Parameters"])
+                    zero_init = gr.CheckboxGroup(["zero.Init"], value=True, label="Initialization")
+            
+            num_gpus = gr.Number(label="GPUs per node", value=1, min=1, step=1)
+            num_nodes = gr.Number(label="Nodes", value=1, min=1, step=1)
         with gr.Row():
             btn = gr.Button("Calculate Memory Usage")
             post_to_hub = gr.Button(
                 value="Report results in this model repo's discussions!\n(Will open in a new tab)", visible=False
             )
 
+    def change_zero_settings(evt: gr.SelectData):  # SelectData is a subclass of EventData
+        if evt.index == 0:
+            return [gr.update(visible = False), gr.update(visible = False)]
+        if evt.index == 1 or evt.index == 2:
+            return [gr.update(choices=["Optimizer"], visible=True), gr.update(visible = False)]
+        if evt.index == 3:
+            return [gr.update(choices=["Optimizer", "Parameters"], visible=True), gr.update(visible = True)]
+    
+    def change_zero_description(evt: gr.SelectData):  # SelectData is a subclass of EventData
+        if evt.index == 0:
+            return gr.update(value=None)
+        if evt.index == 1:
+            return gr.update(value=["Optimizer state"])
+        if evt.index == 2:
+            return gr.update(value=["Optimizer state", "Gradients"])
+        if evt.index == 3:
+            return gr.update(value=["Optimizer state", "Gradients", "Parameters"])
+    
+    def change_offloading(evt: gr.SelectData, zero_stage):  # SelectData is a subclass of EventData
+        
+        if evt.value == "Optimizer" and evt.selected == False:
+            return gr.CheckboxGroup.update(choices=["Optimizer"], value=[])
+        
+        if evt.value == "Optimizer" and evt.selected == True:
+            if zero_stage in [1, 2]:
+                return gr.CheckboxGroup.update(choices=["Optimizer"], value=["Optimizer"])
+            elif zero_stage == 3:
+                return gr.CheckboxGroup.update(choices=["Optimizer", "Parameters"], value=["Optimizer"])
+
+        if evt.value == "Parameters" and evt.selected == False:
+            return gr.CheckboxGroup.update(value=["Optimizer"])
+
+        if evt.value == "Parameters" and evt.selected == True:
+
+            return gr.CheckboxGroup.update(value=["Optimizer", "Parameters"])
+
+
+
+    zero_stage.select(change_zero_settings, None, [offloading, zero_init])
+    zero_stage.select(change_zero_description, None, zero_description)
+    offloading.select(change_offloading, zero_stage, offloading)
+
+
     btn.click(
         get_results,
-        inputs=[inp, library, options, access_token],
+        inputs=[inp, library, precision, training, access_token, zero_stage, num_nodes, num_gpus, offloading, zero_init],
         outputs=[out_text, out, post_to_hub],
     )
 

src/model_utils.py

@@ -6,9 +6,12 @@ import torch
 from accelerate.commands.estimate import check_has_model, create_empty_model
 from accelerate.utils import calculate_maximum_sizes, convert_bytes
 from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
+from parallelism_utils import estimate_zero1_model_states_mem_needs, estimate_zero2_model_states_mem_needs, estimate_zero3_model_states_mem_needs
 
 
 DTYPE_MODIFIER = {"float32": 1, "float16/bfloat16": 2, "int8": 4, "int4": 8}
+PRECISION_FACTOR = {"Mixed precision": 2, "Single precision": 4}
+DTYPE_FACTOR = {"float32": 4, "float16/bfloat16": 2}
 
 
 def extract_from_url(name: str):
@@ -74,12 +77,13 @@ def get_model(model_name: str, library: str, access_token: str):
     return model
 
 
-def calculate_memory(model: torch.nn.Module, options: list):
+def calculate_memory(model: torch.nn.Module, options: dict):
     "Calculates the memory usage for a model init on `meta` device"
     total_size, largest_layer = calculate_maximum_sizes(model)
+    total_params = model.num_parameters()
 
     data = []
-    for dtype in options:
+    for dtype in options["precision"]:
         dtype_total_size = total_size
         dtype_largest_layer = largest_layer[0]
 
@@ -87,15 +91,28 @@ def calculate_memory(model: torch.nn.Module, options: list):
         dtype_total_size /= modifier
         dtype_largest_layer /= modifier
 
-        dtype_training_size = convert_bytes(dtype_total_size * 4)
-        dtype_total_size = convert_bytes(dtype_total_size)
         dtype_largest_layer = convert_bytes(dtype_largest_layer)
+
+        precision_fac = PRECISION_FACTOR[options["training_regime"]]
+        params_fac = DTYPE_FACTOR[dtype]
+
+        if options["zero_stage"] == 0:
+            cpu_mem = dtype_total_size * 4
+            gpu_mem = cpu_mem
+        elif options["zero_stage"] == 1:
+            cpu_mem, gpu_mem = estimate_zero1_model_states_mem_needs(total_params, options["num_gpus_per_node"], options["num_nodes"], options["cpu_offload"], precision_fac, params_fac)
+        elif options["zero_stage"] == 2:
+            cpu_mem, gpu_mem = estimate_zero2_model_states_mem_needs(total_params, options["num_gpus_per_node"], options["num_nodes"], options["cpu_offload"], precision_fac, params_fac)
+        elif options["zero_stage"] == 3:
+            cpu_mem, gpu_mem, largest_layer_memory = estimate_zero3_model_states_mem_needs(total_params, largest_layer[0], options["num_gpus_per_node"], options["num_nodes"], options["cpu_offload"], options["cpu_offload_params"], options["zero_init"], precision_fac, params_fac)
         data.append(
-            {
-                "dtype": dtype,
-                "Largest Layer or Residual Group": dtype_largest_layer,
-                "Total Size": dtype_total_size,
-                "Training using Adam": dtype_training_size,
-            }
-        )
+                {
+                    "dtype": dtype,
+                    "Largest Layer or Residual Group": dtype_largest_layer,
+                    "Total Size": convert_bytes(dtype_total_size),
+                    "per CPU": convert_bytes(cpu_mem),
+                    "per GPU (Adam)": convert_bytes(gpu_mem),
+                }
+            )
+
     return data

src/parallelism_utils.py

@@ -0,0 +1,120 @@
+# Zero Redundancy Optimizer (ZeRO)
+def estimate_zero1_model_states_mem_needs(total_params,
+                                          num_gpus_per_node=1,
+                                          num_nodes=1,
+                                          cpu_offload=True,
+                                          additional_buffer_factor=1.5,
+                                          precision_fac = 2, # half precision
+                                          params_fac = 4 # 4 bytes per float32 model parameter type
+                                          ):
+   
+    # TODO: check if params_fac is needed during full fp32 training.
+    # Normally, mixed precision training results in 1.5x memory compared to FP32.
+    # Currently, we are assuming 2x memory for FP32, as deepspeed's ZeRO-2 is optimized for FP16 training.
+    
+    total_gpus = num_nodes * num_gpus_per_node
+
+    master_params_fac = 4
+    variance_fac = 4
+    momentum_fac = 4
+    grads_fac = 4
+    optimizer_fac = variance_fac + momentum_fac # Adam optimizer
+
+    total_gpus = num_nodes * num_gpus_per_node
+
+    if cpu_offload:
+        gpu_mem = (precision_fac * total_params) + (precision_fac * total_params)
+        cpu_mem = total_params * max(params_fac * total_gpus, (master_params_fac+optimizer_fac+grads_fac)) * additional_buffer_factor
+    else:
+        gpu_mem = (precision_fac * total_params) + (precision_fac * total_params) + int((precision_fac + optimizer_fac + master_params_fac + precision_fac) * total_params / total_gpus)
+        cpu_mem = total_params * params_fac * num_gpus_per_node * additional_buffer_factor
+
+    return int(cpu_mem), int(gpu_mem)
+
+def estimate_zero2_model_states_mem_needs(total_params,
+                                          num_gpus_per_node=1,
+                                          num_nodes=1,
+                                          cpu_offload=True,
+                                          additional_buffer_factor=1.5,
+                                          precision_fac = 2, # half precision
+                                          params_fac = 4 # 4 bytes per float32 model parameter type
+                                          ):
+   
+    # TODO: check if params_fac is needed during full fp32 training.
+    # Normally, mixed precision training results in 1.5x memory compared to FP32.
+    # Currently, we are assuming 2x memory for FP32, as deepspeed's ZeRO-2 is optimized for FP16 training.
+    
+    total_gpus = num_nodes * num_gpus_per_node
+
+    master_params_fac = 4
+    variance_fac = 4
+    momentum_fac = 4
+    grads_fac = 4
+    optimizer_fac = variance_fac + momentum_fac # Adam optimizer
+
+    total_gpus = num_nodes * num_gpus_per_node
+
+    if cpu_offload:
+        gpu_mem = precision_fac * total_params
+        cpu_mem = total_params * max(params_fac * total_gpus, (master_params_fac+optimizer_fac+grads_fac)) * additional_buffer_factor
+    else:
+        gpu_mem = precision_fac * total_params + int((precision_fac + grads_fac + optimizer_fac + master_params_fac + precision_fac) * total_params / total_gpus)
+        cpu_mem = total_params * params_fac * num_gpus_per_node * additional_buffer_factor
+
+    return int(cpu_mem), int(gpu_mem)
+
+
+def estimate_zero3_model_states_mem_needs(total_params,
+                                          largest_layer_params,
+                                          num_gpus_per_node=1,
+                                          num_nodes=1,
+                                          cpu_offload=True,
+                                          cpu_offload_params=True,
+                                          zero_init=True,
+                                          additional_buffer_factor=1.5,
+                                          precision_fac = 2, # half precision
+                                          params_fac = 4 # 4 bytes per float32 model parameter type
+                                          ):
+
+    # TODO: check if params_fac is needed during full fp32 training.
+    # Normally, mixed precision training results in 1.5x memory compared to FP32.
+    # Currently, we are assuming 2x memory for FP32, as deepspeed's ZeRO-2 is optimized for FP16 training.
+        
+    total_gpus = num_nodes * num_gpus_per_node
+    gpus_factor = 1 / num_nodes
+    master_params_fac = 4
+    variance_fac = 4
+    momentum_fac = 4
+    grads_fac = 4
+    optimizer_fac = variance_fac + momentum_fac # Adam optimizer
+
+    largest_layer_memory = (2 * precision_fac) * largest_layer_params # params + grads = (2 * modifier)
+
+    if cpu_offload:
+        if cpu_offload_params:
+            gpu_mem = largest_layer_memory
+
+            if zero_init:
+                cpu_mem = total_params * (master_params_fac + grads_fac + optimizer_fac + params_fac) * gpus_factor * additional_buffer_factor
+            else:
+
+                cpu_mem = total_params * max(params_fac * num_gpus_per_node, (master_params_fac + grads_fac + optimizer_fac + params_fac) * gpus_factor) * additional_buffer_factor
+        else:
+            gpu_mem = largest_layer_memory + int(precision_fac * total_params / total_gpus)
+
+            if zero_init:
+                cpu_mem = total_params * (master_params_fac + grads_fac + optimizer_fac) * gpus_factor * additional_buffer_factor 
+            else:
+                cpu_mem = total_params * max(params_fac * num_gpus_per_node, (master_params_fac + grads_fac + optimizer_fac) * gpus_factor) * additional_buffer_factor
+    else:
+        gpu_mem = largest_layer_memory + int((master_params_fac + grads_fac + optimizer_fac + precision_fac) * total_params / total_gpus)
+        # 2b for fp16 params, 4b master params, 4b grads, 4b momentum and 4b variance per parameter = 18
+
+        if zero_init:
+            cpu_mem = largest_layer_params * params_fac * num_gpus_per_node * additional_buffer_factor
+        else:
+            cpu_mem = total_params * params_fac * num_gpus_per_node * additional_buffer_factor
+
+    return int(cpu_mem), int(gpu_mem), largest_layer_memory
+
+