Инфраструктура для глубокого обучения: из чего собрать стек

Где обучение ломается на практике

Модель на PyTorch запускается в двадцать строк. Проблемы начинаются, когда датасет перестаёт помещаться в память, одна видеокарта считает эпоху восемь часов, а чекпоинт на 40 гигабайт некуда положить. Инфраструктура для глубокого обучения — это не про то, чтобы код работал, а про то, чтобы он работал на нужном масштабе и не простаивал, пока вы платите за GPU по часам.

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

Вычисления: GPU, память и её нехватка

Главный ресурс — не число ядер, а объём памяти ускорителя. Модель, оптимизатор и активации должны поместиться в VRAM. Грубая прикидка для обучения в смешанной точности: на каждый миллиард параметров при использовании Adam уходит порядка 16–20 ГБ памяти (веса, градиенты и два момента оптимизатора), плюс активации, которые зависят от размера батча и длины последовательности.

Отсюда типичная развилка:

• Одна GPU (24–80 ГБ). Файнтюнинг, эксперименты, модели до нескольких миллиардов параметров с приёмами экономии памяти.
• Несколько GPU на одном узле. Соединены по NVLink или PCIe. Data parallelism ¹ заводится почти без усилий.
• Кластер из узлов. Здесь начинается взрослая распределённая тренировка, и узким местом становится сеть между машинами.

Приёмы, которые снимают давление на память, стоит знать до того, как покупать более дорогую карту:

• Mixed precision (bf16 / fp16) — снижает объём активаций примерно вдвое.
• Gradient checkpointing — не хранит все активации, а пересчитывает их на обратном проходе. Меняете память на время.
• Gradient accumulation — имитирует большой батч серией маленьких, когда батч целиком не влезает.
• ZeRO / FSDP — шардирование состояния оптимизатора, градиентов и весов между устройствами.

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

Данные: pipeline важнее модели

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

Что проверить в первую очередь

• Данные лежат на быстром локальном NVMe, а не тянутся по сети на каждой итерации.
• Препроцессинг (декодирование, аугментации) распараллелен по CPU-воркерам и не блокирует основной поток.
• Мелкие файлы упакованы в шардированные форматы (WebDataset, TFRecord, паркет), а не читаются по одному — миллион отдельных jpeg убивает файловую систему.

Минимальный пример загрузчика с параллельными воркерами и предзагрузкой:

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=256,
    num_workers=8,        # параллельная подготовка на CPU
    pin_memory=True,      # быстрее копирование в GPU
    prefetch_factor=4,    # сколько батчей готовить заранее
    persistent_workers=True,
)

Если после этого утилизация GPU не поднялась — узкое место в самом хранилище или в тяжёлом препроцессинге, и его надо выносить в офлайн-этап.

Аренда против своего железа

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

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

Оркестрация и воспроизводимость

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

1. Контейнеры. Зафиксируйте версии CUDA, драйверов и библиотек в образе. «У меня работало» между двумя машинами с разными версиями CUDA — классический источник потерянного дня.
2. Планировщик. Slurm в академической среде, Kubernetes с GPU-операторами — в продакшене. Он ставит задачи в очередь и не даёт двум тренировкам подраться за одну карту.
3. Трекинг экспериментов. Логируйте гиперпараметры, метрики и версию данных. Без этого через месяц вы не вспомните, чем чекпоинт A отличается от чекпоинта B.
4. Чекпоинтинг. Сохраняйте состояние регулярно. Обучение на спотовых инстансах, которые провайдер может отобрать в любой момент, без чекпоинтов бессмысленно.

Сеть: скрытый лимит распределённого обучения

При data parallelism после каждого шага узлы обмениваются градиентами — операция all-reduce. Её объём растёт с числом параметров, и обычный Ethernet на 10 Гбит/с становится тормозом: карты считают быстрее, чем успевают синхронизироваться. Для серьёзных многоузловых тренировок используют InfiniBand или высокоскоростной Ethernet с RDMA, чтобы обмен градиентами не съедал выигрыш от дополнительных GPU.

Проверить, упёрлись ли вы в сеть, просто: если добавление второго узла ускоряет обучение заметно меньше, чем в полтора раза, — коммуникация стала узким местом.

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