Add variance calculation from FusedAdam optimizer states

tests/pytorch/test_fused_optimizer.py

@@ -570,7 +570,7 @@ def forward(self, x):
         return y
 
 
-class AdamTest:
+class TestAdam:
 
     def setup_method(self, *, seed: int = 0) -> None:
         torch.manual_seed(seed)
@@ -792,6 +792,34 @@ def test_native(self):
 
             self.model_.load_state_dict(copy.deepcopy(self.model.state_dict()))
 
+    def test_optimizer_get_moments_and_variance(self):
+        params_ = [p for p in self.model_.parameters() if p.requires_grad]
+        optimizer_ = te.optimizers.FusedAdam(params_, lr=self.lr, capturable=False)
+
+        for i in range(100):
+            x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last)
+            x_ = x.clone()
+            gt = torch.rand([32, 10]).cuda()
+            gt_ = gt.clone()
+
+            # DUT
+            y = self.model_(x)
+            loss_ = ((gt_ - y) ** 2).mean()
+
+            loss_.backward()
+            optimizer_.step()
+
+            # Init for next iteration
+            optimizer_.zero_grad()
+        for param in [p for p in self.model_.parameters() if p.requires_grad]:
+            first_moment = optimizer_.get_unscaled_state(param, "exp_avg")
+            second_moment = optimizer_.get_unscaled_state(param, "exp_avg_sq")
+            assert first_moment is not None
+            assert second_moment is not None
+            expected_variance = second_moment - torch.square(first_moment)
+            variance = optimizer_.get_grad_variance_from_state(param)
+            torch.testing.assert_close(variance, expected_variance)
+
     @largeTensorTest("60GB", "cuda")
     def test_large_tensor(self):
         t = torch.zeros(2359332864, dtype=torch.half, device="cuda")

transformer_engine/pytorch/optimizers/fused_adam.py

@@ -358,6 +358,19 @@ def set_scaled_state(self, param, state_name, unscaled_state):
         else:
             state[state_name].copy_(unscaled_state)
 
+    def get_grad_variance_from_state(self, param):
+        """Return the unscaled state corresponding to the input `param` and `state_name`.
+
+        Arguments:
+            param (torch.nn.Parameter): One of parameters in this optimizer.
+            state_name (string): Name of optimizer states, can be one of 'exp_avg', 'exp_avg_sq',
+                and 'master_param`.
+        """
+        first_moment = self.get_unscaled_state(param, "exp_avg")
+        second_moment = self.get_unscaled_state(param, "exp_avg_sq")
+        variance = second_moment - torch.square(first_moment)
+        return variance
+
     def _initialize_state(
         self, param, state_name, zero_buffer: bool, store_param_remainders: bool = False
     ):