math_test.go

package llmgo

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const delta = 1e-5

func TestEncoderForward(t *testing.T) {
	type args struct {
		out []float32
		inp []int32
		wte []float32
		wpe []float32
		B   int
		T   int
		C   int
	}
	tests := []struct {
		name    string
		args    args
		wantOut []float32
	}{
		{
			name: "",
			args: args{
				inp: []int32{1, 0}, // [1 -> wte (2, 3), wpe(4, 5)] [0 -> wte (0, 1), wpe(6, 7)]
				wte: []float32{0, 1, 2, 3},
				wpe: []float32{4, 5, 6, 7},
				B:   1, // Batch size
				T:   1, // Sequence Len
				C:   2, // Dimensions
			},
			wantOut: []float32{6, 8},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			out := make([]float32, len(tt.args.inp))
			encoderForward(out, tt.args.inp, tt.args.wte, tt.args.wpe, tt.args.B, tt.args.T, tt.args.C)
			assert.Equal(t, tt.wantOut, out)
		})
	}
}

func TestEncoderBackward(t *testing.T) {
	type args struct {
		out  []float32
		inp  []int32
		dwte []float32
		dwpe []float32
		dout []float32
		B    int
		T    int
		C    int
	}
	tests := []struct {
		name     string
		args     args
		wantdwte []float32
		wantdwpe []float32
	}{
		{
			name: "",
			args: args{
				inp:  []int32{1}, //  [0 -> wte (3, 4), wpe(6, 7) (position 0)]
				dwte: []float32{1, 2, 3, 4},
				dwpe: []float32{6, 7, 8, 9},
				dout: []float32{1, 2, 3, 4}, // contains the diff that will be applied to wte and
				B:    1,                     // Batch size
				T:    1,                     // Sequence Len
				C:    2,                     // Dimensions
			},
			wantdwte: []float32{1, 2, 4, 6}, // 3, 4 (wte[inp[0]]) + dout[0]
			wantdwpe: []float32{7, 9, 8, 9},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			encoderBackward(tt.args.dwte, tt.args.dwpe, tt.args.dout, tt.args.inp, tt.args.B, tt.args.T, tt.args.C)
			assert.Equal(t, tt.wantdwpe, tt.args.dwpe)
			assert.Equal(t, tt.wantdwte, tt.args.dwte)
		})
	}
}

func TestLayernormForward(t *testing.T) {
	type args struct {
		inp    []float32
		weight []float32
		bias   []float32
		B      int
		T      int
		C      int
	}
	tests := []struct {
		name     string
		args     args
		wantOut  []float32
		wantMean []float32
		wantRstd []float32
	}{
		{
			name: "",
			args: args{
				inp:    []float32{0.2, 0.1, 0.3, 0.5, 0.1, 0.1},
				weight: []float32{1, 1, 1, 1, 1, 1},
				bias:   []float32{0, 0, 0, 0, 0, 0},
				B:      2,
				T:      1,
				C:      3,
			},
			wantOut:  []float32{0, -1.2238272, 1.2238274, 1.4140146, -0.70700747, -0.70700747},
			wantMean: []float32{0.2, 0.23333335},
			wantRstd: []float32{12.238273, 5.302555},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			out, mean, rstd := make([]float32, len(tt.args.inp)), make([]float32, tt.args.B*tt.args.T), make([]float32, tt.args.B*tt.args.T)
			layernormForward(out, mean, rstd, tt.args.inp, tt.args.weight, tt.args.bias, tt.args.B, tt.args.T, tt.args.C)
			require.InDeltaSlice(t, tt.wantOut, out, delta)
			require.InDeltaSlice(t, tt.wantMean, mean, delta)
			require.InDeltaSlice(t, tt.wantRstd, rstd, delta)
		})
	}
}

func TestLayernormBackward(t *testing.T) {
	type args struct {
		dinp    []float32
		dweight []float32
		dbias   []float32
		dout    []float32
		inp     []float32
		weight  []float32
		mean    []float32
		rstd    []float32
		B       int
		T       int
		C       int
	}
	tests := []struct {
		name string
		args args
	}{
		// TODO: Add test cases.
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			layernormBackward(tt.args.dinp, tt.args.dweight, tt.args.dbias, tt.args.dout, tt.args.inp, tt.args.weight, tt.args.mean, tt.args.rstd, tt.args.B, tt.args.T, tt.args.C)
		})
	}
}

func TestMatmulForward(t *testing.T) {
	type args struct {
		inp    []float32
		weight []float32
		bias   []float32
		B      int
		T      int
		C      int
		OC     int
	}
	tests := []struct {
		name    string
		args    args
		wantOut []float32
	}{
		{
			name: "simple",
			args: args{
				weight: []float32{ // OC (3) * C(2)
					1, 2,
					3, 4,
					5, 6,
				},
				inp: []float32{ // B(1) * T(1) * T(1) * C(2)
					1,
					2,
				},
				bias: []float32{1, 2, 3}, // OC
				// WEIGHT * INP + BIAS
				B:  1,
				T:  1,
				C:  2,
				OC: 3,
			},
			wantOut: []float32{
				6,
				13,
				20,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			out := make([]float32, tt.args.OC)
			matmulForward(out, tt.args.inp, tt.args.weight, tt.args.bias, tt.args.B, tt.args.T, tt.args.C, tt.args.OC)
			assert.Equal(t, tt.wantOut, out)
		})
	}
}

func TestAttentionForward(t *testing.T) {
	type args struct {
		inp []float32
		B   int
		T   int
		C   int
		NH  int
	}
	tests := []struct {
		name       string
		args       args
		wantOut    []float32
		wantPreatt []float32
		wantAtt    []float32
	}{
		{
			name: "Small Input Test",
			args: args{
				inp: []float32{1, 2, 3, 4, 5, 6},
				B:   1,
				T:   1,
				C:   2,
				NH:  1,
			},
			wantOut:    []float32{5, 6},
			wantPreatt: []float32{7.7781744},
			wantAtt:    []float32{1},
		},
		{
			name: "Larger Input Test",
			args: args{
				inp: []float32{ // (B, T, C3)
					/* B = 1 */
					/* T =  0 */
					/*qry*/ 1, 2, 3, // query compared against (4, 5, 6) but not (13, 14, 15) because it's in the future (t=1)
					/*key*/ 4, 5, 6,
					/*val*/ 7, 8, 9,
					/* T =  1 */
					/*qry*/ 10, 11, 12, // will be compared against (4, 5, 6) (t-1) and (13, 14, 15)
					/*key*/ 13, 14, 15,
					/*val*/ 16, 17, 18, // vals are updated to
				},
				B:  1,
				T:  2,
				C:  3,
				NH: 1,
			},
			wantOut: []float32{ // (B, T, C)
				/*      B = 0       */
				/*      T = 0       */
				/* C =  0    1    2 */
				/*  */ 7, 8, 9,
				/* T = 1 */
				/* C =  0    1    2 */
				/*  */ 16, 17, 18,
			},
			wantPreatt: []float32{ // (B, NH, T, T)
				/* B =  0    */
				/* NH = 0    */
				/*T =   1  2 */
				/*T=1*/ 18.475208, 0, // preatt: 18 -> 1, 0 -> 0
				/*T=2*/ 96.417496, 267.89053, // 96 -> 9, 267 -> 1
			},
			wantAtt: []float32{ // (B, NH, T, T)
				/* B = 0     */
				/* NH = 0    */
				/*T =   1  2 */
				/*T=1*/ 1, 0,
				/*T=2*/ 0, 1,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			out, preatt, att := make([]float32, len(tt.wantOut)), make([]float32, len(tt.wantPreatt)), make([]float32, len(tt.wantAtt))
			attentionForward(out, preatt, att, tt.args.inp, tt.args.B, tt.args.T, tt.args.C, tt.args.NH)
			assert.InDeltaSlice(t, tt.wantOut, out, 1e-4, fmt.Sprintf("want: %v got: %v", tt.wantOut, out))
			assert.InDeltaSlice(t, tt.wantPreatt, preatt, 1e-4, fmt.Sprintf("want: %v got: %v", tt.wantPreatt, preatt))
			assert.InDeltaSlice(t, tt.wantAtt, att, 1e-4, fmt.Sprintf("want: %v got: %v", tt.wantAtt, att))
		})
	}
}

func FuzzGeluInverse(f *testing.F) {
	f.Fuzz(func(t *testing.T, inb []byte) {
		inp := make([]float32, len(inb)/4)
		binary.Read(bytes.NewReader(inb), binary.LittleEndian, &inp)
		out := make([]float32, len(inp))
		geluForward(out, inp, len(inp))
		for i, v := range inp {
			got := out[i]
			switch {
			case v > 10:
				// large input values preserve their values
				assert.InDelta(t, v, got, 1e-3)
			case v == Inf(-1):
				assert.True(t, IsNaN(got), got)
			case v < 0:
				// negative input values get mapped close to 0
				assert.LessOrEqual(t, got, float32(0)) // out is less than zero
				assert.Greater(t, got, float32(-1.0))  // out is greater than -1
			}
		}
	})
}

func BenchmarkInference(b *testing.B) {
	randomText := "Kathleen Mary Ferrier CBE (22 April 1912 – 8 October 1953)[1] was an English contralto singer who achieved an international reputation as a stage, concert and recording artist, with a repertoire extending from folksong and popular ballads to the classical works of Bach, Brahms, Mahler and Elgar. Her death from cancer, at the height of her fame, was a shock to the musical world and particularly to the general public, which was kept in ignorance of the nature of her illness until after her death.  The daughter of a Lancashire village schoolmaster, Ferrier showed early talent as a pianist, and won numerous amateur piano competitions while working as a telephonist with the General Post Office. She did not take up singing seriously until 1937, when after winning a prestigious singing competition at the Carlisle Festival she began to receive offers of professional engagements as a vocalist. Thereafter she took singing lessons, first with J. E. Hutchinson and later with Roy Henderson. After the outbreak of the Second World War Ferrier was recruited by the Council for the Encouragement of Music and the Arts (CEMA), and in the following years sang at concerts and recitals throughout the UK. In 1942 her career was boosted when she met the conductor Malcolm Sargent, who recommended her to the influential Ibbs and Tillett concert management agency. She became a regular performer at leading London and provincial venues, and made numerous BBC radio broadcasts.  In 1946 Ferrier made her stage debut in the Glyndebourne Festival premiere of Benjamin Britten's opera The Rape of Lucretia. A year later she made her first appearance as Orfeo in Gluck's Orfeo ed Euridice, a work with which she became particularly associated. By her own choice, these were her only two operatic roles. As her reputation grew, Ferrier formed close working relationships with major musical figures, including Britten, Sir John Barbirolli, Bruno Walter and the accompanist Gerald Moore. She became known internationally through her three tours to the United States between 1948 and 1950 and her many visits to continental Europe.  Ferrier was diagnosed with breast cancer in March 1951. In between periods of hospitalisation and convalescence she continued to perform and record; her final public appearance was as Orfeo, at the Royal Opera House in February 1953, eight months before her death. Among her many memorials, the Kathleen Ferrier Cancer Research Fund was launched in May 1954. The Kathleen Ferrier Scholarship Fund, administered by the Royal Philharmonic Society, has since 1956 made annual awards to aspiring young professional singers."
	model, err := LoadGPT2Model("./gpt2_124M.bin", "./gpt2_tokenizer.bin")
	require.NoError(b, err)
	for range b.Name() {
		output, err := model.Inference(randomText, 1, 9)
		require.NoError(b, err)
		b.Log(output)
	}
}