DeepSeek V3 в ~200 строках чистого JAX: анатомия легковесно…

Введение: Вызов монолитным архитектурам и магия JAX

Когда на продакшн-контуре требуется развернуть экспериментальную архитектуру, а стандартные фреймворки тянут за собой гигабайты зависимостей и вызывают спонтанные утечки памяти, разработчик начинает искать более прозрачные инструменты. Зачем тратить часы на отладку тяжеловесных графов (и молиться, чтобы всё не упало на `OOM Killed`), если можно заглянуть под капот передовых моделей и собрать их на чистой математике? Сегодня мы разберем именно такой сценарий: превратим громоздкую концепцию в компактный и понятный инструмент.

В мире современной генеративной искусственной интуиции доминируют гигантские фреймворки, скрывающие под капотом терабайты зависимостей, сложнейшие графы вычислений и многоуровневые абстракции. Когда разработчики слышат про архитектуру уровня DeepSeek V3, перед глазами сразу встают кластеры из тысяч GPU, сложнейшая балансировка нагрузки и распределенный код на тысячах строк PyTorch или Triton. Однако концептуальная чистота современных нейросетевых архитектур позволяет взглянуть на них под совершенно другим углом.

Что если взять ключевые инновации DeepSeek V3 — такие как Mixture of Experts (MoE) с глубоким разделением экспертов, Multi-head Latent Attention (MLA) и оптимизированные механизмы маршрутизации — и реализовать их в пределах всего двухсот строк кода? Использовать для этого не тяжеловесные экосистемы, а чистый, функциональный и математически прозрачный JAX. Это не просто упражнение в минимализме, а способ понять саму физику современных языковых моделей, убрав весь процедурный шум.

В этой статье мы подробно разберем, как спроектировать и написать компактную, но функциональную реплику DeepSeek V3 на чистом JAX. Мы пройдем путь от тензорных операций до реализации хитрой маршрутизации токенов, доказывая, что сложнейшие архитектурные паттерны можно выразить элегантно и лаконично.

Архитектурные основы DeepSeek V3: что мы сжимаем до 200 строк?

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

Ключевые особенности архитектуры, которые мы перенесем в нашу реализацию:

Sparse Mixture of Experts (MoE): Вместо того чтобы прогонять каждый токен через все веса сети, мы используем роутер (router), который направляет каждый токен только к подмножеству наиболее релевантных экспертов.
Multi-head Latent Attention (MLA): Сжатие ключей и значений (K/V) в латентное пространство для радикального сокращения потребления памяти во время инференса.
Функциональный подход JAX: Полное отсутствие мутабельного состояния внутри слоев. Все веса передаются явно через чистые функции, что идеально ложится на парадигму jax.jit и jax.grad.

JAX выбран неслучайно. Благодаря концепциям vmap, grad и jit, написание кастомных слоев маршрутизации экспертов становится тривиальной задачей линейной алгебры, избавленной от головной боли с написанием кастомных CUDA-ядер.

Подготовка окружения и базовые тензорные структуры

Начнем с импорта необходимых модулей и инициализации базовых компонентов. Нам понадобятся сам jax, его библиотека работы с массивами jax.numpy (jnp), а также утилиты для генерации случайных чисел через jax.random. В чистом JAX управление состоянием генератора случайных чисел — это зона ответственности разработчика, что делает код детерминированным и легко тестируемым.

Посмотрим на базовую инициализацию модели и создание конфигурации:

import jax
import jax.numpy as jnp
from typing import NamedTuple

class Config(NamedTuple):
    vocab_size: int = 32000
    d_model: int = 512
    n_layers: int = 4
    n_heads: int = 8
    d_ff: int = 2048
    n_experts: int = 8
    n_active_experts: int = 2
    max_len: int = 512

config = Config()

Эта компактная конфигурация задает все параметры нашей мини-модели. Обратите внимание на параметры n_experts и n_active_experts. Именно они определяют поведение нашей MoE-подсистемы: каждый токен будет обрабатываться только двумя экспертами из восьми доступных.

Теперь создадим функцию инициализации весов. В чистом JAX веса обычно представляют собой обычные словари (nested dictionaries) с массивами jnp.array. Это избавляет нас от сложных объектно-ориентированных иерархий.

Реализация Sparse MoE слоя на чистом JAX

Сердце DeepSeek V3 — это его эффективная система маршрутизации экспертов. В традиционных MoE возникают проблемы с балансировкой нагрузки: некоторые эксперты перегружаются, в то время как другие простаивают. В нашей компактной реализации на 200 строк мы сосредоточимся на базовой, но математически точной логике гейтинга (routing) с использованием функции softmax и выбора топ-K экспертов.

Давайте напишем код слоя Mixture of Experts:

def expert_layer(x, params, key):
    # x shape: [batch, seq_len, d_model]
    batch_size, seq_len, d_model = x.shape
    
    # Роутер определяет веса для экспертов
    router_logits = jnp.dot(x, params['router']) # [B, S, n_experts]
    router_probs = jax.nn.softmax(router_logits, axis=-1)
    
    # Выбираем топ-K экспертов
    top_k_probs, top_k_indices = jax.lax.top_k(router_probs, config.n_active_experts)
    top_k_probs = top_k_probs / jnp.sum(top_k_probs, axis=-1, keepdims=True)
    
    # Инициализируем выходной тензор
    out = jnp.zeros_like(x)
    
    # Логика применения экспертов
    for i in range(config.n_experts):
        expert_w1 = params['experts'][f'w1_{i}']
        expert_w2 = params['experts'][f'w2_{i}']
        
        # Прогон через FFN конкретного эксперта
        expert_out = jax.nn.silu(jnp.dot(x, expert_w1))
        expert_out = jnp.dot(expert_out, expert_w2)
        
        # Маскирование и суммирование в соответствии с выбором роутера
        mask = (top_k_indices == i).any(axis=-1, keepdims=True)
        out += expert_out * mask * jnp.mean(top_k_probs, axis=-1, keepdims=True)
        
    return out

Этот фрагмент кода демонстрирует всю мощь функционального подхода. Мы вычисляем вероятности роутера, находим индексы наиболее подходящих экспертов и применяем к ним нелинейные преобразования, избегая создания лишних мутабельных состояний. Такой код компилируется в эффективное машинное представление с помощью jax.jit.

Интеграция внимания и сборка полной архитектуры

Успешно разобравшись с маршрутизацией экспертов, переходим к следующему фундаментальному блоку — механизму внимания, который связывает разрозненные токены в единый смысловой контекст.

Чтобы наша модель полноценно функционировала как языковой блок, помимо MoE слоя нам необходим механизм внимания (Attention) и нормализации. В DeepSeek V3 активно используется концепция сжатия ключей и значений, но в нашей компактной реализации мы реализуем стандартное масштабированное dot-произведение внимания (Scaled Dot-Product Attention), обернутое в чистые функции JAX.

Соберем все блоки вместе в единую функцию форварда сети:

def attention_layer(x, params):
    B, S, D = x.shape
    H = config.n_heads
    d_k = D // H
    
    q = jnp.dot(x, params['q_proj']).reshape(B, S, H, d_k).swapaxes(1, 2)
    k = jnp.dot(x, params['k_proj']).reshape(B, S, H, d_k).swapaxes(1, 2)
    v = jnp.dot(x, params['v_proj']).reshape(B, S, H, d_k).swapaxes(1, 2)
    
    scores = jnp.matmul(q, k.swapaxes(-1, -2)) / jnp.sqrt(d_k)
    
    # Примитивная причинная маска (causal mask)
    mask = jnp.tril(jnp.ones((S, S))) == 0
    scores = jnp.where(mask, -1e9, scores)
    
    attn_weights = jax.nn.softmax(scores, axis=-1)
    context = jnp.matmul(attn_weights, v)
    context = context.swapaxes(1, 2).reshape(B, S, D)
    
    return jnp.dot(context, params['o_proj'])

def transformer_block(x, params, key):
    # Нормализация и внимание с residual connection
    norm_x = jax.nn.layer_norm(x, axis=-1)
    x = x + attention_layer(norm_x, params['attn'])
    
    # Нормализация и MoE с residual connection
    norm_x = jax.nn.layer_norm(x, axis=-1)
    x = x + expert_layer(norm_x, params['moe'], key)
    
    return x

Получившийся код поражает своей лаконичностью. Всего несколько десятков строк чистой математики реализуют полноценный трансформерный блок с разреженными экспертами. Никаких скрытых слоев регистрации параметров, менеджеров устройств или кастомных классов контекста.

Компиляция и инференс с помощью JAX Transformations

Главное преимущество написания кода на чистом JAX заключается в том, что написанные нами функции сразу готовы к использовании с мощнейшими трансформациями: jax.jit, jax.vmap и jax.grad. Нам не нужно писать отдельный код для распределенного вычисления градиентов или оптимизации графа — экосистема делает это автоматически.

Посмотрим, как инициализировать модель целиком и запустить один шаг прогона:

def init_model(key):
    k1, k2 = jax.random.split(key)
    params = {
        'embed': jax.random.normal(k1, (config.vocab_size, config.d_model)) * 0.02,
        'layers': []
    }
    
    for _ in range(config.n_layers):
        layer_params = {
            'attn': {
                'q_proj': jax.random.normal(k2, (config.d_model, config.d_model)) * 0.02,
                'k_proj': jax.random.normal(k2, (config.d_model, config.d_model)) * 0.02,
                'v_proj': jax.random.normal(k2, (config.d_model, config.d_model)) * 0.02,
                'o_proj': jax.random.normal(k2, (config.d_model, config.d_model)) * 0.02,
            },
            'moe': {
                'router': jax.random.normal(k2, (config.d_model, config.n_experts)) * 0.02,
                'experts': {
                    f'w1_{i}': jax.random.normal(k2, (config.d_model, config.d_ff)) * 0.02
                    for i in range(config.n_experts)
                } | {
                    f'w2_{i}': jax.random.normal(k2, (config.d_ff, config.d_model)) * 0.02
                    for i in range(config.n_experts)
                }
            }
        }
        params['layers'].append(layer_params)
    return params

def forward(params, input_ids, key):
    x = params['embed'][input_ids]
    for layer_params in params['layers']:
        key, subkey = jax.random.split(key)
        x = transformer_block(x, layer_params, subkey)
    logits = jnp.dot(x, params['embed'].T)
    return logits

# Тестовый запуск
key = jax.random.PRNGKey(0)
params = init_model(key)
dummy_input = jnp.array([[1, 5, 12, 42]], dtype=jnp.int32)
logits = jax.jit(forward)(params, dummy_input, key)
print("Logits shape:", logits.shape)

«Чистый JAX стирает границу между исследовательской математикой и продакшн-кодом. То, что в других фреймворках требует написания кастомных распределенных операторов, здесь решается комбинацией стандартных функциональных примитивов».

Заключение

Создание реплики архитектуры уровня DeepSeek V3 в пределах ~200 строк чистого JAX доказывает удивительную вещь: за сложнейшими маркетинговыми названиями и гигантскими масштабами современных LLM скрываются элегантные и понятные математические концепции (хотя заставить это работать на легаси-инфраструктуре без боли всё равно не выйдет). Использование функционального подхода JAX позволяет инженерам экспериментировать с передовыми идеями вроде Sparse MoE без погружения в дебри тяжеловесных фреймворков.

Конечно, наша реализация опускает промышленные оптимизации: квантование, тензорный параллелизм или сложное кеширование KV. Однако она служит идеальным фундаментом для глубокого понимания того, как устроены современные разреженные языковые модели, и позволяет легко модифицировать код под собственные исследовательские задачи.

Попробуйте прямо сейчас: скопируйте этот сниппет в блокнот Google Colab, запустите его на беспилотном GPU (ведь «работает на моей машине» в мире ML обычно означает наличие у вас хотя бы пары H100) и попробуйте добавить кастомную функцию потерь или изменить количество экспертов — вы удивитесь, насколько легко поддается модификации чистый функциональный код.

DeepSeek V3 в ~200 строках чистого JAX: анатомия легковесного MoE