Add AdamW optimizer to extension/training

[BryanBradfo]

Adds AdamW to the training optimizer extension. It's 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 gets decayed directly instead of mixing the decay into the gradient). Matches torch.optim.AdamW with default settings.

Fixes #18766

Scope

C++ only for this PR. 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. Happy to follow up with that. amsgrad and maximize are also left out, both rarely used and easy to add later if needed.

Test plan

Six new gtests pass, and the SGD regression stays green:

$ buck2 test //extension/training/optimizer/test:adamw_test
[  PASSED  ] 6 tests.

$ buck2 test //extension/training/optimizer/test:sgd_test
[  PASSED  ] 5 tests.

Output was also cross-checked against torch.optim.AdamW on four small cases (simple convergence, decoupled weight decay, multi-parameter). All four match to six decimal places.

cc @JacobSzwejbka

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit e0d1fac with merge base 9f209bd ():

NEW FAILURE - The following job has failed:

• pull / test-qnn-delegate-linux / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

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

[BryanBradfo]

@pytorchbot label "release notes: training"

[Copilot]

Pull request overview

Adds an AdamW optimizer implementation to ExecuTorch’s on-device training extension, aligning behavior with torch.optim.AdamW (decoupled weight decay) and integrating it into the existing C++ optimizer build/test setup.

Changes:

• Introduces AdamW optimizer implementation (adamw.{h,cpp}) and exposes it as a training optimizer target.
• Adds new gtests for AdamW and wires them into the optimizer test targets.
• Updates training extension build source lists and documentation to reflect AdamW availability.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
shim_et/xplat/executorch/build/build_variables.bzl	Adds AdamW source to extension training sources list.
extension/training/optimizer/targets.bzl	Defines a new adamw C++ library target.
extension/training/optimizer/adamw.h	Declares AdamW API, options, param group, and state types.
extension/training/optimizer/adamw.cpp	Implements AdamW step logic and state allocation/freeing.
extension/training/optimizer/test/targets.bzl	Adds adamw_test target.
extension/training/optimizer/test/adamw_test.cpp	New unit tests for AdamW behavior and defaults.
extension/training/README.md	Updates optimizer list to include AdamW.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

In portable (non-USE_ATEN_LIB) mode, the state tensors' TensorImpl are constructed with sizes/dim_order pointers taken from the gradient tensor (g.sizes().data()/g.dim_order().data()). TensorImpl does not own these arrays, and the optimizer state outlives the gradient tensor, so this can become a use-after-free if the gradient's metadata storage doesn't live for the lifetime of the optimizer. Allocate and store your own sizes/dim_order (and strides if needed) with the state, or use a construction path that ensures the metadata buffers are owned/stable for the optimizer lifetime.

[BryanBradfo]

Good catch in principle. This matches the existing pattern in SGD (sgd.cpp:119-125) which constructs TensorImpl with the same d_p.sizes().data() / d_p.dim_order().data() pointers from the gradient. In practice, parameter shapes are stable across training steps in ExecuTorch's captured graph model. Happy to address both optimizers together in a follow-up if the maintainer prefers.

[Copilot]

AdamWOptions parameters are used in bias-correction divisors (1 - beta^step) and sqrt(bias_correction2). With beta1==1 or beta2==1 (or beta2<0), step() will divide by zero / take sqrt of a negative, producing NaNs/inf. Consider validating lr/beta1/beta2/eps ranges (e.g., lr>=0, 0<=beta<1, eps>0) either in AdamWOptions construction or at the start of step() and return Error::InvalidArgument on invalid values.

[BryanBradfo]

SGD also has no parameter validation (sgd.h), and PyTorch's C++ AdamW documents 0 <= beta < 1 as a precondition rather than a runtime check. Adding validation only to AdamW would be inconsistent. Happy to add range checks across both optimizers in a follow-up if the maintainer wants.

[Copilot]

AdamW::step immediately treats both parameters and gradients as float tensors and iterates assuming matching shapes. If a caller passes a non-float tensor or a gradient whose shape/numel differs from the parameter, this can lead to undefined behavior or out-of-bounds reads/writes (e.g., in addcmul_sq_out_hack/adamw_update_hack). Add explicit scalar_type/shape checks (at least dtype==Float and p.numel()==g.numel()) and return Error::InvalidArgument when unsupported.

Suggested change

	if (
	p.scalar_type() != executorch::aten::ScalarType::Float ||
	g.scalar_type() != executorch::aten::ScalarType::Float) {
	return Error::InvalidArgument;
	}
	if (p.numel() != g.numel()) {
	return Error::InvalidArgument;
	}

[Copilot]

The per-parameter state allocation uses malloc() but never checks for allocation failure. On memory-constrained targets this can turn into a null dereference in memset()/from_blob/TensorImpl construction. Check m_buf_ptr/v_buf_ptr for nullptr, free any partially-allocated buffers, and return Error::MemoryAllocationFailed.

Suggested change

	void* v_buf_ptr = malloc(g.nbytes());
	void* v_buf_ptr = malloc(g.nbytes());
	if (m_buf_ptr == nullptr || v_buf_ptr == nullptr) {
	free(m_buf_ptr);
	free(v_buf_ptr);
	return Error::MemoryAllocationFailed;
	}

[JacobSzwejbka]

You can ignore copilot probably its noise ratio is pretty bad. Ive been trying to figure out how to turn it off for the repo.

[JacobSzwejbka]

ahh my legacy... @manuelcandales do you remember how to call executorch ops outside of the interpreter. I cant recall

[Copilot]

Copilot was unable to review this pull request because the user who requested the review is ineligible. To be eligible to request a review, you need a paid Copilot license, or your organization must enable Copilot code review.

[linux-foundation-easycla]

• ✅ login: JacobSzwejbka / name: Jacob Szwejbka (fad7bc6)
• ✅ login: psiddh / name: Siddartha Pothapragada (525f89c, 7259664, 7c76127, 8e39c49, e0d1fac, f675af0)
• ❌ - login: @BryanBradfo / name: Bryan. The commit (5f1c90c, bda405f) is not authorized under a signed CLA. Please click here to be authorized. For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.
• ❌ - login: @BryanBradfo / name: Bryan Chen. The commit (525f89c) is not authorized under a signed CLA. Please click here to be authorized. For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

[psiddh]

@BryanBradfo Can you complete the EasyCLA CI step please ? Once completed, we can merge this PR.
Thank you again for your contribution

[psiddh]

@BryanBradfo Pinging again. Could you pls take care of easy CLA ?

[Copilot]

🟡 Changes recommended

The new AdamW implementation has concrete crash/UB risks (unchecked malloc failures and unconditional float/contiguous assumptions) that should be guarded before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

extension/training/optimizer/adamw.cpp:164

• malloc() can return nullptr; the subsequent memset() would then dereference a null pointer and crash. Check allocation failures and return Error::MemoryAllocationFailed (freeing any partially allocated buffer) before continuing.

        void* m_buf_ptr = malloc(g.nbytes());
        void* v_buf_ptr = malloc(g.nbytes());
        std::memset(m_buf_ptr, 0, g.nbytes());
        std::memset(v_buf_ptr, 0, g.nbytes());

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

[psiddh]

Few compilation issues..

/__w/executorch/executorch/pytorch/executorch/extension/training/optimizer/adamw.cpp:168:22: error: no matching function for call to 'make_tensor_ptr'
auto m_ptr = make_tensor_ptr(
^~~~~~~~~~~~~~~
/__w/executorch/executorch/pytorch/executorch/../executorch/extension/tensor/tensor_ptr.h:81:18: note: candidate function not viable: no known conversion from 'executorch::runtime::TensorShapeDynamism' to 'executorch::aten::Device' (aka 'executorch::runtime::etensor::Device') for 4th argument
inline TensorPtr make_tensor_ptr(
^

[Copilot]

🟡 Changes recommended

A critical compilation issue and two runtime correctness and memory-safety issues remain unresolved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (2)

extension/training/optimizer/adamw.cpp:36

• These helpers unconditionally reinterpret every tensor as float, but the public AdamW API accepts executorch::aten::Tensor and step() never checks its scalar type. Passing a valid double/half/bfloat16 parameter or gradient therefore uses 4-byte indexing against data with a different element size and corrupts the update instead of performing AdamW. Dispatch over supported floating types or reject non-float inputs before updating.

  auto a_ptr = a.const_data_ptr<float>();
  auto b_ptr = b.const_data_ptr<float>();
  auto out_ptr = out.mutable_data_ptr<float>();

extension/training/optimizer/adamw.cpp:164

• The state allocations are unchecked. If either malloc returns null, the following memset dereferences it, and if only the second allocation fails the first buffer is leaked; this is especially problematic for on-device training where step() already has an error-return path. Check both pointers (while allowing zero-byte tensors), free any partial allocation, and return Error::MemoryAllocationFailed before clearing the buffers.

        void* m_buf_ptr = malloc(g.nbytes());
        void* v_buf_ptr = malloc(g.nbytes());
        std::memset(m_buf_ptr, 0, g.nbytes());
        std::memset(v_buf_ptr, 0, g.nbytes());

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

[psiddh]

Pushed two commits to this branch (maintainer_can_modify was on, so no new PR needed): a merge of current main, and a fix for the build break. Your commits are untouched.

The compile error. make_tensor_ptr gained a Device parameter after this branch 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. Both moment buffers come from malloc, so they get an explicit CPU device. I also added the null check Copilot flagged on those same two allocations — they were used unconditionally, so a failure dereferenced null instead of returning an error. The other Copilot points (dtype/shape validation, parameter range checks) I left alone; your SGD-parity argument holds and @JacobSzwejbka said to discount Copilot.

Verification. I reproduced the exact CI error locally, confirmed the fix clears it, and lintrunner is clean. I could not run adamw_test here — 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. Please still run buck2 test //extension/training/optimizer/test:adamw_test on your side to confirm the real binary.

Still blocking: EasyCLA. Commits 5f1c90c and bda405f are not authorized under a signed CLA, so this cannot merge regardless of CI. @BryanBradfo that one needs you — https://api.easycla.lfx.linuxfoundation.org/v2/repository-provider/github/sign/29311126/463628495/18848/#/?version=2

cc @JacobSzwejbka

The fix commit was authored with Claude Code.

[Copilot]

🟡 Changes recommended

Three moderate correctness and robustness issues in the AdamW implementation remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (3)

extension/training/optimizer/adamw.cpp:155

• This mutates p before the only allocation failure that step() reports. If either malloc below returns null, the call returns MemoryAllocationFailed after applying weight decay; retrying the step then decays this parameter again even though its moment update never ran. Allocate/initialize the per-parameter state before changing p (or preflight all state allocations before the update pass).

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

extension/training/optimizer/adamw.cpp:147

• The update helpers always access tensors as float, but this path constructs state using g.scalar_type() and never rejects other dtypes. A Double parameter/gradient is therefore partially written/read as 32-bit values, while Half can be read past its buffer. Validate both tensors as Float before applying weight decay, or implement dtype-specific updates.

      auto g = named_gradient->second;
      auto p = param_iter->second;

extension/training/optimizer/adamw.cpp:167

• malloc(0) is permitted to return nullptr, and ExecuTorch supports valid zero-element tensors. With an empty parameter/gradient this branch can incorrectly return MemoryAllocationFailed; if zero bytes are allowed, the following memset calls must also be skipped rather than passing a null pointer. Treat allocation as failed only when g.nbytes() != 0 and guard zero-byte initialization.

        void* m_buf_ptr = malloc(g.nbytes());
        void* v_buf_ptr = malloc(g.nbytes());
        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: 2
• Review effort level: Lite

[Copilot]

🟡 Changes recommended

Address the AdamW input-safety, allocation-ordering, empty-tensor, and test-coverage issues.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (3)

extension/training/optimizer/adamw.cpp:166

• g.nbytes() is zero for valid zero-element tensors, which ExecuTorch supports. On an allocator where malloc(0) returns nullptr, this branch reports MemoryAllocationFailed even though no state storage is needed; the zero-length memset should also be skipped. Gate allocation failure and initialization on a nonzero byte count so empty parameters remain a no-op.

        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:154

• Weight decay mutates p before the lazy state allocation below. If either allocation fails, step() returns MemoryAllocationFailed after changing this parameter (and possibly earlier ones), so retrying can apply decay twice and leave the step partially applied. Initialize the state before mutating parameters, then apply decay immediately before the moment update.

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

extension/training/optimizer/test/adamw_test.cpp:88

• These behavioral tests use only constant ±1 gradients and a 0.1 tolerance. That makes the normalized update nearly sign-only, so substantial errors in the moment recurrence, bias correction, or epsilon placement can still pass; the claimed six-decimal PyTorch cross-check is not encoded here. Add a varying-gradient multi-step reference case with tight expected values.

  EXPECT_NEAR(p1[0], 2.0, 0.1);

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

[psiddh]

Opened #22857 carrying this work forward, with credit to @BryanBradfo in both the commit message and the PR description.

To be explicit about why it is a squash under a different author rather than a re-post of your commits: EasyCLA keys off commit authorship, so any PR containing 5f1c90c and bda405f fails the same check. That is the only reason the authorship differs — the implementation, the tests and the design calls are all yours.

@BryanBradfo if you sign the CLA, this PR is the better one to land and I will close #22857. The build fix you need is already pushed here, and CI is green apart from the CLA gate:
https://api.easycla.lfx.linuxfoundation.org/v2/repository-provider/github/sign/29311126/463628495/18848/#/?version=2

cc @JacobSzwejbka — your call which one you would rather take.

[Copilot]

🔵 Needs a closer look

Address the three moderate correctness and failure-handling issues in adamw.cpp.

Review details

Suppressed comments (3)

Previously missed (1) — in code that hasn't changed since the last review.

extension/training/optimizer/adamw.cpp:169

• malloc(0) is allowed to return nullptr, so an empty parameter/gradient can incorrectly return MemoryAllocationFailed even though no state storage is needed. ExecuTorch's allocation paths treat null data as valid when nbytes() == 0; handle zero-byte buffers separately and only treat null as failure for nonzero allocations.

extension/training/optimizer/adamw.cpp:153

• Weight decay mutates p before lazy state allocation. If either allocation fails at lines 163-166, step() returns MemoryAllocationFailed after changing the parameter, so retrying applies decay again without a moment update and a multi-parameter step can be partially applied. Allocate/validate the state before mutating parameters, or make this failure path transactional.

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

extension/training/optimizer/adamw.cpp:146

• These helpers unconditionally read and write float elements, but this path accepts arbitrary parameter/gradient dtypes; it also allocates state from g.nbytes() while the final update iterates p.numel(). A non-Float tensor or a gradient with a different numel can therefore produce incorrect results or out-of-bounds reads. Reject unsupported dtypes and mismatched sizes before applying decay.

      auto g = named_gradient->second;
      auto p = param_iter->second;

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

[Copilot]

🔵 Needs a closer look

Two moderate issues remain in adamw.cpp involving dtype validation and zero-element tensor allocation.

Review details

Suppressed comments (2)

Previously missed (1) — in code that hasn't changed since the last review.

extension/training/optimizer/adamw.cpp:164

• malloc(0) is permitted to return nullptr, and ExecuTorch supports zero-numel tensors. For an empty parameter/gradient this branch can therefore return MemoryAllocationFailed before the otherwise no-op step, depending on the allocator/platform. Gate the allocation failure check and memset on g.nbytes() != 0 (or skip state allocation for empty tensors).

extension/training/optimizer/adamw.cpp:146

• These helpers unconditionally reinterpret tensors as float, but step() does not validate the parameter or gradient dtype and allocates the state using g.scalar_type(). A Half/BFloat16 input makes the state buffer smaller than the subsequent mutable_data_ptr<float>() writes, while Double or integer inputs are read with the wrong representation, causing memory corruption or incorrect updates. Reject non-Float tensors here (and document the supported dtype), or add dtype-specific kernels.

      auto g = named_gradient->second;
      auto p = param_iter->second;

• Files reviewed: 7/7 changed files
• Comments generated: 0 new
• 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.