An optimizer keeps param_groups, lists of parameters with their settings such as lr, and state, a dictionary of buffers for each parameter. step() reads each gradient and changes the parameter in place, under torch.no_grad() so the update itself is not recorded.

import torch

class PlainSGD(torch.optim.Optimizer):
    def __init__(self, params, lr):
        super().__init__(params, {"lr": lr})

    @torch.no_grad()
    def step(self):
        for g in self.param_groups:
            for p in g["params"]:
                if p.grad is not None:
                    p -= g["lr"] * p.grad

w = torch.tensor([3.0], requires_grad=True)
opt = PlainSGD([w], lr=0.1)
for _ in range(50):
    opt.zero_grad()
    (w ** 2).sum().backward()
    opt.step()
print(w)

Momentum keeps a velocity per parameter: v = momentum * v + grad, then p -= lr * v. Steps that agree build up speed, steps that alternate cancel out. The velocity must live in self.state[p], which survives from one step to the next and is saved with the optimizer's state_dict.

Write Momentum(params, lr, momentum=0.9), a torch.optim.Optimizer subclass whose step() keeps each parameter's velocity in self.state[p]["velocity"] (starting from zero), updates it to momentum * v + grad and subtracts lr * v, skipping parameters without a gradient. It must match torch.optim.SGD(params, lr, momentum).

The explanation opens here once every check passes.