[Перевод] Обучение ИИ-моделей на обычном ЦП

Ещё раз о ЦП для машинного обучения в эпоху дефицита мощностей GPU

Недавние успехи в области ИИ зачастую связывают с появлением и эволюцией графических процессоров (GPU). Архитектура GPU, как правило, включает тысячи ядер для многопроцессорной обработки, высокоскоростную память, выделенные тензорные ядра и пр. Такая архитектура особенно хорошо подходит для рабочих нагрузок, связанных с ИИ и машинным обучением, которые отличаются высоким потреблением ресурсов. К сожалению, из-за резкого роста разработок в области ИИ также наблюдается всплеск потребности в GPU, из-за чего возник их дефицит. В результате разработчики систем машинного обучения всё активнее изыскивают альтернативные аппаратные платформы, на которых можно было бы обучать и эксплуатировать модели. В таком качестве используются, например, выделенные специализированные интегральные схемы (ASIC) для работы с искусственным интеллектом, такие как облачные тензорные процессоры Google, Haban Gaudi и AWS Trainium. Притом, что эти варианты позволяют значительно сэкономить, они подходят для работы не со всеми моделями машинного обучения и, подобно GPU, также остаются дефицитными. В этом посте мы вновь обратимся к старым добрым классическим процессорам (CPU) и вновь поговорим о том, насколько они адекватны применительно к современным ML-моделям. Безусловно, ЦП обычно не так хороши для обслуживания связанных с машинным обучением рабочих нагрузок как графические процессоры, зато их гораздо проще приобрести. Если бы удалось гонять на ЦП (хотя бы некоторые) из таких рабочих нагрузок, то продуктивность разработки в целом удалось бы значительно повысить.

В других статьях (напр., здесь) подчёркивается, насколько важно анализировать и оптимизировать производительность рабочих нагрузок ИИ/машинного обучения именно во время выполнения — чтобы таким образом ускорить разработку и минимизировать затраты. Притом, что это принципиально важно независимого от того, какой именно вычислительный движок вы используете, инструменты профилирования и приёмы оптимизации могут значительно отличаться от платформы к платформе. В этой статье мы обсудим некоторые варианты оптимизации производительности, актуальные именно при работе с ЦП. Все примеры будем разбирать на процессорах Intel® Xeon® (с арохитектурой Intel® AVX-512) и с применением фреймворка PyTorch (версия 2.4). Правда, схожие приёмы оптимизации вполне будут работать и с другими ЦП и фреймворками. Если ещё конкретнее, все описанные здесь эксперименты проводились на инстансе Amazon EC2 c7i, на котором развёрнут образ AWS для глубокого обучения (Deep Learning AMI). Пожалуйста: какие бы варианты облачной платформы, версии ЦП, фреймворки для машинного обучения или любые другие библиотеки или инструменты мы ни упоминали, не считайте, что мы их рекламируем на фоне имеющихся альтернатив.

Наша цель — продемонстрировать, что, хотя, разработка систем машинного обучения на ЦП явно останется лишь резервным вариантом, всегда есть возможность «смягчить удар», а иногда даже выстроить на основе ЦП жизнеспособную альтернативу.

Дисклеймеры

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

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

Важно, что влияние обсуждаемых нами оптимизаций на производительность во время выполнения, вероятно, будет значительно варьироваться в зависимости от выбранной модели и от деталей окружения (см., например, степень изменчивости от модели к модели, приведённую в официальном дашборде PyTorch TouchInductor CPU Inference Performance). Сравнительные показатели производительности, которые мы приводим, специфичны для построенной нами «игрушечной» модели и именно той среды выполнения, которой мы пользуемся. Обязательно оценивайте все предлагаемые оптимизации заново именно на вашей модели и в вашем рабочем окружении.

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

Игрушечный пример - ResNet-50

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

import torch import torchvision from torch.utils.data import Dataset, DataLoader import time # датасет со случайными изображениями и метками class FakeDataset(Dataset): def __len__(self): return 1000000 def __getitem__(self, index): rand_image = torch.randn([3, 224, 224], dtype=torch.float32) label = torch.tensor(data=index % 10, dtype=torch.uint8) return rand_image, label train_set = FakeDataset() batch_size=128 num_workers=0 train_loader = DataLoader( dataset=train_set, batch_size=batch_size, num_workers=num_workers ) model = torchvision.models.resnet50() criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters()) model.train() t0 = time.perf_counter() summ = 0 count = 0 for idx, (data, target) in enumerate(train_loader): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() batch_time = time.perf_counter() - t0 if idx > 10: # skip first steps summ += batch_time count += 1 t0 = time.perf_counter() if idx > 100: break print(f'average step time: {summ/count}') print(f'throughput: {count*batch_size/summ}')

При выполнении этого скрипта на процессоре c7i.2xlarge (с 8 виртуальными вычислительными ядрами) и ЦП-версией PyTorch 2.4 получаем пропускную способность в 9,12 образцов в секунду. Для сравнения отметим, что пропускная способность такого же (неоптимизированного) скрипта на инстансе Amazon EC2 g5.2xlarge (с 1 GPU и 8 виртуальными вычислительными ядрами) составляет 340 образцов в секунду. Принимая во внимание сравнительные расходы для двух этих типов инстансов ($0,357 в час для c7i.2xlarge и $1,212 в час для g5.2xlarge, по состоянию на момент подготовки оригинала этой статьи), находим, что обучение на GPU-инстансе даёт примерно одиннадцатикратный (!!) выигрыш в соотношении цена/качество. Естественно, исходя из этих результатов, гораздо выгоднее тренировать модели машинного обучения именно на GPU. Давайте рассмотрим некоторые возможности, помогающие сузить этот пробел.

Оптимизация производительности PyTorch

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

Размер батча

Увеличивая размер учебного батча, потенциально можно повысить производительность, если снизить частоту обновления параметров модели. (На GPU в таком случае мы дополнительно выигрываем, так как нам удаётся сократить издержки, связанные с транзакциями на уровне ЦП-GPU, например, с загрузкой ядра). Тем не менее, тогда как на GPU мы подбирали такой размер батча, который позволил бы максимально задействовать при обработке память видеокарты, именно эта стратегия пагубно отразится на производительности при работе на ЦП. По причинам, рассмотрение которых выходит за рамки этого поста, память на ЦП устроена сложнее, и надёжнее всего оптимальный размер батча определяется методом проб и ошибок. Учитывайте, что, если изменить размер батча, это может отразиться на сходимости при обучении.

В следующей таблице обобщены данные по пропускной способности нашей учебной рабочей нагрузки для нескольких (произвольно выбранных) размеров батча:

Пропускная способность при обучении как функция от размера батча (данные от автора)

В отличие от того, что показывает работа с GPU, при использовании инстансов типа c7i.2xlarge наша модель предпочитает батчи размером поменьше.

Многопроцессная загрузка данных

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

Пропускная способность при обучении как функция от количества потоков, загружающих данные (от автора)

Смешанная точность

Также популярна следующая техника: использовать при работе с динамическим диапазоном, например, с torch.bfloat16, обычно считаемым более располагающим к обучению для целей ML, такие типы данных с плавающей точкой, у которых точность ниже обычного — таковы, например, torch.float16 или torch.bfloat16. Естественно, если понизить точность типа данных, это может отрицательно повлиять на сходимость, поэтому делать такое следует аккуратно. В PyTorch есть пакет torch.amp для автоматического подбора смешанной точности, позволяющий оптимизировать использование таких типов данных. В Intel® AVX-512 предлагается поддержка типа данных bfloat16. Ниже показано, как выглядит шаг обучения в изменённом виде:

for idx, (data, target) in enumerate(train_loader): optimizer.zero_grad() with torch.amp.autocast('cpu',dtype=torch.bfloat16): output = model(data) loss = criterion(output, target) loss.backward() optimizer.step()

После такой оптимизации пропускная способность составляет 24,34 образца в секунду, то есть, возрастает на 86%!!

Формат памяти Channels Last

На момент подготовки оригинала этой статьи формат памяти Channels last был оптимизацией, находившейся в стадии «бета» и применялся преимущественно в зрительных моделях. В этом формате можно хранить в памяти четырёхмерные тензоры (NCHW) так, что последнее измерение соответствует каналам. В результате все данные по каждому пикселю хранятся компактно, друг с другом. Считается, что такой формат наиболее удобен при работе с платформами семейства и якобы увеличивает производительность ResNet-50 на ЦП Intel® Xeon®. Вот как выглядит шаг обучения с учётом этой коррекции:

for idx, (data, target) in enumerate(train_loader): data = data.to(memory_format=torch.channels_last) optimizer.zero_grad() with torch.amp.autocast('cpu',dtype=torch.bfloat16): output = model(data) loss = criterion(output, target) loss.backward() optimizer.step()

В результате получаем пропускную способность 37,93 образцов в секунду — улучшение ещё на 56% и в совокупности на 415% по сравнению с нашим исходным экспериментом. Нам явно везёт!

Компиляция Torch

В другом посте были рассмотрены сильные стороны PyTorch в качестве средства для поддержки компиляции графов, а также рассказано, как это может повлиять на производительность во время выполнения. В противовес выбираемому по умолчанию жадному режиму, в котором каждая операция выполняется независимо (это и означает «жадно»), API компиляции преобразует модель в промежуточный вычислительный граф, который затем динамически компилируется в низкоуровневый машинный код именно так, как это наиболее подходит для взаимодействия с основополагающим движком обучения. Этот API поддерживает компиляцию с использованием различных бэкенд-библиотек и со множеством конфигурационных опций. Здесь рассмотрим лишь вычисления с привлечением заданного по умолчанию бэкенда (TorchInductor), а также с бэкендом ipex из расширения Intel® для работы с PyTorch — это библиотека, специально оптимизированная для работы с аппаратным обеспечением Intel. Пожалуйста, посмотрите в документации, как её правильно устанавливать и использовать. Ниже приведено обновлённое определение модели:

import intel_extension_for_pytorch as ipex model = torchvision.models.resnet50() backend='inductor' # опционально меняется на 'ipex' model = torch.compile(model, backend=backend)

Применительно к нашей игрушечной модели эффект от компиляции torch заметен лишь при отключенной оптимизации «channels last» (рост ~27% для каждого из рассмотренных бэкендов). Когда задействуется «channels last», производительность даже падает. Поэтому в последующих экспериментах мы такую оптимизацию применять не будем.

Оптимизация памяти и потоков

Существуют некоторые возможности оптимизировать использование ресурсов того ЦП, на котором идёт работа. В частности, речь об оптимизации управления памятью и выделения потоков с учётом аппаратной структуры того процессора, с которым мы работаем. Улучшить управление памятью можно при помощи продвинутых аллокаторов памяти (таких как Jemalloc и TCMalloc) и/или сократив количество таких обращений к памяти, которые выполняются сравнительно медленно (например, на уровне узлов NUMA). Выделение потоков можно улучшить, подобрав подходящую конфигурацию потоковой библиотеки OpenMP и/или воспользоваться библиотекой Open MP от Intel.

В целом оптимизации такого рода требуют глубоко понимать архитектуру ЦП и возможности программного стека, обеспечивающего её работу. Чтобы упростить ситуацию, в PyTorch предлагается скрипт torch.backends.xeon.run_cpu, помогающий автоматически сконфигурировать память и потоковые библиотеки так, чтобы оптимизировать память во время исполнения. Ниже приведена команда, которая подготавливает работу с выделенной памятью и потоковыми библиотеками. Далее мы сначала поговорим о том, как организовать распределённое обучение, а потом вернёмся к теме узлов NUMA.

Нужно проверить, правильно ли установлены TCMalloc (conda install conda-forge::gperftools) и библиотека Open MP от Intel (pip install intel-openmp), а потом выполнить следующую команду.

python -m torch.backends.xeon.run_cpu train.py

Благодаря скрипту run_cpu, удаётся ещё сильнее повысить производительность во время исполнения — обработать 39,05 образцов в секунду. Обратите внимание: в скрипте run_cpu предусмотрено множество возможностей для дальнейшей настройки производительности. Обязательно изучите документацию по нему, чтобы использовать его максимально эффективно.

Расширение для PyTorch от Intel

Расширение Intel® для PyTorch предоставляет дополнительные возможности оптимизации обучения, реализованные в функции ipex.optimize. Ниже показано, как оно используется по умолчанию. Полностью возможности этого расширения описаны в документации.

model = torchvision.models.resnet50() criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters()) model.train() model, optimizer = ipex.optimize( model, optimizer=optimizer, dtype=torch.bfloat16 )

В сочетании с описанными выше оптимизациями памяти и потоков удаётся довести пропускную способность до 40,73 образцов в секунду. Обратите внимание: близкий к этому результат достигается при отключении конфигурации «channels last».

Распределённое обучение на ЦП

В конструкции процессоров Intel® Xeon® предусмотрен неоднородный доступ к памяти (NUMA). При таком подходе память ЦП делится на группы, также называемые «узлами NUMA», и каждому ядру ЦП присваивается отдельный узел. Хотя, любое ядро ЦП может обращаться к памяти любого узла NUMA, обращения к «своему» узлу (то есть, к локальной памяти) происходит гораздо быстрее. Здесь мы подходим к феномену обучения, распределённого по узлам NUMA, где ядра ЦП, присвоенные каждому из узлов NUMA, действуют как отдельный процесс в группе распределённых процессов. При этом распределением данных по узлам управляет Intel® oneCCL — библиотека Intel, специально отведённая под коллективные коммуникации.

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

import os, time import torch from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler import torch.distributed as dist import torchvision import oneccl_bindings_for_pytorch as torch_ccl import intel_extension_for_pytorch as ipex os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = "29500" os.environ["RANK"] = os.environ.get("PMI_RANK", "0") os.environ["WORLD_SIZE"] = os.environ.get("PMI_SIZE", "1") dist.init_process_group(backend="ccl", init_method="env://") rank = os.environ["RANK"] world_size = os.environ["WORLD_SIZE"] batch_size = 128 num_workers = 0 # определяем датасет и загрузчик данных class FakeDataset(Dataset): def __len__(self): return 1000000 def __getitem__(self, index): rand_image = torch.randn([3, 224, 224], dtype=torch.float32) label = torch.tensor(data=index % 10, dtype=torch.uint8) return rand_image, label train_dataset = FakeDataset() dist_sampler = DistributedSampler(train_dataset) train_loader = DataLoader( dataset=train_dataset, batch_size=batch_size, num_workers=num_workers, sampler=dist_sampler ) # определяем артефакты модели model = torchvision.models.resnet50() criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters()) model.train() model, optimizer = ipex.optimize( model, optimizer=optimizer, dtype=torch.bfloat16 ) # конфигурируем DDP model = torch.nn.parallel.DistributedDataParallel(model) # запускаем цикл обучения # уничтожаем группу процессов dist.destroy_process_group()

К сожалению, на момент подготовки оригинала этой статьи в семействе инстансов Amazon EC2 c7i нет такого, который поддерживал бы работу со множественными областями NUMA. Чтобы протестировать наш скрипт для распределённого обучения, вернёмся к инстансу Amazon EC2 c6i.32xlarge, на котором 64 виртуальных вычислительных ядра и 2 узла NUMA. Проверим установленный у нас экземпляр Intel® oneCCL Bindings для PyTorch и выполним следующую команду (так, как описано в документации здесь):

source $(python -c "import oneccl_bindings_for_pytorch as torch_ccl;print(torch_ccl.cwd)")/env/setvars.sh # В этой команде, приведённой в качестве примера, будут задействованы все сокеты numa, имеющиеся в процессоре, и каждый сокет будет взят в качестве ранга. ipexrun cpu --nnodes 1 --omp_runtime intel train.py

В следующей таблице сравниваются показатели производительности, полученные на инстансе c6i.32xlarge при распределённом обучении и без него:

Распределённое обучение на узлах NUMA (от автора)

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

Обучение на ЦП с применением Torch/XLA

В других статьях (напр., здесь) рассматривается библиотека PyTorch/XLA и обсуждается, как в ней задействуется XLA-компиляция. Такой подход позволяет выполнять обучение с помощью PyTorch на устройствах с поддержкой ускоренной линейной алгебры (XLA), в том числе, на тензорных процессорах (TPU), GPU и ЦП. При XLA, подобно torch-компиляции, методом графовой компиляции генерируется машинный код, оптимизированный для того устройства, на котором он может работать. В рамках проекта OpenXLA, в частности, ставилась цель поддерживать высокопроизводительные вычисления на любых аппаратных бэкендах, в том числе, на ЦП (см. соответствующий RFC здесь). В следующем листинге продемонстрированы корректировки, внесённые в наш исходный (неоптимизированный) скрипт. Они нужны, чтобы задействовать при обучении PyTorch/XLA:

import torch import torchvision import timeimport torch_xla import torch_xla.core.xla_model as xm device = xm.xla_device() model = torchvision.models.resnet50().to(device) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters()) model.train() for idx, (data, target) in enumerate(train_loader): data = data.to(device) target = target.to(device) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() xm.mark_step()

К сожалению (на момент подготовки оригинала этого поста) на нашей игрушечной модели результаты с применением XLA получаются гораздо посредственнее, чем в рассмотренном выше неоптимизированном примере (— приблизительно в 7 раз хуже). Со временем они должны улучшиться, по мере того, как будет доработана поддержка PyTorch/XLA на ЦП.

Результаты

Ниже приводится таблица, в которой обобщены результаты всего подмножества наших экспериментов. Для более полного сравнения добавим результаты обучения нашей модели на инстансе Amazon EC2 g5.2xlarge, развёрнутом на GPU с применением оптимизаций, рассмотренных в этой статье. Показатель образцов в пересчёте на доллар основан на ценах, приведённых на странице Amazon EC2 On-demand pricing ($0,357 в час на c7i.2xlarge и $1,212 в час на g5.2xlarge, на момент подготовки оригинала этой статьи).

Результаты оптимизации производительности (от автора)

Хотя нам и удалось значительно (на 446%) увеличить производительность обучения нашей модели на обычном ЦП, она по-прежнему значительно проигрывает в производительности (оптимизированному) инстансу на GPU. Исходя из наших результатов, обучение на GPU обходится приблизительно в 6,7 раз дешевле. Вероятно, что при дальнейшей настройке производительности и/или при применении дополнительных стратегий оптимизации этот разрыв удалось бы сократить ещё сильнее. Ещё раз подчеркнём, что достигнутые нами сравнительные показатели производительности характерны сугубо для этой модели и этого окружения, в котором идёт исполнение.

Как можно сэкономить благодаря спотовым инстансам Amazon EC2

Чем шире становятся доступны различные типы облачных инстансов на основе ЦП (по сравнению с инстансами на основе GPU), тем проще, соответственно, становится приобрести вычислительные мощности со скидкой. Такие предложения называются «спотовыми инстансами». Amazon EC2 Spot Instances — это инстансы, формируемые из избыточных облачных мощностей, которые могут предлагаться со скидкой до 90% от ценника, выставляемого при приобретении мощностей «по требованию». В обмен на такую скидку AWS оставляет за собой право вытеснить инстанс (практически) без предупреждения. Учитывая высокий спрос на GPU, бывает проще приобрести спотовые инстансы на основе ЦП, нежели их аналоги на GPU.

Притом, что показатели производительности при обучении нашей оптимизированной игрушечной модели на ЦП (именно в подобранном нами облачном окружении) уступают результатам, которые можно получить на GPU, весьма вероятно, что описанные нами шаги оптимизации применительно к другим архитектурам моделей (например, не поддерживаемым GPU) позволят сравниться в производительности между ЦП и GPU или даже превзойти на обычном процессоре графический. Даже в тех случаях, когда полностью устранить такой пробел не удаётся, возможны случаи, в которых из-за дефицита вычислительных мощностей на GPU окажется целесообразно перебросить на ЦП часть рабочих нагрузок, связанных с машинным обучением.

Итоги

Учитывая повсеместное распространение ЦП, важно уметь эффективно использовать их для машинного обучения и/или эксплуатации таких рабочих нагрузок. Это может коренным образом повлиять на продуктивность при разработке и на стратегию развёртывания конечного продукта. Действительно, архитектура ЦП по природе своей хуже подходит для многих задач машинного обучения, чем GPU. Тем не менее, есть множество инструментов и приёмов, позволяющих повысить производительность в таких сценариях. Некоторые из таких инструментов и приёмов были рассмотрены в этой статье.