rwkv_module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include <initializer_list>
#include <vector>
 
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class RWKVModule : public Module {
public:
    RWKVModule(int64_t d_model, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Recurrent; }
 
    void set_W_r(std::initializer_list<float> values);
    void set_W_k(std::initializer_list<float> values);
    void set_W_v(std::initializer_list<float> values);
    void set_W_o(std::initializer_list<float> values);
    void set_w(std::initializer_list<float> values);
    void set_u(std::initializer_list<float> values);
    void set_mu_r(std::initializer_list<float> values);
    void set_mu_k(std::initializer_list<float> values);
    void set_mu_v(std::initializer_list<float> values);
 
    void set_W_r(const std::vector<float>& values);
    void set_W_k(const std::vector<float>& values);
    void set_W_v(const std::vector<float>& values);
    void set_W_o(const std::vector<float>& values);
    void set_w(const std::vector<float>& values);
    void set_u(const std::vector<float>& values);
    void set_mu_r(const std::vector<float>& values);
    void set_mu_k(const std::vector<float>& values);
    void set_mu_v(const std::vector<float>& values);
 
    [[nodiscard]] const Tensor& W_r() const { return w_r_; }
    [[nodiscard]] const Tensor& W_k() const { return w_k_; }
    [[nodiscard]] const Tensor& W_v() const { return w_v_; }
    [[nodiscard]] const Tensor& W_o() const { return w_o_; }
    [[nodiscard]] const Tensor& w() const { return w_; }
    [[nodiscard]] const Tensor& u() const { return u_; }
    [[nodiscard]] const Tensor& mu_r() const { return mu_r_; }
    [[nodiscard]] const Tensor& mu_k() const { return mu_k_; }
    [[nodiscard]] const Tensor& mu_v() const { return mu_v_; }
 
    [[nodiscard]] const Tensor& W_r_grad() const { return w_r_grad_; }
    [[nodiscard]] const Tensor& W_k_grad() const { return w_k_grad_; }
    [[nodiscard]] const Tensor& W_v_grad() const { return w_v_grad_; }
    [[nodiscard]] const Tensor& W_o_grad() const { return w_o_grad_; }
    [[nodiscard]] const Tensor& w_grad() const { return w_grad_; }
    [[nodiscard]] const Tensor& u_grad() const { return u_grad_; }
    [[nodiscard]] const Tensor& mu_r_grad() const { return mu_r_grad_; }
    [[nodiscard]] const Tensor& mu_k_grad() const { return mu_k_grad_; }
    [[nodiscard]] const Tensor& mu_v_grad() const { return mu_v_grad_; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override {
        return {
            {"w_r", {&w_r_, &w_r_grad_}},
            {"w_k", {&w_k_, &w_k_grad_}},
            {"w_v", {&w_v_, &w_v_grad_}},
            {"w_o", {&w_o_, &w_o_grad_}},
            {"w", {&w_, &w_grad_}},
            {"u", {&u_, &u_grad_}},
            {"mu_r", {&mu_r_, &mu_r_grad_}},
            {"mu_k", {&mu_k_, &mu_k_grad_}},
            {"mu_v", {&mu_v_, &mu_v_grad_}},
        };
    }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return backend_->device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    int64_t d_model_;
    DeviceBackend* backend_;
    Tensor w_r_;   // (d_model, d_model)
    Tensor w_k_;   // (d_model, d_model)
    Tensor w_v_;   // (d_model, d_model)
    Tensor w_o_;   // (d_model, d_model)
    Tensor w_;     // (d_model,); decay = exp(-w)
    Tensor u_;     // (d_model,)
    Tensor mu_r_;  // (d_model,)
    Tensor mu_k_;  // (d_model,)
    Tensor mu_v_;  // (d_model,)
    Tensor w_r_grad_;
    Tensor w_k_grad_;
    Tensor w_v_grad_;
    Tensor w_o_grad_;
    Tensor w_grad_;
    Tensor u_grad_;
    Tensor mu_r_grad_;
    Tensor mu_k_grad_;
    Tensor mu_v_grad_;
    Tensor last_input_;   // (N, L, d_model)
    Tensor last_xr_;      // (N, L, d_model); token-shifted receptance input
    Tensor last_xk_;      // (N, L, d_model)
    Tensor last_xv_;      // (N, L, d_model)
    Tensor last_r_;       // (N, L, d_model); sigmoid POST-activation (its own derivative source)
    Tensor last_k_;       // (N, L, d_model)
    Tensor last_v_;       // (N, L, d_model)
    Tensor last_e_;       // (N, L, d_model); exp(u + k_t)
    Tensor last_kk_;      // (N, L, d_model); exp(k_t)
    Tensor last_num_;     // (N, L, d_model)
    Tensor last_den_;     // (N, L, d_model)
    Tensor last_wkv_;     // (N, L, d_model)
    Tensor last_a_;       // (N, L+1, d_model); index 0 along dim 1 = a_0 = 0
    Tensor last_b_;       // (N, L+1, d_model); index 0 along dim 1 = b_0 = 0
    int64_t last_L_ = 0;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix