diff --git a/chat.f90 b/chat.f90 index 9fb6683..b64c7ea 100644 --- a/chat.f90 +++ b/chat.f90 @@ -1,5 +1,7 @@ program chatgpt2 use driver, only: chat +use tokenizer, only: string implicit none -call chat() +type(string), allocatable :: inputs(:) +call chat(.false., inputs) end program diff --git a/driver.f90 b/driver.f90 index 476bd95..9c92238 100644 --- a/driver.f90 +++ b/driver.f90 @@ -1,270 +1,283 @@ module driver -use gpt2_mod, only: generate, model_t -use tokenizer, only: encode, decode, string -use omp, only: omp_get_wtime -implicit none + use gpt2_mod, only : generate, model_t + use tokenizer, only : encode, decode, string + use omp, only : omp_get_wtime + implicit none -integer, parameter :: sp = kind(0.0) -integer, parameter :: dp = kind(0.d0) -character(1), parameter :: LF = achar(10) + integer, parameter :: sp = kind(0.0) + integer, parameter :: dp = kind(0.d0) + character(1), parameter :: LF = achar(10) contains -subroutine load_input(filename, input_txt, n_tokens_to_generate) -! Load the input from a namelist `filename` -character(*), intent(in) :: filename -character(:), allocatable, intent(out) :: input_txt -integer, intent(out) :: n_tokens_to_generate -character(1024) :: input_txt2 -integer :: u, ios -namelist / input_fastGPT / n_tokens_to_generate -allocate(character(0) :: input_txt) -input_txt = "" -open(newunit=u, file=filename, status="old") -read(u, input_fastGPT) -do - read(u, "(a)", iostat=ios) input_txt2 - if (ios /= 0) exit - if (len(input_txt) > 0) input_txt = input_txt // char(10) - input_txt = input_txt // trim(input_txt2) -end do -close(u) -end subroutine + subroutine load_input(filename, input_txt, n_tokens_to_generate) + ! Load the input from a namelist `filename` + character(*), intent(in) :: filename + character(:), allocatable, intent(out) :: input_txt + integer, intent(out) :: n_tokens_to_generate + character(1024) :: input_txt2 + integer :: u, ios + !namelist / input_fastGPT / n_tokens_to_generate + allocate(character(0) :: input_txt) + n_tokens_to_generate = 20 + input_txt = "Alan Turing theorized that computers would one day become very powerful, but even he could not imagine" + !open(newunit=u, file=filename, status="old") + !read(u, input_fastGPT) + !do + ! read(u, "(a)", iostat=ios) input_txt2 + ! if (ios /= 0) exit + ! if (len(input_txt) > 0) input_txt = input_txt // char(10) + ! input_txt = input_txt // trim(input_txt2) + !end do + !close(u) + end subroutine -subroutine load_model(filename, m) -character(*), intent(in) :: filename -type(model_t), intent(out) :: m -! We use the following fastGPT model type number -! fastGPT (digits look similar to the letters they represent) -! 0xfa51697 = 262477463 -integer, parameter :: current_model_mark = 262477463 -integer, parameter :: current_model_version = 1 -integer :: model_mark -integer :: u -open(newunit=u, file=filename, form="unformatted", access="stream", status="old") -read(u) model_mark -if (model_mark /= current_model_mark) then - print *, "Found:", model_mark - print *, "Expected:", current_model_mark - error stop "Invalid fastGPT model file" -end if -read(u) m%model_file_version -if (m%model_file_version /= current_model_version) then - print *, "Found:", m%model_file_version - print *, "Expected:", current_model_version - error stop "Incompatible model version" -end if -read(u) m%n_vocab, m%n_ctx, m%n_embd, m%n_layer, m%n_head, m%n_decoder_idx, & - m%n_decoder_txt, m%n_vocab_idx, m%n_vocab_txt, m%n_byte_encoder -allocate(m%wte(m%n_embd,m%n_vocab), m%wpe(m%n_embd,m%n_ctx), & - m%mlp_fc_w(4*m%n_embd,m%n_embd,m%n_layer), m%mlp_fc_b(4*m%n_embd,m%n_layer), & - m%mlp_proj_w(m%n_embd,4*m%n_embd,m%n_layer), m%mlp_proj_b(m%n_embd,m%n_layer), & - m%attn_w(3*m%n_embd,m%n_embd,m%n_layer), m%attn_b(3*m%n_embd,m%n_layer), & - m%attn_proj_w(m%n_embd,m%n_embd,m%n_layer), m%attn_proj_b(m%n_embd,m%n_layer), & - m%ln1_b(m%n_embd,m%n_layer), m%ln1_g(m%n_embd,m%n_layer), & - m%ln2_b(m%n_embd,m%n_layer), m%ln2_g(m%n_embd,m%n_layer), & - m%lnf_b(m%n_embd), m%lnf_g(m%n_embd), & - m%decoder_idx(0:m%n_decoder_idx-1), m%decoder_txt(m%n_decoder_txt), & - m%vocab_idx(0:m%n_vocab_idx-1), m%vocab_txt(m%n_vocab_txt), & - m%byte_encoder(0:m%n_byte_encoder-1)) -read(u) m%wte, m%wpe, & - m%mlp_fc_w, m%mlp_fc_b, & - m%mlp_proj_w, m%mlp_proj_b, & - m%attn_w, m%attn_b, & - m%attn_proj_w, m%attn_proj_b, & - m%ln1_b, m%ln1_g, & - m%ln2_b, m%ln2_g, & - m%lnf_b, m%lnf_g, & - m%decoder_idx, m%decoder_txt, & - m%vocab_idx, m%vocab_txt, & - m%byte_encoder -close(u) -end subroutine + subroutine load_model(filename, m) + character(*), intent(in) :: filename + type(model_t), intent(out) :: m + ! We use the following fastGPT model type number + ! fastGPT (digits look similar to the letters they represent) + ! 0xfa51697 = 262477463 + integer, parameter :: current_model_mark = 262477463 + integer, parameter :: current_model_version = 1 + integer :: model_mark + integer :: u + open(newunit = u, file = filename, form = "unformatted", access = "stream", status = "old") + read(u) model_mark + if (model_mark /= current_model_mark) then + print *, "Found:", model_mark + print *, "Expected:", current_model_mark + error stop "Invalid fastGPT model file" + end if + read(u) m%model_file_version + if (m%model_file_version /= current_model_version) then + print *, "Found:", m%model_file_version + print *, "Expected:", current_model_version + error stop "Incompatible model version" + end if + read(u) m%n_vocab, m%n_ctx, m%n_embd, m%n_layer, m%n_head, m%n_decoder_idx, & + m%n_decoder_txt, m%n_vocab_idx, m%n_vocab_txt, m%n_byte_encoder + allocate(m%wte(m%n_embd, m%n_vocab), m%wpe(m%n_embd, m%n_ctx), & + m%mlp_fc_w(4 * m%n_embd, m%n_embd, m%n_layer), m%mlp_fc_b(4 * m%n_embd, m%n_layer), & + m%mlp_proj_w(m%n_embd, 4 * m%n_embd, m%n_layer), m%mlp_proj_b(m%n_embd, m%n_layer), & + m%attn_w(3 * m%n_embd, m%n_embd, m%n_layer), m%attn_b(3 * m%n_embd, m%n_layer), & + m%attn_proj_w(m%n_embd, m%n_embd, m%n_layer), m%attn_proj_b(m%n_embd, m%n_layer), & + m%ln1_b(m%n_embd, m%n_layer), m%ln1_g(m%n_embd, m%n_layer), & + m%ln2_b(m%n_embd, m%n_layer), m%ln2_g(m%n_embd, m%n_layer), & + m%lnf_b(m%n_embd), m%lnf_g(m%n_embd), & + m%decoder_idx(0:m%n_decoder_idx - 1), m%decoder_txt(m%n_decoder_txt), & + m%vocab_idx(0:m%n_vocab_idx - 1), m%vocab_txt(m%n_vocab_txt), & + m%byte_encoder(0:m%n_byte_encoder - 1)) + read(u) m%wte, m%wpe, & + m%mlp_fc_w, m%mlp_fc_b, & + m%mlp_proj_w, m%mlp_proj_b, & + m%attn_w, m%attn_b, & + m%attn_proj_w, m%attn_proj_b, & + m%ln1_b, m%ln1_g, & + m%ln2_b, m%ln2_g, & + m%lnf_b, m%lnf_g, & + m%decoder_idx, m%decoder_txt, & + m%vocab_idx, m%vocab_txt, & + m%byte_encoder + close(u) + end subroutine -subroutine gpt2_driver(input, output, m) -integer, allocatable, intent(out) :: input(:), output(:) -type(model_t), intent(out) :: m -character(:), allocatable :: input_txt -integer :: n_tokens_to_generate -real(dp) :: t1, t2 -call load_input("input", input_txt, n_tokens_to_generate) + subroutine gpt2_driver(input, output, m) + integer, allocatable, intent(out) :: input(:), output(:) + type(model_t), intent(out) :: m + character(:), allocatable :: input_txt + integer :: n_tokens_to_generate + real(dp) :: t1, t2 + call load_input("input", input_txt, n_tokens_to_generate) -! Load the model -print "(a)", "Loading the model..." -call cpu_time(t1) -call load_model("model.dat", m) -call cpu_time(t2) -print "(a,f8.3,a,i2)", " done. Time:", t2-t1, "s, Model file version:", m%model_file_version -print * -print "(a)", "Model parameters:" -print "(a,i6)", "n_vocab =", m%n_vocab -print "(a,i6)", "n_ctx =", m%n_ctx -print "(a,i6)", "n_embd =", m%n_embd -print "(a,i6)", "n_layer =", m%n_layer -print "(a,i6)", "n_head =", m%n_head -print * + ! Load the model + print "(a)", "Loading the model..." + call cpu_time(t1) + call load_model("model.dat", m) + call cpu_time(t2) + print "(a,f8.3,a,i2)", " done. Time:", t2 - t1, "s, Model file version:", m%model_file_version + print * + print "(a)", "Model parameters:" + print "(a,i6)", "n_vocab =", m%n_vocab + print "(a,i6)", "n_ctx =", m%n_ctx + print "(a,i6)", "n_embd =", m%n_embd + print "(a,i6)", "n_layer =", m%n_layer + print "(a,i6)", "n_head =", m%n_head + print * -call gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output) -endsubroutine + call gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output) + endsubroutine -subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output) -character(*), intent(in) :: input_txt -integer, intent(in) :: n_tokens_to_generate -type(model_t), intent(in) :: m -integer, allocatable, intent(out) :: input(:), output(:) -integer, allocatable :: byte_decoder(:) -integer :: n_seq -character(:), allocatable :: output_txt -real(dp) :: t1, t2, t1o, t2o -integer :: i -logical :: use_cache + subroutine gpt2_driver2(input_txt, n_tokens_to_generate, m, input, output) + character(*), intent(in) :: input_txt + integer, intent(in) :: n_tokens_to_generate + type(model_t), intent(in) :: m + integer, allocatable, intent(out) :: input(:), output(:) + integer, allocatable :: byte_decoder(:) + integer :: n_seq + character(:), allocatable :: output_txt + real(dp) :: t1, t2, t1o, t2o + integer :: i + logical :: use_cache -! Compute byte_decoder: -allocate(byte_decoder(0:maxval(m%byte_encoder))) -byte_decoder = 0 -do i = 0, size(m%byte_encoder)-1 - byte_decoder(m%byte_encoder(i)) = i -end do + ! Compute byte_decoder: + allocate(byte_decoder(0:maxval(m%byte_encoder))) + print "(a)", "size(m%byte_encoder)" + print *, size(m%byte_encoder) + print "(a)", "byte encoder max" + print *, maxval(m%byte_encoder) + print "(a)", "size(byte_decoder)" + print *, size(byte_decoder) +! print "(a)", "byte encoder" +! print *, m%byte_encoder + byte_decoder = 0 + do i = 0, size(m%byte_encoder) - 1 + byte_decoder(m%byte_encoder(i)) = i + end do +! print "(a)", "byte decoder" +! print *, byte_decoder -print "(a)", "Input text" -print "(a)", input_txt + print "(a)", "Input text" + print *, input_txt -print * -print "(a)", "Encoding: tokenizing input text into tokens (currently slow)..." -call cpu_time(t1) -input = encode(input_txt, m%decoder_idx, m%decoder_txt, m%vocab_idx, m%vocab_txt, & - m%byte_encoder) -call cpu_time(t2) -n_seq = size(input) -print "(a,f8.3,a)", " done. Time:", t2-t1, "s" -print * -print "(a)", "Input parameters:" -print "(a,i4)", "n_seq =", n_seq -print "(a,i4)", "n_tokens_to_generate =", n_tokens_to_generate -print * -print "(a)", "Input tokens:" -print "(1000(i6))", input -print * + print * + print "(a)", "Encoding: tokenizing input text into tokens (currently slow)..." + call cpu_time(t1) + input = encode(input_txt, m%decoder_idx, m%decoder_txt, m%vocab_idx, m%vocab_txt, & + m%byte_encoder) + call cpu_time(t2) + n_seq = size(input) + print "(a,f8.3,a)", " done. Time:", t2 - t1, "s" + print * + print "(a)", "Input parameters:" + print "(a,i4)", "n_seq =", n_seq + print "(a,i4)", "n_tokens_to_generate =", n_tokens_to_generate + print * + print "(a)", "Input tokens:" + !print "(1000(i6))", input + print *, input + print * -if (n_seq + n_tokens_to_generate >= m%n_ctx) then - print *, "The maximum sequence length of the model was surpassed." - print *, "Make the input and/or number of tokens to generate shorter." - error stop -end if + if (n_seq + n_tokens_to_generate >= m%n_ctx) then + print *, "The maximum sequence length of the model was surpassed." + print *, "Make the input and/or number of tokens to generate shorter." + error stop + end if -print "(a)", "Decoded input as text:" -!print "(a)", decode(input, decoder_idx, decoder_txt, byte_decoder) -allocate(character(0) :: output_txt) ! Fix GFortran warning -output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder) -print "(a)", output_txt -print * + print "(a)", "Decoded input as text:" + !print "(a)", decode(input, decoder_idx, decoder_txt, byte_decoder) + allocate(character(0) :: output_txt) ! Fix GFortran warning + output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder) + print *, output_txt + print * -if (input_txt /= output_txt) then - error stop "The decoded input text does not agree with the input text" -end if + if (input_txt /= output_txt) then + error stop "The decoded input text does not agree with the input text" + end if -allocate(output(n_tokens_to_generate)) -print "(a)", "Running model..." -call cpu_time(t1) -t1o = omp_get_wtime() -use_cache = .true. -output = generate(n_tokens_to_generate, m, size(input), input, use_cache, & - byte_decoder) -print * -t2o = omp_get_wtime() -call cpu_time(t2) -print "(a,f8.3,a,f4.2,a)", " done. Time:", t2o-t1o, "s (", (t2-t1)/(t2o-t1o), "x)" -print * -print "(a)", "Output tokens:" -print "(1000(i6))", output -output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder) -print * -print "(a)", "Decoded output as text:" -print "(a)", output_txt -end subroutine + allocate(output(n_tokens_to_generate)) + print "(a)", "Running model..." + call cpu_time(t1) + t1o = omp_get_wtime() + use_cache = .true. + call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, & + byte_decoder) + print * + t2o = omp_get_wtime() + call cpu_time(t2) + print "(a,f8.3,a,f8.2,a)", " done. Time:", t2o - t1o, "s (", (t2 - t1) / (t2o - t1o), "x)" + print * + print "(a)", "Output tokens:" + print *, output + output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder) + print * + print "(a)", "Decoded output as text:" + print *, output_txt + end subroutine -subroutine gpt2_driver3(input_txt, n_tokens_to_generate, stop_text, m, output_txt) -character(*), intent(in) :: input_txt, stop_text -integer, intent(in) :: n_tokens_to_generate -type(model_t), intent(in) :: m -integer, allocatable :: input(:), output(:) -integer, allocatable :: byte_decoder(:) -integer :: n_seq -character(:), allocatable, intent(out) :: output_txt -integer :: i -logical :: use_cache -! TODO: move the decoder into model_t -! Compute byte_decoder: -allocate(byte_decoder(0:maxval(m%byte_encoder))) -byte_decoder = 0 -do i = 0, size(m%byte_encoder)-1 - byte_decoder(m%byte_encoder(i)) = i -end do -input = encode(input_txt, m%decoder_idx, m%decoder_txt, m%vocab_idx, m%vocab_txt, & - m%byte_encoder) -n_seq = size(input) -if (n_seq + n_tokens_to_generate >= m%n_ctx) then - print *, "The maximum sequence length of the model was surpassed." - print *, "Make the input and/or number of tokens to generate shorter." - error stop -end if -allocate(character(0) :: output_txt) ! Fix GFortran warning -output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder) -if (input_txt /= output_txt) then - error stop "The decoded input text does not agree with the input text" -end if -use_cache = .true. -output = generate(n_tokens_to_generate, m, size(input), input, use_cache, & - byte_decoder, stop_text) -output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder) -end subroutine + subroutine gpt2_driver3(input_txt, n_tokens_to_generate, stop_text, m, output_txt) + character(*), intent(in) :: input_txt, stop_text + integer, intent(in) :: n_tokens_to_generate + type(model_t), intent(in) :: m + integer, allocatable :: input(:), output(:) + integer, allocatable :: byte_decoder(:) + integer :: n_seq + character(:), allocatable, intent(out) :: output_txt + integer :: i + logical :: use_cache + ! TODO: move the decoder into model_t + ! Compute byte_decoder: + allocate(byte_decoder(0:maxval(m%byte_encoder))) + byte_decoder = 0 + do i = 0, size(m%byte_encoder) - 1 + byte_decoder(m%byte_encoder(i)) = i + end do + input = encode(input_txt, m%decoder_idx, m%decoder_txt, m%vocab_idx, m%vocab_txt, & + m%byte_encoder) + n_seq = size(input) + if (n_seq + n_tokens_to_generate >= m%n_ctx) then + print *, "The maximum sequence length of the model was surpassed." + print *, "Make the input and/or number of tokens to generate shorter." + error stop + end if + allocate(character(0) :: output_txt) ! Fix GFortran warning + output_txt = decode(input, m%decoder_idx, m%decoder_txt, byte_decoder) + if (input_txt /= output_txt) then + error stop "The decoded input text does not agree with the input text" + end if + use_cache = .true. + call generate(output, n_tokens_to_generate, m, size(input), input, use_cache, & + byte_decoder, stop_text) + output_txt = decode(output, m%decoder_idx, m%decoder_txt, byte_decoder) + end subroutine -function get_prompt() result(input) -character(:), allocatable :: input -character(1024) :: tmp -integer ::ios -read(*,"(a)",iostat=ios) tmp -if (ios == 0) then - input = trim(tmp) -else - input = "" -end if -end function + function get_prompt() result(input) + character(:), allocatable :: input + character(1024) :: tmp + integer :: ios + read(*, "(a)", iostat = ios) tmp + if (ios == 0) then + input = trim(tmp) + else + input = "" + end if + end function -subroutine chat(inputs) -type(string), optional, intent(in) :: inputs(:) -type(model_t) :: m -character(:), allocatable :: prompt, input, output -integer :: i, n_prompts -call load_model("model.dat", m) -prompt = "Your name is fastGPT and you are an AI bot. The user will ask you & -&questions and you answer in a nice, truthful, short way." // LF // "& -&User: What is the capital of Czechia?" // LF // "& -&fastGPT: Prague." // LF // "& -&User: How many legs does a dog have?" // LF // "& -&fastGPT: Four." // LF // "& -&User:" -write(*,"(a)",advance="no") prompt -if (present(inputs)) then - n_prompts = size(inputs) -else - n_prompts = 1024 -end if -do i = 1, n_prompts - write(*,"(a)",advance="no") " " - if (present(inputs)) then - input = inputs(i)%s - write(*,"(a)") input - else - input = get_prompt() - if (input == "") exit - end if - write(*,"(a)",advance="no") "fastGPT:" - prompt = prompt // " " // input // LF // "fastGPT:" - call gpt2_driver3(prompt, 200, "User:", m, output) - prompt = prompt // output -end do -print * -end subroutine + subroutine chat(present_inputs, inputs) + logical, intent(in) :: present_inputs + type(string), optional, intent(in) :: inputs(:) + type(model_t) :: m + character(:), allocatable :: prompt, input, output + integer :: i, n_prompts + call load_model("model.dat", m) + prompt = "Your name is fastGPT and you are an AI bot. The user will ask you & + &questions and you answer in a nice, truthful, short way." // LF // "& + &User: What is the capital of Czechia?" // LF // "& + &fastGPT: Prague." // LF // "& + &User: How many legs does a dog have?" // LF // "& + &fastGPT: Four." // LF // "& + &User:" + write(*, "(a)", advance = "no") prompt + if (present_inputs) then + n_prompts = size(inputs) + else + n_prompts = 1024 + end if + do i = 1, n_prompts + write(*, "(a)", advance = "no") " " + if (present_inputs) then + input = inputs(i)%s + write(*, "(a)") input + else + input = get_prompt() + if (input == "") exit + end if + write(*, "(a)", advance = "no") "fastGPT:" + prompt = prompt // " " // input // LF // "fastGPT:" + call gpt2_driver3(prompt, 200, "User:", m, output) + prompt = prompt // output + end do + print * + end subroutine end module diff --git a/gpt2.f90 b/gpt2.f90 index a318e5b..731259b 100644 --- a/gpt2.f90 +++ b/gpt2.f90 @@ -1,295 +1,380 @@ module gpt2_mod -use linalg, only: matmul_2d, matmul_2d_t -use tokenizer, only: decode -implicit none + use linalg, only : matmul_2d, matmul_2d_t + use tokenizer, only : decode + implicit none -integer, parameter :: sp = kind(0.0) -real(sp), parameter :: pi = 3.14159265358979323846_sp + integer, parameter :: sp = kind(0.0) + real(sp), parameter :: pi = 3.14159265358979323846_sp -! This derived type contains all the data of the GPT-2 model, including all -! weights, model parameters, and encoder/decoder data -type :: model_t - integer :: n_vocab, n_ctx, n_embd, n_layer, n_head, & - n_decoder_idx, n_decoder_txt, & - n_vocab_idx, n_vocab_txt, n_byte_encoder - real(sp), allocatable :: wte(:,:), wpe(:,:), & - mlp_fc_w(:,:,:), mlp_fc_b(:,:), & - mlp_proj_w(:,:,:), mlp_proj_b(:,:), & - attn_w(:,:,:), attn_b(:,:), & - attn_proj_w(:,:,:), attn_proj_b(:,:), & - ln1_b(:,:), ln1_g(:,:), & - ln2_b(:,:), ln2_g(:,:), & - lnf_b(:), lnf_g(:) - integer, allocatable :: decoder_idx(:), vocab_idx(:), byte_encoder(:) - character, allocatable :: decoder_txt(:), vocab_txt(:) - integer :: model_file_version -end type + ! This derived type contains all the data of the GPT-2 model, including all + ! weights, model parameters, and encoder/decoder data + type :: model_t + integer :: n_vocab, n_ctx, n_embd, n_layer, n_head, & + n_decoder_idx, n_decoder_txt, & + n_vocab_idx, n_vocab_txt, n_byte_encoder + real(sp), allocatable :: wte(:, :), wpe(:, :), & + mlp_fc_w(:, :, :), mlp_fc_b(:, :), & + mlp_proj_w(:, :, :), mlp_proj_b(:, :), & + attn_w(:, :, :), attn_b(:, :), & + attn_proj_w(:, :, :), attn_proj_b(:, :), & + ln1_b(:, :), ln1_g(:, :), & + ln2_b(:, :), ln2_g(:, :), & + lnf_b(:), lnf_g(:) + integer, allocatable :: decoder_idx(:), vocab_idx(:), byte_encoder(:) + integer(1), allocatable :: decoder_txt(:), vocab_txt(:) + integer :: model_file_version + end type contains -elemental real(sp) function fast_tanh(x) result(y) -real(sp), intent(in) :: x -real(sp) :: x2 -if (x > 5) then - y = 1 -elseif (x < -5) then - y = -1 -else - x2 = x*x - y = x * (0.98569772605911309407 + x2 *(-0.2794500993392901382 & - + x2 * (6.8280504526399188164e-2 + x2 * (-1.0972014877337651823e-2 & - + x2 * (1.1132367134444316902e-3 + x2 * (-7.018851897305717565e-5 & - + x2 * (2.656616768082727089e-6 + x2 * (-5.5138381821615909058e-8 & - + x2 * 4.8162484477588665996e-10)))))))) -end if -end function + elemental real(sp) function fast_tanh(x) result(y) + real(sp), intent(in) :: x + real(sp) :: x2 + if (x > 5) then + y = 1 + elseif (x < -5) then + y = -1 + else + x2 = x * x + y = x * (0.98569772605911309407 + x2 * (-0.2794500993392901382 & + + x2 * (6.8280504526399188164e-2 + x2 * (-1.0972014877337651823e-2 & + + x2 * (1.1132367134444316902e-3 + x2 * (-7.018851897305717565e-5 & + + x2 * (2.656616768082727089e-6 + x2 * (-5.5138381821615909058e-8 & + + x2 * 4.8162484477588665996e-10)))))))) + end if + end function -elemental real(sp) function gelu(x) result(y) -real(sp), intent(in) :: x -y = 0.5_sp * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715_sp * x**3))) -end function + elemental real(sp) function gelu(x) result(y) + real(sp), intent(in) :: x + y = 0.5_sp * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715_sp * x**3))) + end function -function softmax(x) result(y) -real(sp), intent(in) :: x(:,:) -real(sp) :: y(size(x,1),size(x,2)) -integer :: i -do i = 1, size(x,2) - y(:,i) = exp(x(:,i) - maxval(x(:,i))) - y(:,i) = y(:,i) / sum(y(:,i)) -end do -end function + function softmax(x) result(y) + real(sp), intent(in) :: x(:, :) + real(sp) :: y(size(x, 1), size(x, 2)) + integer :: i, j + real(sp) :: s + do i = 1, size(x, 2) + s = -1e10 + do j = 1, size(x, 1) + if (x(j, i) > s) s = x(j, i) + end do + do j = 1, size(x, 1) + y(j, i) = exp(x(j, i) - s) + end do + s = 0 + do j = 1, size(x, 1) + s = s + y(j, i) + end do + do j = 1, size(x, 1) + y(j, i) = y(j, i) / s + end do + end do + end function -function layer_norm(x, g, b, eps) result(y) -real(sp), intent(in) :: x(:,:), g(:), b(:), eps -real(sp) :: y(size(x,1),size(x,2)) -real(sp) :: mean(size(x,2)), variance(size(x,2)) -integer :: i -do i = 1, size(x,2) - mean(i) = sum(x(:,i)) / size(x,1) - variance(i) = sum((x(:,i) - mean(i))**2) / size(x,1) -end do -!do i = 1, size(x,1) -! y(i,:) = (x(i,:) - mean(:)) / sqrt(variance(:) + eps) -! y(i,:) = g(i) * y(i,:) + b(i) -!end do -do i = 1, size(x,2) - y(:,i) = (x(:,i) - mean(i)) / sqrt(variance(i) + eps) - y(:,i) = g(:) * y(:,i) + b(:) -end do -end function + subroutine layer_norm(y, x, g, b, eps) + real(sp), intent(in) :: x(:, :), g(:), b(:), eps + real(sp), intent(out) :: y(size(x, 1), size(x, 2)) + real(sp) :: mean(size(x, 2)), variance(size(x, 2)) + real(sp) :: xi(size(x, 1)) + integer :: i, j + do i = 1, size(x, 2) + do j = 1, size(x, 1) + xi(j) = x(j, i) + end do + mean(i) = sum(xi) / size(x, 1) + do j = 1, size(x, 1) + xi(j) = (xi(j) - mean(i))**2 + end do + variance(i) = sum(xi) / size(x, 1) + end do + !do i = 1, size(x,1) + ! y(i,:) = (x(i,:) - mean(:)) / sqrt(variance(:) + eps) + ! y(i,:) = g(i) * y(i,:) + b(i) + !end do + do i = 1, size(x, 2) + do j = 1, size(x, 1) + y(j, i) = (x(j, i) - mean(i)) / sqrt(variance(i) + eps) + y(j, i) = g(j) * y(j, i) + b(j) + end do + end do + end subroutine -function linear(x, w, b) result(y) -real(sp), intent(in) :: x(:,:), w(:,:), b(:) -real(sp) :: y(size(b,1),size(x,2)) -integer :: i -!y = matmul(w, x) + spread(b, 2, size(x,2)) -!y = matmul(w, x) -call matmul_2d(w, x, y) -do i = 1, size(y,2) - y(:,i) = y(:,i) + b(:) -end do -end function + function linear(x, w, b) result(y) + real(sp), intent(in) :: x(:, :), w(:, :), b(:) + real(sp) :: y(size(b, 1), size(x, 2)) + integer :: i, j + !y = matmul(w, x) + spread(b, 2, size(x,2)) + !y = matmul(w, x) + call matmul_2d(w, x, y) + do i = 1, size(y, 2) + do j = 1, size(y, 1) + y(j, i) = y(j, i) + b(j) + end do + end do + end function -function ffn(x, fc_w, fc_b, proj_w, proj_b) result(y) -real(sp), intent(in) :: x(:,:), fc_w(:,:), fc_b(:), proj_w(:,:), proj_b(:) -real(sp) :: y(size(x,1),size(x,2)) -!real(sp) :: a(4*size(x,1),size(x,2)) -!a = gelu(linear(x, fc_w, fc_b)) -y = linear(gelu(linear(x, fc_w, fc_b)), proj_w, proj_b) -end function + subroutine ffn(y, x, fc_w, fc_b, proj_w, proj_b) + real(sp), intent(in) :: x(:, :), fc_w(:, :), fc_b(:), proj_w(:, :), proj_b(:) + real(sp), intent(inout) :: y(size(x, 1), size(x, 2)) + real(sp) :: yy(size(x, 1), size(x, 2)) + real(sp) :: a(4 * size(x, 1), size(x, 2)) + real(sp) :: aa(4 * size(x, 1), size(x, 2)) + aa = linear(x, fc_w, fc_b) + a = gelu(aa) + yy = linear(a, proj_w, proj_b) + y = y + yy + end subroutine -function attention(n_embd_head,n_seq,n_seq_x, q, k, v, mask) result(y) -integer, intent(in) :: n_embd_head, n_seq, n_seq_x -real(sp), intent(in) :: q(n_embd_head,n_seq_x), k(n_embd_head,n_seq), v(n_embd_head,n_seq), mask(n_seq,n_seq_x) -real(sp) :: y(n_embd_head,n_seq_x) -real(sp) :: tmp(n_seq,n_seq_x) -!tmp = matmul(transpose(k), q) -!call matmul_2d(transpose(k), q, tmp) -call matmul_2d_t(k, q, tmp) -call matmul_2d(v, softmax(tmp / sqrt(real(n_embd_head,sp)) + mask), y) -end function + subroutine attention(y, n_embd_head, n_seq, n_seq_x, q, k, v, mask) + integer, intent(in) :: n_embd_head, n_seq, n_seq_x + real(sp), intent(in) :: q(n_embd_head, n_seq_x), k(n_embd_head, n_seq), v(n_embd_head, n_seq), mask(n_seq, n_seq_x) + real(sp), intent(out) :: y(n_embd_head, n_seq_x) + real(sp) :: tmp(n_seq, n_seq_x) + integer :: i, j + !tmp = matmul(transpose(k), q) + !call matmul_2d(transpose(k), q, tmp) + call matmul_2d_t(k, q, tmp) + do i = 1, n_seq_x + do j = 1, n_seq + tmp(j, i) = tmp(j, i) / sqrt(real(n_embd_head, sp)) + mask(j, i) + end do + end do + tmp = softmax(tmp) + call matmul_2d(v, tmp, y) + end subroutine -function mha(n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head, & - use_kv_cache, kv_cache) & - result(y) -integer, intent(in) :: n_seq, n_seq_x, n_embd -real(sp), intent(in) :: x(n_embd,n_seq_x), & - attn_w(3*n_embd,n_embd), attn_b(3*n_embd), & - proj_w(n_embd,n_embd), proj_b(n_embd) -real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2) -integer, intent(in) :: n_head -logical, intent(in) :: use_kv_cache -real(sp) :: y(n_embd,n_seq_x) -real(sp) :: causal_mask(n_seq,n_seq_x) -real(sp) :: x2(3*n_embd,n_seq_x) -integer :: i, j -! Mask -if (use_kv_cache) then - causal_mask = 0 -else - do j = 1, n_seq - do i = 1, n_seq - if (i > j) then - causal_mask(i,j) = -1e10_sp + subroutine mha(y, n_seq, n_seq_x, n_embd, x, attn_w, attn_b, proj_w, proj_b, n_head, & + use_kv_cache, kv_cache) + integer, intent(in) :: n_seq, n_seq_x, n_embd + real(sp), intent(in) :: x(n_embd, n_seq_x), & + attn_w(3 * n_embd, n_embd), attn_b(3 * n_embd), & + proj_w(n_embd, n_embd), proj_b(n_embd) + real(sp), intent(inout) :: kv_cache(n_embd, n_seq, 2) + integer, intent(in) :: n_head + logical, intent(in) :: use_kv_cache + real(sp), intent(out) :: y(n_embd, n_seq_x) + real(sp) :: causal_mask(n_seq, n_seq_x) + real(sp) :: x2(3 * n_embd, n_seq_x) + real(sp) :: q(n_embd / n_head, n_seq_x), k(n_embd / n_head, n_seq), v(n_embd / n_head, n_seq) + real(sp) :: yy(n_embd / n_head, n_seq_x) + integer :: i, j, l + ! Mask + if (use_kv_cache) then + causal_mask = 0 + else + do j = 1, n_seq + do i = 1, n_seq + if (i > j) then + causal_mask(i, j) = -1e10_sp + else + causal_mask(i, j) = 0 + end if + end do + end do + end if + x2 = linear(x, attn_w, attn_b) + if (use_kv_cache) then + do j = 1, n_embd + kv_cache(j, n_seq, 1) = x2((2 - 1) * n_embd + j, 1) + kv_cache(j, n_seq, 2) = x2((3 - 1) * n_embd + j, 1) + end do else - causal_mask(i,j) = 0 + do i = 1, n_seq + do j = 1, n_embd + kv_cache(j, i, 1) = x2((2 - 1) * n_embd + j, i) + kv_cache(j, i, 2) = x2((3 - 1) * n_embd + j, i) + end do + end do end if - end do - end do -end if -x2 = linear(x, attn_w, attn_b) -associate ( & - q => x2((1-1)*n_embd+1:1*n_embd,:), & - k => x2((2-1)*n_embd+1:2*n_embd,:), & - v => x2((3-1)*n_embd+1:3*n_embd,:) & - ) - if (use_kv_cache) then - kv_cache(:,n_seq,1) = k(:,1) - kv_cache(:,n_seq,2) = v(:,1) - else - kv_cache(:,:,1) = k - kv_cache(:,:,2) = v - end if -end associate -associate ( & - q => x2((1-1)*n_embd+1:1*n_embd,:), & - k => kv_cache(:,:,1), & - v => kv_cache(:,:,2) & - ) - ! Perform attention over each head - do i = 1, n_head - y((i-1)*n_embd/n_head+1:i*n_embd/n_head,:) = attention( & - n_embd/n_head, n_seq, n_seq_x, & - q((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), & - k((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), & - v((i-1)*n_embd/n_head+1:i*n_embd/n_head,:), & - causal_mask) - end do -end associate -! Out projection -y = linear(y, proj_w, proj_b) -end function + ! Perform attention over each head + do l = 1, n_head + do i = 1, n_seq_x + do j = 1, n_embd / n_head + q(j, i) = x2((l - 1) * n_embd / n_head + j, i) + end do + end do + do i = 1, n_seq + do j = 1, n_embd / n_head + k(j, i) = kv_cache((l - 1) * n_embd / n_head + j, i, 1) + v(j, i) = kv_cache((l - 1) * n_embd / n_head + j, i, 2) + end do + end do + call attention(yy, n_embd / n_head, n_seq, n_seq_x, q, k, v, causal_mask) + do i = 1, n_seq_x + do j = 1, n_embd / n_head + y((l - 1) * n_embd / n_head + j, i) = yy(j, i) + end do + end do + end do + ! Out projection + y = linear(y, proj_w, proj_b) + end subroutine -function transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, & - attn_w, attn_b, attn_proj_w, attn_proj_b, ln1_g, ln1_b, ln2_g, ln2_b, & - n_head, use_kv_cache, kv_cache) result(y) -real(sp), intent(in) :: x(n_embd,n_seq_x), & - mlp_fc_w(:,:), mlp_fc_b(:), & - mlp_proj_w(:,:), mlp_proj_b(:), & - attn_w(:,:), attn_b(:), attn_proj_w(:,:), attn_proj_b(:), & - ln1_g(:), ln1_b(:), ln2_g(:), ln2_b(:) -integer, intent(in) :: n_head -integer, intent(in) :: n_seq, n_seq_x, n_embd -real(sp) :: y(n_embd,n_seq_x) -logical, intent(in) :: use_kv_cache -real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2) -y = x + mha(n_seq, n_seq_x, n_embd, layer_norm(x, ln1_g, ln1_b, 1e-5_sp), & - attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache) -y = y + ffn(layer_norm(y, ln2_g, ln2_b, 1e-5_sp), & - mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b) -end function + subroutine transformer_block(n_seq, n_seq_x, n_embd, x, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, & + attn_w, attn_b, attn_proj_w, attn_proj_b, ln1_g, ln1_b, ln2_g, ln2_b, & + n_head, use_kv_cache, kv_cache) + real(sp), intent(inout) :: x(n_embd, n_seq_x) + real(sp), intent(in) :: & + mlp_fc_w(:, :), mlp_fc_b(:), & + mlp_proj_w(:, :), mlp_proj_b(:), & + attn_w(:, :), attn_b(:), attn_proj_w(:, :), attn_proj_b(:), & + ln1_g(:), ln1_b(:), ln2_g(:), ln2_b(:) + integer, intent(in) :: n_head + integer, intent(in) :: n_seq, n_seq_x, n_embd + logical, intent(in) :: use_kv_cache + real(sp) :: y(n_embd, n_seq_x) + real(sp) :: yy(n_embd, n_seq_x) + real(sp), intent(inout) :: kv_cache(n_embd, n_seq, 2) + call layer_norm(y, x, ln1_g, ln1_b, 1e-5_sp) + call mha(yy, n_seq, n_seq_x, n_embd, y, & + attn_w, attn_b, attn_proj_w, attn_proj_b, n_head, use_kv_cache, kv_cache) + x = x + yy + !print *, "In: ", x(1,1), ln2_g(1), ln2_g(size(ln2_g)), ln2_b(1), ln2_b(size(ln2_b)) + call layer_norm(y, x, ln2_g, ln2_b, 1e-5_sp) + !print *, "Out1:", y(1,1) + !x = y + !print *, "Out2:", x(1,1) + call ffn(x, y, mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b) + end subroutine -function gpt2(n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, & - wte, wpe, & - mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, & - attn_w, attn_b, attn_proj_w, attn_proj_b, & - ln1_g, ln1_b, ln2_g, ln2_b, lnf_g, lnf_b, & - use_kv_cache, kv_cache) result(y) -integer, intent(in) :: n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head -integer, intent(in) :: input(n_seq) -real(sp), intent(in) :: wte(n_embd,n_vocab), wpe(n_embd,n_ctx), & - mlp_fc_w(4*n_embd,n_embd,n_layer), mlp_fc_b(4*n_embd,n_layer), & - mlp_proj_w(n_embd,4*n_embd,n_layer), mlp_proj_b(n_embd,n_layer), & - attn_w(3*n_embd,n_embd,n_layer), attn_b(3*n_embd,n_layer), & - attn_proj_w(n_embd,n_embd,n_layer), attn_proj_b(n_embd,n_layer), & - ln1_b(n_embd,n_layer), ln1_g(n_embd,n_layer), & - ln2_b(n_embd,n_layer), ln2_g(n_embd,n_layer), & - lnf_b(n_embd), lnf_g(n_embd) -logical, intent(in) :: use_kv_cache -real(sp), intent(inout) :: kv_cache(n_embd,n_seq,2,n_layer) -real(sp) :: y(n_vocab,n_seq_x) -real(sp) :: x(n_embd,n_seq_x) -integer :: i -if (use_kv_cache) then - i = n_seq - x(:,1) = wte(:,input(i)+1) + wpe(:,i) -else - do i = 1, n_seq - x(:,i) = wte(:,input(i)+1) + wpe(:,i) - end do -end if -do i = 1, n_layer - x = transformer_block(n_seq, n_seq_x, n_embd, x, & - mlp_fc_w(:,:,i), mlp_fc_b(:,i), & - mlp_proj_w(:,:,i), mlp_proj_b(:,i), & - attn_w(:,:,i), attn_b(:,i), attn_proj_w(:,:,i), attn_proj_b(:,i), & - ln1_g(:,i), ln1_b(:,i), ln2_g(:,i), ln2_b(:,i), & - n_head, use_kv_cache, kv_cache(:,:,:,i)) -end do -x = layer_norm(x, lnf_g, lnf_b, 1e-5) -!y = matmul(transpose(wte), x) -call matmul_2d_t(wte, x, y) -end function + subroutine gpt2(y, n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head, input, & + wte, wpe, & + mlp_fc_w, mlp_fc_b, mlp_proj_w, mlp_proj_b, & + attn_w, attn_b, attn_proj_w, attn_proj_b, & + ln1_g, ln1_b, ln2_g, ln2_b, lnf_g, lnf_b, & + use_kv_cache, kv_cache) + integer, intent(in) :: n_vocab, n_ctx, n_seq, n_seq_x, n_embd, n_layer, n_head + integer, intent(in) :: input(n_seq) + real(sp), intent(in) :: wte(n_embd, n_vocab), wpe(n_embd, n_ctx), & + mlp_fc_w(4 * n_embd, n_embd, n_layer), mlp_fc_b(4 * n_embd, n_layer), & + mlp_proj_w(n_embd, 4 * n_embd, n_layer), mlp_proj_b(n_embd, n_layer), & + attn_w(3 * n_embd, n_embd, n_layer), attn_b(3 * n_embd, n_layer), & + attn_proj_w(n_embd, n_embd, n_layer), attn_proj_b(n_embd, n_layer), & + ln1_b(n_embd, n_layer), ln1_g(n_embd, n_layer), & + ln2_b(n_embd, n_layer), ln2_g(n_embd, n_layer), & + lnf_b(n_embd), lnf_g(n_embd) + logical, intent(in) :: use_kv_cache + real(sp), intent(inout) :: kv_cache(n_embd, n_seq, 2, n_layer) + real(sp), intent(out) :: y(n_vocab, n_seq_x) + real(sp) :: x(n_embd, n_seq_x) + real(sp) :: yy(n_embd, n_seq_x) + integer :: i, j + if (use_kv_cache) then + i = n_seq + do j = 1, n_embd + x(j, 1) = wte(j, input(i) + 1) + wpe(j, i) + end do + else + do i = 1, n_seq + do j = 1, n_embd + x(j, i) = wte(j, input(i) + 1) + wpe(j, i) + end do + end do + end if + !print *, "It fails below:" + do j = 1, n_layer + i = j + !i = 1 + ! print *, i ! Never gets printed + call transformer_block(n_seq, n_seq_x, n_embd, x, & + mlp_fc_w(:, :, i), mlp_fc_b(:, i), & + mlp_proj_w(:, :, i), mlp_proj_b(:, i), & + attn_w(:, :, i), attn_b(:, i), attn_proj_w(:, :, i), attn_proj_b(:, i), & + ln1_g(:, i), ln1_b(:, i), ln2_g(:, i), ln2_b(:, i), & + n_head, use_kv_cache, kv_cache(:, :, :, i)) + ! print *, x(1,1) + end do + call layer_norm(yy, x, lnf_g, lnf_b, 1e-5) + x = yy + !y = matmul(transpose(wte), x) + call matmul_2d_t(wte, x, y) + end subroutine -function generate(n_tokens_to_generate, m, & - n_seq, input, & - use_cache, & - byte_decoder, stop_text) result(output) -integer, intent(in) :: n_seq, n_tokens_to_generate -type(model_t), intent(in) :: m -integer, intent(in) :: input(n_seq) -logical, intent(in) :: use_cache -integer, intent(in) :: byte_decoder(:) -character(*), intent(in), optional :: stop_text ! Stop if you see this text -integer, allocatable :: output(:) -real(sp), allocatable :: logits(:,:) -integer :: i -integer :: n_seq2, n_seq_x -integer :: next_id -integer, allocatable :: input2(:) -logical :: use_kv_cache -real(sp) :: kv_cache(m%n_embd,n_seq+n_tokens_to_generate,2,m%n_layer) -character(:), allocatable :: output_txt, last_token -allocate(input2(size(input))) -if (present(stop_text)) then - output_txt = "" -end if -input2 = input -do i = 1, n_tokens_to_generate - if (use_cache) then - use_kv_cache = (i > 1) ! Use cache for subsequent tokens - else - use_kv_cache = .false. - end if - n_seq2 = size(input2) - if (use_kv_cache) then - n_seq_x = 1 - else - n_seq_x = n_seq2 - end if - allocate(logits(m%n_vocab, n_seq_x)) - logits = gpt2(m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, & - m%n_head, & - input2, & - m%wte, m%wpe, & - m%mlp_fc_w, m%mlp_fc_b, m%mlp_proj_w, m%mlp_proj_b, & - m%attn_w, m%attn_b, m%attn_proj_w, m%attn_proj_b, & - m%ln1_g, m%ln1_b, m%ln2_g, m%ln2_b, m%lnf_g, m%lnf_b, use_kv_cache, kv_cache(:,:n_seq2,:,:)) - next_id = maxloc(logits(:,n_seq_x), dim=1)-1 - input2 = [input2, next_id] - last_token = decode([next_id], m%decoder_idx, & - m%decoder_txt, byte_decoder) - write(*, fmt="(a)", advance="no") last_token - if (present(stop_text)) then - output_txt = output_txt // last_token - if (output_txt(len(output_txt)-len(stop_text)+1:len(output_txt)) == stop_text) then - exit + subroutine generate(output, n_tokens_to_generate, m, & + n_seq, input, & + use_cache, & + byte_decoder, stop_text) + integer, intent(in) :: n_seq, n_tokens_to_generate + type(model_t), intent(in) :: m + integer, intent(in) :: input(n_seq) + logical, intent(in) :: use_cache + integer, intent(in) :: byte_decoder(:) + character(*), intent(in), optional :: stop_text ! Stop if you see this text + integer, allocatable, intent(out) :: output(:) + real(sp), allocatable :: logits(:, :) + integer :: i, i1, i2, i3, i4 + integer :: n_seq2, n_seq_x + integer :: next_id + integer :: input2(size(input) + n_tokens_to_generate) + logical :: use_kv_cache + real(sp) :: kv_cache(m%n_embd, n_seq + n_tokens_to_generate, 2, m%n_layer) + real(sp), allocatable :: kv_cache2(:, :, :, :) + character(:), allocatable :: output_txt, last_token + if (present(stop_text)) then + output_txt = "" end if - end if - deallocate(logits) -end do -output = input2(n_seq+1:) -end function + input2(:n_seq) = input + do i = 1, n_tokens_to_generate + if (use_cache) then + use_kv_cache = (i > 1) ! Use cache for subsequent tokens + else + use_kv_cache = .false. + end if + n_seq2 = n_seq + i - 1 + if (use_kv_cache) then + n_seq_x = 1 + else + n_seq_x = n_seq2 + end if + allocate(logits(m%n_vocab, n_seq_x)) + allocate(kv_cache2(m%n_embd, n_seq2, 2, m%n_layer)) + do i4 = 1, m%n_layer + do i3 = 1, 2 + do i2 = 1, n_seq2 + do i1 = 1, m%n_embd + kv_cache2(i1, i2, i3, i4) = kv_cache(i1, i2, i3, i4) + end do + end do + end do + end do + call gpt2(logits, m%n_vocab, m%n_ctx, n_seq2, n_seq_x, m%n_embd, m%n_layer, & + m%n_head, & + input2(:n_seq2), & + m%wte, m%wpe, & + m%mlp_fc_w, m%mlp_fc_b, m%mlp_proj_w, m%mlp_proj_b, & + m%attn_w, m%attn_b, m%attn_proj_w, m%attn_proj_b, & + m%ln1_g, m%ln1_b, m%ln2_g, m%ln2_b, m%lnf_g, m%lnf_b, use_kv_cache, & + kv_cache2) + do i4 = 1, m%n_layer + do i3 = 1, 2 + do i2 = 1, n_seq2 + do i1 = 1, m%n_embd + kv_cache(i1, i2, i3, i4) = kv_cache2(i1, i2, i3, i4) + end do + end do + end do + end do + deallocate(kv_cache2) + next_id = maxloc(logits(:, n_seq_x), dim = 1) - 1 + input2(n_seq2 + 1) = next_id + last_token = decode([next_id], m%decoder_idx, & + m%decoder_txt, byte_decoder) + write(*, fmt = "(a)", advance = "no") last_token + if (present(stop_text)) then + output_txt = output_txt // last_token + if (output_txt(len(output_txt) - len(stop_text) + 1:len(output_txt)) == stop_text) then + exit + end if + end if + deallocate(logits) + end do + ! output = input2(n_seq+1:n_seq2+1) + allocate(output(n_seq2 - n_seq + 1)) + do i = 1, n_seq2 - n_seq + 1 + output(i) = input2(n_seq + i) + end do + end subroutine end module diff --git a/linalg_c.f90 b/linalg_c.f90 index 790bb0e..db03267 100644 --- a/linalg_c.f90 +++ b/linalg_c.f90 @@ -36,6 +36,7 @@ subroutine matmul_2d_t(A, B, C) ! C = matmul(transpose(A), B) real(sp), intent(in) :: A(:,:), B(:,:) real(sp), intent(out) :: C(:,:) + !print *, "matmul_2d_t" call acc_sgemm_t(size(A,2), size(B,2), size(A,1), A, B, C) end subroutine diff --git a/linalg_openblas.c b/linalg_openblas.c index 63637c8..2c0c04a 100644 --- a/linalg_openblas.c +++ b/linalg_openblas.c @@ -14,5 +14,8 @@ void acc_sgemm_t(int m, int n, int k, float *A, float *B, float *C) { //A[k][m] (to be transposed) //B[k][n] //C[m][n] +// printf("acc_sgemm_t: %d %d %d\n", m, n, k); +// printf("Values: %f %f %f\n", A[0], B[0], C[0]); cblas_sgemm(CblasColMajor, CblasTrans, CblasNoTrans, m, n, k, 1.0, A, k, B, k, 0.0, C, m); +// printf("BLAS done\n"); } diff --git a/lp-fastgpt/README.md b/lp-fastgpt/README.md new file mode 100644 index 0000000..e69de29 diff --git a/lp-fastgpt/create_model.py b/lp-fastgpt/create_model.py new file mode 100644 index 0000000..0d3cb83 --- /dev/null +++ b/lp-fastgpt/create_model.py @@ -0,0 +1,758 @@ +"""This script loads the specified GPT-2 model from OpenAI +using TensorFlow, converts it into our custom format and saves +it to `model.dat`, which contains everything (all the +parameters, all the weights, encoding/decoding information). + +Parts of this script were taken from the picoGPT project: +https://github.com/jaymody/picoGPT + +Those are licensed as: + +MIT License + +Copyright (c) 2023 Jay Mody + +Permission is hereby granted, free of charge, to any person +obtaining a copy of this software and associated documentation +files (the "Software"), to deal in the Software without +restriction, including without limitation the rights to use, +copy, modify, merge, publish, distribute, sublicense, and/or +sell copies of the Software, and to permit persons to whom the +Software is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the +Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY +KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE +WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR +PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +""" + + +from time import monotonic as clock +import os +import json +import re + + +import numpy as np +import requests +import tensorflow as tf +from tqdm import tqdm + + +from typing import Optional, Union, Any + + +def download_gpt2_files(model_size, model_dir): + assert model_size in ["124M", "355M", "774M", "1558M"] + for filename in [ + "checkpoint", + "encoder.json", + "hparams.json", + "model.ckpt.data-00000-of-00001", + "model.ckpt.index", + "model.ckpt.meta", + "vocab.bpe", + ]: + url = "https://openaipublic.blob.core.windows.net/gpt-2/models" + r = requests.get(f"{url}/{model_size}/{filename}", stream=True) + r.raise_for_status() + + with open(os.path.join(model_dir, filename), "wb") as f: + file_size = int(r.headers["content-length"]) + chunk_size = 1000 + with tqdm( + ncols=100, + desc="Fetching " + filename, + total=file_size, + unit_scale=True, + ) as pbar: + # 1k for chunk_size, since Ethernet packet size is around 1500 bytes + for chunk in r.iter_content(chunk_size=chunk_size): + f.write(chunk) + pbar.update(chunk_size) + + +def load_gpt2_params_from_tf_ckpt(tf_ckpt_path, hparams): + def set_in_nested_dict(d, keys, val): + if not keys: + return val + if keys[0] not in d: + d[keys[0]] = {} + d[keys[0]] = set_in_nested_dict(d[keys[0]], keys[1:], val) + return d + + init_vars = tf.train.list_variables(tf_ckpt_path) + params = {"blocks": [{} for _ in range(hparams["n_layer"])]} + for name, _ in init_vars: + array = np.squeeze(tf.train.load_variable(tf_ckpt_path, name)) + name = name.removeprefix("model/") + if name.startswith("h"): + m = re.match(r"h([0-9]+)/(.*)", name) + n = int(m[1]) + sub_name = m[2] + set_in_nested_dict(params["blocks"][n], sub_name.split("/"), array) + else: + set_in_nested_dict(params, name.split("/"), array) + + return params + + +def tf_load_encoder_hparams_and_params(model_size, models_dir): + + assert model_size in ["124M", "355M", "774M", "1558M"] + + model_dir = os.path.join(models_dir, model_size) + + tf_ckpt_path = tf.train.latest_checkpoint(model_dir) + + if not tf_ckpt_path: # download files if necessary + os.makedirs(model_dir, exist_ok=True) + download_gpt2_files(model_size, model_dir) + tf_ckpt_path = tf.train.latest_checkpoint(model_dir) + + hparams = json.load(open(os.path.join(model_dir, "hparams.json"))) + + params = load_gpt2_params_from_tf_ckpt(tf_ckpt_path, hparams) + + return hparams, params + + +from dataclasses import dataclass + + +# === Magic (Unexplained) Numbers ========================= + +DecoderIdxType = np.ndarray +DecoderIdxShape = (50_258,) + +DecoderTxtType = str +DecoderTxtAsciiLen = 320_827 +DecoderTxtUtf8Len = 356_735 # It's checked somewhere below. + +VocabIdxType = np.ndarray +VocabIdxShape = (50_002,) + +VocabTxtType = str +VocabTxtAsciiLen = 370_558 +VocabTxtUtf8Len = 406_304 # It's checked somewhere below. +NVocab = 50_257 + +DecoderShape = (256,) +DecoderLen = DecoderShape[0] + +NBlocks = 12 + +ModelType = 0xfa51697 +ModelVersion = 1 + +NCtx = 1024 +NEmbed = 768 + +Magic2304 = 2304 +Magic3072 = 3072 + + +# === Block-Attention components ====================== +CAttnBType = np.ndarray +CAttnBShape = ( Magic2304,) +CAttnWType = np.ndarray +CAttnWShape = (NEmbed, Magic2304,) +CAttnType = dict[str, Union[CAttnBType, CAttnWType]] + +CProjBType = np.ndarray +CProjBShape = ( NEmbed,) +CProjWType = np.ndarray +CProjWShape = (NEmbed, NEmbed,) +CProjType = dict[str, Union[CProjBType, CProjWType]] +# --- top level +BlockAttnType = dict[str, Union[CAttnType, CProjType]] +# === BlockLn components ============================== +BlockLnBType = np.ndarray +BlockLnBShape = (NEmbed,) +BlockLnGType = np.ndarray +BlockLnGShape = (NEmbed,) +# --- top level +BlockLnType = dict[str, Union[BlockLnBType, BlockLnGType]] +# === BlockMlp components ============================= +MlpCFcBType = np.ndarray +MlpCFcBShape = ( Magic3072,) +MlpCFcWType = np.ndarray +MlpCFcWShape = (NEmbed, Magic3072,) +MlpCFcType = dict[str, Union[MlpCFcBType, MlpCFcWType]] + +MlpCProjBType = np.ndarray +MlpCProjBShape = ( NEmbed,) +MlpCProjWType = np.ndarray +MlpCProjWShape = (Magic3072, NEmbed,) +MlpCProjType = dict[str, Union[MlpCProjBType, MlpCProjWType]] +# --- top level +BlockMlpType = dict[str, Union[MlpCFcType, MlpCProjType]] +# ===================================================== +ParamsBlockType = dict[str, Union[BlockAttnType, + BlockLnType, # two of these + BlockMlpType]] + +ParamsBlocksType = list[ParamsBlockType] +ParamsLnFType = dict[str, np.ndarray] +ParamsLnFValShape = BlockLnBShape +ParamsWpeType = np.ndarray +ParamsWpeShape = ( NCtx, NEmbed,) +ParamsWteType = np.ndarray +ParamsWteShape = (NVocab, NEmbed,) +ParamsType = dict[str, Union[ParamsBlocksType, + ParamsLnFType, + ParamsWpeType, + ParamsWteType]] + +ModelMetadataType = np.ndarray +ModelMetadataShape = (12,) + +HParamsType = dict[str, int] + + +@dataclass +class Model: + # integer metadata + model_type : int + model_version : int + n_vocab : int + n_ctx : int + n_embd : int + n_layer : int + n_head : int + decoder_idx_len : int + decoder_txt_len : int + vocab_idx_len : int + vocab_txt_len : int + byte_decoder_len : int + # check shapes in asserts for now; type system too weak. + mlp_fc_w : np.ndarray + mlp_fc_b : np.ndarray + mlp_proj_w : np.ndarray + mlp_proj_b : np.ndarray + attn_w : np.ndarray + attn_b : np.ndarray + attn_proj_w : np.ndarray + attn_proj_b : np.ndarray + ln1_g : np.ndarray + ln1_b : np.ndarray + ln2_g : np.ndarray + ln2_b : np.ndarray + wte : np.ndarray + wpe : np.ndarray + lnf_g : np.ndarray + lnf_b : np.ndarray + # auxiliary matrices and texts + decoder_idx : np.ndarray + decoder_txt : str + vocab_idx : np.ndarray + vocab_txt : str + byte_decoder : np.ndarray + + +def convert(params, + n_head, + n_ctx, + decoder_idx, + decoder_txt, + vocab_idx, + vocab_txt, + byte_decoder) -> Model: + + t1 = clock() + + # must predefine just to get shapes ... + blocks : ParamsBlocksType = params["blocks"] + nblocks = len(blocks) + assert nblocks == NBlocks + + n_embd = blocks[0]["ln_1"]["b"].size + n_layer = nblocks + assert n_layer == NBlocks + + n_vocab = ParamsWteShape[0] # np.size(mo.wte, 0) + model_type = ModelType + model_version = ModelVersion + + mo : Model = make_empty_model_with_metadata( + model_type = model_type, + model_version = model_version, + n_vocab = n_vocab, + n_ctx = n_ctx, + n_embd = n_embd, + n_layer = n_layer, + n_head = n_head, + decoder_idx_len = DecoderIdxShape[0], + # It's useful to check magic numbers against computations: + decoder_txt_len = len(decoder_txt.encode("utf-8")), + vocab_idx_len = VocabIdxShape[0], + vocab_txt_len = len(vocab_txt.encode("utf-8")), + byte_decoder_len = len(byte_decoder), + ) + + for i, block in enumerate(blocks): + mo.mlp_fc_w[i, :, :] = block["mlp"]["c_fc"]["w"] + mo.mlp_fc_b[i, :] = block["mlp"]["c_fc"]["b"] + mo.mlp_proj_w[i, :, :] = block["mlp"]["c_proj"]["w"] + mo.mlp_proj_b[i, :] = block["mlp"]["c_proj"]["b"] + mo.attn_w[i, :, :] = block["attn"]["c_attn"]["w"] + mo.attn_b[i, :] = block["attn"]["c_attn"]["b"] + mo.attn_proj_w[i, :, :] = block["attn"]["c_proj"]["w"] + mo.attn_proj_b[i, :] = block["attn"]["c_proj"]["b"] + + mo.ln1_g[i, :] = block["ln_1"]["g"] + mo.ln1_b[i, :] = block["ln_1"]["b"] + mo.ln2_g[i, :] = block["ln_2"]["g"] + mo.ln2_b[i, :] = block["ln_2"]["b"] + + mo.wte = params["wte"] + mo.wpe = params["wpe"] + + mo.lnf_g = params["ln_f"]["g"] + mo.lnf_b = params["ln_f"]["b"] + + mo.decoder_idx = decoder_idx + mo.decoder_txt = decoder_txt + mo.vocab_idx = vocab_idx + mo.vocab_txt = vocab_txt + mo.byte_decoder = byte_decoder + + t2 = clock() + print("Transform time: ", t2 - t1) + + check_model(mo) + + # Save the model + t1 = clock() + save_model(mo) + t2 = clock() + print("Save time: ", t2 - t1) + + # Round-trip the model: check the stored version. + + t1 = clock() + m = restore_model() + + check_model(m) + assert np.all(m.wte == mo.wte) + assert np.all(m.wpe == mo.wpe) + assert np.all(m.mlp_fc_w == mo.mlp_fc_w) + assert np.all(m.mlp_fc_b == mo.mlp_fc_b) + assert np.all(m.mlp_proj_w == mo.mlp_proj_w) + assert np.all(m.mlp_proj_b == mo.mlp_proj_b) + assert np.all(m.attn_w == mo.attn_w) + assert np.all(m.attn_b == mo.attn_b) + assert np.all(m.attn_proj_w == mo.attn_proj_w) + assert np.all(m.attn_proj_b == mo.attn_proj_b) + assert np.all(m.ln1_b == mo.ln1_b) + assert np.all(m.ln1_g == mo.ln1_g) + assert np.all(m.ln2_b == mo.ln2_b) + assert np.all(m.ln2_g == mo.ln2_g) + assert np.all(m.lnf_b == mo.lnf_b) + assert np.all(m.lnf_g == mo.lnf_g) + assert np.all(m.decoder_idx == mo.decoder_idx) + assert m.decoder_txt == mo.decoder_txt + assert np.all(m.vocab_idx == mo.vocab_idx) + assert m.vocab_txt == mo.vocab_txt + assert np.all(m.byte_decoder == mo.byte_decoder) + t2 = clock() + print("Restore time: ", t2 - t1) + + return m + + +def save_model(mo : Model): + with open("model.dat", "w") as f: + model_metadata = np.array( + [ + mo.model_type, + mo.model_version, + mo.n_vocab, + mo.n_ctx, + mo.n_embd, + mo.n_layer, + mo.n_head, + len(mo.decoder_idx), + len(mo.decoder_txt.encode("utf-8")), + len(mo.vocab_idx), + len(mo.vocab_txt.encode("utf-8")), + len(mo.byte_decoder)], + dtype=np.int32) + + assert model_metadata.shape == ModelMetadataShape + + model_metadata.tofile(f) + + mo.wte.tofile(f) + mo.wpe.tofile(f) + + mo.mlp_fc_w.tofile(f) + mo.mlp_fc_b.tofile(f) + + mo.mlp_proj_w.tofile(f) + mo.mlp_proj_b.tofile(f) + + mo.attn_w.tofile(f) + mo.attn_b.tofile(f) + + mo.attn_proj_w.tofile(f) + mo.attn_proj_b.tofile(f) + + mo.ln1_b.tofile(f) + mo.ln1_g.tofile(f) + mo.ln2_b.tofile(f) + mo.ln2_g.tofile(f) + + mo.lnf_b.tofile(f) + mo.lnf_g.tofile(f) + + mo.decoder_idx.tofile(f) + f.write(mo.decoder_txt) + mo.vocab_idx.tofile(f) + f.write(mo.vocab_txt) + mo.byte_decoder.tofile(f) + + +def restore_model() -> Model: + m : Model = make_empty_model_with_metadata( + model_type = 0, + model_version = 0, + n_vocab = 0, + n_ctx = 0, + n_embd = 0, + n_layer = 0, + n_head = 0, + decoder_idx_len = 0, + decoder_txt_len = 0, + vocab_idx_len = 0, + vocab_txt_len = 0, + byte_decoder_len = 0, + ) + + floff : int = 0 + metadata : np.ndarray = ( + np.fromfile("model.dat", + dtype=np.int32, + count=ModelMetadataShape[0], + offset=floff)) + + m.model_type = metadata[ 0] + m.model_version = metadata[ 1] + m.n_vocab = metadata[ 2] + m.n_ctx = metadata[ 3] + m.n_embd = metadata[ 4] + m.n_layer = metadata[ 5] + m.n_head = metadata[ 6] + m.decoder_idx_len = metadata[ 7] + m.decoder_txt_len = metadata[ 8] + m.vocab_idx_len = metadata[ 9] + m.vocab_txt_len = metadata[10] + m.byte_decoder_len = metadata[11] + + check_model_metadata(m) + + floff += ModelMetadataShape[0] * BYTES_PER_INT32 + floff, m.wte = restore_floats(ParamsWteShape , floff) + floff, m.wpe = restore_floats(ParamsWpeShape , floff) + floff, m.mlp_fc_w = restore_floats((NBlocks,) + MlpCFcWShape , floff) + floff, m.mlp_fc_b = restore_floats((NBlocks,) + MlpCFcBShape , floff) + floff, m.mlp_proj_w = restore_floats((NBlocks,) + MlpCProjWShape , floff) + floff, m.mlp_proj_b = restore_floats((NBlocks,) + MlpCProjBShape , floff) + floff, m.attn_w = restore_floats((NBlocks,) + CAttnWShape , floff) + floff, m.attn_b = restore_floats((NBlocks,) + CAttnBShape , floff) + floff, m.attn_proj_w = restore_floats((NBlocks,) + CProjWShape , floff) + floff, m.attn_proj_b = restore_floats((NBlocks,) + CProjBShape , floff) + floff, m.ln1_b = restore_floats((NBlocks,) + BlockLnBShape , floff) + floff, m.ln1_g = restore_floats((NBlocks,) + BlockLnBShape , floff) + floff, m.ln2_b = restore_floats((NBlocks,) + BlockLnBShape , floff) + floff, m.ln2_g = restore_floats((NBlocks,) + BlockLnBShape , floff) + floff, m.lnf_b = restore_floats(ParamsLnFValShape , floff) + floff, m.lnf_g = restore_floats(ParamsLnFValShape , floff) + floff, m.decoder_idx = restore_ints(DecoderIdxShape , floff) + + with open("model.dat", "rb") as f: + f.seek(floff) + decoder_txti_ub : bytes = f.read(DecoderTxtUtf8Len) + m.decoder_txt = decoder_txti_ub.decode("utf-8") + assert len(m.decoder_txt) == DecoderTxtAsciiLen + floff += DecoderTxtUtf8Len + + floff, m.vocab_idx = restore_ints(VocabIdxShape, floff) + + with open("model.dat", "rb") as f: + f.seek(floff) + vocab_txt_ub : bytes = f.read(VocabTxtUtf8Len) + m.vocab_txt = vocab_txt_ub.decode("utf-8") + assert len(m.vocab_txt) == VocabTxtAsciiLen + floff += VocabTxtUtf8Len + + floff, m.byte_decoder = restore_ints(DecoderShape, floff) + + return m + + +def prod_tuple(t : tuple[int]) -> int: + result : int = 1 + for e in t: + result *= e + return result + + +BYTES_PER_INT32 = 4 +BYTES_PER_FLOAT32 = 4 + + +def restore_floats(shape : tuple, offset : int) -> tuple[int, np.ndarray]: + """agnostic to length of shape; TODO impossible to statically type""" + result : np.ndarray + count = prod_tuple(shape) + result = np.fromfile("model.dat", + dtype=np.float32, + count=count, + offset=offset) + result = np.reshape(result, shape) + new_offset : int = offset + (count * BYTES_PER_FLOAT32) + return new_offset, result + + +def restore_ints(shape : tuple, offset : int) -> tuple[int, np.ndarray]: + """agnostic to length of shape; TODO impossible to statically type""" + result : np.ndarray + count = prod_tuple(shape) + result = np.fromfile("model.dat", + dtype=np.int32, + count=count, + offset=offset) + result = np.reshape(result, shape) + new_offset : int = offset + (count * BYTES_PER_INT32) + return new_offset, result + + +def make_empty_model_with_metadata( + model_type : int, + model_version : int, + n_vocab : int, + n_ctx : int, + n_embd : int, + n_layer : int, + n_head : int, + decoder_idx_len : int, + decoder_txt_len : int, + vocab_idx_len : int, + vocab_txt_len : int, + byte_decoder_len : int,) -> Model: + + mo: Model = Model( + model_type = model_type, + model_version = model_version, + n_vocab = n_vocab, + n_ctx = n_ctx, + n_embd = n_embd, + n_layer = n_layer, + n_head = n_head, + decoder_idx_len = decoder_idx_len, + decoder_txt_len = decoder_txt_len, + vocab_idx_len = vocab_idx_len, + vocab_txt_len = vocab_txt_len, + byte_decoder_len = byte_decoder_len, + + mlp_fc_w = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32), + mlp_fc_b = np.empty((n_layer, 4 * n_embd) , dtype=np.float32), + mlp_proj_w = np.empty((n_layer, 4 * n_embd, n_embd) , dtype=np.float32), + mlp_proj_b = np.empty((n_layer, n_embd) , dtype=np.float32), + attn_w = np.empty((n_layer, n_embd, 3 * n_embd) , dtype=np.float32), + attn_b = np.empty((n_layer, 3 * n_embd) , dtype=np.float32), + attn_proj_w = np.empty((n_layer, n_embd, n_embd) , dtype=np.float32), + attn_proj_b = np.empty((n_layer, n_embd) , dtype=np.float32), + ln1_g = np.empty((n_layer, n_embd) , dtype=np.float32), + ln1_b = np.empty((n_layer, n_embd) , dtype=np.float32), + ln2_g = np.empty((n_layer, n_embd) , dtype=np.float32), + ln2_b = np.empty((n_layer, n_embd) , dtype=np.float32), + wte = np.empty(0 , dtype=np.float32), + wpe = np.empty(0 , dtype=np.float32), + lnf_g = np.empty(0 , dtype=np.float32), + lnf_b = np.empty(0 , dtype=np.float32), + + decoder_idx = np.empty(0 , dtype=np.int32), + decoder_txt = '', + vocab_idx = np.empty(0 , dtype=np.int32), + vocab_txt = '', + byte_decoder = np.empty(0 , dtype=np.int32), + ) + return mo + + +def check_model(mo : Model) -> None: + check_model_metadata(mo) + + assert mo.mlp_fc_w.shape == (NBlocks,) + MlpCFcWShape + assert mo.mlp_fc_b.shape == (NBlocks,) + MlpCFcBShape + assert mo.mlp_proj_w.shape == (NBlocks,) + MlpCProjWShape + assert mo.mlp_proj_b.shape == (NBlocks,) + MlpCProjBShape + assert mo.attn_w.shape == (NBlocks,) + CAttnWShape + assert mo.attn_b.shape == (NBlocks,) + CAttnBShape + assert mo.attn_proj_w.shape == (NBlocks,) + CProjWShape + assert mo.attn_proj_b.shape == (NBlocks,) + CProjBShape + assert mo.ln1_g.shape == (NBlocks,) + BlockLnGShape + assert mo.ln1_b.shape == (NBlocks,) + BlockLnBShape + assert mo.ln2_g.shape == (NBlocks,) + BlockLnGShape + assert mo.ln2_b.shape == (NBlocks,) + BlockLnBShape + assert mo.wte.shape == ParamsWteShape + assert mo.wpe.shape == ParamsWpeShape + assert mo.lnf_g.shape == BlockLnGShape + assert mo.lnf_b.shape == BlockLnBShape + + assert mo.n_vocab == ParamsWteShape[0] + assert np.size(mo.wte, 1) == NEmbed + assert mo.decoder_idx.shape == (mo.decoder_idx_len,) + assert mo.vocab_idx.shape == (mo.vocab_idx_len,) + assert mo.byte_decoder.shape == (mo.byte_decoder_len,) + +def check_model_metadata(mo): + assert mo.model_type == ModelType + assert mo.model_version == ModelVersion + assert mo.n_vocab == NVocab + assert mo.n_ctx == NCtx + assert mo.n_embd == NEmbed + assert mo.n_layer == NBlocks + assert mo.n_head == NBlocks + + assert mo.decoder_idx_len == DecoderIdxShape[0] + assert mo.decoder_txt_len == DecoderTxtUtf8Len + assert mo.vocab_idx_len == VocabIdxShape[0] + assert mo.vocab_txt_len == VocabTxtUtf8Len + assert mo.byte_decoder_len == DecoderShape[0] + + +def load_decoder(filename): + D = json.load(open(filename)) + D2 = {v: k for k, v in D.items()} + i = 0 + decoder = [] + while True: + if i not in D2: + break + decoder.append(D2[i]) + i += 1 + return decoder + + +def load_vocab(filename): + D = open(filename).read() + D = D.split("\n") + D = D[1:] + return D + + +def load_decoder_idx(decoder): + i = 0 + idx = np.empty(len(decoder) + 1, dtype=np.int32) + idx[0] = i + for n, t in enumerate(decoder): + i += len(t.encode("utf-8")) + idx[n + 1] = i + assert idx[-1] == len("".join(decoder).encode("utf-8")) + return idx + + +def bytes_to_unicode() -> np.ndarray: + bs = list(range(ord("!"), ord("~") + 1)) + \ + list(range(ord("¡"), ord("¬") + 1)) + \ + list(range(ord("®"), ord("ÿ") + 1)) + cs = bs[:] + n = 0 + for b in range(2 ** 8): + if b not in bs: + bs.append(b) + cs.append(2 ** 8 + n) + n += 1 + cs = [chr(n) for n in cs] + btu = dict(zip(bs, cs)) + byte_decoder = {v: k for k, v in btu.items()} + bd = np.zeros(324, dtype=np.int32) + for y in byte_decoder: + x = ord(y) + bd[x] = byte_decoder[y] + bd2 = np.zeros(DecoderShape, dtype=np.int32) + for i in range(np.size(bd)): + bd2[bd[i]] = i + return bd2 + + +def main(model_size: str = "124M", + models_dir: str = "models") -> Model: + + # ================================================================ + # load encoder, hparams, and params from the released open-ai gpt-2 files + print("Loading model") + t1 = clock() + + hparams : HParamsType + params : ParamsType + + hparams, params = \ + tf_load_encoder_hparams_and_params(model_size, models_dir) + + decoder : list[str] = \ + load_decoder(os.path.join(models_dir, model_size, "encoder.json")) + + assert len(decoder) == NVocab # TODO: ??? !!! ??? + + vocab : list[str] = \ + load_vocab(os.path.join(models_dir, model_size, "vocab.bpe")) + + assert len(vocab) == 50_001 + + t2 = clock() + print(" Done. Loading time: ", t2 - t1) + # ================================================================ + # generate output ids + print("Converting model, saving to `model.dat`") + t1 = clock() + + decoder_idx : np.ndarray = load_decoder_idx(decoder) + assert decoder_idx.shape == DecoderIdxShape + + decoder_txt : str = "".join(decoder) + assert len(decoder_txt) == DecoderTxtAsciiLen + + vocab_idx = load_decoder_idx(vocab) + assert vocab_idx.shape == VocabIdxShape + + vocab_txt = "".join(vocab) + assert len(vocab_txt) == VocabTxtAsciiLen + + byte_decoder = bytes_to_unicode() + assert byte_decoder.shape == DecoderShape + + m : Model = \ + convert(params, + hparams["n_head"], + hparams["n_ctx"], + decoder_idx, + decoder_txt, + vocab_idx, + vocab_txt, + byte_decoder) + + t2 = clock() + print(" Done. Time: ", t2 - t1) + # ================================================================ + + return m + + +if __name__ == "__main__": + import fire + + fire.Fire(main) diff --git a/lp-fastgpt/main.py b/lp-fastgpt/main.py new file mode 100644 index 0000000..87bc256 --- /dev/null +++ b/lp-fastgpt/main.py @@ -0,0 +1,6 @@ +from restore_model import Model, restore_model + + +if __name__ == '__main__': + print("hello lp_fastgpt") + m : Model = restore_model() diff --git a/lp-fastgpt/requirements.txt b/lp-fastgpt/requirements.txt new file mode 100644 index 0000000..e69de29 diff --git a/lp-fastgpt/restore_model.py b/lp-fastgpt/restore_model.py new file mode 100644 index 0000000..f9721d3 --- /dev/null +++ b/lp-fastgpt/restore_model.py @@ -0,0 +1,517 @@ +"""This script restores the model from `model.dat`, which +contains everything (all the parameters, all the weights, +encoding/decoding information). + +Parts of this script were taken from the picoGPT project: +https://github.com/jaymody/picoGPT + +Those are licensed as: + +MIT License + +Copyright (c) 2023 Jay Mody + +Permission is hereby granted, free of charge, to any person +obtaining a copy of this software and associated documentation +files (the "Software"), to deal in the Software without +restriction, including without limitation the rights to use, +copy, modify, merge, publish, distribute, sublicense, and/or +sell copies of the Software, and to permit persons to whom the +Software is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the +Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY +KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE +WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR +PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +""" + +import numpy as np +from typing import Union + +from dataclasses import dataclass + +from timer import Timer + + +# === Magic (Unexplained) Numbers ========================= +DecoderIdxType = np.ndarray +DecoderIdxShape = (50_258,) + +DecoderTxtType = str +DecoderTxtAsciiLen = 320_827 +DecoderTxtUtf8Len = 356_735 # It's checked somewhere below. + +VocabIdxType = np.ndarray +VocabIdxShape = (50_002,) + +VocabTxtType = str +VocabTxtAsciiLen = 370_558 +VocabTxtUtf8Len = 406_304 # It's checked somewhere below. +NVocab = 50_257 + +DecoderShape = (256,) +DecoderLen = DecoderShape[0] + +NBlocks = 12 + +ModelType = 0xfa51697 +ModelVersion = 1 + +NCtx = 1024 +NEmbed = 768 + +Magic2304 = 2304 +Magic3072 = 3072 + +# === Block-Attention components ====================== +CAttnBType = np.ndarray +CAttnBShape = ( Magic2304,) +CAttnWType = np.ndarray +CAttnWShape = (NEmbed, Magic2304,) +CAttnType = dict[str, Union[CAttnBType, CAttnWType]] + +CProjBType = np.ndarray +CProjBShape = ( NEmbed,) +CProjWType = np.ndarray +CProjWShape = (NEmbed, NEmbed,) +CProjType = dict[str, Union[CProjBType, CProjWType]] +# --- top level +BlockAttnType = dict[str, Union[CAttnType, CProjType]] + +# === BlockLn components ============================== +BlockLnBType = np.ndarray +BlockLnBShape = (NEmbed,) +BlockLnGType = np.ndarray +BlockLnGShape = (NEmbed,) +# --- top level +BlockLnType = dict[str, Union[BlockLnBType, BlockLnGType]] + +# === BlockMlp components ============================= +MlpCFcBType = np.ndarray +MlpCFcBShape = ( Magic3072,) +MlpCFcWType = np.ndarray +MlpCFcWShape = (NEmbed, Magic3072,) +MlpCFcType = dict[str, Union[MlpCFcBType, MlpCFcWType]] + +MlpCProjBType = np.ndarray +MlpCProjBShape = ( NEmbed,) +MlpCProjWType = np.ndarray +MlpCProjWShape = (Magic3072, NEmbed,) +MlpCProjType = dict[str, Union[MlpCProjBType, MlpCProjWType]] +# --- top level +BlockMlpType = dict[str, Union[MlpCFcType, MlpCProjType]] + +# ===================================================== +ParamsBlockType = dict[str, Union[BlockAttnType, + BlockLnType, # two of these + BlockMlpType]] + +ParamsBlocksType = list[ParamsBlockType] +ParamsLnFType = dict[str, np.ndarray] +ParamsLnFValShape = BlockLnBShape +ParamsWpeType = np.ndarray +ParamsWpeShape = ( NCtx, NEmbed,) +ParamsWteType = np.ndarray +ParamsWteShape = (NVocab, NEmbed,) +ParamsType = dict[str, Union[ParamsBlocksType, + ParamsLnFType, + ParamsWpeType, + ParamsWteType]] + +ModelMetadataType = np.ndarray +ModelMetadataShape = (12,) + +NTokensToGenerate = 20 +MaxTokens = 2048 + +@dataclass +class Model: + # integer metadata + model_type : int + model_version : int + n_vocab : int + n_ctx : int + n_embd : int + n_layer : int + n_head : int + decoder_idx_len : int + decoder_txt_len : int + vocab_idx_len : int + vocab_txt_len : int + byte_encoder_len : int + # check shapes in asserts for now; type system too weak. + mlp_fc_w : np.ndarray + mlp_fc_b : np.ndarray + mlp_proj_w : np.ndarray + mlp_proj_b : np.ndarray + attn_w : np.ndarray + attn_b : np.ndarray + attn_proj_w : np.ndarray + attn_proj_b : np.ndarray + ln1_g : np.ndarray + ln1_b : np.ndarray + ln2_g : np.ndarray + ln2_b : np.ndarray + wte : np.ndarray + wpe : np.ndarray + lnf_g : np.ndarray + lnf_b : np.ndarray + # auxiliary matrices and texts + decoder_idx : np.ndarray + decoder_txt : str + vocab_idx : np.ndarray + vocab_txt : str + byte_encoder : np.ndarray + + +def restore_model() -> Model: + m : Model = empty_model_with_metadata( + model_type = 0, + model_version = 0, + n_vocab = 0, + n_ctx = 0, + n_embd = 0, + n_layer = 0, + n_head = 0, + decoder_idx_len = 0, + decoder_txt_len = 0, + vocab_idx_len = 0, + vocab_txt_len = 0, + byte_encoder_len = 0, + ) + + floff : int = 0 + metadata : np.ndarray = ( + np.fromfile("model.dat", + dtype=np.int32, + count=ModelMetadataShape[0], + offset=floff)) + + m.model_type = metadata[ 0] + m.model_version = metadata[ 1] + m.n_vocab = metadata[ 2] + m.n_ctx = metadata[ 3] + m.n_embd = metadata[ 4] + m.n_layer = metadata[ 5] + m.n_head = metadata[ 6] + m.decoder_idx_len = metadata[ 7] + m.decoder_txt_len = metadata[ 8] + m.vocab_idx_len = metadata[ 9] + m.vocab_txt_len = metadata[10] + m.byte_encoder_len = metadata[11] + + check_model_metadata(m) + + floff += ModelMetadataShape[0] * BYTES_PER_INT32 + floff, m.wte = restore_floats(ParamsWteShape , floff) + floff, m.wpe = restore_floats(ParamsWpeShape , floff) + floff, m.mlp_fc_w = restore_floats((NBlocks,) + MlpCFcWShape , floff) + floff, m.mlp_fc_b = restore_floats((NBlocks,) + MlpCFcBShape , floff) + floff, m.mlp_proj_w = restore_floats((NBlocks,) + MlpCProjWShape , floff) + floff, m.mlp_proj_b = restore_floats((NBlocks,) + MlpCProjBShape , floff) + floff, m.attn_w = restore_floats((NBlocks,) + CAttnWShape , floff) + floff, m.attn_b = restore_floats((NBlocks,) + CAttnBShape , floff) + floff, m.attn_proj_w = restore_floats((NBlocks,) + CProjWShape , floff) + floff, m.attn_proj_b = restore_floats((NBlocks,) + CProjBShape , floff) + floff, m.ln1_b = restore_floats((NBlocks,) + BlockLnBShape , floff) + floff, m.ln1_g = restore_floats((NBlocks,) + BlockLnBShape , floff) + floff, m.ln2_b = restore_floats((NBlocks,) + BlockLnBShape , floff) + floff, m.ln2_g = restore_floats((NBlocks,) + BlockLnBShape , floff) + floff, m.lnf_b = restore_floats(ParamsLnFValShape , floff) + floff, m.lnf_g = restore_floats(ParamsLnFValShape , floff) + floff, m.decoder_idx = restore_ints(DecoderIdxShape , floff) + + with open("model.dat", "rb") as f: + f.seek(floff) + decoder_txti_ub : bytes = f.read(DecoderTxtUtf8Len) + m.decoder_txt = decoder_txti_ub.decode("utf-8") + assert len(m.decoder_txt) == DecoderTxtAsciiLen + floff += DecoderTxtUtf8Len + + floff, m.vocab_idx = restore_ints(VocabIdxShape, floff) + + with open("model.dat", "rb") as f: + f.seek(floff) + vocab_txt_ub : bytes = f.read(VocabTxtUtf8Len) + m.vocab_txt = vocab_txt_ub.decode("utf-8") + assert len(m.vocab_txt) == VocabTxtAsciiLen + floff += VocabTxtUtf8Len + + floff, m.byte_encoder = restore_ints(DecoderShape, floff) + + return m + + +def prod_tuple(t : tuple[int]) -> int: + result : int = 1 + for e in t: + result *= e + return result + + +BYTES_PER_INT32 = 4 +BYTES_PER_FLOAT32 = 4 + + +def restore_floats(shape : tuple, offset : int) -> tuple[int, np.ndarray]: + """agnostic to length of shape; TODO impossible to statically type""" + result : np.ndarray + count = prod_tuple(shape) + result = np.fromfile("model.dat", + dtype=np.float32, + count=count, + offset=offset) + result = np.reshape(result, shape) + new_offset : int = offset + (count * BYTES_PER_FLOAT32) + return new_offset, result + + +def restore_ints(shape : tuple, offset : int) -> tuple[int, np.ndarray]: + """agnostic to length of shape; TODO impossible to statically type""" + result : np.ndarray + count = prod_tuple(shape) + result = np.fromfile("model.dat", + dtype=np.int32, + count=count, + offset=offset) + result = np.reshape(result, shape) + new_offset : int = offset + (count * BYTES_PER_INT32) + return new_offset, result + + +def empty_model_with_metadata( + model_type : int, + model_version : int, + n_vocab : int, + n_ctx : int, + n_embd : int, + n_layer : int, + n_head : int, + decoder_idx_len : int, + decoder_txt_len : int, + vocab_idx_len : int, + vocab_txt_len : int, + byte_encoder_len : int,) -> Model: + + mo: Model = Model( + model_type = model_type, + model_version = model_version, + n_vocab = n_vocab, + n_ctx = n_ctx, + n_embd = n_embd, + n_layer = n_layer, + n_head = n_head, + decoder_idx_len = decoder_idx_len, + decoder_txt_len = decoder_txt_len, + vocab_idx_len = vocab_idx_len, + vocab_txt_len = vocab_txt_len, + byte_encoder_len = byte_encoder_len, + + mlp_fc_w = np.empty((n_layer, n_embd, 4 * n_embd) , dtype=np.float32), + mlp_fc_b = np.empty((n_layer, 4 * n_embd) , dtype=np.float32), + mlp_proj_w = np.empty((n_layer, 4 * n_embd, n_embd) , dtype=np.float32), + mlp_proj_b = np.empty((n_layer, n_embd) , dtype=np.float32), + attn_w = np.empty((n_layer, n_embd, 3 * n_embd) , dtype=np.float32), + attn_b = np.empty((n_layer, 3 * n_embd) , dtype=np.float32), + attn_proj_w = np.empty((n_layer, n_embd, n_embd) , dtype=np.float32), + attn_proj_b = np.empty((n_layer, n_embd) , dtype=np.float32), + ln1_g = np.empty((n_layer, n_embd) , dtype=np.float32), + ln1_b = np.empty((n_layer, n_embd) , dtype=np.float32), + ln2_g = np.empty((n_layer, n_embd) , dtype=np.float32), + ln2_b = np.empty((n_layer, n_embd) , dtype=np.float32), + wte = np.empty(0 , dtype=np.float32), + wpe = np.empty(0 , dtype=np.float32), + lnf_g = np.empty(0 , dtype=np.float32), + lnf_b = np.empty(0 , dtype=np.float32), + + decoder_idx = np.empty(0 , dtype=np.int32), + decoder_txt = '', + vocab_idx = np.empty(0 , dtype=np.int32), + vocab_txt = '', + byte_encoder = np.empty(0 , dtype=np.int32), + ) + return mo + + +def check_model(mo : Model) -> None: + check_model_metadata(mo) + + assert mo.mlp_fc_w.shape == (NBlocks,) + MlpCFcWShape + assert mo.mlp_fc_b.shape == (NBlocks,) + MlpCFcBShape + assert mo.mlp_proj_w.shape == (NBlocks,) + MlpCProjWShape + assert mo.mlp_proj_b.shape == (NBlocks,) + MlpCProjBShape + assert mo.attn_w.shape == (NBlocks,) + CAttnWShape + assert mo.attn_b.shape == (NBlocks,) + CAttnBShape + assert mo.attn_proj_w.shape == (NBlocks,) + CProjWShape + assert mo.attn_proj_b.shape == (NBlocks,) + CProjBShape + assert mo.ln1_g.shape == (NBlocks,) + BlockLnGShape + assert mo.ln1_b.shape == (NBlocks,) + BlockLnBShape + assert mo.ln2_g.shape == (NBlocks,) + BlockLnGShape + assert mo.ln2_b.shape == (NBlocks,) + BlockLnBShape + assert mo.wte.shape == ParamsWteShape + assert mo.wpe.shape == ParamsWpeShape + assert mo.lnf_g.shape == BlockLnGShape + assert mo.lnf_b.shape == BlockLnBShape + + assert mo.n_vocab == ParamsWteShape[0] + assert np.size(mo.wte, 1) == NEmbed + assert mo.decoder_idx.shape == (mo.decoder_idx_len,) + assert mo.vocab_idx.shape == (mo.vocab_idx_len,) + assert mo.byte_encoder.shape == (mo.byte_encoder_len,) + + +def check_model_metadata(mo): + assert mo.model_type == ModelType + assert mo.model_version == ModelVersion + assert mo.n_vocab == NVocab + assert mo.n_ctx == NCtx + assert mo.n_embd == NEmbed + assert mo.n_layer == NBlocks + assert mo.n_head == NBlocks + + assert mo.decoder_idx_len == DecoderIdxShape[0] + assert mo.decoder_txt_len == DecoderTxtUtf8Len + assert mo.vocab_idx_len == VocabIdxShape[0] + assert mo.vocab_txt_len == VocabTxtUtf8Len + assert mo.byte_encoder_len == DecoderShape[0] + + +def tokenize_word(input_ : str, i : int) -> tuple[str, int]: + result = ('', i) + i0 : int = i + if input_ == ' ': # only one leading space ? + i += 1 + while True: + ci: str = input_[i] + if i >= len(input_) or ci == ' ' or ci == '.' or ci == ',': + result = (input_[i0:i], i) + return result + i += 1 + + +def next_token(input_ : str, i : int) -> tuple[str, int]: + result = ('', i) + ci : str = input_[i] + if i >= len(input_): + return result + elif ci == ' ': + result = tokenize_word(input_, i) + elif ci == ',' or ci == '.': + i += 1 + result = (ci, i) + else: + result = tokenize_word(input_, i) + return result + + +def merge_pair(intokens : bytearray, idx : int) -> bytearray: + tokens : bytearray + merged_token : bytearray = bytearray(2) + i : str + + merged_token[0] = intokens[idx] + merged_token[1] = intokens[idx + 1] + tokens = intokens.copy() + tokens + return tokens + + +def merge_utf8_pairs(token : str) -> str: + tokens : bytearray + i : int + j : int + ic : int + one_more_pass : bool = True + + tokens = bytearray(token.encode("utf-8")) + j = 0 + while one_more_pass: + one_more_pass = False + for i in range(j, len(tokens)): + ic : int = tokens[i] + if ic < 0: + ic += 256 + if ic >= 128: + tokens = merge_pair(tokens, i) + one_more_pass = True + j = i + 1 + return tokens.decode("utf-8") + + +def bpe(m : Model, token : str) -> str: + tokens : str = '' + pair_scores = [] + not_found : int = 0 + merge_pair_idx : int = 0 + i : int = 0 + MAGIC : int = 10 + + not_found = len(m.vocab_idx) + MAGIC + tokens = merge_utf8_pairs(token) + + return tokens + + +def encode(m : Model, input_ : str, byte_decoder : np.ndarray) -> np.ndarray: + """Compare to the fortran function in tokenizer.f90.""" + # reshape this later after counting tokens + tokens2 : np.ndarray = np.zeros(MaxTokens, dtype=np.int32) + n_tokens : int = 0 + i : int = 0 # fortran counts from 1 + # Python does not have \p for punctuation. + # rex = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""") + # rex.match(input_) + while True: + tmp : str + tmp, i = next_token(input_, i) + t : str + for t in tmp: + c : int = ord(t) + e = m.byte_encoder[c] + tmp2 : bytes = t.encode("utf-8") + tmp3 : int = ord(tmp2) + if tmp == '': + continue + return tokens2 + + +def gpt2_driver(m : Model, input_ : str) -> str: + print(f"Input to the model =\n{input_}") + n_tokens_to_generate : int = NTokensToGenerate + n_seq : int = len(input_) + byte_encoder_max : int = np.max(m.byte_encoder) + byte_decoder : np.ndarray = \ + np.zeros((byte_encoder_max + 1,), dtype=np.int32) + for i, e in enumerate(m.byte_encoder): + byte_decoder[e] = i + # check against fortran: + # print(f'byte_decoder = \n{byte_decoder}') + encoded : np.ndarray = encode(m, input_, byte_decoder) + result = '' + return result + + +def main() -> Model: + + with Timer(text="Restored the model in {:0.4f} seconds."): + m : Model = restore_model() + + input : str = """Alan Turing theorized that computers would one day become very powerful, but even he could not imagine""" + with Timer(text="Ran the model in {:0.6f} seconds."): + result : str = gpt2_driver(m, input) + + return m + + +if __name__ == "__main__": + main() diff --git a/lp-fastgpt/timer.py b/lp-fastgpt/timer.py new file mode 100644 index 0000000..123c3f7 --- /dev/null +++ b/lp-fastgpt/timer.py @@ -0,0 +1,36 @@ +"""See https://realpython.com/python-timer/""" + + +from dataclasses import dataclass, field +import time +from typing import Callable, Optional + + +class TimerError(Exception): + """custom Exception for Timer errors""" + + +@dataclass +class Timer: + """nestable""" + text : str = "Elapsed time: {:0.4f} seconds" + logger : Callable[[str], None] = print + _start_time : Optional[float] = \ + field(default=None, init=False, repr=False) + + def start(self) -> None: + self._start_time = time.perf_counter() + + def stop(self) -> float: + elapsed_time = time.perf_counter() - self._start_time + self._start_time = None + self.logger(self.text.format(elapsed_time)) + return elapsed_time + + # Protocol methods for context manager: + def __enter__(self): + self.start() + return self + + def __exit__(self, *exc_info): + self.stop() diff --git a/main.f90 b/main.f90 index dcb8ff1..d551d2f 100644 --- a/main.f90 +++ b/main.f90 @@ -1,7 +1,7 @@ program gpt2 -use driver, only: gpt2_driver, model_t -implicit none -integer, allocatable :: input(:), output(:) -type(model_t) :: m -call gpt2_driver(input, output, m) + use driver, only : gpt2_driver, model_t + implicit none + integer, allocatable :: input(:), output(:) + type(model_t) :: m + call gpt2_driver(input, output, m) end program diff --git a/tests/test_chat.f90 b/tests/test_chat.f90 index 41173aa..59fced7 100644 --- a/tests/test_chat.f90 +++ b/tests/test_chat.f90 @@ -19,5 +19,5 @@ program test_chat string("What color is snow?"), & string("What color do plants usually have?") & ] -call chat(inputs(:3)) +call chat(.true., inputs(:3)) end program diff --git a/tests/test_more_inputs.f90 b/tests/test_more_inputs.f90 index 7f3a586..80d5f35 100644 --- a/tests/test_more_inputs.f90 +++ b/tests/test_more_inputs.f90 @@ -11,19 +11,19 @@ program test_more_inputs call load_model("model.dat", m) -call gpt2_driver2("Ondřej Čertík was born in ", 13, m, input, output) +call gpt2_driver2("Ondřej Čertík was born in", 13, m, input, output) print * print *, "TESTS:" call test(input, input_ref, "Input") call test(output, output_ref, "Output") -call gpt2_driver2("San Francisco is ", 8, m, input, output) +call gpt2_driver2("San Francisco is", 8, m, input, output) print * print *, "TESTS:" call test(input, [15017, 6033, 318], "Input") call test(output, [257, 1748, 286, 517, 621, 352, 1510, 661], "Output") -call gpt2_driver2("Cars are ", 13, m, input, output) +call gpt2_driver2("Cars are", 13, m, input, output) print * print *, "TESTS:" call test(input, [34, 945, 389], "Input") diff --git a/tokenizer.f90 b/tokenizer.f90 index 6bb564c..8c705dd 100644 --- a/tokenizer.f90 +++ b/tokenizer.f90 @@ -1,245 +1,290 @@ module tokenizer -implicit none + implicit none -type :: string - character(:), allocatable :: s -end type + type :: string + character(:), allocatable :: s + end type contains -function c2s(x) result(y) -character, intent(in) :: x(:) -character(:), allocatable :: y -integer :: i -allocate(character(size(x)) :: y) -do i = 1, size(x) - y(i:i) = x(i) -end do -end function + function c2s(x) result(y) + integer(1), intent(in) :: x(:) + integer(1) :: xx(size(x)) + character(:), allocatable :: y + integer :: i + xx = x + allocate(character(size(x)) :: y) + do i = 1, size(x) + y(i:i) = char(int(xx(i), 4)) + end do + end function -function next_token(input, i) result(y) -! TODO: tokenize exactly according to this regex: -! re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""") -! Right now we are more greedy, but the bpe() tokenizer seems to still return -! exactly the same tokens for most inputs (it is not clear if for all inputs). -character(*), intent(in) :: input -integer, intent(inout) :: i -character(:), allocatable :: y -if (i > len(input)) then - y = "" -else if (input(i:i) == " ") then - y = tokenize_word(input, i) -else if (input(i:i) == "," .or. input(i:i) == ".") then - y = input(i:i) - i = i + 1 -else - y = tokenize_word(input, i) -end if -end function + function next_token(input, i) result(y) + ! TODO: tokenize exactly according to this regex: + ! re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""") + ! Right now we are more greedy, but the bpe() tokenizer seems to still return + ! exactly the same tokens for most inputs (it is not clear if for all inputs). + character(*), intent(in) :: input + integer, intent(inout) :: i + character(:), allocatable :: y + if (i > len(input)) then + y = "" + else if (input(i:i) == " ") then + y = tokenize_word(input, i) + else if (input(i:i) == "," .or. input(i:i) == ".") then + y = input(i:i) + i = i + 1 + else + y = tokenize_word(input, i) + end if + end function -function tokenize_word(input, i) result(y) -character(*), intent(in) :: input -integer, intent(inout) :: i -character(:), allocatable :: y -integer :: i0 -i0 = i -if (input(i:i) == " ") then - i = i + 1 -end if -do - if (i > len(input)) then - y = input(i0:i-1) - exit - end if - if (input(i:i) == " " .or. input(i:i) == "," .or. input(i:i) == ".") then - y = input(i0:i-1) - exit - end if - i = i + 1 -end do -end function + function tokenize_word(input, i) result(y) + character(*), intent(in) :: input + integer, intent(inout) :: i + character(:), allocatable :: y + integer :: i0 + i0 = i + if (input(i:i) == " ") then + i = i + 1 + end if + do + if (i > len(input)) then + y = input(i0:i - 1) + exit + end if + if (input(i:i) == " " .or. input(i:i) == "," .or. input(i:i) == ".") then + y = input(i0:i - 1) + exit + end if + i = i + 1 + end do + end function -function word_idx(word, idx, decoder_txt) result(token) -character(*), intent(in) :: word -integer, intent(in) :: idx(0:) -character, intent(in) :: decoder_txt(:) -integer :: token -integer :: i -! This is O(n) search instead of O(1) lookup in a dictionary, so it is slow -do i = 0, ubound(idx,1)-1 - if (c2s(decoder_txt(idx(i)+1:idx(i+1))) == word) then - token = i - return - end if -end do -token = -1 -end function + function word_idx(word, idx, decoder_txt) result(token) + character(*), intent(in) :: word + integer, intent(in) :: idx(0:) + integer(1), intent(in) :: decoder_txt(:) + integer :: token + integer :: i + ! This is O(n) search instead of O(1) lookup in a dictionary, so it is slow + do i = 0, ubound(idx, 1) - 1 + if (c2s(decoder_txt(idx(i) + 1:idx(i + 1))) == word) then + token = i + return + end if + end do + token = -1 + end function -subroutine codepoint_to_utf8(s, c) -! UTF-32 -> UTF-8 -character(:), allocatable, intent(inout) :: s -integer, intent(in) :: c -integer :: d1, d2 -if (c < 128) then - s = s // achar(c) -else if (c < 2048) then - d1 = ior(ishft(c, -6), 192) - d2 = iand(ior(c, 128), 191) - s = s // achar(d1) // achar(d2) -else - error stop "UTF-32 range not supported" -end if -end subroutine + subroutine codepoint_to_utf8(s, c) + ! UTF-32 -> UTF-8 + character(:), allocatable, intent(inout) :: s + integer, intent(in) :: c + integer :: d1, d2 + if (c < 128) then + s = s // achar(c) + else if (c < 2048) then + d1 = ior(ishft(c, -6), 192) + d2 = iand(ior(c, 128), 191) + s = s // achar(d1) // achar(d2) + else + error stop "UTF-32 range not supported" + end if + end subroutine -function utf8_to_codepoint(s, i) result(c) -! UTF-8 -> UTF-32 -character(*), intent(in) :: s -integer, intent(inout) :: i -integer :: c, d -c = iachar(s(i:i)) -if (c >= 128) then - i = i + 1 - d = iachar(s(i:i)) - c = ior(ishft(iand(c, 31), 6), iand(d, 63)) -end if -if (c >= 2048) then - error stop "UTF-8 range not supported" -end if -end function + function utf8_to_codepoint(s, i) result(c) + ! UTF-8 -> UTF-32 + character(*), intent(in) :: s + integer, intent(inout) :: i + integer :: c, d + c = iachar(s(i:i)) + if (c >= 128) then + i = i + 1 + d = iachar(s(i:i)) + c = ior(ishft(iand(c, 31), 6), iand(d, 63)) + end if + if (c >= 2048) then + error stop "UTF-8 range not supported" + end if + end function -function merge_pair(intokens, idx) result(tokens) -! Merge the pair `idx` -type(string), intent(in) :: intokens(:) -integer, intent(in) :: idx -type(string), allocatable :: tokens(:) -type(string) :: merged_token -merged_token%s = intokens(idx)%s // intokens(idx+1)%s -tokens = [intokens(:idx-1), merged_token, intokens(idx+2:)] -end function + function merge_pair(intokens, idx) result(tokens) + ! Merge the pair `idx` + type(string), intent(in) :: intokens(:) + integer, intent(in) :: idx + type(string), allocatable :: tokens(:) + type(string) :: merged_token + integer :: i + merged_token%s = intokens(idx)%s // intokens(idx + 1)%s + allocate(tokens(size(intokens) - 1)) + do i = 1, idx - 1 + tokens(i) = intokens(i) + end do + tokens(idx) = merged_token + do i = idx + 2, size(intokens) + tokens(i - 1) = intokens(i) + end do + end function -function merge_utf8_pairs(intokens) result(tokens) -! Merge all UTF-8 character pairs -type(string), intent(in) :: intokens(:) -type(string), allocatable :: tokens(:) -integer :: i, j -logical :: one_more_pass -tokens = intokens -one_more_pass = .true. -!print *, "merge_utf8_pairs:", size(tokens) -!print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens)) -j = 1 -do while(one_more_pass) - one_more_pass = .false. - do i = j, size(tokens)-1 - if (len(tokens(i)%s) == 1 .and. iachar(tokens(i)%s(1:1)) >= 128) then - tokens = merge_pair(tokens, i) - one_more_pass = .true. - j = i + 1 -! print *, "pass" - exit - end if - end do -end do -!print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens)) -end function + function merge_utf8_pairs(intokens) result(tokens) + ! Merge all UTF-8 character pairs + type(string), intent(in) :: intokens(:) + type(string), allocatable :: tokens(:) + integer :: i, j, ic + logical :: one_more_pass + allocate(tokens(size(intokens))) + tokens = intokens + one_more_pass = .true. + !print *, "merge_utf8_pairs:", size(tokens) + !print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens)) + j = 1 + do while(one_more_pass) + one_more_pass = .false. + do i = j, size(tokens) - 1 + if (len(tokens(i)%s) == 1) then + ic = iachar(tokens(i)%s(1:1)) + if (ic < 0) ic = ic + 256 + if (ic >= 128) then + tokens = merge_pair(tokens, i) + one_more_pass = .true. + j = i + 1 + exit + end if + end if + end do + end do + !print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens)) + end function -function bpe(token, vocab_idx, vocab_txt) result(tokens) -! Takes a token as a string, and returns bpe tokens as an array of strings -character(*), intent(in) :: token -integer, intent(in) :: vocab_idx(0:) -character, intent(in) :: vocab_txt(:) -type(string), allocatable :: tokens(:) -integer, allocatable :: pair_scores(:) -integer :: not_found, merge_pair_idx -integer :: i -not_found = size(vocab_idx) + 10 -allocate(tokens(len(token))) -do i = 1, len(token) - tokens(i)%s = token(i:i) -end do -tokens = merge_utf8_pairs(tokens) -do - !print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens)) - if (size(tokens) == 1) then - ! The token pairs were either all merged into one word, or the input - ! token was a one character word, either way we are done: - exit - end if - allocate(pair_scores(size(tokens)-1)) - ! Loop over pairs - do i = 1, size(tokens)-1 - pair_scores(i) = word_idx(tokens(i)%s // " " // tokens(i+1)%s, vocab_idx, vocab_txt) - if (pair_scores(i) == -1) pair_scores(i) = not_found - end do - merge_pair_idx = minloc(pair_scores, 1) - if (pair_scores(merge_pair_idx) == not_found) then - ! No token pair can be merged, so we are done: - exit - end if - !print *, pair_scores - !print *, merge_pair_idx, pair_scores(merge_pair_idx) - tokens = merge_pair(tokens, merge_pair_idx) - deallocate(pair_scores) -end do -!print *, "final tokens = ", (tokens(i)%s // " ", i=1,size(tokens)) -end function + function bpe(token, vocab_idx, vocab_txt) result(tokens) + ! Takes a token as a string, and returns bpe tokens as an array of strings + character(*), intent(in) :: token + integer, intent(in) :: vocab_idx(0:) + integer(1), intent(in) :: vocab_txt(:) + type(string), allocatable :: tokens(:) + integer, allocatable :: pair_scores(:) + integer :: not_found, merge_pair_idx + integer :: i + not_found = size(vocab_idx) + 10 + allocate(tokens(len(token))) + do i = 1, len(token) + tokens(i)%s = token(i:i) + print "(a)", "i" + print *, i + print "(a)", "tokens(i)%s" + print *, tokens(i)%s + print "(a)", "token(i:i)" + print *, token(i:i) + end do + tokens = merge_utf8_pairs(tokens) + do + !print *, "tokens = ", (tokens(i)%s // " ", i=1,size(tokens)) + if (size(tokens) == 1) then + ! The token pairs were either all merged into one word, or the input + ! token was a one character word, either way we are done: + exit + end if + allocate(pair_scores(size(tokens) - 1)) + ! Loop over pairs + do i = 1, size(tokens) - 1 + pair_scores(i) = word_idx(tokens(i)%s // " " // tokens(i + 1)%s, vocab_idx, vocab_txt) + if (pair_scores(i) == -1) pair_scores(i) = not_found + end do + merge_pair_idx = minloc(pair_scores, 1) + if (pair_scores(merge_pair_idx) == not_found) then + ! No token pair can be merged, so we are done: + exit + end if + !print *, pair_scores + !print *, merge_pair_idx, pair_scores(merge_pair_idx) + tokens = merge_pair(tokens, merge_pair_idx) + deallocate(pair_scores) + end do + !print *, "final tokens = ", (tokens(i)%s // " ", i=1,size(tokens)) + end function -function encode(input, idx, decoder_txt, vocab_idx, vocab_txt, byte_encoder) & - result(tokens) -character(*), intent(in) :: input -integer, intent(in) :: idx(0:), vocab_idx(0:), byte_encoder(0:) -character, intent(in) :: decoder_txt(:), vocab_txt(:) -integer, allocatable :: tokens(:) -character(:), allocatable :: tmp, tmp2 -type(string), allocatable :: bpe_tokens(:) -integer :: i, j, c -i = 1 -allocate(tokens(0)) -do - tmp = next_token(input, i) - if (tmp == "") exit - tmp2 = "" - do j = 1, len(tmp) - c = iachar(tmp(j:j)) - c = byte_encoder(c) - ! c is UTF-32 (4 bytes), but only the range [0, 324] is used - ! Encode c from UTF-32 to UTF-8. Due to the limited range - ! either one or two bytes of UTF-8 are appended to tmp2: - call codepoint_to_utf8(tmp2, c) - end do - bpe_tokens = bpe(tmp2, vocab_idx, vocab_txt) - do j = 1, size(bpe_tokens) - tokens = [tokens, word_idx(bpe_tokens(j)%s, idx, decoder_txt)] - end do - deallocate(tmp2) -end do -end function + function encode(input, idx, decoder_txt, vocab_idx, vocab_txt, byte_encoder) & + result(tokens2) + character(*), intent(in) :: input + integer, intent(in) :: idx(0:), vocab_idx(0:), byte_encoder(0:) + integer(1), intent(in) :: decoder_txt(:), vocab_txt(:) + integer, parameter :: max_tokens = 2048 + integer :: tokens(max_tokens) + integer, allocatable :: tokens2(:) + character(:), allocatable :: tmp, tmp2 + type(string), allocatable :: bpe_tokens(:) + integer :: i, j, c, n_tokens + n_tokens = 0 + i = 1 + do + tmp = next_token(input, i) + if (tmp == "") exit + tmp2 = "" + do j = 1, len(tmp) + c = iachar(tmp(j:j)) + c = byte_encoder(c) + ! c is UTF-32 (4 bytes), but only the range [0, 324] is used + ! Encode c from UTF-32 to UTF-8. Due to the limited range + ! either one or two bytes of UTF-8 are appended to tmp2: + call codepoint_to_utf8(tmp2, c) + end do + bpe_tokens = bpe(tmp2, vocab_idx, vocab_txt) + do j = 1, size(bpe_tokens) + n_tokens = n_tokens + 1 + if (n_tokens > max_tokens) error stop "exceeded max_tokens" + tokens(n_tokens) = word_idx(bpe_tokens(j)%s, idx, decoder_txt) + end do + deallocate(tmp2) + end do + allocate(tokens2(n_tokens)) + do i = 1, n_tokens + tokens2(i) = tokens(i) + end do + end function -function decode(tokens, idx, decoder_txt, byte_decoder) result(output) -integer, intent(in) :: tokens(:), idx(0:), byte_decoder(0:) -character, intent(in) :: decoder_txt(:) -character(:), allocatable :: output -character(:), allocatable :: output2, tmp -integer :: i, c -allocate(character(0) :: output2) ! Fix GFortran warning -output2 = "" -do i = 1, size(tokens) - if (tokens(i) < 0) error stop "tokens(i) < 0" - output2 = output2 // c2s(decoder_txt(idx(tokens(i))+1:idx(tokens(i)+1))) -end do -i = 1 -output = "" -do - ! Decode UTF-8 (one or more bytes) to UTF-32 code point (always 4 bytes), - ! However for GPT-2 it seems only range 0-323 is used from UTF-32. - c = utf8_to_codepoint(output2, i) - ! [0,324] -> [0,255] - if (c < 0 .or. c > ubound(byte_decoder,1)) error stop "Codepoint out of range for byte decoder" - tmp = achar(byte_decoder(c)) - output = output // tmp - if (i == len(output2)) exit - i = i + 1 -end do -end function + function decode(tokens, idx, decoder_txt, byte_decoder) result(output) + integer, intent(in) :: tokens(:), idx(0:), byte_decoder(0:) + integer(1), intent(in) :: decoder_txt(:) + character(:), allocatable :: output + character(:), allocatable :: output2, tmp + integer, parameter :: max_len = 4096 + integer(1) :: output3(max_len) + integer(1), allocatable :: output4(:) + integer :: i, j, c, pos + pos = 0 + do i = 1, size(tokens) + if (tokens(i) < 0) error stop "tokens(i) < 0" + ! output2 = output2 // (decoder_txt(idx(tokens(i))+1:idx(tokens(i)+1))) + do j = idx(tokens(i)) + 1, idx(tokens(i) + 1) + pos = pos + 1 + output3(pos) = decoder_txt(j) + end do + end do + allocate(character(0) :: output2) ! Fix GFortran warning + allocate(output4(pos)) + do j = 1, pos + output4(j) = output3(j) + end do + output2 = c2s(output4) + i = 1 + output = "" + do + ! Decode UTF-8 (one or more bytes) to UTF-32 code point (always 4 bytes), + ! However for GPT-2 it seems only range 0-323 is used from UTF-32. + c = utf8_to_codepoint(output2, i) + ! [0,324] -> [0,255] + if (c < 0 .or. c > ubound(byte_decoder, 1)) then + print *, "Codepoint out of range for byte decoder:", c, ubound(byte_decoder, 1) + ! We skip this, to allow LFortran to run + error stop "Codepoint out of range for byte decoder" + else + tmp = achar(byte_decoder(c)) + output = output // tmp + end if + if (i == len(output2)) exit + i = i + 1 + end do + end function end module