[PyTorch] Advance FusedAdam step counter for empty param groups

tests/pytorch/test_fused_optimizer.py

@@ -166,6 +166,31 @@ def test_frozen_model(self):
 
             torch.testing.assert_close(ref_param, tst_param)
 
+    @pytest.mark.parametrize("capturable", [False, True])
+    def test_empty_param_group_advances_step(self, capturable):
+        # An empty param group must advance its step counter like a populated one.
+        # The same group can be empty on one data-parallel rank and populated on
+        # another, and "step" is part of the checkpoint, so a group that stops
+        # counting makes a resumed run apply a stale bias correction.
+        param = torch.nn.Parameter(torch.rand(4, dtype=torch.float, device="cuda"))
+        tst_optim = self.fused_optim(
+            [{"params": [param]}, {"params": []}], capturable=capturable, **self.options
+        )
+
+        num_steps = 3
+        for _ in range(num_steps):
+            param.grad = torch.rand_like(param)
+            tst_optim.step()
+
+        populated_step, empty_step = (int(g["step"]) for g in tst_optim.param_groups)
+        assert populated_step == num_steps
+        assert empty_step == populated_step
+
+        # The counter is checkpointed through the param groups, so the empty group
+        # has to round trip with the same value as the populated one.
+        checkpoint = tst_optim.state_dict()
+        assert int(checkpoint["param_groups"][1]["step"]) == num_steps
+
     def test_empty_param_at_end_of_group(self):
         tensors = [
             torch.ones(4, dtype=torch.float, device="cuda"),

transformer_engine/pytorch/optimizers/fused_adam.py

@@ -557,23 +557,37 @@ def step(self, closure=None, grad_scaler=None):
             loss = closure()
 
         for group in self.param_groups:
-            if len(group["params"]) == 0:
-                continue
-            device = group["params"][0].device
-            bias_correction = 1 if group["bias_correction"] else 0
-            beta1, beta2 = group["betas"]
-
+            # Advance the step counter before skipping empty groups. A param group can be
+            # empty on some data-parallel ranks and populated on others, so incrementing
+            # only for populated groups desynchronizes "step" across the ranks that share
+            # an optimizer state shard. A rank then resumes from a checkpoint written by a
+            # rank where the group was empty and applies a stale bias correction.
             # assume same step across group now to simplify things
             # per parameter step can be easily support by making it tensor, or pass list into kernel
             if "step" in group:
                 group["step"] += (
                     1 if not self.capturable else (self._dummy_overflow_buf != 1).to(torch.int)
                 )
             else:
+                # Empty groups have no parameter to take the device from, so fall back to
+                # the device of the optimizer's own scratch buffer.
+                step_device = (
+                    group["params"][0].device
+                    if len(group["params"]) > 0
+                    else self._dummy_overflow_buf.device
+                )
                 group["step"] = (
-                    1 if not self.capturable else torch.tensor([1], dtype=torch.int, device=device)
+                    1
+                    if not self.capturable
+                    else torch.tensor([1], dtype=torch.int, device=step_device)
                 )
 
+            if len(group["params"]) == 0:
+                continue
+            device = group["params"][0].device
+            bias_correction = 1 if group["bias_correction"] else 0
+            beta1, beta2 = group["betas"]
+
             # create lists for multi-tensor apply
             p_main_of_fp8_model = []
             p_main_of_f16_model = []