Задача
R&D-команда обучала большую визуальную foundation-модель на огромном датасете. Наивный data-parallel упирался в межузловые коммуникации уже после сотни карт: GPU простаивали в ожидании all-reduce, а прогон длиной в недели падал из-за единственного сбойного узла и откатывался почти к началу.
Нужно было обучать на сотнях GPU с близким к линейному масштабированием и так, чтобы выпавший узел не убивал многодневный прогон.
Подход
Перешли на 3D-параллелизм: data-, tensor- и pipeline-параллелизм вместе, с шардингом состояний оптимизатора (ZeRO/FSDP). Коммуникации перекрыли с вычислением, чтобы карты не ждали сети, и затюнили коллективы под топологию InfiniBand. Профилировали не loss, а реальную утилизацию: где простой — там и узкое место.
Чекпойнтинг сделали асинхронным и отказоустойчивым: состояние пишется в фоне, а при выпадении узла прогон эластично пересобирается и продолжается с последнего чекпойнта, а не с нуля. Под нестабильные ноды это превращает многодневное обучение из лотереи в управляемый процесс.
Стек
PyTorch FSDP · NCCL + InfiniBand · 3D parallelism · async checkpointing · Slurm/Kubernetes
Результат
- 512 GPU в одном прогоне с линейностью масштабирования около 91%.
- ×4 короче цикл обучения до целевой сходимости.
- 0 рестартов с нуля — выпавший узел больше не теряет многодневный прогон.
- Утилизация кластера выросла за счёт перекрытия коммуникаций с вычислением.