Как собрать NVIDIA Apex с поддержкой fused‑ядр и замерить эффект на обучении трансформера

Руководство показывает, как собрать NVIDIA Apex из исходников с включением CUDA и C++ расширений, проверить наличие fused‑ядер и сравнить их влияние на пропускную способность обучения трансформера с современным torch.amp.

В руководстве подробно показано, что сборка NVIDIA Apex из исходников с включёнными CUDA и C++ расширениями даёт реальные преимущества при обучении трансформеров: при наличии fused‑реализаций и AMP увеличивается пропускная способность, а при «Python‑only» установке ускорения не будет. Это важно для инженеров и исследователей, которые хотят получить ускорение без перехода на другую библиотеку.

Практические шаги включают проверку доступности CUDA, установку вспомогательных пакетов (ninja, packaging), клонирование репозитория apex и установку в окружении, где выставлены переменные APEX_CPP_EXT и APEX_CUDA_EXT (и другие необходимые переменные) для сборки C++/CUDA расширений. Без этой компиляции библиотека может казаться установленной, но не содержать высокопроизводительных fused‑модулей. Инструкция содержит скрипт для проверки собранных модулей: он ищет amp_C и fused_layer_norm_cuda и фиксирует наличие оптимизированных реализаций. Скрипт умеет обнаруживать FusedAdam и FusedLayerNorm, а также опционально FusedRMSNorm, что позволяет точно понять, какие ускорения доступны на текущей системе.

Авторы прогоняют бенчмарки, сравнивая оптимизатор FusedAdam с PyTorch AdamW, FusedLayerNorm и FusedRMSNorm со стандартными слоями нормализации, а также приводят примеры работы с legacy apex.amp и современным torch.amp. Основная цель этих замеров — измерить реальное влияние объединённых (fused) реализаций и AMP на пропускную способность при тренировке трансформера. Практическая рекомендация: собирать Apex с поддержкой CUDA/C++ и обязательно проверять наличие fused‑ядер после установки; если компиляция не удалась, останется «Python‑only» режим без высокопроизводительных ядер. Для оценки эффекта авторы советуют прогнать небольшой трансформер — сравнить чистый FP32 путь и путь с Apex+AMP, чтобы понять влияние на пропускную способность, время обучения и затраты ресурсов.