Tiny Universe, часть 2: Qwen-blog — путь тензора через Qwen2

О переводе

Полный перевод файла content/Qwen-blog/readme.md проекта tiny-universe (CC BY-NC-SA 4.0 → эта страница тоже под CC BY-NC-SA 4.0). В оригинале почти все схемы и формулы — картинки; картинки мы не копируем, а формулы записали заново в стандартном виде и дали ссылки на рисунки. Код взят авторами из реализации Qwen2 в библиотеке Hugging Face Transformers (версии 2024 года, файл modeling_qwen2.py) и сокращён до главного. Наши пояснения отмечены как «Примечание переводчика». В авторской нумерации разделов есть сбои (два раздела «1.1.2») — мы их исправили.

Общее устройство Qwen

Эфир с разбором архитектуры Qwen: запись трансляции.

Общая архитектура Qwen похожа на Llama2 — см. рисунок (img/framework.JPG).

flowchart TB
T["Текст"] --> TOK["tokenizer: текст → id из словаря"]
TOK --> EMB["embedding: id → вектор"]
EMB --> L["N × Qwen2DecoderLayer<br/>(RMSNorm → Attention → +residual<br/>RMSNorm → MLP → +residual)"]
M["attention_mask"] --> L
L --> N["финальная RMSNorm"]
N --> H1["Linear → логиты следующего токена<br/>(CausalLM)"]
N --> H2["Linear → классы<br/>(SequenceClassification)"]

Схема переводчика по рисунку framework.JPG

Здесь:

• tokenizer превращает текст в числа — номера токенов в словаре.
• Числа проходят через embedding и превращаются во взаимно однозначно соответствующие им векторы.
• attention_mask задаёт, что модель «видит»: только левую часть, правую, обе стороны и т. п.
• Для разных прикладных задач — Causal (генерация), seqcls (классификация последовательностей) и других — к базовой модели model по сути просто приставляется свой слой Linear и своя функция потерь.

1. Qwen2Config

Qwen2Config хранит набор настраиваемых гиперпараметров: vocab_size, hidden_size, num_hidden_layers, num_attention_heads и т. д. К ним можно обращаться как к полям словаря: config.pad_token_id.

1.1 Qwen2Model

1.1.1 Инициализация

• Задаются два атрибута модели: padding_idx (индекс токена-заполнителя) и vocab_size (размер словаря).
• Инициализируются слой эмбеддингов, слои декодера и слой нормализации.
• Слой эмбеддингов (nn.Embedding): модель отображает входные токены в плотные векторные представления.
• Слои декодера (nn.ModuleList()): модель состоит из нескольких слоёв декодера, каждый задан классом Qwen2DecoderLayer.
• Слой нормализации Qwen2RMSNorm: используется Root Mean Square Layer Normalization (RMSNorm).
• Задаётся, использовать ли gradient_checkpoint — в основном для экономии видеопамяти.
• Вызывается post_init() — код, который завершает инициализацию и делает подготовительные проверки.

class Qwen2Model(Qwen2PreTrainedModel):
    def __init__(self, config: Qwen2Config):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        self.gradient_checkpointing = False
        # Инициализируем веса и делаем финальную обработку
        self.post_init()

Функция post_init в основном инициализирует параметры и настраивает контрольные точки градиента (gradient checkpointing):

def post_init(self):
    """
    Метод, который выполняется в конце инициализации каждой модели Transformer, — для кода, которому нужны
    уже правильно созданные модули модели (например, для инициализации весов).
    """
    self.init_weights()
    self._backward_compatibility_gradient_checkpointing()

Примечание переводчика: что такое gradient checkpointing

При обучении PyTorch хранит промежуточные активации всех слоёв, чтобы потом посчитать градиенты. Gradient checkpointing (в книгах — «пересчёт активаций») хранит только часть из них, а остальное пересчитывает на обратном проходе: памяти нужно меньше, времени — примерно на треть больше. Подробнее — в уроке 5.4.3.

1.1.2 Forward

Разберём только главный путь:

inputs_embeds = self.embed_tokens(input_ids)
# эмбеддинги позиций здесь не добавляются — позиция войдёт через RoPE внутри внимания
hidden_states = inputs_embeds

for idx, decoder_layer in enumerate(self.layers):
    # складываем все hidden_states в кортеж
    if output_hidden_states:
        all_hidden_states += (hidden_states,)
    # передаём hs в очередной слой декодера
    layer_outputs = decoder_layer(
        hidden_states,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_value=past_key_value,
        output_attentions=output_attentions,
        use_cache=use_cache,
    )
    # берём hs — выход этого слоя — и подаём в следующий слой;
    # нужен только первый элемент, второй — объект кэша
    hidden_states = layer_outputs[0]

# нормализуем hidden_states после последнего слоя
hidden_states = self.norm(hidden_states)

# добавляем hidden_states последнего слоя
if output_hidden_states:
    all_hidden_states += (hidden_states,)

• Если включено сохранение output_hidden_states, то в кортеже окажутся: первым — эмбеддинги input_ids, затем выходы hs слоёв декодера до n-1-го включительно и, наконец, выход последнего слоя после norm.
• Возвращается всё в виде объекта BaseModelOutputWithPast.

Примечание переводчика

В оригинале комментарий к строке hidden_states = inputs_embeds звучит как «embed positions» — это остаток от GPT-2, где здесь прибавлялись позиционные эмбеддинги. В Qwen2 (как и в LLaMA) отдельного слоя позиций нет: позиция кодируется вращением векторов $q$ и $k$ (RoPE, раздел 1.3.3.2). Мы поправили комментарий.

1.2 Qwen2DecoderLayer

Рисунок в оригинале: img/decoderlayer.png — блок декодера: RMSNorm → самовнимание → сложение с остатком → RMSNorm → MLP → сложение с остатком.

1.2.1 Инициализация

«Три кита» слоя: attn + MLP + norm.

QWEN2_ATTENTION_CLASSES = {
    "eager": Qwen2Attention,  # обычный случай
    "flash_attention_2": Qwen2FlashAttention2,
    "sdpa": Qwen2SdpaAttention,
}

class Qwen2DecoderLayer(nn.Module):
    def __init__(self, config: Qwen2Config, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = QWEN2_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)

        self.mlp = Qwen2MLP(config)
        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

input_layernorm и post_attention_layernorm устроены одинаково — различается только место, где они применяются.

Примечание переводчика

В оригинале у __init__ пропущен аргумент layer_idx, хотя он используется в теле, — мы его добавили, как в настоящем modeling_qwen2.py. Три реализации внимания дают одинаковый результат и различаются скоростью: eager — «ручные» матричные умножения (их и разбираем ниже), sdpa — встроенная функция PyTorch scaled_dot_product_attention, flash_attention_2 — FlashAttention из урока 5.7.2.

1.2.2 Forward

Читать удобнее вместе с рисунком:

• Сначала копируем hidden_states в residual, затем пропускаем hidden_states через Norm и модуль attn.
• Получив выход attn, прибавляем residual; снова сохраняем копию в residual, пропускаем hidden_states через Norm и mlp и опять складываем с residual. Результат — и есть выходной hidden_states.

residual = hidden_states
# нормализуем и подаём во внимание
hidden_states = self.input_layernorm(hidden_states)  # нормализация RMSNorm
# Self Attention
hidden_states, self_attn_weights, present_key_value = self.self_attn(
    hidden_states=hidden_states,
    attention_mask=attention_mask,
    position_ids=position_ids,
    past_key_value=past_key_value,
    output_attentions=output_attentions,
    use_cache=use_cache,
    **kwargs,
)

# складываем остаток с новыми hidden_states
hidden_states = residual + hidden_states

# Полносвязная часть
residual = hidden_states
# та же нормализация RMSNorm
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states

outputs = (hidden_states,)

return outputs

Примечание переводчика

Это схема Pre-Norm: нормализация стоит до подслоя, а остаточная связь обходит и нормализацию, и подслой. В исходном Transformer 2017 года было наоборот (Post-Norm). Почему Pre-Norm устойчивее при глубоких сетях — в уроках 3.4.2 и 4.5.3.

1.3 Qwen2Attention

Рисунок в оригинале: img/Qwen2Attention.png — схема модуля внимания: проекции Q/K/V, RoPE, повтор K/V по группам, скалярное произведение, маска, softmax, умножение на V, выходная проекция.

1.3.1 Инициализация

Главные параметры:

• num_key_value_heads — число голов для ключей и значений;
• num_key_value_groups — число групп ключей-значений, считается как num_heads // num_key_value_heads. Это и есть реализация GQA!
• q_proj, k_proj, v_proj, o_proj — четыре слоя Linear. Именно их обычно и «режет» LoRA при дообучении.

class Qwen2Attention(nn.Module):
    """Многоголовое внимание из статьи 'Attention Is All You Need'"""

    def __init__(self, config: Qwen2Config, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.max_position_embeddings = config.max_position_embeddings
        self.rope_theta = config.rope_theta
        self.is_causal = True
        self.attention_dropout = config.attention_dropout

        if (self.head_dim * self.num_heads) != self.hidden_size:
            raise ValueError(
                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                f" and `num_heads`: {self.num_heads})."
            )
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)

        self.rotary_emb = Qwen2RotaryEmbedding(
            self.head_dim,
            max_position_embeddings=self.max_position_embeddings,
            base=self.rope_theta,
        )

• Смысл параметров из config смотрите в описании Qwen2Config:

max_position_embeddings (`int`, *optional*, defaults to 32768):
            Максимальная длина последовательности, с которой модель вообще может работать.

rope_theta (`float`, *optional*, defaults to 10000.0):
            Базовый период эмбеддингов RoPE.

Примечание переводчика: числа для настоящей модели

У Qwen2-7B, например, hidden_size = 3584, num_attention_heads = 28, num_key_value_heads = 4 (значения из её config.json на Hugging Face). Тогда head_dim = 3584 / 28 = 128, а num_key_value_groups = 28 / 4 = 7: семь голов запросов делят одну пару K/V. KV-кэш при этом в 7 раз меньше, чем у обычного многоголового внимания.

1.3.2 Forward

• Сначала пропускаем hidden_states через слои Linear и получаем query, key и value.
• Применяем вращательное позиционное кодирование rotary_emb: берём его косинусную и синусную части, умножаем на query и key и складываем результаты — так и получается эффект вращательного кодирования.
• Повторяем key_states и value_states group раз и выполняем скалярное внимание (dot attn).
• После скалярного произведения получаем attn_weights, прибавляем attention_mask — так скрываются позиции, которые видеть нельзя, — затем применяем softmax и умножаем на value_states. Получаем attn_output.
• Меняем форму attn_output (reshape), подаём в o_proj и получаем окончательный выход.

# Размеры: на входе hidden_states формы (bs, T, hd)
bsz, q_len, _ = hidden_states.size()

# Линейными слоями получаем query, key и value
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)

# Делим на головы — (bs, T, heads, hd_d), затем переставляем оси — (bs, heads, T, hd_d)
query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

# Применяем вращательное кодирование к запросам и ключам: умножаем их на косинусную и синусную части
# и складываем результаты — так получается вращение
cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

# Повторяем key_states и value_states num_key_value_groups раз
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)

# Скалярное внимание: q·kᵀ / sqrt(hd_d)
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

# Прибавляем attention_mask — задаём порядок «чтения»
attn_weights = attn_weights + attention_mask

# softmax + dropout + умножение на value_states
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
attn_output = torch.matmul(attn_weights, value_states)

# Переставляем оси обратно и меняем форму
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

# Последний шаг — o_proj
attn_output = self.o_proj(attn_output)

# Возвращаем результат
return attn_output, attn_weights, past_key_value

Примечание переводчика

Фрагмент сокращён авторами: переменная kv_seq_len вычисляется в полном коде (длина текущей последовательности плюс длина уже накопленного KV-кэша past_key_value), а обновление кэша между RoPE и repeat_kv опущено. Формула $\mathrm{softmax}(QK^\top/\sqrt{d})V$ — та же, что в уроке 4.4.2; softmax считается в float32 даже при обучении в bfloat16 — ради численной устойчивости.

1.3.3 Детали под отладчиком

1.3.3.1 GQA

Рисунок в оригинале: img/GQA.png — сравнение MHA, GQA и MQA: у MHA своя пара K/V на каждую голову запроса, у GQA одна пара K/V на группу голов, у MQA одна пара на все головы.

Главная мысль: при GQA и MQA во время инференса не нужно хранить столько KV-кэша, значит, кэш занимает меньше видеопамяти, и LLM-сервис может одновременно обслужить больше запросов.

1. Задаём исходные тензоры:

import math
import torch

## форма: (batch, seq_len, head, head_dim)
query = torch.randn(10, 128, 8, 128)
key = torch.randn(10, 128, 2, 128)
value = torch.randn(10, 128, 2, 128)

## здесь число групп равно 4
groups = query.shape[-2] // key.shape[-2]
head_dim = query.shape[-1]  # 128 (добавлено переводчиком: переменная нужна ниже)

2. Затем расширяем key и value. В GQA у key и value голов в group раз меньше, чем у query. Чтобы потом было удобно перемножать матрицы, мы сначала с помощью expand расширяем головы key и value до той же размерности, что у query.

# На входе x; n_rep — сколько раз повторить, обычно это число групп
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:

    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    # повторять не нужно — это обычное многоголовое внимание
    if n_rep == 1:
        return hidden_states
    # сначала расширяем x до (bs, head, group, seq_len, head_dim)
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    # reshape превращает head -> head * group
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

3. Перемножаем матрицы и получаем score и output. Дальше — обычное перемножение q, k, v:

# (bs, head, seq_len, head_dim)
query = query.transpose(1, 2)
key = repeat_kv(key.transpose(1, 2), 4)
value = repeat_kv(value.transpose(1, 2), 4)
scores = torch.matmul(query, key.transpose(2, 3)) / math.sqrt(head_dim)
scores = torch.nn.functional.softmax(scores, dim=-1)

out = torch.matmul(scores, value)
# выше мы переставили оси — переставляем обратно
out = out.transpose(1, 2)
print(out.shape)  # torch.Size([10, 128, 8, 128]) (строка добавлена переводчиком)

Дополнение. Почему сначала expand, а потом reshape, а не сразу встроенный repeat тензора?

• expand расширяет тензор, не выделяя новой памяти: возвращённый тензор делит данные с исходным.
• repeat расширяет тензор, по-настоящему копируя данные: новый тензор не делит данные с исходным и занимает больше памяти.

Примечание переводчика

Строго говоря, последующий reshape расширенного тензора всё равно создаёт копию (данные уже не лежат в памяти подряд). Выигрыш expand в том, что копия делается один раз и сразу в нужной форме, без промежуточного тензора. В новых версиях Transformers и в torch.nn.functional.scaled_dot_product_attention (аргумент enable_gqa) повтор K/V вообще не материализуется.

1.3.3.2 apply_rotary_pos_emb

Позиционное кодирование придаёт каждой координате (dim) каждого токена свою информацию о позиции.

Определение (в оригинале — рисунок img/ROPE1.png; ниже — стандартная запись из статьи RoFormer):

$$q_t = R_{\Theta,t}\,W_q x_t,\qquad k_s = R_{\Theta,s}\,W_k x_s,\qquad A_{t,s} = q_t^\top k_s = x_t^\top W_q^\top R_{\Theta,\,s-t}\,W_k x_s .$$

Словами: запрос токена на позиции $t$ поворачивается на угол, пропорциональный $t$, ключ на позиции $s$ — на угол, пропорциональный $s$; в их скалярное произведение входит только разность углов, то есть относительное расстояние $s-t$.

Идея: вращательное кодирование даёт каждому токену и относительную, и абсолютную позицию.

• Его можно понимать как поправку к матрице самовнимания $A_{t,s}$, которая напрямую отражает относительное положение двух токенов, — и в то же время разложить на векторы $q_t$ и $k_s$, где кодируется абсолютная позиция каждого токена.
• По сути RoPE — это вращение векторов признаков. Возьмём для примера двумерный вектор признаков: для пары соседних координат используется один и тот же угол $\theta$, а вращение задаётся так (в оригинале — img/ROPE2.png):

$$\begin{pmatrix} q'_0 \\ q'_1 \end{pmatrix} = \begin{pmatrix} \cos t\theta & -\sin t\theta \\ \sin t\theta & \cos t\theta \end{pmatrix} \begin{pmatrix} q_0 \\ q_1 \end{pmatrix} = \begin{pmatrix} q_0 \\ q_1 \end{pmatrix}\cos t\theta + \begin{pmatrix} -q_1 \\ q_0 \end{pmatrix}\sin t\theta .$$

Отсюда видно главное: результат поворота $q_t$ (или $k_s$) — это сам вектор, умноженный на косинус, плюс вектор с переставленными координатами, у одной из которых сменён знак, умноженный на синус.

• Для многомерных векторов чередовать чётные и нечётные координаты парами неудобно, поэтому сейчас реализацию упрощают: делят координаты пополам и образуют пары «$i$-я координата первой половины ↔ $i$-я координата второй половины» (рисунок img/ROPE3.png). Операция над первой и второй половинами и есть rotate_half.

Реализация в коде. Сначала задаём углы поворота:

class Qwen2RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        # начальные значения
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        # частоты поворота
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Строим кэш здесь, чтобы работал `torch.jit.trace`.
        self._set_cos_sin_cache(
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
        )
    # для каждого токена последовательности получаем свой уникальный набор углов (внешнее произведение)
    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)

        freqs = torch.outer(t, self.inv_freq)
        # получаем углы (через механизм буферов создаём self.cos_cached и self.sin_cached)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype),
            self.sin_cached[:seq_len].to(dtype=x.dtype),
        )

Сначала вычисляем базовые углы:

$$\theta_n = \frac{1}{10000^{2n/d}},$$

где $n$ — номер пары координат, принимает значения $0, 1, \dots, d/2-1$ (рисунок img/ROPE4.png).

Затем умножаем эти углы на номер каждой позиции, чтобы различать слова последовательности. Это равносильно

$$\theta_{i,n} = \frac{i}{10000^{2n/d}},$$

где $i$ — номер строки, то есть позиции (рисунок img/ROPE5.png).

emb склеивает две копии — получается размерность dim, причём значения повторяются с периодом dim/2 (рисунок img/ROPE6.png).

Когда нужно достать позиционную информацию cos и sin, из таблицы вырезается часть длиной seq: заранее задано 1024 — максимальная позиция, а при использовании берётся только текущая seq_len. Внешнее произведение выше нужно как раз для того, чтобы у каждого слова последовательности был свой из 1024 вариантов позиционного кода (рисунок img/ROPE7.png).

Примечание переводчика: посчитаем руками

Пусть $d = 4$ (две пары координат). Тогда $\theta_0 = 1/10000^{0} = 1$, $\theta_1 = 1/10000^{2/4} = 0{,}01$. Токен на позиции $i = 3$ поворачивает первую пару на 3 радиана (почти полоборота), а вторую — на 0,03 радиана (почти не поворачивает). Быстрые пары различают соседние позиции, медленные — далёкие. Таблица emb для позиции 3 выглядит так: [3, 0.03, 3, 0.03] — это и есть «две копии с периодом dim/2». И ещё: «1024» в тексте — пример авторов; в коде максимальная длина задаётся max_position_embeddings (по умолчанию 2048 в этом классе и 32768 в конфиге Qwen2).

Применяем вращательное кодирование:

# Вторая половина координат меняется местами с первой, и у неё меняется знак.
def rotate_half(x):
    """Поворачивает половину скрытых координат входа."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    """Применяет вращательное позиционное кодирование (RoPE) к тензорам запросов и ключей.

    Возвращает query и key, повёрнутые с помощью RoPE.
    """
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

Соответствующая формула (в оригинале — img/ROPE8.png), для вектора $q$ размерности $d$ на позиции $t$, где $\theta_j$ повторяются с периодом $d/2$:

$$q' = q \odot \cos(t\theta) + \mathrm{rotate\_half}(q) \odot \sin(t\theta),\qquad \mathrm{rotate\_half}(q) = \left(-q_{d/2}, \dots, -q_{d-1},\; q_0, \dots, q_{d/2-1}\right).$$

Здесь индекс $t$ означает, что речь идёт об одной строке, то есть об одной позиции в seq_len; для двух элементов, образующих пару… (на этом месте фраза в оригинале обрывается).

1.3.3.3 attention_mask: порядок чтения

На первом шаге модель видит только себя, на втором — позиции 0 и 1, и т. д.; всё остальное — минус бесконечность (рисунок img/Mask1.png).

После softmax веса на позициях с минус бесконечностью становятся нулевыми — так и получается чтение только слева направо (рисунок img/Mask2.png).

Примечание переводчика: маска на числах

Для трёх токенов маска и веса выглядят так (допустим, скоры $q\cdot k$ после деления на $\sqrt{d}$ в строке 3 равны $1, 2, 0$):

$$M = \begin{pmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{pmatrix},\qquad \mathrm{softmax}(1, 2, 0) \approx (0{,}24;\ 0{,}67;\ 0{,}09).$$

В первой строке после маски остаётся одно число, и softmax даёт вес 1 самому первому токену.

• attn имеет форму (bs, heads, T, T), values — (bs, heads, T, hd); в итоге всё сводится к умножению матриц [30, 30] × [30, 128], где 30 — длина фразы в шагах, то есть число слов.
• Каждое слово в value описывается 128 координатами. Для первого слова, поскольку attn нижнетреугольная, каждая координата умножается только на первый ненулевой элемент — единицу, а всё остальное умножается на 0.
  • Во второй строке уже два разных веса: каждая из 128 координат value берётся по этим двум весам из первых двух строк — так вычисляются 128 координат второго слова выхода (второго шага, второй строки). Такое взвешивание и выражает учёт связи с предыдущими словами.
• На шаге $n$ соответственно $n$ весов: они описывают связи между шагами с 1-го по $n$-й, и по ним вычисляется каждая координата.
• Результат каждого умножения матриц — значение координаты для следующего слова; scores по маске определяют, по каким предыдущим словам можно взвешивать значения этой координаты, чтобы получить её значение у следующего слова (рисунок img/Mask3.png).
• При инференсе запросы разной длины не проблема: все веса модели имеют размер dim × dim, а матрица attention score — seq × seq, так что веса от длины последовательности не зависят. Они относятся к координатам dim.
• При инференсе attention_mask можно и не задавать: токены выдаются по одному, каждый новый дописывается (cat) к входу следующего шага, и каждый раз берётся последний выход — это и есть предсказание следующего токена.

Примечание переводчика

Последний пункт верен для генерации с KV-кэшем, когда на каждом шаге в модель подаётся один новый токен: ему и так можно смотреть на всё прошлое. Но на этапе префилла (обработка всей подсказки за один проход) и при батче из подсказок разной длины маска нужна — и причинная, и маска заполнителей. Подробнее — в уроке 7.1.1 «Префилл и декодирование».

1.4 Qwen2 MLP

Рисунок в оригинале: img/MLP1.png.

Входной hidden_state параллельно подаётся в два слоя Linear; выход одного пропускается через функцию активации и умножается на выход другого, а результат проходит ещё через один Linear — это и есть итоговый выход.

class Qwen2MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        # эти двое в пояснениях не нуждаются
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size

        # три полносвязных слоя
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj

Примечание переводчика

У Qwen2 hidden_act = "silu", так что это SwiGLU: $\mathrm{MLP}(x) = W_{\text{down}}\big(\mathrm{SiLU}(W_{\text{gate}}x) \odot W_{\text{up}}x\big)$ — та же формула, что в уроке 4.5.3.

1.5 Qwen2RMSNorm

Формула (в оригинале — рисунок img/RMSNorm_formulation.jpg):

$$\mathrm{RMSNorm}(x)_i = w_i \cdot \frac{x_i}{\sqrt{\dfrac{1}{n}\sum_{j=1}^{n} x_j^2 + \epsilon}}$$

Здесь:

• $x$ — входной hidden_state слоя;
• $w_i$ — значение, относящееся к последнему измерению hidden_state;
• $n$ — число элементов в последнем измерении входа;
• $\epsilon$ — очень маленькое число, чтобы не делить на ноль.

class Qwen2RMSNorm(nn.Module):  # слой нормализации
    def __init__(self, hidden_size, eps=1e-6):
        """
        Qwen2RMSNorm эквивалентна T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

• torch.rsqrt — это «производная от корня из входа».
• .pow(2).mean(-1, keepdim=True) — возвести в квадрат и усреднить по последнему измерению.

Примечание переводчика: две неточности оригинала

1. $w_i$ — это не «значение последнего измерения», а обучаемый вес (self.weight, вначале все единицы), по одному на каждую координату: он позволяет модели заново растянуть координаты после нормализации.
2. torch.rsqrt(x) — это обратная величина корня, $1/\sqrt{x}$, а не производная. В китайском тексте, видимо, опечатка: по-китайски «обратная величина» и «производная» звучат почти одинаково (оба слова — «даошу», разница только в тоне).

Пример: $x = (3, 4)$, $\epsilon \approx 0$. Среднее квадратов $= (9 + 16)/2 = 12{,}5$, корень $\approx 3{,}54$, результат $\approx (0{,}85;\ 1{,}13)$ — при весах $w = (1, 1)$.

От переводчика: что сделать руками

1. Откройте modeling_qwen2.py в установленной у вас библиотеке transformers (python -c "import transformers, os; print(os.path.dirname(transformers.__file__))", далее models/qwen2/modeling_qwen2.py) и найдите все пять классов из статьи. В новых версиях код отрефакторен (eager_attention_forward, класс кэша), но формулы те же.
2. Загрузите маленькую модель Qwen (например, Qwen2.5-0.5B-Instruct), поставьте точку останова в Qwen2Attention.forward и проследите формы тензоров, как делают авторы в разделе «Детали под отладчиком».
3. Посчитайте, во сколько раз GQA сокращает KV-кэш у Qwen2.5-0.5B: возьмите num_attention_heads и num_key_value_heads из её config.json. Сверьтесь с уроком 5.2.1.

Словарик

Русский	Пояснение
слой эмбеддингов (embedding)	id токена → вектор
слой (блок) декодера (decoder layer)	внимание + MLP + нормализации
остаточная связь (residual)	x + f(x)
нормализация RMSNorm	делим на среднеквадратичное
вращательное позиционное кодирование (RoPE)	позиция = угол поворота $q$ и $k$
внимание с групповыми запросами (GQA)	одна пара K/V на группу голов
KV-кэш (KV cache)	сохранённые ключи и значения прошлых токенов
маска внимания (attention mask)	запрещает смотреть «в будущее»
пересчёт активаций (gradient checkpointing)	экономия памяти за счёт времени
гейт, «вентиль» (gate_proj)	ветвь MLP, которая умножается на другую