AI Research Foundations, курс 7, часть 2: bfloat16 и накопление градиентов

Перевод лабораторных 7.5–7.6 курса Google DeepMind «Accelerate Your Model»: загрузка Gemma 4B в bfloat16 и дообучение с LoRA на одной GPU, затем нехватка памяти при большом батче и накопление градиентов в Keras.

Вторая часть курса «Accelerate Your Model» («Ускорьте свою модель»). В первой части вы посчитали, что Gemma-4B в 32-битном формате не помещается даже для инференса на 16-гигабайтную T4, а обучение требует около 62 ГБ. Здесь два приёма, которые эту стену обходят: bfloat16 вдвое уменьшает каждое число, а накопление градиентов (gradient accumulation) позволяет получить эффект большого батча, не держа его в памяти целиком.

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

Лабораторная	Время	Когда делать	Что нужно
7.5. Дообучите модель в bfloat16	25 мин + ≈20 мин обучения	Вместе с уроком «Смешанная точность — FP32, BF16, FP8»	GPU, Kaggle
7.6. Примените накопление градиентов	30 мин + ≈20 мин обучения	После урока «Цикл обучения LLM» (там накопление градиентов на PyTorch)	GPU, Kaggle

Гиперпараметры оптимизатора в этих лабораторных — AdamW с $\beta_2 = 0{,}95$ , обрезкой нормы градиента и weight decay — стандартный рецепт устойчивого обучения больших моделей; почему именно так, объясняет урок «Muon и стабильность обучения» и Tuning Playbook.

flowchart LR
A["Gemma-4B в FP32: 15 ГБ весов"] -->|"dtype=bfloat16"| B["7,5 ГБ весов"]
B -->|"+ LoRA rank 4"| C["Дообучение с батчем 1 помещается"]
C -->|"batch_size=4"| D["OOM: активации ×4"]
D -->|"gradient_accumulation_steps=8"| E["Эффективный батч 8 при физическом 1"]

Как лабораторные 7.5 и 7.6 обходят стену памяти

Лабораторная 7.5. Дообучите модель в bfloat16

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

Узнайте, как успешно обучить большую модель на одной GPU с помощью типа данных bfloat16. 25 минут.

Обзор

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

В этой лабораторной вы научитесь загружать параметры модели в типе данных bfloat16. Объединив этот приём с LoRA, вы сможете дообучить модель на 4B на доступной вам GPU.

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

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

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

Задачи

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

загрузите модель Gemma-4B с точностью bfloat16;
дообучите модель генерировать учебные карточки, как в курсе 5 «Дообучите свою модель»;
поразмышляете о качестве дообученной модели.

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

Импорты

В этой лабораторной вы будете загружать и дообучать модель 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_ALLOCATOR"] = "platform"

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.

Загружаем и обрабатываем данные

Как и в прошлых лабораторных, следующая ячейка определяет функцию format_question, которая форматирует промпт. Она также загружает набор Africa Galore QA и обрабатывает вопросы так, чтобы на них можно было дообучить модель генерировать ответы для карточек, как в прошлых курсах.

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

В прошлой лабораторной, «Упираемся в стену», загрузить Gemma-4B не удалось: параметры модели (в 32-битном формате по умолчанию) были слишком велики для памяти GPU.

Решение — загрузить модель в более экономном числовом формате. Вы воспользуетесь bfloat16: он использует 16 бит вместо 32 и, по сути, вдвое уменьшает память, нужную под веса модели.

В следующей ячейке загрузите параметры модели Gemma-4B в формате bfloat16.

В Keras параметры загружаются в bfloat16 так:

model = keras_hub.models.Gemma3CausalLM.from_preset(preset=<MODEL_NAME>, dtype="bfloat16")

Для модели Gemma-4B значение <MODEL_NAME> — "gemma3_4b_text".

Загрузка модели займёт 2–3 минуты.

# Загружаем модель Gemma3-4B для Keras с точностью bfloat16.
model = ... # Добавьте сюда свой код.
model.summary()

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

# Берём первый блок трансформера.
first_transformer_block = model.backbone.get_layer("decoder_block_0")

# Берём слой внимания.
attention_layer = first_transformer_block.attention

# Берём матрицу весов проекции запросов и проверяем её тип данных.
query_weights = attention_layer.query_dense.kernel

print(f"The model's weight precision is {query_weights.dtype}.")

Ожидаемый вывод: The model's weight precision is bfloat16.

Включаем LoRA

Даже если хранить параметры модели в bfloat16, полностью дообучить её на доступной GPU не получится. Поэтому модель снова придётся обучать с LoRA. Вспомните курс 5 «Дообучите свою модель»: при дообучении с LoRA обновляются только параметры низкоранговых матриц. Они составляют малую долю всех параметров. Значит, в памяти GPU нужно хранить градиенты, активации и состояние оптимизатора лишь для небольшого числа параметров. Это резко снижает общие требования к памяти. Именно благодаря этому модель на 4B можно дообучить на T4.

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

model.backbone.enable_lora(rank=4)
model.summary()

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

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

Обучение в формате пониженной точности вроде bfloat16 часто требует более аккуратного подбора гиперпараметров, чем обычное 32-битное. Числа менее точны, поэтому обучение может быть чувствительнее, и важно выбрать настройки, при которых оно идёт устойчиво.

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

# Задаём гиперпараметры.
optimizer = keras.optimizers.AdamW(
    learning_rate=5e-5,
    weight_decay=0.01,
    beta_1=0.9,
    beta_2=0.95,
    epsilon=1e-6,
    clipnorm=0.5,
)

# Задаём число эпох.
num_epochs = 4

# Задаём максимальную длину.
model.preprocessor.sequence_length = 400

# Компилируем модель с оптимизатором.
model.compile(
    optimizer=optimizer,
)

weight_decay=0.01 — небольшое «стягивание» весов к нулю (регуляризация). beta_1, beta_2 — скорости забывания скользящих средних градиента и его квадрата в Adam; $\beta_2 = 0{,}95$ вместо стандартных 0,999 делает оценку разброса градиентов «короче памятью» и обучение устойчивее к редким всплескам. epsilon=1e-6 — маленькая добавка в знаменателе шага Adam, которая защищает от деления на почти ноль (чуть больше, чем 1e-7 по умолчанию в Keras, — для запаса устойчивости). clipnorm=0.5 обрезает градиент, если его норма больше 0,5, и гасит отдельные «взрывы». Подробнее — в уроках «Momentum и Adam» и «Muon и стабильность обучения».

Следим за ходом обучения

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

Следующая ячейка создаёт callback Keras — вспомогательную функцию, которая автоматически выполняет код в конце каждой эпохи. Здесь она генерирует и печатает ответы на три контрольных промпта. Так можно следить за прогрессом модели в реальном времени.

Три контрольных промпта проверяют разные стороны обучения:

«What is Kente cloth?» («Что такое ткань кенте?») — выучивает ли модель конкретные знания из набора для дообучения;
«What is Kilimanjaro?» («Что такое Килиманджаро?») — обобщение: в наборе для дообучения есть «Mount Kilimanjaro» («гора Килиманджаро»), а здесь слово «Mount» опущено, чтобы усложнить задачу;
«What is Tokyo?» («Что такое Токио?») — о Токио и Японии в наборе для дообучения нет ничего, поэтому промпт проверяет, сохраняет ли модель знания из предобучения, выучивая новую задачу.

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

# Список промптов, которые проверяем после каждой эпохи.
test_prompts = [
    "What is Kente cloth?",
    "What is Kilimanjaro?",
    "What is Tokyo?"
]

class GenerationMonitor(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        print(f"\n--- Generations after epoch {epoch + 1} ---")
        for prompt in test_prompts:
            # Правильно форматируем промпт для модели.
            formatted_prompt = format_question(prompt)

            # Генерируем и печатаем ответ.
            output = self.model.generate(formatted_prompt, max_length=150)
            print(output)
            print("-" * 20)

Теперь запустите следующие ячейки, чтобы дообучить модель с LoRA в формате bfloat16. Обратите внимание: ошибки нехватки памяти больше нет.

Сравните ответы после каждой эпохи. Как они выглядят?

На T4 одна эпоха займёт около пяти минут, четыре эпохи — около 20 минут.

# Создаём экземпляр callback для наблюдения.
generation_callback = GenerationMonitor()

# Обучаем модель.
history = model.fit(
    data,
    epochs=num_epochs,
    batch_size=1,
    callbacks=[generation_callback],
)

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

В прошлой лабораторной, «Упираемся в стену», уже попытка загрузить модель в полной 32-битной точности сразу закончилась нехваткой памяти. Теперь одно-единственное изменение — загрузка модели в bfloat16 вместо 32-битных чисел — позволяет не только загрузить модель, но и дообучить её с LoRA.

Вывод обучения показывает и то, как модель через несколько эпох научилась отвечать на вопросы в нужном формате карточки. Обратите внимание на неустойчивость в эпохе 1. На промпт про кенте модель повторяет случайный токен («KMnO4»). На промпт про Токио модель всё время генерирует списки «Category:». Но после этого генерации резко улучшаются. Модель выдаёт всё более связный текст в формате карточки. Итоговые генерации в эпохе 4 качественные, фактически точные и правильные по стилю. Это показывает, что bfloat16 не портит качество обученной модели.

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

Итоги

Загрузка параметров модели в формате пониженной точности вроде bfloat16 позволяет использовать и дообучать более крупные модели при ограниченных ресурсах GPU. Как вы увидели, с представлением bfloat16 удалось загрузить в память все параметры Gemma-4B и даже дообучить модель на T4, доступной в Colab. Получилась модель, которая выдаёт качественные ответы благодаря мощной базовой модели на 4 миллиарда параметров, взятой за основу дообучения.

Решения

Задание 1

model = keras_hub.models.Gemma3CausalLM.from_preset("gemma3_4b_text", dtype="bfloat16")

В Hugging Face Transformers аналог — AutoModelForCausalLM.from_pretrained(..., torch_dtype=torch.bfloat16), а LoRA включается через библиотеку PEFT (get_peft_model(model, LoraConfig(r=4, ...))). Целиком этот путь пройден в уроке «Практика — LoRA-дообучение Qwen». Учтите: чистое bfloat16 для всех весов, как здесь, нормально работает с LoRA, потому что замороженные веса не обновляются; для полного обучения используют смешанную точность с мастер-весами в FP32 (урок «Смешанная точность»).

Лабораторная 7.6. Примените накопление градиентов

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

Научитесь имитировать большой размер батча на GPU с ограниченной памятью с помощью накопления градиентов. 30 минут.

Обзор

Раньше вы видели, что больший размер батча может сделать обучение устойчивее и повысить эффективность использования GPU. Но вы также убедились, что память GPU — критическое узкое место. Что же делать, если хочется получить выгоды большого батча, а памяти GPU на него не хватает?

В этой лабораторной вы столкнётесь именно с этой проблемой. Сначала вы попробуете обучить модель Gemma-4B с большим батчем. Затем научитесь использовать накопление градиентов (gradient accumulation), чтобы добиться того же результата на своей GPU с ограниченными ресурсами.

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

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

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

Задачи

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

попробуете дообучить Gemma-4B с размером батча 2 и увидите ошибку памяти (в самом задании ниже и в решении батч равен 4 — ориентируйтесь на него; примечание переводчика);
перенастроите оптимизатор на накопление градиентов;
успешно дообучите модель, имитируя больший размер батча — 8.

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

Импорты и данные

Ячейки импорта и подготовки данных — те же, что в лабораторной 7.5 выше: те же пакеты, те же переменные окружения (KERAS_BACKEND, XLA_FLAGS, XLA_PYTHON_CLIENT_ALLOCATOR), та же функция format_question и тот же словарь data с ключами prompts и responses.

Задание 1. Дообучаем с большим размером батча

Как вы уже знаете, загрузка данных батчами позволяет:

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

Что будет, если увеличить размер батча при обучении Gemma-4B?

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

model.backbone.enable_lora(rank=4)
model.summary()

# Задаём гиперпараметры.
optimizer = keras.optimizers.AdamW(
    learning_rate=1e-4,
    weight_decay=0.01,
    beta_1=0.9,
    beta_2=0.95,
    epsilon=1e-6,
    clipnorm=0.5,
)

# Задаём число эпох.
num_epochs = 4

# Задаём максимальную длину.
model.preprocessor.sequence_length = 400

# Компилируем модель с оптимизатором.
model.compile(
    optimizer=optimizer,
)

# Список промптов, которые проверяем после каждой эпохи.
test_prompts = [
    "What is Kente cloth?",
    "What is Kilimanjaro?",
    "What is Tokyo?"
]

class GenerationMonitor(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        print(f"\n--- Generations after epoch {epoch + 1} ---")
        for prompt in test_prompts:
            # Правильно форматируем промпт для модели.
            formatted_prompt = format_question(prompt)

            # Генерируем и печатаем ответ.
            output = self.model.generate(formatted_prompt, max_length=150)
            print(output)
            print("-" * 20)

# Создаём экземпляр callback для наблюдения.
generation_callback = GenerationMonitor()

# Обучаем модель.
history = model.fit(
    data,
    epochs=num_epochs,
    batch_size=..., # Добавьте сюда свой код.
    callbacks=[generation_callback],
)

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

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

Это снова ошибка нехватки памяти. Она означает, что при выполнении шага обучения у GPU закончилась её собственная память.

Ключевая строка — в конце:

Out of memory while trying to allocate 52428800 bytes.

Она говорит, что для этого шага GPU понадобилось выделить ещё больше 52 МБ памяти, но свободного места не хватило.

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

Перезапустите сеанс Colab

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

Прежде чем переходить к заданию 2, перезапустите сеанс Colab:

В верхнем меню нажмите Runtime («Среда выполнения»).
Выберите Restart session («Перезапустить сеанс»).
Во всплывающем окне подтвердите кнопкой Yes.

Сеанс Colab перезапустится с чистым окружением. Это стирает все переменные и импортированные библиотеки, поэтому перед продолжением нужно снова выполнить импорты и подготовительные ячейки. В ноутбуке они для удобства скопированы ниже, так что можно запускать все ячейки начиная с раздела «Задание 2». (В Jupyter на своей машине — меню Kernel → Restart Kernel; повторяющиеся ячейки импорта и данных здесь не дублируются.)

Задание 2. Дообучаем с накоплением градиентов

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

Приём работает так: прямой и обратный проходы выполняются на нескольких маленьких батчах, а их градиенты накапливаются (суммируются). Только обработав заданное число маленьких батчей, оптимизатор обновляет веса модели по объединённому градиенту. Это даёт такой же устойчивый сигнал обучения, как большой батч, но без больших затрат памяти.

flowchart LR
B1["Батч 1: forward + backward"] --> S["Сумма градиентов"]
B2["Батч 2: forward + backward"] --> S
B3["…"] --> S
B8["Батч 8: forward + backward"] --> S
S --> U["Один шаг AdamW"]
U --> Z["Обнулить сумму"]

Накопление градиентов за 8 шагов: в памяти всё время активации только одного маленького батча

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

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

model.backbone.enable_lora(rank=4)
model.summary()

В следующей ячейке настройте оптимизатор AdamW на накопление градиентов.

Для этого добавьте прямо в оптимизатор аргумент gradient_accumulation_steps. Настройте его так, чтобы имитировать размер батча 8 (то есть в 8 раз больше физического размера батча).

optimizer = keras.optimizers.AdamW(
   learning_rate=1e-4,
   ...,
   gradient_accumulation_steps=<NUM_STEPS>
)

# Задаём гиперпараметры.
optimizer = keras.optimizers.AdamW(
    # Умножаем learning rate под новый эффективный размер батча.
    learning_rate=8e-4,
    weight_decay=0.01,
    beta_1=0.9,
    beta_2=0.95,
    epsilon=1e-6,
    clipnorm=0.5,
    # Добавьте сюда свой код.
)

# Задаём число эпох.
num_epochs = 4

# Задаём максимальную длину.
model.preprocessor.sequence_length = 400

# Компилируем модель с оптимизатором.
model.compile(
    optimizer=optimizer,
)

Дальше — та же ячейка с callback GenerationMonitor, что выше, и обучение с физическим батчем 1:

# Создаём экземпляр callback для наблюдения.
generation_callback = GenerationMonitor()

# Обучаем модель.
history = model.fit(
    data,
    epochs=num_epochs,
    batch_size=1,
    callbacks=[generation_callback],
)

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

По ходу эпох генерации, скорее всего, становились всё более связными и по теме. После первой эпохи ещё есть некоторая бессвязность. Но заметьте: модель не зацикливается на одном токене, как «KMnO4» в прошлом запуске из лабораторной «Дообучите модель в bfloat16». К концу второй эпохи ответы уже беглые и фактически верные.

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

Это «правило линейного масштабирования»: если эффективный батч вырос в $k$ раз, шум градиента уменьшился, и можно делать шаг примерно в $k$ раз больше — здесь $1 \cdot 10^{-4} \times 8 = 8 \cdot 10^{-4}$ . Правило приближённое и работает не всегда; при смене размера батча гиперпараметры, строго говоря, нужно подбирать заново — подробно об этом говорит Tuning Playbook. Заметьте ещё, что сравнение с лабораторной 7.5 не совсем честное: там learning rate был $5 \cdot 10^{-5}$ , а здесь — $8 \cdot 10^{-4}$ , то есть отличается не только размер батча.

Итоги

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

Решения

Задание 1

# Обучаем модель.
history = model.fit(
    data,
    epochs=num_epochs,
    batch_size=4,
    callbacks=[generation_callback],
)

Задание 2

# Задаём гиперпараметры.
optimizer = keras.optimizers.AdamW(
    learning_rate=8e-4, # Умножаем learning rate под новый эффективный размер батча.
    weight_decay=0.01,
    beta_1=0.9,
    beta_2=0.95,
    epsilon=1e-6,
    clipnorm=0.5,
    gradient_accumulation_steps=8,
)

То же на PyTorch

Не часть оригинала. В PyTorch накопление градиентов пишут прямо в цикле обучения: loss.backward() сам складывает градиенты в .grad, нужно только делить потери на число шагов и вызывать optimizer.step() раз в accum_steps батчей. Подробно — в уроке «Цикл обучения LLM».

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Linear(16, 1)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
accum_steps = 8                                   # эффективный батч = 8 × физический

X, y = torch.randn(64, 16), torch.randn(64, 1)
opt.zero_grad()
for step in range(64):                            # физический батч = 1 пример
    xb, yb = X[step:step + 1], y[step:step + 1]
    loss = nn.functional.mse_loss(model(xb), yb) / accum_steps  # делим, чтобы получить среднее
    loss.backward()                               # градиенты складываются в .grad
    if (step + 1) % accum_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)  # как clipnorm=0.5
        opt.step()                                # один шаг на 8 примеров
        opt.zero_grad()
print("шагов оптимизатора:", 64 // accum_steps)   # 8

На одной RTX 4080 с 32 ГБ Gemma-4B в bfloat16 с LoRA поместится и с batch_size=4 — «стену» из лабораторной 7.6 вы, скорее всего, не увидите. Чтобы её воспроизвести, увеличьте batch_size до 16–32 или sequence_length до 2048. С четырьмя картами появляется третий способ увеличить батч — data parallelism: каждая карта считает свой кусок батча, градиенты усредняются между картами (урок «Четыре видеокарты: DDP и FSDP на практике»). Накопление градиентов и DDP часто совмещают: эффективный батч = физический × шаги накопления × число карт.

Словарик

English	Русский
bfloat16 (brain floating point)	bfloat16, 16-битный формат с плавающей точкой
precision, low-precision format	точность, формат пониженной точности
dtype	тип данных
memory footprint	объём занимаемой памяти
resource-constrained GPU	GPU с ограниченными ресурсами
weight decay	затухание весов (weight decay)
gradient clipping (`clipnorm`)	обрезка градиента по норме
callback	функция обратного вызова (callback)
physical / effective batch size	физический / эффективный размер батча
gradient accumulation	накопление градиентов
restart session	перезапуск сеанса
linear scaling rule	правило линейного масштабирования

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