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 на этой библиотеке