Спланируйте воспроизводимый запуск PyTorch

Докажите один полный шаг перед масштабированием

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

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

Измерьте репрезентативный шаг оптимизации

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

torch.cuda.reset_peak_memory_stats()
optimizer.zero_grad(set_to_none=True)
loss = model(**batch).loss
loss.backward()
optimizer.step()
torch.cuda.synchronize()
peak_gib = torch.cuda.max_memory_allocated() / 2**30
print(f"Peak tensor memory: {peak_gib:.2f} GiB")

Здесь отображается пиковое выделение тензоров, а не полный объем памяти устройства. Также проверьте зарезервированную память, накладные расходы среды выполнения и отчет системы о памяти устройства.

Выбирайте DDP или шардирование по правильной причине

DistributedDataParallel хранит реплику модели в каждом процессе и синхронизирует градиенты. Это полезно, когда полное состояние обучения помещается на каждом GPU и требуется обрабатывать больше данных параллельно; он не объединяет VRAM в один непрерывный пул.

Fully Sharded Data Parallel может распределять параметры, градиенты и состояние оптимизатора. Используйте его, когда давление на память оправдывает шардирование, и учитывайте дополнительные затраты на коммуникацию и сложность распределенных контрольных точек.

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

Запустите пилот на одном узле

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

torchrun --standalone --nproc-per-node=4 \
  train.py --config configs/pilot.yaml

Запускайте это только на выделенном узле. Локальная демонстрация AnchorGPU не выполняет обучение и не выделяет удалённые GPU.

Выберите оборудование по результатам пилотного запуска

Используйте A100 как более доступный базовый вариант с 80 GB CUDA, а H100 тестируйте, если у рабочей нагрузки есть полезный путь, специфичный для Hopper. H200 предлагает конфигурацию памяти CUDA большего объёма. MI300X предлагает большую ёмкость на ускоритель в этом каталоге при условии, что всё приложение квалифицировано на ROCm.

Для AMD проверяйте пользовательские расширения, ядра, версии пакетов и протестированный с ROCm образ PyTorch как единый комплект. Совместимость на уровне фреймворка не подтверждает каждую необязательную операцию.

Проверяйте восстановление, а не только сохранение

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

Используйте метод контрольных точек, соответствующий распределенному состоянию. Проверьте, что сохраненные артефакты можно прочитать независимо от исходного процесса и что результат оценки остается в пределах ожидаемого допуска.

Планируйте бюджет на весь конвейер

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

Срок 7 дней полезен для ограниченной проверки совместимости и профилирования. Выбирайте резервирование на 30 дней только после того, как сможете оценить полезный объем работы и операционный запас. Большее количество GPU помогает лишь тогда, когда дополнительный объем вычислений превышает ограничения коммуникаций и входного конвейера.