Почему я не смог обучить Gemma 3n локально (и почему я использую Vertex AI вместо этого)
Я работал над приложением для коучинга произношения, которому требуется мультимодальный ИИ, способный понимать как текст, так и аудио. В настоящее время приложение использует собственные аудиомодели Google Gemini, но я хотел дообучить специализированную модель для коучинга произношения. Google Gemma 3n показалась идеальной благодаря нативной поддержке аудио через кодировщик Universal Speech Model. У меня MacBook Pro M4 Pro с 48 ГБ ОЗУ. Этого, безусловно, должно хватить для локального обучения модели с 5 миллиардами параметров, верно?
Спойлер: нет. Вот что я узнал из своих неудач.
Почему я пробовал локальное обучение
Привлекательность была очевидна:
- Отсутствие затрат: Обучение на облачных GPU стоит $2–5 за попытку на таких сервисах, как RunPod или Vast.ai.
- Конфиденциальность: Мои обучающие данные остаются локальными.
- Быстрая итерация: Не нужно загружать наборы данных или настраивать облачные среды.
- Практическое обучение: Я хотел поэкспериментировать со скриптами обучения PyTorch, понять методы оптимизации памяти и научиться методом проб и ошибок. Локально я могу изменить код и повторить попытку за секунды, в отличие от развертывания в облаке и ожидания выполнения скриптов.
У меня было 988 обучающих примеров, охватывающих ошибки произношения, анализ манеры речи и обучающие диалоги. Это казалось идеальным кандидатом для ночного локального обучения.
Что на самом деле означает обучение модели?
Краткое уточнение: когда люди говорят «обучение», они могут иметь в виду две совершенно разные вещи.
Обучение с нуля — это то, что делают Google и OpenAI. Вы начинаете со случайных чисел и учите нейронную сеть всему с нуля. Это требует миллиардов текстовых примеров, сотен GPU и месяцев вычислительного времени. Это дорого и медленно.
Дообучение (Fine-tuning) — это взятие уже обученной модели и обучение ее вашей конкретной задаче. Самая сложная работа уже сделана. Модель знает язык и логику. Вы просто показываете ей свой конкретный формат и предметную область. Это как нанять кого-то, кто уже умеет кодировать, и научить его нашей кодовой базе, вместо того чтобы учить кого-то кодированию с нуля.
Для моего случая мне нужно дообучение. Модель уже понимает английский язык и произношение. Мне просто нужно, чтобы она давала обратную связь в моем формате.
Как дообучать с помощью Hugging Face
Hugging Face делает весь этот процесс намного проще, чем должен быть. Они размещают тысячи предварительно обученных моделей и предоставляют вам библиотеки Python для их дообучения. Вам не нужна докторская степень, чтобы это сделать.
Вот основной рабочий процесс:
from transformers import AutoProcessor, AutoModelForImageTextToText
# Загрузка модели и процессора
processor = AutoProcessor.from_pretrained("google/gemma-3n-E2B-it")
model = AutoModelForImageTextToText.from_pretrained("google/gemma-3n-E2B-it")
Две строки кода. Это загружает на ваш ноутбук ИИ-модель с 5,4 миллиардами параметров. processor преобразует ваши данные в нужный формат, а model — это сама нейронная сеть.
Затем вы используете класс Trainer для фактического запуска обучения:
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir="outputs/gemma3n",
num_train_epochs=3,
per_device_train_batch_size=1,
learning_rate=3e-5,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=my_dataset,
)
trainer.train() # Начать обучение!
Это здорово, потому что вы пропускаете все низкоуровневые вещи, такие как написание циклов градиентного спуска и управление контрольными точками. Недостаток? Вы не совсем понимаете, что происходит под капотом, особенно сколько памяти все потребляет. Что и стало моей проблемой.
Проблема с памятью
Обучение большой языковой модели требует одновременного размещения нескольких элементов в памяти:
- Веса модели: Для Gemma 3n E2B (5,4 млрд параметров) это около 11 ГБ в формате float16.
- Градиенты: Еще около 11 ГБ (того же размера, что и веса).
- Состояния оптимизатора: Adam хранит 2 состояния на параметр, добавляя около 22 ГБ.
- Активации: Промежуточные результаты вычислений.
- Данные пакета (Batch data): Обрабатываемые обучающие примеры.
Итого: 50–60 ГБ для меньшей модели E2B, 70–80 ГБ для E4B. Моих 48 ГБ ОЗУ не хватило.
Что я пробовал
Попытка 1: Полное дообучение с E4B
Ошибка нехватки памяти (Out of memory, OOM) при инициализации оптимизатора. Модель загрузилась (16 ГБ), но оптимизатор не смог выделить свое состояние.
Попытка 2: Меньшие размеры пакетов (batch sizes)
Уменьшение batch_size с 2 до 1 лишь отсрочило ошибку OOM до обратного прохода (backward pass).
Попытка 3: LoRA с E4B
LoRA (Low-Rank Adaptation) обучает только небольшие адаптерные слои вместо всей модели, уменьшая количество обучаемых параметров с 7,8 млрд до 40 млн (0,5%). Но замороженная базовая модель все равно должна поместиться в памяти: 59 ГБ только для E4B.
Попытка 4: LoRA с E2B
Это почти сработало. Меньшая модель E2B с LoRA потребляла 40–45 ГБ, оставляя запас в 3–8 ГБ. Обучение началось и продолжалось несколько часов, пока не возникли ошибки памяти при обработке определенных пакетов с более длинными последовательностями.
Почему в итоге не получилось
Даже при агрессивной оптимизации (LoRA, batch_size=1, без градиентного чекпоинтинга) 48 ГБ не хватило для стабильного обучения. Использование памяти варьировалось в зависимости от примера:
- Короткие примеры: 38–42 ГБ (нормально)
- Длинные примеры: 46–50 ГБ (сбои)
Я, вероятно, мог бы заставить это работать, отфильтровав более длинные примеры или уменьшив max_length до 1024 токенов, но в таком случае я бы пожертвовал возможностями модели ради ограничений оборудования.
Переход на Vertex AI
После трех дней безуспешных попыток я переключаюсь на обучение в Google Cloud Vertex AI:
- GPU A100 40GB: Более чем достаточно памяти для модели E4B с LoRA.
- Управляемая инфраструктура: Никакой отладки памяти, никаких сбоев.
- Стоимость: Около $3–5 за полный цикл обучения (3 эпохи).
- Время: 2–3 часа против моих 6–9 часов локальных попыток.
Я напишу отдельную статью о процессе обучения в Vertex AI, включая:
- Настройка пользовательских задач обучения с Gemma 3n.
- Загрузка наборов данных в Cloud Storage.
- Мониторинг обучения с помощью TensorBoard.
- Развертывание обученной модели в конечных точках Vertex AI.
Когда можно обучать локально?
Локальное обучение на Apple Silicon целесообразно, если:
- Ваша модель меньше: Модели с 1–3 млрд параметров с LoRA легко помещаются в 48 ГБ.
- У вас больше ОЗУ: M4 Max со 128 ГБ легко справился бы с Gemma 3n E2B.
- Вы используете квантование: 4-битное или 8-битное квантование дополнительно снижает потребление памяти, но усложняет процесс.
Ключевые выводы
- Память — это узкое место: 48 ГБ унифицированной памяти звучит как много, пока вы не попытаетесь обучить модели с 5+ млрд параметров. Только состояния оптимизатора могут удвоить или утроить ваши требования к памяти.
- LoRA помогает, но не является волшебством: Он уменьшает количество обучаемых параметров на 99%, но веса замороженной базовой модели все равно должны помещаться в памяти во время прямого и обратного проходов.
- Облачное обучение экономически выгодно: $3–5 за цикл обучения — это разумная цена по сравнению с днями отладки ошибок OOM и неудачных экспериментов.
- Знайте пределы своего оборудования: Унифицированная память Apple Silicon отлично подходит для инференса и обучения небольших моделей (1–3 млрд параметров), но для надежного обучения моделей с 5+ млрд параметров требуется выделенная VRAM GPU.