Работаем с PyTorch на CPU

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

Оглавление

0. Зачем нужен PyTorch на CPU?
1. Железо
  1.1 Процессор
  1.2 Память
  1.3 Детали о моей сборке (шум, материнская плата)
2. Софт
  2.1 Pytorch и CPU
  2.2 Как всё работает по умолчанию
  2.3 DNN-библиотеки для AMD
  2.4 DNN-библиотеки для Intel
  2.5 Многопроцессность в Python
    2.5.0 Fork или Spawn?
    2.5.1 Загрузчики данных
    2.5.2 Хитрости, без которых многопроцессность в PyTorch не работает или жутко глючит
    2.5.3 Используем pool в PyTorch. Отличие map от imap()
3. Пара слов о быстродействии. Итоги

0. Зачем нужен PyTorch на CPU?

Если вы хотите прикоснуться к миру машинного обучения, то вам не нужна дискретная видеокарта. Достаточно того, что у вас есть итак — обычный CPU. А ещё лучше — многоядерный процессор, который имеет и массу других полезных применений.

Сейчас покупка мощной видеокарты Nvidia требует серьёзных денег и намерений. Топовые AMD-карты, которые поддерживает PyTorch, тоже весьма дорогие.

А если у вас есть видеокарта, подходящая для PyTorch, то вы сможете использовать инфу из этой статьи для ускорения даталоадеров или параллельно считать и на CPU, и на видеокарте.

Есть ещё пара неочевидных причин. Первая. Для некоторых больших LLM-моделей требуется огромное количество оперативной памяти, а видеокарты с таким объёмом памяти стоят совсем фантастических денег. А вот 256ГБ серверной памяти для ПК могут стоить всего, например, 44т.

Вторая. При обучении модели требуется памяти значительно больше, чем занимают её веса. Там требуется хранить градиенты (дельты изменения весов), и чем больше batch, тем больше требуется места. Батч — это набор данных из обучающей выборки, который за один раз пересылается в устройство, обучающее нейросеть. Чем он больше, тем более точно и быстро обучается нейросеть. Обычно используются батчи от 64 до 256. И многие даже считают, что больше и не нужно, так как есть мнение, что небольшие батчи «шумят» и этим улучшают генерализацию обучения.

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

Например, батч размером 25000 немыслим для видеокарты с несколькими гигами при обучении модели уже с несколькими тысячами нейронов, хотя в некоторых задачах он будет оптимален. На CPU и при наличии RAM он будет посчитан без проблем.

1. Железо

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

▍ 1.1 Процессор

Есть 4 главных принципа, по которым выбирается процессор для обучения нейросетей:

1. Цена.
2. Многоядерность.
3. Размер кеша L3.
4. Количество каналов памяти.

С ценой всё ясно. Чем дешевле в пересчёте на ядро и мегагерцы — тем лучше. Количество ядер приоритетнее мегагерцев.

С точки зрения максимального количества ядер и цены обращают на себя внимание процессоры от AMD серий Threadripper и Epyc, а также интересно подумать и о Xeon V3 и V4. Правда, сравнив производительность последних с эпиками второй серии, я понял, что это неоптимальный выбор. Эпики можно апгрейдить на том же сокете, а вот ксеоны только на свалку/барахолку.

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

Процессоры Epyc интересны тем, что, фактически это SoC (система-на-чипе). Амд сделала так, что теоретически в плату из-под самого первого Epyc можно вставить эпик третьего поколения. То есть чипсет вмонтирован в сам процессор. Поэтому им нужен минимум драйверов.

Причём плата первого поколения стоит очень дёшево. Например, H11-SSLi. Она поддерживает 7xx1 и 7xx2 поколения (первое и второе). Но технически (при правильном биосе) могла бы поддерживать и третье.

Но жадные производители системных плат искусственно ограничивают работоспособность плат со следующими поколениями. Или не заявляют о такой возможности, или не обновляют биосы. Я не стал пытаться взломать систему и купил на барахолке H12-SSLi (за 33т в начале 2023 года). Она поддерживает второе и третье поколение. Третье поколение Epyc интересно тем, что поддерживает процессоры с 768МБ кеша L3.

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

Рост производительности при увеличении размера L3 проявляется практически во всех задачах. Вот данные о том, что Ryzen 5800X3D (8 ядер, L3 96MB) побеждает 5950X (16 ядер, L3 128MB) в задачах машинного обучения. А при равном числе ядер 5800X3D будет на 13.6% быстрее, чем 5800X. Разница только в размере L3 кеша.

По соображениям соотношения цены и производительности я купил на барахолке Epyc 7542 в декабре 2023 года за 33т.р. Это Epyc 2 поколения, 32 ядра и 64 потока, 128MB L3, 8 каналов памяти ECC Reg до 3200MHz, 2.9 ГГц или до 3.4 ГГц в режиме Max Boost Clock.

В дальнейшем я хочу перейти на 64х-ядерный CPU серии 7xx3, а потом на 64 ядра + 768мб кеш.

Наличие развитой многоканальности (например, серверный Epyc имеет 8 каналов памяти против 2 у десктопного Ryzen) позволяет более эффективно использовать большое число ядер, расширяя бутылочное горло по доступу процессора к данным.

▍ 1.2 Память

С памятью обнаружилось 2 интересных момента:

1. Для максимального быстродействия нужно использовать все каналы. В моих экспериментах с PyTorch разница оказалась существенной. Около +20% в 8-канальном режиме против 4-канального режима.
2. Нужно использовать 2-ранговую память 2R, а не одноранговую память 1R. Дело в том, что на доставание данных из памяти тратится довольно много циклов, а контроллер памяти из двухранговой может параллельно доставать данные их 2 банков на одном канале. Считается, что такая память быстрее на 5%.

▍ 1.3 Детали о моей сборке (шум, материнская плата)

Серверная материнская плата H12-SSLi имеет ATX-формат, поэтому встанет во многие большие корпуса. Единственная сложность заключалась в том, что мне пришлось разделить проводок Power Led, так как рядом на гребёнке не нашлось GRD (земли). Рядом с гребёнкой нашлась земля в другом разъёме, куда я её и воткнул.

Я очень переживал, что Epyc в товер-корпусе будет страшно шуметь. Однако с выбранным мною кулером Epyc Arctic Freezer 4U SP3 (до 300W TDP) все беспокойства оказались совершенно напрасными. Даже с нагрузкой на все ядра комп довольно тихий. При частичной нагрузке до 70% вообще не слышно кулера. Сам кулер заказал через посредника из Средней Азии.

Windows 10 LTSC поддерживает на практике.

2. Софт

▍ 2.1 Pytorch и CPU

Со страницы загрузки можно скачать специальную версию PyTorch для CPU. Она будет весить значительно меньше, чем версия с поддержкой Cuda.

▍ 2.2 Как всё работает по умолчанию

Когда я запустил свою первую модель на CPU, то был удивлён, что в списке процессов вижу только экземпляры Python. Я ожидал, что Pytorch скомпилирует мой код в некий исполняемый файл, который я и увижу, набрав top (у меня Linux).

Однако я видел только множество экземпляров Python. Оказалось, что PyTorch компилирует код в разделяемые библиотеки, которые потом загружаются в пространство процессов Python.

Что мне не понравилось — загрузка процессора, как бы я ни старался, не доходила до 100%. Дело в том, что стандартный бэкенд для Python генерирует слишком универсальный код для всех возможных процессоров, и он неэффективный, даже если везде использовать функцию compile() из PyTorch.

▍ 2.3 Специализированные библиотеки для CPU

Поэтому я стал копать в сторону специализированных библиотек от производителей процессоров для нейросетей. Для AMD я нашёл библиотеку ZenDNN. И, да, она стоила того!

Загрузка процессора при обучении сразу подскочила до 100%, и PyTorch стал намного производительнее. AMD её позиционирует как библиотеку для Epyc 4-го поколения, но на практике гарантированно работает, начиная с поколения 7xx2. Я не удивлюсь, если заработает и с настольными процессорами Ryzen, начиная со второго поколения. Но сам не проверял. Абсолютно точно, что не работает с FX-процессорами AMD.

AMD уже приготовила версию ZenDNN, которая включает в себя PyTorch, поэтому имеет смысл ставить сразу ZenDNN в голое виртуальное окружение.

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

#Используем фиксированный seed
torch.backends.zendnn.deterministics = True

Что интересно: даже предыдущая версия ZenDnn, которая поддерживала только PyTorch 1.3 и Python 3.10, была примерно на 40% быстрее, чем связка из PyTorch 2 + Python 3.11 + обычный бэкенд для CPU. Это несмотря на то, что в PyTorch 2.x появилась функция compile().

В мае 2024 года произошёл релиз новой версии ZenDNN, которая поддерживает PyTorch 2.1 и Python 3.11.

▍ 2.4 DNN-библиотеки для Intel

Intel имеет собственный набор библиотек для PyTorch. Как для CPU, так и для GPU.

Вот для CPU.

▍ 2.5 Многопроцессность в Python

Несмотря на то что, как правило, DNN-библиотеки умеют правильно использовать многоядерные процессоры, в некоторых случаях улучшить производительность помогает многопроцессность в Python.

В принципе Python умеет поддерживать и многотредовость, но это на порядок более глючно, чем многопроцессность. Так как написать thread-save-приложение значительно сложнее, чем однопроцессное или даже чем многопроцессное.

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

▍ 2.5.0 Fork или Spawn?

В Линукс есть уникальный системный вызов fork(), который клонирует процесс, а его данные делает общими по схеме copy-on-write. Это классная штука, и если у вас Линукс, проще всего использовать её. И fork быстрее spawn.

Обычно метод запуска fork используется по умолчанию. Но мы можем сами назначать метод создания воркеров.

from torch import multiprocessing as mp
import torch

mp.set_start_method('fork')

При методе запуска процесса Spawn запускается отдельная копия Python, в неё передаётся текущий интерпретируемый файл. В Windows доступен только Spawn. В коде нужно как-то отличать главный файл от таких же, но побочных, запущенных только для параллелизации и ускорения вычислений. Обычно это делается так:

# всякие функции и объекты
if __name__ == '__main__':
  # это главная ветка программы
  # ... инициализируем и запускаем модели
  exit()
else:
  # это spawn-ветка
  # тут ничего не происходит
  # так как вызываемые функции находятся в самом верху файла

▍ 2.5.1 Загрузчики данных

В PyTorch есть класс DataLoader, наследуя от которого можно сделать свои загрузчики.

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
           batch_sampler=None, num_workers=0, collate_fn=None,
           pin_memory=False, drop_last=False, timeout=0,
           worker_init_fn=None, *, prefetch_factor=2,
           persistent_workers=False)

Поскольку загрузчики, как правило, пишутся на Python, то работают они медленно, поэтому, если структура данных позволяет, лучше использовать многопоточные загрузчики.

Нас интересуют параметры num_workers — число потоков загрузчика, worker_init_fn — функция инициализации загрузчика в отдельном потоке, persistent_workers — оставлять ли воркеры в памяти. В большинстве случаев достаточно указать только num_workers, и загрузка данных в модель ускорится, если ваш не заточен под однопроцессность.

▍ 2.5.2 Хитрости, без которых многопроцессность не работает или жутко глючит

Самая главная хитрость — это предварять любой многопроцессный код командой отключения многопоточности. Скомпилированный код PyTorch, который Python загружает в своё пространство, по умолчанию активно использует многопоточность и в сочетании с многопроцессностью всё начинает активно глючить.

torch.set_num_threads(1)
.... многопроцессный код
torch.set_num_threads(PROCESS_NUM)

▍ 2.5.3 Используем pool в PyTorch. Отличие map от imap()

Самый простой способ ускорить Python-код в PyTorch — это распараллелить его. А самый простой способ распараллелить — это использовать pool() из модуля multiprocessing от PyTorch.

PyTorch имеет свою обёртку над питоновской стандартной библиотекой многопроцессности. Эта обёртка умеет передавать тензоры в другие потоки. Делается это через разделяемую память (shared mem) в Linux. Однако если вам известно, что тензор будет передаваться в другой поток, то лучше его перегнать в разделяемую память заранее, так заметно быстрее.

Пример многопроцессного кода, который генерирует датасет, который затем будет использоваться в загрузчике.

# мультипроцессинг для работы с тензорами
from torch import multiprocessing as mp
# мультипроцессинг для всего остального
import multiprocessing as mpPython
import torch

# заполняем кусок тензора в 1 потоке
def fill_mp_tensor(args):
  (tData, tVer, gSeed, min, max) = args
  
  for i in range(min, max):
    tData[i] = create_data_value(gSeed, i, device=tData.device)
    tVer[i] = create_verification_value(gSeed, i, device=tVer.device)

# создаём тензоры мультипоточным образом
def createTensorMP(tData, tVer, gSeed, INTERVAL, PROCESSES_NUM):
  
  if PROCESSES_NUM < INTERVAL:
    nTask = PROCESSES_NUM
  else:
    nTask = 1
  
  chunk = INTERVAL // nTask
  
  if chunk * nTask < INTERVAL:
    nTask += 1
  
  torch.set_num_threads(1)
  
  with mp.Pool(PROCESSES_NUM) as pool:
    pool.map(
      fill_mp_tensor,
      [(tData, tVer, gSeed, i*chunk, min((i+1)*chunk, INTERVAL)) for i in range(nTask)]
    )
    
  torch.set_num_threads(PROCESSES_NUM)

# Тут код определения сети, функции потерь, оптимайзера
# ......

# чтобы меньше занять shared mem, создаём вектора однобайтными
tData = torch.zeros(RND_DS_SIZE, 512, dtype=torch.uint8).share_memory_()
tVer = torch.zeros(RND_DS_SIZE, 2, dtype=torch.uint8).share_memory_()

# генерируем данные для датасета в мультипоточном режиме      
createATensorMP(tData, tVer, SEED_NUMBER, RND_DS_SIZE, PROCESSES_NUM)
# PyTorch умеет тренировать сетки в fp32
tData = tData.float()
tVer = tVer.float()

ds = TensorDataset(tData, tVer)
# загрузчик работает в мультипроцессном режиме и каждый раз перемешивает данные
dl = DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=LOADER_NUM)

# тренируем сеть, пока не достигнем нужных параметров аккуратности и потерь      
# повторная тренировка, если плохая точность
while fAccuracy*100 < 50.0 or fAvgLoss > 0.693:  
  fAccuracy, fAvgLoss = train(NN, dl, loss_fn, optimizer)

Теперь об отличии map() от imap(). Первый должен обработать всё множество и собрать в результаты в массив.

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

with mp.Pool(PROCESSES_NUM) as pool:
    iRes = pool.imap(
      mp_func,
      [(COUNT_BF, i, i*g) for i in range(nTask)]
      [(arg1, arg2, i*chunk, min((i+1)*chunk, INTERVAL)) for i in range(nTask)]
    )
    
    for (res1, res2, res3) in iRes:
      work(res1, res2, res3)

В моих проектах в большинстве случаев достаточно pool.map().

3. Пара слов о быстродействии. Итоги

На сравнительно малых и небольших моделях (тысячи нейронов) я получал примерно такое же быстродействие на моём процессоре (32 ядра), как и на ускорителе Tesla T4 GPU 16ГБ из Google Colab.

Работать на своём компьютере значительно удобнее, чем на Google Colab. По памяти и процессору нет практических ограничений. И скорость всего, что связано с процессором и мультипроцессностью, значительно выше на домашнем компьютере.

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

Безусловно, всё зависит от задач. И на нейросетях с миллионами параметров по идее даже супермногоядерные процессоры должны сдуваться, но для образовательных задач в PyTorch и небольших проектов CPU вполне достаточно.