Fix LAMB norm reduction with ZeRO CPU offload

colossalai/nn/optimizer/distributed_lamb.py

@@ -158,8 +158,13 @@ def step(self, closure=None):
                     g_sum = (update**2).sum()
                     if self.dp_size > 1 and self.is_zero:
                         # ZeRO 2 doesn't shard param. Compute full param norm w/o communication.
-                        dist.all_reduce(g_sum, group=self.dp_group)
+                        # The master parameter and its update can be offloaded to CPU, while
+                        # the working parameter remains on the accelerator.  NCCL cannot
+                        # reduce a CPU tensor, so use the working parameter's device for the
+                        # norm reduction as well as the following stacked TP reduction.
                         p_local = self.shard_to_working_param[id(p)]
+                        g_sum = g_sum.to(p_local.device)
+                        dist.all_reduce(g_sum, group=self.dp_group)
 
                     w_sum = (p_local**2).sum()
                     sums = torch.stack([w_sum, g_sum])

tests/test_optimizer/test_dist_lamb.py

@@ -254,6 +254,35 @@ def run_dist_lamb_fwd_bwd(
         raise e
 
 
+@clear_cache_before_run()
+def run_dist_lamb_cpu_offload() -> None:
+    """Regression test for LAMB norm reductions with CPU-offloaded master parameters."""
+    rank = dist.get_rank()
+    seed_all(_SEED)
+    model = Net(_IN_DIM, _HID_DIM).to(rank).half()
+    optim = DistributedLamb(model.parameters(), lr=1e-3)
+    optim = LowLevelZeroOptimizer(
+        optim,
+        cpu_offload=True,
+        initial_scale=2**20,
+        partition_grad=True,
+        verbose=True,
+    )
+    optim.optim.setup_distributed(
+        dp_group=dist.group.WORLD,
+        shard_to_working_param=optim.get_master_to_working_map(),
+        is_zero=True,
+    )
+
+    x = data_gen().cuda().half()
+    output = model(x)
+    optim.backward(output.float().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)
+
+
 def check_dist_lamb(rank, world_size, port):
     disable_existing_loggers()
     colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
@@ -266,6 +295,9 @@ def check_dist_lamb(rank, world_size, port):
     run_dist_lamb_fwd_bwd()
     coordinator.print_on_master("Forward-backward tests passed")
 
+    run_dist_lamb_cpu_offload()
+    coordinator.print_on_master("CPU-offload test passed")
+
     run_bert_test(optim_class=Lamb, sharded_optim_class=Lamb)
     print(f"rank {rank} tests passed :)")