Add unittests for AnyPrecisionOptimizer (#62)

* AnyPrecision unittests

Author: rohan-varma

src/python/torchdistx/optimizers/__init__.py

@@ -0,0 +1 @@
+from .anyprecision_optimizer import AnyPrecisionAdamW

src/python/torchdistx/optimizers/anyprecision_optimizer.py

@@ -46,7 +46,8 @@ def __init__(
             momentum_dtype = dtype for momentum  (default: BFloat32)
             variance_dtype = dtype for uncentered variance (default: BFloat16)
             compensation_buffer_dtype  = dtype for Kahan summation
-                                         buffer (default: BFloat16)
+                                         buffer (default: BFloat16). Only used if
+                                         ``use_kahan_summation=True``.
 
             # Usage
             This optimizer implements optimizer states, and Kahan summation

tests/python/test_anyprecision_optimizer.py

@@ -0,0 +1,83 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from copy import deepcopy
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.testing._internal.common_utils import (
+    TestCase,
+    instantiate_parametrized_tests,
+    parametrize,
+    run_tests,
+)
+
+from torchdistx.optimizers import AnyPrecisionAdamW
+
+
+class TestAnyPrecisionOptimizer(TestCase):
+    def _test_adam_equivalence(self, model, model_clone):
+        # Test non-default options
+        betas = (0.8, 0.88)
+        weight_decay = 0.03
+
+        adam_opt = optim.AdamW(
+            model_clone.parameters(), betas=betas, weight_decay=weight_decay
+        )
+        anyprecision_adam = AnyPrecisionAdamW(
+            model.parameters(),
+            variance_dtype=torch.float32,
+            betas=betas,
+            weight_decay=weight_decay,
+        )
+
+        # Verify params are equal initially
+        model_orig_params = [p.clone() for p in model.parameters()]
+        for p1, p2 in zip(model_clone.parameters(), model_orig_params):
+            self.assertEqual(p1, p2)
+
+        for i in range(6):
+            adam_opt.zero_grad()
+            anyprecision_adam.zero_grad()
+            inp = torch.randn(5, 5, device=next(model.parameters()).device)
+            model(inp).sum().backward()
+            model_clone(inp).sum().backward()
+            adam_opt.step()
+            anyprecision_adam.step()
+
+            # Ensure params are modified from original
+            if i == 0:
+                for p1, p2 in zip(model.parameters(), model_orig_params):
+                    self.assertNotEqual(p1, p2)
+
+            for p1, p2 in zip(model.parameters(), model_clone.parameters()):
+                self.assertEqual(p1, p2)
+
+    @parametrize("device", ["cpu", "cuda"])
+    def test_adam_equivalence(self, device):
+        """
+        Tests that AnyPrecisionAdamW is equivalent to AdamW when
+        kahan summation and different dtypes for momentum, variance,
+        and compensation buffer are turned off (i.e. all float32).
+        """
+        if device == "cuda" and not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA not available")
+
+        model = nn.Sequential(nn.Linear(5, 5), nn.Linear(5, 5), nn.Linear(5, 5))
+        if device == "cuda":
+            model.cuda()
+
+        model_clone = deepcopy(model)
+
+        self._test_adam_equivalence(model, model_clone)
+
+
+instantiate_parametrized_tests(TestAnyPrecisionOptimizer)
+
+if __name__ == "__main__":
+    run_tests()