fix: functional-closure passthrough works for per-candidate optimizers

DE/GA evaluate one candidate at a time by calling closure() directly as a
zero-argument function. The vmap passthrough previously set _current_closure
to the functional closure (which requires a flat_params argument), so DE
crashed with 'missing 1 required positional argument' under
set_functional_closure. Now set_functional_closure also stores a zero-arg
plain equivalent (loss_fn(model), evaluating the model's current weights),
and step() installs that as the passthrough. Population-synchronous
optimizers (PSO/GWO/WOA/CMA-ES) still take the batched vmap path inside
_evaluate_fitness; per-candidate optimizers (DE/GA) now run correctly on the
sequential path. All 5 algorithms verified with the functional closure.

gpu_nn: construct optimizers by signature introspection (swarm_size vs
population_size) and wrap each cell in try/except so one failure can't kill
the sweep or lose the report.

Author: hallelx2

swarmtorch/base/swarm_optimizer.py

@@ -181,8 +181,18 @@ def loss_fn(forward):
         if model is None or loss_fn is None:
             self._functional_closure = None
             self._functional_model = None
+            self._plain_from_functional = None
             return
 
+        # A zero-argument plain closure equivalent: evaluate the model in its
+        # *current* parameter state. Optimizers that evaluate one candidate at
+        # a time (e.g. DE, GA) call ``closure()`` directly after writing a
+        # candidate into the model; they use this, while population-synchronous
+        # optimizers (PSO, GWO, ...) take the batched vmap path inside
+        # ``_evaluate_fitness``. Passing the live ``model`` as the ``forward``
+        # callable makes ``loss_fn(model)`` evaluate the current weights.
+        self._plain_from_functional = lambda: loss_fn(model)
+
         from torch.func import functional_call
 
         param_specs: list[tuple[str, torch.Size, int]] = [
@@ -227,15 +237,15 @@ def step(self, closure: Any = None) -> Any:
             here — ``_update_positions`` evaluates it once per particle.
         """
         # Stash the closure so subclasses can pull it from
-        # ``self._current_closure`` inside ``_update_positions``.
-        # If the user has registered a functional closure (vmap fast path)
-        # but didn't supply a plain closure, install a passthrough
-        # sentinel so the existing ``if closure is None: return`` guards
-        # in subclass ``_update_positions`` don't short-circuit. The
-        # actual fitness evaluation still goes through the vmap path
-        # inside ``_evaluate_fitness``.
+        # ``self._current_closure`` inside ``_update_positions``. If the user
+        # registered a functional closure (vmap fast path) but didn't supply a
+        # plain closure, install a *zero-argument* plain equivalent so that
+        # (a) the ``if closure is None: return`` guards don't short-circuit and
+        # (b) optimizers that call ``closure()`` directly per candidate (DE,
+        # GA) still work. Population-synchronous optimizers ignore this and take
+        # the batched vmap path inside ``_evaluate_fitness``.
         if closure is None and getattr(self, "_functional_closure", None) is not None:
-            self._current_closure = self._functional_closure
+            self._current_closure = self._plain_from_functional
         else:
             self._current_closure = closure
 

swarmtorch/benchmark/gpu_nn.py

@@ -70,11 +70,21 @@ def _run_one(algo_cls, hidden, swarm_size, max_fe, seed, device, use_vmap):
     x = torch.randn(batch, in_dim, device=device)
     y = torch.randint(0, out_dim, (batch,), device=device)
 
-    try:
-        opt = algo_cls(model.parameters(), swarm_size=swarm_size, device=device,
-                       init_strategy="model")
-    except TypeError:
-        opt = algo_cls(model.parameters(), swarm_size=swarm_size, device=device)
+    # Build kwargs by introspection: some optimizers name the population
+    # ``swarm_size`` (PSO/GWO/WOA/CMAES) and others ``population_size`` (DE/GA).
+    import inspect
+
+    sig = inspect.signature(algo_cls.__init__)
+    params = sig.parameters
+    accepts_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
+    kwargs: dict[str, Any] = {"device": device}
+    if "swarm_size" in params or accepts_kw:
+        kwargs["swarm_size"] = swarm_size
+    elif "population_size" in params:
+        kwargs["population_size"] = swarm_size
+    if "init_strategy" in params or accepts_kw:
+        kwargs["init_strategy"] = "model"
+    opt = algo_cls(model.parameters(), **kwargs)
 
     if use_vmap:
         opt.set_functional_closure(model, lambda fwd: F.cross_entropy(fwd(x), y))
@@ -130,10 +140,15 @@ def run_nn_speedup(
                 for seed in seeds:
                     for vname, dev, use_vmap in variants:
                         cell += 1
-                        # warm-up (untimed)
-                        _run_one(ALGOS[algo], h, ss, ss, seed, dev, use_vmap)
-                        wall, best, npar, fe = _run_one(
-                            ALGOS[algo], h, ss, max_fe, seed, dev, use_vmap)
+                        try:
+                            # warm-up (untimed)
+                            _run_one(ALGOS[algo], h, ss, ss, seed, dev, use_vmap)
+                            wall, best, npar, fe = _run_one(
+                                ALGOS[algo], h, ss, max_fe, seed, dev, use_vmap)
+                        except Exception as e:  # one bad cell must not kill the sweep
+                            print(f"[gpu_nn] {cell}/{total} {algo} {vname} FAILED: {e}",
+                                  flush=True)
+                            continue
                         rec = NNBenchResult(
                             algorithm=algo, n_params=npar, swarm_size=ss,
                             variant=vname, device=dev, seed=seed,