[Анализ исходного кода] PipeDream (1) --- Этап профиля параллельного конвейера глубокого обучения

машинное обучение глубокое обучение

0x00 сводка

Вслед за GPipe мы открыли новую серию параллельных тренингов по конвейеру, представив PipeDream от Microsoft. Эта статья знакомит с его общей идеей, архитектурой и стадией профиля.

Конвейер Gpipe имеет две проблемы: низкое использование оборудования и большое использование памяти. Итак, в другой параллельной статье Microsoft PipeDream предложила улучшенный метод для этих проблем, который представляет собой стратегию 1F1B (один прямой проход, за которым следует один обратный проход). Эта улучшенная стратегия может решить проблему количества активаций кеша, так что количество кешей активации связано только с количеством этапов, тем самым дополнительно экономя видеопамять и обучая более крупные модели.

PipeDream можно разделить на 4 этапа: профиль, вычислительный раздел, преобразование модели, время выполнения и другие четыре этапа.

  • Этап профиля: время обучения DNN выводится из профиля мини-пакетных данных.
  • Этап Compute Partition: определите время работы всех слоев в соответствии с результатами профиля, а затем выполните оптимизацию.Оптимизатор возвращает аннотированный граф операторов, и каждый слой модели сопоставляется с идентификатором этапа.
  • Стадия преобразования модели. Выполните BFS-обход графа операторов, создав отдельный код torch.nn.Module для каждой стадии. PipeDream приказывает операторам на каждом этапе убедиться, что они остаются совместимыми с зависимостями ввода-вывода исходного графа модели PyTorch.
  • Стадии выполнения: среда выполнения PipeDream назначает каждый этап (включая реплики этапа репликации) одному рабочему процессу в соответствии со своей политикой планирования 1F1B-RR.

В этой статье сначала рассматриваются общая идея, архитектура и этап профиля PipeDream.

0x01 Обзор

1.1 Предыдущий обзор

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

  • Параллельные конвейеры, особенно как автоматически устанавливать конвейеры;
  • накопление градиента;
  • обратный пересчет;
  • стратегия 1F1B;

В предыдущих статьях мы представили, как Gpipe работает с первыми тремя методами. В начале этой статьи мы рассмотрим, как он реализует конвейерный параллелизм и стратегию 1F1B с помощью PipeDream, распределенной системы обучения DNN от Microsoft.

1.2 Текущие проблемы

Обучение DNN характеризуется двунаправленным обучением. Обучение вычисляется итеративно в прямом и обратном каналах. Два распространения проходят через одни и те же слои в обратном порядке. На каждой итерации процесс обучения проходит через мини-пакет входных данных и обновляет параметры модели. .

1.2.1 Параллелизм данных

Наиболее распространенным способом распараллеливания обучения DNN является распараллеливание данных, при котором входные данные распределяются между рабочими процессами (каждый рабочий процесс владеет полной моделью) для запуска. К сожалению, несмотря на некоторый прогресс в оптимизации производительности для ускорения параллелизма данных, обучение работе с облачной инфраструктурой сопряжено со значительными затратами на связь. Более того, с быстрым увеличением скорости вычислений на графическом процессоре трудоемкое узкое место обучения всех моделей в дальнейшем обратится к каналу связи.

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

1.2.2 Параллелизм модели

Распараллеливание моделей — это еще одна форма распараллеленного обучения, при которой операторы распределяются между рабочими для вычислений, что часто используется для обучения больших моделей DNN. В этой статье под параллелизмом модели понимается размещение разных слоев модели на разных машинах, а не разбиение одного и того же слоя на разные машины.

На рисунке ниже показана распараллеливание модели, показывающая временную шкалу вычислений, в примере четыре машины и один конвейер.

  • На прямом этапе каждый этап выполняет прямой проход мини-пакета слоев на этом этапе и отправляет результат на следующий этап. Выходной каскад вычисляет потери минипакета после завершения прямого прохода.
  • На обратном этапе каждый этап формирует обратный канал для передачи потерь на предыдущий этап один за другим.

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

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

1.2.3 Gpipe

В дополнение к общей проблеме, параллельная стратегия конвейера GPipe также имеет проблему с памятью: необходимо кэшировать несколько активаций.

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

0x02 Документы

Для этих задач PipeDream предлагает улучшенный метод 1F1B. PipeDream — первая система, объединяющая конвейерный параллелизм, параллелизм моделей и параллелизм данных автоматизированным и универсальным способом. Сначала PipeDream разделяет DNN, используя модельный параллелизм, и назначает подмножество каждого уровня каждому рабочему процессу. Но в отличие от традиционного параллелизма моделей, PipeDream передает в конвейер небольшие пакеты данных, что позволяет создать потенциальную схему параллелизма в конвейере. В любой момент разные воркеры обрабатывают разные входы, что обеспечивает полную загрузку пайплайна и параллельного BSP.

Майкрософт в газетеPipeDream: Fast and Efficient Pipeline Parallel DNN TrainingPipeDream подробно объясняется, поэтому мы проанализируем его на основе этой статьи.

2.1 Обзор программы

2.1.1 Параллельный режим

Базовой единицей модели PipeDream является слой, и PipeDream делит эти слои DNN на несколько этапов. Каждый этап состоит из набора последовательных слоев модели.

Основной параллельный метод PipeDream заключается в размещении разных слоев модели на разных этапах, при этом разные этапы развертываются на разных машинах, а прямые и обратные вычисления выполняются последовательно для формирования конвейера.

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

Для использованияпараллелизм данныхНа этапе используется метод round-robin для назначения задач каждому устройству, это необходимо для того, чтобы пакет данных происходил на одной и той же машине в прямом и обратном направлениях.

2.1.2 1F1B

Так как активацию прямого расчета нужно выпустить после завершения соответствующего обратного расчета (независимо от того, используется технология Checkpointing или нет), при параллельном конвейере, если вы хотите максимально сохранить количество кэшированных активаций, вы должны попытаться сократить каждую активацию.Экономия времени заключается в том, чтобы выпустить каждую активацию как можно раньше, поэтому данные каждой микропартии должны быть рассчитаны как можно раньше.Назад с меньшими метками пакетов выполняются в первую очередь, чем вперед с более крупными микропакетами -пакетные этикетки. Следовательно, если мы позволим последнему этапу выполнять обратный расчет этого микропакета сразу после завершения прямого и обратного микропакета, то мы можем позволить другим этапам начать обратный расчет как можно раньше, что является стратегией 1F1B.

Режим планирования 1F1B (один вперед-один-назад) будет попеременно выполнять прямое и обратное вычисление небольших пакетов данных на каждой рабочей машине, при этом гарантируя, что эти небольшие пакеты могут быть направлены на «вперед» во время «обратного распространения». .к тому же работнику, который размножается».

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

0x03 Конвейер

Конвейерный параллелизм PipeDream (PP) — это новая стратегия распараллеливания, сочетающая внутрипакетный параллелизм с межпакетным параллелизмом.

3.1 Улучшения конвейера

Давайте сначала посмотрим на улучшения, которые конвейеры вносят в модель параллелизма.

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

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

По сравнению с обычным параллельным обучением между слоями конвейеры имеют два основных преимущества:

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

На рисунке ниже показана реализация конвейера 1F1B. Машина 1 сначала вычисляет синий 1, затем отправляет синий 1 на машину 2, чтобы продолжить расчет. Затем машина 1 вычисляет синий 2. Подмножество только модели между Machine 1 и Machine 2. Расчет и связь могут быть параллельными.

3.2 Проблемы

Цель PipeDream — объединить конвейерный параллелизм, параллелизм моделей и параллелизм данных таким образом, чтобы минимизировать общее время обучения. Однако, чтобы этот подход был эффективным для больших моделей DNN и чтобы воспользоваться потенциальными преимуществами конвейерного распараллеливания обучения, PipeDream должен преодолеть несколько серьезных проблем:

  • Как эффективно разделить трубопровод. Подобно конвейерам в процессорах, DNN необходимо эффективно и правильно разделить на несколько «этапов» (последовательностей слоев), при этом каждый этап развертывается для выполнения на другом рабочем потоке.

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

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

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

    • Одна проблема с конвейером заключается в том, что существует много версий веса. Если при обратном распространении используется более высокая версия веса, чем при прямом распространении, качество обучающей модели будет снижено.
    • PipeDream управляет версией веса в обратном канале и решает эту проблему, поддерживая номер версии для каждой небольшой партии гирь, так что версия веса, используемая в обратном канале, совпадает с версией, используемой в прямом канале, и, таким образом, численно Уметь правильно рассчитывать градиент (мы объясним это в следующей статье).

3.4 Алгоритм разделения конвейера

PipeDream автоматически разделяет слои DNN на основе результатов краткосрочного анализа и использует алгоритмы для балансировки вычислительной нагрузки между различными этапами на основе результатов анализа при минимизации связи. Общая цель алгоритма автоматического разделения PipeDream — создать сбалансированный конвейер, гарантирующий, что каждый этап выполняет примерно одинаковый общий объем работы. Также необходимо обеспечить, чтобы объем данных, передаваемых между этапами, был как можно меньше, чтобы избежать прерывания связи. Алгоритм следующий:

  • Разделите слой DNN на этапы так, чтобы каждый этап выполнялся примерно с одинаковой скоростью, т. е. занимал примерно одинаковое количество времени вычислений.
  • Постарайтесь свести к минимуму взаимодействие между рабочими процессами с учетом топологии (например, по возможности отправляйте большие выходные данные на каналы с более высокой пропускной способностью).
  • Поскольку DNN не всегда могут быть равномерно распределены между доступными рабочими процессами, для дальнейшего улучшения балансировки нагрузки PipeDream позволяет дублировать этап, т. е. использовать несколько рабочих процессов на этом этапе для параллелизма данных. Таким образом, несколько рабочих могут быть назначены на один и тот же этап конвейера для параллельной обработки различных мини-партий партии, что повышает эффективность обработки. Поскольку параллелизм данных использует RR, эта стратегия также называется 1F1B-RR (один-вперед-ни-назад-циклический перебор).

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

Подробности следующие:

3.5 Profile

Одной из характеристик обучения DNN является то, что время вычислений для разных входных данных практически не меняется. Таким образом, PipeDream в полной мере использует этот факт, учитывая DNN с N слоями и M доступными машинами, PipeDream сначала анализирует модель на одной машине, записывая время вычислений, затраченное на прямой и обратный процесс, размер выходного слоя и размер соответствующих параметров каждого слоя и, наконец, вывод в виде файла результатов.

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

Итак, общий алгоритм примерно такой:

Поскольку PipeDream заимствует много идей у ​​GPipe, вы можете видеть его прогресс по сравнению с Gpipe.

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

0x04 Стадия профиля

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

Это усовершенствование PipeDream для GPipe, оба из которых оценивают время работы каждого слоя, а затем разделяют модель.

  • GPipe использует эмпирические или математические методы для оценки времени работы.
  • PipeDream оценивает время выполнения на основе результатов профилирования.

PipeDream является более точным и продвинутым, поскольку поддерживается фактическими данными.

4.1 Идеи

Механизм оценки использует преимущество того факта, что обучение DNN имеет небольшие различия во времени вычислений и связи. Таким образом, мы можем вывести время обучения DNN по профилю мини-пакетных данных. Чтобы определить время выполнения всех слоев, PipeDream профилирует краткосрочные (минуты) запуски модели DNN, используя 1000 мини-пакетов на одной из машин.

4.1.1 Как рассчитать

часы работы

Для времени работы каждого слоя мы можем получить его как время работы = время вычисления + время связи.

  • Время расчета — это время расчета прямого и обратного направлений каждого слоя, которое можно получить из профиля.
  • Время связи необходимо оценивать в соответствии с размером модели.PipeDream оценивает время, необходимое для связи, как «количество передаваемых данных», деленное на «пропускную способность канала связи».

время связи

На конвейере большинство коммуникаций состоит из трех этапов:

1) На отправляющей машине переместите данные с GPU на CPU.

2) Отправка данных от отправителя к получателю по сети.

3) На принимающей стороне данные перемещаются из ЦП в ГП.

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

  • PipeDream оценивает на основе «значения активации» время передачи значения активации с уровня i на уровень i + 1.

  • Если сконфигурирован параллелизм данных (для слоя i для параллелизма данных используется m воркеров), время на синхронизацию весов оценивается с помощью «веса»:

    • При использовании распределенного сервера параметров количество весов оценивается как 4 x ( m - 1 ) x | w i | / m.
    • Если используется all_reduce, каждый рабочий процесс отправляет ( m - 1 ) x | w i | / m байтов другим рабочим процессам и получает такое же количество байтов.

4.1.2 Содержание профиля

Таким образом, PipeDream записывает три величины для каждого слоя i в профиле:

  • Ti — сумма времени прямого и обратного вычисления слоя i на GPU, то есть время прямого и обратного вычисления каждого слоя;
  • ai размер выходной активации слоя i (и размер входного градиента в обратном процессе) в байтах, то есть размер выхода каждого слоя;
  • wi, размер параметра веса слоя i (в байтах), то есть размер параметра слоя каждого слоя;

4.2 Код

Разные модели или разные поля имеют разные профили.

Мы используем profiler/translation/train.py в качестве записи для анализа.

4.2.1 Сценарий обучения

Ниже мы опускаем посторонний код.

4.2.1.1 Процесс обучения
class Seq2SeqTrainer:
​
    def feed_data(self, data_loader, training=True):
        """
        Runs training or validation on batches from data_loader.
​
        :param data_loader: data loader
        :param training: if True runs training else runs validation
        """
        # 白名单
        module_whitelist = ["EmuBidirLSTM", "RecurrentAttention", "Classifier"]
        
        # 样本集
        for i, (src, tgt) in enumerate(data_loader):
            break
        (src, src_length) = src
        (tgt, tgt_length) = tgt
        src_length = torch.LongTensor(src_length).cuda()
        src = src.cuda()
        tgt = tgt.cuda()
        model_input = (src, src_length, tgt[:-1])
        
        # 使用torchsummary计算网络的计算参数等信息
        summary = torchsummary.summary(model=self.model, module_whitelist=module_whitelist,
                                       model_input=model_input, verbose=True)
​
         for i, (src, tgt) in enumerate(data_loader):
​
            if training and i in eval_iters:
                test_bleu, _ = self.translator.run(calc_bleu=True,
                                                   epoch=self.epoch,
                                                   iteration=i)
                # 训练模型
                self.model.train()
                self.preallocate(data_loader, training=True)
                        
        # 从模型建立图      
        if training:
            create_graph(self.model, module_whitelist, (src, tgt), summary,
                         os.path.join("profiles", self.arch))  
4.2.1.2 Расчетные параметры

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

import torch
import torch.nn as nn
from torchsummary import summary
​
class SimpleConv(nn.Module):
    def __init__(self):
        super(SimpleConv, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )
​
    def forward(self, x, y):
        x1 = self.features(x)
        x2 = self.features(y)
        return x1, x2
    
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleConv().to(device)
​
summary(model, [(1, 16, 16), (1, 28, 28)])

Он печатается следующим образом:

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [-1, 1, 16, 16]              10
              ReLU-2            [-1, 1, 16, 16]               0
            Conv2d-3            [-1, 1, 28, 28]              10
              ReLU-4            [-1, 1, 28, 28]               0
================================================================
Total params: 20
Trainable params: 20
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.77
Forward/backward pass size (MB): 0.02
Params size (MB): 0.00
Estimated Total Size (MB): 0.78
----------------------------------------------------------------
4.2.1.3 Создание графика

Роль create_graph заключается в использовании torchgraph.GraphCreator для создания графа, который можно понимать как граф DAG внутри модели, и каждый узел записывает следующую информацию.

node10 -- Dropout(p=0.2) -- forward_compute_time=0.064, backward_compute_time=0.128, activation_size=6291456.0, parameter_size=0.000

Конкретный код выглядит следующим образом:

def create_graph(model, module_whitelist, model_input, summary, directory):
    """Given a model, creates and visualizes the computation DAG
       of the model in the passed-in directory."""
    # 创建图
    graph_creator = torchgraph.GraphCreator(model, summary, module_whitelist
    # 构建hook                                        
    graph_creator.hook_modules(model, root=True) 
    (src, tgt) = model_input
    (src, src_length) = src
    (tgt, tgt_length) = tgt
    src_length = torch.LongTensor(src_length).cuda()
    src = src.cuda()
    tgt = tgt.cuda()
    # 运行以得到profile                                        
    model(src, src_length, tgt[:-1])
    graph_creator.unhook_modules()
    # 输出profile结果                                        
    graph_creator.persist_graph(directory)

4.2.2 Создание графика

Создание графиков в основном выполняется в GraphCreator.

class GraphCreator(object):
    def __init__(self, model, summary, module_whitelist):
        if isinstance(model, torch.nn.Module) is False:
            raise Exception("Not a valid model, please provide a 'nn.Module' instance.")
​
        self.model = model
        self.module_whitelist = module_whitelist
        self.summary = copy.deepcopy(summary)
        self.forward_original_methods = {}
        self.graph = graph.Graph()
        self.inputs = {}
4.2.2.1 Установка оболочки

Роль hook_modules состоит в том, чтобы установить оболочку для прямой функции модели и установить обход для подмодуля, чтобы можно было отслеживать соединение между моделями во время работы модели.

def hook_modules(self, module, root=False):
    this_creator = self
    sub_modules = module.__dict__['_modules']
​
    # Wrapper function to "forward()", keeping track of dependencies.
    def forward_wrapper(self, *wrapped_inputs):
        input = []
        wrapped_inputs_list = list(wrapped_inputs)
        for i in range(len(wrapped_inputs_list)): # 遍历输入
            if isinstance(wrapped_inputs_list[i], TensorWrapper):
                # 如果已经被包装,则插入input
                input.append(wrapped_inputs_list[i].tensor)
            else:
                key = wrapped_inputs_list[i]
                if key in this_creator.inputs: # 如果是原始输入,则不进行包装
                    wrapped_inputs_list[i] = this_creator.inputs[key]
                else:
                    j = len(this_creator.inputs)
                    # 如果没有被wrap, 则构建一个TensorWrapper进行包装
                    wrapped_inputs_list[i] = TensorWrapper(wrapped_inputs_list[i],
                                                           "Input%d" % j, this_creator)
                    this_creator.inputs[key] = wrapped_inputs_list[i]
                input.append(wrapped_inputs_list[i].tensor) # 则插入input
        result = this_creator.forward_original_methods[self](*input)
        # 对结果进行包装
        wrapped_result = TensorWrapper(result, str(self), this_creator)
        
        # 把边添加进入图
        for wrapped_input in wrapped_inputs_list:
            this_creator.graph.add_edge(wrapped_input.node(), wrapped_result.node())
​
        return wrapped_result
​
    # Wrapper function to "forward()", keeping track of dependencies.
    def forward_wrapper_root(self, *wrapped_inputs):
        input = []
        wrapped_inputs_list = list(wrapped_inputs)
        for i in range(len(wrapped_inputs_list)):
            if isinstance(wrapped_inputs_list[i], TensorWrapper):
                input.append(wrapped_inputs_list[i].tensor)
            else:
                key = wrapped_inputs_list[i]
                if key in this_creator.inputs:
                    wrapped_inputs_list[i] = this_creator.inputs[key]
                else:
                    j = len(this_creator.inputs)
                    wrapped_inputs_list[i] = TensorWrapper(wrapped_inputs_list[i],
                                                           "Input%d" % j, this_creator)
                    this_creator.inputs[key] = wrapped_inputs_list[i]
                input.append(wrapped_inputs_list[i].tensor)
        result = this_creator.forward_original_methods[self](*input)
​
        return result
​
    # 遍历子模块,递归设置wrapper  
    for name, sub_module in sub_modules.items():
        # nn.Module is the only thing we care about.
        if sub_module is None or isinstance(sub_module, torch.nn.Module) is False:
            break
​
        sub_module_name = sub_module.__class__.__name__
        sub_sub_modules = sub_module.__dict__['_modules']
        if len(sub_sub_modules) == 0 or sub_module_name in self.module_whitelist:
            sub_module.reset_hooks()
            #
            # Hook nn.Module with no descendants.
            #
​
            # Replace "forward" with "wrapped_forward".
            # 使用wrapped_forward替换forward
            if sub_module not in this_creator.forward_original_methods:
                this_creator.forward_original_methods.update({sub_module:
                                                               sub_module.forward})
                sub_module.forward = forward_wrapper.__get__(sub_module, sub_module.__class__)
​
        if len(sub_sub_modules) >forward_compute_time 0 and sub_module_name not in self.module_whitelist:
            #
            # Recursively visit this module's descendants.
            # 递归设置wrapper
            self.hook_modules(sub_module)
    if root: # 对于root进行处理
        this_creator.forward_original_methods.update({module: module.forward})
        module.forward = forward_wrapper_root.__get__(module, module.__class__)
4.2.2.2 TensorWrapper

TensorWrapper реализует функцию-оболочку, а graph_creator.summary — это сеть и другая информация, полученная torchsummary.summary ранее. Вы можете видеть, что этот класс будет проходить сводку, вычислять forward_compute_time и другую информацию и, наконец, строить узел.

Примечание: размеры Activation_Sizes рассчитываются на основе output_shape.

class TensorWrapper(object):
    def __init__(self, tensor, node_desc, graph_creator, activation_size=None):
        self.tensor = tensor
        global object_id
        self.object_id = object_id
        object_id += 1
        self.node_desc = node_desc
​
        i = 0
        for i in range(len(graph_creator.summary)):
            if str(graph_creator.summary[i]['layer_name']) == node_desc:
                break
​
        if i < len(graph_creator.summary) and node_desc == str(graph_creator.summary[i]['layer_name']):
            summary_elem = graph_creator.summary.pop(i)
            forward_compute_time = summary_elem['forward_time']
            backward_compute_time = summary_elem['backward_time']
            if isinstance(summary_elem['output_shape'][0], list):
                activation_sizes = [4.0 * functools.reduce(lambda x, y: x * y, elem)
                                    for elem in summary_elem['output_shape']]
            else:
                activation_sizes = 4.0 * functools.reduce(lambda x, y: x * y, summary_elem['output_shape'])
            parameter_size = 4.0 * float(summary_elem['nb_params'])
            self._node = graph.Node("node%d" % object_id, node_desc=node_desc,
                                    forward_compute_time=forward_compute_time,
                                    backward_compute_time=backward_compute_time,
                                    activation_size=activation_sizes,
                                    parameter_size=parameter_size)
        elif activation_size is not None:
            self._node = graph.Node("node%d" % object_id, node_desc=node_desc,
                                    activation_size=activation_size)
        else:
            self._node = graph.Node("node%d" % object_id, node_desc=node_desc)
        self.graph_creator = graph_creator

Для некоторых встроенных методов это также будет обрабатываться соответствующим образом, как показано ниже.

def __iadd__(self, other):
    self_activation_size = self.node().activation_size
    other_activation_size = other.node().activation_size
    assert(self_activation_size == other_activation_size)
    wrapped_result = TensorWrapper(self.tensor, "Add(inplace)", self.graph_creator,
                                   activation_size=self_activation_size)
    self.tensor += other.tensor
    self.graph_creator.graph.add_edge(self._node, wrapped_result.node())
    self.graph_creator.graph.add_edge(other.node(), wrapped_result.node())
    return wrapped_result

Окончательная переписка:

node58 -- Add(inplace) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=102760448.000, parameter_size=0.000

4.2.3 Постоянство

persist_graph предназначен для вывода результатов профиля в файл.

def persist_graph(self, directory):
    self.graph.to_dot(os.path.join(directory, "graph.dot"))
    with open(os.path.join(directory, "graph.txt"), 'w') as f:
        f.write(str(self.graph))
    self.graph.render_bar_graphs_and_cdfs(directory)

В частности, функция graph.py вызывается для завершения Вот выдержка из функции to_dot:

def to_dot(self, arch):
    dot = graphviz.Digraph()
    for node in self.nodes.values():
        node_desc = "%s\n[forward_compute_time=%.3f,backward_compute_time=%.3f,activation_size=%s,parameter_size=%.1f]" % (
            node.node_desc, node.forward_compute_time, node.backward_compute_time,
            node.activation_size, node.parameter_size)
        if node.stage_id is not None:
            color = self._colors[node.stage_id % len(self._colors)]
            dot.node(node.node_id, node_desc,
               color=color, style='filled')
        else:
            dot.node(node.node_id, node_desc)
    for node in self.nodes.values():
        if node.node_id not in self.edges:
            continue
        for out_node in self.edges[node.node_id]:
            dot.edge(node.node_id, out_node.node_id)
    dot.render(arch)

4.3 Результаты

Мы используем результаты в исходном коде в качестве примера pipedream-pipedream/profiler/translation/profiles/gnmt/graph.txt, чтобы показать вам конкретные результаты.

node1 -- Input0 -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node4 -- Embedding(32320, 1024, padding_idx=0) -- forward_compute_time=0.073, backward_compute_time=6.949, activation_size=6291456.0, parameter_size=132382720.000
node5 -- EmuBidirLSTM(  (bidir): LSTM(1024, 1024, bidirectional=True)  (layer1): LSTM(1024, 1024)  (layer2): LSTM(1024, 1024)) -- forward_compute_time=5.247, backward_compute_time=0.016, activation_size=12582912.0, parameter_size=67174400.000
node2 -- Input1 -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node6 -- Dropout(p=0.2) -- forward_compute_time=0.077, backward_compute_time=0.196, activation_size=12582912.0, parameter_size=0.000
node7 -- LSTM(2048, 1024) -- forward_compute_time=3.190, backward_compute_time=5.348, activation_size=[6291456.0; 131072.0; 131072.0], parameter_size=50364416.000
node8 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node9 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node10 -- Dropout(p=0.2) -- forward_compute_time=0.064, backward_compute_time=0.128, activation_size=6291456.0, parameter_size=0.000
node11 -- LSTM(1024, 1024) -- forward_compute_time=2.491, backward_compute_time=4.203, activation_size=[6291456.0; 131072.0; 131072.0], parameter_size=33587200.000
node12 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node13 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node14 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node15 -- Dropout(p=0.2) -- forward_compute_time=0.059, backward_compute_time=0.121, activation_size=6291456.0, parameter_size=0.000
node16 -- LSTM(1024, 1024) -- forward_compute_time=2.492, backward_compute_time=4.201, activation_size=[6291456.0; 131072.0; 131072.0], parameter_size=33587200.000
node17 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node18 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node19 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6291456.0, parameter_size=0.000
node3 -- Input2 -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node21 -- Embedding(32320, 1024, padding_idx=0) -- forward_compute_time=0.066, backward_compute_time=0.328, activation_size=6291456.0, parameter_size=132382720.000
node20 -- hidden -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node22 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node23 -- RecurrentAttention(  (rnn): LSTM(1024, 1024)  (attn): BahdanauAttention(    (linear_q): Linear(in_features=1024, out_features=1024, bias=False)    (linear_k): Linear(in_features=1024, out_features=1024, bias=False)    (dropout): Dropout(p=0)  )  (dropout): Dropout(p=0)) -- forward_compute_time=4.546, backward_compute_time=6.141, activation_size=[6160384.0; 131072.0; 131072.0; 6160384.0; 288768.0], parameter_size=41979904.000
node24 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node25 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node26 -- __getitem__(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node27 -- __getitem__(3) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node28 -- Dropout(p=0.2) -- forward_compute_time=0.058, backward_compute_time=0.176, activation_size=6160384.0, parameter_size=0.000
node29 -- Concat(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node30 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node31 -- LSTM(2048, 1024) -- forward_compute_time=3.151, backward_compute_time=5.288, activation_size=[6160384.0; 131072.0; 131072.0], parameter_size=50364416.000
node32 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node33 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node34 -- Dropout(p=0.2) -- forward_compute_time=0.061, backward_compute_time=0.174, activation_size=6160384.0, parameter_size=0.000
node35 -- Concat(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node36 -- __getitem__(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node37 -- LSTM(2048, 1024) -- forward_compute_time=3.145, backward_compute_time=5.306, activation_size=[6160384.0; 131072.0; 131072.0], parameter_size=50364416.000
node38 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node39 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node40 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node41 -- Dropout(p=0.2) -- forward_compute_time=0.055, backward_compute_time=0.198, activation_size=6160384.0, parameter_size=0.000
node42 -- Concat(2) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node43 -- __getitem__(3) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=0.0, parameter_size=0.000
node44 -- LSTM(2048, 1024) -- forward_compute_time=3.149, backward_compute_time=15.883, activation_size=[6160384.0; 131072.0; 131072.0], parameter_size=50364416.000
node45 -- __getitem__(0) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node46 -- __getitem__(1) -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=131072.0, parameter_size=0.000
node47 -- Add -- forward_compute_time=0.000, backward_compute_time=0.000, activation_size=6160384.0, parameter_size=0.000
node48 -- Classifier(  (classifier): Linear(in_features=1024, out_features=32320, bias=True)) -- forward_compute_time=5.609, backward_compute_time=1.227, activation_size=194437120.0, parameter_size=132512000.000
   node1 -- node4
   node4 -- node5
   node2 -- node5
   node5 -- node6
   node6 -- node7
   node7 -- node8
   node7 -- node9
   node8 -- node10
   node10 -- node11
   node11 -- node12
   node11 -- node13
   node12 -- node14
   node8 -- node14
   node14 -- node15
   node15 -- node16
   node16 -- node17
   node16 -- node18
   node17 -- node19
   node14 -- node19
   node3 -- node21
   node20 -- node22
   node21 -- node23
   node22 -- node23
   node19 -- node23
   node2 -- node23
   node23 -- node24
   node23 -- node25
   node23 -- node26
   node23 -- node27
   node24 -- node28
   node28 -- node29
   node26 -- node29
   node20 -- node30
   node29 -- node31
   node30 -- node31
   node31 -- node32
   node31 -- node33
   node32 -- node34
   node34 -- node35
   node26 -- node35
   node20 -- node36
   node35 -- node37
   node36 -- node37
   node37 -- node38
   node37 -- node39
   node38 -- node40
   node32 -- node40
   node40 -- node41
   node41 -- node42
   node26 -- node42
   node20 -- node43
   node42 -- node44
   node43 -- node44
   node44 -- node45
   node44 -- node46
   node45 -- node47
   node40 -- node47
   node47 -- node48

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

В следующей статье мы разберем, как рассчитать автоматическое разбиение.

0xEE Личная информация

★★★★★★Думая о жизни и технологиях★★★★★★

Публичный аккаунт WeChat:мысли Росси

ссылка 0xFF

Woohoo.Microsoft.com/En-US/Горячий цвет…

lingvo framework день чтение заметок

Tensorflow понимает, что градиенты нескольких мини-пакетных вычислений сначала накапливаются, а затем распространяются обратно.

Накопление градиента с помощью tensorflow2

В десять раз время расчета модели увеличилось всего на 20%: плагин для замены градиента с открытым исходным кодом OpenAI

PipeDream: Fast and Efficient Pipeline Parallel DNN Training

Paper Interpretation Series 5: Microsoft Stanford и другие PipeDream быстро обучают крупномасштабные нейронные сети

На данный момент 231 you.GitHub.IO/neural-net…

Блог Woohoo.cn на.com/geek found/afraid/14…

Технология оптимизации видеопамяти во время обучения - слияние ОП и контрольная точка градиента

Pytorch Notes 04 - Пользовательский torch.autograd.Function

Учебное пособие по Autograd для PyTorch

Пользовательское расширение pytorch (3) - простое определение и случай torch.autograd.Function

Пользовательское расширение pytorch (2) - torch.autograd.Function завершает пользовательский слой

torch.autograd интерпретации исходного кода PyTorch: подробное объяснение расчета градиента

Обратное распространение

Перевод примечаний к курсу CS231n: примечания по обратному распространению

Самая большая обратная цепь частично упорядоченного множества [двудольный граф]

Топологическая сортировка