Low-bit memory optimizers, CpuOffloadOptimizer, Memory Reports (#3)

* update * update * update * update * update * update * update * update * update
2026-10-18 00:33:52 +00:00 · 2024-10-04 18:18:54 +05:30
parent 512a8a6854
commit d81209bf49
11 changed files with 434 additions and 314 deletions
@@ -1,4 +1,5 @@
 import gc
+import inspect
 from typing import Optional, Tuple, Union

 import torch
@@ -22,7 +23,11 @@ def get_optimizer(
    prodigy_use_bias_correction: bool = False,
    prodigy_safeguard_warmup: bool = False,
    use_8bit: bool = False,
+    use_4bit: bool = False,
+    use_torchao: bool = False,
    use_deepspeed: bool = False,
+    use_cpu_offload_optimizer: bool = False,
+    offload_gradients: bool = False,
 ) -> torch.optim.Optimizer:
    optimizer_name = optimizer_name.lower()

@@ -38,18 +43,32 @@ def get_optimizer(
            weight_decay=weight_decay,
        )

+    if use_8bit and use_4bit:
+        raise ValueError("Cannot set both `use_8bit` and `use_4bit` to True.")
+
+    if (use_torchao and (use_8bit or use_4bit)) or use_cpu_offload_optimizer:
+        try:
+            import torchao
+
+            torchao.__version__
+        except ImportError:
+            raise ImportError(
+                "To use optimizers from torchao, please install the torchao library: `USE_CPP=0 pip install torchao`."
+            )
+
+    if not use_torchao and use_4bit:
+        raise ValueError("4-bit Optimizers are only supported with torchao.")
+
    # Optimizer creation
-    supported_optimizers = ["adam", "adamw", "prodigy"]
+    supported_optimizers = ["adam", "adamw", "prodigy", "came"]
    if optimizer_name not in supported_optimizers:
        logger.warning(
            f"Unsupported choice of optimizer: {optimizer_name}. Supported optimizers include {supported_optimizers}. Defaulting to `AdamW`."
        )
        optimizer_name = "adamw"

-    if use_8bit and optimizer_name not in ["adam", "adamw"]:
-        logger.warning(
-            f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was set to {optimizer_name}."
-        )
+    if (use_8bit or use_4bit) and optimizer_name not in ["adam", "adamw"]:
+        raise ValueError("`use_8bit` and `use_4bit` can only be used with the Adam and AdamW optimizers.")

    if use_8bit:
        try:
@@ -60,24 +79,32 @@ def get_optimizer(
            )

    if optimizer_name == "adamw":
-        optimizer_class = bnb.optim.AdamW8bit if use_8bit else torch.optim.AdamW
+        if use_torchao:
+            from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit

-        optimizer = optimizer_class(
-            params_to_optimize,
-            betas=(beta1, beta2),
-            eps=epsilon,
-            weight_decay=weight_decay,
-        )
+            optimizer_class = AdamW8bit if use_8bit else AdamW4bit if use_4bit else torch.optim.AdamW
+        else:
+            optimizer_class = bnb.optim.AdamW8bit if use_8bit else torch.optim.AdamW
+
+        init_kwargs = {
+            "betas": (beta1, beta2),
+            "eps": epsilon,
+            "weight_decay": weight_decay,
+        }

    elif optimizer_name == "adam":
-        optimizer_class = bnb.optim.Adam8bit if use_8bit else torch.optim.Adam
+        if use_torchao:
+            from torchao.prototype.low_bit_optim import Adam4bit, Adam8bit

-        optimizer = optimizer_class(
-            params_to_optimize,
-            betas=(beta1, beta2),
-            eps=epsilon,
-            weight_decay=weight_decay,
-        )
+            optimizer_class = Adam8bit if use_8bit else Adam4bit if use_4bit else torch.optim.Adam
+        else:
+            optimizer_class = bnb.optim.Adam8bit if use_8bit else torch.optim.Adam
+
+        init_kwargs = {
+            "betas": (beta1, beta2),
+            "eps": epsilon,
+            "weight_decay": weight_decay,
+        }

    elif optimizer_name == "prodigy":
        try:
@@ -92,17 +119,43 @@ def get_optimizer(
                "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
            )

-        optimizer = optimizer_class(
-            params_to_optimize,
-            lr=learning_rate,
-            betas=(beta1, beta2),
-            beta3=beta3,
-            weight_decay=weight_decay,
-            eps=epsilon,
-            decouple=prodigy_decouple,
-            use_bias_correction=prodigy_use_bias_correction,
-            safeguard_warmup=prodigy_safeguard_warmup,
+        init_kwargs = {
+            "lr": learning_rate,
+            "betas": (beta1, beta2),
+            "beta3": beta3,
+            "eps": epsilon,
+            "weight_decay": weight_decay,
+            "decouple": prodigy_decouple,
+            "use_bias_correction": prodigy_use_bias_correction,
+            "safeguard_warmup": prodigy_safeguard_warmup,
+        }
+
+    elif optimizer_name == "came":
+        try:
+            import came_pytorch
+        except ImportError:
+            raise ImportError("To use CAME, please install the came-pytorch library: `pip install came-pytorch`")
+
+        optimizer_class = came_pytorch.CAME
+
+        init_kwargs = {
+            "lr": learning_rate,
+            "eps": (1e-30, 1e-16),
+            "betas": (beta1, beta2, beta3),
+            "weight_decay": weight_decay,
+        }
+
+    if use_cpu_offload_optimizer:
+        from torchao.prototype.low_bit_optim import CPUOffloadOptimizer
+
+        if "fused" in inspect.signature(optimizer_class.__init__).parameters:
+            init_kwargs.update({"fused": True})
+
+        optimizer = CPUOffloadOptimizer(
+            params_to_optimize, optimizer_class=optimizer_class, offload_gradients=offload_gradients, **init_kwargs
        )
+    else:
+        optimizer = optimizer_class(params_to_optimize, **init_kwargs)

    return optimizer