Migrate optimizer state across devices [GH-1615]#1616
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (5)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughAdam and SGD ChangesOptimizer state device migration
Estimated code review effort: 2 (Simple) | ~12 minutes Related Issues: Suggested labels: bug, optimizer Suggested reviewers: none 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Greptile SummaryThis PR fixes a bug where
Confidence Score: 5/5Safe to merge — the fix is a minimal two-line addition per optimizer that only activates when shape and dtype already match, with no risk to the existing allocation or zero-initialization paths. The change is surgical and well-contained: an elif branch guarded by a device inequality check, delegating to warp.array.to() which is already used throughout the codebase and internally calls warp.clone() for value-preserving cross-device copies. Existing behavior for shape/dtype mismatches is untouched. The new tests exercise both migration directions (CPU↔CUDA) for both optimizers, verify value preservation, confirm unmoved-parameter identity is retained, and validate a full optimizer step post-migration. No files require special attention. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["set_params(params) called"] --> B["For each param[i]"]
B --> C{"m[i] is None OR\nshape mismatch OR\ndtype mismatch?"}
C -- Yes --> D["Allocate fresh zeros on param.device\n(state reset)"]
C -- No --> E{"m[i].device\n!= param.device?"}
E -- Yes --> F["m[i] = m[i].to(param.device)\nwarp.clone() copies values\nacross devices"]
E -- No --> G["No-op: reuse existing buffer\n(same device, shape, dtype)"]
D --> H["Same check for v[i] / b[i]"]
F --> H
G --> H
H --> I["step() runs kernel\non param.device\nwith co-located state"]
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
A["set_params(params) called"] --> B["For each param[i]"]
B --> C{"m[i] is None OR\nshape mismatch OR\ndtype mismatch?"}
C -- Yes --> D["Allocate fresh zeros on param.device\n(state reset)"]
C -- No --> E{"m[i].device\n!= param.device?"}
E -- Yes --> F["m[i] = m[i].to(param.device)\nwarp.clone() copies values\nacross devices"]
E -- No --> G["No-op: reuse existing buffer\n(same device, shape, dtype)"]
D --> H["Same check for v[i] / b[i]"]
F --> H
G --> H
H --> I["step() runs kernel\non param.device\nwith co-located state"]
Reviews (4): Last reviewed commit: "Migrate optimizer state across devices (..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
warp/tests/test_adam.py (1)
162-208: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing CUDA→CPU and CUDA→CUDA migration coverage.
test_adam_set_params_migrates_stateonly exercises CPU→CUDA migration (plus an unmoved CPU parameter). The linked issue explicitly asks for coverage of CPU-to-CUDA, CUDA-to-CPU, and CUDA-to-CUDA cases, but only CPU-to-CUDA (and the mixed case) is tested here. This gap means a regression in the CUDA→CPU path, or the potential unsynchronized peer-to-peer copy issue on CUDA→CUDA (flagged inadam.py), would go undetected.Consider adding symmetric variants (starting the parameter on
deviceand moving it to"cpu", and — when more than one CUDA device is available — moving between two CUDA devices) to close this gap.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@warp/tests/test_adam.py` around lines 162 - 208, The Adam migration test only covers CPU-to-CUDA, so expand test_adam_set_params_migrates_state to also verify the symmetric CUDA-to-CPU path and, when multiple CUDA devices are available, CUDA-to-CUDA migration. Reuse the existing Adam setup and state assertions around warp.optim.Adam.set_params, but add separate cases that start a parameter on device, move it to "cpu", and move it between two CUDA devices to ensure m/v buffers follow correctly in all directions.warp/_src/optim/adam.py (1)
136-145: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMigration logic duplicated across
m/v(here) andbinsgd.py.The reuse/reallocate/migrate pattern is repeated almost verbatim three times (Adam's
m, Adam'sv, and SGD'sbinwarp/_src/optim/sgd.py). Extracting a small shared helper would reduce duplication and ensure any future fix (e.g. the sync issue above) only needs to be applied once.♻️ Suggested helper extraction
def _migrate_or_reset_state(buffer, shape, dtype, device): if buffer is None or buffer.shape != shape or buffer.dtype != dtype: return wp.zeros(shape=shape, dtype=dtype, device=device) if buffer.device != device: migrated = wp.empty(shape=shape, dtype=dtype, device=device) wp.copy(migrated, buffer) return migrated return bufferUsage:
self.m[i] = _migrate_or_reset_state(self.m[i], param.shape, dtype, param.device), similarly forvand SGD'sb.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@warp/_src/optim/adam.py` around lines 136 - 145, The state reset/migration logic for Adam’s `m` and `v` is duplicated and matches the `b` handling in `sgd.py`, so extract the repeated reuse/reallocate/migrate pattern into a small shared helper. Add a helper like `_migrate_or_reset_state` in the optimizers module and use it from `Adam`’s update path for both `self.m[i]` and `self.v[i]`, and from SGD for `self.b[i]`, so the shape/dtype/device checks and migration behavior live in one place.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@warp/_src/optim/adam.py`:
- Around line 136-145: Synchronize the source CUDA device before migrating Adam
moments in set_params() so cross-GPU copies do not race with in-flight updates
from step(). In the migration branches that use wp.copy for self.m[i] and
self.v[i], add a source-device sync or event barrier before copying into
migrated_m and migrated_v, then continue assigning the migrated buffers as
before.
---
Nitpick comments:
In `@warp/_src/optim/adam.py`:
- Around line 136-145: The state reset/migration logic for Adam’s `m` and `v` is
duplicated and matches the `b` handling in `sgd.py`, so extract the repeated
reuse/reallocate/migrate pattern into a small shared helper. Add a helper like
`_migrate_or_reset_state` in the optimizers module and use it from `Adam`’s
update path for both `self.m[i]` and `self.v[i]`, and from SGD for `self.b[i]`,
so the shape/dtype/device checks and migration behavior live in one place.
In `@warp/tests/test_adam.py`:
- Around line 162-208: The Adam migration test only covers CPU-to-CUDA, so
expand test_adam_set_params_migrates_state to also verify the symmetric
CUDA-to-CPU path and, when multiple CUDA devices are available, CUDA-to-CUDA
migration. Reuse the existing Adam setup and state assertions around
warp.optim.Adam.set_params, but add separate cases that start a parameter on
device, move it to "cpu", and move it between two CUDA devices to ensure m/v
buffers follow correctly in all directions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Enterprise
Run ID: 55ddbbeb-927e-41a9-90ab-e16cf00568f4
📒 Files selected for processing (5)
CHANGELOG.mdwarp/_src/optim/adam.pywarp/_src/optim/sgd.pywarp/tests/test_adam.pywarp/tests/test_sgd.py
| 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 |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== Relevant code in warp/_src/optim/adam.py ==\n'
ast-grep outline warp/_src/optim/adam.py --view expanded || true
sed -n '1,240p' warp/_src/optim/adam.py | cat -n
printf '\n== Search for wp.copy / synchronize_device semantics ==\n'
rg -n "synchronize_device|wp\.copy\(|copy\(" warp/_src docs README.md -g '!**/*.pyc' || true
printf '\n== Locate copy implementation and any device-ordering docs ==\n'
fd -a "copy*.py" warp/_src || true
fd -a "*.md" . | sed -n '1,120p'Repository: NVIDIA/warp
Length of output: 24857
🌐 Web query:
Warp documentation wp.copy CUDA CUDA peer access source device synchronize_device multi-device copy semantics
💡 Result:
In NVIDIA Warp, the wp.copy function handles data movement between arrays, including cross-device transfers [1][2]. Its behavior regarding multi-device copy semantics, synchronization, and peer access is summarized below. Multi-Device Copy Semantics and Synchronization Copying between CUDA arrays on different devices is asynchronous with respect to the host [3][4]. Because CUDA devices operate asynchronously relative to each other, peer-to-peer transfers require careful synchronization to ensure correctness [3]. When a transfer occurs between GPUs, the destination device typically performs the copy, necessitating that any prior work on the source device is complete before the transfer begins [3][4]. To manage this, users should synchronize devices explicitly: 1. Use wp.synchronize_device(source_device) before the copy to ensure the source array is ready [3][4]. 2. Use wp.synchronize_device(dest_device) after the copy to ensure the transfer has finished [3][4]. By default, if the stream argument is omitted, Warp selects a stream for the copy based on these rules: 1. If the destination array is on a CUDA device, it uses the current stream on that destination device [1]. 2. Otherwise, if the source array is on a CUDA device, it uses the current stream on that source device [1]. Users can override this by explicitly providing a stream, allowing the copy to be performed on a stream from any device [3]. CUDA Peer Access CUDA peer access enables direct memory access (DMA) between GPUs, which significantly accelerates cross-device transfers compared to using CPU staging buffers [5][6]. - Support: You can check if peer access is supported between two devices using wp.is_peer_access_supported(device1, device2) [5][6]. - Configuration: Peer access is directional and can be enabled using wp.set_peer_access_enabled(target_device, peer_device, True) or managed temporarily with the wp.ScopedPeerAccess(target_device, peer_device, True) context manager [5][7][6]. - Scope: Peer access settings apply to all allocations made with the default CUDA allocator [7]. - Note on Memory Pools: Peer access does not accelerate transfers between arrays allocated via stream-ordered memory pools (introduced in Warp 0.14.0); for those, memory-pool access must be enabled instead [5][8]. As with all asynchronous CUDA operations, wp.synchronize_device is the primary tool for ensuring memory consistency between the host and specific GPU devices [6].
Citations:
- 1: https://nvidia.github.io/warp/stable/api_reference/_generated/warp.copy.html
- 2: https://nvidia.github.io/warp/stable/user_guide/runtime.html
- 3: https://nvidia.github.io/warp/stable/deep_dive/concurrency.html
- 4: https://github.com/NVIDIA/warp/blob/7ddcbbb2/docs/modules/concurrency.rst
- 5: https://nvidia.github.io/warp/stable/user_guide/devices.html
- 6: https://github.com/NVIDIA/warp/blob/56019eda/docs/user_guide/devices.rst
- 7: https://nvidia.github.io/warp/api_reference/_generated/warp.ScopedPeerAccess.html
- 8: https://nvidia.github.io/warp/stable/deep_dive/memory_access.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== warp/_src/optim/adam.py excerpt ==\n'
sed -n '120,170p' warp/_src/optim/adam.py | cat -n
printf '\n== repo references to synchronize_device ==\n'
rg -n "synchronize_device" warp/_src | sed -n '1,120p' || true
printf '\n== repo references to wp.copy ==\n'
rg -n "wp\.copy\(" warp/_src | sed -n '1,120p' || trueRepository: NVIDIA/warp
Length of output: 7182
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== docs/deep_dive/concurrency.rst excerpt ==\n'
sed -n '90,170p' docs/deep_dive/concurrency.rst | cat -n
printf '\n== docs/user_guide/devices.rst excerpt ==\n'
sed -n '60,210p' docs/user_guide/devices.rst | cat -n
printf '\n== warp/_src/context.py copy implementation excerpt ==\n'
sed -n '12526,12870p' warp/_src/context.py | cat -nRepository: NVIDIA/warp
Length of output: 31026
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== docs/deep_dive/concurrency.rst around cross-device copy ==\n'
sed -n '120,155p' docs/deep_dive/concurrency.rst | cat -n
printf '\n== docs/deep_dive/concurrency.rst around stream/device sync discussion ==\n'
sed -n '240,270p' docs/deep_dive/concurrency.rst | cat -n
printf '\n== warp/_src/context.py around wp.copy device selection and CUDA copy paths ==\n'
sed -n '12526,12740p' warp/_src/context.py | cat -nRepository: NVIDIA/warp
Length of output: 16610
Synchronize the source CUDA device before migrating Adam moments across devices. wp.copy(migrated_m, self.m[i]) / wp.copy(migrated_v, self.v[i]) enqueue the peer copy on the destination stream, but cross-GPU copies need outstanding source-device work to finish first. set_params() can run immediately after step(), so this can race with the previous update and copy stale m/v state; add a source-device sync or event barrier before the copy.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@warp/_src/optim/adam.py` around lines 136 - 145, Synchronize the source CUDA
device before migrating Adam moments in set_params() so cross-GPU copies do not
race with in-flight updates from step(). In the migration branches that use
wp.copy for self.m[i] and self.v[i], add a source-device sync or event barrier
before copying into migrated_m and migrated_v, then continue assigning the
migrated buffers as before.
9d08214 to
a3342ed
Compare
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>
a3342ed to
964aa34
Compare
1 participant
Description
Closes #1615.
Adam and SGD reused shape- and dtype-compatible optimizer state buffers
when replacement parameters moved to another device. This could leave the
parameters and their state on different devices. Migrate compatible state
buffers to the replacement parameter's device while preserving their values.
Shape or dtype changes continue to allocate zeroed state.
Changes
Checklist
Unreleasedsection.Validation summary
test_adam_set_params_migrates_statepreserves both momentbuffers and reuses the state of an unmoved parameter.
test_sgd_set_params_migrates_statepreserves the momentumbuffer and reuses the state of an unmoved parameter.
passed.
Bug fix
Without this change, the optimizer state remains on the CPU after replacing
the parameter with one on CUDA:
Summary by CodeRabbit
set_params()calls, and that unchanged parameters’ state buffers keep the same object identity.