Optimizer CPU offload for single GPU training

benchmarks/benchmark_low_bit_adam.py

@@ -32,7 +32,7 @@
 
 # lpmm doesn't have Adam, only AdamW
 OPTIM_MAP = dict(
-    Adam=torch.optim.Adam,
+    Adam=partial(torch.optim.Adam, fused=True),
     Adam8bitBnb=bnb.optim.Adam8bit,
     Adam8bitAo=low_bit_optim.Adam8bit,
     AdamFp8Ao=low_bit_optim.AdamFp8,
@@ -90,6 +90,7 @@ def get_parser():
     parser.add_argument("--lr", type=float, default=1e-4)
     parser.add_argument("--weight_decay", type=float, default=0)
     parser.add_argument("--cosine_lr_scheduler", action="store_true")
+    parser.add_argument("--optim_cpu_offload", action="store_true")
 
     parser.add_argument("--project")
     parser.add_argument("--run_name", default="debug")
@@ -177,6 +178,8 @@ def evaluate_model(model, args):
     print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
 
     optim = OPTIM_MAP[args.optim](model.parameters(), args.lr, weight_decay=args.weight_decay)
+    if args.optim_cpu_offload:
+        optim = low_bit_optim.CPUOffloadOptimizer(optim)
     lr_schedule = CosineSchedule(args.lr, len(dloader) * args.n_epochs)
 
     grad_scaler = torch.amp.GradScaler("cuda", enabled=args.amp == "fp16")

torchao/prototype/low_bit_optim/__init__.py

@@ -1,2 +1,3 @@
 from .adam import Adam8bit, Adam4bit, AdamFp8
 from .adamw import AdamW8bit, AdamW4bit, AdamWFp8
+from .cpu_offload import CPUOffloadOptimizer

torchao/prototype/low_bit_optim/cpu_offload.py

@@ -0,0 +1,34 @@
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class CPUOffloadOptimizer:
+    def __init__(self, base_optimizer: Optimizer) -> None:
+        self.optim = base_optimizer
+        self.param_cpu2cuda_map = dict()
+
+        # swap param in param_groups with CPU param
+        for param_group in base_optimizer.param_groups:
+            for i, p in enumerate(param_group["params"]):
+                p_cpu = p.detach().cpu().pin_memory()
+                param_group["params"][i] = p_cpu
+                self.param_cpu2cuda_map[p_cpu] = p
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        # copy gradients from CUDA to CPU
+        for p_cpu, p_cuda in self.param_cpu2cuda_map.items():
+            if p_cuda.grad is not None:
+                p_cpu.grad = p_cuda.grad.to("cpu", non_blocking=True)
+                p_cuda.grad = None
+        torch.cuda.synchronize()
+
+        self.optim.step(closure)
+
+        # copy updated param from CPU to CUDA
+        for p_cpu, p_cuda in self.param_cpu2cuda_map.items():
+            p_cuda.copy_(p_cpu, non_blocking=True)
+
+    # redirect calls to base optimizer
+    def __getattr__(self, name: str):
+        return getattr(self.optim, name)