AI Research Foundations, курс 7, часть 1: FLOPs, память и «стена»

Перевод лабораторных 7.1–7.4 курса Google DeepMind «Accelerate Your Model»: сравнение Gemma 1B и 4B, оценка вычислений на обучение по формуле 6PN, нехватка памяти при загрузке Gemma 4B и калькулятор памяти GPU.

Курс 7 программы AI Research Foundations называется «Accelerate Your Model» («Ускорьте свою модель»). Он о том, во что обходится модель: сколько операций нужно на обучение, сколько памяти видеокарты занимают параметры, градиенты, состояние оптимизатора и активации, и что делать, когда памяти не хватает. В этой части — четыре лабораторные, которые ставят проблему и учат её считать. Во второй части — два приёма, которые её решают: формат bfloat16 и накопление градиентов.

Все учебные тексты, задания, справки и решения лабораторных 7.1–7.4; комментарии и строки документации в коде — на русском, печатаемые строки и промпты — как в оригинале. Служебные разделы про Colab и GPU — на странице «Обзор программы», настройка Kaggle — в части 2 курса 5. Курс 7 ссылается на статьи платформы Google Skills («Что происходит в GPU при обучении и инференсе», «Умные числа»), которых нет в репозитории, — их тему закрывают уроки ChinaAI, указанные ниже.

Как проходить вместе с курсом ChinaAI

Лабораторная	Время	Когда делать	Что нужно
7.1. Сравните модели разного размера	15 мин	Вместе с уроком «Prefill и decode, узкие места» или после «Законов масштабирования»	GPU
7.2. Оцените FLOPs обучения	15 мин	Вместе с уроком «Законы масштабирования» (раздел про $C \approx 6ND$ )	CPU
7.3. Упираемся в стену	15 мин	Перед уроком «Смешанная точность»	GPU, Kaggle
7.4. Оцените память GPU	15 мин	Вместе с уроком «Смешанная точность» (раздел про байты на параметр)	CPU

В уроке 5.1.1 курса ChinaAI та же формула записана как $C \approx 6ND$ : $N$ — параметры, $D$ — токены. В лабораторной 7.2 буквы другие — $P$ и $N$ , — но смысл тот же. А урок «Как устроена видеокарта и куда уходит время» продолжает тему: не только «сколько операций», но и «почему видеокарта считает их не на пиковой скорости».

flowchart LR
A["7.1: больше модель — лучше, но медленнее"] --> B["7.2: сколько операций — 6·P·N"]
B --> C["7.3: Gemma 4B не загружается — OOM"]
C --> D["7.4: куда уходит память"]
D --> E["Часть 2: bfloat16 и накопление градиентов"]

Логика первой части курса 7

Лабораторная 7.1. Сравните модели разного размера

Ноутбук: gdm_lab_7_1_compare_models_of_different_sizes.ipynb · открыть в Colab

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

Обзор

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

Чему вы научитесь

К концу лабораторной вы сможете:

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

Задачи

В этой лабораторной вы:

напишете код, чтобы посмотреть на генерации трансформеров Gemma-1B и Gemma-4B;
сравните время генерации и ответы обеих моделей, чтобы лучше понять компромиссы между маленькими и большими моделями.

Лабораторную нужно запускать на GPU. Выберите T4 GPU.

Импорты

В этой лабораторной вы в основном работаете с пакетом ai_foundations, который сделан специально для курса. За кулисами он загружает модели Gemma-1B и Gemma-4B и задаёт им промпты через пакет gemma.

%%capture
!pip install jax[cuda12]==0.7.2
# Ставим специальный пакет курса.
!pip install "git+https://github.com/google-deepmind/ai-foundations.git@main"

import os # Чтобы задавать системные переменные.
import time # Чтобы засекать, сколько длятся генерации.

# Чтобы красиво выводить генерации модели.
from IPython.display import display, HTML
from ai_foundations import generation # Чтобы генерировать тексты с Gemma.

# Разрешаем JAX использовать всю память GPU.
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1.00"

Генерации Gemma-1B

Чтобы подготовиться к сравнению генераций Gemma-1B и Gemma-4B, запустите следующую ячейку. Она загрузит модель Gemma-1B и сгенерирует продолжение промпта «Jide was hungry so she went looking for» («Джиде проголодалась и пошла искать…»), который встречался вам в прошлых курсах программы. Учтите: это исходный предобученный вариант Gemma-1B, а не версия, дообученная следовать инструкциям и оптимизированная для диалога.

# Загружаем модель Gemma-1B.
print("Loading Gemma-1B model...")
gemma_1b_model = generation.load_gemma(model_name="Gemma-1B")
print("Loaded Gemma-1B model.\n\n")

# Генерируем продолжение.
prompt = "Jide was hungry so she went looking for"
output_text_transformer, next_token_logits, tokenizer = (
    generation.prompt_transformer_model(
        prompt, max_new_tokens=10, loaded_model=gemma_1b_model
    )
)

print(f"Generation by Gemma-1B:\n{output_text_transformer}")

Задание 1. Промпты Gemma-1B

Промпты (по порядку): «Самая высокая гора Африки — это», «Главные ингредиенты и способ приготовления традиционного нигерийского риса джолоф — это», «Дикий кофе появился в», «Национальный цветок ЮАР — это», «Значение города Марракеш в том, что», «Самая населённая страна Африки — это», «Самая глубокая река мира — это», «Народ масаи происходит из», «Африканская пословица „Рычащий лев добычи не убьёт“ означает», «Браай — это».

list_of_prompts = [
    "The tallest mountain in Africa is",
    "The key ingredients and preparation method for making traditional Nigerian Jollof rice are",
    "Wild coffee plants originated in",
    "The national flower of South Africa is",
    "The significance of the city of Marrakesh is",
    "The most populous African nation is",
    "The world's deepest river is",
    "The Maasai ethnic group originate from ",
    "The African proverb 'A roaring lion kills no game' means",
    "Braai is",
]

generations_gemma_1b = []

before_loop = time.time_ns()

# Добавьте сюда свой код.

after_loop = time.time_ns()

# Длительность в наносекундах.
duration_ns = after_loop - before_loop

# Переводим в секунды.
total_seconds = (after_loop - before_loop) / 1_000_000_000

# Считаем минуты и оставшиеся секунды.
minutes = int(total_seconds // 60)
seconds = total_seconds % 60

if minutes > 0:
    if minutes == 1:
        minutes_str = "minute"
    else:
        minutes_str = "minutes"
    if seconds == 1:
        seconds_str = "second"
    else:
        seconds_str = "seconds"
    print(f"Duration: {minutes} {minutes_str} {seconds:.2f} {seconds_str}")
else:
    print(f"Duration: {seconds:.2f} seconds")

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

for i in range(len(list_of_prompts)):
    display(HTML(f"<h3>Prompt:</h3><p>{list_of_prompts[i]}</p>"))
    display(
        HTML(
            f"<blockquote><b>Gemma-1B generation:</b><br>{generations_gemma_1b[i]}</blockquote>"
        )
    )
    display(HTML("<hr>"))

Генерации Gemma-4B

Запустите следующую ячейку, чтобы загрузить модель Gemma-4B и сгенерировать продолжение того же промпта «Jide was hungry so she went looking for».

# Удаляем модель Gemma-1B, чтобы освободить память.
del gemma_1b_model

# Загружаем модель Gemma-4B.
print("Loading Gemma-4B model...")
gemma_4b_model = generation.load_gemma(model_name="Gemma-4B")
print("Loaded Gemma-4B model.\n\n")

# Генерируем продолжение.
prompt = "Jide was hungry so she went looking for"
output_text_transformer, next_token_logits, tokenizer = (
    generation.prompt_transformer_model(
        prompt, max_new_tokens=10, loaded_model=gemma_4b_model
    )
)

print(f"Generation by Gemma-4B:\n{output_text_transformer}")

Задание 2. Промпты Gemma-4B

Ячейка та же, что в задании 1 (тот же список промптов и тот же замер времени), только список результатов называется generations_gemma_4b, а после неё — такая же ячейка вывода с подписью «Gemma-4B generation».

Сравните генерации

Теперь, когда есть ответы обеих моделей, сравните и оцените генерации.

Скорость. Сначала сравните, сколько времени каждой модели понадобилось на генерацию. Какая модель быстрее?

Качество. Затем просмотрите ответы на каждый промпт. Какая модель в целом даёт более точные, подробные и связные ответы?

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

# Эта ячейка показывает ответы обеих моделей на каждый промпт.
for i in range(len(list_of_prompts)):
    display(HTML(f"<h3>Prompt:</h3><p>{list_of_prompts[i]}</p>"))
    display(
        HTML(
            f"<blockquote><b>Gemma-1B generation:</b>"
            f"<br>{generations_gemma_1b[i]}</blockquote>"
        )
    )
    display(
        HTML(
            f"<blockquote><b>Gemma-4B generation:</b>"
            f"<br>{generations_gemma_4b[i]}</blockquote>"
        )
    )
    display(HTML("<hr>"))

Что вы заметили?

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

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

И наоборот, маленькая Gemma-1B может оказаться лучшим выбором там, где критична скорость, а качество «достаточно хорошее» приемлемо, — например, в чат-боте, отвечающем в реальном времени.

Итоги

Эта лабораторная показала на практике компромиссы между языковыми моделями разного размера. Главные выводы:

Качество против эффективности. Большие модели вроде Gemma-4B обычно пишут более качественный, точный и связный текст. Но платят за это заметно меньшей скоростью.
Выбор модели зависит от задачи. Маленькие модели вроде Gemma-1B намного быстрее, поэтому подходят там, где важны скорость и эффективность. «Лучшая» модель — всегда баланс между нужным качеством и имеющимися ресурсами.
Интуиция о масштабировании. Вы на практике увидели, как рост числа параметров влияет и на возможности модели, и на её скорость.

Решения

Задание 1

# Добавьте этот код в ячейку выше.
for prompt in list_of_prompts:
    output_text_transformer, next_token_logits, tokenizer = (
        generation.prompt_transformer_model(
            prompt,
            max_new_tokens=50,
            loaded_model=gemma_1b_model,
            sampling_mode="greedy",
        )
    )
    generations_gemma_1b.append(output_text_transformer)

Задание 2

# Добавьте этот код в ячейку выше.
for prompt in list_of_prompts:
    output_text_transformer, next_token_logits, tokenizer = (
        generation.prompt_transformer_model(
            prompt,
            max_new_tokens=50,
            loaded_model=gemma_4b_model,
            sampling_mode="greedy",
        )
    )
    generations_gemma_4b.append(output_text_transformer)

Лабораторная 7.2. Оцените FLOPs обучения

Ноутбук: gdm_lab_7_2_estimate_training_flops.ipynb · открыть в Colab

Научитесь оценивать затраты на вычисления в FLOPs. 15 минут.

Обзор

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

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

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

В этой лабораторной вы научитесь оценивать вычислительную стоимость обучения языковой модели. Главная мера этой стоимости — FLOPs, от FLoating-point OPerations, «операции с плавающей точкой».

Чему вы научитесь

К концу лабораторной вы сможете:

объяснить, что такое FLOPs и почему ими измеряют объём вычислений;
применить простую формулу, чтобы оценить FLOPs, нужные для обучения.

Задачи

В этой лабораторной вы:

напишете функцию, которая вычисляет число FLOPs;
исследуете, как размер модели и набора данных влияет на число FLOPs, нужных для обучения.

Импорты

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

%%capture
# Ставим специальный пакет курса.
!pip install "git+https://github.com/google-deepmind/ai-foundations.git@main"

# Используемые пакеты.
import pandas as pd # Чтобы показывать таблицы.
from ai_foundations.utils import formatting # Чтобы форматировать большие числа.
from ai_foundations.feedback.course_7 import flops # Чтобы проверять решения.

# Чтобы красиво выводить результаты.
from IPython.display import display, HTML

Что такое FLOPs?

Вспомните: каждое умножение матриц в блоке трансформера состоит из умножений и сложений. Каждая такая операция — это операция с плавающей точкой, или FLOP, потому что в ней участвуют числа с плавающей точкой.

Чтобы измерить общую вычислительную стоимость задачи, например обучения модели, можно сложить все эти отдельные операции. Общее число называют числом FLOPs. В примере из предыдущей статьи «Что происходит в GPU при обучении и при инференсе?» простое умножение матрицы $2 \times 3$ на матрицу $3 \times 2$ потребовало 20 операций. А обучение языковой модели может потребовать миллиардов, триллионов и ещё больше таких операций.

Вам часто будут встречаться две очень похожие аббревиатуры: FLOPs и FLOPS. Регистр последней буквы S принципиален — он меняет смысл.

FLOPs (floating point operations) — множественное число от FLOP. Это общее число операций, нужное для конкретного процесса обучения или для инференса на конкретном примере.
- Например: «Обучение этой модели потребует 500 триллионов FLOPs».
FLOPS (floating point operations per second) — скорость, единица измерения производительности. Она показывает, сколько операций с плавающей точкой (FLOPs) конкретное оборудование (например, GPU) выполняет за секунду. Заглавная S здесь означает «per Second» — «в секунду».
- Например: «Производительность этой GPU — 200 терафлопс ( $200 \times 10^{12}$ FLOPS)».

Связь между ними: если обучение требует 1 000 FLOPs, а GPU выполняет 100 FLOPS, обучение займёт 10 секунд — ведь GPU делает 100 операций в секунду.

Оцениваем FLOPs

Широко распространённая формула «на салфетке» для числа FLOPs на одну эпоху обучения (то есть на один полный проход по набору данных):

\text{FLOPs} \approx 6 \times P \times N

где $P$ — общее число обучаемых параметров модели, а $N$ — общее число токенов в обучающем наборе. Читается так: «число операций примерно равно шести, умноженным на число параметров и на число токенов».

Множитель 6 — широко используемое практическое правило, которое учитывает работу на одном шаге обучения:

Прямой проход стоит примерно $2 \times P \times N$ FLOPs. Двойка — потому что на каждый параметр приходится примерно одно умножение и одно сложение (операция «умножить и накопить», multiply-accumulate).
Обратный проход стоит примерно $4 \times P \times N$ FLOPs: обратный проход, в котором вычисляются градиенты, примерно вдвое дороже прямого [1].

Сложив обе части, получаем итоговую оценку $6 \times P \times N$ .

Подробнее о способах оценки FLOPs — в статье Google DeepMind Training Compute-Optimal Large Language Models [2].

Задание 1. Реализуйте формулу оценки FLOPs

Чтобы считать FLOPs, нужна функция, которая возвращает их число.

def compute_num_flops(param_count: float, num_tokens: float) -> float:
    """Оценивает FLOPs на одну эпоху обучения.

    Args:
      param_count: Общее число обучаемых параметров модели.
      num_tokens: Общее число токенов в обучающем наборе.

    Returns:
      Оценка общего числа операций с плавающей точкой.
    """
    total_flops = ... # Добавьте сюда свой код.

    return total_flops

# @title Запустите эту ячейку, чтобы проверить реализацию.
flops.test_compute_num_flops(compute_num_flops)

Задание 2. Оцените стоимость обучения

Следующая задача — оценить FLOPs обучения для моделей и наборов данных разного размера. Все примеры — настоящие модели и наборы данных, созданные в Google.

В первой ячейке посчитайте оценку FLOPs для каждого из примеров ниже с помощью вашей функции compute_num_flops.
Впишите эти значения в словарь во второй ячейке.

Примеры

Сценарий	Параметры модели	Токены в наборе
BERT-base [3]	110 млн	3,3 млрд
T5-large [4]	770 млн	1 трлн
Gemma-1B и набор Africa Galore	1 млрд	30 тыс.
PaLM [5]	540 млрд	780 млрд

# @markdown Введите значения сценария, чтобы посчитать FLOPs.

param_count = 1e9  # @param {type:"number"}
num_tokens = 30e3  # @param {type:"number"}

# Считаем и показываем результат вашей функцией.
total_flops = compute_num_flops(param_count, num_tokens)

output_string = f"""
  <p>
      For a model with <b>{formatting.format_large_number(param_count)}
      parameters</b>
      trained on <b>{formatting.format_large_number(num_tokens)} tokens</b>:
  </p>
  <hr>
  <p><b>Estimated training FLOPs:</b> {formatting.format_flops(total_flops)}</p>
"""
display(HTML(output_string))

Впишите посчитанные значения FLOPs в следующую ячейку.

# Замените все '...' на посчитанные значения.
final_results_data = [
    {
        "Scenario": "BERT-Base",
        "Parameters": formatting.format_large_number(110e6),
        "Tokens": formatting.format_large_number(3.3e9),
        "FLOPs": ...,
    },
    {
        "Scenario": "T5-Large",
        "Parameters": formatting.format_large_number(770e6),
        "Tokens": formatting.format_large_number(1e12),
        "FLOPs": ...,
    },
    {
        "Scenario": "Gemma-1B (Africa Galore)",
        "Parameters": formatting.format_large_number(1e9),
        "Tokens": formatting.format_large_number(30e3),
        "FLOPs": ...,
    },
    {
        "Scenario": "PaLM",
        "Parameters": formatting.format_large_number(540e9),
        "Tokens": formatting.format_large_number(780e9),
        "FLOPs": ...,
    },
]

# Проверяем, верны ли результаты.
if flops.test_final_results(final_results_data):

    for scenario in final_results_data:
        scenario["FLOPs"] = formatting.format_flops(scenario["FLOPs"])

    # Строим и показываем таблицу.
    df = pd.DataFrame(final_results_data)
    display(df)

Что вы заметили?

Посмотрите на результаты в таблице. Обратите внимание, как стремительно растёт вычислительная стоимость с ростом моделей и наборов данных. Обучение Gemma-1B на наборе Africa Galore требует ничтожной доли вычислений, нужных для обучения гораздо большей модели вроде PaLM.

Упражнение показывает прямую связь между масштабом модели и огромными аппаратными ресурсами, которые нужны для обучения.

Задание 3. Исследуйте сами

С помощью ползунков в следующей ячейке исследуйте, как размер модели и набора данных влияет на вычислительную стоимость.

Эти числа растут очень быстро, поэтому ползунки меняют показатель степени (логарифмическая шкала). Каждый шаг ползунка увеличивает значение в 10 раз — так можно перейти от миллионов к миллиардам и триллионам параметров и токенов.

Например, для ползунка log_param_count: значение 6 означает $10^{6}$ — миллион параметров; 9 — $10^{9}$ , миллиард; 12 — $10^{12}$ , триллион.

# @title Оцените вычислительную стоимость
# @markdown Двигайте ползунки и смотрите, как растут FLOPs.

log_param_count = 9  # @param {type:"slider", min:6, max:12, step:1}
log_num_tokens = 14  # @param {type:"slider", min:9, max:14, step:1}

param_count = 10**log_param_count
num_tokens = 10**log_num_tokens

# Считаем и показываем результат.
total_flops = compute_num_flops(param_count, num_tokens)

output_string = f"""
    <p>
        <b>Model with:</b> {formatting.format_large_number(param_count)} parameters<br>
        <b>Dataset with:</b> {formatting.format_large_number(num_tokens)} tokens
    </p>
    <hr>
    <p><b>Estimated training FLOPs:</b> {formatting.format_flops(total_flops)}</p>
"""

display(HTML(output_string))

Заметьте: если увеличить на 1 любой из ползунков, log_param_count или log_num_tokens, показатель степени в итоговой «Estimated Training FLOPs» тоже вырастет примерно на 1. Шкала логарифмическая, поэтому +1 в показателе — это рост самого числа в 10 раз.

Отсюда видна прямая пропорциональная зависимость: если увеличить число параметров в 10 раз, нужный объём вычислений тоже вырастет в 10 раз. Такую предсказуемую зависимость «один к одному» называют линейной. Вот почему обучение моделей с триллионом параметров требует астрономического объёма вычислений по сравнению с меньшими моделями.

Итоги

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

В следующем задании (статья на платформе) речь пойдёт о затратах энергии на обучение моделей.

Решения

Задание 1

# Полная реализация compute_num_flops.
def compute_num_flops(param_count, num_tokens) -> float:
    """Оценивает FLOPs на одну эпоху обучения.

    Args:
      param_count: Общее число обучаемых параметров модели.
      num_tokens: Общее число токенов в обучающем наборе.

    Returns:
      Оценка общего числа операций с плавающей точкой.
    """

    total_flops = 6 * param_count * num_tokens

    return total_flops

Задание 2

final_results_data = [
    {
        "Scenario": "BERT-Base",
        "Parameters": formatting.format_large_number(110e6),
        "Tokens": formatting.format_large_number(3.3e9),
        "FLOPs": 2.18e18,
    },
    {
        "Scenario": "T5-Large",
        "Parameters": formatting.format_large_number(770e6),
        "Tokens": formatting.format_large_number(1e12),
        "FLOPs": 4.62e21,
    },
    {
        "Scenario": "Gemma-1B (Africa Galore)",
        "Parameters": formatting.format_large_number(1e9),
        "Tokens": formatting.format_large_number(30e3),
        "FLOPs": 1.80e14,
    },
    {
        "Scenario": "PaLM",
        "Parameters": formatting.format_large_number(540e9),
        "Tokens": formatting.format_large_number(780e9),
        "FLOPs": 2.53e24,
    },
]

Чтобы перевести FLOPs во время, разделите их на реальную скорость видеокарты. Паспортный пик RTX 4080 в BF16 — 97,5 TFLOPS (около $10^{14}$ FLOPS, см. урок «Как устроена видеокарта»); на практике при обучении достигается 30–50% пика (эту долю называют MFU, model FLOPs utilization — урок «Законы масштабирования»). Пример: модель на 125 млн параметров на 2,5 млрд токенов — это $6 \times 1{,}25 \cdot 10^{8} \times 2{,}5 \cdot 10^{9} \approx 1{,}9 \cdot 10^{18}$ FLOPs. При $4 \cdot 10^{13}$ полезных FLOPS на одну карту — около $4{,}7 \cdot 10^{4}$ секунд, то есть порядка 13 часов; на четырёх картах — примерно втрое-вчетверо быстрее.

Литература

Dzmitry Bahdanau. 2022. The FLOPs Calculus of Language Model Training. Medium. https://medium.com/@dzmitrybahdanau/the-flops-calculus-of-language-model-training-3b19c1f025e4
Jordan Hoffmann et al. 2022. Training Compute-Optimal Large Language Models. arXiv:2203.15556.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv:1810.04805.
Colin Raffel et al. 2019. Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. arXiv:1910.10683.
Aakanksha Chowdhery et al. 2022. PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311.

Лабораторная 7.3. Упираемся в стену

Ноутбук: gdm_lab_7_3_hitting_a_wall.ipynb · открыть в Colab

Исследуйте аппаратные ограничения обучения больших моделей и сами получите ошибку нехватки памяти. 15 минут.

Обзор

Ранее в этом курсе вы видели, что большие модели вроде Gemma-4B могут писать текст качественнее маленьких вроде Gemma-1B. Что понадобится, чтобы самому дообучить большую модель вроде Gemma-4B?

В прошлом курсе вы успешно дообучили Gemma-1B с LoRA. В этой лабораторной вы попробуете применить тот же процесс дообучения к более крупной модели Gemma-4B. Упражнение задумано так, чтобы вы на собственном опыте столкнулись со «стеной памяти» (memory wall). Это важнейшее аппаратное ограничение, которое, как вы увидите, и заставляет использовать приёмы повышения эффективности, которые вы изучаете в этом курсе.

Чему вы научитесь

К концу лабораторной вы сможете:

объяснить, почему даже с LoRA нельзя дообучить более крупные модели;
объяснить, зачем нужны приёмы экономии памяти.

Задачи

В этой лабораторной вы:

попробуете дообучить базовую модель Gemma-4B;
получите ошибку нехватки памяти (OOM) и подумаете о её причинах.

Лабораторную нужно запускать на GPU. Выберите T4 GPU. Также нужен аккаунт Kaggle, чтобы скачать веса Gemma 3 (инструкция — в части 2 курса 5).

Импорты

В этой лабораторной вы будете загружать модель Gemma через Keras, а набор данных — через pandas.

%%capture
!pip install jax[cuda12]==0.7.2
# Ставим специальный пакет курса.
!pip install "git+https://github.com/google-deepmind/ai-foundations.git@main"

import os # Чтобы задавать системные переменные.

from google.colab import userdata # Чтобы пользоваться секретами Colab.

os.environ["KAGGLE_USERNAME"] = userdata.get("KAGGLE_USERNAME")
os.environ["KAGGLE_KEY"] = userdata.get("KAGGLE_KEY")

os.environ["KERAS_BACKEND"] = "jax"  # Вычисления Keras — на JAX.
# Отключаем заблаговременное резервирование буферов команд, чтобы освободить память.
os.environ["XLA_FLAGS"] = "--xla_gpu_enable_command_buffer="
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="0.9"

import keras # Чтобы задавать и обучать модели.
import keras_hub # Чтобы загрузить реализацию Gemma на Keras.
import pandas as pd # Чтобы загрузить набор данных.
from textwrap import fill # Чтобы форматировать длинные абзацы.
from ai_foundations import formatting # Чтобы форматировать обучающие данные.

keras.utils.set_random_seed(812)  # Для слоёв Keras.

Дообучаем Gemma-4B

В прошлом курсе вы научились дообучать Gemma-1B с LoRA, чтобы построить генератор учебных карточек для повторения. В этой лабораторной вы примените те же приёмы, чтобы дообучить Gemma-4B.

Запустите следующую ячейку, чтобы:

определить функцию format_question, которая форматирует входы модели;
подготовить набор Africa Galore к дообучению.

def format_question(
    question: str,
    sot = "<start_of_turn>",
    eot = "<end_of_turn>"
) -> str:
    """
    Форматирует вопрос для промпта модели и добавляет специальные разделители
    в начале и в конце вопроса.

    Args:
      text: Вопрос, который нужно отформатировать.
      sot: Токен начала реплики.
      eot: Токен конца реплики.

    Returns:
      Отформатированная строка с вопросом.
    """

    formatted_q = f"{sot}user\n{question}{eot}\n"

    return formatted_q

# Загружаем набор вопросов и ответов.
africa_galore_qa = pd.read_json(
    "https://storage.googleapis.com/dm-educational/assets/ai_foundations/africa_galore_qa_v2.json"
)

questions = []  # Список отформатированных вопросов.
answers = []  # Список отформатированных ответов.

for idx, row in africa_galore_qa.iterrows():
    # Форматируем вопрос и ответ функцией format_qa из прошлой лабораторной.
    question, answer = formatting.format_qa(row)
    questions.append(question)
    answers.append(answer)

# Показываем первую пару «вход — выход».
print(questions[0])
print(fill(answers[0], replace_whitespace=False))

# Готовим словарь данных для дообучения Gemma.
data = {
    "prompts": questions,
    "responses": answers
}

Задание 1. Загрузите Gemma-4B

# Загружаем модель Gemma3-4B для Keras.
model = keras_hub.models.Gemma3CausalLM.from_preset(preset="gemma3_4b_text")
model.summary()

Что вы заметили?

Вы должны были получить ошибку RESOURCE_EXHAUSTED и сообщение о том, сколько памяти процесс пытался занять.

Ошибки нехватки памяти

RESOURCE_EXHAUSTED — техническое название ошибки нехватки памяти (out-of-memory, OOM). Оно означает, что при попытке загрузить модель у GPU закончилась память.

Ошибка показывает: модель такого размера при стандартных настройках нельзя даже загрузить. Но эта проблема не непреодолима. Именно её и решают современные приёмы повышения эффективности. В следующей статье вы узнаете, что именно занимает столько памяти GPU. А позже увидите, как приёмы вроде смешанной точности помогают преодолеть эту стену памяти.

Если вам интересно, почему в первой лабораторной этого курса Gemma-4B загрузить удалось: в той лабораторной за кулисами уже работали некоторые приёмы экономии памяти.

Итоги

Эта лабораторная наглядно показала аппаратные ограничения обучения больших моделей. Вы снова упёрлись в «стену памяти». Вы также увидели, что попытка дообучить большую модель может упасть с нехваткой памяти уже на этапе загрузки весов. В следующем задании вы подробнее узнаете о главных потребителях памяти GPU.

4 млрд параметров × 4 байта ≈ 15 ГиБ только на веса. На T4 с 15 ГБ это не помещается, а на RTX 4080 с 32 ГБ загрузка пройдёт — «стена» сдвинется дальше: полное обучение в float32 потребует около 62 ГБ (см. расчёт в лабораторной 7.4), а это уже больше одной карты. Чтобы воспроизвести опыт лабораторной, попробуйте загрузить gemma3_12b_text (12 млрд × 4 байта ≈ 45 ГиБ) — ошибка будет той же. Модель скачивается с Kaggle и занимает место на диске, так что, если не хочется тратить трафик, просто прикиньте числа.

Лабораторная 7.4. Оцените память GPU

Ноутбук: gdm_lab_7_4_estimate_gpu_memory.ipynb · открыть в Colab

Научитесь оценивать, сколько памяти GPU нужно для обучения модели. 15 минут.

В лабораторной «Оцените FLOPs обучения» вы выясняли, сколько вычислений нужно для обучения модели. Но есть ещё один критический ресурс — память GPU. Прежде чем выполнить хоть одно вычисление, все нужные данные надо загрузить в память видеокарты. Это параметры модели, входные данные и временные величины вроде градиентов.

В этой лабораторной вы соберёте калькулятор памяти, чтобы оценить, сколько всего памяти GPU нужно для обучения модели.

Чему вы научитесь

К концу лабораторной вы сможете:

вычислять, сколько памяти нужно для обучения модели;
объяснять, почему обучать большие модели на слабом оборудовании так трудно.

Задачи

В этой лабораторной вы:

напишете функции, которые считают память под параметры, входные данные, градиенты, состояние оптимизатора и активации;
с их помощью выясните, можно ли обучить модель на 4 миллиарда параметров на обычной GPU с 16 ГБ.

Импорты

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

%%capture
# Ставим специальный пакет курса.
!pip install "git+https://github.com/google-deepmind/ai-foundations.git@main"

from ai_foundations.utils import formatting # Чтобы форматировать оценки памяти.
# Чтобы проверять ваши решения.
from ai_foundations.feedback.course_7 import memory as feedback

# Чтобы форматировать результаты вычислений.
from IPython.display import display, HTML

Следующая ячейка задаёт все константы лабораторной. Они описывают воображаемый сценарий обучения модели Gemma-4B: PARAM_COUNT (число параметров), BYTES_PER_PARAMETER (байт на параметр — для 32-битной точности) и другие детали вроде BATCH_SIZE и MAX_LENGTH, которые нужны, чтобы оценить объём памяти под активации.

Запустите ячейку, чтобы эти константы были доступны в остальной лабораторной.

# Параметры модели (4 миллиарда).
PARAM_COUNT = 4e9

# Точность (байт на параметр).
# В 32-битном формате (FP32) каждый параметр занимает 4 байта.
BYTES_PER_PARAMETER = 4

# Константы для формулы активаций.
BATCH_SIZE = 8
MAX_LENGTH = 1024 # Максимальная длина последовательности.
NUM_LAYERS = 32
EMBEDDING_DIM = 2560

Задание 1. Память под параметры

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

def calculate_param_memory(param_count: float, bytes_per_param: int) -> float:
    """Считает, сколько ГБ памяти нужно, чтобы хранить параметры модели.

    Args:
      param_count: Общее число параметров модели.
      bytes_per_param: Сколько байт занимает один параметр.

    Returns:
      Общий объём памяти под параметры в гигабайтах.
    """
    total_bytes = ...  # Добавьте сюда свой код.

    return formatting.bytes_to_gb(total_bytes)

# @title Запустите эту ячейку, чтобы проверить реализацию
feedback.test_calculate_param_memory(calculate_param_memory)

Задание 2. Память под входные данные

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

Допишите функцию calculate_input_data_memory ниже.

Она должна возвращать число байт, нужное, чтобы хранить вход модели.

Подсказки:

GPU обрабатывает весь батч обучающих примеров сразу. Чтобы найти память под весь входной батч, нужно определить, сколько байт в батче.
В батче batch_size примеров длиной max_length токенов. Вспомните: длина каждой последовательности — max_length, потому что мы дополняем и обрезаем их до max_length.
Один входной токен занимает bytes_per_token_id байт памяти.

Реализовав функцию, запустите две следующие ячейки, чтобы определить её и проверить код.

def calculate_input_data_memory(
    batch_size: int, max_length: int, bytes_per_token_id: int
) -> float:
    """Считает, сколько ГБ памяти нужно, чтобы хранить батч id входных токенов.

    Args:
      batch_size: Число последовательностей в одном батче.
      max_length: Длина каждой последовательности в токенах после дополнения.
      bytes_per_token_id: Сколько байт занимает id одного токена.

    Returns:
      Общий объём памяти под батч входных данных в гигабайтах.
    """
    total_bytes = ...  # Добавьте сюда свой код.

    return formatting.bytes_to_gb(total_bytes)

# @title Запустите эту ячейку, чтобы проверить реализацию.
feedback.test_calculate_input_data_memory(calculate_input_data_memory)

Задание 3. Память под градиенты

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

def calculate_gradient_memory(param_count: float, bytes_per_param: int) -> float:
    """Считает, сколько ГБ памяти нужно, чтобы хранить градиенты.

    Args:
      param_count: Общее число параметров модели.
      bytes_per_param: Сколько байт занимает один параметр.

    Returns:
      Общий объём памяти под градиенты в гигабайтах.
    """
    total_bytes = ...  # Добавьте сюда свой код.

    return formatting.bytes_to_gb(total_bytes)

# @title Запустите эту ячейку, чтобы проверить реализацию.
feedback.test_calculate_gradient_memory(calculate_gradient_memory)

Задание 4. Память под оптимизатор

Современные оптимизаторы вроде Adam хранят ещё и сведения, по которым вычисляется скорость обучения для каждого отдельного параметра. Эти сведения тоже нужно хранить на GPU.

Допишите функцию calculate_optimizer_memory в следующей ячейке.

Для каждого параметра модели оптимизатор вроде Adam хранит сведения, по которым вычисляется скорость обучения именно этого параметра. Значит, нужная память — общее число параметров, умноженное на число байт, которое занимают эти сведения. Adam хранит два значения состояния на параметр, поэтому памяти под эти сведения нужно вдвое больше, чем байт на хранение одного параметра.

Реализовав функцию, запустите две следующие ячейки, чтобы определить её и проверить код.

def calculate_optimizer_memory(param_count: float, bytes_per_param: int) -> float:
    """Считает, сколько ГБ памяти занимает состояние оптимизатора Adam.

    Args:
      param_count: Общее число параметров модели.
      bytes_per_param: Сколько байт занимает один параметр.

    Returns:
      Общий объём памяти под состояние оптимизатора в гигабайтах.
    """
    total_bytes = ...  # Добавьте сюда свой код.

    return formatting.bytes_to_gb(total_bytes)

# @title Запустите эту ячейку, чтобы проверить расчёт для оптимизатора
feedback.test_calculate_optimizer_memory(calculate_optimizer_memory)

Задание 5. Память под активации

При прямом проходе модель должна хранить на GPU результаты каждого вычисления — их называют активациями.

Активация одного токена после прохода через один слой трансформера — это вектор. Его размер равен размерности эмбеддинга.

$\text{память на активацию одного токена} = \text{размерность эмбеддинга} \times \text{байт на параметр}$

При прямом проходе такой вектор активации хранится для каждого токена последовательности.

$\text{память на одну последовательность} = \text{максимальная длина последовательности} \times \text{память на активацию одного токена}$

При обучении активации всех слоёв держат в памяти, потому что они нужны для обратного прохода.

$\text{память на все слои} = \text{число слоёв} \times \text{память на одну последовательность}$

Наконец, всё это происходит одновременно для каждого примера в батче. GPU должна держать активации их всех сразу.

$\text{вся память под активации} = \text{размер батча} \times \text{память на все слои}$

def calculate_activation_memory(
    batch_size: int,
    max_length: int,
    num_layers: int,
    embedding_dim: int,
    bytes_per_param: int,
) -> float:
    """Оценивает по упрощённой формуле, сколько ГБ памяти нужно под активации.

    Args:
      batch_size: Число последовательностей в батче.
      max_length: Длина каждой последовательности в токенах после дополнения.
      num_layers: Число слоёв трансформера в модели.
      embedding_dim: Размер скрытого слоя (размерность эмбеддинга) модели.
      bytes_per_param: Сколько байт занимает одно значение активации.

    Returns:
      Оценка общего объёма памяти под активации в гигабайтах.
    """
    total_bytes = ... # Добавьте сюда свой код.

    return formatting.bytes_to_gb(total_bytes)

# @title Запустите эту ячейку, чтобы проверить реализацию.
feedback.test_calculate_activation_memory(calculate_activation_memory)

Считаем всю память

Теперь, когда у вас есть функции для каждой составляющей, запустите следующую ячейку, чтобы посчитать, сколько всего памяти GPU нужно для Gemma-4B. Ячейка вызывает все функции, которые вы реализовали, и выводит разбивку требований к памяти по составляющим.

# Считаем память каждой составляющей вашими функциями.
param_mem = calculate_param_memory(PARAM_COUNT, BYTES_PER_PARAMETER)
input_data_mem = calculate_input_data_memory(
    BATCH_SIZE, MAX_LENGTH, BYTES_PER_PARAMETER
)
grad_mem = calculate_gradient_memory(PARAM_COUNT, BYTES_PER_PARAMETER)
optim_mem = calculate_optimizer_memory(PARAM_COUNT, BYTES_PER_PARAMETER)
activ_mem = calculate_activation_memory(
    BATCH_SIZE, MAX_LENGTH, NUM_LAYERS, EMBEDDING_DIM, BYTES_PER_PARAMETER
)

total_inference_memory = param_mem + input_data_mem + activ_mem

# Складываем.
total_training_memory = total_inference_memory + grad_mem + optim_mem

# Показываем результаты.
display(HTML("<h3>--- GPU memory consumption breakdown ---</h3>"))
display(HTML("<h4>------ During training and inference ------</h4>"))

# Для каждой составляющей — готовая функция вывода.
formatting.display_memory("model parameters", param_mem)
formatting.display_memory("input data batch", input_data_mem, decimal_places=6)
formatting.display_memory("activations", activ_mem)
display(HTML("<h4>------ During training only ------</h4>"))
formatting.display_memory("gradients", grad_mem)
formatting.display_memory("optimizer states (Adam)", optim_mem)

# Выводим разделитель и итог.
display(
    HTML(
        f"<h3>Total estimated GPU memory required during inference: "
        f"{total_inference_memory:.2f} GB</h3>"
    )
)
display(
    HTML(
        f"<h3>Total estimated GPU memory required during training or "
        f"fine-tuning: {total_training_memory:.2f} GB</h3>"
    )
)

Посмотрите на итоговую оценку памяти GPU в ячейке выше. В Google Colab обычно доступна GPU, у которой памяти немногим меньше 16 ГБ.

Как показывает ваш расчёт, и для инференса модели Gemma на 4B, и для обучения или дообучения модели Gemma на 4 миллиарда параметров с такими настройками нужно больше 16 ГБ. В нынешнем виде все нужные данные на эту GPU не поместить. Как вы видели в прошлой лабораторной, попытка загрузить или обучить эту модель с такими настройками или полностью её дообучить сразу упадёт с ошибкой нехватки памяти.

К счастью, такое ограничение памяти — частая проблема, и для неё придумано несколько решений. В следующих заданиях вы познакомитесь с несколькими приёмами, которые повышают эффективность использования оборудования. Они позволят дообучать более крупные модели на той GPU, что у вас есть.

Функция bytes_to_gb делит на $1024^3$ , поэтому «ГБ» здесь — гибибайты. При 4 байтах на число:

Составляющая	Формула	ГБ
Параметры	$4 \cdot 10^9 \times 4$	14,90
Входной батч	$8 \times 1024 \times 4$	0,000031
Активации	$8 \times 1024 \times 32 \times 2560 \times 4$	2,50
Градиенты	$4 \cdot 10^9 \times 4$	14,90
Состояние Adam	$2 \times 4 \cdot 10^9 \times 4$	29,80
Инференс	параметры + вход + активации	17,40
Обучение	всё вместе	62,10

Обратите внимание: на обучение уходит примерно $4 + 4 + 8 = 16$ байт на параметр, а самая большая статья — не веса, а состояние оптимизатора. Ровно эту бухгалтерию (с мастер-весами для смешанной точности) ведёт урок «Смешанная точность». Формула активаций здесь сильно упрощена: в настоящем трансформере на слой хранится не один вектор на токен, а несколько (входы внимания, MLP, нормализаций), так что реальная цифра больше.

Снижаем требования к памяти с bfloat16

В статье «Умные числа» вы узнали о формате bfloat16, который придумали специально, чтобы снизить требования моделей глубокого обучения к памяти почти без потери качества. В последнем задании этой лабораторной посчитайте требования к памяти модели Gemma-4B, если параметры, градиенты и активации хранятся как числа bfloat16, а не как 32-битные числа с плавающей точкой.

# Число байт на параметр — 2: bfloat16 хранит каждый параметр, градиент
# и активацию в двух байтах.
BYTES_PER_PARAMETER = 2

# ... дальше — те же вызовы функций и тот же вывод, что в ячейке выше.

Как показывает этот расчёт, представление чисел в bfloat16 вдвое сокращает требования к памяти. Например, это позволяет запускать инференс модели на 4 миллиарда параметров на GPU с 16 ГБ. Вы воспользуетесь этим приёмом в следующей лабораторной: загрузите и дообучите модель Gemma-4B, храня параметры, градиенты и состояние оптимизатора в этом более экономном формате.

Все строки ровно вдвое меньше: параметры 7,45; вход 0,000015; активации 1,25; градиенты 7,45; Adam 14,90 ГБ. Инференс — 8,70 ГБ (помещается в 16 ГБ), обучение всех параметров — 31,05 ГБ (не помещается в 16 ГБ и впритык для карты на 32 ГБ, без запаса на служебные буферы). Поэтому в лабораторной 7.5 к bfloat16 добавляют ещё и LoRA. На практике моменты Adam обычно оставляют в FP32 ради устойчивости — это тоже обсуждается в уроке о смешанной точности.

Итоги

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

параметры модели;
входные данные;
градиенты;
состояние оптимизатора;
активации.

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

Решения

Задания 1–5

def calculate_param_memory(param_count: int, bytes_per_param: int) -> float:
    """Память под параметры модели в ГБ."""
    total_bytes = param_count * bytes_per_param
    return formatting.bytes_to_gb(total_bytes)


def calculate_input_data_memory(
    batch_size: int, max_length: int, bytes_per_token_id: int
) -> float:
    """Память под батч id входных токенов в ГБ (по умолчанию 4 байта на id)."""
    total_bytes = batch_size * max_length * bytes_per_token_id
    return formatting.bytes_to_gb(total_bytes)


def calculate_gradient_memory(param_count: int, bytes_per_param: int) -> float:
    """Память под градиенты в ГБ."""
    total_bytes = param_count * bytes_per_param
    return formatting.bytes_to_gb(total_bytes)


def calculate_optimizer_memory(param_count: int, bytes_per_param: int) -> float:
    """Память под состояние оптимизатора Adam в ГБ."""
    total_bytes = 2 * param_count * bytes_per_param
    return formatting.bytes_to_gb(total_bytes)


def calculate_activation_memory(
    batch_size: int,
    max_length: int,
    num_layers: int,
    embedding_dim: int,
    bytes_per_param: int,
) -> float:
    """Память под активации в ГБ (упрощённая формула)."""
    total_bytes = batch_size * max_length * num_layers * embedding_dim * bytes_per_param
    return formatting.bytes_to_gb(total_bytes)

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

Словарик

English	Русский
trade-off	компромисс
greedy sampling	жадная генерация (самый вероятный токен)
FLOP / FLOPs	операция с плавающей точкой / число таких операций
FLOPS	операций в секунду (скорость)
multiply-accumulate	«умножить и накопить»
forward / backward pass	прямой / обратный проход
back-of-the-envelope estimate	оценка «на салфетке»
scientific notation	экспоненциальная запись
logarithmic scale	логарифмическая шкала
memory wall	стена памяти
out-of-memory (OOM), `RESOURCE_EXHAUSTED`	нехватка памяти
parameters, gradients, optimizer states, activations	параметры, градиенты, состояние оптимизатора, активации
precision, 32-bit floating point (FP32)	точность, 32-битное число с плавающей точкой
bfloat16	bfloat16 (16-битный формат «мозговой» точности)

Перевод выполнен по лицензии оригинала (CC-BY-4.0 (тексты), Apache-2.0 (код)); автор оригинала — Google DeepMind и UCL. Нашли неточность — сверяйтесь с оригиналом по ссылке выше. · обновлено 2026-10-05