Skip to content

ROADMAP.md

Latest commit

 

History

History
232 lines (202 loc) · 14.1 KB

ROADMAP.md

File metadata and controls

232 lines (202 loc) · 14.1 KB
Raw

pycograd — Roadmap

Working name. A small, readable reverse-mode automatic-differentiation library for inventing neural-net architectures, grown out of the autodiff example in pyccolo (pyccolo/examples/autodiff.py).

Vision

A NumPy-backed autograd you can actually read, that scales down to "explain backprop on one slide" and up, eventually, to training real models. The reverse-mode core is already correct and tiny; pycograd is about building the framework around it without losing the legibility.

Where this comes from (current state)

The starting point is the pyccolo autodiff example, which already has:

  • A reverse-mode tape node (Var) wrapping a NumPy array, with operator overloading, broadcasting (gradients reduced back over broadcast axes), and VJP rules for the common elementwise / reduction / linear-algebra / shape ops.
  • Indexed reads (gather forward, scatter-add backward), where/clip, var/std, max/min with subgradients, concatenate, detach.
  • value_and_grad / grad, a plain-SGD gradient_descent, and demos that train logistic regression, an MLP, an MLP+LayerNorm+Dropout, and a single-head Transformer block — each gradient-checked against finite differences.
  • Transparent interception of numpy/math calls via pyccolo's before_call (so you write ordinary np.exp(x) and it differentiates), helper instrumentation on demand, and a pipescript |> integration.

So the autograd core ships working. Everything below is the surrounding system.

Guiding decision: drop the interception layer from the core

pyccolo's before_call interception is a great teaching device — it shows that operator overloading alone can't make np.exp(x) differentiable and that a tracer can fix that transparently. But it is the wrong foundation for a real library: it only works inside an instrumented function under the tracer, and it adds per-call overhead.

pycograd's Tensor (the productized Var) should instead implement NumPy's own extension protocols — __array_ufunc__ and __array_function__ — so a Tensor is differentiable everywhere, with no tracer required. (The example deliberately sets __array_ufunc__ = None to motivate pyccolo; pycograd reverses that.)

pyccolo can remain an optional "transparent mode" for niceties the protocols can't reach — routing scalar math.* through tensor ops, differentiating through un-annotated helper bodies, and the pipescript |> syntax — but nothing in the core should depend on it.

The existing demos (logistic regression → MLP → LayerNorm/Dropout → Transformer) become the first integration tests / examples and the bar for "don't regress."

Roadmap

Phased so each layer rests on the one before it. Each item is tagged with how crucial it is to the goal (developing novel architectures and/or training at scale) and a rough feasibility / effort. Note the tension: the single most crucial thing for scale (hardware + performance) is intentionally last because it is gated by the foundation.

Phase 0 — Extraction & foundation

  • Carve out into pycograd. Lift Var + the VJP rules out of the pyccolo example into a standalone package (pycograd.tensor, pycograd.ops), with its own tests/CI. Crucial: prerequisite. Feasibility: high (days).
  • A real Tensor type via NumPy protocols. Implement __array_ufunc__ / __array_function__ so np.* on a Tensor dispatches to our VJPs natively; drop the __array_ufunc__ = None fail-loud stance. Crucial: foundational. Feasibility: medium (the protocols are fiddly but well-documented).
  • dtype & a device seam. The array-backend seam is done: device("cupy") swaps the array library the tape, primitives, and optimizers compute with (NumPy default, CuPy for GPU), so a net trains on-device unchanged. Still open: dtype — stop forcing float64, track dtype, default float32, support bf16/float16. Crucial: high (memory + speed + the GPU on-ramp). Feasibility: medium.
  • Keep pyccolo as optional transparent mode (math.* routing, helper instrumentation, pipescript |>), gated behind an extra. Crucial: low. Feasibility: high (already built).

Phase 1 — A usable research framework (highest impact-per-effort)

  • Module / pytree parameter abstraction. Named, nested params + state; stop threading bare arrays positionally (the Transformer demo takes 14 positional args). Crucial: highest for developing architectures. Feasibility: high (pure Python).
  • Real optimizers. Adam/AdamW, SGD+momentum, weight decay, gradient clipping, LR schedules, zero_grad. Crucial: high. Feasibility: high.
  • Fused, numerically-stable primitives. Landed. log_softmax, logsumexp, softmax, cross_entropy, the activation family, and the normalization layers are first-class ops in pycograd/functional.py: layer_norm, rms_norm, group_norm, instance_norm, and batch_norm (with running stats, state-in/state-out (y, new_mean, new_var)). batch_norm's running stats ride a new mutable-buffer leaf (Param.mutable + buffer[...]
    • ParamDict.update_buffers): non-trainable, optimizer-skipped, advanced out-of-band. Also landed: scaled_dot_product_attention, multi_head_attention, embedding, linear, dropout, and conv_transpose2d / upsample_nearest2d. Open: a forward scatter-add primitive (x.at[i].add(...), Phase 2) — until then conv_transpose2d is built from an einsum input-dilation + flipped conv rather than a col2im scatter.
  • Broader op coverage. einsum (general contractions / attention variants), gather/scatter + embeddings, sort/argsort/topk, cumsum, padding, one-hot. Crucial: high for novel architectures. Feasibility: medium (einsum is real work).
  • Convolution / pooling. Conv1d/2d, pooling. Crucial: high (CNNs). Feasibility: medium — naive is easy but slow; fast conv = im2col/FFT.
  • RNG management. Landed (pycograd/random.py): a pure, splittable JAX-style key API — key / split / fold_in plus bernoulli / uniform / normal / randint, built on numpy's counter-based SeedSequence + Philox. dropout takes an explicit key (or rng) and no longer falls back to a hidden global. Keys are sampled host-side (a per-backend RNG seam is still open); vmap over a per-row key is not supported (host RNG can't consume a batched-tracer key), but a single key already gives independent per-sample masks over a (B, …) batch.
  • Training glue. Data loading/batching/shuffling, checkpoint save/load (params + optimizer state), basic metrics/logging. Crucial: medium. Feasibility: high (orthogonal).

Phase 2 — Depth & memory (correctness at real model sizes)

  • Gradient checkpointing. Landed (pycograd/checkpoint.py): checkpoint(f) wraps a segment so its intermediate activations are dropped on the forward and rematerialized in backward by re-running f -- trading ~one extra forward for a peak-memory drop from "every segment at once" to "one segment at a time". A single boundary node stands in for outputs = f(inputs); its value is the flat concatenation of the output leaves and the user-visible outputs are real slice+reshape views, so the multi-output cotangent join rides the existing VJPs. Backward lifts the saved input/weight values into fresh leaves, re-runs f, contracts each output leaf with its cotangent, and scatters the recomputed grads onto the boundary's parents. Works with positional grad/value_and_grad and the ambient weights.grad path (ambient weights enter by globals and the live binding is gone by backward time, so checkpoint discovers the weight Vars a segment touches via a small active-bindings registry and re-binds them for the remat), over arbitrary pytree outputs, and nests. It also saves memory under vmap: checkpoint lowers the batch into the boundary (vmap(checkpoint(f)) == checkpoint(vmap(f)) — physicalize the batched inputs, build the boundary one level down, re-wrap the batched outputs), so grad(vmap(checkpoint(f))), the per-sample vmap(grad(checkpoint(f))) (including per-sample gradients of a shared param), and nested vmap(vmap(...)) all rematerialize the batched activations. Under a live jvp the inputs are tangent tracers and a Var boundary can't be built without dropping the tangent axis, so checkpoint is transparent there (correct grads, no memory saving in that case); reverse-over-reverse of a checkpointed segment (grad(grad) / jacrev of its gradient) raises with a pointer to the forward-over-reverse jacfwd(grad) Hessian. Constraint: f must be deterministic in its inputs+weights (RNG/dropout-in-f replay is a follow-up).
  • Explicit tape lifetime. Free graphs deterministically; no_grad context; guard against the tape/_INSTRUMENTED caches growing unbounded. Crucial: medium. Feasibility: medium.
  • In-place / scatter updates. __setitem__, masked assignment, scatter-add for embeddings / KV-caches / efficient optimizer steps. Prefer a functional x.at[i].add(...) form to sidestep aliasing hazards (true in-place + autograd needs version tracking). Crucial: medium. Feasibility: medium.
  • Mixed-precision training. Loss scaling, autocast. Crucial: medium-high for scale. Feasibility: medium.

Phase 3 — Advanced autodiff transforms

  • Higher-order gradients. Landed. Var.backward has a differentiable path: when an enclosing transform is live it accumulates cotangents as level-connected Vars (a _VJP_FOR rule table riding bind) instead of mutating numpy .grad, while a single top-level grad/training run keeps the original raw path byte-for-byte (gated by trace.num_transform_levels()). This gives Hessians/HVPs via forward-over-reverse (jvp(grad(f)), jacfwd(grad(f))) and literal reverse-over-reverse (grad of a scalarized inner grad, jacrev(grad(f))) — reverse and forward Hessians agree and match finite differences; gradient-penalty losses work. Per-sample Hessians/HVPs also compose with vmap — vmap(jacfwd(grad(f)))/vmap(jacrev(grad(f))) and batched jvp(grad(f)) give per-example (B,n,n)/(B,n) results (realized as reverse-over-reverse over one batched forward, sidestepping the _unbroadcast/_KEEP_BATCH_AXES seam). Constraints, as in JAX: the per-example f is scalar-output and single-array (multi-arg/shared-param per-sample Hessians not yet wired); grad requires scalar output (grad(grad(f)) on a non-scalar inner is ill-posed and raises a clear error).
  • Forward-mode / jvp. Landed (pycograd/forward.py): jvp(f, primals, tangents)
    • jacfwd as a JVPTrace level riding bind; composes with vmap both ways and with itself (jvp(jvp(f)), second-order forward).
  • vmap (auto-batching) / per-sample gradients. Landed — including the composable v2. A trace-level interpreter stack (pycograd/trace.py: bind over a stack of Trace/Tracer levels, fed by pyccolo's operator/subscript interception) lets transforms be simultaneously live. So beyond single-level vmap/grad(vmap(f)), these all now work: vmap(vmap(f)) (a BatchTracer per level, batch axis materialized and moved-to-front per level), per-sample gradients of a shared parameter (vmap(grad(f)) returns (B, *w.shape) via a batched-cotangent backward / the shared param tiled per example), and gather both from batched data (vmap(lambda x, i: x[i])) and from a shared table (vmap(lambda i: E[i]), riding Var's scatter-add backward). The same bind/Trace machinery is the substrate for the trace-and-compile work below; the one piece deliberately left on the old path is ShapedArray/eval_shape (it still rides the backend swap — reframing it as a Tracer is a natural, separate follow-up). Forward-mode jvp + higher-order grad (next bullet) would slot in as another Trace level.

Phase 4 — Scale (the hard ceiling)

  • GPU / accelerator backend. The execution model is CPU-bound and allocates a Python Var per op (~2–4 orders of magnitude too slow/heavy for real training). A CuPy backend (its NumPy-mirroring API means much "just works") gets GPU at moderate effort. Crucial: highest for scale. Feasibility: medium for CuPy; the per-op Python overhead is not fixed by this.
  • Graph capture + compilation / fusion. Eliminate per-op Python overhead by tracing to a graph IR and compiling/fusing (XLA-style). This is a fundamentally different execution model from the eager tape and likely a separate engine. Crucial: highest for throughput. Feasibility: low (research-grade). The abstract-shape engine (shapes.py) is a down payment: its per-primitive abstract_eval rules already size every node without data, carry symbolic dims (_dims.Dim) for data-dependent shapes, and now run as a first-class AbstractTrace level on the same stack (ShapedArray is a Tracer) — so a capture tracer is just another Trace/Tracer reusing these rules, with vmap's batch axis one more dim in the same algebra.
  • Distributed / multi-device. Data / model / pipeline parallelism, collectives. Crucial: only for truly large scale. Feasibility: very low; pointless before the above.

Non-goals (for now)

  • A PyTorch-compatible API surface. Legibility over familiarity.
  • Beating PyTorch/JAX on speed. The point is a system you can fully read.
  • Being a hard dependency of pyccolo or pipescript — the relationship is reversed: pycograd may optionally use pyccolo for transparent mode.

Open questions

  • Eager tape vs. trace-and-compile: Phase 4 (and vmap) may not be reachable from the eager design without a second execution mode. Decide early whether pycograd is "always eager, readable" or grows a compiled path.
  • How much of the pyccolo transparent-interception story to keep front-and-center vs. relegate to an optional extra once Tensor works natively with NumPy.
  • Backend abstraction: design the array seam now (Phase 0) so CuPy/other backends drop in later without touching the VJP rules.