Анализ и оптимизация производительности модели PyTorch — Часть 8
Делиться

Узкие места производительности в конвейере ввода данных модели машинного обучения, работающей на GPU, могут быть особенно раздражающими. В большинстве рабочих нагрузок хост (CPU) и устройство (GPU) работают в тандеме: CPU отвечает за подготовку и подачу данных, в то время как GPU выполняет тяжелую работу — выполнение модели, выполнение обратного распространения во время обучения и обновление весов.
В идеальной ситуации мы хотим, чтобы GPU — самый дорогой компонент нашей инфраструктуры AI/ML — был максимально загружен. Это приводит к более быстрым циклам разработки, более низким затратам на обучение и уменьшению задержек при развертывании. Чтобы достичь этого, GPU должен непрерывно снабжаться входными данными. В частности, мы хотели бы предотвратить возникновение «голодания GPU» — ситуации, в которой наш самый дорогой ресурс простаивает, ожидая входных данных. К сожалению, «голодание GPU» из-за узких мест в конвейере входных данных встречается довольно часто и может значительно снизить эффективность системы. Таким образом, разработчикам AI/ML важно иметь надежные инструменты и стратегии для диагностики и решения таких проблем.
Этот пост — восьмой в нашей серии на тему Анализа и оптимизации производительности модели PyTorch — представляет простую стратегию кэширования для выявления узких мест в конвейере ввода данных. Как и в предыдущих постах, мы стремимся усилить две ключевые идеи:
- Разработчики искусственного интеллекта/машинного обучения должны нести ответственность за производительность своих моделей во время выполнения.
- Вам не нужно быть экспертом в области CUDA или систем, чтобы реализовать существенную оптимизацию производительности.
Начнем с описания некоторых распространенных причин нехватки GPU. Затем мы представим нашу основанную на кэшировании стратегию для выявления и анализа проблем производительности входного конвейера. Завершим обзором набора практических инструментов, приемов и методов (TTT) для преодоления узких мест производительности в входном конвейере данных.
Для облегчения обсуждения мы определим игрушечную модель PyTorch и связанный с ней конвейер ввода данных. Код, которым мы поделимся, предназначен для демонстрационных целей — пожалуйста, не полагайтесь на его правильность или оптимальность. Кроме того, пожалуйста, не упоминайте какой-либо инструмент или метод как одобрение его использования.
Игрушечная модель PyTorch
Мы определяем простую модель классификации изображений на основе PyTorch:
неопределенный
Мы определяем синтетический набор данных с рядом преобразований — намеренно разработанных для включения серьезного узкого места входного конвейера. Для получения более подробной информации об определении набора данных см. этот пост.
import numpy as np from PIL import Image from torchvision.datasets.vision import VisionDataset import torchvision.transforms as T class FakeDataset(VisionDataset): def __init__(self, transform): super().__init__(root=None, transform=transform) self.size = 10000 def __getitem__(self, index): # создать случайное изображение 1024×1024 img = Image.fromarray(np.random.randint( low=0, high=256, size=(input_img_size[0], input_img_size[1], 3), dtype=np.uint8 )) # создать случайную метку target = np.random.randint( low=0, high=num_classes, dtype=np.uint8).item() # применить преобразования img = self.transform(img) return img, target def __len__(self): return self.size class RandomMask(torch.nn.Module): def __init__(self, ratio=0.25): super().__init__() self.ratio=ratio def dilate_mask(self, mask): # выполнить расширение 4 соседей по маске from scipy.signal import convolve2d dilated = convolve2d(mask, [[0, 1, 0], [1, 1, 1], [0, 1, 0]], mode='same').astype(bool) return dilated def forward(self, img): mask = np.random.uniform(size=(img_size, img_size)) < self.ratio dilated_mask = torch.unsqueeze(torch.tensor(self.dilate_mask(mask)),0) dilated_mask = dilated_mask.expand(3,-1,-1) img[dilated_mask] = 0. return img class ConvertColor(torch.nn.Module): def __init__(self): super().__init__() self.A=torch.tensor( [[0.299, 0.587, 0.114], [-0.16874, -0.33126, 0.5], [0.5, -0.41869, -0.08131]] ) self.b=torch.tensor([0.,128.,128.]) def forward(self, img): img = img.to(dtype=torch.get_default_dtype()) img = torch.matmul(self.A,img.view([3,-1])).view(img.shape) img = img + self.b[:,None,None] return img class Scale(object): def __call__(self, img): return img.to(dtype=torch.get_default_dtype()).div(255) transform = T.Compose([T.PILToTensor(), T.RandomCrop(img_size), RandomMask(), ConvertColor(), Scale()]) train_set = FakeDataset(transform=transform) train_loader = torch.utils.data.DataLoader(train_set, batch_size=256, num_workers=4, pin_memory=True)
Далее мы определяем модель, функцию потерь, оптимизатор, шаг обучения и цикл обучения, которые мы заключаем в контекстный менеджер PyTorch Profiler для сбора данных о производительности.
из статистики импорт среднее, дисперсия из времени импорт время устройство = torch.device(«cuda:0») модель = Net().cuda(device) критерий = nn.CrossEntropyLoss().cuda(device) оптимизатор = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9) def train_step(model, criteria, optimizer, inputs, labels): выходы = model(inputs) потери = criteria(outputs, labels) optimizer.zero_grad(set_to_none=True) потери.backward() оптимизатор.step() модель.train() t0 = time() времена = [] с torch.profiler.profile( schedule=torch.profiler.schedule(wait=10, warmup=2, active=10, repeat=1), on_trace_ready=torch.profiler.tensorboard_trace_handler('/tmp/prof'), record_shapes=True, profile_memory=True, with_stack=True ) as prof: for step, data in enumerate(train_loader): # копирование данных на устройство inputs = data[0].to(device=device, non_blocking=True) labels = data[1].to(device=device, non_blocking=True) # запуск шага обучения train_step(model, criteria, optimizer, inputs, labels) prof.step() times.append(time()-t0) t0 = time() if step >= 100: break print(f'average time: {mean(times[1:])}, variance: {variance(times[1:])}')
Для наших экспериментов мы используем экземпляр Amazon EC2 g5.xlarge (содержащий NVIDIA A10G GPU и 4 vCPU) с запущенным PyTorch (2.6) Deep Learning AMI (DLAMI). Запуск нашего игрушечного скрипта в этой среде приводит к средней пропускной способности 0,89 шагов в секунду, неудовлетворительной загрузке GPU в 22% и к следующему профилирующему следу:

Как подробно обсуждалось в предыдущем посте, трассировка профилирования показывает четкую картину голодания GPU — когда GPU проводит большую часть своего времени в ожидании данных от PyTorch DataLoader. Это говорит о том, что в конвейере ввода данных есть узкое место производительности, которое не позволяет достаточно быстро подготовить пакеты ввода, чтобы полностью загружать GPU. Важно отметить, что проблемы с производительностью конвейера ввода могут возникать из разных источников. В случае нашего игрушечного примера причина узкого места не очевидна из трассировки, полученной выше.
Краткое примечание для читателей/разработчиков, которые (несмотря на все наши лекции) по-прежнему негативно относятся к использованию PyTorch Profiler: метод, основанный на кэшировании данных, который мы обсудим ниже, предоставит альтернативный способ определения нехватки ресурсов графического процессора — так что не отчаивайтесь.
Нехватка ресурсов графического процессора — поиск первопричины
В этом разделе мы кратко рассмотрим распространенные причины возникновения узких мест в производительности конвейера входных данных.
Напомним, что типичный поток выполнения модели:
- Необработанные данные загружаются или передаются в потоковом режиме из хранилища (например, локальной оперативной памяти или диска, удаленной сетевой файловой системы или облачного хранилища объектов, такого как Amazon S3 или Google Cloud Storage).
- Затем он предварительно обрабатывается на ЦП.
- Наконец, обработанные данные копируются в графический процессор для вывода или обучения.
Соответственно, узкие места могут возникнуть на каждом из следующих этапов:
- Медленное извлечение данных: существует множество факторов, которые могут ограничивать скорость извлечения необработанных данных процессором, в том числе: выбор хранилища (например, облачное хранилище или локальный SSD), доступная пропускная способность сети, формат данных и многое другое.
- Истощение или неправильное использование ресурсов ЦП : задачи предварительной обработки, такие как увеличение данных, преобразование изображений или декомпрессия, могут быть ресурсоемкими для ЦП. Когда количество или сложность этих операций превышает доступную мощность ЦП или если ресурсы ЦП управляются неэффективно (например, неоптимальный выбор количества рабочих процессов), может возникнуть узкое место. Стоит отметить, что ЦП также отвечают за другие обязанности, связанные с моделью, такие как загрузка ядер ГП, управление памятью, отчетность по метрикам и многое другое.
- Узкие места передачи данных с хоста на устройство : после обработки данных их необходимо передать на графический процессор. Это может стать узким местом, если пакеты данных велики относительно пропускной способности памяти CPU-GPU или если копирование памяти выполняется неэффективно (например, копируются отдельные образцы, а не полные пакеты).
Ограничение профилировщиков производительности
Распространенным способом определения узких мест конвейера данных является использование профилировщика производительности. В части 4 этой серии, Решение узких мест на конвейере ввода данных с помощью профилировщика PyTorch и TensorBoard, мы продемонстрировали, как это сделать с помощью встроенного профилировщика PyTorch. Однако, учитывая, что конвейер входных данных работает на ЦП, можно использовать любой профилировщик Python.
Проблема этого подхода в том, что мы обычно используем несколько рабочих процессов для загрузки данных, что делает профилирование производительности особенно сложным. В нашем предыдущем посте мы преодолели это, запустив загрузку данных и выполнение модели в одном процессе (т. е. мы установили аргумент num_workers конструктора DataLoader равным нулю). Однако это очень навязчивое изменение конфигурации, которое может оказать существенное влияние на общую производительность нашей модели.
Метод на основе кэширования, который мы представляем в этом посте, направлен на выявление источника узкого места производительности гораздо менее навязчивым образом. В частности, он позволит нам измерить производительность модели, не изменяя поведение загрузки данных многопользовательской системой.
Обнаружение узких мест с помощью кэширования
В этом разделе мы предлагаем многошаговый подход для анализа производительности конвейера входных данных. Мы покажем, как этот метод можно применить к нашей игрушечной учебной рабочей нагрузке для выявления причин голодания GPU.
Шаг 1: Кэширование пакета на устройстве
Мы начинаем с создания одного входного пакета, копируем его в GPU, а затем измеряем производительность выполнения модели при итерации только этого пакета. Это обеспечивает теоретическую верхнюю границу пропускной способности модели — т. е. максимальную пропускную способность, достижимую, когда GPU не испытывает нехватки данных.
В следующем блоке кода мы изменяем цикл обучения нашего игрушечного скрипта так, чтобы он выполнялся на одном пакете, который кэшируется на графическом процессоре:
данные = next(iter(train_loader)) входы = data[0].to(device=device, non_blocking=True) метки = data[1].to(device=device, non_blocking=True) t0 = time() раз = [] для шага в диапазоне (100): train_step(модель, критерий, оптимизатор, входы, метки) раз.append(time()-t0) t0 = time()
Результирующая средняя пропускная способность составляет 3,45 шага в секунду — почти в 4 раза выше нашего базового результата. Это не только подтверждает существенное узкое место конвейера данных, но и количественно определяет его влияние.
Дополнительный совет: профилирование и оптимизация с использованием кэшированных данных устройства
Запуск профилировщика на одном пакете, кэшированном на GPU, изолирует выполнение модели от входного конвейера. Это помогает вам выявить неэффективность в необработанном вычислительном пути модели. В идеале загрузка GPU здесь должна приближаться к 100%. В нашем случае загрузка составляет около 95%, что приемлемо.
Шаг 2: Кэширование пакета на хосте (ЦП)
Далее мы кэшируем один входной пакет на хосте (CPU) вместо устройства. Теперь каждый шаг включает как копирование памяти из CPU в GPU, так и выполнение модели.
Поскольку закрепление памяти PyTorch допускает асинхронную передачу данных, мы ожидаем, что копирование памяти хоста на устройство для пакета N+1 будет перекрываться с выполнением модели на пакете N. Следовательно, мы ожидаем, что пропускная способность будет примерно такой же, как и в случае кэширования на устройстве. Если нет, это будет явным признаком узкого места в копировании памяти хоста на устройство.
Следующий блок кода содержит применение этого шага к нашей игрушечной модели:
данные = next(iter(train_loader)) t0 = time() times = [] для шага в диапазоне (100): входные данные = data[0].to(device=device, non_blocking=True) метки = data[1].to(device=device, non_blocking=True) train_step(модель, критерий, оптимизатор, входные данные, метки) times.append(time()-t0) t0 = time()
Результирующая пропускная способность после этого изменения составляет 3,33 шага в секунду — незначительное падение по сравнению с предыдущим результатом — указывая на то, что передача с хоста на устройство не является узким местом. Нам нужно продолжать искать источник нашего узкого места производительности.
Шаги 3 и далее: кэширование на промежуточных этапах конвейера данных
Мы продолжаем наш поиск, «поднимаясь» вверх по конвейеру ввода данных, кэшируя в различных промежуточных точках, чтобы точно определить узкое место. Точное применение этого процесса будет зависеть от деталей конвейера. Предположим, что конвейер можно разбить на K этапов. Если кэширование после этапа N дает значительно худшую пропускную способность при кэшировании после этапа N+1, мы можем сделать вывод, что включение обработки этапа N+1 — это то, что нас замедляет.
Шаг 3а: Кэширование одного обработанного образца
В блоке кода ниже мы изменяем наш набор данных для кэширования одного полностью обработанного образца. Это имитирует конвейер, включающий сопоставление данных и копирование данных из ЦП в ГП.
class FakeDataset(VisionDataset): def __init__(self, transform): super().__init__(root=None, transform=transform) self.size = 10000 self.cache = None def __getitem__(self, index): if self.cache is None: # создать случайное изображение 1024×1024 img = Image.fromarray(np.random.randint( low=0, high=256, size=(input_img_size[0], input_img_size[1], 3), dtype=np.uint8 )) # создать случайную метку target = np.random.randint( low=0, high=num_classes, dtype=np.uint8).item() # применить преобразования img = self.transform(img) self.cache = img, target return self.cache
Результирующая пропускная способность составляет 3,23 шага в секунду — все еще намного выше нашего базового значения 0,89. Мы все еще не нашли виновника.
Шаг 3б: Кэширование необработанных данных (до преобразования)
Далее мы изменяем набор данных, чтобы кэшировать необработанные данные (например, необработанные файлы изображений). Конвейер входных данных теперь включает преобразования данных, сопоставление данных и копирование данных с CPU на GPU.
class FakeDataset(VisionDataset): def __init__(self, transform): super().__init__(root=None, transform=transform) self.size = 10000 self.cache = None def __getitem__(self, index): if self.cache is None: # создаем случайное изображение 1024×1024 img = Image.fromarray(np.random.randint( low=0, high=256, size=(input_img_size[0], input_img_size[1], 3), dtype=np.uint8 )) # создаем случайную метку target = np.random.randint(low=0, high=num_classes, dtype=np.uint8).item() self.cache = img, target # Применяем преобразования img = self.transform(self.cache[0]) return img, self.cache[1]
На этот раз пропускная способность резко падает — вплоть до 1,72 шагов в секунду. Мы нашли нашего первого виновника: функцию преобразования данных.
Промежуточные результаты
Вот краткое изложение проведенных на данный момент экспериментов:

Результаты указывают на значительное замедление, вызванное этапом преобразования данных. Разрыв между результатом кэширования необработанных данных и базовым значением также предполагает, что загрузка необработанных данных может быть еще одним виновником. Начнем с узкого места обработки данных.
Оптимизация преобразования данных
Теперь мы приступим к нашему новому открытию узкого места производительности в функции обработки данных. Следующим логическим шагом было бы разбить функцию преобразования на отдельные компоненты и применить нашу технику кэширования к каждому из них, чтобы получить больше информации о точных источниках нашего голодания GPU. Для краткости мы пропустим вперед и применим оптимизации обработки данных, обсуждавшиеся в нашей предыдущей статье, Решение узких мест в конвейере ввода данных с помощью PyTorch Profiler и TensorBoard. Подробности см. там.
После оптимизации преобразования данных пропускная способность эксперимента с кэшированными необработанными данными взлетает до 3,23. Мы устранили узкое место в функции обработки данных.
Однако наша новая базовая пропускная способность (без кэширования) становится 1,28 шагов в секунду, что указывает на то, что в загрузке сырых данных остается узкое место. Это похоже на конечный результат, которого мы достигли в нашем предыдущем посте.

Оптимизация загрузки необработанных данных
Чтобы устранить оставшееся узкое место, мы имитируем оптимизацию, продемонстрированную в части 5 этой серии, Как оптимизировать ваш DL Data-Input Pipeline с помощью пользовательского оператора PyTorch. Мы делаем это, уменьшая размер нашего начального случайного изображения с 1024×1024 до 256×256. После этого изменения сквозной (некэшированный) шаг обучения увеличивается до 3,23. Проблема решена!
Важные предостережения
В заключение сделаем несколько важных замечаний и оговорок.
- Падение пропускной способности в результате включения определенного шага обработки данных в конвейер данных не обязательно означает, что именно этот шаг требует оптимизации. Вполне возможно, что это еще один шаг, загрузка ЦП которого близка к пределу, и новый шаг просто перевернул его.
- Если размер входных данных различается, пропускная способность одного кэшированного образца данных или пакета образцов может не отражать реальную производительность.
- То же предостережение применяется, если модель ИИ включает динамические, зависящие от данных функции, например, если компоненты графа модели зависят от входных данных.
Советы, приемы и методы устранения узких мест на конвейере ввода данных
Мы завершаем этот пост списком советов, приемов и методов оптимизации конвейера ввода данных моделей ИИ на основе PyTorch. Этот список ни в коем случае не является исчерпывающим — существуют многочисленные дополнительные оптимизации в зависимости от вашего конкретного варианта использования и инфраструктуры. Мы разделяем оптимизации на три категории:
- Оптимизация ввода/извлечения необработанных данных
- Оптимизация обработки данных
- Оптимизация передачи данных с хоста на устройство
Оптимизация ввода/извлечения необработанных данных
Эффективная загрузка данных начинается с быстрого и надежного доступа к необработанным данным. Следующие советы могут помочь:
- Выберите тип экземпляра с достаточной пропускной способностью сетевого входа.
- Используйте быстрое и экономичное решение для хранения данных. Локальные SSD-накопители быстрые, но дорогие. Облачные решения, такие как S3, предлагают масштабируемость, но могут вызывать задержку.
- Максимизируйте выход сети хранения. Рассмотрите возможность секционирования наборов данных в S3 или настройки параллельных загрузок для снижения дросселирования.
- Рассмотрите возможность сжатия необработанных данных. Сжатие файлов может сократить время передачи, но будьте осторожны с увеличением нагрузки на процессор во время распаковки.
- Группируйте небольшие образцы в более крупные файлы. Это может сократить накладные расходы, связанные с открытием и закрытием множества файлов.
- Используйте оптимизированные инструменты передачи данных. Например, s5cmd может значительно превзойти AWS CLI для массовых загрузок S3.
- Настройте параметры извлечения данных. Настройка размера фрагмента или параметров параллелизма может существенно повлиять на производительность чтения.
Устранение узких мест в обработке данных
- Настройте количество рабочих процессов загрузки данных и коэффициент предварительной выборки.
- По возможности перенесите обработку данных на этап подготовки данных.
- Выберите тип экземпляра с оптимальным соотношением вычислительных мощностей ЦП/ГП.
- Оптимизируйте порядок преобразований. Например, применение кадрирования перед размытием будет быстрее, если размыть полноразмерное изображение и только потом кадрировать.
- Используйте библиотеки ускорения Python. Например, Numba и JAX могут ускорить чистые операции Python с помощью JIT-компиляции.
- При необходимости создавайте пользовательские операторы ЦП PyTorch (например, см. здесь).
- Рассмотрите возможность добавления вспомогательных ЦП (серверов данных) — (например, см. здесь).
- Переместить преобразования, дружественные GPU, в граф GPU. Некоторые преобразования (например, нормализация) могут быть выполнены после загрузки на GPU для лучшего перекрытия.
- Настройте конфигурации потоков и памяти на уровне ОС.
Оптимизация копирования данных с хоста на устройство
- Используйте закрепление памяти и неблокирующие копии данных для предварительной выборки данных непосредственно на GPU. Также см. специальный CudaDataPrefetcher, предлагаемый TorchTNT.
- Отложите преобразование типов данных int8 в float32 на GPU, чтобы уменьшить нагрузку на копирование памяти в 4 раза.
- Если в вашей модели используются числа с плавающей точкой меньшей точности (например, fp16/bfloat16), передайте их на ЦП, чтобы уменьшить полезную нагрузку вдвое.
- Отложите распаковку векторов one-hot на GPU — т.е. сохраните их в качестве идентификаторов меток до самого последнего момента.
- Если у вас много двоичных значений, рассмотрите возможность использования битовых масок для сжатия полезной нагрузки. Например, если у вас 8 двоичных карт, рассмотрите возможность сжатия их в один uint8.
- Если ваши входные данные разрежены, рассмотрите возможность использования разреженных представлений данных.
- Избегайте ненужного дополнения. Хотя дополнение нулями является популярным методом работы с входными образцами переменного размера, оно может значительно увеличить размер копии памяти. Рассмотрите альтернативные варианты (например, см. здесь).
- Убедитесь, что вы не копируете данные, которые вам на самом деле не нужны на GPU!!
Краткое содержание
Хотя графические процессоры считаются необходимыми для современной разработки AI/ML, они стоят дорого. Как только вы решите сделать необходимые инвестиции в их приобретение, вам захочется убедиться, что они используются как можно чаще. Последнее, чего вы хотите, — чтобы ваш графический процессор простаивал, ожидая входных данных из-за предотвратимого узкого места в другом месте конвейера.
К сожалению, такие неэффективности встречаются слишком часто. В этой статье мы представили простую методику диагностики этих проблем путем итеративного кэширования данных на разных этапах входного конвейера. Изолируя влияние времени выполнения каждого компонента конвейера, этот метод помогает выявить конкретные узкие места — будь то загрузка необработанных данных, предварительная обработка или передача с хоста на устройство.
Конечно, точная реализация будет различаться в зависимости от проекта и конвейера, но мы надеемся, что эта стратегия обеспечит полезную основу для диагностики и решения проблем производительности в ваших собственных рабочих процессах ИИ/МО.
Источник: towardsdatascience.com



























