Update paged AdEMAMix

matthewdouglas · matthewdouglas · commit 1a3d3c6dea88 · 2024-09-20T14:54:32.000-04:00
diff --git a/bitsandbytes/optim/ademamix.py b/bitsandbytes/optim/ademamix.py
@@ -260,7 +260,7 @@ def update_step(self, group, p, gindex, pindex):
             )
 
     def _get_state_double_buffer(self, p, dtype=torch.float32):
-        if not self.is_paged or p.numel() < 0.5e5:
+        if not self.is_paged or p.numel() < 1e5:
             return torch.zeros((2, *p.size()), dtype=dtype, device=p.device)
         else:
             buff = F.get_paged(*(2, *p.size()), dtype=dtype, device=p.device)

Original file line number	Diff line number	Diff line change
`@@ -260,7 +260,7 @@ def update_step(self, group, p, gindex, pindex):`
`260`	`260`	`)`
`261`	`261`
`262`	`262`	`def _get_state_double_buffer(self, p, dtype=torch.float32):`
`263`		`- if not self.is_paged or p.numel() < 0.5e5:`
	`263`	`+ if not self.is_paged or p.numel() < 1e5:`
`264`	`264`	`return torch.zeros((2, *p.size()), dtype=dtype, device=p.device)`
`265`	`265`	`else:`
`266`	`266`	`buff = F.get_paged((2, p.size()), dtype=dtype, device=p.device)`