A reverse-mode automatic-differentiation engine with custom gradient operators and CUDA kernels for neural-network training.
When
Feb 2024 – Mar 2024
Type
Side build
Reverse mode on L(w·x + b): one forward pass, then gradients back along every edge
What I built8
Built the tape: every forward operation records its inputs and a local gradient rule, so backward is a reverse walk over the recorded graph rather than a hand-derived formula per model.
Implemented the operator set with its adjoints - add, multiply, matmul, transpose, reshape, ReLU, softmax and cross-entropy - each one a forward rule plus a vector-Jacobian product.
Got broadcasting right, which is where a hand-rolled engine usually breaks: a gradient flowing back into a broadcast dimension has to be summed over that axis or the shapes silently stop matching.
Wrote CUDA kernels for the matmul and elementwise paths, with the reduction in the backward pass done in shared memory rather than with an atomic per element.
Verified every adjoint against central-difference numerical gradients before trusting a single training run, since a wrong gradient trains to a worse loss instead of crashing.
Accumulated rather than overwrote gradients at nodes with multiple consumers, the bug that makes a network with any weight reuse train subtly wrong.
Trained MLPs end to end against a reference implementation, matching loss curves step for step on a fixed seed as the correctness bar.
Freed the tape after backward so a training loop does not retain every intermediate for the whole run, which is the difference between a toy and something that finishes an epoch.