Fix GaLore 8-bit updates with ZeRO CPU offload

colossalai/nn/optimizer/distributed_galore.py

@@ -9,6 +9,7 @@
 import torch.nn.functional as F
 from bitsandbytes.optim.optimizer import Optimizer2State
 
+from colossalai.accelerator import get_accelerator
 from colossalai.interface.optimizer import DistributedOptim
 from colossalai.tensor.d_tensor import get_shard_dim_1d, is_distributed_tensor
 
@@ -163,6 +164,17 @@ def step(self, closure=None):
                     continue
                 state = self.state[p]
 
+                # bitsandbytes only provides CUDA update kernels.  ZeRO CPU
+                # offload keeps the master shard, gradient, and optimizer state
+                # on CPU, so stage one parameter at a time on the accelerator
+                # for the update and move it back afterwards.
+                cpu_offload = p.device.type == "cpu"
+                if cpu_offload:
+                    compute_device = get_accelerator().get_current_device()
+                    p.data = p.data.to(compute_device)
+                    p.grad = p.grad.to(compute_device)
+                    self._move_state(state, compute_device)
+
                 if "step" not in state:
                     state["step"] = 0
 
@@ -259,12 +271,34 @@ def step(self, closure=None):
                     group["weight_decay"] = group["weight_decay_saved"]
                     del group["weight_decay_saved"]
 
+                if cpu_offload:
+                    p.data = p.data.cpu()
+                    p.grad = p.grad.cpu()
+                    self._move_state(state, torch.device("cpu"))
+                    if hasattr(p, "saved_data"):
+                        del p.saved_data
+
         if self.is_paged:
             # all paged operation are asynchronous, we need
             # to sync to make sure all tensors are in the right state
             torch.cuda.synchronize()
         return loss
 
+    @staticmethod
+    def _move_state(state, device) -> None:
+        """Move non-paged bitsandbytes and GaLore projector state to ``device``."""
+        for key, value in state.items():
+            if isinstance(value, torch.Tensor) and not getattr(value, "is_paged", False):
+                state[key] = value.to(device)
+
+        projector = state.get("projector")
+        if projector is not None:
+            ortho_matrix = projector.ortho_matrix
+            if isinstance(ortho_matrix, torch.Tensor):
+                projector.ortho_matrix = ortho_matrix.to(device)
+            elif isinstance(ortho_matrix, list):
+                projector.ortho_matrix = [matrix.to(device) for matrix in ortho_matrix]
+
     def to_master_shape(self, data, padding):
         """Pad to master (optimizer) param shape"""
         if not self.is_zero:

tests/test_optimizer/test_dist_galore.py

@@ -266,6 +266,41 @@ def run_dist_galore_fwd_bwd(p_g_dtype: tuple[torch.dtype, torch.dtype], tp_zero_
         raise e
 
 
+def run_dist_galore_cpu_offload() -> None:
+    """Regression test for bitsandbytes updates on CPU-offloaded master shards."""
+    rank = dist.get_rank()
+    seed_all(_SEED)
+    model = Net(_IN_DIM, _HID_DIM).to(rank)
+    optim = DistGaloreAwamW(model.parameters(), lr=lr)
+    optim = LowLevelZeroOptimizer(
+        optim,
+        cpu_offload=True,
+        partition_grad=False,
+        initial_scale=128,
+    )
+    optim.optim.setup_distributed(
+        dp_group=dist.group.WORLD,
+        shard_to_working_param=optim.get_master_to_working_map(),
+        padding_map=optim.get_param_padding_map(),
+        is_zero=False,
+    )
+
+    for _ in range(2):
+        optim.zero_grad()
+        x = data_gen().cuda()
+        output = model(x)
+        optim.backward(output.square().mean())
+        optim.step()
+
+    for master_params in optim._master_param_groups_of_current_rank.values():
+        assert all(param.device.type == "cpu" for param in master_params)
+        for param in master_params:
+            assert all(
+                not isinstance(value, torch.Tensor) or value.device.type == "cpu"
+                for value in optim.optim.state[param].values()
+            )
+
+
 def check_dist_galore(rank, world_size, port):
     disable_existing_loggers()
     colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
@@ -279,6 +314,9 @@ def check_dist_galore(rank, world_size, port):
     # run_dist_galore_fwd_bwd()
     # _COORDINATOR.print_on_master("Forward-backward tests passed")
 
+    run_dist_galore_cpu_offload()
+    coordinator.print_on_master("CPU-offload test passed")
+
     coordinator.print_on_master(
         "Running bert tests, which are expected to produce minor errors due to instability in SVD convergence. \
             For example, a 1e-9 grad diff causes drastic difference in SVD output."