Fix for weights-only load (pytorch#1228)

stack-info: PR: pytorch#1228, branch: drisspg/stack/19
nighting0le01 · Nov 7, 2024 · 27b1d16 · 27b1d16
1 parent 668a4ba
commit 27b1d16
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/test/prototype/test_low_bit_optim.py b/test/prototype/test_low_bit_optim.py
@@ -345,8 +345,9 @@ def test_optim_bf16_stochastic_round_correctness(self):
             optim2.step()
             optim2.zero_grad()
 
-            torch.testing.assert_close(loss1, loss2, msg=lambda msg: f"Iteration {idx}. {msg}")
-
+            torch.testing.assert_close(
+                loss1, loss2, msg=lambda msg: f"Iteration {idx}. {msg}"
+            )
 
     @pytest.mark.skipif(not TORCH_VERSION_AT_LEAST_2_3, reason="requires PyTorch >= 2.3")
     @parametrize("optim_name", ["Adam8bit", "AdamW8bit", "Adam4bit", "AdamW4bit", "AdamFp8", "AdamWFp8"])

diff --git a/torchao/prototype/low_bit_optim/adam.py b/torchao/prototype/low_bit_optim/adam.py
@@ -13,7 +13,17 @@
 
 class _AdamBase(Optimizer):
     def __init__(
-        self, params, lr, betas, eps, weight_decay, amsgrad, *, block_size, bf16_stochastic_round, is_adamw, exclude_low_bit_optim_params=None
+        self,
+        params,
+        lr,
+        betas,
+        eps,
+        weight_decay,
+        amsgrad,
+        *,
+        block_size,
+        bf16_stochastic_round,
+        is_adamw,
     ) -> None:
         if not 0.0 <= lr:
             raise ValueError("Invalid learning rate: {}".format(lr))