transformer_block.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include "pulsatrix/module.hpp"
#include "pulsatrix/multihead_attention_module.hpp"
#include "pulsatrix/rms_norm_module.hpp"
#include "pulsatrix/swiglu_module.hpp"
 
namespace pulsatrix {
 
class TransformerBlock : public Module {
public:
    TransformerBlock(int64_t d_model, int64_t num_heads, int64_t d_ff, DeviceBackend* backend, bool use_rope = true,
                      bool use_qk_norm = false);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Elementwise; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override;
 
    void set_training(bool training) override;
 
    [[nodiscard]] int64_t d_model() const { return mha_.d_model(); }
 
    [[nodiscard]] RMSNormModule& norm1() { return norm1_; }
    [[nodiscard]] MultiHeadAttentionModule& mha() { return mha_; }
    [[nodiscard]] RMSNormModule& norm2() { return norm2_; }
    [[nodiscard]] SwiGLUModule& swiglu() { return swiglu_; }
 
 
    [[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_;
 
    RMSNormModule norm1_;
    MultiHeadAttentionModule mha_;
    RMSNormModule norm2_;
    SwiGLUModule swiglu_;
 
    // Forward caches -- the four operands the two residual splits need.
    Shape last_input_shape_ = Shape({0});
    Tensor last_x_;         
    Tensor last_attn_out_;  
    Tensor last_y1_;        
    Tensor last_ffn_out_;   
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix