Edenzzzz's fix for min_8bit_size functionality in Optimizer base clas…

…ses (#1286) * fix min_8bit_size invalid bug * Apply same fix to other optimizer base class --------- Co-authored-by: Edenzzzz <[email protected]>
bitsandbytes-foundation · Jul 22, 2024 · 5212a0f · 5212a0f
1 parent 0bdd57c
commit 5212a0f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/bitsandbytes/optim/optimizer.py b/bitsandbytes/optim/optimizer.py
@@ -437,7 +437,7 @@ def init_state(self, group, p, gindex, pindex):
         state = self.state[p]
         state["step"] = 0
 
-        if dtype == torch.float32 or (dtype == torch.uint8 and p.numel() < 4096):
+        if dtype == torch.float32:
             state["state1"] = self.get_state_buffer(p, dtype=torch.float32)
             state["state2"] = self.get_state_buffer(p, dtype=torch.float32)
         elif dtype == torch.uint8:
@@ -656,7 +656,7 @@ def init_state(self, group, p, gindex, pindex):
         state = self.state[p]
         state["step"] = 0
 
-        if dtype == torch.float32 or (dtype == torch.uint8 and p.numel() < 4096):
+        if dtype == torch.float32:
             state["state1"] = self.get_state_buffer(p, dtype=torch.float32)
         elif dtype == torch.uint8:
             if state["step"] == 0: