[egeozkoc]

Paged optimizers on MPS could reach cget_managed_ptr and fail with a CPU-library error instead of allocating regular state buffers. AdEMAMix's separate double-buffer allocator also bypassed the fallback on both CPU and MPS for parameters with at least 100,000 elements.

Both allocation paths now use regular tensors on devices outside CUDA/ROCm/XPU, with the existing once-per-optimizer warning naming the device. CUDA/ROCm/XPU keep the managed-memory path. The existing test_optimizer32bit paged-case skip is widened to non-paging devices; the extra regression test was removed as requested.

Validation on Apple M5/macOS 26.0, Python 3.12.13, both PyTorch 2.12.1 and 2.15.0, with this patch applied to current main (a616886):

• Existing optimizer selection: 72 passed, 72 skipped on each of CPU and MPS (108 deselected, including benchmarks).
• Separate local integration checks: 48 state-allocation cases covering all 12 paged classes, small/large buffers, FP32/uint8 state and once-only warnings; 24 optimizer-step cases comparing PagedAdam/AdamW/Lion/AdEMAMix to their non-paged counterparts across FP32/FP16/BF16 on CPU/MPS. All pass with exact parameter/state parity. The 8-bit checks validate allocation only, not MPS 8-bit optimizer updates.
• Before the AdEMAMix follow-up, four large-buffer allocation cases failed (PagedAdEMAMix and PagedAdEMAMix8bit on CPU/MPS); all 48 allocation cases now also pass on the exact PR head.
• Applicable pre-commit hooks and git diff --check pass. The local validation scripts are kept outside the PR test suite.

CUDA/ROCm/XPU hardware was unavailable locally. Upstream CI remains required: the previous PR Tests run expired while awaiting maintainer approval. This change does not implement MPS 8-bit optimizers.

Ref #2011.

🤖 Generated with Claude Code; follow-up fix and validation assisted by Codex.

[matthewdouglas]

Thanks for the PR. Given that paged optimizers are only actually supported on CUDA/ROCm/XPU, it seems superfluous. But considering we'd later implement optimizers for MPS or fallback, I suppose it's reasonable in that respect.

With that said, the extra test for regression against previous behavior, literally saying "previously" in the comments, IMHO seems just unnecessary. I do not understand why the agents love adding these pointless tests so much, but please remove that.

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[egeozkoc]

Removed the test, and dropped a now-stale reference to it from the skip message. The remaining test change is just widening the existing paged_* skip in test_optimizer32bit from CPU to any non-paging device.

[egeozkoc]

@matthewdouglas, I’ve pushed 80a30b6 to cover a missed allocation path: large PagedAdEMAMix/PagedAdEMAMix8bit double buffers still called cget_managed_ptr on CPU/MPS. The helper now uses the same device allow-list as the base allocator. The removed regression test remains removed.

Refreshed validation on Apple M5/macOS 26.0, PyTorch 2.12.1 and 2.15.0, with the patch applied to current main (a616886):

• 72 existing tests passed on each of CPU and MPS, with 72 expected skips per selection.
• 72 separate local allocation/update checks passed per version: 48 state-allocation cases and 24 paged-versus-non-paged step comparisons. All 48 allocation cases also pass on the exact PR head; four AdEMAMix cases failed before this fix.
• Applicable pre-commit hooks pass. These validation scripts are outside the PR test suite.

CUDA/ROCm/XPU hardware was unavailable locally; MPS 8-bit checks cover allocation only.

The fresh PR Tests, Lint, Python package, and documentation workflows are awaiting approval. Could you approve those runs and review the updated PR? The description now includes the fix and validation details.