Class: NanoGPT::Layers::Block

Inherits:
Torch::NN::Module
  • Object
show all
Defined in:
lib/nano_gpt/layers/block.rb

Overview

Transformer block: LayerNorm -> Attention -> LayerNorm -> MLP

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(config) ⇒ Block

Returns a new instance of Block.



9
10
11
12
13
14
15
 
# File 'lib/nano_gpt/layers/block.rb', line 9

def initialize(config)
  super()
  @ln_1 = LayerNorm.new(config.n_embd, bias: config.bias)
  @attn = CausalSelfAttention.new(config)
  @ln_2 = LayerNorm.new(config.n_embd, bias: config.bias)
  @mlp = MLP.new(config)
end

Instance Attribute Details

#attn ⇒ Object (readonly)

Returns the value of attribute attn.



7
8
9
 
# File 'lib/nano_gpt/layers/block.rb', line 7

def attn
  @attn
end

Instance Method Details

#crop_mask(block_size) ⇒ Object 



25
26
27
 
# File 'lib/nano_gpt/layers/block.rb', line 25

def crop_mask(block_size)
  @attn.crop_mask(block_size)
end

#forward(x) ⇒ Object 



17
18
19
20
21
22
23
 
# File 'lib/nano_gpt/layers/block.rb', line 17

def forward(x)
  x = x + @attn.call(@ln_1.call(x))
  x = x + @mlp.call(@ln_2.call(x))
  # Trigger GC to free intermediate tensors (critical for torch.rb memory management)
  GC.start(full_mark: false, immediate_sweep: true)
  x
end