Update calculations

src/app.py

@@ -9,7 +9,7 @@ from huggingface_hub.utils import HfHubHTTPError
 MODEL = None
 
 
-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):
+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, additional_buffer_factor: float):
     global MODEL
     MODEL = get_model(model_name, library, access_token)
     try:
@@ -26,6 +26,7 @@ def get_results(model_name: str, library: str, precision: list, training: list,
         "num_nodes": num_nodes,
         "num_gpus_per_node": num_gpus,
         "training_regime": training,
+        "additional_buffer_factor": additional_buffer_factor
     }
     data = calculate_memory(MODEL, options)
     
@@ -36,7 +37,7 @@ def get_results(model_name: str, library: str, precision: list, training: list,
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown(
-            """<img src="https://huggingface.co/spaces/hf-accelerate/model-memory-usage/resolve/main/measure_model_size.png" style="float: left;" width="250" height="250"><h1>🤗 Model Memory Calculator</h1>
+            """<img src="https://huggingface.co/spaces/andstor/model-memory-usage/resolve/main/measure_model_size.png" style="float: left;" width="250" height="250"><h1>🤗 DeepSpeed Model Memory Calculator</h1>
 
     This tool will help you calculate how much vRAM is needed to train and perform big model inference
     on a model hosted on the 🤗 Hugging Face Hub. The minimum recommended vRAM needed for a model
@@ -74,16 +75,21 @@ with gr.Blocks() as demo:
                 label="Training Paradigm",
             )
             access_token = gr.Textbox(label="API Token", placeholder="Optional (for gated models)")
-        with gr.Row():
-            with gr.Column():
+            num_gpus = gr.Number(label="GPUs per node", value=4, minimum=1, step=1)
+            num_nodes = gr.Number(label="Nodes", value=1, minimum=1, step=1)
+        with gr.Column(variant="panel"):
+            with gr.Row(equal_height=True):
+
                 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(equal_height=True):
+                #with gr.Column():
+                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=["zero.Init"], label="Initialization")
+                
+                #with gr.Column():
+                additional_buffer_factor = gr.Number(label="Additional Buffer Factor", value=1.5, minimum=1, step=0.1)
         with gr.Row():
             btn = gr.Button("Calculate Memory Usage")
             post_to_hub = gr.Button(
@@ -135,7 +141,7 @@ with gr.Blocks() as demo:
 
     btn.click(
         get_results,
-        inputs=[inp, library, precision, training, access_token, zero_stage, num_nodes, num_gpus, offloading, zero_init],
+        inputs=[inp, library, precision, training, access_token, zero_stage, num_nodes, num_gpus, offloading, zero_init, additional_buffer_factor],
         outputs=[out_text, out, post_to_hub],
     )
 

src/model_utils.py

@@ -10,8 +10,8 @@ from parallelism_utils import estimate_zero1_model_states_mem_needs, estimate_ze
 
 
 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}
+PRECISION = {"Mixed precision": "mixed", "Single precision": "single"}
+DTYPE = {"float32": torch.float32, "float16/bfloat16": torch.float16}
 
 
 def extract_from_url(name: str):
@@ -93,25 +93,25 @@ def calculate_memory(model: torch.nn.Module, options: dict):
 
         dtype_largest_layer = convert_bytes(dtype_largest_layer)
 
-        precision_fac = PRECISION_FACTOR[options["training_regime"]]
-        params_fac = DTYPE_FACTOR[dtype]
+        precision = PRECISION[options["training_regime"]]
+        model_dtype = DTYPE[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)
+            cpu_mem, gpu_mem = estimate_zero1_model_states_mem_needs(total_params, options["num_gpus_per_node"], options["num_nodes"], options["cpu_offload"], options["additional_buffer_factor"], precision, model_dtype)
         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)
+            cpu_mem, gpu_mem = estimate_zero2_model_states_mem_needs(total_params, options["num_gpus_per_node"], options["num_nodes"], options["cpu_offload"], options["additional_buffer_factor"], precision, model_dtype)
         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)
+            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"], options["additional_buffer_factor"], precision, model_dtype)
         data.append(
                 {
-                    "dtype": dtype,
+                    "Model dtype": dtype,
                     "Largest Layer or Residual Group": dtype_largest_layer,
-                    "Total Size": convert_bytes(dtype_total_size),
+                    "Model Size": convert_bytes(dtype_total_size),
                     "per CPU": convert_bytes(cpu_mem),
-                    "per GPU (Adam)": convert_bytes(gpu_mem),
+                    "per GPU": convert_bytes(gpu_mem),
                 }
             )
 

src/parallelism_utils.py

@@ -1,65 +1,90 @@
-# Zero Redundancy Optimizer (ZeRO)
+import torch
+
+def get_precision_fac(precision: str):
+    if precision == "mixed":
+        return 2
+    elif precision == "single":
+        return 4
+    else:
+        raise ValueError("Precision must be either 'mixed' or 'single'")
+
+
+def get_params_fac(model_dtype: torch.dtype):
+    if model_dtype == torch.float16:
+        return 2
+    elif model_dtype == torch.float32:
+        return 4
+    else:
+        raise ValueError("Model dtype must be either torch.float16 or torch.float32")
+
+
+
+####################### Zero Redundancy Optimizer (ZeRO) #######################
+
+VARIANCE_FACTOR = 4
+MOMENTUM_FACTOR = 4
+OPTIMIZER_FACTOR = VARIANCE_FACTOR + MOMENTUM_FACTOR # Adam optimizer
+FP32_GRADS_FACTOR = 4
+FP32_PARAM_FACTOR = 4
+MASTER_PARAMS_FACTOR = FP32_PARAM_FACTOR
+
+
+# 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.
+    
+
+
 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
+                                          precision="mixed",
+                                          model_dtype = torch.float16,
                                           ):
    
-    # 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
+    
+    precision_fac = get_precision_fac(precision)
+    params_fac = get_params_fac(model_dtype)
 
     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
+        gpu_mem = (precision_fac * total_params) # + (grads_fac * total_params)
+        cpu_mem = total_params * max(params_fac * total_gpus, (MASTER_PARAMS_FACTOR + OPTIMIZER_FACTOR + FP32_GRADS_FACTOR)) * 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)
+        if precision == "mixed":
+            gpu_mem = (precision_fac * total_params) + (FP32_GRADS_FACTOR * total_params) + int((OPTIMIZER_FACTOR + FP32_PARAM_FACTOR) * total_params / total_gpus)
+        else:
+            gpu_mem = (precision_fac * total_params) + (FP32_GRADS_FACTOR * total_params) + int(OPTIMIZER_FACTOR * 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
+                                          precision="mixed",
+                                          model_dtype = torch.float16,
                                           ):
-   
-    # 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
+    
+    precision_fac = get_precision_fac(precision)
+    params_fac = get_params_fac(model_dtype)
 
     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
+        gpu_mem = precision_fac * total_params # Negligible memory usage for partitioned gradients
+        cpu_mem = total_params * max(params_fac * total_gpus, (MASTER_PARAMS_FACTOR + OPTIMIZER_FACTOR + FP32_GRADS_FACTOR)) * 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
+        if precision == "mixed":
+            gpu_mem = precision_fac * total_params + int((FP32_GRADS_FACTOR + OPTIMIZER_FACTOR + FP32_PARAM_FACTOR) * total_params / total_gpus)
+        else:
+            gpu_mem = precision_fac * total_params + int((FP32_GRADS_FACTOR + OPTIMIZER_FACTOR) * total_params / total_gpus)
+        cpu_mem =  params_fac * total_params * num_gpus_per_node * additional_buffer_factor
 
     return int(cpu_mem), int(gpu_mem)
 
@@ -72,43 +97,48 @@ def estimate_zero3_model_states_mem_needs(total_params,
                                           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
+                                          precision="mixed",
+                                          model_dtype = torch.float16,
                                           ):
 
-    # 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)
+    precision_fac = get_precision_fac(precision)
+    params_fac = get_params_fac(model_dtype)
+    grads_fac = precision_fac
+
+    largest_layer_memory = (grads_fac + precision_fac) * largest_layer_params
 
     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
+                cpu_mem = total_params * (MASTER_PARAMS_FACTOR + FP32_GRADS_FACTOR + OPTIMIZER_FACTOR + 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
+                cpu_mem = total_params * max(params_fac * num_gpus_per_node, (MASTER_PARAMS_FACTOR + FP32_GRADS_FACTOR + OPTIMIZER_FACTOR + params_fac) * gpus_factor) * additional_buffer_factor
+        
         else:
-            gpu_mem = largest_layer_memory + int(precision_fac * total_params / total_gpus)
+            gpu_mem = max(
+                largest_layer_memory,
+                int((precision_fac) * total_params / total_gpus) # No need for gradients: ZeRO-Offload can transfer these gradients for each parameter individually or in small groups to the CPU memory immediately after they are computed
+            )
 
             if zero_init:
-                cpu_mem = total_params * (master_params_fac + grads_fac + optimizer_fac) * gpus_factor * additional_buffer_factor 
+                cpu_mem = total_params * (MASTER_PARAMS_FACTOR + FP32_GRADS_FACTOR + OPTIMIZER_FACTOR) * 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
+                cpu_mem = total_params * max(params_fac * num_gpus_per_node, (MASTER_PARAMS_FACTOR + FP32_GRADS_FACTOR + OPTIMIZER_FACTOR) * 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 precision == "mixed":
+            gpu_mem = max(
+                int((precision_fac + FP32_GRADS_FACTOR + OPTIMIZER_FACTOR + FP32_PARAM_FACTOR) * largest_layer_params),
+                int((precision_fac + FP32_GRADS_FACTOR + OPTIMIZER_FACTOR + FP32_PARAM_FACTOR) * total_params / total_gpus)
+            )
+        else:
+            gpu_mem = max(
+                int((precision_fac + FP32_GRADS_FACTOR + OPTIMIZER_FACTOR) * largest_layer_params),
+                int((precision_fac + FP32_GRADS_FACTOR + OPTIMIZER_FACTOR) * total_params / total_gpus)
+            )
 
         if zero_init:
             cpu_mem = largest_layer_params * params_fac * num_gpus_per_node * additional_buffer_factor