[Ciko 0]

Ця програма є вільною від залежностей реалізацію GPT-2. Вона завантажує матрицю ваг і файл BPE з оригінальних файлів TensorFlow, токенізує висновок за допомогою простого енкодера, що працює за принципом частотного кодування, реалізує базовий пакет для лінійної алгебри, в якому укладені математичні операції над матрицями, визначає архітектуру трансформера, виконує інференс трансформера, а потім Все це приблизно в 3000 байт на C.

Код досить ефективно оптимізований настільки, що малий GPT-2 на будь-якій сучасній машині видає відгук всього за кілька секунд. Щоб цього досягти, я реалізував KV-кешування та застосував ефективний алгоритм перемноження матриць, а також додав опціональний OMP-паралелізм.

Взявши це за основу, можна створити якийсь аналог ChatGPT - за умови, що вас не дуже хвилює якість виводу (об'єктивно кажучи, висновок виходить просто жахливий... але рішення працює). Тут є деякі глюки (особливо з обробкою символів кодування UTF-8), а для експлуатації моделі розміром XL з широким контекстним вікном може знадобитися ~100 ГБ оперативної пам'яті. Але якщо ви просто набираєте текст у кодуванні ASCII за допомогою малого GPT2, то така модель має нормально працювати приблизно скрізь.

Я виклав весь код на GitHub , тому можете вільно брати його там і експериментувати з ним.

Ця програма складається з наступних основних блоків (кожен з них супроводжується відповідним кодом): Бібліотека для базової матричної математики (700 байт) Швидке перемноження матриць ( 300 байт) Шари нейронної мережі ( 300 байт) Модель-трансформер (600 байт4 ) байт) Завантаження ваги (300 байт) Завантаження даних для частотного кодування (300 байт)

#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<math.h>
int U,C,K,c,d,S,zz;char*bpe;typedef struct{float*i;int j,k;} A;void*E,*n;A*f;FILE*fp;
#define N(i,j)for(int i=0; i<j; i++)
A o(int j,int k,int i){float*a=E;E+=S=4*j*k;memset(a,0,S*i);A R={ a,j,k} ;return R;}
#define I(R,B)A R(A a,float k){ N(i,a.j*a.k){ float b=a.i[i]; a.i[i]=B; } return a; }
I(l,b/k)I(q,b+k)I(u,1./sqrt(b))I(z,exp(b))I(r,a.i[(i/a.k)*a.k])I(P,(i/k<i%(int)k)?0:exp(b/8))I(Q,b/2*(1+tanh(.7978845*(b+.044715*b*b*b))))
#define F(R,B)A R(A a,A b){ N(i,a.j*a.k){ a.i[i]=a.i[i]B b.i[i]; } return a; }
F(V,+)F(v,*)F(H,/)F(at,+b.i[i%a.k];)F(mt,*b.i[i%a.k];)A X(A a){A R=o(a.j,a.k,1);N(i,a.j*a.k)R.i[(i/a.k)*a.k]+=a.i[i];r(R,0);return R;}A p(A a){A R=o(a.k,a.j,1);N(i,a.j*a.k)R.i[i%a.k*a.j+i/a.k]=a.i[i];return R;}
A g(A a,A b){A R=o(a.j,b.j,!c);{for(int i=c;i<d;i++){for(int j=0;j<b.j;j+=4){for(int k=0;k<a.k;k+=4){N(k2,4)N(j2,4)R.i[i*b.j+j+j2]+=a.i[i*a.k+k+k2]*b.i[(j+j2)*b.k+k+k2];}}}}return
V(o(R.j,R.k,1),R);}A J(A a,int b,int j,int k){A R={ a.i+b*j,j,k} ;return R;}A s(A a,int i){A b=V(a,l(X(a),-a.k));A k=l(X(v(V(o(b.j,b.k,1),b),b)),b.k-1);A R=at(mt(v(V(o(b.j,b.k,1),b),u(q(k,1e-5),0)),f[i+1]),f[i]);return R;}
#define G(a,i)at(g(a,f[i+1]),f[i])
A m(int j,int k){j+=!j;k+=!k;A a=o(j,k,1);fread(a.i,S,1,fp);return p(a);}
int t;int Y(char*R){if(!*R)return 0;int B=1e9,r;N(i,5e4){if(bpe[999*i]&&strncmp(bpe+999*i,R,S=strlen(bpe+999*i))==0){int k=Y(R+S)+i+1e7;if(k<B){B=k;r=i;}}}t=r;return B;}int *w(char*q,int*B){char R[1000];int i=0;while(q[i]){int j=i++;while(47<q[i]&&q[i]<58||64<q[i]){fflush(stdout);i++;}strcpy(R,q+j);R[i-j]=0;fflush(stdout);int k=0;while(R[k]){Y(R+k);char*M=bpe+t*999;k+=strlen(M);*B++=t;}}return B;}
int main(int S,char**D){S=D[1][5]+3*D[1][7]+3&3;K=12+4*S+(S>2);U=K*64;C=12*S+12;zz=atoi(D[4]);E=malloc(2LL*U*U*C*zz);
bpe=malloc(1e9);fp=fopen(D[2],"r");unsigned char a[S=999],b[S];N(i,5e4){int k=i*S;if(i<93){bpe[k]=i+33;bpe[k+1]=0;} else if(i>254){fscanf(fp,"%s %s",a,b);strcat((char*)a,(char*)b);int j=0;N(i,a[i])bpe[k+j++]=a[i]^196?a[i]:a[++i]-128;bpe[k+j++]=0;} else if(i>187){bpe[k]=i-188;bpe[k+1]=0;}}int e[1024];d=w(D[3],e)-e;int h;N(i,d){if(e[i]==18861)h=i+1;}printf("AI");N(i,d-h)printf("%s",bpe+e[i+h]*999);
fp=fopen(D[1],"r");A\
 x[999];A*R=x;N(i,C){N(j,12)*R++=m(U+U*(j?j^8?j^11?0:3:3:2),U*((j%8==3)+3*(j%8==1)+(j==9)));}*R++=m(U,1);*R++=m(U,1);A QA=m(1024,U),Z=p(m(5e4,U));
while(1){char W[1000]={ 0} ;int T;strcat(W,"\nAlice: ");printf("\n%s: ",bpe+20490*999);fflush(stdout);fgets(W+8,1000,stdin);printf("AI:");strcat(W,"\nBob:");d=w(W,e+d)-e;n=E;c=0;
while(1){E=n;T=d+32-d%32;c*=!!(d%32);A O=o(T,U,1);N(i,d){N(j,U)O.i[i*U+j]=Z.i[e[i]*U+j]+QA.i[j*1024+i];}N(i,C){int y;S=0;N(j,10){if(j==i)y=S;S++;N(k,10*(j>0)){if(j*10+k<C&&S++&&i==j*10+k)y=S;}}f=x+12*y;A QB=p(J(G(s(O,4),0),0,T*3,U));A B=o(U,T,1);N(k,K){A L=p(J(QB,k*3,64*T,3)),a=P(g(p(J(L,0,64,T)),p(J(L,T,64,T))),T),R=p(g(H(a,X(a)),J(L,T*2,64,T)));memcpy(B.i+64*T*k,R.i,64*T*4);}O=V(O,G(p(B),2));O=V(O,G(Q(G(s(O,6),8),0),10));}f=x;O=s(O,12*C);c=0;int S=d;d=1;A B=g(p(J(O,S-1,U,1)),Z);c=d=S;S=0;N(i,5e4){if(B.i[i]>B.i[S])S=i;}if(d==zz){memcpy(e,e+zz/2,S*2);d-=zz/2;c=0;}e[d++]=S;
if(bpe[S*999]==10)break;printf("%s",bpe+S*999);fflush(stdout);}}}

Контекст: ChatGPT та трансформери

Тим, хто в танку – нагадаю, що ChatGPT – це така програма. У ньому ви можете спілкуватися з так званою «(великою) мовною моделлю» як із співрозмовником-людиною. Вона напрочуд добре підтримує розмову, а GPT-4, новітня модель, що лежить в основі ChatGPT, взагалі неймовірно вражає.

У цій програмі на C відтворено поведінку ChatGPT, але за допомогою набагато слабшої моделі GPT-2, що з'явилася ще в 2019 році. Незважаючи на те, що номер версії у неї всього на 2 менший, ніж у GPT-4, можливості їх просто непорівнянні — проте модель GPT-2 є опенсорною. Тому з нею і працюватимемо.

GPT-2- це тип моделей машинного навчання, так званий"трансформер". Такі нейронні мережі приймають на вхід фіксовану послідовність слів, після чого прогнозують слово, яке має бути наступним. Знову і знову повторюючи цю процедуру можна генерувати за допомогою трансформера лексичні послідовності довільної довжини.

У цьому пості я не збираюся робити настільки повний екскурс у машинне навчання, після якого ви б усвідомили,чомутрансформер спроектований так, а не інакше. Далі просто буде описано, як саме працює наведений вище код на C.

Розбір коду C

Почнемо: матрична математика (700 байт)

Вся картина світу нейронної мережі укладена в матричних операціях. Тому для початку нам потрібно буде спорудити бібліотеку для роботи з матрицями, витративши на це мінімум байт.

Ось мінімальне визначення матриці:

typedef struct {
  float* dat;
  int rows, cols;
} Matrix;

Для початку відзначимо, що, нехай нам і потрібно реалізувати цілу низку різних операцій, всі вони діляться на два принципові «типу»:

1. Матрично-константні операції (напр., додати 7 до кожної з записів в матриці) 2.

Матрично-матричні операції (напр., скласти відповідні матричні

витягнути деяку загальну логіку метапроцедуру, в якій прописано, наприклад, як поводитися з парами матриць. У разі деталі конкретних операцій залежатимуть від реалізації.

Щоб зробити це на C, визначимо функцію

#define BINARY(function, operation)

як

Matrix FUNCTION(Matrix a, Matrix b) {
  for (int i = 0; i < a.rows; i++) {
    for (int j = 0; j < a.cols; j++) {
      a[i*a.cols + j] = a[i*a.cols + j] OPERATION b[i*a.cols+j];
    }
  }
  return a;
}

Що, наприклад, дозволить нам написати
BINARY(matrix_elementwise_add, +); BINARY(matrix_elementwise_multiply, *);

та передбачити можливість розширення до повної операції, в рамках якої відбувалося б поелементне додавання або перемноження двох матриць. Визначу ще кілька операцій, зрозуміти які не складає труднощів:

Тепер зупинимося на прийнятих у C#defines. По суті, це просто розфуфирені регулярні вирази. Тому коли запустимо такий #define, у програмі насправді відбудеться

a[i*a.cols + j] = a[i*a.cols + j] OPERATION b[i*a.cols+j];

що у випадку з множенням розшириться до

a[i*a.cols + j] = a[i*a.cols + j] * b[i*a.cols+j];

На перший погляд, цей код виглядає досить заплутано — наприклад, що тут робить ця точка з комою? Але, якщо ви просто заміните код регулярним виразом, то побачите, як він розширюється до

a[i*a.cols + j] = a[i*a.cols + j] + b.dat[i%a.cols] ; b[i*a.cols+j];

Оскільки другий вираз тут нічого не робить, цей код фактично еквівалентний
a[ia.cols + j] = a[ia.cols + j] + b.dat[i%a.cols] ; b[i*a.cols+j];

(КОРИСТУЙТЕСЯ МОВАМИ З ЯКІСНИМИ МАКРОСАМИ. LISP НЕ ЗАВЖДИ КРАЩЕ C!)

Швидке перемноження матриць (300 байт)

Базова реалізація перемноження матриць зовсім проста: ми всього лише реалізуємо тривіальні цикли з кубічною обчислювальною складністю. (У моєму прикладі з перемноженням матриць немає нічого особливого. Якщо ви вмієте швидко перемножувати матриці, то просто стежте за кодом).

Matrix matmul(Matrix a, Matrix b) {
  Matrix out = NewMatrix(a.rows, b.rows);
  for (int i = 0; i < a.rows; i++)
    for (int j = 0; j < b.rows; j++)
      for (int k = 0; k < a.cols; k++)
        out.dat[i * b.rows + j] += a.dat[i * a.cols + k+k2] * b.dat[(j+j2) * b.cols + k];


  return out;
}

На щастя, цю процедуру можна значно прискорити, зробивши її трохи розумнішою.
Matrix matmul_t_fast(Matrix a, Matrix b) {
Matrix out = NewMatrix(a.rows, b.rows);
for (int i = 0; i < a.rows; i++)
for (int j = 0; j < b.rows; j += 4)
for (int k = 0; k < a.cols; k += 4)
for (int k2 = 0; k2 < 4; k2 += 1)
for (int j2 = 0; j2 < 4; j2 += 1)
out.dat[i b.rows + j+j2] += a.dat[i a.cols + k+k2] b.dat[(j+j2) b.cols + k+k2];


return out;
}

Пізніше ми внесемо ще одну зміну до механізму логічного висновку, а також додамо до перемноження матриць ще один параметр, який дозволить нам лише частково множити матрицю A на матрицю B. Це корисно у випадках, коли ми встигли попередньо обчислити частину твору.

Шари нейронної мережі (300 байт)

Щоб написати трансформер, потрібно визначити кілька специфічних шарів нейронної мережі. Один з них - це функція активації GELU , яку можете сприймати як чаклунську.

UNARY(GELU, b / 2 * (1 + tanh(.7978845 * (b + .044715 * b * b * b))))

Також я реалізую функцію, що задає нижню діагональ матриці (після зведення значень у ступінь).

UNARY(tril, (i/k<i%(int)k) ? 0 : exp(b/8))

Нарешті, нам знадобиться функція нормалізації шарів (ще одна магічна річ, про яку можете самі почитати докладніше, якщо захочете. По суті вона нормалізує середнє і дисперсію кожного шару).

Matrix LayerNorm(Matrix a, int i) {
  Matrix b = add(a, divide_const(sum(a), -a.cols));
  Matrix k = divide_const(sum(multiply(
    add(NewMatrix(b.rows,b.cols,1),b), b)), b.cols-1);
  Matrix out = add_tile(multiply_tile(
    multiply(add(NewMatrix(b.rows,b.cols,1),b),
    mat_isqrt(add_const(k, 1e-5),0)), layer_weights[i+1]),
                        layer_weights[i]);


  return out;
}

Останній елемент моделі - це лінійна функція, що просто перемножує матриці і плюсує зрушення (з розбиттям на макрооперації - тайлінгом).

#define Linear(a, i) add_tile(matmul_t_fast(a, layer_weights[i+1]), layer_weights[i])

Архітектура трансформера (600 байт)

Вирішивши всі ці питання, ми, нарешті, зможемо реалізувати наш трансформер лише у 600 байтах.

for (int i = 0; i < NLAYER; i++) {
  layer_weights = weights + 12*permute;


  // Вычисляем ключи, запросы и значение — всё за одну большую операцию умножения
  Matrix qkv = transpose(slice(Linear(LayerNorm(line, 4), 0), 0, T*3, DIM));


  // Освобождаем место для вывода из вычислений
  Matrix result = NewMatrix(DIM, T, 1);


  for (int k = 0; k < NHEAD; k++) {
    // Распределяем qkv на три вычислительные головы
    Matrix merge = transpose(slice(qkv, k*3, 64*T, 3)),
      // Получаем произведение запросов и ключей, а затем возводим результат в степень
      a = tril(matmul_t_fast(transpose(slice(merge, 0, 64, T)),
                             transpose(slice(merge, T, 64, T))), T),
      // наконец, умножаем вывод softmax (a/sum(a)) на матрицу значеемй
      out = transpose(matmul_t_fast(divide(a, sum(a)), slice(merge, T*2, 64, T)));
    // и копируем вывод в ту часть результирующей матрицы, где он должен находиться
    memcpy(result.dat+64*T*k, out.dat, 64*T*4);
  }


  // Остаточная связь
  line = add(line,Linear(transpose(result), 2));


  // Функция активации и остаточная связь
  line = add(line, Linear(GELU(Linear(LayerNorm(line, 6), 8), 0), 10));
}


// Сбросить веса слоёв так, чтобы последний слой можно было взять за норму
layer_weights = weights;
line = LayerNorm(line, 12*NLAYER);


Matrix result = matmul_t_fast(transpose(slice(line, tmp-1, DIM, 1)), wte);

Тепер озвучу один момент про логічний висновок у трансформерах, для вас, можливо, цілком очевидний. Коли ви вже викликали модель, наказавши їй згенерувати один токен, вам не доводиться перерахувати всю функцію для створення наступного токена. Насправді, для генерації кожного наступного токена потрібно виконати лише мінімум роботи.

Справа в тому, що як тільки ви обчислили висновок трансформера для всіх токенів аж до N-го, ви можете повторно використовувати майже весь цей висновок для обчислення N+1го токена (виконавши ще трохи роботи.)

Щоб все це реалізувати, я побудував всі операції виділення пам'яті в коді послідовно, в межах одного і того ж блоку пам'яті. Так гарантується, що при будь-якій операції матричного множення буде задіяна та сама пам'ять. Відповідно, на кожній ітерації циклу я можу не обнулювати пам'ять перед тим, як задіяти її на наступній ітерації, і в пам'яті вже буде результат попередньої ітерації. Мені просто потрібно виконати обчислення для N + рядка.

Частотне кодування (400 байт)

Найпростіше побудувати мовну модель з урахуванням послідовності слів. Але, оскільки загальна безліч слів, по суті, не обмежена, для будь-якої мовної моделі потрібно введення фіксованого розміру, де знадобиться замінити досить рідкісні слова спеціальним токеном [OUT OF DISTRIBUTION] (поза навчальним розподілом). Це недобре.

Так, існує найпростіший засіб, що дозволяє з цим впоратися використовувати моделі, що працюють «на рівні символів» і тому знають тільки окремі літери. Але тут виникає проблема: фактично, такої моделі доведеться вивчати значення кожного слова з чистого аркуша. Також через це звузиться реальний розмір контекстного вікна мовної моделі, і коефіцієнт такого зниження дорівнює середній довжині слова.

Щоб уникнути таких проблем, моделі на кшталт GPT-2 створюють токени з «підслів». Деякі слова можуть бути токенами власними силами, але рідкісні слова додатково дробляться. Наприклад, слово "nicholas" можна поділити на "nich", "o", "las".

Реалізувати загальний алгоритм для цієї мети нескладно: беремо слово, яке хочемо токенізувати, і спочатку поділяємо його на окремі символи. Потім шукаємо пари таких суміжних токенів, які можна було б поєднати, і якщо знаходимо - поєднуємо. Повторюємо процедуру доти, доки варіантів подальшого злиття не залишиться.

При всій простоті цей алгоритм, на жаль, дуже важко реалізувати на C, тому що для нього потрібно багаторазово виділяти пам'ять та відстежувати розвиток деревоподібної структури токенів.

Тому в такому випадку ми перетворюємо досить простий алгоритм з лінійною складністю на алгоритм з потенційно експоненційною складністю, зате пишемо набагато менше коду. Базова ідея розкриватиметься приблизно так, як показано в цьому C-подібному псевдокоді:

word_tokenize(word) {
  if len(word) == 0 { return (0, 0); }
  result = (1e9, -1);
  for (int i = 0; i < VOCAB_LEN; i++) {
    if (is_prefix(bpe[i]), word) {
      sub_cost = word_tokenize(word+len(bpe[i]))[0] + i + 1e7;
      result = min(result, (sub_cost, i));
    }
  }
  return result;
}

Тобто, щоб токенізувати слово, потрібно перевірити всі можливі слова зі словника та дізнатися, чи воно не є префіксом актуального слова. Якщо є, то ми візьмемо його як перший токен, а потім спробуємо так само рекурсивно токенізувати все слово. Ми стежитимемо, який варіант токенізації виходить найкращим (про це судимо за довжиною, нічийні результати розподіляємо за індексом токена у словнику) — і саме його повертаємо.

Завантаження ваги (300 байт)

Майже готово! Останнє, що нам залишилося зробити, — завантажити з диска в нейронну мережу фактичні ваги. Це насправді не складно, оскільки ваги зберігаються у простому двійковому форматі, який легко зчитується в C. Фактично це абсолютно плоска серіалізація 32-розрядних чисел з плаваючою точкою.

Єдине, що потрібно дізнатися, наскільки великі різні матриці. На щастя, це також легко з'ясувати. Який би не був розмір моделі GPT-2, архітектура у них у всіх однакова, і ваги зберігаються в тому самому порядку. Тому нам лише потрібно прочитати з диска правильно оформлені матриці.

Але наприкінці ложка дьогтю. Шари нейронної мережі не зберігаються на диску в тому порядку, в якому слід було б очікувати: спочатку шар 0, потім шар 1, далі шар 2. Насправді першим йде шар 0, за ним шар 1, а потім шар .... ДЕСЯТЬ! (І далі шар 11, а за ним 12.) Справа в тому, що при збереженні ваги сортуються за лексикографічним принципом. А лексикографічно "0" передує "1", але "10" передує "2". Тому нам потрібно трохи попрацювати, щоб переставити ваги в правильному порядку. Для цього напишемо наступний код

int permute;
tmp=0;
for (int j = 0; j < 10; j++) {
  if (j == i) {
    permute = tmp;
  }
  tmp++;
  for (int k = 0; k < 10*(j>0); k++) {
    if (j*10+k < NLAYER && tmp++ && i == j*10+k) {
      permute = tmp;
    }
  }
}

Завантаження даних для частотного кодування (300 байт)

Щоб здійснити частотне кодування, ми спочатку повинні завантажити з диска словник з відповідними байтовими парами. В ідеалі хотілося б мати список усіх слів зі словника, збережених у якомусь осудному C-читаному форматі. Але оскільки вихідний файл (a) розрахований на читання в Python і (b) не призначений для легкого синтаксичного аналізу мінімальних байтових фрагментів, нам тут доведеться попрацювати.

Логічно припустити, що формат файлу передбачає просто список слів, що йдуть один за одним, але насправді інформація в ньому закодована у вигляді списку байтових пар. Тобто, ми зможемо прочитати як один токен не «Hello», а рядок «H» «ello». Таким чином, нам слід об'єднати токени H і ello в один токен Hello.

Інша проблема полягає в тому, що файл представлений в кодуванні UTF-8, що згладжує (з застереженнями) і ... на те є причина. Усі символи ascii, які можна вивести на друк, кодуються самі по собі, а символи 0-31, що не виводяться, кодуються у форматі 188+символ. Так, наприклад, пропуск кодується у вигляді токена «Ġ». Але проблема: в кодуванні UTF8 на диску символу «Ġ» відповідає 0xc4 0xa0. Тому при зчитуванні нам доведеться докласти зусиль, щоб перетворити цей символ назад на прогалину.

При тому, що все це, в принципі, ніщо з цього не складно, для всіх операцій потрібно написати багато коду. Це, звичайно, дратує, коли прагнеш максимальної компактності.

unsigned char a[tmp=999],b[tmp];
for (int i = 0; i < 5e4; i++) {
  int k = i*tmp;
  if (i < 93) {
    // Первые 92 токена – это просто символы ascii, выводимые на печать
    bpe[k] = i + 33;
    bpe[k+1] = 0;
  } else if (i > 254) {
    // Те, что сверх 254, взяты из файла BPE. Загружаем их
    fscanf(fp, "%s %s", a, b);
    strcat((char*)a, (char*)b);
    int j = 0;
  for (int i = 0; a[i]; i++) {
    // Кодировка UTF8 усложняет жизнь, поэтому обрабатываем её здесь
      bpe[k+j++] = a[i] ^ 196 ? a[i] : a[++i]-128;
    }
    bpe[k+j++] = 0;
  } else if (i > 187) {
    // Токены выше 187 — это не выводимые на печать символы asii в диапазоне 0-32
    bpe[k] = i-188;
    bpe[k+1] = 0;
  }
}

Висновок

Справді, примітно, як можна сконцентрувати кілька тисяч байт цілі десятиліття у розвитку машинного навчання. По суті тут є все необхідне (крім фактичних ваг моделі), щоб ви могли запустити будь-яку сучасну нейронну мережу. При тому, що я реалізував цей проект здебільшого з інтересу, він добре демонструє, наскільки справдіпростінейронні мережі.