netcl.optim — Optimizers & Schedules

netcl.optim provides the parameter-update machinery that turns accumulated gradients into actual Tensor mutations. It mirrors torch.optim: a step() call applies the update rule, a zero_grad() clears the gradient buffers, and a separate scheduler family adjusts the learning rate over epochs.

All public symbols are re-exported from the package root:

from netcl.optim import (
    SGD, Adam, AdamW, Momentum, RMSProp,
    CosineAnnealingLR, ReduceLROnPlateau, WarmupCosine,
    clip_grad_norm, clip_grad_norm_device,
    AMPGradScaler,
)

Optimizer Index

SGD

from netcl.optim import SGD

opt = SGD(
    params,
    lr=1e-2,
    momentum=0.0,
    weight_decay=0.0,
)

Update rule (per parameter θ, with gradient g):

• If weight_decay > 0: g ← g + weight_decay * θ
• If momentum > 0: v ← momentum * v + g g ← v
• θ ← θ - lr * g

The default (momentum=0.0) reduces to vanilla gradient descent.

Momentum

Classic Polyak momentum. The implementation is identical to SGD with momentum > 0 and dampening=0.0, but the constructor only exposes the momentum-style knobs:

from netcl.optim import Momentum

opt = Momentum(params, lr=1e-2, momentum=0.9, weight_decay=1e-4)

Adam / AdamW

Adam keeps two exponential moving averages per parameter — the first moment m and the uncentered second moment v — and computes a bias-corrected adaptive step.

from netcl.optim import Adam, AdamW

opt  = Adam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0)
optw = AdamW(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01)

Per parameter θ with gradient g and step t:

• m ← β1 * m + (1 - β1) * g
• v ← β2 * v + (1 - β2) * g²
• m̂ ← m / (1 - β1^t)
• v̂ ← v / (1 - β2^t)
• θ ← θ - lr * m̂ / (sqrt(v̂) + eps)

In Adam, weight_decay is applied as an L2 penalty on the gradient (the original formulation); in AdamW the decay is decoupled from the gradient update — it is applied directly to the weights as θ ← θ - lr * (m̂ / (sqrt(v̂) + eps) + weight_decay * θ). Decoupled weight decay is the modern default for transformer-style training and is the recommended variant unless you have a specific reason to use the L2-in-gradient form.

Fused Clip + Step: AdamW.step(max_norm=...)

AdamW.step accepts an optional max_norm keyword. When given, gradient-norm clipping is fused into the same device-side pass as the parameter update, instead of being a separate clip_grad_norm_device call before step():

opt = AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

with ag.Tape() as tape:
    loss = ...
tape.backward(loss)
opt.step(max_norm=1.0)   # clip-then-update, single fused dispatch
opt.zero_grad()

This is the recommended pattern over clip_grad_norm_device(params, max_norm=1.0); opt.step() whenever it applies — it is strictly faster and never less correct, because the norm reduction and the update read the same gradient buffer without an intervening host round-trip.

The fused path activates automatically, with no code change beyond passing max_norm=, when:

• the fused path is enabled (NETCL_ADAMW_FUSED=1, the default), and
• at least NETCL_ADAMW_FUSED_MIN_PARAMS parameters have a gradient (default 8), and
• the total element count across those parameters is at least NETCL_ADAMW_FUSED_MIN_ELEMS (default 16384).

Below those thresholds (e.g. a tiny model, or a handful of parameters), step() transparently falls back to the per-parameter kernel path — max_norm still clips correctly, it is just not worth gathering into a flat buffer for a handful of small tensors. Trainer detects at construction time whether the optimizer's step accepts max_norm (via inspect.signature) and prefers this fused path automatically when grad_clip is set.

RMSProp

Per-parameter adaptive learning rate using a moving average of squared gradients:

from netcl.optim import RMSProp

opt = RMSProp(
    params,
    lr=1e-3,
    alpha=0.99,
    eps=1e-8,
    weight_decay=0.0,
    momentum=0.0,
)

Per parameter θ with gradient g:

• v ← alpha * v + (1 - alpha) * g²
• θ ← θ - lr * g / (sqrt(v) + eps)

Set momentum > 0 to add a momentum term on top of the RMSProp update.

Gradient Clipping

Two helpers, both in netcl.optim.clip, that operate in place on the Tensor.grad fields of the parameters you pass in. Call them after tape.backward(loss) and before opt.step().

from netcl.optim import clip_grad_norm, clip_grad_norm_device

# Host-side: pull grads to NumPy, scale, copy back.
clip_grad_norm(parameters, max_norm=1.0)

# Device-side: runs as an OpenCL kernel — avoids the H2D/D2H round-trip.
clip_grad_norm_device(parameters, max_norm=1.0)

clip_grad_norm computes the total L2-norm of the concatenated gradient vector and scales every gradient by min(1, max_norm / total_norm). clip_grad_norm_device does the same computation entirely on the device and is faster for large parameter counts.

LR Schedules

netcl schedulers are not tied to an optimizer — they are standalone objects that return the new learning rate from step(). You assign the returned value to opt.lr yourself:

opt.lr = sched.step()

CosineAnnealingLR

Smoothly anneals from max_lr down to min_lr over T_max epochs. Optionally starts with a linear warm-up phase via warmup_epochs.

from netcl.optim import Adam, CosineAnnealingLR

opt   = Adam(model.parameters(), lr=3e-4)
sched = CosineAnnealingLR(max_lr=3e-4, min_lr=1e-6, T_max=50)

for epoch in range(50):
    train(...)
    opt.lr = sched.step()   # returns the lr for the next epoch

The curve is lr = min_lr + (max_lr - min_lr) * (1 + cos(π * t / T_max)) / 2.

WarmupCosine

Linear warm-up for warmup_epochs epochs followed by cosine decay to min_lr. Call sched.lr(epoch) with the current epoch number (0-based):

from netcl.optim import Adam, WarmupCosine

opt   = Adam(model.parameters(), lr=3e-4)
sched = WarmupCosine(base_lr=3e-4, max_epochs=50, warmup_epochs=5, min_lr=1e-6)

for epoch in range(50):
    opt.lr = sched.lr(epoch)   # set lr before the training step
    train(...)

ReduceLROnPlateau

The only metric-based scheduler. It watches a scalar (typically validation loss) and reduces the LR when the metric has stopped improving for patience epochs. step(metric) returns the (possibly reduced) learning rate.

from netcl.optim import Adam, ReduceLROnPlateau

opt     = Adam(model.parameters(), lr=1e-3)
plateau = ReduceLROnPlateau(
    base_lr=1e-3,
    factor=0.5,      # new_lr = old_lr * factor
    patience=5,      # epochs without improvement before decay
    min_lr=1e-6,
    threshold=1e-4,
    mode="min",      # "min" or "max"
)

for epoch in range(epochs):
    val_loss = validate(...)
    opt.lr = plateau.step(val_loss)

Set mode="max" for accuracy-like metrics.

AMPGradScaler & AMP

Mixed-precision training scales the loss to keep fp16 gradients representable. AMPGradScaler is re-exported from netcl.optim; the underlying class lives in netcl.amp.

from netcl.optim import AMPGradScaler
from netcl.amp import autocast
import netcl.autograd as ag

scaler = AMPGradScaler()

for x, y in loader:
    with ag.Tape() as tape:
        with autocast():
            logits = model(x)
            loss = loss_fn(logits, y)
    tape.backward(scaler.scale(loss))
    scaler.step(opt)
    scaler.update()
    opt.zero_grad()

The full AMP contract, including the autocast context manager and the device support matrix, is documented on its own page.

Full Training Loop

import netcl.autograd as ag
from netcl.nn import functional as F
from netcl.optim import AdamW, CosineAnnealingLR, clip_grad_norm

opt   = AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
sched = CosineAnnealingLR(max_lr=3e-4, min_lr=1e-6, T_max=20)

for epoch in range(20):
    for x, y in loader:          # x, y are Tensors
        with ag.Tape() as tape:
            logits = model(ag.tensor(x))   # wrap input for autograd
            loss = F.cross_entropy(logits, y)
        tape.backward(loss)
        clip_grad_norm(model.parameters(), max_norm=1.0)
        opt.step()
        opt.zero_grad()
    opt.lr = sched.step()
    print(f"epoch {epoch+1}  loss={loss.value.to_host()[0]:.4f}  lr={opt.lr:.2e}")

Distributed Notes

Each worker in a Data Parallel setup holds its own copy of every optimizer. The distributed API takes care of gradient all-reduce before the optimizer's step() is called, so a parameter update is identical to the single-process case. Parameter sharding (ZeRO-style) is not part of the optimizer contract; if you need it, wrap the parameters in sharded views before handing them to the optimizer constructor.

For an end-to-end multi-device example, see Data Parallel.

See also

• MNIST with MLP — end-to-end training loop using AdamW, gradient clipping, and a cosine schedule.
• Data Parallel — using these optimizers in a multi-process setup.
• autograd API — the Tape that produces the gradients passed into opt.step().
• amp API — the GradScaler and autocast context that wrap these optimizers for mixed precision.
• distributed API — gradient all-reduce in front of opt.step().
• Tensor — the actual Tensor objects that the optimizers mutate in place.
• Trainer — the training-loop wrapper that prefers AdamW's fused step(max_norm=...) path automatically when available.