Советы по ускорению ИИ/машинного обучения на ЦП — Часть 2
Делиться
Разработка и развертывание моделей ИИ/машинного обучения может быть чрезвычайно дорогостоящим мероприятием. Многие из наших статей были посвящены широкому спектру советов, приемов и методов анализа и оптимизации производительности рабочих нагрузок ИИ/машинного обучения во время их выполнения. Наш аргумент состоит из двух частей:
- Анализ и оптимизация производительности должны быть неотъемлемой частью каждого проекта по разработке ИИ/машинного обучения, и,
- Для достижения существенного повышения производительности и снижения затрат не требуется высокая степень специализации. С этим справится любой разработчик ИИ/машинного обучения. И каждый разработчик ИИ/машинного обучения должен это уметь.
В недавней публикации мы рассмотрели проблему оптимизации рабочей нагрузки машинного обучения на процессоре Intel® Xeon®. Мы начали с обзора ряда сценариев, в которых процессор может быть наилучшим выбором для вывода ИИ/машинного обучения даже в эпоху множества специализированных чипов для вывода ИИ. Затем мы представили простую модель классификации изображений на PyTorch и продемонстрировали множество методов повышения ее производительности на экземпляре Amazon EC2 c7i.xlarge, работающем на процессорах Intel Xeon Scalable 4-го поколения. В этой публикации мы расширяем наше обсуждение на собственные процессоры AWS на базе архитектуры Arm — Graviton. Мы вернемся ко многим оптимизациям, которые мы обсуждали в наших предыдущих публикациях — некоторые из которых потребуют адаптации к процессору Arm — и оценим их влияние на ту же простую модель. Учитывая существенные различия между процессорами Arm и Intel, пути к оптимальной конфигурации могут быть разными.
АВС Гравитон
AWS Graviton — это семейство процессоров на базе Arm Neoverse, специально разработанных и созданных AWS для обеспечения оптимального соотношения цены, производительности и энергоэффективности. Их специализированные механизмы для векторной обработки (NEON и SVE/SVE2) и умножения матриц (MMLA), а также поддержка операций с плавающей запятой типа Bfloat16 (начиная с Graviton3) делают их привлекательным вариантом для выполнения ресурсоемких вычислительных задач, таких как вывод данных в области ИИ/машинного обучения. Для обеспечения высокопроизводительной работы ИИ/машинного обучения на Graviton весь программный стек был оптимизирован для его использования:
- Ядра низкоуровневых вычислений из библиотеки Arm Compute Library (ACL) оптимизированы для использования аппаратных ускорителей Graviton (например, SVE и MMLA).
- Библиотеки промежуточного программного обеспечения для машинного обучения, такие как oneDNN и OpenBLAS, направляют операции глубокого обучения и линейной алгебры к специализированным ядрам ACL.
- Фреймворки для искусственного интеллекта и машинного обучения, такие как PyTorch и TensorFlow, скомпилированы и настроены для использования этих оптимизированных бэкэндов.
В этой статье мы будем использовать экземпляр Amazon EC2 c8g.xlarge, работающий на базе четырех процессоров AWS Graviton4 и образа AWS ARM64 PyTorch Deep Learning AMI (DLAMI).
Цель этой статьи — продемонстрировать советы по повышению производительности на экземпляре AWS Graviton. Важно отметить, что мы не ставим перед собой цель сравнивать AWS Graviton с альтернативными процессорами и не выступаем за использование одного процессора вместо другого. Оптимальный выбор процессора зависит от множества факторов, выходящих за рамки этой статьи. Одним из важных факторов будет максимальная производительность вашей модели на каждом процессоре. Другими словами: насколько эффективно мы можем использовать имеющиеся ресурсы? Таким образом, принятие обоснованного решения о выборе наилучшего процессора является одной из причин оптимизации производительности на каждом из них.
Еще одна причина оптимизации производительности нашей модели для различных устройств вывода — повышение ее портативности. Сфера ИИ/машинного обучения чрезвычайно динамична, и устойчивость к меняющимся обстоятельствам имеет решающее значение для успеха. Нередко вычислительные ресурсы определенных типов внезапно становятся недоступными или дефицитными. И наоборот, увеличение мощности экземпляров AWS Graviton может означать их доступность со значительными скидками, например, на рынке спотовых экземпляров Amazon EC2, что открывает возможности для экономии средств, которые вы не захотите упустить.
Отказ от ответственности
Представленный фрагмент кода, шаги по оптимизации и полученные результаты призваны продемонстрировать преимущества оптимизации производительности машинного обучения на экземпляре AWS Graviton. Эти преимущества могут значительно отличаться от результатов, которые вы можете получить со своей собственной моделью и средой выполнения. Пожалуйста, не полагайтесь на точность или оптимальность информации в этом посте. Не следует воспринимать упоминание какой-либо библиотеки, фреймворка или платформы как рекомендацию к их использованию.
Оптимизация вывода в AWS Graviton
Как и в нашей предыдущей статье, мы продемонстрируем этапы оптимизации на примере простой модели классификации изображений:
import torch, torchvision import time def get_model(channels_last=False, compile=False): model = torchvision.models.resnet50() if channels_last: model= model.to(memory_format=torch.channels_last) model = model.eval() if compile: model = torch.compile(model) return model def get_input(batch_size, channels_last=False): batch = torch.randn(batch_size, 3, 224, 224) if channels_last: batch = batch.to(memory_format=torch.channels_last) return batch def get_inference_fn(model, enable_amp=False): def infer_fn(batch): with torch.inference_mode(), torch.amp.autocast( 'cpu', dtype=torch.bfloat16, enabled=enable_amp ): output = model(batch) return output return infer_fn def benchmark(infer_fn, batch): # разминка for _ in range(20): _ = infer_fn(batch) iters = 100 start = time.time() for _ in range(iters): _ = infer_fn(batch) end = time.time() return (end — start) / iters batch_size = 1 model = get_model() batch = get_input(batch_size) infer_fn = get_inference_fn(model) avg_time = benchmark(infer_fn, batch) print(f»nСреднее количество выборок в секунду: {(batch_size/avg_time):.2f}»)
Начальная пропускная способность составляет 12 выборок в секунду (SPS).
Обновите PyTorch до последней версии.
В то время как в нашем DLAMI установлена версия PyTorch 2.8, на момент написания этой статьи последняя версия PyTorch — 2.9. Учитывая стремительные темпы развития в области ИИ/машинного обучения, настоятельно рекомендуется использовать самые актуальные пакеты библиотек. В качестве первого шага мы обновимся до PyTorch 2.9, который включает в себя ключевые обновления для его бэкенда Arm.
pip3 install -U torch torchvision —index-url https://download.pytorch.org/whl/cpu
В случае нашей модели в ее исходной конфигурации обновление версии PyTorch не оказывает никакого эффекта. Однако этот шаг имеет решающее значение для получения максимальной отдачи от методов оптимизации, которые мы будем оценивать.
Пакетный вывод
Чтобы уменьшить накладные расходы на запуск и повысить эффективность использования аппаратных ускорителей, мы группируем образцы и применяем пакетный вывод. В таблице ниже показано, как производительность модели изменяется в зависимости от размера пакета:
Оптимизация памяти
Для оптимизации выделения и использования памяти мы применяем ряд методов из нашей предыдущей публикации. К ним относятся формат памяти channels-last, автоматическая смешанная точность с типом данных bfloat16 (поддерживается начиная с Graviton3), библиотека выделения памяти TCMalloc и выделение больших страниц памяти. Подробности см. в предыдущей публикации. Мы также включаем быстрый математический режим ядер ACL GEMM и кэширование примитивов ядра — две оптимизации, которые описаны в официальных рекомендациях по запуску вывода PyTorch на Graviton.
Ниже приведены инструкции командной строки, необходимые для включения этих оптимизаций:
# Установка TCMalloc sudo apt-get install google-perftools # Программирование использования TCMalloc export LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libtcmalloc.so.4 # Включение выделения памяти для больших страниц export THP_MEM_ALLOC_ENABLE=1 # Включение быстрого математического режима ядер GEMM export DNNL_DEFAULT_FPMATH_MODE=BF16 # Установка емкости LRU-кэша для кэширования примитивов ядра export LRU_CACHE_CAPACITY=1024
В следующей таблице показано влияние последовательно применяемых оптимизаций памяти:
В случае нашей модельной программы наибольший эффект оказали оптимизации channels-last и bfloat16-mixed precision. После применения всех оптимизаций памяти средняя пропускная способность составила 53,03 SPS.
Компиляция модели
Поддержка компиляции PyTorch для AWS Graviton — это направление, которому команда AWS Graviton уделяет особое внимание. Однако в случае нашей тестовой модели это приводит к неболькому снижению пропускной способности: с 53,03 SPS до 52,23.
Многопроцессный вывод
Хотя обычно этот подход применяется в системах с гораздо большим количеством виртуальных процессоров, мы демонстрируем реализацию многопоточного вывода, модифицировав наш скрипт для поддержки закрепления ядер:
if __name__ == '__main__': # привязываем процессоры к рангу рабочего процесса import os, psutil rank = int(os.environ.get('RANK','0')) world_size = int(os.environ.get('WORLD_SIZE','1')) cores = list(range(psutil.cpu_count(logical=True))) num_cores = len(cores) cores_per_process = num_cores // world_size start_index = rank * cores_per_process end_index = (rank + 1) * cores_per_process pid = os.getpid() p = psutil.Process(pid) p.cpu_affinity(cores[start_index:end_index]) batch_size = 8 model = get_model(channels_last=True) batch = get_input(batch_size, channels_last=True) infer_fn = get_inference_fn(model, enable_amp=True) avg_time = benchmark(infer_fn, batch) print(f»nСреднее количество выборок в секунду: {(batch_size/avg_time):.2f}»)
Мы отмечаем, что в отличие от других типов экземпляров AWS EC2 CPU, каждый виртуальный процессор Graviton напрямую соответствует одному физическому ядру ЦП. Мы используем утилиту torchrun для запуска четырех рабочих процессов, каждый из которых работает на одном ядре ЦП:
export OMP_NUM_THREADS=1 # установить один поток OpenMP на каждый рабочий процесс torchrun —nproc_per_node=4 main.py
В результате достигается пропускная способность 55,15 SPS, что на 4% выше нашего предыдущего лучшего результата.
Квантование INT8 для ARM
Еще одна область активной разработки и постоянного совершенствования на Arm — это квантизация INT8. Инструменты квантизации INT8 обычно тесно связаны с целевым типом экземпляра. В нашей предыдущей статье мы продемонстрировали экспорт квантизации PyTorch 2 с бэкендом X86 через Inductor, используя библиотеку TorchAO (0.12.1). К счастью, в последних версиях TorchAO есть специальный квантизатор для Arm. Обновленная последовательность квантизации показана ниже. Как и в нашей предыдущей статье, нас интересует только потенциальное влияние на производительность. На практике квантизация INT8 может существенно повлиять на качество модели и потребовать более сложной стратегии квантизации.
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e import torchao.quantization.pt2e.quantizer.arm_inductor_quantizer as aiq def quantize_model(model): x = torch.randn(4, 3, 224, 224).contiguous( memory_format=torch.channels_last) example_inputs = (x,) batch_dim = torch.export.Dim(«batch») with torch.no_grad(): exported_model = torch.export.export( model, example_inputs, dynamic_shapes=((batch_dim, torch.export.Dim.STATIC, torch.export.Dim.STATIC, torch.export.Dim.STATIC), ) ).module() quantizer = aiq.ArmInductorQuantizer() quantizer.set_global(aiq.get_default_arm_inductor_quantization_config()) prepared_model = prepare_pt2e(exported_model, quantizer) prepared_model(*example_inputs) converted_model = convert_pt2e(prepared_model) optimized_model = torch.compile(converted_model) return optimized_model batch_size = 8 model = get_model(channels_last=True) model = quantize_model(model) batch = get_input(batch_size, channels_last=True) infer_fn = get_inference_fn(model, enable_amp=True) avg_time = benchmark(infer_fn, batch) print(f»nСреднее количество выборок в секунду: {(batch_size/avg_time):.2f}»)
В результате пропускная способность составляет 56,77 SPS, что на 7,1% выше, чем у решения на основе bfloat16.
Компиляция AOT с использованием ONNX и OpenVINO
В нашей предыдущей публикации мы рассмотрели методы предварительной компиляции моделей (AOT) с использованием Open Neural Network Exchange (ONNX) и OpenVINO. Обе библиотеки включают специальную поддержку для работы на AWS Graviton (например, см. здесь и здесь). Для экспериментов в этом разделе требуется установка следующих библиотек:
pip install onnxruntime onnxscript openvino nncf
Следующий фрагмент кода демонстрирует компиляцию и выполнение модели на Arm с использованием ONNX:
def export_to_onnx(model, onnx_path=»resnet50.onnx»): dummy_input = torch.randn(4, 3, 224, 224) batch = torch.export.Dim(«batch») torch.onnx.export( model, dummy_input, onnx_path, input_names=[«input»], output_names=[«output»], dynamic_shapes=((batch, torch.export.Dim.STATIC, torch.export.Dim.STATIC, torch.export.Dim.STATIC), ), dynamo=True ) return onnx_path def onnx_infer_fn(onnx_path): import onnxruntime as ort sess = ort.InferenceSession( onnx_path, providers=[«CPUExecutionProvider»] ) sess_options = ort.SessionOptions() sess_options.add_session_config_entry( «mlas.enable_gemm_fastmath_arm64_bfloat16», «1») input_name = sess.get_inputs()[0].name def infer_fn(batch): result = sess.run(None, {input_name: batch}) return result return infer_fn batch_size = 8 model = get_model() onnx_path = export_to_onnx(model) batch = get_input(batch_size).numpy() infer_fn = onnx_infer_fn(onnx_path) avg_time = benchmark(infer_fn, batch) print(f»nСреднее количество выборок в секунду: {(batch_size/avg_time):.2f}»)
Следует отметить, что среда выполнения ONNX поддерживает выделенный ACL-ExecutionProvider для работы на Arm, но для этого требуется специальная сборка ONNX (на момент написания этой статьи), что выходит за рамки данной публикации.
В качестве альтернативы, мы можем скомпилировать модель с помощью OpenVINO. Приведенный ниже блок кода демонстрирует его использование, включая опцию квантования INT8 с использованием NNCF:
import openvino as ov import nncf def openvino_infer_fn(compiled_model): def infer_fn(batch): result = compiled_model([batch])[0] return result return infer_fn class RandomDataset(torch.utils.data.Dataset): def __len__(self): return 10000 def __getitem__(self, idx): return torch.randn(3, 224, 224) quantize_model = False batch_size = 8 model = get_model() calibration_loader = torch.utils.data.DataLoader(RandomDataset()) calibration_dataset = nncf.Dataset(calibration_loader) if quantize_model: # квантизация модели PyTorch model = nncf.quantize(model, calibration_dataset) ovm = ov.convert_model(model, example_input=torch.randn(1, 3, 224, 224)) ovm = ov.compile_model(ovm) batch = get_input(batch_size).numpy() infer_fn = openvino_infer_fn(ovm) avg_time = benchmark(infer_fn, batch) print(f»nСреднее количество выборок в секунду: {(batch_size/avg_time):.2f}»)
В случае нашей модельной программы компиляция с использованием OpenVINO приводит к дополнительному увеличению пропускной способности до 63,48 SPS, но квантование NNCF разочаровывает, приводя к показателю всего 55,18 SPS.
Результаты
Результаты наших экспериментов суммированы в таблице ниже:
Как и в нашем предыдущем посте, мы повторно провели эксперименты на второй модели — Vision Transformer (ViT) из библиотеки timm — чтобы продемонстрировать, как влияние обсуждаемых нами оптимизаций во время выполнения может варьироваться в зависимости от деталей модели. Результаты представлены ниже:
Краткое содержание
В этой статье мы рассмотрели ряд относительно простых методов оптимизации и применили их к двум простым моделям PyTorch. Как показали результаты, влияние каждого шага оптимизации может сильно различаться в зависимости от деталей модели, и путь к максимальной производительности может быть разным. Шаги, представленные в этой статье, были лишь затравкой; несомненно, существует множество других методов оптимизации, которые могут обеспечить еще большую производительность.
В процессе работы мы отметили множество библиотек для ИИ/машинного обучения, которые обеспечили глубокую поддержку архитектуры Graviton, а также, по-видимому, непрерывные усилия сообщества по постоянной оптимизации. Достигнутые нами улучшения производительности в сочетании с этой очевидной преданностью делу доказывают, что AWS Graviton прочно занимает место в «высшей лиге» в плане выполнения ресурсоемких задач ИИ/машинного обучения.
Источник: towardsdatascience.com
