Optimizer CPU offload for single GPU training

benchmarks/benchmark_low_bit_adam.py

@@ -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.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)