Add AdamW optimizer to extension/training (carries #18848 by @BryanBradfo)

[psiddh]

Attribution

This is Bryan's work, not mine. The AdamW implementation, the test suite, and the
design decisions in this PR are all from @BryanBradfo in #18848, opened 13 April
2026. I am carrying it forward only because that PR cannot merge: two of its commits
(5f1c90c, bda405f) are not covered by a signed Linux Foundation CLA, and the
contributor has not responded to requests to sign it (pinged 2 Sept, 8 Sept, and 15 Sept).

EasyCLA keys off commit authorship, so re-opening with their commits intact fails the same
check. Squashing under my authorship is what makes the check pass — maintainers should be
aware that this is the only reason the commit says what it says, and should decide
whether they are comfortable landing it on that basis rather than waiting for
@BryanBradfo to sign. I have no objection to this being closed in favour of #18848 if the
CLA gets signed.

Original PR: #18848

Fixes #18766

Summary

Adds AdamW to the training optimizer extension. It is a port of the existing SGD
implementation at extension/training/optimizer/sgd.{h,cpp}, with the main algorithmic
difference being decoupled weight decay — the parameter is decayed directly instead of
mixing the decay into the gradient. Matches torch.optim.AdamW with default settings.

Scope

C++ only. Python bindings are left out on purpose: the pybindings file has a TODO to build
a generic optimizer interface first, so copying PySGD to PyAdamW now would just add
duplication. amsgrad and maximize are also left out — both rarely used and easy to add
later.

Changes on top of #18848

Two things, both mine:

• Build fix. make_tensor_ptr gained a Device parameter after Add AdamW optimizer to extension/training #18848 was written, so
  the two state-buffer calls were passing TensorShapeDynamism::STATIC where a Device is
  expected and the deleter where the dynamism is expected. This is what broke every build
  job on that PR. Both moment buffers come from malloc, so they get an explicit CPU device.
• malloc null check before memset, as flagged in review on Add AdamW optimizer to extension/training #18848. The two
  allocations were used unconditionally, so a failure on a memory-constrained target
  dereferenced null instead of returning Error::MemoryAllocationFailed.

The other review points on #18848 (dtype/shape validation, hyperparameter range checks) are
deliberately not addressed here. @BryanBradfo's argument was that SGD has neither and adding
them to AdamW alone would be inconsistent, and @JacobSzwejbka noted the Copilot review is
noisy. Both optimizers can be hardened together in a follow-up.

Verification

lintrunner is clean, and adamw.cpp compiles against current main — I reproduced the
exact CI error from #18848 locally and confirmed this clears it.

I could not execute adamw_test (this machine SIGKILLs freshly built binaries), so instead
I transcribed AdamW::step's arithmetic verbatim, including the float round-trips, and
diffed it against torch.optim.AdamW over six configurations:

case                                adamw.cpp        torch     absdiff
Simple  (wd=0, g=-1, 10 steps)      2.0000002    2.0000002    0.00e+00
DecoupledWD (g=0, wd=0.5, 1)        0.9500000    0.9500000    0.00e+00
MultiParam p1 (g=-1, 5 steps)       1.5000001    1.5000001    0.00e+00
MultiParam p2 (g=+1, 5 steps)       1.4999999    1.4999999    0.00e+00
Both wd and grad nonzero            2.9613018    2.9613023    4.77e-07
Nondefault betas                   -0.2605234   -0.2605233    2.98e-08

The decoupled-weight-decay ordering matches and all four gtest expectations agree with
PyTorch independently. @BryanBradfo reported buck2 test //extension/training/optimizer/test:adamw_test
passing 6/6 and the SGD regression staying green on #18848; someone with a working Buck
setup should confirm that still holds, since the tests are a Buck-only target and CI does
not run them.

cc @JacobSzwejbka

---

The build fix and this PR were prepared with Claude Code.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22857

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit 9d1ccdc with merge base 81a2379 ():

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / unittest / macos / macos-job (gh)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Copilot]

🟡 Changes recommended

The critical tensor validation issue remains unresolved, along with two nits.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds a C++ AdamW optimizer with decoupled weight decay to ExecuTorch training.

Changes:

• Implements AdamW state, options, and updates.
• Registers library and test targets.
• Adds behavioral tests and documentation.

Review findings:

• adamw.cpp:145, 152, 163 — critical (1 vote): validate dtype/layout and matching shapes before updates.
• README.md:11 — nit (3 votes): use C++ instead of Cpp.
• adamw_test.cpp:88 — nit (3 votes): add changing-gradient multi-step coverage with tighter expectations.

File summaries

File	Summary
shim_et/xplat/executorch/build/build_variables.bzl	Registers the AdamW source.
extension/training/README.md	Documents AdamW availability.
extension/training/optimizer/test/targets.bzl	Defines the AdamW test target.
extension/training/optimizer/test/adamw_test.cpp	Adds AdamW unit tests.
extension/training/optimizer/targets.bzl	Defines the AdamW library target.
extension/training/optimizer/adamw.h	Declares the AdamW API and state.
extension/training/optimizer/adamw.cpp	Implements AdamW updates and state allocation.

Review details

Suppressed comments (2)

extension/training/optimizer/adamw.cpp:166

• Valid zero-element tensors are supported by ExecuTorch, but this treats malloc(0) returning nullptr as an allocation failure. On such allocators an empty parameter/gradient incorrectly returns MemoryAllocationFailed; only treat null as failure for nonzero storage and skip zero-length initialization.

        if (m_buf_ptr == nullptr || v_buf_ptr == nullptr) {
          free(m_buf_ptr);
          free(v_buf_ptr);
          return Error::MemoryAllocationFailed;

extension/training/optimizer/adamw.cpp:153

• Weight decay is applied before the lazy state allocation below. If either malloc fails, step() returns MemoryAllocationFailed after mutating p but before updating its moments; retrying decays it again and can leave a partially applied step. Allocate and initialize the state before mutating parameters, or preflight all allocations for the step.

      if (weight_decay != 0.0) {
        add_out_hack(p, p, -lr * weight_decay, p);

• Files reviewed: 7/7 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[psiddh]

Dispositioning the three Copilot findings. All three were also raised on #18848; the make_tensor_ptr and malloc ones it flagged there are gone because those are fixed.

Cpp -> C++ in the README — fixed, it was right.

Constant-gradient tests — fair, and fixed. A constant ±1 gradient makes Adam's normalized update very nearly sign-only, so the existing three tests would have passed with a broken moment recurrence or bias correction. Added AdamWOptimizerVaryingGradient: an eight-step sequence {0.5, -1.5, 2.0, -0.25, 1.0, 0.75, -2.5, 0.1}, run with and without decoupled weight decay, expectations from torch.optim.AdamW at 1e-6. Adam with coupled L2 lands on 0.79996 against AdamW's 0.77575 on the decay case, so the tolerance also pins the AdamW-vs-Adam distinction under a non-trivial gradient.

dtype/layout/shape validation (flagged High) — not fixed here, deliberately. The report is technically correct: const_data_ptr<T>() is an unchecked static_cast with no dtype enforcement, so a non-Float tensor would reinterpret bytes, and a p.numel() != g.numel() mismatch would read past the state buffers. But sgd.cpp's add_out_hack / mul_out_hack / clone_out_hack have exactly the same hazard on the same public step() surface — this is a property of the optimizer extension as a whole, not something AdamW introduces. @BryanBradfo made this argument on #18848 and I agree with it. Hardening one optimizer and not the other would leave the codebase in a worse-documented state than either fixing both or fixing neither. Happy to do both in a follow-up if @JacobSzwejbka wants it.

Content is byte-identical to #18848 (tree b580cf4), which has the same two commits pushed to it.

[Copilot]

🟡 Changes recommended

Allocation failure can mutate parameters before state initialization, causing repeated weight decay on retry.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

extension/training/optimizer/adamw.cpp:166

• g.nbytes() may be zero for a valid empty tensor, and malloc(0) is allowed to return nullptr. This condition then returns MemoryAllocationFailed (and the following zero-length memset calls still receive null pointers), so AdamW cannot no-op on empty parameters. Treat zero-byte state as valid and only allocate/initialize when the byte count is nonzero.

        if (m_buf_ptr == nullptr || v_buf_ptr == nullptr) {
          free(m_buf_ptr);
          free(v_buf_ptr);
          return Error::MemoryAllocationFailed;

• Files reviewed: 7/7 changed files
• Comments generated: 1
• Review effort level: Lite

[psiddh]

Copilot's follow-up on the allocation-failure path was right, and it was on code I added rather than @BryanBradfo's. My malloc null check returned MemoryAllocationFailed after the in-place weight decay had already changed p, while its moments and step counter stayed put — so a caller retrying the failed step would decay the parameter a second time. Before my change there was no early return there at all, so I introduced that window.

Fixed by moving the state lookup and allocation above the decay, so nothing mutates p until the fallible part has succeeded. Decay still precedes the moment update, so the arithmetic is untouched and the reference values in AdamWOptimizerVaryingGradient are unchanged.

Same commit is on #18848; trees still identical at 9b843c2.

[Copilot]

🟡 Changes recommended

Zero-element tensors can incorrectly trigger MemoryAllocationFailed during state allocation.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 7/7 changed files
• Comments generated: 1
• Review effort level: Lite

[psiddh]

Upgraded verification: the real C++ now runs, not just a transcription

Earlier I could only say I had transcribed AdamW::step()'s arithmetic and diffed it against
torch.optim.AdamW, because the gtest binary was being SIGKILLed in my environment. I got the
actual compiled ExecuTorch code running, so this is now direct evidence.

case                                    got     expected     |diff|  result
Simple (wd=0,g=-1,x10)           2.00000024   2.00000000   2.38e-07  PASS
DecoupledWD (g=0,wd=.5,x1)       0.94999999   0.94999999   0.00e+00  PASS
MultiParam p1 (g=-1,x5)          1.50000012   1.50000000   1.19e-07  PASS
MultiParam p2 (g=+1,x5)          1.49999988   1.50000000   1.19e-07  PASS
Varying wd=0 (8 steps)           0.84547687   0.84547687   0.00e+00  PASS
Varying wd=0.1 (8 steps)         0.77574724   0.77574724   0.00e+00  PASS
6/6 passed

The two varying-gradient values are the ones taken from torch.optim.AdamW, and the real
implementation reproduces them exactly.

Per-parameter state isolation, which the cases above do not cover because they use one optimizer
each, checked separately with two parameters in a single optimizer:

two params, one optimizer, 5 steps:
  a: 1.50000012 (expect ~1.5) PASS
  b: 1.49999988 (expect ~1.5) PASS
  partial gradients: b unchanged = PASS (1.49999988 -> 1.49999988)

That last line also covers the named_gradients.find(...) == end() skip path, and the whole run
exercises the state-allocation reorder from the third commit, since state is allocated on step 1
and reused on steps 2-5.

What this still does not cover

The gtest harness itself (TensorFactory, the ET_EXPERIMENTAL wrappers) and the Buck target. My
driver calls AdamW::step() directly against TensorImpl. The gtest binary is still SIGKILLed
here by what looks like an endpoint-security rule — a trivial binary linking the same objects runs
fine, so it is not the code. Someone should still run
buck2 test //extension/training/optimizer/test:adamw_test before merge.