x-evolution 0.1.25__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
@@ -93,11 +94,17 @@ class EvoStrategy(Module):
         verbose = True,
         accelerator: Accelerator | None = None,
         accelerate_kwargs: dict = dict(),
-        reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None
+        reject_generation_fitnesses_if: Callable[[Tensor], bool] | None = None,
+        vectorized = False,
+        vector_size: int | None = None,
+        sync_on_init = True
     ):
         super().__init__()
         self.verbose = verbose
 
+        self.vectorized = vectorized
+        self.vector_size = vector_size
+
         if not exists(accelerator):
             accelerator = Accelerator(cpu = cpu, **accelerate_kwargs)
 
@@ -122,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
-
-        wrapped_model = accelerator.prepare(model)
+        # maybe sync model params and buffers
 
-        with torch.no_grad():
-            environment(wrapped_model)
+        if sync_on_init:
+            self.sync_model_params_and_buffers_()
 
         # get param dictionary
 
@@ -249,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
@@ -475,24 +491,28 @@ class EvoStrategy(Module):
                     fitnesses.append([0., 0.] if self.mirror_sampling else 0.)
                     continue
 
-                individual_param_seeds = with_seed(individual_seed)(randint)(0, MAX_SEED_VALUE, (self.num_params,))
-
-                noise_config = dict(zip(self.param_names_to_optimize, individual_param_seeds.tolist()))
-
-                # determine noise scale, which can be fixed or learned
+                def get_fitness(negate = False):
+                    individual_param_seeds = with_seed(individual_seed.item())(randint)(0, MAX_SEED_VALUE, (self.num_params,))
+                    noise_config = dict(zip(self.param_names_to_optimize, individual_param_seeds.tolist()))
 
-                noise_config_with_scale = dict()
+                    noise_config_with_scale = dict()
+                    for param_name, seed in noise_config.items():
+                        noise_scale = self._get_noise_scale(param_name)
+                        noise_config_with_scale[param_name] = (seed, noise_scale)
 
-                for param_name, seed in noise_config.items():
+                    with model.temp_add_noise_(noise_config_with_scale, negate = negate):
+                        fitness = with_seed(maybe_rollout_seed)(self.environment)(model)
 
-                    noise_scale = self._get_noise_scale(param_name)
+                    if isinstance(fitness, Tensor) and fitness.numel() > 1:
+                        fitness = fitness.mean().item()
+                    elif isinstance(fitness, Tensor):
+                        fitness = fitness.item()
 
-                    noise_config_with_scale[param_name] = (seed, noise_scale)
+                    return fitness
 
-                # maybe roll out with a fixed seed
+                # evaluate
 
-                with model.temp_add_noise_(noise_config_with_scale):
-                    fitness = with_seed(maybe_rollout_seed)(rollout_for_fitness)()
+                fitness = get_fitness(negate = False)
 
                 if not self.mirror_sampling:
                     fitnesses.append(fitness)
@@ -500,8 +520,7 @@ class EvoStrategy(Module):
 
                 # handle mirror sampling
 
-                with model.temp_add_noise_(noise_config_with_scale, negate = True):
-                    fitness_mirrored = with_seed(maybe_rollout_seed)(rollout_for_fitness)()
+                fitness_mirrored = get_fitness(negate = True)
 
                 fitnesses.append([fitness, fitness_mirrored])
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.25
+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
@@ -38,7 +38,7 @@ Requires-Dist: accelerate
 Requires-Dist: beartype
 Requires-Dist: einops>=0.8.0
 Requires-Dist: torch>=2.4
-Requires-Dist: x-mlps-pytorch>=0.1.31
+Requires-Dist: x-mlps-pytorch>=0.2.0
 Requires-Dist: x-transformers>=2.11.23
 Provides-Extra: examples
 Requires-Dist: gymnasium[box2d]>=1.0.0; extra == 'examples'

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.25.dist-info/RECORD

@@ -1,6 +0,0 @@
-x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
-x_evolution/x_evolution.py,sha256=lvN3ePqD6a5dW1gOv0d1I9yQ4rdv6OuIVvKvXa0yRBM,19126
-x_evolution-0.1.25.dist-info/METADATA,sha256=1fAtssjj_t76rXwLu728z_ohlRZKQDhCm0oOl2eeIxA,5854
-x_evolution-0.1.25.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-x_evolution-0.1.25.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-x_evolution-0.1.25.dist-info/RECORD,,