spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/tensor.h"56#include <functional>7#include <memory>8#include <vector>910namespace forge {1112// A Var is a Tensor plus (lazily allocated) gradient storage. Copies share13// the same impl, so parameters handed to modules, the optimizer and the14// tape all see one .grad — which is also what makes tied embeddings and15// gradient accumulation across micro-batches work for free.16class Var {17public:18    Var() = default;19    explicit Var(Tensor value, bool requires_grad = false)20        : impl_(std::make_shared<Impl>(Impl{std::move(value), Tensor{}, requires_grad})) {}2122    bool defined() const { return impl_ != nullptr; }23    // Stable identity of the underlying storage — two Vars with the same id24    // are the same parameter (tied embeddings/lm_head).25    const void* id() const { return impl_.get(); }26    Tensor& value() const { return impl_->value; }27    bool requires_grad() const { return impl_ && impl_->requires_grad; }2829    // Gradient, allocated as zeros on first touch. Grads are f32 always.30    Tensor& grad() const {31        if (!impl_->grad.defined()) {32            impl_->grad = Tensor::zeros(impl_->value.shape(), DType::F32);33        }34        return impl_->grad;35    }36    bool has_grad() const { return impl_ && impl_->grad.defined(); }37    void zero_grad() const {38        if (impl_ && impl_->grad.defined()) impl_->grad.zero_();39    }4041    // Shape-only view: shares BOTH value and grad storage with this Var, so42    // gradient flow needs no tape node.43    Var reshaped(std::vector<int64_t> shape) const {44        Var out(value().view(shape), requires_grad());45        if (requires_grad()) out.impl_->grad = grad().view(std::move(shape));46        return out;47    }4849private:50    struct Impl {51        Tensor value;52        Tensor grad;53        bool requires_grad;54    };55    std::shared_ptr<Impl> impl_;56};5758// Dynamic tape. Ops that produce grad-requiring outputs push a backward59// lambda; backward() runs them in reverse and clears the tape. Not60// thread-safe by design: one training thread.61class Tape {62public:63    static Tape& get();6465    bool enabled() const { return enabled_; }66    void set_enabled(bool e) { enabled_ = e; }6768    void record(std::function<void()> backward_fn) {69        if (enabled_) nodes_.push_back(std::move(backward_fn));70    }7172    // Seeds d(loss)/d(loss) = 1 and walks the tape in reverse.73    void backward(const Var& loss);7475    void clear() { nodes_.clear(); }76    size_t size() const { return nodes_.size(); }7778private:79    Tape() = default;80    std::vector<std::function<void()>> nodes_;81    bool enabled_ = true;82};8384// RAII scope that disables tape recording (inference / generation).85struct NoGrad {86    NoGrad() : prev_(Tape::get().enabled()) { Tape::get().set_enabled(false); }87    ~NoGrad() { Tape::get().set_enabled(prev_); }88    NoGrad(const NoGrad&) = delete;89    NoGrad& operator=(const NoGrad&) = delete;9091private:92    bool prev_;93};9495} // namespace forge96