add adamw-8bit

app.py

@@ -10,15 +10,12 @@ from vram_helpers import ModelConfig, TrainingConfig, filter_params_for_dataclas
 
 ZERO_STAGES = [0, 1, 2, 3]
 BATCH_SIZES = [1, 2, 4, 8, 16, 32, 64]
-OPTIMIZERS = ["adam", "adamw", "sgd"]
+OPTIMIZERS = ["adam", "adamw", "adamw_8bit", "sgd"]
 HUGGINGFACE_URL_CONFIG = "https://huggingface.co/{}/resolve/main/config.json"
 
 
-
-
 def parse_args():
     parser = argparse.ArgumentParser(description="Parser for VRAM estimator")
-
     parser.add_argument("--repo_id", type=str, default=None, help="HuggingFace repo id to automatically determine model settings")
     parser.add_argument("--model_size", type=float, default=7, help="Model size (in billion parameters)")
     parser.add_argument("--hidden_size", type=int, default=4096, help="Hidden size")

vram_helpers.py

@@ -91,7 +91,7 @@ def optimizer_memory(parameters, optimizer= "adamw", precision = "fp32"):
     optimizer_choices = {"adam": 3,    # Adam: stores precision copies of the optimizer parameters, momentum, and variance -> 4 + 4 + 4 = 12 bytes per model parameter
                          "adamw": 3,   # AdamW: Same for Adam
                          "sgd": 2,      # For SGD: optimier parameters and gradients -> 4 + 4 = 8 bytes per model parameter
-                         "adam-8bit": 1.5, # Adam 8-bit: same for Adam-> 2 + 2 + 2 = 6 bytes per model parameter
+                         "adamw_8bit": 1.5, # Adam 8-bit: same for Adam-> 2 + 2 + 2 = 6 bytes per model parameter
                          }
     return optimizer_choices[optimizer] * parameters * PRECISION_TO_BYTES[precision]