Remember the first PID gradient for the next derivative

[shaneraphel]

Summary

The first step of PID did this:

g_buf = param_state["grad_buffer"] = torch.zeros_like(p)
g_buf = d_p

The second assignment rebinds the local name. The buffer stored in the state stays zero, and the derivative on that step is d_p - d_p = 0. The following step still reads the zero buffer, so its derivative is taken against 0 rather than against the gradient from the previous step. The same two lines are in the cited reference, tensorboy/PIDOptimizer. Copying them preserved the missed store.

The gradient is now copied into the buffer. The first derivative stays 0, because there is no earlier gradient. With momentum 0.5, integral 0, derivative 10 and learning rate 1, gradients 1 then 3 move a parameter that starts at 0 to -1 and then to -14. Against a zero previous gradient the second step would have been -19.

Test plan

• pytest tests/test_pid_grad_buffer.py — after the first step the buffer equals 1 and the parameter is -1; after the second step the parameter is -14.

This change was produced with assistance from an AI coding tool. I traced the rebinding, checked the second step against the finite difference of the two gradients, and ran the test locally.