Компания NVIDIA представила значительное обновление для библиотек NVIDIA cuML и NVIDIA cuVS версии 25.06. Ключевым нововведением стала поддержка распределенных вычислений (multi-GPU) для критически важного этапа алгоритма UMAP — построения графа k-ближайших соседей (all-neighbors kNN graph). Это решение позволяет обрабатывать датасеты объемом до 870 ГБ за 8 минут на кластере из восьми GPU NVIDIA H100, обеспечивая ускорение до 74 раз по сравнению с проекциями на CPU.
Техническая суть: как это работает
Проблема масштабирования UMAP заключалась в том, что этап обучения (training) ограничивался одним GPU, хотя этап трансформации уже поддерживал несколько устройств. Новая архитектура решает это путем разбиения датасета на сбалансированные кластеры. Каждый GPU независимо вычисляет локальные графы kNN для своего кластера, после чего они объединяются в глобальный граф. Такой подход исключает необходимость в дорогостоящей коммуникации типа "all-to-all" между всеми узлами, сохраняя качество вложения (embedding) даже на огромных объемах данных.
Ключевые метрики производительности
Эксперименты проводились на наборах данных MIRACL и Wiki. Результаты демонстрируют радикальное сокращение времени обучения:
| Параметр | Значение / Результат |
|---|---|
| Объем обрабатываемых данных | 870 ГБ векторов |
| Время выполнения (multi-GPU) | 8 минут |
| Ускорение относительно CPU | до 74x |
| Аппаратная платформа | 8x NVIDIA H100 |
| Качество (Trustworthiness) | Сохранено на высоком уровне |
Настройка гиперпараметров для баланса скорости и качества
Для эффективной работы с multi-GPU введены два ключевых параметра, позволяющих управлять компромиссом между памятью, временем и точностью:
- knn_n_clusters: Количество кластеров, на которые разбивается данные. Увеличение этого числа снижает нагрузку на память каждого отдельного GPU, распределяя векторы более равномерно.
- knn_overlap_factor: Фактор перекрытия между кластерами. Определяет, сколько ближайших кластеров учитывается для каждой точки. Значение по умолчанию для качественного эмбеддинга — 2. Увеличение этого параметра (например, до 4 или 6) повышает точность поиска соседей, но требует пропорционального увеличения knn_n_clusters для сохранения стабильного потребления памяти.
Разработчики рекомендуют начинать с knn_overlap_factor=2 и увеличивать его небольшими шагами для умеренных масштабов, или большими шагами для датасетов с количеством кластеров >100. Это делает интерактивный анализ данных и тонкую настройку моделей feasible для задач, которые ранее требовали дней вычислений.
Источник: NVIDIA dev blog ↗