x-evolution 0.1.17__py3-none-any.whl → 0.1.22__py3-none-any.whl

x_evolution/x_evolution.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from functools import partial
 
 import torch
+import torch.distributed as dist
 from torch import tensor, Tensor, stack, is_tensor, arange, randint
 from torch.nn import Module, ModuleList, Parameter, ParameterList
 from torch.optim import SGD, Adam, Optimizer
@@ -96,19 +97,38 @@ class EvoStrategy(Module):
         reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None
     ):
         super().__init__()
+        self.verbose = verbose
 
         if not exists(accelerator):
             accelerator = Accelerator(cpu = cpu, **accelerate_kwargs)
 
         self.accelerate = accelerator
 
+        # environment - with optional init
+
+        self.environment = environment
+
+        if accelerator.is_main_process:
+            if hasattr(environment, 'pre_main_callback') and callable(environment.pre_main_callback):
+                self.print('pre_main_callback detected on environment passed in and is invoked')
+                environment.pre_main_callback()
+
+        accelerator.wait_for_everyone()
+
+        # take care of model and parameters
+
         if isinstance(model, list):
             model = ModuleList(model)
 
         self.model = model
         self.noisable_model = Noisable(model, low_rank = noise_low_rank)
 
-        self.environment = environment
+        # use prepare and run through environment once to sync params
+
+        wrapped_model = accelerator.prepare(model)
+        environment(wrapped_model)
+
+        # get param dictionary
 
         named_parameters_dict = dict(model.named_parameters())
 
@@ -217,10 +237,6 @@ class EvoStrategy(Module):
 
         self.reject_generation_fitnesses_if = reject_generation_fitnesses_if
 
-        # verbose
-
-        self.verbose = verbose
-
         # checkpointing
 
         self.checkpoint_every = checkpoint_every

{x_evolution-0.1.17.dist-info → x_evolution-0.1.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.17
+Version: 0.1.22
 Summary: x-evolution
 Project-URL: Homepage, https://pypi.org/project/x-evolution/
 Project-URL: Repository, https://github.com/lucidrains/x-evolution
@@ -153,3 +153,4 @@ $ accelerate launch train.py
 }
 ```
 
+*Nothing makes sense except in the light of evolution* - Theodosius Dobzhansky

x_evolution-0.1.22.dist-info/RECORD

@@ -0,0 +1,6 @@
+x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
+x_evolution/x_evolution.py,sha256=ui4HgIcJ_qM0JtL-HzsiwWNF6r5Ybo9c2DG0mWTXo4w,19124
+x_evolution-0.1.22.dist-info/METADATA,sha256=FrF0ebgJVX2NrZkjgcH4QHozWfXMhGtTGMpXlRnP-cE,5728
+x_evolution-0.1.22.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+x_evolution-0.1.22.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+x_evolution-0.1.22.dist-info/RECORD,,

x_evolution-0.1.17.dist-info/RECORD

@@ -1,6 +0,0 @@
-x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
-x_evolution/x_evolution.py,sha256=UaDdGLvd2xURQp2SnBjLZkdyOPQ41Y55djo00jnKCNE,18503
-x_evolution-0.1.17.dist-info/METADATA,sha256=KSqDsQ5whxAfgPE6is0xQZTDPAPMJFKSCn6lSGd1Ork,5649
-x_evolution-0.1.17.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-x_evolution-0.1.17.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-x_evolution-0.1.17.dist-info/RECORD,,