Migrate optimizer state across devices (GH-1615)

Adam and SGD previously reused compatible optimizer state buffers even
when replacement parameters moved to another device. This left kernels
with state and parameters on different devices.

Migrate compatible state buffers while preserving their accumulated
values. Keep incompatible shape or dtype changes using freshly zeroed
buffers, and cover mixed moved and unmoved parameters.

Signed-off-by: Eric Shi <ershi@nvidia.com>

Author: shi-eric

CHANGELOG.md

@@ -204,10 +204,11 @@
   capture, which APIC cannot record for replay ([GH-1559](https://github.com/NVIDIA/warp/issues/1559)).
 - Fix APIC graph capture for `warp.fem` workflows by making partition node counts graph-capture safe
   ([GH-1407](https://github.com/NVIDIA/warp/issues/1407)).
-- Fix `warp.optim.Adam.set_params()` re-zeroing the optimizer moment buffers on every call for `fp16`
-  parameters, which silently discarded accumulated state. The reuse check now compares against the
-  moment-buffer dtype rather than the parameter dtype
-  ([GH-1593](https://github.com/NVIDIA/warp/issues/1593)).
+- Fix optimizer state reuse in `warp.optim.Adam.set_params()` and `warp.optim.SGD.set_params()`. Adam no longer
+  re-zeros moment buffers on every call for `fp16` parameters, and compatible moment and momentum buffers migrate
+  when replacement parameters move devices, preserving accumulated state and preventing invalid cross-device reuse
+  ([GH-1593](https://github.com/NVIDIA/warp/issues/1593),
+  [GH-1615](https://github.com/NVIDIA/warp/issues/1615)).
 
 ### Documentation
 

warp/_src/optim/adam.py

@@ -133,8 +133,16 @@ def set_params(self, params):
                 # buffer dtype), not ``param.dtype``, otherwise the buffers are re-zeroed on every call.
                 if self.m[i] is None or self.m[i].shape != param.shape or self.m[i].dtype != dtype:
                     self.m[i] = wp.zeros(shape=param.shape, dtype=dtype, device=param.device)
+                elif self.m[i].device != param.device:
+                    migrated_m = wp.empty(shape=param.shape, dtype=dtype, device=param.device)
+                    wp.copy(migrated_m, self.m[i])
+                    self.m[i] = migrated_m
                 if self.v[i] is None or self.v[i].shape != param.shape or self.v[i].dtype != dtype:
                     self.v[i] = wp.zeros(shape=param.shape, dtype=dtype, device=param.device)
+                elif self.v[i].device != param.device:
+                    migrated_v = wp.empty(shape=param.shape, dtype=dtype, device=param.device)
+                    wp.copy(migrated_v, self.v[i])
+                    self.v[i] = migrated_v
 
     def reset_internal_state(self):
         """Reset moment buffers and timestep to zero."""

warp/_src/optim/sgd.py

@@ -82,6 +82,10 @@ def set_params(self, params):
                 param = params[i]
                 if self.b[i] is None or self.b[i].shape != param.shape or self.b[i].dtype != param.dtype:
                     self.b[i] = wp.zeros_like(param)
+                elif self.b[i].device != param.device:
+                    migrated_b = wp.empty_like(param)
+                    wp.copy(migrated_b, self.b[i])
+                    self.b[i] = migrated_b
                 # Overload the kernel for each parameter so we can precompile the SGD kernel
                 if param is not None:
                     wp.overload(sgd_step_kernel, {"g": param, "b": param, "params": param})

warp/tests/test_adam.py

@@ -159,6 +159,46 @@ def test_adam_set_params_preserves_fp16_state(test, device):
                 test.assertTrue((opt.v[0].numpy() == 1.0).all(), f"second moment reset for {param_dtype}")
 
 
+def test_adam_set_params_migrates_state(test, device):
+    """Verify compatible moment buffers follow parameters that move devices."""
+    moved_param = wp.zeros(4, dtype=wp.float32, device="cpu")
+    unmoved_param = wp.zeros(4, dtype=wp.float32, device="cpu")
+    opt = warp.optim.Adam([moved_param, unmoved_param], lr=0.02)
+
+    moved_m, moved_v = opt.m[0], opt.v[0]
+    unmoved_m, unmoved_v = opt.m[1], opt.v[1]
+    moved_m.fill_(1.0)
+    moved_v.fill_(2.0)
+
+    expected_m = moved_m.numpy().copy()
+    expected_v = moved_v.numpy().copy()
+    replacement = wp.zeros(4, dtype=wp.float32, device=device)
+    opt.set_params([replacement, unmoved_param])
+
+    test.assertEqual(opt.m[0].device, device)
+    test.assertEqual(opt.v[0].device, device)
+    np.testing.assert_array_equal(opt.m[0].numpy(), expected_m)
+    np.testing.assert_array_equal(opt.v[0].numpy(), expected_v)
+    test.assertIs(opt.m[1], unmoved_m)
+    test.assertIs(opt.v[1], unmoved_v)
+
+    opt.step([wp.ones_like(replacement), wp.ones_like(unmoved_param)])
+    replacement.numpy()
+    unmoved_param.numpy()
+
+    expected_m = opt.m[0].numpy().copy()
+    expected_v = opt.v[0].numpy().copy()
+    replacement = wp.zeros(4, dtype=wp.float32, device="cpu")
+    opt.set_params([replacement, unmoved_param])
+
+    test.assertEqual(opt.m[0].device, wp.get_device("cpu"))
+    test.assertEqual(opt.v[0].device, wp.get_device("cpu"))
+    np.testing.assert_array_equal(opt.m[0].numpy(), expected_m)
+    np.testing.assert_array_equal(opt.v[0].numpy(), expected_v)
+    test.assertIs(opt.m[1], unmoved_m)
+    test.assertIs(opt.v[1], unmoved_v)
+
+
 devices = get_test_devices()
 
 
@@ -172,6 +212,12 @@ class TestAdam(unittest.TestCase):
 add_function_test(
     TestAdam, "test_adam_set_params_preserves_fp16_state", test_adam_set_params_preserves_fp16_state, devices=devices
 )
+add_function_test(
+    TestAdam,
+    "test_adam_set_params_migrates_state",
+    test_adam_set_params_migrates_state,
+    devices=get_cuda_test_devices(),
+)
 
 
 if __name__ == "__main__":

warp/tests/test_sgd.py

@@ -3,6 +3,8 @@
 
 import unittest
 
+import numpy as np
+
 import warp as wp
 import warp.optim
 from warp.tests.unittest_utils import *
@@ -205,6 +207,37 @@ def test_sgd_reset_internal_state(test, device):
         test.assertAlmostEqual(result2[0], 0.9, places=5)
 
 
+def test_sgd_set_params_migrates_state(test, device):
+    """Verify compatible momentum buffers follow parameters that move devices."""
+    moved_param = wp.zeros(4, dtype=wp.float32, device="cpu")
+    unmoved_param = wp.zeros(4, dtype=wp.float32, device="cpu")
+    opt = warp.optim.SGD([moved_param, unmoved_param], lr=0.02, momentum=0.9)
+
+    moved_b = opt.b[0]
+    unmoved_b = opt.b[1]
+    moved_b.fill_(1.0)
+
+    expected_b = moved_b.numpy().copy()
+    replacement = wp.zeros(4, dtype=wp.float32, device=device)
+    opt.set_params([replacement, unmoved_param])
+
+    test.assertEqual(opt.b[0].device, device)
+    np.testing.assert_array_equal(opt.b[0].numpy(), expected_b)
+    test.assertIs(opt.b[1], unmoved_b)
+
+    opt.step([wp.ones_like(replacement), wp.ones_like(unmoved_param)])
+    replacement.numpy()
+    unmoved_param.numpy()
+
+    expected_b = opt.b[0].numpy().copy()
+    replacement = wp.zeros(4, dtype=wp.float32, device="cpu")
+    opt.set_params([replacement, unmoved_param])
+
+    test.assertEqual(opt.b[0].device, wp.get_device("cpu"))
+    np.testing.assert_array_equal(opt.b[0].numpy(), expected_b)
+    test.assertIs(opt.b[1], unmoved_b)
+
+
 devices = get_test_devices()
 
 
@@ -219,6 +252,12 @@ class TestSGD(unittest.TestCase):
 add_function_test(TestSGD, "test_sgd_combined_features", test_sgd_combined_features, devices=devices)
 add_function_test(TestSGD, "test_sgd_no_momentum", test_sgd_no_momentum, devices=devices)
 add_function_test(TestSGD, "test_sgd_reset_internal_state", test_sgd_reset_internal_state, devices=devices)
+add_function_test(
+    TestSGD,
+    "test_sgd_set_params_migrates_state",
+    test_sgd_set_params_migrates_state,
+    devices=get_cuda_test_devices(),
+)
 
 
 if __name__ == "__main__":