Молекулярная динамика (MD) — это не просто вычислительная задача, это один из самых ресурсоемких процессов в современной науке. Используя эти симуляции, исследователи могут наблюдать за поведением атомов с невероятной детализацией: от сворачивания белков до открытия новых лекарств и материалов. Однако цена такого детального моделирования высока. Симуляции обрабатывают от сотен тысяч до десятков миллионов атомов, продвигаясь на фемтосекунды за каждый шаг через миллиарды шагов. Поскольку размер задачи обычно фиксирован (мы хотим изучить конкретную молекулу), параллелизм должен распределять работу по всем доступным ресурсам одновременно. Этот режим известен как сильное масштабирование (strong scaling).
Пакет GROMACS является одним из самых популярных инструментов в этой области. Современные аппаратные обеспечения и гетерогенная архитектура CPU-GPU позволили достичь производительности в субмиллисекундном диапазоне — 100–200 микросекунд на временной шаг на нескольких GPU. Однако достижение сильного масштабирования на этом уровне требует экстремально низкой задержки ядер (kernel latency). Здесь на сцену выходит фундаментальная проблема: коммуникация между GPU остается главным узким местом при масштабировании симуляций на кластеры.
Большинство крупных приложений HPC, включая GROMACS, используют интерфейс Message Passing Interface (MPI) для межпроцессного взаимодействия. Но MPI был разработан для CPU-центричного выполнения. Когда GROMACS работает на GPU, рабочий процесс коммуникации заставляет GPU приостанавливаться, пока CPU оркестрирует передачу данных, прежде чем сигнализировать GPU о возобновлении работы. В алгоритме GROMACS, известном как halo exchange (обмен данными о граничных атомах между соседними GPU-доменами), эта передача повторяется во всех трех пространственных измерениях, потребляя более 50% общего времени CPU на пиковых скоростях инициации и ограничивая масштабируемость.
01Проблема CPU-оркестрации в традиционных MPI-решениях
Ключ к ускорению GROMACS заключается в устранении передач данных между CPU и GPU с помощью нативной коммуникации GPU. Традиционный подход, используемый в GROMACS, использует механизм поэтапной пересылки: данные передаются через промежуточные ранги (процессы) в каждом пространственном измерении за один или несколько шагов коммуникации, называемых «импульсами» (pulses). Каждый импульс начинается с ядра упаковки (pack kernel), которое собирает граничные атомы в непрерывный буфер отправки. Это создает цепочку зависимостей между фазами коммуникации (Z → Y → X), которую традиционная реализация обрабатывает через грубую сериализацию на уровне фаз.
Рассмотрим, как выглядит этот процесс в коде. Ниже представлен пример ядра GPU для упаковки координат граничных атомов в непрерывный буфер отправки:
// GPU kernel: pack boundary atom coordinates into a contiguous send buffer
__global__ void packSendBufKernel(float3* dataPacked, // packed coordinate buffer
const float3* data, // full coordinate array
const int* map, // index map: which atoms to pack
const int mapSize)
{
int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
if (threadIndex < mapSize)
{
dataPacked[threadIndex] = data[map[threadIndex]];
}
}Однако проблема кроется не в самом ядре, а в хост-коде (CPU), который его вызывает. Вот как выглядит типичный цикл обмена гало через MPI:
// CPU host code: coordinate halo exchange (MPI path)
for each dimension d in [Z, Y, X]:
for each pulse p in dimension d:
// 1. Launch GPU pack kernel on non-local stream
packSendBufKernel<<<grid, block, 0, nonLocalStream>>(
sendBuf, coords, indexMap, mapSize);
// 2. CPU blocks until GPU pack completes
cudaStreamSynchronize(nonLocalStream); // *** BLOCKING ***
// 3. CPU orchestrates MPI transfer
MPI_Isend(sendBuf, sendSize, sendRank, ...);
MPI_Irecv(coords + atomOffset, recvSize, recvRank, ...);
MPI_Waitall(...); // *** BLOCKING ***
// Received data lands directly
// in the coordinate buffer at the correct offset and is consumed
// by the non-local non-bonded force kernel as-is.Хотя эти ядра просты, узким местом является управляющий поток на стороне CPU, который их оборачивает. Каждый импульс требует, чтобы CPU блокировался на GPU, запускал передачу MPI и сигнализировал GPU о возобновлении работы в сериализованной последовательности для каждого импульса во всех трех измерениях. Обмен силами (force halo exchange) следует тому же паттерну, но в обратном направлении (X → Y → Z). Ранг, который отправил координаты, теперь получает силы, вычисленные на этих граничных атомах. Силы требуют шага распаковки (scatter-unpack) через ядро unpackRecvBufKernel, так как они должны аккумулироваться в правильных позициях полного массива сил.
Каждый импульс требует двух блокирующих синхронизаций CPU–GPU: одну перед вызовом MPI, чтобы гарантировать завершение ядра упаковки, и одну, навязанную перед тем, как следующий импульс сможет потребить пересланные данные. При 3D-декомпозиции и одном импульсе на измерение это шесть блокирующих ожиданий на временной шаг для координат и шесть для сил — всего 12. Каждое добавляет задержку на критическом пути, влияя на скорость итераций GROMACS.
02Решение: GPU-initiated коммуникация с NVSHMEM
NVIDIA NVSHMEM — это библиотека для реализации удаленного доступа к памяти (Remote Memory Access, RMA) на основе модели распределенного глобального адресного пространства OpenSHMEM. NVSHMEM позволяет ядрам GPU инициировать передачи данных напрямую, исключая CPU из критического пути и обеспечивая лучшее перекрытие вычислений и коммуникации.
Естественным началом является использование API, управляемого потоками (stream-triggered API, например, nvshmemx_put_signal_nbi_on_stream), которое устраняет барьеры синхронизации CPU-GPU с минимальной перестройкой кода. Однако у него есть два ограничения: коммуникация может начинаться только на границе ядра, что предотвращает перекрытие упаковки и коммуникации внутри импульса, а также последовательность, управляемая потоками, означает, что цепочка зависимой пересылки (Z → Y → X) не может быть слита (fused).
Мы решаем обе проблемы с помощью тонкозернистых зависимостей на уровне импульса, заменяя хост-коммуникацию и обеспечивая параллелизм между импульсами. Вместо того чтобы полагаться на CPU для синхронизации, мы используем зависимости,aware kernel fusion. Мы разбиваем карту индексов атомов каждого импульса на независимые и зависимые подмножества, позволяя немедленно упаковывать и передавать большинство данных и выполнять тонкозернистую сигнализацию на уровне импульса для «хвоста» зависимостей. Это снижает количество запусков ядер с шести до одного на временной шаг.
Замена хост-коммуникации на NVSHMEM
В новой архитектуре упаковка, удаленная запись (remote put) и ожидание завершения сливаются в единое ядро для каждого импульса коммуникации. Хост по-прежнему инициирует запуски в порядке импульсов, но он отказывается от барьеров CPU-GPU между импульсами. Правильный порядок выполнения теперь гарантируется самим потоком CUDA. Хост-вызов слитного ядра packSendBufAndPutNvshmemKernel выглядит так:
// CPU host code: one kernel launch per pulse, no CPU-GPU sync in between
int pulseOffset = 0;
for each dimension d in [Z, Y, X]:
for each pulse p in dimension d:
// Pack + send + wait are fused into a single GPU kernel
packSendBufAndPutNvshmemKernel<<<grid, block, 0, nonLocalStream>>(
coords, dataPacked, indexMap, sendSize,
sendRank, atomOffsetInSendRank,
signalReceiverRank + pulseOffset, signalCounter, bar, recvSize);
pulseOffset++;Вся операция pack–put–wait слита в единый запуск на устройстве. Адаптированная из реализации GROMACS, реализация ядра выглядит следующим образом:
__global__ void packSendBufAndPutNvshmemKernel(
float3* coords, float3* dataPacked, int* map,
int sendSize, int sendRank, int atomOffsetInSendRank,
uint64_t* signalReceiverRank, uint64_t signalCounter,
cuda::barrier<cuda::thread_scope_device>* bar, int recvSize)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (sendSize > 0)
{
// Pack: grid-stride gather into contiguous send buffer
for (int i = tid; i < sendSize; i += blockDim.x * gridDim.x)
dataPacked[i] = coords[map[i]];
// Device-scoped barrier: all CTAs must finish packing before put
auto token = bar->arrive();
if (blockIdx.x == 0)
{
bar->wait(std::move(token));
// Put packed data into peer's coordinate array
nvshmemx_float_put_signal_nbi_block(
&coords[atomOffsetInSendRank], // peer destination
dataPacked, // local source
sendSize * 3, // float count
signalReceiverRank, signalCounter,
NVSHMEM_SIGNAL_SET, sendRank);
}
}
// Wait for incoming data from recvRank before exiting
if (tid == 0 && recvSize > 0)
nvshmem_signal_wait_until(signalReceiverRank, NVSHMEM_CMP_EQ, signalCounter);
}Эта реализация устраняет всю синхронизацию CPU-GPU из критического пути. CPU выполняет только запуски ядер, которые могут перекрываться с выполнением GPU. Все потоки GPU (CTA) прибывают к барьеру на уровне устройства после упаковки, а блок 0 ожидает этот барьер, чтобы убедиться, что все потоки завершили работу, прежде чем вызывать nvshmemx_float_put_signal_nbi_block. Этот вызов передает упакованные данные на пир и устанавливает signalReceiverRank на пире для подтверждения завершения; получатель ожидает этот сигнал, прежде чем потреблять данные.
03Оптимизация для NVLink и RDMA
В описанном дизайне данные всегда маршрутизируются через транспорт NVSHMEM (nvshmemx_float_put_signal_nbi_block), что работает для любого интерконнекта. Однако, когда пир подключен через NVLink, GPU могут загружать/сохранять данные напрямую в память друг друга. Вместо упаковки в локальный буфер и отправки put, мы можем упаковывать данные напрямую в массив координат пира, устраняя отдельный круговой путь через глобальную память.
Функция nvshmem_ptr(remotePtr, peerRank) возвращает ненулевой указатель устройства, когда пир достижим через NVLink, и нуль в противном случае. Мы запрашиваем это один раз в ядре и ветвимся соответственно:
__global__ void packSendBufAndPutNvshmemKernel(
float3* coords, float3* dataPacked, int* map,
int sendSize, int sendRank, int atomOffsetInSendRank,
uint64_t* signalReceiverRank, uint64_t signalCounter,
cuda::barrier<cuda::thread_scope_device>* bar, int recvSize)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (sendSize > 0)
{
// Probe for NVLink: non-null means direct store is possible
float3* remotePtr = (float3*)nvshmem_ptr(coords, sendRank);
bool isNVLink = (remotePtr != nullptr);
// NVLink: pack directly into peer's coordinate buffer
// IB: pack into local staging buffer for later put
float3* dest = isNVLink ? (remotePtr + atomOffsetInSendRank)
: dataPacked;
for (int i = tid; i < sendSize; i += blockDim.x * gridDim.x)
dest[i] = coords[map[i]];
auto token = bar->arrive();
if (blockIdx.x == 0)
{
bar->wait(std::move(token));
if (isNVLink)
{
// Data already in peer memory — just signal completion
uint64_t* peerSignal = (uint64_t*)nvshmem_ptr(
signalReceiverRank, sendRank);
storeReleaseSysAsm(peerSignal, signalCounter);
}
else
{
// IB: combined data transfer + signal notification
nvshmemx_float_put_signal_nbi_block(
(float*)&coords[atomOffsetInSendRank],
(float*)dataPacked,
sendSize * 3,
signalReceiverRank, signalCounter,
NVSHMEM_SIGNAL_SET, sendRank);
}
}
}
if (tid == 0 && recvSize > 0)
nvshmem_signal_wait_until(signalReceiverRank, NVSHMEM_CMP_EQ, signalCounter);
}На пире NVLink данные уже находятся в памяти пира — остается только сигнализировать о завершении с помощью storeReleaseSysAsm. Для InfiniBand (IB) используется комбинированная передача данных и уведомление о сигнале через NVSHMEM put. Этот подход также использует движок TMA (Tensor Memory Accelerator) NVIDIA Hopper для эффективной массовой удаленной записи.
04Результаты масштабирования
Бенчмаркинг на суперкомпьютере NVIDIA Eos и кластерах NVIDIA GB200 NVL72 продемонстрировал улучшение сильного масштабирования intra- и inter-node до 2 раз по сравнению с GPU-aware MPI, особенно для систем, ограниченных задержкой. Подход обобщается на любое приложение HPC с паттернами обмена гало, хотя стандартизация нативных примитивов коммуникации GPU остается открытой проблемой.
05Что это значит на практике
Для исследователей и разработчиков HPC-приложений внедрение NVSHMEM означает не просто небольшое ускорение, а качественный скачок в эффективности использования ресурсов. В условиях, когда доступ к мощным GPU-кластерам в России и мире становится все более конкурентным, возможность запускать более крупные симуляции за меньшее время критически важна.
Локальный запуск таких оптимизированных версий GROMACS требует наличия GPU с поддержкой NVLink для максимальной эффективности (внутренние узлы), но даже через InfiniBand (межузловые соединения) прирост производительности остается значительным. Важно отметить, что для достижения максимального эффекта необходимо правильно настраивать потоки CUDA и барьеры, чтобы избежать простоев GPU.
Эта архитектура открывает путь к симуляциям, которые ранее были невозможны из-за ограничений масштабируемости. Исследователи могут моделировать более сложные биомолекулярные системы, такие как целые вирусные частицы или крупные белковые комплексы, с атомарной точностью, не тратя недели на ожидание завершения итераций. Стандартные библиотеки MPI, хотя и надежны, уступают место более гибким и производительным решениям, таким как NVSHMEM, которые учитывают специфику современных гетерогенных архитектур.
nvshmemx_put_signal_nbi_block.Таким образом, GPU-initiated communication через NVSHMEM — это не просто оптимизация, а необходимость для следующего поколения вычислительной биологии и материаловедения. Она устраняет исторические ограничения CPU-центричных подходов, позволяя GPU делать то, для чего они созданы: параллельно вычислять и передавать данные без посредников.
Источник: NVIDIA Developer ↗