x-evolution 0.1.26__py3-none-any.whl → 0.1.27__py3-none-any.whl

x_evolution/x_evolution.py

@@ -11,6 +11,7 @@ from torch.nn import Module, ModuleList, Parameter, ParameterList
 from torch.optim import SGD, Adam, Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
+import torch.distributed as dist
 import torch.nn.functional as F
 
 from beartype import beartype
@@ -95,7 +96,8 @@ class EvoStrategy(Module):
         accelerate_kwargs: dict = dict(),
         reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None,
         vectorized = False,
-        vector_size: int | None = None
+        vector_size: int | None = None,
+        sync_on_init = True
     ):
         super().__init__()
         self.verbose = verbose
@@ -127,12 +129,10 @@ class EvoStrategy(Module):
         self.model = model
         self.noisable_model = Noisable(model, low_rank = noise_low_rank)
 
-        # use prepare and run through environment once to sync params
+        # maybe sync model params and buffers
 
-        wrapped_model = accelerator.prepare(model)
-
-        with torch.no_grad():
-            environment(wrapped_model)
+        if sync_on_init:
+            self.sync_model_params_and_buffers_()
 
         # get param dictionary
 
@@ -254,6 +254,17 @@ class EvoStrategy(Module):
     def device(self):
         return self.accelerate.device
 
+    @torch.no_grad()
+    def sync_model_params_and_buffers_(self):
+        if not self.accelerate.num_processes > 1:
+            return
+
+        for param in self.model.parameters():
+            dist.broadcast(param, src = 0)
+
+        for buffer in self.model.buffers():
+            dist.broadcast(buffer, src = 0)
+
     def print(self, *args, **kwargs):
         if not self.verbose:
             return

{x_evolution-0.1.26.dist-info → x_evolution-0.1.27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.26
+Version: 0.1.27
 Summary: x-evolution
 Project-URL: Homepage, https://pypi.org/project/x-evolution/
 Project-URL: Repository, https://github.com/lucidrains/x-evolution

x_evolution-0.1.27.dist-info/RECORD

@@ -0,0 +1,6 @@
+x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
+x_evolution/x_evolution.py,sha256=-G5qXGMjwVzdkxIDR6xL_YGium4KfKC0cnlY76Upy0o,19799
+x_evolution-0.1.27.dist-info/METADATA,sha256=hj0MUpIGVWoOY5wHsoy_ZF_cx7s48_HZicd4IgNUFEo,5853
+x_evolution-0.1.27.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+x_evolution-0.1.27.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+x_evolution-0.1.27.dist-info/RECORD,,

x_evolution-0.1.26.dist-info/RECORD

@@ -1,6 +0,0 @@
-x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
-x_evolution/x_evolution.py,sha256=Jln3wpkIQIp7xwa3KMibh0kSuob1NIi7Aj7Miz8RJdY,19491
-x_evolution-0.1.26.dist-info/METADATA,sha256=7zamSGrDtvOUQzpyYCkrt42FmWC4GPYpVSWUAO8a6OA,5853
-x_evolution-0.1.26.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-x_evolution-0.1.26.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-x_evolution-0.1.26.dist-info/RECORD,,