Пишем свой PyTorch на NumPy. Финал. Запускаем GPT-2

PyTorch — это мощный и гибкий фреймворк для машинного обучения, широко используемый для создания нейронных сетей. Он особенно популярен благодаря простоте использования, динамическим вычислительным графам и богатой экосистеме инструментов для обучения моделей. Для использования этого фреймворка, часто достаточно поверхностно понимать работу алгоритмов машинного обучения.

Но Андрей Карпаты, известный исследователь в области ИИ, считает, что реализация алгоритмов с нуля позволяет понять их суть и детали работы, что сложно осознать, используя только готовые библиотеки. Это помогает развить интуицию для дальнейшего применения и улучшения методов. Андрей посвящает много собственного времени, чтобы объяснять ключевые принципы работы нейросетей в своих блогах и на своём ютуб-канале. Он также не раз подчеркивал, что на его курсе в Cтэнфорде есть задачи по реализации различных алгоритмов, например, обратное распространение.

Я хотел бы посвятить данную статью этой идеи, потому что мне самому особенно интересно копаться в алгоритмах глубокого обучения. Эта статья продолжение третьей статьи

Итак, как я и сказал в предыдущей статье, у нас достаточно знаний, чтобы собрать их в целую библиотеку. Так я и сделал, реализовав pycandle.
Там находиться всё то, что мы изучали на протяжении 3 частей. В этой части мы будем писать инференс код для GPT2 на собственной библиотеке!
Начнём с импорта библиотек

import torch import torch.nn as nn

Ой, не то. Я имел в виду

import candle import candle.nn as nn

Также воспользуемся токенизатором от OpenAI, который они используют в своих моделях. Сделаем все необходимые импорты

import tiktoken from candle import Tensor from dataclasses import dataclass

Определим конфигурацию нашей модели с помощью dataclasses

@dataclass class GPTConfig:     block_size: int = 1024 # размер окна вниманиә     vocab_size: int = 50257 # размер словаря BPE     n_layer: int = 12 # количество слоёв     n_head: int = 12 # количество голов в механизме внимания     n_embd: int = 768 # размерность вектора эмбеддингов

Определим линейный слой модели

class MLP(nn.Module):     def __init__(self, config):         super().__init__()         self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)         self.gelu = nn.GeLU()         self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)      def forward(self, x):         x = self.c_fc(x)         x = self.gelu(x)         return self.c_proj(x)

Обратите внимание тут candle.nn, а не torch.nn. С линейными слоями мы уже знакомы, а GeLU(), это что такое?

Исследователи из OpenAI, выяснили, что в их моделях такая функция работает лучше всего, посмотрим на её реализацию

class GeLU:     def __init__(self):         Parameter([self,[]])      def __call__(self, x):         return Tensor.gelu(x)         #return 0.5 * x * (1 + Tensor.tanh(0.79788456 * (x + 0.044715 * (x ** 3))))

Я закомментировал вторую строчку и вместо этого использую встроенный метод Tensor.gelu(), но никто не запрещает и честно считать по формуле. Заглянем в этот метод

from scipy.stats import norm @classmethod def gelu(cls, x): return x * norm.cdf(x.value) # формула из картинки

Теперь в линейном слое нам все понятно, идём дальше!

class CausalSelfAttetion(nn.Module):     def __init__(self, config):         super().__init__()         self.config = config         self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=True)         self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=True)      def forward(self, x):         B, T, C = x.shape         qkv = self.c_attn(x)         q, k, v = Tensor.split(qkv, 3, axis=2)         n_head = self.config.n_head         q = q.reshape(B, T, n_head, C // n_head).transpose(0, 2, 1, 3)         k = k.reshape(B, T, n_head, C // n_head).transpose(0, 2, 1, 3)         v = v.reshape(B, T, n_head, C // n_head).transpose(0, 2, 1, 3)         att = q @ k.transpose(0, -3, -1, -2) * (k.shape[-1] ** -0.5)         lg = att.local_gradients         att = Tensor.tril(att)         att[att == 0] = float('-inf')         att.local_gradients = lg         probs = Tensor.softmax(att, axis=-1)         y = probs @ v         y = y.transpose(0, 2, 1, 3).reshape(B, T, C)         y = self.c_proj(y)         return y

Слой внимания, я предполагаю, что уже видели подобную реализацию, поэтому много уделять внимания не буду. Видим, незнакомые нам методы Tensor.split(), Tensor.tril, они делают тоже самое, что и их аналоги в pytorch

@classmethod def tril(cls, input, diagonal=0): value = np.tril(input.value, k=diagonal) local_gradients = ( ('tril', input, lambda x: x * np.tril(np.ones_like(input.value), k=diagonal)), ) return cls(value, local_gradients=local_gradients)  @classmethod def split(cls, array, split_size_or_sections, axis=0): value = np.split(array.value, split_size_or_sections, axis=axis) return cls(value, requires_grad=False)

Обратите внимание, я не определил вычисление градиента для второй операции. Значит, либо мне придётся его определить, либо я просто не смогу обучать модели!

class Block(nn.Module):     def __init__(self, config):         super().__init__()         self.ln_1 = nn.LayerNorm(config.n_embd)         self.attn = CausalSelfAttetion(config)         self.ln_2 = nn.LayerNorm(config.n_embd)         self.mlp = MLP(config)      def forward(self, x):         x = x + self.attn(self.ln_1(x))         x = x + self.mlp(self.ln_2(x))         return x

Тут я использую готовый слой candle.nn.LayerNorm(), заглянем в него!

class LayerNorm:     def __init__(self, dim, eps=1e-5):         super().__init__()         self.param = None         self.gamma = Tensor.ones(dim, requires_grad=True)         self.beta = Tensor.zeros(dim, requires_grad=True)         self.eps = eps         self.all_layers = [self.gamma, self.beta]         self.grad = None      def __call__(self, x):         xmean = Tensor.mean(x, axis=2, keepdims=True)         xstd = Tensor.std(x, axis=2, keepdims=True)         x = (x - xmean) / (xstd + self.eps)         return self.gamma * x + self.beta

В целом ничего сложного, если вы идейно знакомы с методами нормализации. Тут я использую Tensor.mean(), Tensor.std(). Посмотрим на их реализацию

@classmethod def mean(cls, array, axis=None, keepdims=False): if axis == None: return Tensor.sum(array, axis=None, keepdims=keepdims) / np.size(array.value) else: delimeter = 1 if not isinstance(axis, int): for ax in axis: delimeter = delimeter * array.shape[ax] else: delimeter = array.shape[axis]  return Tensor.sum(array, axis=axis, keepdims=keepdims) / delimeter  @classmethod def std(cls, array, axis=None, keepdims=False):  if axis == None or axis == 0: mean = Tensor.mean(array, axis=axis, keepdims=False) sub = array - mean squared = sub ** 2 scaled_sum = Tensor.mean(squared, axis=axis, keepdims=keepdims) std = Tensor.sqrt(scaled_sum) return std  elif axis >= 1: mean = Tensor.mean(array, axis=axis, keepdims=True) sub = array - mean squared = sub ** 2 scaled_sum = Tensor.mean(squared, axis=axis, keepdims=keepdims) std = Tensor.sqrt(scaled_sum) out = Tensor(std, local_gradients=None) out.local_gradients = (('std', std, lambda x: x * array.shape[axis] / (array.shape[axis] - 1)),) # array.shape[axis] / (array.shape[axis] - 1) additional multiplier due to dissimilarity return out 

class GPT(nn.Module):     def __init__(self, config):         super().__init__()         self.config = config         self.wte = nn.Embedding(config.vocab_size, config.n_embd)         self.wpe = nn.Embedding(1024, config.n_embd)         self.h = nn.ModuleList([Block(config) for _ in range(config.n_layer)])         self.ln_f = nn.LayerNorm(config.n_embd)         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)      def forward(self, x):         B, T = x.shape         assert T >= self.config.block_size, f"Cannot forward the sequence of length {T}, block_size is smaller"         pos = Tensor.arange(T).reshape(1, -1)         pos_embd = self.wpe(pos)         tok_embd = self.wte(x)         x = pos_embd + tok_embd          for block in self.h:             x = block(x)         x = self.ln_f(x)         logits = self.lm_head(x)         return logits 

Разберемся как работают candle.nn.Embedding() и candle.nn.ModuleList().

class Embedding:      def __init__(self, num_emb, emb_dim):         self.w = Tensor.randn((num_emb, emb_dim), requires_grad=True)         self.w.local_gradients = None         self.num_embd = num_emb         self.emb_dim = emb_dim         self.param = None         self.grad = None         self.all_layers = [self.w]      def __call__(self, x):          if self.param:             global Parameter             Parameter = self.param                      def multiply_by_locgrad(path_value):             temp = np.zeros_like(self.w.value)             np.add.at(np.zeros_like(self.w.value), x.value, path_value)             return temp                      x.value = x.value.astype(int)         local_gradients = (('embd', self.w, multiply_by_locgrad),)         return Tensor(self.w.value[x.value], local_gradients=local_gradients)

В целом ничего сложного, просто ожидаем на входе тензор из целых значений и рассматриваем эти значения как индексы для матрицы, хранящей эмбеддинги.

class ModuleList:      def __init__(self, layers):         self.layers = layers         self.index = 0      def __call__(self, x):         for layer in self.layers:             x = layer(x)         return x      def __len__(self):         return len(self.layers)      def __iter__(self):         self.index = 0         return self      def __next__(self):         if self.index < len(self.layers):             result = self.layers[self.index]             self.index += 1             return result         else:             raise StopIteration      def __getitem__(self, index):         return self.layers[index]

Оказывается ModuleListэто просто итератор!
Note! Для корректной работы, нам нужно определить метод Module.__setattr__(), иначе слои которые находятся внутри ModuleList, просто не будут видны нашей глобальной переменной Parameter, я не буду на этом останавливаться, но при желании можно заглянуть в код и разобраться!

class GPT(nn.Module):    @classmethod     def from_pretrained(cls, model_type):         assert model_type in ('gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl')         config_args = {             'gpt2': dict(n_layer=12, n_head=12, n_embd=768),             'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024),             'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280),             'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600)         }[model_type]         config_args['vocab_size'] = 50257         config_args['block_size'] = 1024          config = GPTConfig(**config_args)         model = GPT(config)         model = GPT.get_params(model, model_type)         return model

Здесь мы определяем модель из библиотеки hf.transformers, это самый простой путь для доступа к весам GPT2.

@staticmethod def get_params(model, model_type):  from transformers import GPT2LMHeadModel, logging  logging.set_verbosity_error()  model_hf = GPT2LMHeadModel.from_pretrained(model_type) sd_hf = model_hf.state_dict() ...

Это часть кода, отвечающего за перенос весов с модели huggingface на нашу модель. Там ничего умного, кроме аккуратной работы.

Генерация

model = GPT.from_pretrained(model_type=model_type)  max_length = 100 number_of_examples = 3 starting_sentence = "Hello, I'm a language model,"  enc = tiktoken.get_encoding('gpt2') tokens = enc.encode(starting_sentence) tokens = Tensor(tokens, dtype=int) tokens = Tensor.unsqueeze(tokens, 0) tokens = Tensor.repeat(tokens, number_of_examples, 0) x = tokens topk = 50 # responsible for "creativity" or "adequacy"  for i in tqdm(range(max_length)): logits = model(x) logits = logits[:, -1, :] probs = Tensor.softmax(logits, axis=-1) topk_indices, topk_probs = Tensor.topk(probs, topk) new_token = Tensor.multinomial_from_array(topk_indices, topk_probs, num_samples=1).reshape(-1, 1) x = Tensor.cat([x, new_token], axis=1)  for sample in x:     sample = sample.value.astype(int).tolist()     print(enc.decode(sample), end='\n')     print('---------------') 

Hello, I'm a language model, not a coding model, but a compiler model.  I am thinking of what it means to be a developer and a language model  and to be able to write in some simple yet expressive way code that makes it  possible to use it with our software in any meaningful way to make  applications more readable, better, more readable. I can think of more words to talk about this. In this post, I'm trying to explain some of the main differences --------------- Hello, I'm a language model, based on my personal experience with code and development using C#.  This tutorial will help you create my own virtual language that you can use to implement things like an HTML page in your own app.  Now, what if I were going to learn how to program a website in an Objective-C program and create a web app out of it,  but didn't know how to do that and want to do some extra work to try it out? First I --------------- Hello, I'm a language model, I'm my own language. Well, a language model has two parts: a semantic construct and a structural construct. To understand the relationship with the semantic construct, let's see how all of the variables on a graph come together. Now I know that some of the variables on a graph represent numbers. But when you see the graphs, they're graphs. In other words, you see the connections and they are networks. We have an example

Вы можете поиграться с этой моделью на каггле или скачав к себе из репозитория smolGPT

Вот и подошла к концу моя мини-серия по созданию собственной библиотеки на NumPy.
Благодарю за внимание!

Надеюсь, вы нашли для себя этот цикл статей познавательным!

Если вам понравилось, пожалуйста, поделитесь ими, поставьте upvote и поставьте звёздочки моим реализациям на гитхабе.

Первая версия библиотеки

Вторая версия библиотеки

GPT-2 на этой библиотеке