Обучите трансформер от начала до конца с помощью Flax NNX и Orbax.

1. Введение

Обучение Jax на графическом процессоре. Лабораторная работа 7: Обучение небольшой модели трансформера.

В этом практическом задании вы определите небольшой трансформер-декодер в JAX с использованием Flax NNX , обучите его на тексте Шекспира на обоих графических процессорах вашего узла, сохраните и восстановите его с помощью Orbax , а также сгенерируете новый текст на основе обученных весов.

Модель намеренно миниатюрна: 4 слоя, 256-мерные эмбеддинги, байтовый словарь, благодаря чему она обучается менее чем за минуту на двух графических процессорах L4. Архитектура и шаблоны обучения идентичны тем, которые используются в гораздо более крупных моделях.

Что вы будете делать

  • Определите преобразователь декодера с помощью Flax NNX, используя nnx.Embed , nnx.MultiHeadAttention , nnx.Linear и nnx.LayerNorm
  • Подключите causal jax.nn.dot_product_attention в качестве бэкенда для управления вниманием.
  • Обучение модели TinyShakespeare на байтовом уровне с использованием nnx.Optimizer и Optax AdamW сначала на одном графическом процессоре, а затем на всех.
  • Измерьте пропускную способность в токенах/сек и сравните результаты двух запусков.
  • Сохранение и восстановление параметров модели с помощью Orbax StandardCheckpointer .
  • Сгенерируйте текст в стиле Шекспира на основе обученной модели.

Что вам понадобится

  • Проект Google Cloud с включенной оплатой и кредитами для семинара или резервированием, покрывающим использование графического процессора.
  • Квота на использование как минимум двух видеокарт NVIDIA L4 в выбранном вами регионе ( как проверить квоту на видеокарты )
  • Пройдены лабораторные работы с 1 по 6, или создана эквивалентная среда JAX GPU, включающая как минимум два графических процессора.
  • Доступ к интернету с устройства Pod позволит при первом запуске загрузить текстовый файл TinyShakespeare.

Примерное время выполнения: 70 минут .

Архитектура, которую вы строите

Модель представляет собой стопку трансформаторных блоков. Каждый блок имеет два подслоя, каждый из которых заключен в остаточное соединение:

  1. Самовнимание — каждая позиция обращает внимание на все предыдущие позиции (причинно-следственная маска).
  2. Сеть прямого распространения (FFN) — это два линейных слоя с функцией активации GELU, расширяющие, а затем сжимающие представление.

Оба подслоя используют предварительную нормализацию, при этом LayerNorm применяется перед подслоем. Предварительная нормализация более стабильна для обучения и является стандартом в современных трансформерах.

2. Прежде чем начать

Выберите свой проект

В консоли Google Cloud выберите или создайте проект с включенной функцией выставления счетов.

Открытая облачная оболочка

Чтобы запустить сеанс Cloud Shell , нажмите кнопку «Активировать Cloud Shell» (значок терминала в правом верхнем углу консоли), а затем укажите в качестве исполнителя свой проект:

gcloud config set project <YOUR_PROJECT_ID>

Подготовка среды для работы с графическим процессором.

Выполните следующие команды в Cloud Shell . В Codelab 1 подробно объясняется каждая команда, включая требования к квоте GPU и зоне.

gcloud services enable \
  container.googleapis.com \
  compute.googleapis.com \
  iam.googleapis.com \
  cloudresourcemanager.googleapis.com \
  logging.googleapis.com \
  monitoring.googleapis.com

git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform

cp terraform.tfvars.example terraform.tfvars

Отредактируйте файл terraform.tfvars и установите project_id = " " Затем выполните настройку кластера и разверните project_id = " " :

terraform init
terraform apply
$(terraform output -raw get_credentials_command)

cd ..
kubectl apply -f deploy/jupyter.yaml

terraform apply занимает около 12 минут. После завершения дождитесь завершения работы Pod и LoadBalancer, затем считайте одноразовый токен JupyterLab из лога Pod:

kubectl get pod jax-jupyter -w         # wait for Running, then Ctrl+C
kubectl get svc jax-jupyter-svc -w     # wait for EXTERNAL-IP, then Ctrl+C
kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1

Откройте http:// :8884 , вставьте токен и создайте новый блокнот Python 3 в /workspace . Каждый блок кода из этого практического занятия помещается в ячейку этого блокнота.

Установите все необходимое для этого практического занятия. 

!pip install --quiet flax optax orbax-checkpoint matplotlib

Настройте и проверьте графический процессор.

Импортируйте JAX, Flax NNX, Optax и Orbax и проверьте, сколько графических процессоров может видеть контейнер.

import os

os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get(
    "LD_LIBRARY_PATH", ""
)
import hashlib
import html
import math
import pathlib
import time
import urllib.request
import warnings
from IPython.display import HTML, display
import matplotlib.pyplot as plt
import numpy as np

warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", message=".*ml_dtypes.*")
warnings.filterwarnings("ignore", message=".*JAX_PLATFORMS.*")
import jax
import jax.numpy as jnp
import optax
from flax import nnx
import orbax.checkpoint as ocp
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding

devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]
NUM_DEVICES = len(gpu_devices)
print(f"JAX version:     {jax.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"GPU devices:     {gpu_devices}")
print(f"GPU count:       {NUM_DEVICES}")
assert len(gpu_devices) >= 2, (
    f"This lesson needs at least 2 GPUs. Found {len(gpu_devices)}. "
    f"Available devices: {devices}"
)

def block_tree(tree):
    return jax.block_until_ready(tree)

def show_table(headers, rows, title=None, aligns=None):
    aligns = aligns or ["left"] * len(headers)
    parts = ["<div style='font-family: system-ui; max-width: 980px;'>"]
    if title:
        parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    parts.append(
        "<table style='border-collapse: collapse; width: 100%; font-size: 13px;'>"
    )
    parts.append("<thead><tr>")
    for h, a in zip(headers, aligns):
        parts.append(
            f"<th style='text-align:{a}; border-bottom:1px solid #d0d7de; padding:6px;'>"
            f"{html.escape(str(h))}</th>"
        )
    parts.append("</tr></thead><tbody>")
    for row in rows:
        parts.append("<tr>")
        for cell, a in zip(row, aligns):
            parts.append(
                f"<td style='text-align:{a}; border-bottom:1px solid #eef1f4; padding:6px;'>"
                f"{html.escape(str(cell))}</td>"
            )
        parts.append("</tr>")
    parts.append("</tbody></table></div>")
    display(HTML("".join(parts)))

def show_bars(rows, title, unit="", lower_is_better=False):
    max_value = max(float(value) for _, value in rows) or 1.0
    color = "#1a7f37" if not lower_is_better else "#0969da"
    parts = ["<div style='font-family: Arial, sans-serif; max-width: 760px;'>"]
    parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    for label, value in rows:
        width = max(3, 100 * float(value) / max_value)
        parts.append(
            "<div style='display:grid; grid-template-columns: 190px 1fr 130px; gap: 8px; "
            "align-items:center; margin: 6px 0;'>"
            f"<div style='font-size:13px;'>{html.escape(str(label))}</div>"
            "<div style='background:#f6f8fa; border-radius:6px; overflow:hidden; height:22px;'>"
            f"<div style='height:22px; width:{width:.1f}%; background:{color};'></div></div>"
            f"<div style='font-size:13px; font-variant-numeric: tabular-nums;'>{float(value):,.0f} {html.escape(unit)}</div>"
            "</div>"
        )
    parts.append(
        f"<div style='font-size:12px; color:#57606a;'>"
        f"{'Lower' if lower_is_better else 'Higher'} is better.</div></div>"
    )
    display(HTML("".join(parts)))

Вы должны увидеть версию JAX, gpu в качестве бэкенда по умолчанию, список из двух устройств CUDA и GPU count: 2 Вспомогательные функции block_tree , show_table и show_bars отображают таблицы и гистограммы, используемые на последующих этапах.

3. Подготовка данных TinyShakespeare на уровне байтов.

TinyShakespeare — это один текстовый файл размером около 1 МБ, содержащий объединенные произведения Шекспира. В этом практическом задании используется побайтовая токенизация, при которой каждый байт текста в кодировке UTF-8 становится одним токеном. Это ограничивает словарь 256 возможными значениями и устраняет зависимость от токенизатора.

Текст разбит на непересекающиеся последовательности длиной SEQ_LEN . Каждая последовательность представляет собой один обучающий пример, и модель учится предсказывать следующий байт в каждой позиции.

Запустите код, чтобы загрузить файл, проверить его контрольную сумму, разделить его на обучающую и проверочную последовательности и перемешать обучающий набор:

SHAKESPEARE_URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
SHAKESPEARE_MD5 = "d015dc5942f9b2908e24d4827a3e7a5e"
DATA_DIR = pathlib.Path.home() / ".cache" / "jax-course"
DATA_DIR.mkdir(parents=True, exist_ok=True)
DATA_FILE = DATA_DIR / "tinyshakespeare.txt"

def md5sum(path):
    digest = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()

if DATA_FILE.exists() and md5sum(DATA_FILE) == SHAKESPEARE_MD5:
    print("Using cached tinyshakespeare.txt")
else:
    print(f"Downloading tinyshakespeare.txt")
    urllib.request.urlretrieve(SHAKESPEARE_URL, DATA_FILE)
raw_text = DATA_FILE.read_text()
data = np.frombuffer(raw_text.encode("utf-8"), dtype=np.uint8).astype(np.int32)
print()
VOCAB_SIZE = 256
SEQ_LEN = 256
PER_DEVICE_BATCH = 32
num_sequences = len(data) // SEQ_LEN
data = data[: num_sequences * SEQ_LEN].reshape(num_sequences, SEQ_LEN)
num_train = int(0.9 * num_sequences)
train_data = data[:num_train]
val_data = data[num_train:]
rng = np.random.default_rng(0)
train_data = train_data[rng.permutation(num_train)]

def make_batches(data, batch_size):
    usable = (len(data) // batch_size) * batch_size
    return data[:usable].reshape(-1, batch_size, SEQ_LEN)

show_table(
    ["", "Value"],
    [
        ("Total bytes", f"{len(raw_text):,}"),
        ("Vocabulary", f"{VOCAB_SIZE} (raw bytes)"),
        ("Sequence length", SEQ_LEN),
        ("Training sequences", f"{num_train:,}"),
        ("Validation sequences", f"{len(val_data):,}"),
    ],
    title="TinyShakespeare — byte-level tokenization",
)
print()
print("Sample text (first 200 bytes):")
print(raw_text[:200])

Вы должны увидеть либо сообщение Using cached tinyshakespeare.txt , либо сообщение о загрузке, за которым следует таблица с указанием размера набора данных, общего количества байтов, словаря объемом 256 необработанных байтов, длины последовательности, количества последовательностей в обучающей и валидационной выборках, а затем первые 200 байтов текста, чтобы вы могли увидеть, на чем обучается модель.

Три константы, заданные здесь, имеют значение для остальной части практического задания. VOCAB_SIZE равно 256, поскольку байт имеет 256 возможных значений, а SEQ_LEN равно 256 позициям на один обучающий пример.

PER_DEVICE_BATCH равно 32 и остается фиксированным как для однопроцессорной, так и для многопроцессорной версии. Именно эта фиксированная величина для каждого процессора делает последующее сравнение пропускной способности неэффективным с точки зрения масштабирования.

4. Определите преобразователь с помощью Flax NNX.

Flax NNX — это упрощенный API для нейронных сетей в JAX. Слои определяются как объекты Python, которые отвечают за инициализацию весов и прямой проход. На этом шаге представлены все необходимые для трансформера компоненты NNX, поэтому предполагается отсутствие предварительного опыта работы с NNX.

Модель состоит из четырех основных элементов:

  • nnx.Embed — это справочная таблица, которая сопоставляет индекс токена с вектором.
  • nnx.Linear — это операция умножения плотной матрицы с дополнительным необязательным смещением.
  • nnx.LayerNorm нормализует признаки перед подслоями внимания и FFN.
  • nnx.MultiHeadAttention обрабатывает проекции Q/K/V, механизм внимания и проекцию выходных данных.

Настройте причинно-следственное внимание.

nnx.MultiHeadAttention получает скрытые состояния в формате (B, T, D_MODEL) , создает внутренние объекты Q, K и V и разделяет их на головы. Хук attention_fn управляет только основной операцией внимания, которая выполняется после этих проекций.

NNX передает в attention_fn необязательные аргументы в стиле Flax, такие как генератор случайных чисел для отсечения (dropout rng), тип данных (dtype) и точность (precision). Функция jax.nn.dot_product_attention их не принимает, поэтому приведенная ниже обертка принимает их с помощью универсального обработчика и передает только то, что необходимо функции JAX.

D_MODEL = 256
NUM_HEADS = 4
FFN_DIM = 1024
NUM_LAYERS = 4
MAX_SEQ_LEN = 256
LR = 3e-4
WEIGHT_DECAY = 1e-4

def causal_sdpa(query, key, value, **_):
    return jax.nn.dot_product_attention(query, key, value, is_causal=True)

Определите блок и модель.

Существует два класса. TransformerBlock — это один подслой внимания плюс один подслой FFN, а TinyTransformer размещает num_layers таких слоев между эмбеддингами и заголовком языковой модели. Каждый класс наследует nnx.Module и создает все свои слои в __init__ .

class TransformerBlock(nnx.Module):
    def __init__(self, d_model: int, num_heads: int, ffn_dim: int, rngs: nnx.Rngs):
        self.ln1 = nnx.LayerNorm(d_model, rngs=rngs)
        self.attn = nnx.MultiHeadAttention(
            num_heads=num_heads,
            in_features=d_model,
            decode=False,
            attention_fn=causal_sdpa,
            rngs=rngs,
        )
        self.ln2 = nnx.LayerNorm(d_model, rngs=rngs)
        self.fc_up = nnx.Linear(d_model, ffn_dim, rngs=rngs)
        self.fc_down = nnx.Linear(ffn_dim, d_model, rngs=rngs)

    def __call__(self, x):
        x = x + self.attn(self.ln1(x))
        h = jax.nn.gelu(self.fc_up(self.ln2(x)))
        x = x + self.fc_down(h)
        return x

Предварительная нормализация видна в __call__ . x = x + self.attn(self.ln1(x)) нормализует перед механизмом внимания и добавляет результат обратно в остаточный поток, а ветвь FFN делает то же самое с self.ln2 . 

class TinyTransformer(nnx.Module):
    def __init__(
        self,
        vocab_size: int,
        d_model: int,
        num_heads: int,
        ffn_dim: int,
        num_layers: int,
        max_seq_len: int,
        rngs: nnx.Rngs,
    ):
        self.token_embed = nnx.Embed(vocab_size, d_model, rngs=rngs)
        self.pos_embed = nnx.Embed(max_seq_len, d_model, rngs=rngs)
        self.blocks = nnx.List(
            [
                TransformerBlock(d_model, num_heads, ffn_dim, rngs=rngs)
                for _ in range(num_layers)
            ]
        )
        self.final_norm = nnx.LayerNorm(d_model, rngs=rngs)
        self.lm_head = nnx.Linear(d_model, vocab_size, use_bias=False, rngs=rngs)

    def __call__(self, tokens):
        B, T = tokens.shape
        x = self.token_embed(tokens) + self.pos_embed(jnp.arange(T))
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        return self.lm_head(x)

TinyTransformer.__call__ добавляет векторное представление токена и векторное представление позиции, выполняет блоки по порядку, применяет итоговую норму LayerNorm и проецирует данные в логиты размером vocab_size .

Создайте экземпляр модели и проверьте её.

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

model = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(0),
)
param_count = sum(x.size for x in jax.tree.leaves(nnx.state(model, nnx.Param)))
show_table(
    ["", "Value"],
    [
        ("Architecture", f"Decoder-only transformer"),
        ("Layers", NUM_LAYERS),
        ("Model dimension", D_MODEL),
        ("Attention heads", f"{NUM_HEADS} (head dim = {D_MODEL // NUM_HEADS})"),
        ("FFN dimension", FFN_DIM),
        ("Vocabulary", f"{VOCAB_SIZE} (byte-level)"),
        ("Max sequence length", MAX_SEQ_LEN),
        ("Parameters", f"{param_count:,}"),
    ],
    title="TinyTransformer",
)
logits = model(jnp.zeros((1, 16), dtype=jnp.int32))
print(f"Test forward pass: input (1, 16) \u2192 logits {logits.shape}")

Вы должны увидеть таблицу, описывающую архитектуру: 4 слоя, размерность модели 256, 4 механизма внимания с размерностью головки 64, байтовый словарь и количество параметров. В последней строке отображается результат прямого прохода теста с входными данными формы (1, 16) , генерирующими логиты формы (1, 16, 256) — одно распределение по 256 байтовым значениям для каждой из 16 позиций входных данных.

5. Напишите шаг обучения NNX.

Здесь используется @nnx.jit для автоматической обработки состояния модулей NNX. Он разделяет модули на структуры и массивы для JIT-компиляции, а затем объединяет обновленные массивы обратно, поэтому вы пишете этот шаг так, как если бы модули были обычными объектами Python.

Функция потерь — это предсказание следующего токена. На каждой позиции модель предсказывает следующий токен, поэтому вы сравниваете logits[:, :-1] (предсказания на позициях от 0 до T-2) с tokens[:, 1:] (фактические токены на позициях от 1 до T-1). 

@nnx.jit
def train_step(model, optimizer, tokens):
    def loss_fn(model):
        logits = model(tokens)
        pred = logits[:, :-1].reshape(-1, VOCAB_SIZE)
        target = tokens[:, 1:].reshape(-1)
        return optax.softmax_cross_entropy_with_integer_labels(pred, target).mean()

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return {"loss": loss, "perplexity": jnp.exp(loss)}

Цикл вокруг этого шага один раз прогревается, поэтому время компиляции не измеряется, затем выполняется фиксированное количество шагов, и прошедшее время преобразуется в токены в секунду. 

def train_loop(model, optimizer, batches, steps=1000, log_every=100):
    num_batches = batches.shape[0]
    history = []
    # Warmup: compile the training step
    warmup_metrics = train_step(model, optimizer, batches[0])
    block_tree(warmup_metrics)
    start = time.perf_counter()
    for step in range(steps):
        tokens = batches[step % num_batches]
        metrics = train_step(model, optimizer, tokens)
        if step % log_every == 0 or step == steps - 1:
            metrics = block_tree(metrics)
            history.append(
                {
                    "step": step,
                    "loss": float(metrics["loss"]),
                    "perplexity": float(metrics["perplexity"]),
                }
            )
    block_tree(metrics)
    elapsed = time.perf_counter() - start
    batch_size = int(batches.shape[1])
    tokens_per_step = batch_size * (SEQ_LEN - 1)
    tokens_per_sec = steps * tokens_per_step / elapsed
    return history, elapsed, tokens_per_sec

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

6. Обучение на одном графическом процессоре.

Начните с одного графического процессора, чтобы установить базовый уровень. jax.device_put привязывает весь массив пакетов к gpu_devices[0] , поэтому ничего не распределяется по устройствам до того, как вы доберетесь до целенаправленного сравнения нескольких графических процессоров. 

BENCHMARK_STEPS = 500
STEPS_1GPU = BENCHMARK_STEPS
batches_1gpu = make_batches(train_data, PER_DEVICE_BATCH)
single_device = gpu_devices[0]
batches_1gpu = jax.device_put(batches_1gpu, single_device)
model_1gpu = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(1),
)
optimizer_1gpu = nnx.Optimizer(
    model_1gpu, optax.adamw(LR, weight_decay=WEIGHT_DECAY), wrt=nnx.Param
)
history_1gpu, elapsed_1gpu, tps_1gpu = train_loop(
    model_1gpu, optimizer_1gpu, batches_1gpu, steps=STEPS_1GPU
)
show_table(
    ["Step", "Loss", "Perplexity"],
    [(h["step"], f"{h['loss']:.3f}", f"{h['perplexity']:.1f}") for h in history_1gpu],
    title=f"Single-GPU training \u2014 {tps_1gpu:,.0f} tokens/sec",
    aligns=["right", "right", "right"],
)

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

7. Примените один и тот же шаг ко всем графическим процессорам.

Это шаблон параллельной обработки данных, при котором вы создаете сетку, дублируете модель и разделяете данные по размерности пакета. Код этапа обучения остается неизменным, поскольку @nnx.jit обрабатывает параллелизм в зависимости от расположения массивов.

Для обеспечения объективности сравнения пропускной способности, в однопроцессорном и многопроцессорном режимах используется одинаковое количество шагов по времени. Каждый графический процессор по-прежнему обрабатывает PER_DEVICE_BATCH последовательностей за шаг, поэтому в многопроцессорном режиме обрабатывается больший глобальный пакет.

Для репликации модели NNX требуется три вызова. nnx.state извлекает состояние модуля в виде PyTree, jax.device_put размещает этот PyTree на каждом устройстве с реплицированным шардингом, а nnx.update записывает его обратно в модуль. Аналогичным образом обрабатывается и состояние оптимизатора. 

STEPS_MULTI = BENCHMARK_STEPS
GLOBAL_BATCH = PER_DEVICE_BATCH * NUM_DEVICES
mesh = Mesh(np.array(gpu_devices), ("data",))
replicated = NamedSharding(mesh, P())
data_sharding = NamedSharding(mesh, P(None, "data", None))
batches_multi = make_batches(train_data, GLOBAL_BATCH)
batches_multi = jax.device_put(batches_multi, data_sharding)
model_multi = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(1),
)
optimizer_multi = nnx.Optimizer(
    model_multi, optax.adamw(LR, weight_decay=WEIGHT_DECAY), wrt=nnx.Param
)
# Replicate model and optimizer state across all GPUs
model_state = nnx.state(model_multi)
nnx.update(model_multi, jax.device_put(model_state, replicated))
opt_state = nnx.state(optimizer_multi)
nnx.update(optimizer_multi, jax.device_put(opt_state, replicated))
print(
    f"Global batch: {GLOBAL_BATCH} ({PER_DEVICE_BATCH} per GPU \u00d7 {NUM_DEVICES} GPUs)"
)
print(f"Training batches: {batches_multi.shape}")
print()
history_multi, elapsed_multi, tps_multi = train_loop(
    model_multi, optimizer_multi, batches_multi, steps=STEPS_MULTI
)
show_table(
    ["Step", "Loss", "Perplexity"],
    [(h["step"], f"{h['loss']:.3f}", f"{h['perplexity']:.1f}") for h in history_multi],
    title=f"Multi-GPU training \u2014 {tps_multi:,.0f} tokens/sec",
    aligns=["right", "right", "right"],
)

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

8. Сравните пропускную способность и постройте графики.

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

ms_per_step_1gpu = elapsed_1gpu / STEPS_1GPU * 1e3
ms_per_step_multi = elapsed_multi / STEPS_MULTI * 1e3
speedup = tps_multi / tps_1gpu

show_table(
    ["", "1 GPU", f"{NUM_DEVICES} GPUs", "Ratio"],
    [
        ("Batch size", PER_DEVICE_BATCH, GLOBAL_BATCH, f"{NUM_DEVICES}×"),
        ("Per-GPU batch", PER_DEVICE_BATCH, PER_DEVICE_BATCH, "same"),
        ("Timed steps", STEPS_1GPU, STEPS_MULTI, "same"),
        ("ms/step", f"{ms_per_step_1gpu:.2f}", f"{ms_per_step_multi:.2f}", f"{ms_per_step_1gpu / ms_per_step_multi:.2f}×"),
        ("Tokens/sec", f"{tps_1gpu:,.0f}", f"{tps_multi:,.0f}", f"{speedup:.2f}×"),
    ],
    title="Throughput comparison",
    aligns=["left", "right", "right", "right"],
)

if speedup > NUM_DEVICES * 1.25:
    print(
        f"Note: the measured speedup is superlinear (> {NUM_DEVICES}x). "
        "For this small benchmark, treat that as a measurement artifact rather "
        "than a general hardware-scaling claim."
    )

show_bars(
    [("1 GPU", tps_1gpu), (f"{NUM_DEVICES} GPUs", tps_multi)],
    "Training throughput",
    "tokens/s",
)

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

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

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
for label, history, batch_size in [
    ("1 GPU", history_1gpu, PER_DEVICE_BATCH),
    (f"{NUM_DEVICES} GPUs", history_multi, GLOBAL_BATCH),
]:
    tokens_m = [
        (h["step"] + 1) * batch_size * (SEQ_LEN - 1) / 1e6 for h in history
    ]
    losses = [h["loss"] for h in history]
    perps = [h["perplexity"] for h in history]
    ax1.plot(tokens_m, losses, "o-", label=label, markersize=4)
    ax2.plot(tokens_m, perps, "o-", label=label, markersize=4)
ax1.set_xlabel("Tokens processed (millions)")
ax1.set_ylabel("Loss")
ax1.set_title("Training loss")
ax1.legend()
ax1.grid(True, alpha=0.25)
ax2.set_xlabel("Tokens processed (millions)")
ax2.set_ylabel("Perplexity")
ax2.set_title("Training perplexity")
ax2.legend()
ax2.grid(True, alpha=0.25)
fig.suptitle("Training curves vs tokens processed")
fig.tight_layout()
plt.show()

Вы должны увидеть две панели рядом: слева — потери, справа — замешательство. На каждой панели представлена ​​одна кривая для каждого запуска, обе с нисходящей тенденцией по мере увеличения количества обработанных токенов.

9. Сохранение и восстановление контрольной точки с помощью Orbax

Orbax сохраняет состояние модели в виде каталога файлов массивов. StandardCheckpointer — это его простейший API: один вызов для сохранения, один вызов для восстановления.

Для моделей NNX параметры извлекаются с помощью nnx.state(model, nnx.Param) , сохраняется полученный PyTree, а затем восстанавливается в новой модели с помощью nnx.update . Структура графа ( nnx.GraphDef ) не сохраняется и берется из определения класса Python, поэтому для перестройки модели перед загрузкой весов необходимо наличие TinyTransformer .

Приведённый ниже код выполняет весь жизненный цикл. Он извлекает только параметры обученной модели, сохраняет их на диск, строит дерево ShapeDtypeStruct указывающее Orbax, какие формы и типы данных следует ожидать, восстанавливает модель в заново инициализированный TinyTransformer и проверяет, что восстановленная модель выдаёт те же логиты, что и исходная, на небольшом тестовом входном сигнале. 

ckpt_dir = pathlib.Path("/tmp/jax-course/l7-checkpoints")
# Extract model parameters (not optimizer state)
model_params = nnx.state(model_multi, nnx.Param)
# Save
checkpointer = ocp.StandardCheckpointer()
if (ckpt_dir / "trained").exists():
    import shutil

    shutil.rmtree(ckpt_dir / "trained")
checkpointer.save(ckpt_dir / "trained", model_params)
print(f"Checkpoint saved to {ckpt_dir / 'trained'}")
# Create abstract target for restore
abstract_params = jax.tree.map(
    lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype),
    model_params,
)
# Restore into a fresh model
model_restored = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(99),
)
restored_params = checkpointer.restore(ckpt_dir / "trained", abstract_params)
nnx.update(model_restored, restored_params)
# Test
test_input = jnp.zeros((1, 16), dtype=jnp.int32)
logits_original = model_multi(test_input)
logits_restored = model_restored(test_input)
max_diff = float(jnp.max(jnp.abs(logits_original - logits_restored)))
show_table(
    ["", "Value"],
    [
        ("Checkpoint path", str(ckpt_dir / "trained")),
        ("Parameters saved", f"{sum(x.size for x in jax.tree.leaves(model_params)):,}"),
        ("Max |original \u2212 restored|", f"{max_diff:.2e}"),
        ("Match", "\u2713" if max_diff < 1e-5 else "\u2717"),
    ],
    title="Orbax checkpoint save and restore",
)

Вы должны увидеть путь к контрольной точке, количество сохраненных параметров, максимальную абсолютную разницу между исходными и восстановленными логитами, а также галочку, если эта разница меньше 1e-5.

10. Создайте текст в стиле Шекспира.

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

Функция generate дополняет свои входные данные значением MAX_SEQ_LEN , чтобы JIT-компилируемый прямой проход всегда видел одинаковую форму входных данных без перекомпиляции по мере роста последовательности. При наличии причинно-следственного механизма дополнение после реальных токенов не влияет на выходные данные в более ранних позициях.

nnx.split разделяет модуль на graphdef и state PyTree, чтобы его можно было передать через jax.jit , а nnx.merge перестраивает модуль внутри скомпилированной функции. 

@jax.jit
def get_logits_jit(graphdef, model_state, tokens):
    model = nnx.merge(graphdef, model_state)
    return model(tokens)

def generate(model, prompt_text, max_new_tokens=300, temperature=0.8):
    graphdef, model_state = nnx.split(model)
    tokens = list(prompt_text.encode("utf-8"))
    key = jax.random.key(42)
    for _ in range(max_new_tokens):
        context = tokens[-MAX_SEQ_LEN:]
        padded = context + [0] * (MAX_SEQ_LEN - len(context))
        input_arr = jnp.array([padded], dtype=jnp.int32)
        logits = get_logits_jit(graphdef, model_state, input_arr)
        next_logit = logits[0, len(context) - 1]
        if temperature <= 0:
            next_token = int(jnp.argmax(next_logit))
        else:
            key, subkey = jax.random.split(key)
            next_token = int(jax.random.categorical(subkey, next_logit / temperature))
        tokens.append(next_token)
    return bytes(tokens).decode("utf-8", errors="replace")

# Put model on a single device for generation
gen_model = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(99),
)
nnx.update(gen_model, checkpointer.restore(ckpt_dir / "trained", abstract_params))
print("=== Prompt: 'ROMEO:' | temperature=0.8 ===")
print()
print(generate(gen_model, "ROMEO:", max_new_tokens=300, temperature=0.8))
print()
print("=== Prompt: 'To be, or not' | temperature=0.6 ===")
print()
print(generate(gen_model, "To be, or not", max_new_tokens=300, temperature=0.6))

Вы должны увидеть два сгенерированных блока текста, по одному на каждый запрос.

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

11. Уборка

Удалите рабочую нагрузку Jupyter, включая балансировщик нагрузки и постоянный том: 

kubectl delete -f deploy/jupyter.yaml

Уничтожьте кластер, пул узлов, VPC и учетную запись службы: 

cd terraform
terraform destroy

При появлении запроса введите yes , затем подтвердите, что ничего не осталось: 

gcloud container clusters list
gcloud compute instances list

Оба поля должны быть пустыми для этого проекта. Если вы создали проект только для этой серии, вы можете удалить весь проект из консоли Cloud .

12. Поздравляем!

Вы создали трансформер-декодер с помощью Flax NNX, обучили его на обоих графических процессорах вашего узла, сохранили контрольную точку с помощью Orbax и сгенерировали текст на основе восстановленных весов.

Что вы узнали

  • Как Flax NNX организует модель в многократно используемые модули: nnx.Embed , nnx.Linear , nnx.LayerNorm и nnx.MultiHeadAttention обрабатывают создание параметров и прямой проход, а на этапе обучения используются nnx.value_and_grad с nnx.Optimizer
  • Как механизм причинного внимания с is_causal=True маскирует будущие позиции, так что модель может анализировать только прошлое, и почему хуку attention_fn необходима обертка **_
  • Как параллельное обучение данных воспроизводит модель и распределяет пакет данных по графическим процессорам точно так же, как в практическом задании 6, без изменения этапа обучения.
  • Как Orbax сохраняет и восстанавливает параметры модели, используя nnx.state и nnx.update для обеспечения связи между модулями NNX и обычными PyTrees.
  • Как измерить пропускную способность в токенах/сек, естественной единице измерения для языковых моделей, при этом функции warmup и block_until_ready продолжают обеспечивать точность этих данных.
  • Как генератор текста поочередно передает модели собственные прогнозы, и как температура контролирует случайность этой выборки.

Следующие шаги

  • Codelab 8: Экспорт и запуск обученной модели JAX. Для этого необходимо подготовить модель к запуску с помощью JIT-компиляции, AOT-компиляции и экспорта в портативные форматы.
  • Измените значения D_MODEL (попробуйте 128 или 512) и NUM_LAYERS (попробуйте 2 или 6), затем запустите процесс заново. Большая емкость означает более медленные шаги, а большее количество слоев означает больший объем памяти.
  • Измените температуру генерации (попробуйте 0,0, 0,5 и 1,2). Ноль — это слишком высокая скорость генерации, а значения выше 1 дают более случайный результат.
  • Измените бэкенд механизма внимания: измените тело функции causal_sdpa , чтобы она вызывала jax.nn.dot_product_attention(..., is_causal=True, implementation="cudnn") и используйте активации, совместимые с bf16, а затем понаблюдайте за поведением механизма внимания cuDNN.

Справочная документация