Привет Хабр! В этой статье я попробую немного разобрать код LLM Llama 3. Полностью проанализировать каждую строку кода не получится, но самые важные и базовые концепции мы все-таки разберем насколько это возможно.

Падаем в кроличью нору

Изучать мы будем класс Llama (файл generation.py) и его метод text_completion. На рисунке у нас три больших блока, каждый из которых состоит из более маленьких частей. Подробно рассмотреть все нюансы модели будет довольно трудно, поэтому сосредоточимся на наиболее важных моментах. Начнем мы с входного блока.

Первое, что нужно сделать — это преобразовать входной текст в токены. Много всего было сказано про токены, лишний раз повторяться не буду. Скажу лишь, что токен – это почти слово. В реальности при определенных обстоятельствах одно слово может разбиться на несколько токенов. Для LLM токен – это своего рода стандартная единица измерения текста. В llama 3 используется готовый токенайзер (или токенизатор, кому как угодно) Tiktoken от OpenAI. Принцип работы токенайзера довольно прост: представьте, что у вас есть большой словарь, где все слова пронумерованы. Задача токенайзера сначала разбить входной текст на токены, а потом найти в словаре соответствующую позицию этого токена (токен ID). Эту задачу выполняет функция encode класса Tokenizer (файл tokenizer.py). Результатом работы функции будет список id наших токенов.

В методе text_completion строка 261:

prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]

Это была только часть блока Input, где свою работу выполнил Tokenizer. После получения списка id токенов нужно выполнить так называемую векторизацию – LLM модель оперирует не текстом, а его числовым представлением. Для этого необходимо строку текста представить в виде так называемого вектора – многомерного массива чисел. Такой вектор в пространстве из нескольких измерений обладает рядом характеристик. Чтобы совсем не запутаться проще всего представить каждое измерение, как отдельный признак: различие между частями речи, формами слов и т.д. Операцию преобразование токенов в вектор выполняет эмбеддинг (Embedding).

Пришло время посмотреть на файл model.py. Там сосредоточен код транформера (класс Transformer). Это сердце нашей LLM модели, на схеме он состоит из ряда блоков. К нему мы вернемся чуть позже, а пока посмотрим на строки 258-260:

self.tok_embeddings = VocabParallelEmbedding(
  params.vocab_size, params.dim, init_method=lambda x: x
)

Здесь мы создаем на эмеддинг VocabParallelEmbedding. Теперь посмотрим на класс Transformer – это основной блок архитектуры llama 3, который выполняет основную работу модели. Рассмотрим метод forward строка 280:

h = self.tok_embeddings(tokens)

Здесь и выполняется процедура эмбеддинга. Сами векторы будут созданы чуть раньше в классе Llama метод generate 159-161 строки:

        tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long, device="cuda")
        for k, t in enumerate(prompt_tokens):
            tokens[k, : len(t)] = torch.tensor(t, dtype=torch.long, device="cuda")

Тут стоит обратить внимание на пару моментов – размер сформированного вектора равен 4096. То есть мы проецируем наш токен в 4096-мерное пространство. Звучит безумно, но это факт – для дальнейших вычислений мы должны рассматривать наш токен, как вектор в пространстве с большим количеством измерений. И что самое интересное по ходу вычислений количество измерений в векторе может меняться как в большую, так и в меньшую сторону.

Давайте посмотрим на класс ModelArgs – это параметры модели:

@dataclass
class ModelArgs:
    dim: int = 4096
    n_layers: int = 32
    n_heads: int = 32
    n_kv_heads: Optional[int] = None
    vocab_size: int = -1
    multiple_of: int = 256  # make SwiGLU hidden layer size multiple of large power of 2
    ffn_dim_multiplier: Optional[float] = None
    norm_eps: float = 1e-5
    rope_theta: float = 500000

    max_batch_size: int = 32
    max_seq_len: int = 2048

Количество измерений в dim равно 4096, количество слоев n_layers 32. Это еще один важный момент архитектуры llama 3 – на схеме трансформер состоит из 32 блоков TransfomerBlock. Собственно говоря, это оно и есть - n_layers определяет количество блоков TransfomerBlock. Без преувеличения могу сказать, что TransfomerBlock – это самая важная часть всей архитектуры модели. Выход каждого блока подается на следующий. Давайте посмотрим на TransfomerBlock чуть повнимательней:

class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.head_dim = args.dim // args.n_heads
        self.attention = Attention(args)
        self.feed_forward = FeedForward(
            dim=args.dim,
            hidden_dim=4 * args.dim,
            multiple_of=args.multiple_of,
            ffn_dim_multiplier=args.ffn_dim_multiplier,
        )
        self.layer_id = layer_id
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
        out = h + self.feed_forward(self.ffn_norm(h))
        return out

Довольно лаконично. На самом деле тут сконцентрирована большая часть логики модели.

В 228 строке создается объект Attention:

self.attention = Attention(args)

Давайте представим, что мы подали на вход модели текст «Java – это объектно-ориентированный язык программирования». К чему относится слово «Java»? Может к «язык» или к «программирования»? Как понять связь между токенами в последовательности? Чтобы выявить эти связи и понять на, что стоит обратить внимание при анализе текста был придуман механизм self-attention.

Посмотрим на класс Attention:

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

        self.wq = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wk = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wv = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wo = RowParallelLinear(
            args.n_heads * self.head_dim,
            args.dim,
            bias=False,
            input_is_parallel=True,
            init_method=lambda x: x,
        )

        self.cache_k = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()

Параметр n_heads указывает на количество так называемых «голов» - механизм self-attention работает параллельно, что существенно повышает производительность модели. Количество таких голов рано 32. Суть блока self-attention – вычисление баллов, которые указывают степень внимания, которое стоит уделить остальным частям текста при разборе одного слова. Звучит запутанно, но если вспомнить пример с Java, то возможно станет немного проще. Другими словами, можно описать этот процесс, как выделение самых важные смысловых частей исходного текста.

Для этого исходный вектор преобразуется в сразу в три: вектор запросов (Query), вектор ключей (Key) и вектор значений (Value). Нам надо получить балл для нашего токена. Вычисление балла происходит путем умножения полученных векторов на специальные матрицы query, key и value. Эти матрицы формируются в процессе длительного и ресурсоемкого обучения модели.

Метод forward выполняет логику вычисления баллов – для полученного от эмбеддинга токена в позиции 1 вычисляется произведение query1 на key1. Это первая оценка. Вторая оценка query1 на query 2 и т.д. Для оптимизации вычислений активно используется кэш (KV-cache).

        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)

        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

Далее баллы проходят через операцию softmax:

scores = F.softmax(scores.float(), dim=-1).type_as(xq)

Softmax «нормализует» баллы – делает так, чтобы все они были положительными и в сумме давали 1. Результат softmax и будет определять степень внимания, которое стоит уделить слову в заданной позиции.

На последнем этапе умножаем баллы на вектор values:

output = torch.matmul(scores, values)

Тут мы отсеиваем нерелевантные слова, то есть умножаем их на очень маленькие значения близкие к 0, чтобы сосредоточить максимум внимания на словах с большим баллом.

Каждый раз, когда мы подаем данные на вход self-attention, выполняется процедура нормализации:

h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)

Логика нормализации реализована в классе RMSNorm. Это еще один элемент оптимизации работы LLM. Представьте, что вы в процессе обучения мы использовали набор красных яблок, а в процессе использования еще и зеленые.

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

Процедура RMSNorm будет выполняться как при входе в блок self-attention, так и на его выходе, то есть при в ходе в блок FeedForward. Пришло время его рассмотреть.

class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
        ffn_dim_multiplier: Optional[float],
    ):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        # custom dim factor multiplier
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )
        self.w2 = RowParallelLinear(
            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
        )
        self.w3 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

Классическая нейронная сеть с прямой связью. Если в self-attention мы рассматривали токены в контексте других токенов, то есть учитывая связи между ними, то в случае FeedForward каждый токен рассматривается по отдельности в независимости от других токенов. В FeedForward происходит доуточнение смысла каждого токена. Вернемся к примеру «Java – это объектно-ориентированный язык программирования». Какой вообще контекст у этого предложения – может это «Java» или может быть это «язык»? На этот вопрос – поиск основного контекста исходного текста – ответит блок FeedForward.

Когда мы получили выходной вектор чисел после всех 32 блоков нам надо превратить его в слово. Тут в дело вступает Linear layer – простая связанная нейронная сеть, которая преобразует полученный вектор в вектор с большим пространством – вектор логитов. Да, количество измерений существенно вырастит. Посмотрим на 179 строку класса Llama:

        for cur_pos in range(min_prompt_len, total_len):
            logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
            if temperature > 0:
                probs = torch.softmax(logits[:, -1] / temperature, dim=-1)
                next_token = sample_top_p(probs, top_p)
            else:
                next_token = torch.argmax(logits[:, -1], dim=-1)

Представим, что наша модель содержит 4500 слов, что довольно скромно, но для примера сгодится. Вектор логитов будет иметь ширину 4500 ячеек – каждая ячейка соответствует баллу токена.

Операция softmax все также преобразует все оценки в вероятности. После чего выбирается ячейка с самой высокой вероятностью и соответствующее ей слово подается на выход.

За скобками остался процесс обучения – это мы разберем в следующей статье. Буду рад любым комментариям и не забудьте подписаться на мой телеграм-канал.