feat(core): support AdamW

[tyler-yankee]

See https://docs.pytorch.org/docs/stable/generated/torch.optim.AdamW.html.

[jlm429]

Thanks for putting this together!

Before merging, I’d like to run this through a few tests to make sure everything behaves as expected.

One thought for the longer term: it might be slightly cleaner to expose adamw as a first-class optimizer option rather than modifying Adam via a flag.

In the meantime, if anyone wants to experiment with AdamW behavior right away - you can add a local drop-in patch (for example pyperch_adamw_patch.py) in your working directory:

import torch
from pyperch.train.trainer import Trainer  # adjust if path differs

_original_init = Trainer.__init__

def _patched_init(self, model, loss_fn, config):
    _original_init(self, model, loss_fn, config)

    tc = getattr(config, "torch_config", None)
    if tc and getattr(tc, "optimizer", "").lower() == "adamw":
        if self.grad_params:
            self.torch_optim = torch.optim.AdamW(
                self.grad_params,
                lr=tc.lr,
                betas=tc.betas,
                eps=tc.eps,
                weight_decay=tc.weight_decay,
            )

Trainer.__init__ = _patched_init

Then import before constructing the Trainer:

import pyperch_adamw_patch

I’ll circle back once I’ve had a chance to test this more thoroughly. Thanks again for the contribution!

[tyler-yankee]

One thought for the longer term: it might be slightly cleaner to expose adamw as a first-class optimizer option rather than modifying Adam via a flag.

That's fair. As I have time over the next few days I'll revisit the implementation of this.

you can add a local drop-in patch

Thanks for offering this!

[tyler-yankee]

I've modified the PR to expose torch.optim.AdamW directly, with a new "adamw" option.