Как дообучить модель 70B: требования к GPU и порядок работы

Дообучение (fine-tuning) модели 70B можно провести и на одном GPU, и на целом 8-GPU-сервере, и оба варианта верны: всё зависит от того, обновляете ли вы все веса, обучаете адаптер LoRA поверх 16-битных весов или обучаете его поверх 4-битной базовой модели с помощью QLoRA. Нужный объём видеопамяти (VRAM) — примерно от 40 ГБ до намного больше 1 ТБ. В этом руководстве объясняется, откуда берутся эти цифры, какие GPU нужны для каждого метода и как пройти весь путь от подготовки данных до инференса. В примерах используется Llama 3.1 70B: 80 слоёв, скрытая размерность 8 192 и 8 KV-голов.

Требования к GPU для дообучения модели 70B: коротко о главном

Если говорить о бюджете, эти конфигурации стоят 2 190 $ в месяц за один H100, 3 090 $ за один H200, 4 590 $ за один B200, 5 190 $ за один B300, 16 900 $ за 8× H100, 23 900 $ за 8× H200, 34 900 $ за 8× B200 и 39 900 $ за 8× B300. К тому же месяца хватает на несколько запусков, а дообучение обычно так и проходит. Дальше в руководстве эти цифры разобраны подробно; общий метод, на котором они основаны, описан в нашем руководстве по видеопамяти для LLM (больших языковых моделей).

Полное дообучение: почему нужно больше 1 ТБ

При стандартном обучении со смешанной точностью и AdamW на каждый параметр хранится пять чисел:

Схемы, в которых градиенты или коммуникационные буферы хранятся в FP32, дают ближе к 20 байтам на параметр, то есть 1,4 ТБ. В любом случае состояние обучения превышает 1 ТБ ещё до того, как сохранена хотя бы одна активация.

Столько не вмещает ни один GPU, а обычный параллелизм по данным не помогает, потому что каждый GPU хранит полную копию. Нужно шардирование: PyTorch FSDP или DeepSpeed ZeRO stage 3 распределяют веса, градиенты и состояния оптимизатора по всем GPU и собирают веса каждого слоя только на время его вычисления. На восьми GPU каждый хранит 140 ГБ состояния при 16 байтах на параметр или 175 ГБ при 20. Эта арифметика и объясняет выбор оборудования:

• 8× H100 (всего 640 ГБ) не вмещают это состояние в пределах одного сервера.
• 8× H200 (по 141 ГБ) вмещают по 140 ГБ состояния на GPU, но для активаций не остаётся ничего. Схема работает, только если перенести мастер-копию весов в FP32 и моменты Adam в память CPU (около 840 ГБ — это помещается в 2 ТБ оперативной памяти сервера) или ужать их с помощью 8-битных состояний оптимизатора, если ваша схема шардирования это поддерживает. Выгрузка замедляет каждый шаг, потому что оптимизатор работает на CPU, а состояние передаётся по PCIe.
• 8× B200 (по 180 ГБ) при 16 байтах на параметр оставляют около 40 ГБ на GPU для активаций и буферов — при чекпоинтинге активаций (activation checkpointing) этого достаточно для работы. При схеме ближе к 20 байтам на параметр почти ничего не остаётся, и здесь тоже понадобится выгрузка.
• 8× B300 (по 288 ГБ) оставляют намного больше 100 ГБ на GPU — этого хватает для более длинных последовательностей и более крупных микробатчей без всяких ухищрений.

Ещё один путь — два сервера, соединённые быстрой сетью, но полное дообучение модели 70B редко стоит пробовать первым. Для большинства задач обучения следованию инструкциям (instruction tuning) и адаптации к предметной области начинайте с LoRA, а к полному дообучению переходите, только когда у вас есть данные и подтверждение того, что узкое место — именно адаптер. Если полное дообучение всё же понадобится, наши 8-GPU-серверы — это bare metal с NVLink между всеми восемью GPU.

Память под активации и длина последовательности

Активации — вторая половина счёта за память, и именно она растёт вместе с вашими данными. Во время прямого прохода каждый слой сохраняет промежуточные тензоры для обратного прохода. Для Llama-подобной модели 70B в BF16 это примерно 25 МБ на токен по всем 80 слоям, так что одной последовательности из 8 192 токенов понадобилось бы около 200 ГБ на её GPU. Шардирование этого не уменьшает, потому что каждый GPU хранит активации своего микробатча.

Эту проблему решает чекпоинтинг активаций, который также называют gradient checkpointing. Вместо того чтобы хранить всё, каждый слой сохраняет только свой вход — один вектор из 8 192 значений BF16 (16 КБ) на токен — и пересчитывает остальное во время обратного прохода. Это около 1,3 МБ на токен для 80 слоёв, или примерно 11 ГБ на одну последовательность из 8 192 токенов, плюс несколько гигабайт рабочего пространства для пересчитываемого слоя. Цена — примерно один дополнительный прямой проход вычислений на каждом шаге.

Ещё две вещи часто застают врасплох. Выходные логиты по словарю Llama 3.1 из 128 256 токенов занимают около 4,2 ГБ в FP32 на одну последовательность из 8 192 токенов, не считая их градиентов; реализации функции потерь со слиянием операций (fused) или с обработкой по частям (chunked) позволяют не материализовать их все сразу. А длина последовательности — главный рычаг, которым вы управляете: упаковывайте короткие примеры вместе, держите микробатч на уровне одной-двух последовательностей и используйте накопление градиентов, чтобы получить нужный эффективный размер батча.

LoRA в BF16: нижнюю границу задают замороженные веса

LoRA замораживает базовую модель и обучает небольшие матрицы низкого ранга, добавленные к выбранным проекционным слоям. Базовые веса всё равно должны находиться в видеопамяти: около 140 ГБ в BF16. Адаптеры по сравнению с ними крошечные. LoRA ранга 16 на всех семи проекционных матрицах каждого из 80 слоёв добавляет около 207 млн обучаемых параметров — примерно 0,3% модели, или около 3,3 ГБ вместе с их градиентами и состояниями оптимизатора.

Активации почти не уменьшаются: градиенты всё равно проходят в обратном направлении через замороженные слои, чтобы дойти до адаптеров в первых слоях, поэтому не отключайте чекпоинтинг активаций. Итого для одной последовательности из 8 192 токенов — около 160 ГБ. Отсюда и требования к оборудованию:

• 2× 80 ГБ — нижняя граница: всего 160 ГБ, этого хватает только при коротких последовательностях.
• 8× H100 или 8× H200 — с запасом. При шардировании на восемь GPU замороженные веса занимают по 17,5 ГБ на GPU, так что остаётся место и для длинных последовательностей, и для параллелизма по данным на всех восьми.
• Один GPU тоже подходит — без настройки шардирования, но со скоростью одного GPU: B200 со 180 ГБ вмещает задачу с запасом около 20 ГБ при 8 192 токенах, а NVIDIA B300 с 288 ГБ — с большим запасом.

В Hugging Face PEFT типичная конфигурация для Llama-подобных моделей охватывает все линейные проекции:

from peft import LoraConfig

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    task_type="CAUSAL_LM",
)

Авторы QLoRA обнаружили, что в их экспериментах сравняться с полным дообучением удавалось, только если применять адаптеры ко всем линейным слоям, а не только к слоям внимания. Скорость обучения (learning rate) около 1e-4 — обычная отправная точка для LoRA на больших моделях: это примерно в десять раз больше, чем при полном дообучении.

QLoRA: модель 70B на одном GPU

QLoRA хранит замороженную базовую модель в 4-битном формате NF4 и обучает поверх неё 16-битные адаптеры LoRA, деквантуя веса на лету. Базовая модель 70B тогда занимает около 35 ГБ, а на практике — около 40 ГБ, если добавить константы квантования, а также слой эмбеддингов и выходной слой, которые остаются 16-битными. Добавьте 3–4 ГБ на адаптеры и состояния их оптимизатора и примерно 14 ГБ активаций с чекпоинтингом на последовательность из 8 192 токенов — и задача помещается на один NVIDIA H100 с 80 ГБ. H200 со 141 ГБ даёт место для более длинных последовательностей или более крупных микробатчей.

В оригинальной статье о QLoRA модель 65B дообучили на одном GPU с 48 ГБ. Llama 3.1 70B немного крупнее, а словарь у неё в четыре раза больше, поэтому закладывайте 80 ГБ. С Transformers и bitsandbytes загрузка базовой модели в 4-битном формате выглядит так:

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-70B-Instruct",
    quantization_config=bnb_config,
    device_map="auto",
)

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

Как дообучить модель 70B: практический порядок работы

Оборудование — самая простая часть. Если дообучение не удаётся, то чаще всего из-за данных или оценки, а не из-за памяти.

1. Подготовьте данные. Оформите примеры как диалоги и преобразуйте их в текст с помощью собственного шаблона чата модели (tokenizer.apply_chat_template в Transformers). Удалите дубликаты, отфильтруйте некачественные примеры и уберите персональные данные, на которых обучать модель не следует. Отложите выборку для оценки, на которой модель никогда не обучается. Измерьте длину примеров в токенах: самые длинные из оставленных последовательностей определяют объём памяти под активации.
2. Выберите инструмент. Hugging Face TRL вместе с PEFT даёт готовые тренеры для дообучения с учителем и для обучения на основе предпочтений, а для работы на нескольких GPU использует Accelerate, FSDP или DeepSpeed. Axolotl упаковывает те же компоненты в YAML-конфигурации. torchtune предлагает нативные для PyTorch рецепты полного дообучения, а также LoRA и QLoRA. Все три инструмента поддерживают описанные выше методы; выберите тот, чей стиль конфигураций ваша команда готова поддерживать.
3. Зафиксируйте окружение. Используйте контейнер или виртуальное окружение с фиксированными версиями PyTorch, Transformers, PEFT и вашего тренера, чтобы повторный запуск через месяц вёл себя так же.
4. Начните со smoke-теста в малом масштабе. Прогоните тот же пайплайн с моделью 8B в течение нескольких сотен шагов. Так вы выявите ошибки в шаблоне, маскировании и данных до того, как они обойдутся вам в долгий запуск.
5. Запустите и наблюдайте. Запускайте задачи на нескольких GPU командой torchrun --nproc_per_node=8 train.py или собственным лаунчером вашего инструмента внутри tmux, чтобы обрыв SSH-соединения не прервал запуск. Следите за функцией потерь на обучающей и оценочной выборках, памятью GPU и временем шага; скачок функции потерь или её рост на оценочной выборке — повод остановиться досрочно.

Об обмене данными между GPU на 8-GPU-серверах — в нашем руководстве по работе с несколькими GPU и NCCL.

Чекпоинты на NVMe

Чекпоинты занимают много места, и хранить их лучше на быстрых локальных дисках:

• Полное дообучение: чекпоинт для возобновления обучения, включающий состояния оптимизатора, весит столько же, сколько состояние обучения, — 1,1–1,4 ТБ. Сохраняйте его по шардам: каждый процесс (rank) параллельно записывает свою часть.
• Только веса модели: около 140 ГБ в BF16 — для оценки и инференса.
• Адаптер LoRA: около 0,4 ГБ при ранге 16 в BF16, так что хранить можно много адаптеров.

В наших серверах HGX установлено 8× 3,84 ТБ NVMe — места хватит для датасетов, базовой модели и нескольких полных чекпоинтов. В тарифе с одним H100 — 2 ТБ, этого с лихвой хватает для QLoRA, где самый большой файл — скачанная базовая модель в BF16. Храните последние несколько чекпоинтов локально, лучший копируйте за пределы сервера и помните, что мы не делаем резервных копий локальных дисков, если это не согласовано письменно. Подробности — в руководстве по локальным NVMe-накопителям.

Сначала оценка, потом инференс

Сравнивайте результат с базовой моделью на тех же промптах, с тем же шаблоном чата и теми же настройками сэмплирования:

• Функция потерь на отложенной выборке — чтобы заметить переобучение.
• Метрики задачи на вашем собственном наборе данных для оценки: точное совпадение, оценка по критериям или попарное сравнение.
• Регрессионный набор общих промптов — чтобы заметить утраченные навыки, например следование инструкциям или ответы на других языках.
• Ручной просмотр выборки ответов. Автоматические оценщики помогают, но кое-что упускают.

Для инференса есть два варианта. Можно выполнить слияние адаптера с базовыми весами с помощью merge_and_unload() в PEFT и обслуживать получившуюся модель, квантовав её до FP8, если хотите уместить её на одном H200. Для слияния нужна базовая модель в BF16 — около 140 ГБ, а это помещается в 225 ГБ оперативной памяти нашего тарифа с одним H100. Или можно держать базовую модель загруженной и подключать адаптеры во время работы: vLLM умеет обслуживать несколько адаптеров LoRA на одной базовой модели. Серверная часть описана в нашем руководстве по самостоятельному развёртыванию инференса LLM с vLLM. Если вы планируете запуск, посмотрите наши GPU-серверы для дообучения.

Часто задаваемые вопросы

Сколько видеопамяти нужно для дообучения модели 70B?

Без учёта активаций — около 40–45 ГБ для QLoRA и около 145 ГБ для LoRA на базовой модели в BF16. Полному дообучению с AdamW нужно 1,12–1,4 ТБ только на веса, градиенты и состояния оптимизатора.

Можно ли дообучить модель 70B на одном GPU?

Да, с QLoRA: 4-битная базовая модель занимает около 40 ГБ, поэтому на одном H100 с 80 ГБ или H200 со 141 ГБ остаётся место для адаптеров и активаций. LoRA на базовой модели в BF16 помещается на один B200 со 180 ГБ, а с большим запасом — на B300 с 288 ГБ.

Можно ли провести полное дообучение модели 70B на 8× H100?

На одном сервере со стандартным AdamW и смешанной точностью — нет: состояние обучения превышает 1 ТБ, а у 8× H100 всего 640 ГБ. Используйте LoRA на 8× H100, а для полного дообучения переходите на 8× B300 или на 8× B200 с экономной схемой.

Достаточно ли LoRA по сравнению с полным дообучением?

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

Сколько весит адаптер LoRA для модели 70B?

У адаптера ранга 16 на всех линейных проекциях около 207 млн параметров — примерно 0,4 ГБ в BF16, поэтому можно хранить много версий и переключать их при инференсе.