GPU Memory Swapping миграция страниц между VRAM и оперативной памятью в Linux

Видеопамять на ускорителях давно стала тем ресурсом, который заканчивается раньше всего: 80 гигабайт на флагманском ускорителе выглядят солидно ровно до момента, когда появляется модель, которая весит 120, и к ней просятся два батча данных. Классический ответ, переложить данные в оперативную память и по мере необходимости возвращать их обратно, у GPU работает иначе, чем обычный swap у процессора, зато работает автоматически: страницы видеопамяти мигрируют в RAM и обратно под управлением драйвера и ядра. Механизм этот называется по-разному в разных стеках, у NVIDIA это Unified Memory поверх модуля nvidia-uvm, у AMD механизм поверх HMM ядра Linux, у Intel shared allocations в Level Zero, а смысл один: приложение видит один указатель, а система сама решает, где физически лежат страницы.

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

Unified Memory и миграция страниц между VRAM и RAM

Идея унифицированной памяти у NVIDIA появилась в CUDA 6.0 в 2014 году, но по-настоящему взрослой она стала с поколением Pascal, которое научилось обрабатывать отказы страниц прямо на GPU. Приложение выделяет буфер через cudaMallocManaged, получает обычный указатель, доступный и ядрам на устройстве, и коду на процессоре. Первое обращение к ещё не резидентной странице порождает отказ страницы, драйвер мигрирует данные по PCIe в нужную сторону, и выполнение продолжается. Миграции идут блоками покрупнее обычных страниц, на свежих архитектурах обычно по 2 мегабайта, чтобы снизить число событий, а за счётчиками доступа следит аппаратура, понимая, какие страницы остывают.

Минимальный рабочий фрагмент выглядит так:

float *buf;
size_t bytes = (size_t)N * sizeof(float);
cudaMallocManaged(&buf, bytes);

cudaMemPrefetchAsync(buf, bytes, 0, 0);
cudaMemAdvise(buf, bytes, cudaMemAdviseSetReadMostly, 0);

kernel<<<blocks, threads>>>(buf, N);

Ни одного cudaMemcpy в коде нет, и это не упрощение примера, а суть подхода: движение данных становится свойством адресного пространства, а не обязанностью программиста. Выделение стартует с размещения в оперативной памяти, резидентность в VRAM возникает по мере первых обращений, и именно поэтому первый прогон на свежих данных почти всегда медленнее прогретого: по дороге система собирает статистику accesses и подтягивает горячие диапазоны. Для проверки, что устройство вообще допускает одновременный доступ к управляемой памяти, читается атрибут cudaDevAttrConcurrentManagedAccess, а переменная окружения CUDA_MANAGED_FORCE_DEVICE_ALLOC=1 заставляет размещать управляемые выделения прямо в VRAM там, где это поддержано.

Ручное управление через prefetch и советы

Автоматика миграции умна ровно настолько, насколько предсказуема приложение, поэтому API даёт рычаги ручного управления. cudaMemPrefetchAsync инициирует заблаговременную перегонку диапазона на целевое устройство или обратно на процессор через идентификатор cudaCpuDeviceId, причём операция асинхронная и выполняется в потоке CUDA, то есть встраивается в общий конвейер вычислений. Правильное место для prefetch это пауза между этапами, когда можно перекрыть перенос данных полезной работой, а не момент, когда ядро уже упёрлось в отсутствие страницы.

Второй рычаг это советы через cudaMemAdvise. Совет ReadMostly помечает диапазон как преимущественно читаемый, после чего копии страниц могут жить на нескольких устройствах одновременно, и чтение из разных мест не вызывает пинг-понг миграций. Совет PreferredLocation подсказывает, где диапазон должен проживать по умолчанию, а AccessedBy перечисляет устройства, которым этот диапазон нужен, позволяя драйверу заранее готовить таблицы страниц. Обратная перегонка делается тем же вызовом с идентификатором процессора:

cudaMemPrefetchAsync(buf, bytes, cudaCpuDeviceId, stream);

Такая операция полезна на паузе между этапами, когда результаты нужно обработать на CPU или сохранить на диск, и хочется вывести данные из VRAM до того, как поместится следующий батч. Практика обучения моделей складывается из трёх привычек: веса помечать ReadMostly, активационные буферы держать с PreferredLocation на устройстве, а большие датасеты prefetch-ить батчами чуть вперёд от итерации.

Когда подкачка превращается в тормоз

Соотношение скоростей здесь беспощадное. Локальная полоса HBM у современных ускорителей измеряется терабайтами в секунду, у A100 это порядка 2 ТБ/с, у H100 около 3.35 ТБ/с, а PCIe четвёртого поколения на 16 линий даёт порядка 32 ГБ/с в каждую сторону. Каждый отказ страницы, который не удалось перекрыть вычислениями, значит паузу в микросекунды, а в плохом сценарии, когда рабочий набор систематически больше VRAM, приложение входит в режим перемалывания: ядра стоят в очередях за страницами, миграции съедают шину, и производительность падает в разы и десятки раз.

Арифметика промаха понятная. Перенос страницы на 2 мегабайта по PCIe 4.0 при полосе около 32 ГБ/с занимает порядка 60 микросекунд в одну сторону, и если итерация обучения ловит пару сотен неперекрытых промахов, к шагу добавляются десятки миллисекунд чистого ожидания. Увидеть это глазами можно в профилировщике Nsight Systems: на временной шкале появляются события отказов страниц и миграций Unified Memory, и по их плотности сразу видно, где узкое место. Лечение сводится к трём приёмам. Во-первых, укрупнение и выравнивание миграций, чтобы система двигала мегабайтные блоки вместо россыпи мелких страниц. Во-вторых, prefetch с перекрытием, о котором речь выше. В-третьих, честное перепланирование работы: если рабочий набор не влезает, данные делят на слои или тайлы и обрабатывают порциями, потому что никакие советы не спасут приложение, которое каждый шаг трогает всю память сразу.

Подкачка у других вендоров и в виртуализации

NVIDIA не одинока в этом механизме. У AMD в ROCm унифицированная память строится поверх HMM, механизма гетерогенного управления памятью в ядре Linux, который умеет отражать таблицы страниц устройства и мигрировать страницы между RAM и VRAM по тем же принципам отказов. У Intel в oneAPI есть класс распределений shared USM, работающий аналогично: один указатель, страницы мигрируют по мере обращений, а в драйвере i915 и новом xe это поддержано аппаратными отказами страниц. На Windows механизм давно привычен любому геймеру: менеджер видеопамяти WDDM выгружает страницы из VRAM в системную память под давлением, из-за чего фризы в играх с переполненной видеопамятью выглядят как подгрузка с диска.

Виртуализация добавляет свой слой. Современные версии NVIDIA vGPU умеют переподписывать VRAM между виртуальными машинами, выгружая страницы гостя в память хоста, что превращает подкачку видеопамяти в инструмент плотности размещения, а не только способ уместить большую модель. Для администратора лимит VRAM на виртуальную машину перестаёт быть жёсткой стеной, но цена промаха растёт: переподписанная память живёт на хосте, и всплеск отказов бьёт сразу по нескольким гостям одновременно. В ядре Linux сами миграции видны в счётчиках /proc/vmstat, строки pgmigrate_success и pgmigrate_fail показывают, сколько страниц ядру удалось передвинуть и сколько попыток не состоялось, и на системах с активным HMM-трафиком эти числа живо реагируют на нагрузку.

Что происходит внутри ядра при отказе страницы

Под капотом всей этой автоматики работает связка драйвера и ядра. Модуль nvidia-uvm обрабатывает отказы страниц от GPU и решает, куда мигрировать данные, а ядро Linux со своей стороны начиная с версии 4.14 развивает HMM, гетерогенное управление памятью: механизм зеркалирует таблицы страниц устройства, помечает специальный тип device private pages и предоставляет ядру набор migrate_vma для переноса содержимого страниц между обычной RAM и памятью, доступной только устройству. Драйвер amdgpu строит унифицированную память именно на HMM, и отладка таких систем начинается с проверки, что модуль загружен и не ругался при старте.

lsmod | grep -E 'uvm|amdgpu'
dmesg | grep -iE 'uvm|hmm'

Разница между стеками в деталях управления: проприетарный драйвер NVIDIA ведёт собственную политику миграций внутри uvm, тогда как amdgpu опирается на общую механику ядра, и от этого зависит, какие счётчики и настройки доступны для наблюдения. Общий знаменатель один: таблицы страниц процессора и устройства синхронизируются, а физическое расположение данных становится деталью, которую приложение не обязано знать.

Фреймворки и офлоад в оперативку и NVMe

В глубоких фреймворках подкачка вышла на уровень конфигурации. DeepSpeed реализует офлоад состояний оптимизатора и параметров: ZeRO-Offload выгружает градиенты и состояния Adam в оперативную память, а ZeRO-Infinity умеет писать их даже на NVMe, размечая, что и куда съезжает. Конфигурация при этом читаемая:

{
  "offload_optimizer": {
    "device": "nvme",
    "nvme_path": "/mnt/nvme",
    "pin_memory": true
  },
  "offload_param": {
    "device": "cpu"
  }
}

Логика распределения важна для понимания: первыми уезжают состояния оптимизатора, потому что их объём кратно превышает сами веса и трогаются они редко, потом параметры, а активации стараются держать на устройстве, ибо они нужны каждый шаг. Стадия ZeRO-3 дополнительно делит сами параметры между процессами, и офлоад параметров в такой конфигурации включается осознанно: обращения к ним за шаг слишком частые, чтобы каждый раз платить за перенос. В PyTorch встроенные проверки дают быстрый мониторинг: torch.cuda.mem_get_info возвращает свободную и занятую VRAM, а ручной офлоад тензоров через .to(device) остаётся валютой там, где фреймворк не предлагает автоматики. Отдельная линия развития это NVIDIA GPUDirect Storage, который гонит данные из NVMe прямо в VRAM без промежуточной остановки в RAM, и в связке с офлоадом он укорачивает путь возврата больших чекпоинтов и датасетов.

Диагностика, мониторинг и практический порядок

Быстрая диагностика на работающей системе сводится к короткому набору команд:

nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv
nvidia-smi topo -m
grep pgmigrate /proc/vmstat

Первая строка отвечает на вопрос, сколько памяти занято прямо сейчас, вторая показывает топологию связности GPU с CPU и NVMe, что важно при планировании, куда офлоадить, третья показывает миграции страниц на стороне ядра. Для детального разбора очередь за профилировщиками: Nsight Systems для событий Unified Memory, DCGM для длительного мониторинга фермы, а в самописном коде cudaMemGetInfo и счётчики отказов через CUPTI закрывают вопрос телеметрии. В DeepSpeed дополнительно имеет смысл смотреть логи на предмет соотношения времени вычислений и офлоада: если шаг упирается в обмен с NVMe, очередь за офлоад-большим объёмом видна по падению доли вычислительного времени.

Порядок внедрения для команды, упёршейся в нехватку VRAM, обычно такой:

1. замерить реальный рабочий набор и профиль обращений через профилировщик, отделив горячие данные от холодных;
2. перевести выделения на управляемую память там, где код допускает, и помечать данные советами ReadMostly и PreferredLocation;
3. расставить prefetch по границам этапов, перекрывая миграции вычислениями;
4. вынести состояния оптимизатора в RAM или NVMe средствами фреймворка, начиная с самого холодного;
5. ограничить число одновременных процессов на устройстве, потому что их суммарный рабочий набор и есть причина давления;
6. вернуть в план апгрейд VRAM или переход на модель с меньшей точностью, когда подкачка всё ещё чаще, чем хочется.

Итог у технологии двойственный, и это надо признать честно. Подкачка видеопамяти спасает проект, который иначе не стартует вовсе, и позволяет пережить пики, но постоянный режим жизни с переполненной VRAM всегда дороже, чем влезающие данные, из-за разрыва в полосе между HBM и PCIe. Грамотная эксплуатация выглядит как управление вероятностью: чем точнее советы, крупнее блоки и раньше prefetch, тем реже приложение ловит отказ страницы в самый неподходящий момент, и тем ближе его скорость к работе без всякой подкачки.