Путешествие KT по сокращению времени обучения модели видовых преобразователей с использованием Amazon SageMaker

Как KT сократил время обучения модели видовых преобразователей с помощью Amazon SageMaker увлекательное путешествие

KT Corporation является одним из крупнейших провайдеров телекоммуникационных услуг в Южной Корее, предлагающим широкий спектр услуг, включая проводные телефоны, мобильную связь и интернет, а также искусственный интеллект. AI Food Tag KT – это решение для управления питанием на основе искусственного интеллекта, которое определяет тип и пищевую ценность продуктов на фотографиях с использованием модели компьютерного зрения. Данная модель разработанная KT, полагается на модель, предварительно обученную на большом количестве неразмеченных данных изображений, для анализа пищевой ценности и информации о калорийности различных продуктов. AI Food Tag может помочь пациентам с хроническими заболеваниями, такими как диабет, контролировать свое питание. KT использовала AWS и Amazon SageMaker для тренировки модели AI Food Tag в 29 раз быстрее, чем раньше, и оптимизации ее для развертывания в производственной среде с использованием техники сжатия модели. В этой статье мы расскажем о пути развития и успехе модели KT, использующей SageMaker.

Представление проекта KT и определение проблемы

Модель AI Food Tag, предварительно обученная KT, основана на архитектуре vision transformers (ViT) и имеет больше параметров модели, чем их предыдущая модель компьютерного зрения для повышения точности. Для уменьшения размера модели для производства KT использует технику сжатия знаний (KD), которая позволяет снизить количество параметров модели без значительного влияния на точность. С использованием сжатия знаний, предварительно обученная модель называется учителем, и небольшая легковесная модель является учеником, как показано на следующей схеме. Легковесная модель имеет меньше параметров модели, чем учитель, что снижает требования к памяти и позволяет развертывать ее на более маленьких, менее дорогих экземплярах. Ученик сохраняет приемлемую точность, несмотря на свой меньший размер, учась на основе выходов учительской модели.

Общий процесс обучения с использованием сжатия знаний

Учительская модель остается неизменной во время KD, но ученик обучается с использованием выходных данных учительской модели в качестве меток для расчета потерь. С такой парадигмой KD, как учителя, так и ученика необходимо размещать на одной GPU для обучения. KT изначально использовала две графических карты (A100 80 ГБ) во внутренней среде для обучения ученика, но процесс занимал около 40 дней для охвата 300 эпох. Чтобы ускорить обучение и получить ученикскую модель за меньшее время, KT сотрудничал с AWS. Вместе команды существенно сократили время обучения модели. В этой статье описывается, как команда использовала Amazon SageMaker Training, библиотеку SageMaker Data Parallelism Library, Amazon SageMaker Debugger и Amazon SageMaker Profiler, чтобы успешно разработать легковесную модель AI Food Tag.

Создание распределенной среды обучения с SageMaker

SageMaker Training – это управляемая среда машинного обучения (ML) на AWS, которая предоставляет набор функций и инструментов для упрощения процесса обучения и может быть полезна в распределенных вычислениях, как показано на следующей диаграмме.

Среда распределенного обучения модели с помощью SageMaker Training

Клиенты SageMaker также могут получить доступ к встроенным образам Docker с различными предустановленными глубокими фреймворками обучения моделей и необходимыми пакетами Linux, NCCL и Python. Ученые-исследователи или инженеры по машинному обучению, которые хотят запустить обучение модели, могут сделать это без необходимости настройки инфраструктуры обучения или управления Docker и совместимостью разных библиотек.

Во время 1-дневного семинара нам удалось настроить конфигурацию распределенного обучения на основе SageMaker в учетной записи AWS у KT, ускорить скрипты обучения KT с помощью библиотеки SageMaker Distributed Data Parallel (DDP) и даже протестировать задачу обучения с использованием двух экземпляров ml.p4d.24xlarge. В этом разделе мы описываем опыт KT в работе с командой AWS и использовании SageMaker для разработки их модели.

В рамках доказательства концепции мы хотели ускорить процесс обучения, используя библиотеку SageMaker DDP, которая оптимизирована для инфраструктуры AWS во время распределенного обучения. Чтобы перейти от PyTorch DDP к SageMaker DDP, вам просто нужно объявить пакет torch_smddp и изменить параметр backend на smddp, как показано в следующем коде:

import smdistributed.dataparallel.torch.torch_smddpdist.init_process_group(backend='smddp',rank=args.rank,world_size=args.world_size)

Чтобы узнать больше о библиотеке SageMaker DDP, обратитесь к библиотеке параллельного обучения SageMaker.

Анализ причин медленной скорости обучения с помощью SageMaker Debugger и Profiler

Первый шаг в оптимизации и ускорении рабочей нагрузки обучения состоит в понимании и выявлении мест узких мест. Для задачи обучения KT мы измерили время обучения на итерацию загрузчика данных, прямого прохода и обратного прохода:

1 время итерации – загрузчик данных: 0.00053 сек, прямой проход: 7.77474 сек, обратный проход: 1.58002 сек
2 время итерации – загрузчик данных: 0.00063 сек, прямой проход: 0.67429 сек, обратный проход: 24.74539 сек
3 время итерации – загрузчик данных: 0.00061 сек, прямой проход: 0.90976 сек, обратный проход: 8.31253 сек
4 время итерации – загрузчик данных: 0.00060 сек, прямой проход: 0.60958 сек, обратный проход: 30.93830 сек
5 время итерации – загрузчик данных: 0.00080 сек, прямой проход: 0.83237 сек, обратный проход: 8.41030 сек
6 время итерации – загрузчик данных: 0.00067 сек, прямой проход: 0.75715 сек, обратный проход: 29.88415 сек

Анализируя время в стандартном выводе для каждой итерации, мы увидели, что время выполнения обратного прохода значительно варьируется от итерации к итерации. Эта вариация необычна и может повлиять на общее время обучения. Чтобы найти причину такой несвязанной скорости обучения, мы сначала попытались выявить узкие места ресурсов, используя Монитор Системы (SageMaker Debugger UI), который позволяет отлаживать задачи обучения на SageMaker Training и просматривать состояние ресурсов, таких как ЦПУ, ГПУ, сеть и ввод-вывод в пределах заданного количества секунд.

Интерфейс пользователя SageMaker Debugger предоставляет подробные и необходимые данные, которые могут помочь выявить и диагностировать узкие места в задаче обучения. В частности, линейная диаграмма использования ЦП и таблицы тепловых карт использования ЦП/ГП на каждом экземпляре привлекли наше внимание.

На линейной диаграмме использования ЦП мы заметили, что некоторые ЦП используются на 100%.

Диаграмма использования ЦП с узким местом в виде ЦП

В таблице тепловой карты (где более темные цвета указывают на более высокое использование) мы отметили, что некоторые ядра ЦП имели высокое использование на протяжении всего обучения, в то время как использование ГПУ не было постоянно высоким со временем.

Тепловая карта использования ЦП с узким местом процессора

Отсюда мы начали подозревать, что одной из причин медленной скорости обучения является узкое место в процессоре. Мы просмотрели код скрипта обучения, чтобы увидеть, не вызывает ли что-нибудь узкое место в процессоре. Самой подозрительной частью было большое значение num_workers в загрузчике данных, поэтому мы изменили это значение на 0 или 1 для уменьшения использования процессора. Затем мы снова запустили рабочую нагрузку обучения и проверили результаты.

Ниже приведены скриншоты графика использования CPU, использования GPU и тепловой карты после устранения узкого места в процессоре.

График использования ЦП после устранения узкого места в процессоре

График использования ЦП и график использования ГП после устранения узкого места в процессореТепловая карта использования ЦП после устранения узкого места в процессоре

Просто изменив значение num_workers, мы заметили значительное уменьшение использования ЦП и общий рост использования ГП. Это было важное изменение, которое существенно улучшило скорость обучения. Однако мы хотели узнать, где еще можно оптимизировать использование ГП. Для этого мы использовали SageMaker Profiler.

SageMaker Profiler помогает выявить подсказки для оптимизации, предоставляя видимость использования операций, включая отслеживание метрик использования ЦП и ГП и потребления ядер ЦП/ГП в скриптах обучения. Он помогает пользователям понять, какие операции потребляют ресурсы. Сначала, чтобы использовать SageMaker Profiler, необходимо добавить ProfilerConfig в функцию, которая вызывает рабочую нагрузку обучения с использованием SageMaker SDK, как показано в следующем коде:

from sagemaker import ProfilerConfig, Profilerfrom sagemaker.debugger import (ProfilerRule, rule_configs)rules=[ProfilerRule.sagemaker(rule_configs.ProfilerReport())]profiler_config = ProfilerConfig(profile_params = Profiler(cpu_profiling_duration=3600))from sagemaker.pytorch import PyTorchregion_name = 'us-west-2'image_uri=f'763104351884.dkr.ecr.{region_name}.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker'estimator = PyTorch(entry_point='train.py',source_dir='src',role=role,image_uri=image_uri,instance_count=4,instance_type='ml.p4d.24xlarge',distribution={'smdistributed': {'dataparallel': {'enabled': True}}},profiler_config=profiler_config,hyperparameters=hyperparameters,sagemaker_session=sagemaker_session,)

В SageMaker Python SDK у вас есть гибкость добавлять функции annotate для SageMaker Profiler, чтобы выбрать код или шаги в скрипте обучения, которые требуют профилирования. Вот пример кода, который необходимо объявить для SageMaker Profiler в скриптах обучения:

import smppySMProf = smppy.SMProfiler.instance()config = smppy.Config()config.profiler = {"EnableCuda": "1",}SMProf.configure(config)SMProf.start_profiling()…with smppy.annotate("Forward"):student_out = student_model(inp)with smppy.annotate("Backward"):loss.backward()…SMProf.stop_profiling()

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

Потраченное время всеми ядрами GPU (1)

Когда мы проверили топ-5 потребления времени операций ядра GPU среди результатов SageMaker Profiler, мы обнаружили, что для сценария обучения KT больше всего времени тратится на операцию умножения матриц, которая является общей операцией умножения матриц (GEMM) на GPU. Имея этот важный анализ от SageMaker Profiler, мы начали исследование способов ускорения этих операций и повышения использования GPU.

Ускорение времени обучения

Мы рассмотрели различные способы сокращения времени вычисления умножения матриц и применили две функции PyTorch.

Фрагментация состояния оптимизатора с использованием ZeroRedundancyOptimizer

Если вы посмотрите на Zero Redundancy Optimizer (ZeRO), техника DeepSpeed/ZeRO позволяет эффективно тренировать большую модель с лучшей скоростью тренировки за счет устранения избыточности в памяти, используемой моделью. ZeroRedundancyOptimizer в PyTorch использует технику фрагментации состояния оптимизатора для уменьшения использования памяти для каждого процесса в распределенном параллельном обучении (DDP). DDP использует синхронизированные градиенты на обратном проходе, чтобы все реплики оптимизатора итерировались по одним и тем же параметрам и значениям градиента, но вместо того, чтобы иметь все параметры модели, каждое состояние оптимизатора поддерживается фрагментацией только для разных процессов DDP для уменьшения использования памяти.

Чтобы использовать это, вы можете оставить существующий оптимизатор в optimizer_class и объявить ZeroRedundancyOptimizer с остальными параметрами модели и скоростью обучения в качестве параметров.

student_optimizer = ZeroRedundancyOptimizer(student_model.parameters(),optimizer_class=torch.optim.AdamW,lr=initial_lr)

Автоматическая смешанная точность

Автоматическая смешанная точность (AMP) использует тип данных torch.float32 для некоторых операций и torch.bfloat16 или torch.float16 для других, для удобства быстрого вычисления и снижения использования памяти. В особенности, поскольку модели глубокого обучения обычно более чувствительны к показателям экспоненты, чем к показателям дробной части в их вычислениях, torch.bfloat16 эквивалентен показателям экспоненты torch.float32, позволяя им быстро обучаться с минимальной потерей. torch.bfloat16 работает только на экземплярах с архитектурой NVIDIA A100 (Ampere) или выше, таких как ml.p4d.24xlarge, ml.p4de.24xlarge и ml.p5.48xlarge.

Чтобы применить AMP, вы можете объявить torch.cuda.amp.autocast в скриптах обучения, как показано в приведенном выше коде, и объявить dtype как torch.bfloat16.

with torch.cuda.amp.autocast(dtype="torch.bfloat16"):teacher = teacher_model(input_data)student = student_model(input_data)loss = loss(teacher, student, target)loss.requires_grad_(True)loss.backward()student_optimizer.step()student_optimizer.zero_grad(set_to_none=True)

Результаты в SageMaker Profiler

После применения этих двух функций к скриптам обучения и проведения тренировочного задания для одной эпохи снова, мы проверили топ-5 потребления времени операций ядра GPU в SageMaker Profiler. На следующем рисунке показаны наши результаты.

Потраченное время всеми ядрами GPU (2)

Мы видим, что операция GEMM, которая была первой в списке до применения двух функций Torch, исчезла из топ-5 операций, заменившись операцией ReduceScatter, которая обычно возникает в распределенном обучении.

Результаты скорости обучения модели KT distilled

Мы увеличили размер пакета обучения на 128, чтобы учесть экономию памяти от применения двух функций Torch, в результате чего итоговый размер пакета составил 1152 вместо 1024. Обучение финальной модели студента проходило со скоростью 210 эпох в день; время обучения и ускорение между внутренней средой обучения KT и SageMaker подводятся итогом в следующей таблице.

Среда обучения Видеокарта для обучения Количество видеокарт Время обучения (в часах) Эпоха Часов на эпоху Коэффициент снижения
Внутренняя среда обучения KT A100 (80GB) 2 960 300 3.20 29
Amazon SageMaker A100 (40GB) 32 24 210 0.11 1

Масштабируемость AWS позволила нам завершить задачу обучения в 29 раз быстрее, используя 32 видеокарты вместо 2 в офисе. В результате использования большего количества видеокарт в SageMaker обучение занимало гораздо меньше времени без изменения общих затрат на обучение.

Заключение

Пак Сан-мин (руководитель команды технологии обслуживания Vision AI) из лаборатории AI2XL в центре сходимости технологий KT прокомментировал сотрудничество с AWS по разработке модели AI Food Tag:

“В последнее время, поскольку в области компьютерного зрения стало больше моделей на основе трансформера, увеличиваются параметры моделей и требуемая память видеокарты. Мы используем легковесную технологию для решения этой проблемы, и для обучения требуется много времени, приблизительно месяц. Через это PoC с AWS нам удалось определить проблемы с ресурсами с помощью SageMaker Profiler и Debugger, устранить их, а затем с использованием библиотеки параллелизма данных SageMaker завершить обучение примерно за один день с оптимизированным кодом модели на четырех ml.p4d.24xlarge инстансах.”

SageMaker помогла команде Сан-мина сэкономить недели времени на обучение и разработку модели.

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

Чтобы узнать больше о связанных функциях SageMaker, ознакомьтесь с следующими ресурсами: