Масштабный UMAP на нескольких GPU: скорость и точность

В мире анализа данных и машинного обучения Uniform Manifold Approximation and Projection (UMAP) давно стал золотым стандартом для снижения размерности. Этот алгоритм позволяет визуализировать многомерные данные, сохраняя их локальную и глобальную структуру, что критически важно для исследовательского анализа, тематического моделирования и даже анализа одиночных клеток. Однако по мере роста объемов данных, которые мы собираем и обрабатываем, традиционные методы сталкиваются с серьезными ограничениями. Когда датасеты достигают размеров в десятки или даже сотни миллионов векторов, вычислительная стоимость каждого запуска UMAP возрастает экспоненциально. Это превращает итеративный анализ, требующий частых запусков с разными параметрами, в мучительно медленный процесс, который может занимать часы или даже дни.

Ранее решения, такие как out-of-core подход в NVIDIA cuML, позволяли обрабатывать данные, превышающие объем памяти одного GPU, но ограничивали масштабирование этапа обучения (fit) одним устройством, оставляя этап трансформации (transform) для распределенных вычислений. Это создавало узкое место, не позволяющее полностью раскрыть потенциал современных графических процессоров. Выходом из этой ситуации стала новая функциональность, представленная в NVIDIA cuML и NVIDIA cuVS версии 25.06, которая распределяет самый ресурсоемкий этап — построение графа k-ближайших соседей (kNN) для всех точек — между несколькими GPU. Это изменение не только радикально сокращает время обучения, но и делает возможным работу с данными объемом в сотни гигабайт за считанные минуты, сохраняя при этом высокую точность встраиваний (embeddings).

01Как работает масштабирование UMAP на нескольких GPU?

Ключевая идея, лежащая в основе эффективного масштабирования UMAP, заключается в построении графа всех соседей (all-neighbors kNN graph) без необходимости помещать весь датасет в память GPU одновременно. Этот подход, известный как out-of-core, уже был описан в предыдущих публикациях NVIDIA. Он работает путем разбиения датасета на сбалансированные кластеры. Важным аспектом является перекрытие (overlap) векторов между соседними кластерами, что позволяет сохранить связи ближайших соседей на границах кластеров, не теряя при этом качество итогового встраивания.

Локальные графы kNN вычисляются независимо для каждого кластера, после чего эти локальные графы объединяются в один глобальный граф всех соседей. Поскольку эта задача уже декомпозирована на независимые единицы работы, она естественным образом расширяется на мульти-GPU среду. Каждый кластер может обрабатываться изолированно, что означает, что вычисление локальных графов kNN не требует доступа ко всему набору данных или координации с другими кластерами. В результате кластеры могут быть распределены по разным GPU, где каждый GPU независимо загружает данные для назначенных ему кластеров из памяти CPU.

Каждый GPU затем вычисляет локальные графы всех соседей и объединяет их с глобальным графом kNN. Независимое вычисление этих локальных графов позволяет избежать дорогостоящей коммуникации типа "all-to-all", которая обычно ограничивает масштабируемость распределенных рабочих нагрузок. Этот метод обеспечивает значительное ускорение общей производительности на очень больших наборах данных, делая возможным обработку данных объемом в сотни гигабайт за минуты вместо часов.

02Настройка UMAP для работы с несколькими GPU

Для использования новой функциональности важно понимать два гиперпараметра, которые определяют баланс между пространством (памятью), временем и качеством результата. Реализация multi-GPU UMAP в cuML следует тем же шагам, что и однопоточная версия, но с добавлением следующих параметров:

• knn_n_clusters: Количество кластеров, на которые будет разделены данные. Эти кластеры примерно сбалансированы, что означает равномерное распределение векторов по доступным GPU для обработки. Увеличение этого значения уменьшает количество точек, назначенных каждому кластеру, снижая объем данных, которые должны помещаться в память каждого GPU. Это особенно важно для работы с данными, превышающими объем видеопамяти одного устройства.
• knn_overlap_factor: Общее количество ближайших кластеров, к которым будет назначена каждая точка данных. Этот параметр увеличивает перекрытие точек между кластерами, сохраняя больше истинных ближайших соседей на границах кластеров. Увеличение этого параметра, как правило, улучшает качество графа всех соседей, что, в свою очередь, повышает качество финальных встраиваний UMAP. Однако это улучшение достигается ценой увеличения времени вычислений и использования памяти, так как большему количеству векторов необходимо обрабатываться в рамках каждого кластера.

Вместе knn_overlap_factor и knn_n_clusters обеспечивают контролируемый компромисс между пространством, временем и качеством. Распределенное построение графа всех соседей напрямую экспонирует эти параметры через API всех соседей cuVS. Хотя cuML UMAP использует их внутренне, они также доступны напрямую для приложений, которым требуется самостоятельное построение графа всех соседей.

Пример кода: API всех соседей cuVS с использованием нескольких GPU

Ниже приведен пример использования API всех соседей cuVS с несколькими GPU. Параметры n_clusters и overlap_factor соответствуют аргументам knn_n_clusters и knn_overlap_factor, экспонируемым cuML UMAP, которые передаются в этот API cuVS во время построения графа.

from cuvs.neighbors import all_neighbors
from cuvs.common import MultiGpuResources

params = all_neighbors.AllNeighborsParams(
    algo="nn_descent",
    n_clusters=32,
    overlap_factor=2
)
# Using all GPUs on the system
res = MultiGpuResources()

indices, distances = all_neighbors.build(
    data,
    k,
    params,
    distances=cupy.empty((n_rows, k)),
    resources=res
)

03Управление памятью и производительностью

Для управления компромиссом между памятью и качеством встраивания рекомендуется повышать knn_overlap_factor одновременно с увеличением knn_n_clusters, чтобы поддерживать использование памяти на относительно постоянном уровне. Например, удвоение коэффициента перекрытия удваивает ожидаемое количество точек на кластер, поэтому удвоение количества кластеров сохранит количество точек на кластер неизменным. Рекомендуется использовать достаточно knn_n_clusters, чтобы каждый локальный граф и векторы в этом кластере комфортно помещались в памяти GPU.

Требования к памяти на каждом GPU зависят от набора данных и локального графа kNN, который строится. Для набора данных с N векторами и размерностью D доминирующими компонентами памяти являются:

M_data = N * (knn_overlap_factor / knn_n_clusters) * D * sizeof(data)

M_graph = N * k * (knn_overlap_factor / knn_n_clusters) * (sizeof(index) + sizeof(distance))

На практике требуется дополнительная память для накладных расходов, таких как временные выделения рабочей области во время построения графа. Поэтому фактическое пиковое использование памяти может быть выше оцененного. Мы рекомендуем выбирать knn_overlap_factor и knn_n_clusters с запасом ниже доступной памяти GPU.

Например, при использовании GPU с 80 ГБ памяти и наборе данных с N=100 млн и D=1024 векторов float32 (409 ГБ), вы можете выбрать knn_overlap_factor=2 и knn_n_clusters=24. Это делает размер данных каждого кластера M_data примерно 34 ГБ. Используя тип индекса int64 и тип расстояния float32, и k=15, M_graph составляет около 1,5 ГБ, что в сумме дает 35,5 ГБ. Это оставляет достаточный запас для небольшого дисбаланса кластеров и других накладных расходов.

04Практическое использование multi-GPU UMAP в NVIDIA cuML

Использование multi-GPU UMAP в NVIDIA cuML требует лишь нескольких дополнительных параметров конфигурации по сравнению с предыдущим использованием cuML UMAP. В следующем примере Python показано, как построить встраивание UMAP, используя все доступные GPU, или только определенные GPU (например, с ID 0, 4 и 5).

from cuml.manifold import UMAP

# Using all GPUs on the system
umap = UMAP(
    build_kwds={
        "knn_n_clusters": 32,
        "knn_overlap_factor": 2,
    },
    device_ids="all",
)
embedding = umap.fit_transform(data)

# Using a subset of GPUs
umap = UMAP(
    build_kwds={
        "knn_n_clusters": 32,
        "knn_overlap_factor": 2,
    },
    device_ids=[0, 4, 5],
)
embedding = umap.fit_transform(data)

Аргумент device_ids управляет тем, какие GPU участвуют в рабочей нагрузке, в то время как knn_n_clusters и knn_overlap_factor контролируют построение графа всех соседей. В целом, увеличение количества GPU снижает время выполнения за счет распределения построения всех соседей по устройствам.

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

05Визуализация встраиваний масштабного уровня

Для оценки качества результатов важно сравнить встраивания, сгенерированные на больших наборах данных. На рисунке 1 сравниваются встраивания, полученные на наборе данных MIRACL (106 млн векторов, размерность 2048), используя эталонную реализацию CPU с предварительно вычисленным GPU-графом всех соседей и нативную реализацию GPU UMAP от cuML. Важно отметить, что UMAP инвариантен к масштабу, трансляциям и вращениям, поэтому эти встраивания функционально идентичны, даже если изображения выглядят немного по-разному.

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

06Производительность multi-GPU UMAP на данных, превышающих память одного GPU

Мы оценили влияние использования multi-GPU UMAP на наборах данных, превышающих память одного GPU (Wiki и MIRACL). Все бенчмарки выполнялись на системе NVIDIA DGX с восемью GPU NVIDIA H100 и процессором Intel Xeon 8480CL с 2 ТБ оперативной памяти. Метрика "trustworthiness" (достоверность) используется для оценки качества встраивания, возвращая число от 0 до 1, где более высокое значение означает лучшее сохранение локальной структуры соседства.

На рисунке 2a показано поведение времени выполнения широко используемой эталонной реализации CPU на прогрессивно увеличивающихся подвыборках наборов данных Wiki и MIRACL. По мере увеличения размера набора данных время выполнения быстро растет из-за затрат памяти и вычислений. На полном масштабе реализация CPU не смогла завершить работу из-за чрезмерного потребления памяти, даже на системе с 2 ТБ ОЗУ. Для оценки времени выполнения на этих масштабах мы спроецировали полное время выполнения CPU, экстраполируя измеренный тренд масштабирования от меньших подвыборок.

На рисунке 2b сравнивается спроецированное время выполнения CPU с фактическим временем выполнения cuML multi-GPU UMAP, работающего на восьми GPU NVIDIA H100. На наборе данных MIRACL из 106 млн векторов cuML достигает ускорения до 74 раз по сравнению со спроецированным временем выполнения CPU. Это делает UMAP практичным на ранее недоступных масштабах, завершая полную обработку набора данных из 106 млн векторов всего за 8 минут.

07Масштабирование multi-GPU

На рисунке 3 показано поведение масштабирования cuML multi-GPU UMAP на наборах данных Wiki и MIRACL по мере увеличения количества GPU от одного до восьми H100. Улучшение производительности достигается при сохранении сопоставимого качества встраивания во всех конфигурациях GPU. Производительность продолжает улучшаться по мере добавления большего количества GPU, без ущерба для достоверности (trustworthiness) результатов.

08Что это значит на практике

Внедрение multi-GPU поддержки для построения графа всех соседей в UMAP через NVIDIA cuML и cuVS 25.06 представляет собой значительный прорыв для аналитиков данных и исследователей, работающих с большими данными. Теперь задачи, которые ранее считались невыполнимыми из-за ограничений памяти и времени, такие как обработка сотен миллионов векторов, становятся рутинными. Скорость в 74 раза быстрее, чем у CPU, и время выполнения всего в 8 минут для огромных датасетов открывает новые возможности для итеративного анализа, где исследователи могут быстро тестировать различные параметры и подходы без ожидания.

Для специалистов по данным в России и других регионах, где доступ к облачным ресурсам может быть ограничен, локальный запуск таких вычислений на системах с несколькими GPU (например, NVIDIA H100 или A100) становится еще более актуальным. Это позволяет сохранять конфиденциальность данных, одновременно получая доступ к передовым технологиям обработки. Важно правильно настраивать параметры knn_n_clusters и knn_overlap_factor, чтобы найти оптимальный баланс между использованием памяти и качеством результатов. Следуя рекомендациям NVIDIA, можно достичь максимальной эффективности, используя все доступные ресурсы системы.

Таким образом, NVIDIA продолжает укреплять свои позиции в области ускоренных вычислений для ИИ, предоставляя инструменты, которые не только ускоряют существующие процессы, но и расширяют границы возможного в анализе данных. Интеграция multi-GPU UMAP в стандартные пайплайны Python через cuML делает эти мощные возможности доступными для широкого круга пользователей, от исследователей до инженеров по машинному обучению.

Источник: NVIDIA Developer ↗