Бинцзин Чжан

Процесс обучения модели в машинном обучении больших данных требует больших затрат как вычислений, так и памяти. Многие алгоритмы параллельного машинного обучения состоят из итерации вычислений по набору обучающих данных и обновления связанных параметров модели до тех пор, пока модель не сойдется. В эпоху больших данных и объем набора данных, и количество параметров модели могут быть огромными. Чтобы ускорить выполнение итерационных вычислений, обычно данные обучения и параметры модели кэшируются в памяти. Однако из-за ограничений памяти во многих сценариях это может не подходить для всех. Эффективно запускать алгоритмы машинного обучения с ограничениями памяти - сложная задача.

Существующие решения

В настоящее время существует несколько подходов к решению проблемы ограничения памяти. Эти решения могут быть реализованы либо на уровне приложения, либо на уровне времени выполнения, либо на уровне ОС.

Уровень приложения

Разработав алгоритм с внепрограммным исполнением для конкретного обучающего приложения, можно разделить параметры набора данных / модели так, чтобы часть из них можно было кэшировать в памяти, а остальное можно было поместить на диск. С меньшим рабочим набором программа все еще может выполняться с данными и параметрами модели, которые существуют в памяти. Поскольку это решение в значительной степени зависит от алгоритма, оно возлагает на разработчиков большую нагрузку по реализации правильных алгоритмов.

Другая проблема в этом подходе заключается в том, что в таких алгоритмах, как Скрытое распределение Дирихле с выборкой Гиббса, использование памяти варьируется на разных этапах процесса обучения. Поскольку модель продолжает сжиматься, общий объем используемой памяти также уменьшается, и в результате программа может кэшировать больше параметров набора данных / модели на более позднем этапе, чем она может кэшировать в начале процесса обучения. Если этот тип алгоритма выполнения вне ядра используется для разделения набора данных и параметров модели, он будет тратить память на более позднем этапе выполнения, когда потребление памяти уменьшится.

Уровень Время выполнения

Такие платформы, как Spark, предоставляют набор прозрачных API-интерфейсов набора данных, которые используют как память, так и диски для хранения данных. По сравнению с решениями для конкретных алгоритмов это решение освобождает разработчиков от бремени управления памятью при реализации алгоритмов обучения. Однако выполнение вне ядра не является полностью прозрачным - оно по-прежнему требует конфигурации уровней хранения для всего задания / кластера и имеет дело с постоянным хранением данных, но не с распределением памяти при выполнении задачи.

Уровень ОС

Использование пространства подкачки ОС для решения проблемы ограничения памяти имеет преимущества в том, что оно адаптивно, поскольку оно используется только тогда, когда ресурсы памяти ограничены, и прозрачно для приложения, поскольку нет необходимости в дополнительном программировании. Кроме того, использование пространства подкачки позволяет избежать накладных расходов файловой системы.

Однако это решение также имеет ограничения. Объем пространства подкачки обычно ограничен - его размер сложно динамически перенастроить, требуя привилегированного доступа к операционной системе, и это можно сделать только при отключенном подкачке. Также сложно добиться более высокой производительности. Хотя можно настроить подкачку для повышения производительности, отрегулировав степень детализации подкачки и политику исключения, для этого требуется знание адресов в виртуальной памяти, которые часто скрыты за многими уровнями системных абстракций.

Адаптивное, прозрачное и высокопроизводительное внешнее выполнение

Мы выбрали работу на уровне среды выполнения, чтобы обеспечить прозрачный интерфейс между пользовательским кодом алгоритма обучения и средой выполнения, который скрывает логику выполнения вне ядра от кода приложения. Мы используем абстракцию среды выполнения параллельного выполнения (рисунок 1), в которой каждый исполнитель (слово, которое мы используем для описания параллельных компонентов) кэширует связанную с ним часть набора данных и поддерживает собственное состояние вычислений. Во время итераций эти исполнители либо планируются с вычислительными задачами, либо координируются для синхронизации.

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

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

Наше внешнее исполнение реализовано со следующими функциями:

  1. Адаптация к ограничению памяти: мы устанавливаем предел памяти для программы, построенной на нашей среде выполнения. Изначально все исполнители загружаются в память. Затем, когда достигнут предел памяти, мы сохраняем некоторые исполнители, которые не работают в файловой системе, и позволяем текущему выполнению продолжаться с освобожденной памятью. Поскольку файловая система намного больше, чем пространство подкачки, наше решение не требует предварительной настройки, что обеспечивает большую гибкость. На более поздних этапах выполнения, если использование памяти сокращается, все исполнители автоматически возвращаются в память.
  2. Прозрачный интерфейс: когда разработчики пишут параллельные алгоритмы на основе исполнителей, им не нужно вмешиваться в какую-либо логику, связанную с обнаружением нехватки памяти или заменой / заменой исполнителей. Все это контролируется средой выполнения. Однако разработчикам по-прежнему нужно писать код, чтобы сообщить, как сохранять / загружать исполнителей.
  3. Высокая производительность. Несмотря на то, что выполнение вне ядра на уровне фреймворка исключает оптимизацию, обходящую накладные расходы, связанные с использованием файловой системы, оно позволяет использовать другие мощные средства оптимизации, учитывающие потребности приложений. Знание зависимостей между задачами приложения позволяет среде выполнения выбрать и сохранить исполнителя, которому не назначены задачи или который имеет самый низкий приоритет задачи, чтобы избежать нехватки памяти.

Ловушка OOM и выселение

Чтобы обнаружить нехватку памяти, мы используем ulimit -v, чтобы установить верхний предел виртуальной памяти программы; при достижении предела вызовы выделения памяти возвращают нулевые указатели или коды ошибок. Мы также обертываем API выделения памяти, чтобы мы могли фиксировать ошибки нехватки памяти и добавлять дополнительные процедуры удаления. Эти оболочки используют dlsym для связи с библиотекой распределения памяти для выполнения распределения памяти. Если какое-либо выделение памяти не удается, оболочки выполняют процедуру удаления, которая перемещает исполнителей, которые не выполняются, в файловую систему. Если нет исполнителя для замены, но распределение по-прежнему не удается, программа терпит неудачу (Рисунок 2).

В процессе выселения исполнители сериализуются на диск в виде файлов. Поскольку в этот момент свободная память может быть недоступна, вызовы выделения памяти перенаправляются в зарезервированную область памяти, где мы управляем всеми действиями по выделению памяти во время сериализации исполнителя. Поскольку все потоки, которым не удается выделить память, обращаются к исполнителям одновременно, мы избегаем конкуренции потоков, позволяя потокам, которым не удается получить доступ к исполнителям, переходить в спящий режим. Потоки, которые успешно завершают процедуру выселения, затем уведомляют тех, которые ушли в спящий режим. Поскольку некоторые библиотеки распределения памяти могут сохранять виртуальную память, даже когда они были освобождены, что означает, что любое выделение памяти после вытеснения все еще может завершиться ошибкой, мы используем библиотеку Jemalloc и просим ее вернуть неиспользованную виртуальную память.

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

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

Предварительные результаты

Мы сравнили наш прототип с системой подкачки с использованием мини-пакетных K-средних в облаке AWS с двумя типами инстансов (Рисунок 3): m4.2xlarge с 8 vCPU, 32 ГБ памяти и i3.xlarge с 4 виртуальными ЦП, памятью 32 ГБ. Мы установили количество вычислительных потоков равным восьми и четырем соответственно и использовали тома SSD (gp2) EBS общего назначения объемом 256 ГБ. Мы использовали набор данных из полного набора данных ImageNet ILSVRC2012, который содержит 1140297 изображений, каждое из которых содержит 10K функций. Мы запустили 100 пакетов для обучения 1000 классов, каждый пакет состоял из 40 задач, а каждая задача состояла из 1100 изображений. Мы ограничили размер виртуальной памяти числом, близким к общему размеру физической памяти (ulimit -v 31000000), и позволили ОС планировать 96 ГБ виртуальной памяти между физической памятью и пространством подкачки.

Мы обнаружили, что на i3.xlarge наша система использовала 2,13 часа, а системный своп - 12,30 часа. На m4.2xlarge наша система заняла 1,52 часа, в то время как замена системы заняла 13,10 часа. Системная подкачка была намного медленнее, потому что ее выполнение вызывает перегрузку памяти, что приводит к высоким IOPS (операций ввода-вывода в секунду) и достигает верхнего предела производительности томов EBS.

Эти предварительные результаты показывают, что наш подход к выполнению вне ядра значительно снижает временные затраты по сравнению с методом подкачки системы. Мы ожидаем, что благодаря прозрачному интерфейсу и адаптивному управлению памятью этот метод поможет снизить сложность разработки и развертывания приложений для обучения параллельных моделей в машинном обучении больших данных, сохраняя при этом высокую производительность. Поскольку не каждый пользователь AI / ML имеет доступ к неограниченным вычислительным ресурсам крупных компаний или национальных лабораторий, этот метод помогает демократизировать доступ к приложениям AI / ML для больших данных.