From 3a28f350f167e1d506c11d5b8345986e33a714f0 Mon Sep 17 00:00:00 2001
From: Zhu Lei <woicyou@gmail.com>
Date: Tue, 15 Sep 2026 08:10:28 +0000
Subject: [PATCH] Fix ZeroFlow sparse RNG to follow parameter device

Construct sparse_grad_rng from the model parameter device instead of
hardcoding cuda-if-available, so CPU/MPS models do not pair CUDA
generators with non-CUDA tensors during sparse masking. Also drop the
duplicate SimpleNamespace import and the redundant SAM get_grad_reduce
call already performed by InftyBaseOptimizer.

Fixes #8
---
 src/infty/optim/geometry_reshaping/sam.py     |  1 -
 .../optim/zeroth_order_updates/zeroflow.py    |  8 +--
 tests/optim/test_zeroth_order_updates.py      | 52 +++++++++++++++++++
 3 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/src/infty/optim/geometry_reshaping/sam.py b/src/infty/optim/geometry_reshaping/sam.py
index f020e0e..bd4c7a1 100644
--- a/src/infty/optim/geometry_reshaping/sam.py
+++ b/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)
 
     
     
diff --git a/src/infty/optim/zeroth_order_updates/zeroflow.py b/src/infty/optim/zeroth_order_updates/zeroflow.py
index cd02de5..3a0e6d1 100644
--- a/src/infty/optim/zeroth_order_updates/zeroflow.py
+++ b/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
diff --git a/tests/optim/test_zeroth_order_updates.py b/tests/optim/test_zeroth_order_updates.py
index ec23026..1f4ee51 100644
--- a/tests/optim/test_zeroth_order_updates.py
+++ b/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"