Fix ZeroFlow sparse RNG to follow parameter device

src/infty/optim/geometry_reshaping/sam.py

@@ -19,7 +19,6 @@ def __init__(self, params, base_optimizer, model, args, adaptive=False, perturb_
         self.rho_scheduler = args_ns.rho_scheduler
 
         self.update_rho_t()
-        self.get_grad_reduce(grad_reduce)
 
 
 

src/infty/optim/zeroth_order_updates/zeroflow.py

@@ -8,7 +8,6 @@
 from types import SimpleNamespace
 
 from torch.optim import Optimizer
-from types import SimpleNamespace
 
 class ZeroFlow(Optimizer):
     def __init__(self, params, base_optimizer, model, args, **kwargs):
@@ -36,8 +35,11 @@ def __init__(self, params, base_optimizer, model, args, **kwargs):
         self.gradient_sparsity = args_ns.gradient_sparsity
         self.memory_efficient = args_ns.memory_efficient
 
-        # generator
-        self.sparse_grad_rng = torch.Generator(device='cuda' if torch.cuda.is_available() else 'cpu')
+        # Generator device must match parameter tensors for sparse masking.
+        # torch.Generator only supports cpu/cuda; non-CUDA param devices use cpu.
+        param_device = next(model.parameters()).device
+        gen_device = 'cuda' if param_device.type == 'cuda' else 'cpu'
+        self.sparse_grad_rng = torch.Generator(device=gen_device)
         self.sparse_grad_random_seed = np.random.randint(int(1e9))
 
         # inject model and base_optimizer

tests/optim/test_zeroth_order_updates.py

@@ -33,3 +33,55 @@ def closure():
     after = [p.detach().clone() for p in model.parameters()]
 
     assert any(not torch.equal(a, b) for a, b in zip(before, after))
+
+
+def test_zeroflow_sparse_grad_rng_matches_parameter_device():
+    from infty.optim import ZeroFlow
+
+    torch.manual_seed(0)
+    model = torch.nn.Linear(2, 1)
+    base_optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+    optimizer = ZeroFlow(
+        model.parameters(),
+        base_optimizer=base_optimizer,
+        model=model,
+        args={
+            "q": 1,
+            "zo_eps": 1e-3,
+            "perturbation_mode": "two_side",
+            "inftyopt": "zo_sgd",
+            "gradient_sparsity": 0.5,
+        },
+    )
+
+    param_device = next(model.parameters()).device
+    expected = "cuda" if param_device.type == "cuda" else "cpu"
+    assert str(optimizer.sparse_grad_rng.device) == expected
+
+    # Sparse masking must be usable with the constructed generator device.
+    x = torch.randn(4, 2)
+    y = torch.randn(4, 1)
+
+    def closure():
+        pred = model(x)
+        loss = torch.nn.functional.mse_loss(pred, y)
+        return pred, [loss]
+
+    optimizer.set_closure(closure)
+    optimizer.step()
+
+
+def test_zeroflow_does_not_hardcode_cuda_when_unavailable(monkeypatch):
+    from infty.optim import ZeroFlow
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+
+    model = torch.nn.Linear(2, 1)  # stays on CPU
+    base_optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+    optimizer = ZeroFlow(
+        model.parameters(),
+        base_optimizer=base_optimizer,
+        model=model,
+        args={"q": 1, "inftyopt": "zo_sgd"},
+    )
+    assert str(optimizer.sparse_grad_rng.device) == "cpu"