Class: Ignis::AI::NN::Linear

Inherits:
Module
  • Object
show all
Defined in:
lib/nnw/ai/nn/linear.rb

Overview

Linear layer: y = x @ W^T + b Uses cuBLAS via Ignis::LinAlg::Matmul for the hot path.

Instance Attribute Summary collapse

Attributes inherited from Module

#training

Instance Method Summary collapse

Methods inherited from Module

#call, #eval!, #load_state_dict, #named_parameters, #num_parameters, #parameters, #state_dict, #to, #train!, #zero_grad!

Constructor Details

#initialize(in_features, out_features, bias: true, device_id: 0) ⇒ Linear

Returns a new instance of Linear.

Parameters:

  • in_features (Integer)
  • out_features (Integer)
  • bias (Boolean) (defaults to: true) —

    whether to include bias

  • device_id (Integer) (defaults to: 0) 


21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
 
# File 'lib/nnw/ai/nn/linear.rb', line 21

def initialize(in_features, out_features, bias: true, device_id: 0)
  super()
  @in_features = in_features
  @out_features = out_features

  # Kaiming uniform initialization: bound = sqrt(6 / in_features)
  bound = Math.sqrt(6.0 / in_features)

  weight_nv = Ignis::Shared::NvArray.new(shape: [out_features, in_features],
                                        dtype: :float32, device_id: device_id)
  # Allocate on device (no host Array — for a 128k-vocab tied head that array
  # would be 262M Ruby Floats / ~10GB). The kaiming kernel writes every element.
  weight_nv.to_device

  # Initialize with Kaiming uniform via JIT kernel
  init_kernel = Ignis::JIT::Kernels::Elementwise.kaiming_uniform_init
  n = out_features * in_features
  seed = ::Random.new.rand(2**64)  # stdlib RNG (Ignis::Random is the cuRAND module)
  init_kernel.launch(grid: [(n + 255) / 256], block: [256],
                     args: [weight_nv, bound.to_f, Ignis::JIT::Kernel::U64.new(seed), n])

  @weight = register_parameter("weight",
             Tensor.new(data: weight_nv, requires_grad: true))

  if bias
    bias_nv = Ignis::Shared::NvArray.new(shape: [out_features],
                                        dtype: :float32, device_id: device_id)
    bias_bound = 1.0 / Math.sqrt(in_features)
    bias_nv.to_device
    init_kernel.launch(grid: [(out_features + 255) / 256], block: [256],
                       args: [bias_nv, bias_bound.to_f, Ignis::JIT::Kernel::U64.new(seed + 1), out_features])

    @bias = register_parameter("bias",
             Tensor.new(data: bias_nv, requires_grad: true))
  else
    @bias = nil
  end
end

Instance Attribute Details

#bias ⇒ Tensor? (readonly)

Returns bias vector [out_features].

Returns:

  • (Tensor, nil) —

    bias vector [out_features]



15
16
17
 
# File 'lib/nnw/ai/nn/linear.rb', line 15

def bias
  @bias
end

#weight ⇒ Tensor (readonly)

Returns weight matrix [out_features, in_features].

Returns:

  • (Tensor) —

    weight matrix [out_features, in_features]



12
13
14
 
# File 'lib/nnw/ai/nn/linear.rb', line 12

def weight
  @weight
end

Instance Method Details

#forward(x) ⇒ Tensor

Forward pass: x @ W^T + b

Parameters:

  • x (Tensor) —

    input [*, in_features]

Returns:

  • (Tensor) —

    output [*, out_features]



63
64
65
66
67
68
69
70
71
 
# File 'lib/nnw/ai/nn/linear.rb', line 63

def forward(x)
  # x @ W^T, with cuBLAS doing the transpose in the GEMM (transpose_b) —
  # avoids materializing W^T every forward (the LM head's was a 765ms
  # transpose of a 38M-element weight).
  out = x.matmul(@weight, transpose_b: true)
  # Bias is [out_features]; broadcast-add it across rows of [*, out_features].
  out = out.add_bias(@bias) if @bias
  out
end

#to_s ⇒ String

Returns:

  • (String) 


74
75
76
 
# File 'lib/nnw/ai/nn/linear.rb', line 74

def to_s
  "Linear(in=#{@in_features}, out=#{@out_features}, bias=#{!@bias.nil?})"
end