В мире высокопроизводительных вычислений (HPC) и искусственного интеллекта существует вечный конфликт: с одной стороны, исследователи и инженеры хотят писать чистый, читаемый код на Python, используя знакомые инструменты вроде NumPy и PyTorch. С другой стороны, для достижения максимальной производительности часто требуется погружаться в низкоуровневые API, такие как C++ или CUDA, что замедляет разработку и усложняет поддержку кода. NVIDIA решила эту проблему, выпустив nvmath-python v1.0 — библиотеку, которая действует как мост между миром Python-разработчиков и мощью экосистемы CUDA-X.
Эта библиотека не просто оборачивает существующие функции; она предлагает новую парадигму работы с математическими операциями. Она позволяет запускать вычисления на CPU, одном GPU или распределенных системах из множества узлов, сохраняя при этом привычный интерфейс. В этой статье мы подробно разберем, как nvmath-python меняет правила игры, какие новые возможности она дает для работы с разреженными данными, как оптимизировать затраты на планирование вычислений и почему слияние операций (kernel fusion) критически важно для производительности.
01Что такое nvmath-python и зачем она нужна?
nvmath-python — это Pythonic-абстракция над библиотеками CUDA-X и NVPL. Если говорить простым языком, это слой, который делает доступными такие мощные инструменты, как cuFFT (быстрое преобразование Фурье), cuBLASLt (линейная алгебра), cuDSS (разреженные системы), cuSPARSE, cuTENSOR и cuBLASMp, через интуитивно понятный Python-интерфейс.
Главная ценность библиотеки заключается в ее универсальности. Операции могут выполняться на процессоре (CPU), на графическом ускорителе (GPU) или в распределенной среде (multi-node). При этом разработчику не нужно переписывать код под каждую платформу. Библиотека автоматически определяет, где лучше выполнить задачу, и управляет перемещением данных между памятью CPU и GPU, минимизируя накладные расходы.
Особое внимание уделено интеграции с популярными фреймворками. nvmath-python работает «в паре» с NumPy, CuPy и PyTorch. Это означает, что вы можете передать в функцию библиотеки массив NumPy, получить результат в виде массива NumPy, но при этом использовать всю оптимизацию GPU, если она доступна. Это не замена NumPy, а его мощное расширение для специфических математических задач.
Пример интеграции с NumPy
Рассмотрим простой пример умножения матриц. В стандартном NumPy это делается одной строкой, но nvmath-python позволяет использовать продвинутые опции, недоступные в базовой версии. Обратите внимание, что входные и выходные данные остаются объектами NumPy, что сохраняет совместимость с остальным кодом.
import numpy as np
import nvmath
m, n, k = 10, 40, 100
a = np.random.randn(m, k) # a - массив NumPy
b = np.random.randn(k, n) # b - массив NumPy
# Выполнение умножения с использованием оптимизированных рутин
# Результат c также является массивом NumPy
c = nvmath.linalg.advanced.matmul(a, b)
02Гибкость выбора бэкенда: CPU, GPU и распределенные системы
Одной из самых сильных сторон nvmath-python является абстракция от конкретного устройства. Библиотека поддерживает несколько уровней вычислительных мощностей:
- GPU-библиотеки: cuBLAS, cuFFT и другие, работающие на видеокартах NVIDIA.
- CPU-библиотеки: NVPL (для процессоров NVIDIA Grace на архитектуре ARM) и Intel MKL (для x86-систем). Это позволяет запускать код на серверах без GPU, используя оптимизированные математические ядра.
- Распределенные библиотеки: cuBLASMp, cuSOLVERMp, cuFFTMp. Эти инструменты позволяют масштабировать вычисления на кластеры из множества GPU и узлов, что критически важно для больших моделей ИИ и научных симуляций.
Такая гибкость упрощает миграцию кода. Вы можете разработать алгоритм на ноутбуке с CPU, а затем развернуть его на сервере с мощными GPU, изменив лишь несколько строк или вообще не меняя код, если библиотека сама определит оптимальное пространство выполнения.
Демонстрация работы с разными пространствами памяти
В следующем примере показано, как одна и та же функция fft (быстрое преобразование Фурье) может работать как с данными на GPU (CuPy), так и с данными на CPU (NumPy). Библиотека автоматически определяет пространство выполнения на основе типа входного тензора.
import cupy as cp
import numpy as np
import nvmath
N = 2048
# Данные на GPU
gpu_data = cp.random.randn(N) + 1j * cp.random.randn(N)
# Данные на CPU
cpu_data = np.random.randn(N) + 1j * np.random.randn(N)
# FFT выполняется автоматически в соответствующем пространстве
c_gpu = nvmath.fft.fft(gpu_data)
c_cpu = nvmath.fft.fft(cpu_data)
03Два типа API: Универсальные инструменты и специализированные скальпели
В nvmath-python APIs разделены на два больших класса, что помогает разработчикам выбирать правильный инструмент для задачи.
1. Generic APIs (Общие API)
Это «многофункциональные инструменты». Они обеспечивают единообразный пользовательский опыт для различных пространств выполнения и типов операндов. Однако их конфигурация ограничена базовыми функциями, которые есть у всех бэкендов. Они идеальны для задач, где важна совместимость, а не максимальная производительность, или когда специфическая оптимизация не требуется.
2. Specialized APIs (Специализированные API)
Эти API находятся в подмодулях advanced (например, nvmath.linalg.advanced). Они представляют собой «скальпели» — узкоспециализированные инструменты, которые предоставляют полный контроль над конфигурацией. Например, продвинутое умножение матриц (advanced matmul) реализует композитные операции специально для плотных операндов на GPU, позволяя выжать максимум из аппаратных возможностей.
Выбор зависит от контекста. Если операция является узким местом (bottleneck) в вашем коде и требует аппаратной оптимизации, используйте advanced API. Если же задача не критична к производительности или требует кросс-платформенности, выбирайте общие API.
04Прозрачность вычислений: Логирование и отладка
Одна из проблем при работе с GPU-вычислениями — непонимание того, где именно происходит работа с данными. nvmath-python интегрируется со стандартным модулем logging Python, позволяя отслеживать каждый шаг вычислений. Это особенно полезно для отладки и оптимизации перемещения данных между CPU и GPU.
При включении логирования вы можете увидеть, из какого пространства памяти взяты данные, в какое пространство они отправлены на выполнение и какой тип данных используется. Это помогает выявлять узкие места, связанные с избыточным перемещением данных (data transfer overhead).
Пример логирования умножения матриц
import numpy as np
import nvmath
import logging
# Настройка логирования
logging.basicConfig(level=logging.INFO,
format="%(asctime)s %(levelname)-8s %(message)s", force=True)
logging.disable(logging.NOTSET)
m, n, k = 8000, 2000, 4000
a_cpu = np.random.randn(m, k).astype(np.float32)
b_cpu = np.random.randn(k, n).astype(np.float32)
# Выполнение операции
d_cpu = nvmath.linalg.advanced.matmul(a_cpu, b_cpu)В логе вы увидите сообщения вроде: "The input operands' memory space is cpu, and the execution space is on device 0". Это прямо указывает на то, что данные были перемещены с CPU на GPU для вычислений. Если вы используете FFT с данными на CPU, лог покажет, что выполнение происходит на CPU, избегая лишних пересылок.
05Сила композитных операций и Kernel Fusion
Почему просто цепочка вызовов NumPy может быть медленной? Проблема в так называемой арифметической интенсивности (arithmetic intensity). Когда вы выполняете несколько операций по отдельности (например, умножение матриц, затем умножение на скаляр, затем сложение), каждая операция требует загрузки данных в память GPU и выгрузки результата. Если арифметических операций мало по сравнению с объемом данных, GPU простаивает в ожидании данных.
nvmath-python решает эту проблему через kernel fusion (слияние ядер). Вместо того чтобы запускать три разных ядра CUDA, библиотека создает одно ядро, которое выполняет все операции за один проход по памяти. Это радикально снижает накладные расходы на ввод-вывод.
Пример: Tall-and-Skinny матрицы
Рассмотрим случай умножения матриц, где одна матрица очень высокая и узкая (tall-and-skinny). В NumPy это делается в несколько шагов:
import cupy as cp
import nvmath
m, n, k = 10_000_000, 40, 10
a = cp.random.randn(m, k, dtype=cp.float32)
b = cp.random.randn(k, n, dtype=cp.float32)
c = cp.random.randn(m, n, dtype=cp.float32)
alpha, beta = 1.5, 0.5
# Вариант 1: Множество ядер (медленно из-за пересылок)
d1 = alpha * cp.matmul(a, b) + beta * c
# Вариант 2: Одно слитное ядро (быстро)
d2 = nvmath.linalg.advanced.matmul(a, b, c=c, alpha=alpha, beta=beta)Как показывает практика, вариант с nvmath.linalg.advanced.matmul работает значительно быстрее благодаря использованию cuBLASLt и JIT-компиляции слитного ядра. Это особенно критично для задач глубокого обучения, где такие операции встречаются постоянно.
06Экономия времени: Stateful API и автонастройка
Высокая производительность CUDA-библиотек часто достигается за счет фазы планирования (planning) и автонастройки (autotuning). Библиотека анализирует размер задачи, тип данных и архитектуру GPU, чтобы выбрать оптимальный алгоритм. Этот процесс может занимать время. Если вы используете функциональный (stateless) API, эта плата взимается с вас при каждом вызове функции.
Однако в задачах ИИ одна и та же операция (например, умножение матриц в нейронной сети) выполняется тысячи раз с разными данными. Здесь вступает в игру Stateful API (классовый интерфейс). Вы создаете объект один раз, планируете и настраиваете его, а затем многократно выполняете вычисления, переиспользуя готовый план.
Пример использования Stateful API
import nvmath
from nvmath.linalg.advanced import MatmulEpilog
import cupy as cp
feed_count = 10 # Количество повторений
batch_size = 1024
m, n, k = 1024, 1024, 1024
a = cp.random.rand(batch_size, m, k, dtype=cp.float32)
b = cp.random.rand(batch_size, k, n, dtype=cp.float32)
bias = cp.random.rand(batch_size, m, 1, dtype=cp.float32)
# Создание объекта и выполнение в три этапа
with nvmath.linalg.advanced.Matmul(a, b) as mm:
# 1. Планирование
mm.plan(epilog=MatmulEpilog(MatmulEpilog.RELU_BIAS),
epilog_inputs={"bias": bias})
# 2. Автонастройка (выбор лучшего алгоритма)
mm.autotune(iterations=5)
# 3. Выполнение (многократно, с минимальными затратами)
for i in range(feed_count):
d = mm.execute()
# Обновление операндов для следующей итерации
mm.reset_operands_unchecked(a=d)На графике ниже видно, как стоимость выполнения снижается с увеличением числа итераций при использовании stateful API. Первоначальные затраты на планирование и автонастройку распределяются (амортизируются) на все последующие вызовы. Более того, настроенный план можно сохранить на диск и загрузить в другом сеансе, что экономит время при развертывании моделей.
07Расширенные возможности: Разреженные тензоры и кастомные ядра
nvmath-python v1.0 вводит поддержку Universal Sparse Tensors (UST) через специализированный язык (DSL). Это позволяет пользователям создавать собственные оптимальные форматы разреженных данных, не прибегая к написанию сложного кода на C++. Это открывает новые горизонты для работы с большими моделями, где разреженность данных может значительно ускорить вычисления.
Кроме того, библиотека интегрируется с numba-cuda. Это позволяет писать кастомные функции на Python, которые компилируются в JIT-код и используются как прологи или эпилоги для операций FFT или линейной алгебры. Например, можно применить пользовательский фильтр Гаусса к изображению, используя FFT из nvmath-python, но с кастомной логикой обработки, написанной на Python.
08Установка и развертывание
Установка nvmath-python спроектирована с учетом потребностей современных CI/CD процессов. Поддерживаются популярные менеджеры пакетов: pip, conda, uv, pixi. Вы можете выбрать полную установку со всеми зависимостями или минималистичную версию для сред без GPU. Библиотека позволяет гибко настраивать бэкенды (CPU, GPU, распределенные) и выбирать библиотеки-компаньоны (NumPy, CuPy, PyTorch).
09Что это значит на практике
Для инженеров по машинному обучению и научных вычислителей nvmath-python v1.0 — это не просто еще одна библиотека, а инструмент для преодоления «узких горлышек» производительности. Вот ключевые выводы:
- Единый интерфейс для всего стека: Пишите код один раз, запускайте на CPU, одном GPU или кластере. Это упрощает разработку и тестирование.
- Максимальная скорость через слияние: Используйте
advancedAPI для критичных к производительности операций. Слияние ядер (kernel fusion) экономит время на пересылку данных, что особенно важно для моделей с низкой арифметической интенсивностью. - Амортизация затрат: Не вызывайте функции «в лоб» в циклах. Используйте Stateful API (классы) для планирования и автонастройки один раз, а затем многократно выполняйте вычисления. Это может дать ускорение в разы.
- Прозрачность: Используйте логирование, чтобы понимать, где происходят пересылки данных, и оптимизировать их.
В эпоху, когда модели ИИ становятся все больше, а требования к эффективности вычислений растут, такие инструменты, как nvmath-python, становятся необходимостью. Они позволяют разработчикам сосредоточиться на алгоритмах, а не на низкоуровневой оптимизации, не жертвуя при этом производительностью.
Источник: NVIDIA Developer ↗