mirror of
https://github.com/matrix-construct/construct
synced 2024-11-12 13:01:07 +01:00
1454 lines
38 KiB
Common Lisp
1454 lines
38 KiB
Common Lisp
// Matrix Construct
|
|
//
|
|
// Copyright (C) Matrix Construct Developers, Authors & Contributors
|
|
// Copyright (C) 2016-2021 Jason Volk <jason@zemos.net>
|
|
//
|
|
// Permission to use, copy, modify, and/or distribute this software for any
|
|
// purpose with or without fee is hereby granted, provided that the above
|
|
// copyright notice and this permission notice is present in all copies. The
|
|
// full license for this software is available in the LICENSE file.
|
|
|
|
#pragma clang fp exceptions(ignore)
|
|
#pragma clang fp reassociate(on)
|
|
#pragma clang fp contract(fast)
|
|
|
|
#include <ircd/config.h>
|
|
#include <ircd/portable.h>
|
|
#include <clc/clc.h>
|
|
#include <ircd/simt/simt.h>
|
|
#include <ircd/gpt/vector.h>
|
|
#include <ircd/gpt/opts.h>
|
|
#include <ircd/gpt/ctrl.h>
|
|
#include <ircd/gpt/gpu.h>
|
|
|
|
//
|
|
// head
|
|
//
|
|
|
|
__kernel void
|
|
__attribute__((visibility("protected")))
|
|
ircd_gpt_alloc(__global const void *const restrict model,
|
|
__global void *const restrict master,
|
|
__constant const void *const opts,
|
|
__global void *const restrict ctrl,
|
|
__global void *const restrict frame0,
|
|
__global void *const restrict frame1,
|
|
__global void *const restrict frame2,
|
|
__global void *const restrict frame3,
|
|
__global void *const restrict frame4,
|
|
__global void *const restrict frame5,
|
|
__global void *const restrict frame6,
|
|
__global void *const restrict frame7)
|
|
{
|
|
}
|
|
|
|
__kernel void
|
|
__attribute__((visibility("protected")))
|
|
ircd_gpt_enter(__global const void *const restrict model,
|
|
__global void *const restrict state,
|
|
__global void *const restrict master,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global struct ircd_gpt_ctrl *const restrict ctrl)
|
|
{
|
|
const ushort
|
|
gi = get_global_id(0),
|
|
li = get_local_id(0),
|
|
ln = get_local_size(0),
|
|
cycle = ctrl->clk.cycle;
|
|
|
|
if(li == 0)
|
|
ctrl->prof.entered = ircd_simt_cycles();
|
|
}
|
|
|
|
__kernel void
|
|
ircd_gpt_lm_embed(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global ircd_gpt_vectorv *const restrict accum,
|
|
__global const ircd_gpt_vectorv *const restrict pos,
|
|
__global const ircd_gpt_vectorv *const restrict vocab)
|
|
{
|
|
const ushort
|
|
li = get_local_id(0),
|
|
ln = get_local_size(0);
|
|
|
|
const uint
|
|
wo = get_global_offset(0);
|
|
|
|
assume(ln == 192);
|
|
assume(wo % ln == 0);
|
|
|
|
const ushort
|
|
wi = wo / ln + get_group_id(0);
|
|
|
|
_ircd_gpt_lm_embed(ctrl, opts, accum, pos, vocab, wi, wi, li);
|
|
}
|
|
|
|
static void
|
|
_ircd_gpt_lm_embed(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global ircd_gpt_vectorv *const restrict accum,
|
|
__global const ircd_gpt_vectorv *const restrict pos,
|
|
__global const ircd_gpt_vectorv *const restrict vocab,
|
|
const ushort out_idx,
|
|
const ushort tok_idx,
|
|
const ushort elem_idx)
|
|
{
|
|
const ushort
|
|
token = ctrl->token[tok_idx];
|
|
|
|
const float4
|
|
wpe = pos[tok_idx].elem[elem_idx],
|
|
wte = vocab[token].elem[elem_idx],
|
|
res = wte + wpe;
|
|
|
|
accum[out_idx].elem[elem_idx] = res;
|
|
}
|
|
|
|
//
|
|
// Frontside
|
|
//
|
|
|
|
void
|
|
ircd_gpt_ffnn_fcon_tmul(__constant const struct ircd_gpt_opts *const opts,
|
|
__local ircd_gpt_ffnn_aperaturev *const restrict out,
|
|
__local const ircd_gpt_vectorv *const restrict in,
|
|
__global const ircd_gpt_ffnn_aperaturev *const restrict bias,
|
|
__global const ircd_gpt_ffnn_aperaturev *const restrict weight,
|
|
const uint li)
|
|
{
|
|
const uint
|
|
lanes = 4,
|
|
segs = ircd_gpt_ffnn_segs,
|
|
height = ircd_gpt_vector_elems / lanes;
|
|
|
|
assume(height > 0);
|
|
assume(height % lanes == 0);
|
|
|
|
for(uint x = 0; x < segs; ++x)
|
|
out->proj[x][li] = bias->proj[x][li];
|
|
|
|
for(uint y = 0; y < height; ++y)
|
|
for(uint k = 0; k < lanes; ++k)
|
|
for(uint x = 0; x < segs; ++x)
|
|
{
|
|
const uint
|
|
row = y * lanes + k;
|
|
|
|
out->proj[x][li] += in->elem[y][k] * weight[row].proj[x][li];
|
|
}
|
|
}
|
|
|
|
void
|
|
ircd_gpt_ffnn_fcon(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local ircd_gpt_ffnn_aperaturev *const restrict out,
|
|
__local const ircd_gpt_vectorv *const restrict in,
|
|
__global const ircd_gpt_ffnn_aperaturev *const restrict bias,
|
|
__global const ircd_gpt_ffnn_aperaturev *const restrict weight,
|
|
const uint ln,
|
|
const uint li)
|
|
{
|
|
const uint
|
|
segs = ircd_gpt_ffnn_segs;
|
|
|
|
// Fully connected
|
|
ircd_gpt_ffnn_fcon_tmul
|
|
(
|
|
opts,
|
|
out,
|
|
in,
|
|
bias,
|
|
weight,
|
|
li
|
|
);
|
|
|
|
for(uint i = 0; i < segs; ++i)
|
|
ircd_gpt_ffnn_gelu(out, out, i * ln + li);
|
|
}
|
|
|
|
void
|
|
ircd_gpt_ffnn_proj_tmul(__constant const struct ircd_gpt_opts *const opts,
|
|
__local ircd_gpt_vectorv *const restrict out,
|
|
__local const ircd_gpt_ffnn_aperaturev *const restrict in,
|
|
__global const ircd_gpt_vectorv *const restrict bias,
|
|
__global const ircd_gpt_vectorv *const restrict weight,
|
|
const uint li)
|
|
{
|
|
const uint
|
|
lanes = 4,
|
|
height = ircd_gpt_ffnn_fcon_elems / lanes;
|
|
|
|
assume(height > 0);
|
|
assume(height % lanes == 0);
|
|
|
|
out->elem[li] = bias->elem[li];
|
|
|
|
for(uint y = 0; y < height; ++y)
|
|
for(uint k = 0; k < lanes; ++k)
|
|
{
|
|
const uint
|
|
row = y * lanes + k;
|
|
|
|
out->elem[li] += in->fcon[y][k] * weight[row].elem[li];
|
|
}
|
|
}
|
|
|
|
void
|
|
ircd_gpt_ffnn(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local ircd_gpt_vectorv *const restrict token,
|
|
__local ircd_gpt_ffnn_aperaturev *const restrict buf,
|
|
__local ircd_gpt_vectorv *const restrict tmp,
|
|
__global const ircd_gpt_vectorv *const restrict norm_bias,
|
|
__global const ircd_gpt_vectorv *const restrict norm_weight,
|
|
__global const ircd_gpt_ffnn_aperaturev *const restrict fcon_bias,
|
|
__global const ircd_gpt_ffnn_aperaturev *const restrict fcon_weight,
|
|
__global const ircd_gpt_vectorv *const restrict proj_bias,
|
|
__global const ircd_gpt_vectorv *const restrict proj_weight,
|
|
const uint ln,
|
|
const uint li)
|
|
{
|
|
// Layer re-normalization
|
|
ircd_gpt_norm(token, token, tmp, norm_bias, norm_weight, ln, li);
|
|
|
|
// ln's writes are still pending but fcon reads results across threads.
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
// Fully connected
|
|
ircd_gpt_ffnn_fcon
|
|
(
|
|
ctrl,
|
|
opts,
|
|
buf,
|
|
token,
|
|
fcon_bias,
|
|
fcon_weight,
|
|
ln,
|
|
li
|
|
);
|
|
|
|
// fcon's writes are still pending but proj reads results across threads.
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
// Projection
|
|
ircd_gpt_ffnn_proj_tmul
|
|
(
|
|
opts,
|
|
token,
|
|
buf,
|
|
proj_bias,
|
|
proj_weight,
|
|
li
|
|
);
|
|
}
|
|
|
|
static void
|
|
ircd_gpt_attn_self_samax(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local float self[restrict][12],
|
|
const uint ln,
|
|
const uint li,
|
|
const uint wn,
|
|
const uint wi)
|
|
{
|
|
struct ircd_math_samax samax =
|
|
{
|
|
.mu = -10000.0f,
|
|
.sum = 0.0f,
|
|
};
|
|
|
|
__attribute__((opencl_unroll_hint))
|
|
for(uint i = 0; i < wn; ++i)
|
|
samax.mu = max(samax.mu, self[i][li]);
|
|
|
|
__attribute__((opencl_unroll_hint))
|
|
for(uint i = 0; i < wn; ++i)
|
|
self[i][li] -= samax.mu;
|
|
|
|
for(uint i = 0; i < wn; ++i)
|
|
self[i][li] = native_exp(self[i][li]);
|
|
|
|
__attribute__((opencl_unroll_hint))
|
|
for(uint i = 0; i < wn; ++i)
|
|
samax.sum += self[i][li];
|
|
|
|
samax.sum += FLT_EPSILON;
|
|
samax.lambda = 1.0f / samax.sum;
|
|
|
|
__attribute__((opencl_unroll_hint))
|
|
for(uint i = 0; i < wn; ++i)
|
|
self[i][li] *= samax.lambda;
|
|
}
|
|
|
|
static void
|
|
ircd_gpt_attn_self_keys(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local float self[restrict][ircd_gpt_attn_rank],
|
|
__global const ircd_gpt_attn_qkvv *const restrict token,
|
|
const uint ln,
|
|
const uint li,
|
|
const uint wi,
|
|
const uint kn,
|
|
const uint i)
|
|
{
|
|
assume(i < wi);
|
|
|
|
self[i][li] = 0.0f;
|
|
|
|
__attribute__((opencl_unroll_hint))
|
|
for(uint k = 0; k < kn; ++k)
|
|
{
|
|
float4
|
|
qry = token[wi].qry.attn[li][k],
|
|
key = token[i].key.attn[li][k],
|
|
res = qry * key;
|
|
|
|
self[i][li] += ircd_simt_hadd_f4(res);
|
|
}
|
|
|
|
self[i][li] /= 8.0f;
|
|
}
|
|
|
|
static void
|
|
ircd_gpt_attn_self_vals(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local ircd_gpt_vectorv *const restrict out,
|
|
__local const float self[restrict][ircd_gpt_attn_rank],
|
|
__global const ircd_gpt_attn_qkvv *const restrict token,
|
|
const uint li,
|
|
const uint wi,
|
|
const uint ki,
|
|
const uint ti)
|
|
{
|
|
out->attn[ti][ki] = 0.0f;
|
|
|
|
__attribute__((opencl_unroll_hint))
|
|
for(uint i = 0; i < wi; ++i)
|
|
{
|
|
const float4
|
|
val = token[i].val.attn[ti][ki],
|
|
attn = self[i][ti],
|
|
res = attn * val;
|
|
|
|
out->attn[ti][ki] += res;
|
|
}
|
|
}
|
|
|
|
static void
|
|
ircd_gpt_attn_self(__global struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local ircd_gpt_vectorv *const restrict out,
|
|
__local float self[restrict][ircd_gpt_attn_rank],
|
|
__global float attns[restrict][ircd_gpt_attn_rank],
|
|
__global const ircd_gpt_attn_qkvv *const restrict token,
|
|
const uint ln,
|
|
const uint li,
|
|
const uint wi)
|
|
{
|
|
//assume(opts->attn_rank == sizeof(self[0]) / sizeof(float));
|
|
assume(opts->attn_rank == ircd_gpt_attn_rank);
|
|
assume(ctrl->count < ircd_gpt_context_tokens);
|
|
assume(ctrl->tokens <= ircd_gpt_context_tokens);
|
|
assume(ctrl->tokens > wi);
|
|
assume(ctrl->tokens > 0);
|
|
|
|
const uint
|
|
wn = ctrl->tokens,
|
|
kn = ln / opts->attn_rank,
|
|
ki = li / opts->attn_rank,
|
|
ti = li % opts->attn_rank;
|
|
|
|
// Low-rank mask
|
|
if(li < opts->attn_rank)
|
|
{
|
|
// Left attention
|
|
uint i;
|
|
for(i = 0; i < wi; ++i)
|
|
ircd_gpt_attn_self_keys(ctrl, opts, self, token, ln, li, wi, kn, i);
|
|
|
|
// Future mask
|
|
__attribute__((opencl_unroll_hint))
|
|
while(i < wn)
|
|
self[i++][li] = -10000.0f;
|
|
|
|
// Three-piece softmax
|
|
ircd_gpt_attn_self_samax(ctrl, opts, self, ln, li, wn, wi);
|
|
}
|
|
|
|
// Propagate to full width for value dot prod.
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
ircd_gpt_attn_self_vals(ctrl, opts, out, self, token, li, wi, ki, ti);
|
|
|
|
// Save softmax results for later analysis/observation.
|
|
if(li < opts->attn_rank)
|
|
{
|
|
__attribute__((opencl_unroll_hint))
|
|
for(uint i = 0; i < wn; ++i)
|
|
attns[i][li] = self[i][li];
|
|
}
|
|
}
|
|
|
|
static void
|
|
ircd_gpt_attn_proj_tmul(__constant const struct ircd_gpt_opts *const opts,
|
|
__local ircd_gpt_vectorv *const restrict out,
|
|
__local const ircd_gpt_vectorv *const restrict in,
|
|
__global const ircd_gpt_vectorv *const restrict bias,
|
|
__global const ircd_gpt_vectorv *const restrict weight,
|
|
const uint li)
|
|
{
|
|
const uint
|
|
lanes = 4,
|
|
height = ircd_gpt_vector_elems / 4;
|
|
|
|
assume(height > 0);
|
|
assume(height % lanes == 0);
|
|
|
|
out->elem[li] = bias->elem[li];
|
|
|
|
for(uint y = 0; y < height; ++y)
|
|
for(uint k = 0; k < lanes; ++k)
|
|
{
|
|
const uint
|
|
row = y * lanes + k;
|
|
|
|
const float4
|
|
a = in->elem[y][k],
|
|
b = weight[row].elem[li];
|
|
|
|
out->elem[li] += a * b;
|
|
}
|
|
}
|
|
|
|
__kernel void
|
|
__attribute__((visibility("protected")))
|
|
ircd_gpt_coil(__global struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__private const uint layer,
|
|
__global ircd_gpt_vectorv *const restrict accum,
|
|
__global float attns[restrict][ircd_gpt_attn_rank],
|
|
__global const ircd_gpt_attn_qkvv *const restrict state,
|
|
__global const ircd_gpt_vectorv *const restrict attn_proj_bias,
|
|
__global const ircd_gpt_vectorv *const restrict attn_proj_weight,
|
|
__global const ircd_gpt_vectorv *const restrict ffnn_norm_bias,
|
|
__global const ircd_gpt_vectorv *const restrict ffnn_norm_weight,
|
|
__global const ircd_gpt_ffnn_aperaturev *const restrict ffnn_fcon_bias,
|
|
__global const ircd_gpt_ffnn_aperaturev *const restrict ffnn_fcon_weight,
|
|
__global const ircd_gpt_vectorv *const restrict ffnn_proj_bias,
|
|
__global const ircd_gpt_vectorv *const restrict ffnn_proj_weight)
|
|
{
|
|
const uint
|
|
li = get_local_id(0),
|
|
ln = get_local_size(0),
|
|
wo = get_global_offset(0),
|
|
wi = wo / ln + get_group_id(0);
|
|
|
|
assume(ln == 192);
|
|
assume(wo % ln == 0);
|
|
|
|
__local union
|
|
{
|
|
float
|
|
attn_self[ircd_gpt_context_tokens][ircd_gpt_attn_rank];
|
|
|
|
ircd_gpt_ffnn_aperaturev
|
|
ffnn_fcon[2];
|
|
|
|
ircd_gpt_vectorv
|
|
vector[8];
|
|
}
|
|
buf;
|
|
|
|
__local ircd_gpt_vectorv
|
|
buf0, buf1,
|
|
*const restrict attn_self = &buf1,
|
|
*const restrict token = &buf0,
|
|
*const restrict tmp = &buf1;
|
|
|
|
// Self-attention backend; this computes the self-attention result now
|
|
// that keys and values are globally visible across tokens.
|
|
ircd_gpt_attn_self
|
|
(
|
|
ctrl,
|
|
opts,
|
|
attn_self,
|
|
buf.attn_self,
|
|
attns,
|
|
state,
|
|
ln,
|
|
li,
|
|
wi
|
|
);
|
|
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
// Project result of self-attention.
|
|
ircd_gpt_attn_proj_tmul
|
|
(
|
|
opts,
|
|
token,
|
|
attn_self,
|
|
attn_proj_bias,
|
|
attn_proj_weight,
|
|
li
|
|
);
|
|
|
|
// Frontend accumulation
|
|
{
|
|
const float4
|
|
attn = token->elem[li],
|
|
resid = accum[wi].elem[li],
|
|
result = resid + attn;
|
|
|
|
token->elem[li] = result;
|
|
accum[wi].elem[li] = result;
|
|
}
|
|
|
|
barrier(CLK_GLOBAL_MEM_FENCE);
|
|
|
|
// Backend mlp
|
|
ircd_gpt_ffnn
|
|
(
|
|
ctrl,
|
|
opts,
|
|
token,
|
|
buf.ffnn_fcon,
|
|
tmp,
|
|
ffnn_norm_bias,
|
|
ffnn_norm_weight,
|
|
ffnn_fcon_bias,
|
|
ffnn_fcon_weight,
|
|
ffnn_proj_bias,
|
|
ffnn_proj_weight,
|
|
ln,
|
|
li
|
|
);
|
|
|
|
// Backend accumulation
|
|
{
|
|
const float4
|
|
ffnn = token->elem[li],
|
|
resid = accum[wi].elem[li],
|
|
result = resid + ffnn;
|
|
|
|
accum[wi].elem[li] = result;
|
|
}
|
|
}
|
|
|
|
static void
|
|
ircd_gpt_attn_fcon_tmul(__constant const struct ircd_gpt_opts *const opts,
|
|
__local ircd_gpt_attn_aperaturev *const restrict out,
|
|
__local const ircd_gpt_vectorv *const restrict in,
|
|
__global const ircd_gpt_attn_aperaturev *const restrict bias,
|
|
__global const ircd_gpt_attn_aperaturev *const restrict weight,
|
|
const uint ln,
|
|
const uint li)
|
|
{
|
|
const uint
|
|
lanes = 4,
|
|
segs = ircd_gpt_attn_segs,
|
|
height = ircd_gpt_vector_elems / lanes;
|
|
|
|
assume(height > 0);
|
|
assume(height % segs == 0);
|
|
assume(height % lanes == 0);
|
|
|
|
for(uint x = 0; x < segs; ++x)
|
|
out->proj[x][li] = bias->proj[x][li];
|
|
|
|
for(uint y = 0; y < height; ++y)
|
|
for(uint k = 0; k < lanes; ++k)
|
|
for(uint x = 0; x < segs; ++x)
|
|
{
|
|
const uint
|
|
row = y * lanes + k;
|
|
|
|
const float4
|
|
a = in->elem[y][k],
|
|
b = weight[row].proj[x][li];
|
|
|
|
out->proj[x][li] += a * b;
|
|
}
|
|
}
|
|
|
|
__kernel void
|
|
ircd_gpt_attn_fcon(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__private const uint layer,
|
|
__global ircd_gpt_attn_aperaturev *const restrict state,
|
|
__global const ircd_gpt_vectorv *const restrict accum,
|
|
__global const ircd_gpt_vectorv *const restrict norm_bias,
|
|
__global const ircd_gpt_vectorv *const restrict norm_weight,
|
|
__global const ircd_gpt_attn_aperaturev *const restrict fcon_bias,
|
|
__global const ircd_gpt_attn_aperaturev *const restrict fcon_weight)
|
|
{
|
|
const uint
|
|
li = get_local_id(0),
|
|
ln = get_local_size(0),
|
|
wo = get_global_offset(0),
|
|
wi = wo / ln + get_group_id(0),
|
|
segs = ircd_gpt_attn_segs;
|
|
|
|
assume(ln == 192);
|
|
assume(wo % ln == 0);
|
|
|
|
__local ircd_gpt_attn_aperaturev
|
|
attn;
|
|
|
|
__local ircd_gpt_vectorv
|
|
token, *const restrict tmp = attn.vector;
|
|
|
|
token.elem[li] = accum[wi].elem[li];
|
|
|
|
// Layer re-normalization
|
|
ircd_gpt_norm(&token, &token, tmp, norm_bias, norm_weight, ln, li);
|
|
|
|
// Ln's writes are still pending; fcon requires results across threads.
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
// Fully connected
|
|
ircd_gpt_attn_fcon_tmul
|
|
(
|
|
opts,
|
|
&attn,
|
|
&token,
|
|
fcon_bias,
|
|
fcon_weight,
|
|
ln,
|
|
li
|
|
);
|
|
|
|
// Export queries, keys, and values.
|
|
for(uint x = 0; x < segs; ++x)
|
|
state[wi].proj[x][li] = attn.proj[x][li];
|
|
}
|
|
|
|
__kernel void
|
|
ircd_gpt_lm_norm(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global ircd_gpt_vectorv *const restrict accum,
|
|
__global const ircd_gpt_vectorv *const restrict norm_bias,
|
|
__global const ircd_gpt_vectorv *const restrict norm_weight)
|
|
{
|
|
const uint
|
|
li = get_local_id(0),
|
|
ln = get_local_size(0),
|
|
wo = get_global_offset(0),
|
|
wi = wo / ln + get_group_id(0);
|
|
|
|
assume(ln == 192);
|
|
assume(wo % ln == 0);
|
|
|
|
__local ircd_gpt_vectorv
|
|
tmp, token;
|
|
|
|
token.elem[li] = accum[wi].elem[li];
|
|
|
|
// Final re-normalization
|
|
ircd_gpt_norm(&token, &token, &tmp, norm_bias, norm_weight, ln, li);
|
|
|
|
accum[wi].elem[li] = token.elem[li];
|
|
}
|
|
|
|
__kernel void
|
|
ircd_gpt_lm_logit(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global float *const restrict logit,
|
|
__global const ircd_gpt_vectorv *const restrict accum,
|
|
__global const ircd_gpt_vectorv *const restrict pos,
|
|
__global const ircd_gpt_vectorv *const restrict vocab)
|
|
{
|
|
const uint
|
|
gi = get_global_id(0),
|
|
wi = ctrl->count - 1;
|
|
|
|
assume(opts->embed_width == 192);
|
|
assume(opts->logits <= 65536);
|
|
|
|
if(gi >= opts->logits)
|
|
{
|
|
logit[gi] = -10000.0f;
|
|
return;
|
|
}
|
|
|
|
float acc = 0.0f;
|
|
for(uint j = 0; j < opts->embed_width; ++j)
|
|
{
|
|
const float4
|
|
token = vocab[gi].elem[j],
|
|
in = accum[wi].elem[j],
|
|
wpe = pos[wi].elem[j],
|
|
res = in * token + wpe;
|
|
|
|
acc += ircd_simt_hadd_f4(res);
|
|
}
|
|
|
|
logit[gi] = acc;
|
|
}
|
|
|
|
__kernel void
|
|
__attribute__((reqd_work_group_size(256, 1, 1)))
|
|
ircd_gpt_lm_logsm(__global struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global float logit[restrict 65536])
|
|
{
|
|
const uint
|
|
li = get_local_id(0),
|
|
ln = get_local_size(0),
|
|
//wo = get_global_offset(0),
|
|
//wi = wo / ln + get_group_id(0),
|
|
wn = 50432,
|
|
tn = wn / ln,
|
|
start = tn * li,
|
|
stop = min(start + tn, opts->logits);
|
|
|
|
__local float
|
|
mu[256], sum[256], lambda[256];
|
|
|
|
__local struct ircd_math_samax
|
|
samax;
|
|
|
|
assume(ln == 256);
|
|
|
|
mu[li] = -10000.0f;
|
|
__attribute__((opencl_unroll_hint))
|
|
for(uint ti = start; ti < stop; ++ti)
|
|
mu[li] = max(mu[li], logit[ti]);
|
|
|
|
ircd_simt_reduce_max_flldr(mu, ln, li);
|
|
|
|
if(li == 0)
|
|
samax.mu = mu[li];
|
|
|
|
ircd_simt_broadcast_flldr(mu, ln, li);
|
|
|
|
sum[li] = 0.0f;
|
|
for(uint ti = start; ti < stop; ++ti)
|
|
{
|
|
const float
|
|
sub = logit[ti] - mu[li],
|
|
res = native_exp(sub);
|
|
|
|
sum[li] += res;
|
|
}
|
|
|
|
ircd_simt_reduce_add_flldr(sum, ln, li);
|
|
|
|
if(li == 0)
|
|
sum[li] += FLT_EPSILON,
|
|
samax.sum = sum[li],
|
|
samax.lambda = lambda[li] = 1.0f / sum[li];
|
|
|
|
ircd_simt_broadcast_flldr(lambda, ln, li);
|
|
|
|
for(uint ti = start; ti < stop; ++ti)
|
|
{
|
|
const float
|
|
sub = logit[ti] - mu[li],
|
|
res = lambda[li] * native_exp(sub);
|
|
|
|
logit[ti] = res;
|
|
}
|
|
}
|
|
|
|
void
|
|
ircd_gpt_lm_result_top(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local const ushort *const restrict idx,
|
|
__global const float *const restrict logsm,
|
|
const uint i)
|
|
{
|
|
const ushort
|
|
token = idx[i];
|
|
|
|
const float
|
|
samax = logsm[token];
|
|
|
|
ctrl->top[i].token = token;
|
|
ctrl->top[i].samax = samax;
|
|
}
|
|
|
|
void
|
|
ircd_gpt_lm_result_label_mean(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local struct ircd_math_mean *const mean,
|
|
const float last)
|
|
{
|
|
const uint
|
|
div = mean->div + 1,
|
|
sum_sel = mean->div % 4;
|
|
|
|
const float
|
|
sum = mean->sum[0] + mean->sum[1] + mean->sum[2] + mean->sum[3] + last,
|
|
res = sum / div;
|
|
|
|
mean->sum[sum_sel] += last;
|
|
mean->div = div;
|
|
mean->last = last;
|
|
mean->mean = res;
|
|
}
|
|
|
|
void
|
|
ircd_gpt_lm_result_label(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local struct ircd_gpt_ctrl_label *const label,
|
|
__global const float *const restrict logsm)
|
|
{
|
|
const ushort
|
|
token = label->logit.token;
|
|
|
|
const float
|
|
samax = logsm[token],
|
|
loss = 0.0f - native_log(samax),
|
|
ppl = (1.0f - samax) * native_log2(opts->logits);
|
|
|
|
label->logit.samax = samax;
|
|
ircd_gpt_lm_result_label_mean(ctrl, opts, &label->loss, loss);
|
|
ircd_gpt_lm_result_label_mean(ctrl, opts, &label->ppl, ppl);
|
|
}
|
|
|
|
ushort
|
|
ircd_gpt_lm_result_select(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local const ushort *const restrict idx,
|
|
__global const float *const restrict logsm)
|
|
{
|
|
const ulong
|
|
ent_k = max(opts->top_k, 1U) - 1,
|
|
rnd = ircd_simt_rand_xoshiro256pl(ctrl->rand);
|
|
|
|
const float
|
|
ent_p = min(max(opts->top_p, 0.0f), 1.0f),
|
|
thresh = ent_p;
|
|
|
|
float acc = 1.0f;
|
|
ushort select = 0;
|
|
for(; select < ent_k; ++select)
|
|
if((acc -= logsm[idx[select]]) < thresh)
|
|
break;
|
|
|
|
const ushort
|
|
token = idx[select];
|
|
|
|
return token;
|
|
}
|
|
|
|
static ushort
|
|
ircd_gpt_lm_result(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__local const ushort *const restrict idx,
|
|
__global const float *const restrict logsm)
|
|
{
|
|
const ushort
|
|
token = ircd_gpt_lm_result_select(ctrl, opts, idx, logsm);
|
|
|
|
// Update the dynamic result label.
|
|
ctrl->select.logit.token = token;
|
|
ircd_gpt_lm_result_label(ctrl, opts, &ctrl->select, logsm);
|
|
|
|
// Update the dynamic target label.
|
|
ctrl->target.logit.token = ctrl->count < ctrl->tokens?
|
|
ctrl->token[ctrl->count]:
|
|
ctrl->select.logit.token;
|
|
|
|
ircd_gpt_lm_result_label(ctrl, opts, &ctrl->target, logsm);
|
|
|
|
const bool
|
|
hit = ctrl->select.logit.token == ctrl->target.logit.token;
|
|
|
|
// Update the token context.
|
|
if(ctrl->count == ctrl->tokens)
|
|
ctrl->token[ctrl->tokens++] = ctrl->select.logit.token;
|
|
else
|
|
ctrl->accept = -2;
|
|
|
|
ctrl->miss += !hit;
|
|
ctrl->hit += hit;
|
|
ctrl->count++;
|
|
return token;
|
|
}
|
|
|
|
static void
|
|
ircd_gpt_lm_result_attns(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global const float *const restrict attns,
|
|
const uint ln,
|
|
const uint li)
|
|
{
|
|
const uint
|
|
layer = li / opts->layers,
|
|
head = li % opts->attn_rank,
|
|
base = layer * opts->attn_self_elems;
|
|
|
|
uint best = 0;
|
|
float bestv = 10000.0f;
|
|
for(uint i = 0; i < ctrl->count; ++i)
|
|
{
|
|
const uint
|
|
bx = (((i + 1) * i) / 2) * opts->attn_rank,
|
|
idx = base + bx + i * 12 + head;
|
|
|
|
if(attns[idx] < bestv)
|
|
bestv = attns[idx],
|
|
best = i;
|
|
}
|
|
|
|
ctrl->attn[layer][head] = best;
|
|
}
|
|
|
|
__kernel void
|
|
__attribute__((visibility("protected")))
|
|
__attribute__((reqd_work_group_size(256, 1, 1)))
|
|
ircd_gpt_lm_select(__global struct ircd_gpt_ctrl *const restrict ctrl_,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global const float logsm[restrict 65536],
|
|
__global const float *const restrict attns)
|
|
{
|
|
const uint
|
|
li = get_local_id(0),
|
|
ln = get_local_size(0),
|
|
logits_pad = ln - (opts->logits % ln),
|
|
tn = (opts->logits + logits_pad) / ln,
|
|
start = tn * li,
|
|
stop = min(start + tn, opts->logits);
|
|
|
|
__local ushort idx[256];
|
|
__local struct ircd_gpt_ctrl ctrl;
|
|
__private event_t event[1];
|
|
|
|
assume(ln == 256);
|
|
assume(start < stop);
|
|
|
|
event[0] = async_work_group_copy
|
|
(
|
|
(__local char16 *)&ctrl,
|
|
(__global const char16 *)ctrl_,
|
|
sizeof(struct ircd_gpt_ctrl) / sizeof(char16),
|
|
0
|
|
);
|
|
|
|
idx[li] = start;
|
|
for(uint j = start + 1; j < stop; ++j)
|
|
if(logsm[j] > logsm[idx[li]])
|
|
idx[li] = j;
|
|
|
|
ircd_simt_sort_idx16_flldr(idx, logsm, ln, li);
|
|
wait_group_events(1, event);
|
|
|
|
if(ctrl.count >= opts->buffer_tokens)
|
|
return;
|
|
|
|
if(li < opts->top_n)
|
|
ircd_gpt_lm_result_top(&ctrl, opts, idx, logsm, li);
|
|
|
|
if(li < opts->labels)
|
|
ircd_gpt_lm_result_label(&ctrl, opts, ctrl.label + li, logsm);
|
|
|
|
if(li < opts->layers * opts->attn_rank)
|
|
ircd_gpt_lm_result_attns(&ctrl, opts, attns, ln, li);
|
|
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
if(li == 0)
|
|
ircd_gpt_lm_result(&ctrl, opts, idx, logsm);
|
|
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
event[0] = async_work_group_copy
|
|
(
|
|
(__global char16 *)ctrl_,
|
|
(__local const char16 *)&ctrl,
|
|
sizeof(struct ircd_gpt_ctrl) / sizeof(char16),
|
|
0
|
|
);
|
|
|
|
wait_group_events(1, event);
|
|
}
|
|
|
|
__kernel void
|
|
__attribute__((visibility("protected")))
|
|
__attribute__((reqd_work_group_size(256, 1, 1)))
|
|
ircd_gpt_leave(__global const void *const restrict model,
|
|
__global void *const restrict state,
|
|
__global void *const restrict master,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global struct ircd_gpt_ctrl *const ctrl_,
|
|
__global struct ircd_gpt_ctrl *const frame)
|
|
{
|
|
const ushort
|
|
li = get_local_id(0),
|
|
ln = get_local_size(0);
|
|
|
|
assume(ln == 256);
|
|
|
|
__local struct ircd_gpt_ctrl _ctrl;
|
|
__local struct ircd_gpt_ctrl *const ctrl = &_ctrl;
|
|
|
|
if(li == 0)
|
|
*ctrl = *ctrl_;
|
|
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
if(li == 0 && ctrl->accept < 0)
|
|
ircd_gpt_accept(ctrl, opts);
|
|
|
|
barrier(CLK_LOCAL_MEM_FENCE);
|
|
|
|
const uint
|
|
batch_size = opts->batch_size,
|
|
samps = opts->training_steps + opts->validation_steps + opts->testing_steps,
|
|
steps = samps / batch_size;
|
|
|
|
const bool
|
|
accepting = ctrl->accept >= 0,
|
|
cycling = !accepting,
|
|
sampling = accepting,
|
|
stepping = sampling && (ctrl->clk.samp + 1) >= batch_size,
|
|
epoching = stepping && (ctrl->clk.step + 1) >= steps;
|
|
|
|
if(li == 0)
|
|
ctrl->prof.finished = ircd_simt_cycles();
|
|
|
|
if(li == 0)
|
|
{
|
|
// BARTS won't update the scalar after the copy. In this case we'll
|
|
// be setting magic for all remaining frames :/
|
|
#if defined(__R600__)
|
|
ctrl->magic = 0xC7012C70UL;
|
|
#endif
|
|
|
|
*frame = *ctrl;
|
|
frame->magic = 0xC7012C70UL;
|
|
}
|
|
|
|
if(li == 0 && !accepting)
|
|
{
|
|
ctrl->clk.cycle += cycling;
|
|
ctrl->clk.samp += sampling;
|
|
ctrl->clk.step += stepping;
|
|
ctrl->clk.epoch += epoching;
|
|
}
|
|
|
|
if(li == 0 && accepting)
|
|
ctrl->magic = 0xC7012C70UL;
|
|
|
|
if(li == 0)
|
|
*ctrl_ = *ctrl;
|
|
}
|
|
|
|
void
|
|
ircd_gpt_accept(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts)
|
|
{
|
|
const bool
|
|
unlimited = opts->limit < 0,
|
|
unprocessed = ctrl->accept == -2,
|
|
acceptable = ctrl->accept < 0 && !unprocessed;
|
|
|
|
const uint
|
|
batch_size = opts->batch_size,
|
|
samps = opts->training_steps + opts->validation_steps + opts->testing_steps,
|
|
steps = samps / batch_size,
|
|
unproc = ctrl->tokens - ctrl->count,
|
|
limit = opts->limit;
|
|
|
|
const int
|
|
cycle_remain = opts->context_tokens - (ctrl->clk.cycle + 1), // cycle not yet incr
|
|
token_remain = opts->context_tokens - ctrl->count, // but count already incr
|
|
remain = min(cycle_remain, token_remain);
|
|
|
|
int
|
|
accept = ircd_gpt_accept_check(ctrl, opts),
|
|
dispatch = accept < 0? min(abs(accept), (uint)remain): 0;
|
|
|
|
if(opts->limit > 0 && ctrl->clk.cycle >= limit - 1)
|
|
{
|
|
accept = accept >= 0? accept: -1;
|
|
dispatch = 0;
|
|
}
|
|
|
|
if(opts->limit == 0 && unprocessed)
|
|
{
|
|
accept = accept >= 0? accept: -1;
|
|
dispatch = min((uint)remain, unproc);
|
|
}
|
|
|
|
if(opts->limit != 0 && !acceptable)
|
|
{
|
|
accept = accept >= 0? accept: -1;
|
|
dispatch = max(dispatch, 1);
|
|
}
|
|
|
|
ctrl->accept = accept;
|
|
ctrl->dispatch = min((uint)dispatch, opts->frames);
|
|
}
|
|
|
|
int
|
|
ircd_gpt_accept_check(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts)
|
|
{
|
|
int best = opts->frames;
|
|
for(uint i = 0; i < 4; ++i)
|
|
{
|
|
const int
|
|
remain = ircd_gpt_accept_match(ctrl, opts, i);
|
|
|
|
if(remain == 0)
|
|
return i;
|
|
|
|
if(remain < best)
|
|
best = remain;
|
|
}
|
|
|
|
return -best;
|
|
}
|
|
|
|
uint
|
|
ircd_gpt_accept_match(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
const uint i)
|
|
{
|
|
const uint
|
|
len = ircd_gpt_accept_len(ctrl, opts, i),
|
|
n = min(ctrl->count, len),
|
|
maxlen = opts->frames;
|
|
|
|
uint ret = len?: maxlen;
|
|
for(uint j = 1; j <= n; ++j)
|
|
{
|
|
uint match = 0;
|
|
for(; match < j; ++match)
|
|
{
|
|
const uint
|
|
accept = opts->accept[i][match],
|
|
token = ctrl->token[ctrl->count - j + match];
|
|
|
|
if(token != accept)
|
|
break;
|
|
}
|
|
|
|
if(match >= j)
|
|
if(!(ret = len - match))
|
|
break;
|
|
}
|
|
|
|
ret = max(ret, ctrl->tokens - ctrl->count);
|
|
ret = min(ret, maxlen);
|
|
return ret;
|
|
}
|
|
|
|
uint
|
|
ircd_gpt_accept_len(__local struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
const uint i)
|
|
{
|
|
uint len = 0;
|
|
for(; len < 8; ++len)
|
|
if(opts->accept[i][len] == (ushort)-1U)
|
|
break;
|
|
|
|
return len;
|
|
}
|
|
|
|
//
|
|
// backside
|
|
//
|
|
|
|
void
|
|
__attribute__((always_inline))
|
|
ircd_gpt_prop_elem(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global float4 *const restrict param_,
|
|
__global float4 *const restrict exp_avg_,
|
|
__global float4 *const restrict exp_avg_sqr_)
|
|
{
|
|
const uint
|
|
li = get_local_id(0),
|
|
ts = ctrl->clk.step;
|
|
|
|
const float4
|
|
param = param_[li],
|
|
grad = ctrl->target.loss.mean,
|
|
alpha[2] = { 1.0f - opts->beta[0], 1.0f - opts->beta[1], },
|
|
exp_avg = ts? exp_avg_[li]: 0.0f,
|
|
exp_avg_sqr = ts? exp_avg_sqr_[li]: 0.0f,
|
|
exp_avg_mul = exp_avg * opts->beta[0],
|
|
exp_avg_dot = exp_avg_mul + alpha[0] * grad,
|
|
exp_avg_sqr_mul = exp_avg_sqr * opts->beta[1],
|
|
exp_avg_sqr_dot = exp_avg_sqr_mul + alpha[1] * grad * grad,
|
|
denom = native_sqrt(exp_avg_sqr_dot) + FLT_EPSILON,
|
|
delta = opts->alpha * (exp_avg_dot / denom),
|
|
update = param - delta;
|
|
|
|
param_[li] = update;
|
|
exp_avg_[li] = exp_avg_dot;
|
|
exp_avg_sqr_[li] = exp_avg_sqr_dot;
|
|
}
|
|
|
|
//
|
|
// backpropagations
|
|
//
|
|
|
|
__kernel void
|
|
__attribute__((always_inline))
|
|
ircd_gpt_norm_prop(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global ircd_gpt_vectorv *const restrict bias,
|
|
__global ircd_gpt_vectorv *const restrict bias_m0,
|
|
__global ircd_gpt_vectorv *const restrict bias_m1,
|
|
__global ircd_gpt_vectorv *const restrict weight,
|
|
__global ircd_gpt_vectorv *const restrict weight_m0,
|
|
__global ircd_gpt_vectorv *const restrict weight_m1)
|
|
{
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
bias->elem,
|
|
bias_m0->elem,
|
|
bias_m1->elem
|
|
);
|
|
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
weight->elem,
|
|
weight_m0->elem,
|
|
weight_m1->elem
|
|
);
|
|
}
|
|
|
|
__kernel void
|
|
ircd_gpt_coil_prop_attn(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global ircd_gpt_vectorv *const restrict norm_bias,
|
|
__global ircd_gpt_vectorv *const restrict norm_bias_m0,
|
|
__global ircd_gpt_vectorv *const restrict norm_bias_m1,
|
|
__global ircd_gpt_vectorv *const restrict norm_weight,
|
|
__global ircd_gpt_vectorv *const restrict norm_weight_m0,
|
|
__global ircd_gpt_vectorv *const restrict norm_weight_m1,
|
|
__global ircd_gpt_attn_aperaturev *const restrict fcon_bias,
|
|
__global ircd_gpt_attn_aperaturev *const restrict fcon_bias_m0,
|
|
__global ircd_gpt_attn_aperaturev *const restrict fcon_bias_m1,
|
|
__global ircd_gpt_attn_aperaturev *const restrict fcon_weight,
|
|
__global ircd_gpt_attn_aperaturev *const restrict fcon_weight_m0,
|
|
__global ircd_gpt_attn_aperaturev *const restrict fcon_weight_m1,
|
|
__global ircd_gpt_vectorv *const restrict proj_bias,
|
|
__global ircd_gpt_vectorv *const restrict proj_bias_m0,
|
|
__global ircd_gpt_vectorv *const restrict proj_bias_m1,
|
|
__global ircd_gpt_vectorv *const restrict proj_weight,
|
|
__global ircd_gpt_vectorv *const restrict proj_weight_m0,
|
|
__global ircd_gpt_vectorv *const restrict proj_weight_m1)
|
|
{
|
|
const uint
|
|
fcon_height = opts->embed_elems,
|
|
proj_height = opts->embed_elems,
|
|
segs = 3;
|
|
|
|
ircd_gpt_norm_prop
|
|
(
|
|
ctrl, opts,
|
|
norm_bias,
|
|
norm_bias_m0,
|
|
norm_bias_m1,
|
|
norm_weight,
|
|
norm_weight_m0,
|
|
norm_weight_m1
|
|
);
|
|
|
|
for(uint j = 0; j < segs; ++j)
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
fcon_bias->proj[j],
|
|
fcon_bias_m0->proj[j],
|
|
fcon_bias_m1->proj[j]
|
|
);
|
|
|
|
for(uint i = 0; i < fcon_height; ++i)
|
|
for(uint j = 0; j < segs; ++j)
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
fcon_weight[i].proj[j],
|
|
fcon_weight_m0[i].proj[j],
|
|
fcon_weight_m1[i].proj[j]
|
|
);
|
|
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
proj_bias->elem,
|
|
proj_bias_m0->elem,
|
|
proj_bias_m1->elem
|
|
);
|
|
|
|
for(uint i = 0; i < proj_height; ++i)
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
proj_weight[i].elem,
|
|
proj_weight_m0[i].elem,
|
|
proj_weight_m1[i].elem
|
|
);
|
|
}
|
|
|
|
__kernel void
|
|
ircd_gpt_coil_prop_ffnn(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global ircd_gpt_vectorv *const restrict norm_bias,
|
|
__global ircd_gpt_vectorv *const restrict norm_bias_m0,
|
|
__global ircd_gpt_vectorv *const restrict norm_bias_m1,
|
|
__global ircd_gpt_vectorv *const restrict norm_weight,
|
|
__global ircd_gpt_vectorv *const restrict norm_weight_m0,
|
|
__global ircd_gpt_vectorv *const restrict norm_weight_m1,
|
|
__global ircd_gpt_ffnn_aperaturev *const restrict fcon_bias,
|
|
__global ircd_gpt_ffnn_aperaturev *const restrict fcon_bias_m0,
|
|
__global ircd_gpt_ffnn_aperaturev *const restrict fcon_bias_m1,
|
|
__global ircd_gpt_ffnn_aperaturev *const restrict fcon_weight,
|
|
__global ircd_gpt_ffnn_aperaturev *const restrict fcon_weight_m0,
|
|
__global ircd_gpt_ffnn_aperaturev *const restrict fcon_weight_m1,
|
|
__global ircd_gpt_vectorv *const restrict proj_bias,
|
|
__global ircd_gpt_vectorv *const restrict proj_bias_m0,
|
|
__global ircd_gpt_vectorv *const restrict proj_bias_m1,
|
|
__global ircd_gpt_vectorv *const restrict proj_weight,
|
|
__global ircd_gpt_vectorv *const restrict proj_weight_m0,
|
|
__global ircd_gpt_vectorv *const restrict proj_weight_m1)
|
|
{
|
|
const uint
|
|
fcon_height = opts->embed_elems,
|
|
proj_height = opts->ffnn_elems,
|
|
segs = 4;
|
|
|
|
ircd_gpt_norm_prop
|
|
(
|
|
ctrl, opts,
|
|
norm_bias,
|
|
norm_bias_m0,
|
|
norm_bias_m1,
|
|
norm_weight,
|
|
norm_weight_m0,
|
|
norm_weight_m1
|
|
);
|
|
|
|
for(uint j = 0; j < segs; ++j)
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
fcon_bias->proj[j],
|
|
fcon_bias_m0->proj[j],
|
|
fcon_bias_m1->proj[j]
|
|
);
|
|
|
|
for(uint i = 0; i < fcon_height; ++i)
|
|
for(uint j = 0; j < segs; ++j)
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
fcon_weight[i].proj[j],
|
|
fcon_weight_m0[i].proj[j],
|
|
fcon_weight_m1[i].proj[j]
|
|
);
|
|
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
proj_bias->elem,
|
|
proj_bias_m0->elem,
|
|
proj_bias_m1->elem
|
|
);
|
|
|
|
for(uint i = 0; i < proj_height; ++i)
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
proj_weight[i].elem,
|
|
proj_weight_m0[i].elem,
|
|
proj_weight_m1[i].elem
|
|
);
|
|
}
|
|
|
|
__kernel void
|
|
ircd_gpt_lm_embed_prop(__global const struct ircd_gpt_ctrl *const ctrl,
|
|
__constant const struct ircd_gpt_opts *const opts,
|
|
__global ircd_gpt_vectorv *const restrict pos,
|
|
__global ircd_gpt_vectorv *const restrict pos_m0,
|
|
__global ircd_gpt_vectorv *const restrict pos_m1,
|
|
__global ircd_gpt_vectorv *const restrict token,
|
|
__global ircd_gpt_vectorv *const restrict token_m0,
|
|
__global ircd_gpt_vectorv *const restrict token_m1)
|
|
{
|
|
const uint
|
|
ln = get_local_size(0),
|
|
wi = get_global_offset(0) / ln + get_group_id(0),
|
|
wn = ctrl->count,
|
|
cn = opts->context_tokens / wn,
|
|
ci = cn * wi,
|
|
tn = opts->logits / wn,
|
|
ti = tn * wi;
|
|
|
|
for(uint i = ci; i < ci + cn; ++i)
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
pos[i].elem,
|
|
pos_m0[i].elem,
|
|
pos_m1[i].elem
|
|
);
|
|
|
|
for(uint i = ti; i < ti + tn; ++i)
|
|
ircd_gpt_prop_elem
|
|
(
|
|
ctrl, opts,
|
|
token[i].elem,
|
|
token_m0[i].elem,
|
|
token_m1[i].elem
|
|
);
|
|
}
|
|
|
|
/// Gaussian Error Linear Unit
|
|
void
|
|
ircd_gpt_ffnn_gelu(__local ircd_gpt_ffnn_aperaturev *const out,
|
|
__local const ircd_gpt_ffnn_aperaturev *const in_,
|
|
const uint i)
|
|
{
|
|
const float4
|
|
in = in_->fcon[i];
|
|
|
|
float4 a;
|
|
a = 0.044715f;
|
|
a *= in;
|
|
a *= in;
|
|
a += 1.0f;
|
|
a *= 0.7978845608f;
|
|
a *= in;
|
|
|
|
a = tanh(a);
|
|
a += 1.0f;
|
|
a *= in;
|
|
a *= 0.5f;
|
|
|
|
out->fcon[i] = a;
|
|
}
|
|
|
|
void
|
|
ircd_gpt_norm(__local ircd_gpt_vectorv *const out,
|
|
__local const ircd_gpt_vectorv *const in,
|
|
__local ircd_gpt_vectorv *const restrict tmp,
|
|
__global const ircd_gpt_vectorv *const restrict bias,
|
|
__global const ircd_gpt_vectorv *const restrict weight,
|
|
const uint ln,
|
|
const uint li)
|
|
{
|
|
// Layer re-normalization
|
|
ircd_simt_math_norm_f4lldr(out->elem, in->elem, tmp->elem, ln, li);
|
|
ircd_gpt_norm_fmad(out, out, bias, weight, li);
|
|
}
|
|
|
|
void
|
|
ircd_gpt_norm_fmad(__local ircd_gpt_vectorv *const out,
|
|
__local const ircd_gpt_vectorv *const in,
|
|
__global const ircd_gpt_vectorv *const restrict bias,
|
|
__global const ircd_gpt_vectorv *const restrict weight,
|
|
const uint i)
|
|
{
|
|
out->elem[i] = in->elem[i] * weight->elem[i] + bias->elem[i];
|
|
}
|