Оптимизация 3D-диффузионных моделей на GPU: ускорение в 100 раз без потери качества
Диффузионные модели широко используются для высококачественного синтеза 3D-изображений МРТ, однако их развертывание сдерживается огромными требованиями к ресурсам GPU. Каждая выборка требует сотен вычислений U-Net, а поведение ядер отличается высокой неоднородностью.
В новой статье на arXiv исследователи провели всесторонний анализ производительности современной медицинской диффузионной модели Med-DDPM на трёх поколениях архитектур NVIDIA. Они изучили поядерное разбиение времени выполнения, характеристику набора инструкций, использование памяти, активность варпов и приоритетные оценки профилировщика.
Выяснилось, что обучение доминируется ядрами cuDNN сверток и implicit-GEMM, причём неэффективность возникает из-за паттернов доступа к памяти, преобразований тензорных макетов и ограниченного использования тензорных ядер (Tensor Cores).
Руководствуясь этими наблюдениями, авторы применили две архитектурно-зависимые оптимизации: активацию вычислений TF32 на тензорных ядрах и переход на 3D-формат расположения каналов (channels-last). Результаты оказались впечатляющими: число циклов SM сократилось до 100 раз, количество динамических инструкций уменьшилось на два порядка, загрузка тензорных ядер выросла с 1.45 до 9.98 раз, а количество инструкций на такт (IPC) увеличилось на 7% на графическом ускорителе A100.
Важно, что качество синтеза 3D-МРТ при этом не ухудшилось. Таким образом, предложенные методы позволяют существенно ускорить работу диффузионных моделей без потери точности, что открывает путь к более широкому внедрению таких моделей в медицинскую практику.
Хотя работа выполнена на примере Med-DDPM, авторы отмечают, что выявленные закономерности и оптимизации применимы и к другим генеративным диффузионным моделям для 3D-данных.
