[Анализ исходного кода] Конвейер глубокого обучения, параллельный Gpipe(1) --- Базовая реализация конвейера

0x00 сводка

GPipe — это параллельная библиотека для обучения нейронных сетей, разработанная на основе Lingvo (Lingvo — это вторичная среда разработки Google на основе TensorFlow для последовательных моделей) и поддерживающая сверхкрупномасштабные модели.В этой статье представлены ее основные функции и механизм конвейера.

0x01 Обзор

1.1 Что такое GPipe

GPipe — это фреймворк, основанный на Lingvo (Lingvo — это фреймворк для моделей последовательностей, основанный на вторичной разработке Google TensorFlow).GitHub.com/tensorflow/…) разработала параллельную библиотеку для обучения нейронных сетей, поддерживающую сверхбольшие модели, характеристики которой следующие:

• GPipe делит L-уровневую сеть на K составных слоев. Каждый композитный слой работает на отдельном ядре TPU.
• Составные слои K ядра могут выполняться только последовательно, но GPipe вводит стратегию параллельного конвейера, чтобы облегчить проблему производительности этого последовательного выполнения, и подразделяет мини-пакет на несколько более мелких макропакетов для повышения степени параллелизма.
• GPipe также использует простой и эффективный способ пересчета для уменьшения объема памяти, что позволяет дополнительно обучать более крупные модели.

1.2 Проблемы

Среда глубокого обучения — это, по сути, механизм компиляции и выполнения графа вычислений (Graph), состоящий из выражений вычисления (оператора) между тензорами, предоставляющий ряд определений тензора, унарные операции, бинарные операции и т. д. Математические примитивы и автоматическое получение градиентов и обновления модели на основе обратного распространения. После того, как большой объем данных поступает в вычислительный граф для обучения модели в пакетах, модель может изучить внутреннюю взаимосвязь в данных, чтобы получить способность «интеллектуального» восприятия и суждения в соответствующей сцене.

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

• Статистическая эффективность, количество эпох, необходимое для достижения желаемой точности;
• Аппаратная эффективность, время, необходимое для завершения одной эпохи. Общее время обучения для достижения желаемого уровня точности является просто произведением этих двух показателей;

GPU в основном предоставляет два вида ресурсов:вычислительные ресурсыиресурсы пропускной способности памяти. Таким образом, обучение больших моделей имеет две фундаментальные проблемы:эффективность памятииВычислительная эффективность.

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

0x02 Параллелизм

В отрасли используются некоторые параллельные механизмы для достижения целей оптимизации.

2.1 Классификация механизмов и компромиссы

В этом разделе приведены следующие основные ссылки:

Efficient Large-Scale Language Model Training on GPU Clusters

DeepSpeed: Extreme-scale model training for everyone

DeepSpeed: инструмент обучения гипермасштабируемой модели для всех

PipeDream: Fast and Efficient Pipeline Parallel DNN Training.

В документе «Эффективное обучение крупномасштабным языковым моделям на кластерах графических процессоров» NVIDIA описывает три необходимых метода распараллеливания для распределенного обучения очень больших моделей:

• Параллелизм данных
• Параллелизм тензорной модели
• Параллелизм конвейерной модели

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

Параллелизм данных является наиболее распространенным подходом. Его характеристики следующие:

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

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

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

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

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

• Параллелизм модели включает в себя разделение модели между рабочими процессами таким образом, чтобы каждый рабочий процесс оценивал и обновлял только подмножество параметров модели. Это можно разделить на межуровневый параллелизм и внутриуровневый параллелизм моделей.
• Межуровневый параллелизм модели разделяет слои модели между несколькими рабочими процессами.
• Параллелизм внутриуровневой модели разделяет параметры модели каждого уровня на несколько устройств. Параллелизм внутрислойной модели в некоторых статьях называется «параллелизм модели на уровне тензора», который заключается в сегментации тензора модели определенного слоя (например, переменной в линейном/плотном слое), тем самым разделяя большой тензор модели на несколько относительных Меньший тензор для параллельных вычислений;
• Межслойные значения (активации и градиенты) часто являются единственными параметрами, которые необходимо передавать между машинами.

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

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

• эффективность памяти: Обучение DNN с параллельным моделированием приводит к серьезному недоиспользованию ресурсов графического процессора. Модельный параллелизм сокращает использование памяти пропорционально количеству рабочих процессов за счет разделения активной памяти между модельно-параллельными рабочими процессами. Важно отметить, что это единственный способ уменьшить активную память одного сетевого уровня.
• Вычислительная эффективность: Параллелизм модели вычислительно неэффективен из-за необходимости дополнительных активаций связи при каждом прямом и обратном распространении. Параллелизм модели требует высокой пропускной способности связи и плохо масштабируется для узлов с ограниченной пропускной способностью связи. Кроме того, каждый параллельный рабочий процесс модели снижает объем вычислений, выполняемых между каждым этапом связи, что влияет на эффективность вычислений. Параллелизм моделей часто используется в сочетании с параллелизмом данных для компромисса между памятью и вычислительной эффективностью.
• Эффективность разработки: Бремя разделения модели на несколько графических процессоров возлагается на программиста, и определение того, как лучше всего разделить модель DNN между работниками, является сложной задачей даже для самых опытных специалистов по машинному обучению, что часто приводит к дополнительной неэффективности. В некоторых недавних работах изучалось, как использовать обучение с подкреплением для автоматического определения местоположения устройства для параллелизма модели. К сожалению, такие методы онлайн-принятия решений требуют много времени и ресурсов; они также не сочетают в себе конвейерную обработку, параллелизм данных и параллелизм моделей.

2.1.3 Конвейерный параллелизм

Параллелизм модели конвейера в некоторых статьях называется параллелизмом модели уровня конвейера, и его характеристики:

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

Но есть еще некоторые проблемы с конвейерным параллелизмом:

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

2.2 Как использовать

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

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

• Некоторые операторы (такие как свертка) имеют небольшое количество рабочих параметров, но большое количество промежуточных результатов.Для уменьшения объема передачи наиболее экономичным методом является параллелизм данных. То есть данные сегментированы, разные устройства обрабатывают разные данные, а параметры время от времени синхронизируются между устройствами;
• Для некоторых операторов промежуточные результаты вычислений меньше количества параметров, поэтому они подходят для модельного параллелизма.
• Также есть некоторые операторы, количество сетевых параметров/промежуточных результатов вычислений очень велико, и оптимальным может быть использование конвейерного параллелизма (то есть в виде реле).
• По сравнению с оператором, который использует только один параллельный режим, оператор также может использовать несколько параллельных режимов одновременно, что может еще больше уменьшить объем передачи.Например, если скрытый слой относительно велик, можно сократить матрица данных и модель Матрица сокращений.

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

0x03 Pytorch вручную указывает параллелизм

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

Давайте посмотрим, как Pytorch указывает это вручную, основной отрывок (перевод):

py torch.org/tutorials/i…

3.1 Базовые знания

PyTorch использует Tensor в качестве базовой единицы, что больше соответствует интуиции инженеров-алгоритмов при написании сценариев Python, а также построении и обучении моделей объектно-ориентированным способом. Назначение и нарезка Tensor так же просты в использовании, как и numpy.

PyTorch — это однокарточная перспектива Тензоры и скрипты моделей на одном устройстве не имеют прямого отношения к Тензорам и скриптам моделей на другом устройстве Скрипты моделей на каждом устройстве полностью симметричны (Зеркало) Простейший параллелизм данных Тем не менее, нет очевидных недостатки в дизайне, таком как PyTorch. Сценарий на каждом устройстве запускается для части обновления модели (оптимизатора) одного и того же пакета, а синхронизация модели (операция AllReduce) выполняется единообразно для завершения параллелизма данных.Это модуль DDP (Distributed DataParallel) PyTorch.

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

3.2 Особенности

Параллелизм модели PyTorchРазделите одну модель на разные графические процессоры вместо дублирования всей модели на каждом графическом процессоре.(В частности, если предположить, что модельmСодержит 10 слоев. При использованииDataParallel, каждый GPU имеет копию каждого из 10 слоев, тогда как каждый GPU может содержать 5 слоев, если параллелизм модели используется на двух GPU).

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

3.3 Основное использование

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

import torch
import torch.nn as nn
import torch.optim as optim
​
class ToyModel(nn.Module):
  def __init__(self):
    super(ToyModel, self).__init__()
    self.net1 = torch.nn.Linear(10, 10).to('cuda:0')  # 将net1放置在第1个GPU上
    self.relu = torch.nn.ReLU()
    self.net2 = torch.nn.Linear(10, 5).to('cuda:1')   # 将net2放置在第2个GPU上
​
  def forward(self, x):
    x = self.relu(self.net1(x.to('cuda:0')))
    return self.net2(x.to('cuda:1'))

Обратите внимание, что дляToyModel, в дополнение к пяти для размещения линейных слоев и тензоров на соответствующих устройствахto(device)Кроме того, приведенное выше очень похоже на реализацию функции на одном графическом процессоре. Это единственное место в модели, которое нужно изменить (т.to(device)).backward()иtorch.optimбудет автоматически следоватьградиент(градиенты), как если бы модель была графическим процессором.При вызове функции потерь просто убедитесь, что метка находится на том же устройстве, что и выход.

model = ToyModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.paraeters(), lr=0.001)
​
optimizer.zero_grad()
outputs = model(torch.randn(20, 10))
labels = torch.randn(20, 5).to('cuda:1') # ToyMode 的 output 是在 'cuda:1' 上，此处的 label 也应该置于 'cuda:1' 上
loss_fn(outputs,labels).backward()
optimizer.step()

3.4 Применение параллелизма моделей к существующим модулям

Существующие модули с одним GPU можно запускать на нескольких GPU, изменив всего несколько строк. Следующий код показывает, как разбить егоtorchvision.models.reset50()для двух графических процессоров. идеи исходят из существующихResNetНаследование модулей и разделение слоев на два графических процессора во время сборки. Затем перезапишитеforwardспособ сшить две подсети, соответствующим образом сдвинув промежуточные выходы.

from torchvision.models.resnet import ResNet, Bottleneck
​
num_classes = 1000
​
class ModelParallelResNet50(ResNet):
    def __init__(self, *args, **kwargs):
        super(ModelParallelResNet50, self).__init__(
            Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs)
​
        self.seq1 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,
​
            self.layer1,
            self.layer2
        ).to('cuda:0')  # 放置在第1个GPU上
​
        self.seq2 = nn.Sequential(
            self.layer3,
            self.layer4,
            self.avgpool,
        ).to('cuda:1')  # 放置在第2个GPU上
​
        self.fc.to('cuda:1')
​
    def forward(self, x):
        x = self.seq2(self.seq1(x).to('cuda:1'))
        return self.fc(x.view(x.size(0), -1))

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

заМодель слишком велика, чтобы поместиться на одном графическом процессоре.Приведенная выше реализация решает проблему. Однако, как вы могли заметить,Если модель подходит, то она (параллельная модель) будет работать медленнее, чем работа на одном GPU. Это связано с тем, что в любой момент времени работает только один из двух графических процессоров, а другой там ничего не делает. существуетlayer2иlayer3между ними промежуточный вывод необходимо изменить сcuda:0скопировать вcuda:1, что еще больше ухудшает производительность.

3.5 Ускорение за счет конвейерной обработки входов

В следующих экспериментах мы далее делим каждую партию из 120 изображений на 20 изображений. Когда PyTorch запускает операции CUDA асинхронно, реализация обеспечивает параллелизм без создания нескольких потоков.

class PipelineParallelResNet50(ModelParallelResNet50):
    def __init__(self, split_size=20, *args, **kwargs):
        super(PipelineParallelResNet50, self).__init__(*args, **kwargs)
        self.split_size = split_size
​
    def forward(self, x):
        splits = iter(x.split(self.split_size, dim=0))
        s_next = next(splits)
        s_prev = self.seq1(s_next).to('cuda:1')
        ret = []
​
        for s_next in splits:
            # A. s_prev runs on cuda:1
            s_prev = self.seq2(s_prev)
            ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
​
            # B. s_next runs on cuda:0, which can run concurrently with A
            s_prev = self.seq1(s_next).to('cuda:1')
​
        s_prev = self.seq2(s_prev)
        ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
​
        return torch.cat(ret)
​
​
setup = "model = PipelineParallelResNet50()"
pp_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times)
​
plot([mp_mean, rn_mean, pp_mean],
     [mp_std, rn_std, pp_std],
     ['Model Parallel', 'Single GPU', 'Pipelining Model Parallel'],
     'mp_vs_rn_vs_pp.png')

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

0x04 Ключевые технологии

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

Поэтому технология автоматического параллелизма (как автоматически решить проблему выбора параллельной стратегии на уровне фреймворка) стала предметом исследований.

Автоматическое распараллеливание прогнозирует и выбирает оптимальную параллельную стратегию путем создания модели затрат (нельзя гарантировать, что она будет оптимальной стратегией на данный момент, потому что выбор оптимальной стратегии — это задача NP-Hard).

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

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

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

0x05 Базовые знания и система поддержки

5.1 Платформа Лингво

5.1.1 Основные компоненты

Основные компоненты Lingvo:

• Models: модель — это абстрактная коллекция, содержащая одну или несколько задач. Модель эквивалентна слою оболочки для Задачи. Для многозадачных моделей Модель будет контролировать, какие переменные будут общими между задачами и как каждая задача отбирается во время обучения.
• Tasks: Задача — это описание полной проблемы оптимизации, такой как классификация изображений или распознавание речи. Содержит входной генератор.
• Layers: Слой представляет собой случайную функцию с обучаемыми параметрами. Слой может содержать другие слои в качестве дочерних. Softmax, LSTM, Attension и даже задача — примеры Layer.
• Params: этот объект содержит гиперпараметры модели. Слои, задачи и модели создаются на основе спецификаций в Params. Параметры иерархичны, параметры объекта могут содержать параметры его дочерних объектов.
• NestedMap: структура словаря для передачи данных. Большинство объектов Python в коде являются либо экземплярами Tensor, либо подклассами BaseLayer или NestedMap.

5.1.2 Определение модели

В Lingvo сеть представляет собой вложенную структуру слоев. Большинство классов в Lingvo являются подклассами [Lingvo/core/base_layer.py] BaseLayer.

• Params: Используется для настройки класса, определяет ключи, необходимые для настройки, и эти ключи должны быть определены при создании объекта. Объекты Params также могут содержать объекты Params, используемые для настройки подслоев. Каждый класс слоя будет иметь classmethod параметров, этот метод создаст новый объект параметров и настроит слой с определенными ключами, а также даст некоторые разумные значения по умолчанию для этих ключей.

  Свойства объекта Params включают:

  • cls: класс Python, связанный с объектом tParams. Это можно использовать для создания экземпляров классов;
  • name: имя слоя;
  • dtype: Тип данных по умолчанию для использования при создании переменных.

• __init__конструктор: здесь должны быть созданы все подслои и переменные.

• CreateVariable: Метод для создания переменной. Каждый уровень отвечает за создание и управление своими переменными.

• CreateChild: метод создания подслоев.

• FProp: все слои имеют функцию FProp(), которая реализует прямое распространение слоя и будет вызываться на прямом шаге вычисления. Поскольку это может быть выполнено на разных устройствах во время распределенного обучения, Lingvo получает доступ к переменным через параметр theta вместо self.vars или self.theta из соображений производительности.

• FPropMeta: вернуть слой оFPropВычисляемые метаданные. вmeta.flopsДает предполагаемое количество операций с плавающей запятой при получении входного тензора.

Для реализации алгоритма модели особенно важны два показателя:

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

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

5.2 Вычислительная память

5.2.1 Общий анализ

Мы в основном обращаемся кZeRO: Memory Optimization Towards Training A Trillion Parameter Modelsидея этой статьи.

Во время обучения модели большая часть памяти потребляется одной из трех вещей:

• и) Активировать.
• ii) состояние OGP, то есть тензор, состоящий из состояния оптимизатора, градиентов параметров и самих параметров.
• iii) Временный буфер.

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

Мы анализируем их один за другим.

5.2.2 Функция активации

Для функции активации она имеет следующие характеристики:

• Дополнительная видеопамять, используемая функцией активации, увеличивается с размером пакета.Если для пакета установлено значение 1, обучение модели с триллионом параметров создаст более 1 ТБ видеопамяти для функции активации.
• Существующие решения в отрасли, такие как бумагаTraining deep nets with sublinear memory cost, может устранить почти всю память, необходимую для активации, за счет 33% накладных расходов на повторные вычисления. Этот трюк называется градиентной контрольной точкой, иногда называемой рематериализацией, повторной переадресацией.

5.2.3 Статус OGP

5.2.3.1 Параметры самой модели

Параметры самой модели относятся к Весу и Смещению каждого сетевого слоя, которые будут заняты после загрузки модели. Кроме того, следует отметить, что некоторые слои имеют параметры, такие как CNN, RNN, а некоторые слои не имеют параметров, например, уровень активации, слой пула и т. д.

5.2.3.2 Параметры оптимизатора

Параметры оптимизатора относятся к параметрам, генерируемым моделью в процессе оптимизации, то есть обратного распространения.Эти параметры в основном относятся к dw, то есть к градиенту.В SGD его размер такой же, как и параметры, поэтому во время оптимизации параметры модели совпадают с параметрами Используемая память будет удвоена.

Стоит отметить, что разным оптимизаторам необходимо сохранять разные параметры оптимизации, для Адама, поскольку ему также необходимо сохранять остальные параметры, количество параметров модели будет увеличено в четыре раза на интервале оптимизации.

5.2.3.3 Пример

Для состояния OGP возьмем в качестве примера ADAM. Обучение моделей смешанной точности с параметрами Ψ с использованием ADAM.

• Параметры модели: копия параметров fp16, требования к памяти 2Ψ байт.
• градиент: копия градиента fp16, требуемая память 2Ψ байт.
• Состояние оптимизатора: fp32 копии параметров, Momentum и Variance с требованиями к памяти 4Ψ, 4Ψ и 4Ψ байт соответственно. Обозначим множитель памяти для состояний оптимизатора через K, то есть дополнительная память, необходимая для их хранения, составляет KΨ байт.

В целом, для состояния OGP требуется память 2Ψ+2Ψ+KΨ=16Ψ байт (K=12 для смешанной точности ADAM).

детали следующим образом:

Синий — параметры, оранжевый — градиенты, зеленый — состояния оптимизатора.

В формуле потребления памяти Ψ представляет размер модели (количество параметров), K представляет собой множитель памяти состояния оптимизатора, а Nd представляет параллелизм данных. В этом примере мы предполагаем размер модели Ψ = 7,5 миллиардов, обучение со смешанной точностью на основе оптимизатора Adam, параллелизм данных Nd = 64 (т. е. 64 графических процессора) и K = 12.

Для такой модели, как GPT-2 с 1,5 миллиардами параметров, это приводит к необходимости памяти не менее 24 ГБ, что намного больше, чем 3 ГБ, необходимых для сохранения одних только параметров fp16.

5.2.4 Временный буфер

Временный буфер — это буфер, используемый для хранения временных результатов, например, для модели GPT-2 с 1,5 миллиардами параметров для буфера fp32 потребуется 6 ГБ памяти.

5.3 Вычислительная мощность

5.3.1 Фоновые знания

• FLOPS: Обратите внимание, что все это пишется с заглавной буквы, что является аббревиатурой операций с плавающей запятой в секунду, что означает количество операций с плавающей запятой в секунду, что понимается как скорость вычислений. Это мера производительности оборудования.
• FLOPs: обратите внимание, что s в нижнем регистре, что является аббревиатурой операций с плавающей запятой (s означает множественное число), что означает арифметику с плавающей запятой, которая понимается как количество вычислений. Может использоваться для измерения сложности алгоритма/модели.

Вычислительная мощность, необходимая для прямого распространения, определяется выражениемFLOPsпроявляется, тоFLOPsКак рассчитать?

Мы знаем, что когда модель выполняет прямое распространение, будут выполняться такие операции, как свертка, объединение, BatchNorm, Relu, Upsample и т. д. Выполнение этих операций будет иметь соответствующее потребление вычислительной мощности, среди которых потребление вычислительной мощности, соответствующее свертке, является самым высоким. Итак, давайте возьмем операцию свертки в качестве примера, чтобы увидеть вычислительную мощность, соответствующую свертке.

Процесс получения:Сверточный слой wx + b должен вычислить две части, сначала рассмотрим сумму расчета первой половины wx:

сделать :

• k представляет размер ядра свертки;
• c представляет количество входных карт объектов;

Затем для выходной карты объектов нане замужемЕдиницы имеют:

k * k * c 次乘法，以及 k * k * c - 1 次加法

Если разрешение выходной карты объектов равно H * W и выводятся o карты объектов, общее количество единиц, содержащихся в выходной карте объектов, равно H * W * o.

Следовательно, этот сверточный слой имеет:

k * k * c * H * W * o 次乘法          --（1）
(k * k * c - 1) * H * W * o 次加法    --（2）

Затем рассмотрим объем вычислений, связанных с членом смещения b:

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

H * W * o 次加法      --（3）

Накопление количества вычислений двух частей wx и b сверточного слоя имеет:

Умножение уравнения (1):

k * k * c * H * W * o 次乘法

Уравнение (2) + уравнение (3) дополнение:

(k * k * c - 1) * H * W * o  + H * W * o  = k * k * c * H * W * o

Видно, что формула (2) + формула (3) = формула (1)

Следовательно, для сверточного слоя со смещением потребляемая вычислительная мощность этого слоя составляет:

k * k * c * H * W * o

5.3.2 Реализация в lingvo

В Lingvo удельная вычислительная мощность рассчитывается путем расчета каждого класса.FPropMetaПолные, эти методы реализуются каждым классом в соответствии со своими особенностями. Давайте найдем несколько конкретных примеров, чтобы увидеть, как рассчитать FLOPS.

Conv2DLayerNoPadding рассчитывается следующим образом:

  @classmethod
  def FPropMeta(cls, p, inputs):
    py_utils.CheckShapes((inputs,))
    b, h, w, c = inputs
    fh, fw, ic, oc = p.filter_shape
    assert ic == c
    sh, sw = p.filter_stride
    if p.padding == 'SAME':
      oh = sympy.ceiling(h / sh)
      ow = sympy.ceiling(w / sw)
    else:
      oh = sympy.ceiling((h - fh + 1) / sh)
      ow = sympy.ceiling((w - fw + 1) / sw)
    flops = b * oh * ow * fh * fw * ic * oc * 2  # mul/add counts as 2 flop.
    outputs = tshape.Shape([b, oh, ow, oc])
    return py_utils.NestedMap(flops=flops, out_shapes=(outputs,))

DropoutLayer рассчитывается следующим образом:

​
  @classmethod
  def FPropMeta(cls, p, inputs, *args):
    py_utils.CheckShapes((inputs,))
    flops_per_element = 10  # Approximately 10 flops per element.
    return py_utils.NestedMap(
        flops=inputs.num_elements() * flops_per_element, out_shapes=(inputs,))
​

FLOPS для BatchNormLayer рассчитывается следующим образом.

  @classmethod
  def FPropMeta(cls, p, inputs, padding=None):
    py_utils.CheckShapes((inputs,))
    return py_utils.NestedMap(
        flops=inputs.num_elements() * _BN_FLOPS_PER_ELEMENT,
        out_shapes=(inputs,))

ActivationLayer рассчитывается следующим образом:

  @classmethod
  def FPropMeta(cls, p, inputs):
    py_utils.CheckShapes((inputs,))
    return py_utils.NestedMap(
        flops=inputs.num_elements() * GetFlops(p.activation),
        out_shapes=(inputs,))

0x06 Конвейер

6.1 Базовые знания

6.1.1 Проблемы

проблемы общения

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

Например, при использовании сервера параметров трехэтапный процесс выглядит следующим образом:Pull weight ---> Compute new weight ---> Push new weight.

Если трехстадийный процесссериалДля связи и вычислений, независимо от того, является ли связь быстрой или медленной, накладные расходы на это время заставят каждую итерацию в распределенной среде занять больше времени, чем автономная версия (пропускная способность или задержка Ethernet будут влиять только на продолжительность этого времени, но не сократить это время до нуля). Таким образом, совмещение связи и вычислений с целью «замаскировать» время связи является почти необходимым шагом. Как совместить вычисления и передачу, чтобы улучшить использование устройства, очень сложно.

недоиспользуемый

В принципе, мы можем обучать более крупные модели DNN на GPU или TPU с помощью параллельных вычислений. Однако из-за последовательного характера DNN этот подход может привести к тому, что во время вычислений будет активен только один ускоритель, что не может полностью использовать вычислительную мощность устройства.

6.1.2 Как спроектировать систему

Возвращаясь к процессу обучения нейронной сети, как спроектировать систему, чтобы совмещать вычисления и связь?

Есть две особенности, которые можно использовать при обратном распространении:

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

Поэтому в соответствии с этой особенностью был введен конвейерный параллелизм.

Параллелизм данных и параллелизм моделей позволяют устройствам выполнять вычисления на одном уровне, в то время как конвейерный параллелизм делит задачи на несколько этапов с четкой последовательностью и назначает разные этапы разным вычислительным устройствам, так что одно устройство отвечает только за расчет сети. среднего слоя. Сеть модели асимметрично распределена по устройствам, причем каждое устройство «ретранслирует» часть сети. Между каждым этапом и следующим этапом необходимо передавать только одни смежные тензорные данные, а объем передачи данных на машину не имеет ничего общего с общим размером сети, общим количеством машин и параллельным масштабом.

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

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

6.2 Обзор конвейера Gpipe

6.2.1 Ключевые моменты

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

У GPipe есть несколько ключевых выводов:

• Сетевой раздел: разделите N-уровневую сеть на K разделов, каждый раздел выполняется на отдельном TPU, и между разделами необходимо вставить некоторые операции сетевого взаимодействия.
• Конвейерный параллелизм (конвейерный параллелизм): технология конвейерного параллелизма в ЦП используется в глубоком обучении, в основном для лучшей реорганизации двух операций вычислений и сетевого взаимодействия. То есть обучающие выборки мини-пакета автоматически делятся на более мелкие микропакеты и запускаются в конвейере, чтобы ядра ТПУ могли работать параллельно.
• Накопление градиента: градиенты всегда накапливаются микропакетами, поэтому количество разделов не влияет на качество модели.
• Повторная материализация: повторная материализация означает, что в процессе прямого расчета GPipe записывает только выходные данные деления этапа.При расчете градиента GPipe повторно выполняет логику прямого расчета, чтобы получить значение каждого оператора. , а затем вычислить результат градиента. Это то же самое, что и контрольная точка градиента с открытым исходным кодом OpenAI, за исключением того, что GPipe реализован на TPU, а OpenAI может работать только на GPU.

6.2.2 Иллюстрация

• Левый край изображения ниже — исходная модель.

• Правый конец показывает, что модель нейронной сети GPipe с несколькими упорядоченными слоями разделена на четыре ускорителя. Fk — составная функция прямого вычисления k-го раздела. Bk — соответствующая функция обратного распространения. Bk зависит от промежуточных функций активации Bk+1 и Fk из верхних слоев.

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

6.2.3 Проблемы

По идеям диссертации мы поставили несколько вопросов, и стремимся следовать схеме в дальнейшем.

• Как разделить этапы?

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

• По какому разделу делать поток?

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

• Как сделать рематериализацию?

6.3 Этапы делятся по вычислительной мощности

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

6.3.1 PartitionSequentialLayers

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

• Ввод: параметр слоя или список параметров слоя;
• На выходе: список параметров FeatureExtractionLayer;

Логика такова:

• Если params — это просто слой, то создайте слой в список подслоев, содержащих подслои;

• Используйте FPropMeta, чтобы рассчитать формы и общее количество флопов этого списка подписок, и назначьте их для гисто;

• Используйте histo для расчета нормализованной кумулятивной гистограммы стоимости слоя;

• Создайте переменную частей;

  • Переменная представляет собой массив размером num_partitions, каждый элемент массива также является массивом;
  • По гистограмме сабы разбиваются на каждый пункт по частям, так что каждая часть[i] имеет несколько слоев, а некоторые операторы с небольшой вычислительной мощностью объединены в одну часть, цель сделать конечные части каждой вычислительная мощность изделия максимально возможная;

• Преобразование частей в список параметров FeatureExtractionLayer;

def PartitionSequentialLayers(params, num_partitions, *shapes):
  r"""Partition a layer composed of sequential layers.
​
  This routine strives to partition layers so that each partition costs roughly
  the same flops given the input shapes.
​
  Args:
    params: A layer param or a list of layer param.
    num_partitions: The desired number of partitions.
    *shapes: A tuple of tshape.Shape representing input tensors to the first
      layer.
​
  Returns:
    A list of FeatureExtractionLayer params.
  """
​
  # Recursively concatenate SequentialLayer into a list.
  # SequentialLayer 是一个层，其作用是把若干层按顺序连接起来
  def FlattenSeq(p):
    if isinstance(p, list): # 已经是列表则返回
      return p
    if p.cls not in [builder_layers.SequentialLayer, FeatureExtractionLayer]:
      return [p.Copy()]
    subs = []
    for _ in range(p.repeat): # 把p包含的所有层都组装成一个层列表
      for s in p.sub:
        subs += FlattenSeq(s)
    return subs
​
  # 如果params是一个layer，那么就依据这个layer，构建一个包含sub-layers的新列表subs，如果是列表则直接返回
  subs = FlattenSeq(params)
​
  assert len(shapes) == 1
  tf.logging.info('num_partitions: {} input_shape: {}'.format(
      num_partitions, shapes[0]))
​
  # 利用 FPropMeta 计算出来这个 subs 列表的shapes和总flops，赋值给了 histo
  # Computes the estimate cost for each sub layer.
  # 假设有7个sub-layers，其flops分别是 10，40，30，10，20，50,10
  total, histo, output_shapes = 0, [], []
  for i, s in enumerate(subs):
    s.name = 'cell_%03d' % i
    meta = s.cls.FPropMeta(s, *shapes) # 
    total += meta.flops
    histo.append(total)
    output_shapes.append(meta.out_shapes)
    shapes = meta.out_shapes
  tf.logging.vlog(1, 'len %d histogram = %s', len(subs), histo)
  # 则对应的histo 为：[10，50，80，90，110，160, 170]，total为170
​
  # 利用 histo 计算出来一个层代价（layer's cost）的归一化累积直方图
  # Computes the normalized cumulative histogram of the layer's cost.
  histo_pct = [float(x / total) for x in histo]
  tf.logging.vlog(1, 'cost pct = %s', histo_pct)
  # histo_pct 为 [1/17,5/17,8/17,9/17,11/17,16/17, 1], 
  # 假设 num_partitions = 3
​
  # 构建一个parts变量，该变量是一个num_partitions大小的数组，数组每个item也是一个数组
  # 依据直方图把subs分到parts中的每个item之中，这样每个parts[i]都拥有部分layers，目的是让最终 parts 每个item的算力尽量相同
  # i-th sub layer is put into partition j, where j is roughly i-th cumulative
  # histogram times num_partitions.
​
  parts = [[] for _ in range(num_partitions)]
  parts_cost = [0] * num_partitions
  pre_hist_cost = 0
  for i, s in enumerate(subs):
    # 从histogram数组中找出s对应cost的index，j也就是s对应的partition
    # 对于i,s，则 histo_pct[i] * num_partitions 分别为: [3/17, 15/17, 24/17, 27/17, 33/17, 48/17，3]，j分别为[0,0,1,1,1,2,2]
    j = min(int(histo_pct[i] * num_partitions), num_partitions - 1)
    # The boundary at parts[j] where j > 0
    if j > 0 and not parts[j]:
      parts_cost[j - 1] = histo_pct[i - 1] - pre_hist_cost
      pre_hist_cost = histo_pct[i - 1]
    parts[j].append(s) # 把s加入到对应的partition
    # 三个桶内容分别为:[1,2],[3,4,5],[6,7]
    # 对应每个桶的flops为: [60,280,330]
    
  # 把parts转换成一个 FeatureExtractionLayer 列表
  parts_cost[num_partitions - 1] = 1.0 - pre_hist_cost
  seqs = []
  for i, pa in enumerate(parts):
    tf.logging.info('Partition %d #subs %d #cost %.3f', i, len(pa),
                         parts_cost[i])
    seqs.append(FeatureExtractionLayer.Params().Set(name='d%d' % i, sub=pa))
  return seqs

6.3.2 FeatureExtractionLayer

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

FeatureExtractionLayer извлекает объекты из последовательности слоев со следующими специфическими характеристиками:

• объединить несколько слоев в последовательность;
• Может получать и передавать очки активации;

class FeatureExtractionLayer(base_layer.BaseLayer):
  """A layer that extrac features from a sequence of layers.
​
  FeatureExtractionLayer is a layer which connects a few layers in a sequence.
  It is also capable of fetching and forwarding activation endpoints.
  # TODO(huangyp): Make it a sublayer of builder_layers.SequentialLayer
  """
​
  @classmethod
  def Params(cls):
    p = super().Params()
    p.Define('variable_name_prefix', '',
             'Prefix for variable names in sub layers')
    p.Define('sub', [], 'A list of layers' params.')
    p.Define('num_act_inputs', 0, 'Number of activation inputs.')
    p.Define('num_act_outputs', 0, 'Number of activation outputs.')
    p.Define('act_fetch_layers', [],
             'Names of fetch layers that cached extra activations')
    return p
​
  def __init__(self, params):
    super().__init__(params)
    p = self.params
    assert p.num_act_inputs >= 0
    assert p.num_act_outputs >= 0
    p.act_fetch_layers = p.act_fetch_layers or []
    assert p.num_act_outputs == p.num_act_inputs + len(p.act_fetch_layers)
    self._seq = []
    for sub in p.sub:
      assert sub.name
      sub.name = p.variable_name_prefix + sub.name
      self.CreateChild(sub.name, sub)
      self._seq.append((sub.name, self.children[sub.name])) # 把一些层连接成一个序列
​
  def FProp(self, theta, *args): # 实现该层的前向传播，在计算的前向step时将会被调用
    p = self.params
    assert len(args) > p.num_act_inputs
    out_args = args[:-p.num_act_inputs] if p.num_act_inputs > 0 else args
    extra_args = args[-p.num_act_inputs:] if p.num_act_inputs > 0 else ()
    for (name, ch) in self._seq:
      th = theta[name]
      out_args = _ToTuple(out_args)
      out_args = ch.FProp(th, *out_args)
    # Append fetched activations to fprop outputs.
    for fetch_layer in p.act_fetch_layers:
      assert fetch_layer in self.children
      activation = self.children[fetch_layer].activation # 子层激活点
      if isinstance(activation, (tuple, list)):
        activation = activation[0] # 如果是list，得到相应激活点
      extra_args += (activation,) # 把激活点添加进来
    if extra_args:
      out_args = _ToTuple(out_args) + extra_args # 最终返回所有激活点
    return out_args
​
  @classmethod
  def FPropMeta(cls, p, *args): # 返回该层关于`FProp`计算的元数据
    assert len(args) > p.num_act_inputs
    seq_args = args[:-p.num_act_inputs] if p.num_act_inputs > 0 else args
    extra_args = args[-p.num_act_inputs:] if p.num_act_inputs > 0 else ()
    total = 0
    act_fetch_metas = {}
    for sub in p.sub:
      meta = sub.cls.FPropMeta(sub, *seq_args)
      if sub.name in p.act_fetch_layers:
        act_fetch_metas[sub.name] = meta.out_shapes[0]
      total += meta.flops
      seq_args = meta.out_shapes
    for fetch_layer in p.act_fetch_layers:
      extra_args += (act_fetch_metas[fetch_layer],)
    return py_utils.NestedMap(flops=total, out_shapes=seq_args + extra_args)

Процесс расчета показан на рисунке ниже Конкретные значения см. в примерах в приведенных выше абзацах кода:

  +--------------+   +--------------+   +---------------+
  |              |   |              |   |               |
  |  sub-layer 1 |   |  sub-layer 2 |   |  sub-layer n  |
  |              |   |              |   |               |
  +-------+------+   +-------+------+   +--------+------+
          |                  |                   |
          |FPropMeta         |FPropMeta          |FPropMeta
          |                  |                   |
          v                  v                   v
       flops 1            flops 2             flops n
          +                  +                   +
          |                  |                   |
          |                  |                   |
          +--------------------------------------+
                             |
                             |
                             v
                  for i, s in enumerate(subs):
                     total += meta.flops
                     histo.append(total)
                  histo=[10,50,80,90,110,160,170]
                             +
                             |
                             |
                             v
Computes the normalized cumulative histogram of the layer's cost
        histo_pct = [float(x / total) for x in histo]
       histo_pct=[1/17,5/17,8/17,9/17,11/17,16/17,1]
                             +
                             |
                             |
                             +
           Assign layers to partition based on histogram
                   [1,2],[3,4,5],[6,7]
                             +
                             |
                             |
                             v
      +----------------------+----------------------------+
      | parts                                             |
      |                                                   |
      | +--------------+  +------------+  +-------------+ |
      | | sub-layer 1  |  |sub-layer 3 |  | sub-layer 6 | |
      | |              |  |            |  |             | |
      | | sub-layer 2  |  |sub-layer 4 |  | sub-layer 7 | |
      | |              |  |            |  |             | |
      | |              |  |sub-layer 5 |  |             | |
      | +--------------+  +------------+  +-------------+ |
      +---------------------------------------------------+

6.4 Распределение конвейера

6.4.1 Базовый класс SeqLayer

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

Параметры включают:

• before_tpl : настроить слои CNN, которые выполняются перед конвейером;
• cell_tpl : список FeatureExtractionLayer;

6.4.1.1 Инициализация

Логика функции инициализации такова:

• Перейдите к before_tpl, вызовите CreateChild для каждого элемента, чтобы построить его подслой, и добавьте элемент в _before_layers;
• Пройдите через cell_tpl, вызовите CreateChild для каждого элемента, чтобы построить его подслой, и добавьте элемент в _cells;

  def __init__(self, params):
    super().__init__(params)
    p = self.params
    self._before_layers = []
    self._cells = []
    # 遍历before_tpl，对于每个item调用CreateChild构建其子层，把item添加到 _before_layers 之中
    for l in p.before_tpl:
      self.CreateChild(l.name, l)
      self._before_layers.append((l.name, self.children[l.name]))
    # 遍历cell_tpl，对于每个item调用CreateChild构建其子层，把item添加到 _cells 之中  
    for l in p.cell_tpl:
      self.CreateChild(l.name, l)
      self._cells.append((l.name, self.children[l.name]))

6.4.1.2 _CreateChildrenVariables

Строить переменные. Логика следующая:

• Если используется тпу, то

  • использоватьcluster.WorkerDeviceInModelSplit(0)Для сборки before_tpl_device, то есть использовать первое устройство кластера как before_tpl_device;
  • Пройдитесь по другим устройствам в кластере и назначьте их на cell_devices;

• Пройдите _before_layers и разверните каждую переменную в before_tpl_device;

• Пройдите _cells и разверните каждую переменную в cell_devices;

  def _CreateChildrenVariables(self):
    p = self.params
​
    num_cells = len(p.cell_tpl)
    before_tpl_device = ''
    cell_devices = [''] * num_cells
    if py_utils.use_tpu(): # 如果使用 tpu
      # 利用 `cluster.WorkerDeviceInModelSplit(0)` 来构建 before_tpl_device，即用集群的第一个设备作为 before_tpl_device
      cluster = self.cluster
      before_tpl_device = cluster.WorkerDeviceInModelSplit(0)
      # 遍历集群的其他设备，分配给cell_devices
      cell_devices = [
          cluster.WorkerDeviceInModelSplit(i) for i in range(num_cells)
      ]
​
    # 遍历 _before_layers，把其中每个变量部署在 before_tpl_device
    for unused_name, l in self._before_layers:
      with tf.device(before_tpl_device):
        l.InstantiateVariables()
​
    # 遍历 _cells，把其中每个变量部署在 cell_devices
    for i, (unused_name, l) in enumerate(self._cells):
      with tf.device(cell_devices[i]):
        l.InstantiateVariables()
​
    super()._CreateChildrenVariables()

6.4.1.3 FProp

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

• Пройдите _before_layers и вызовите его FProp для каждого слоя;
• перебирать _cells для каждого слоя вcluster.WorkerDeviceInModelSplit(i)вызовите его FProp выше;

  def FProp(self, theta, *args):
    """Round-robin every children cells in cell_tpl among worker devices.
​
    Args:
      theta: A NestedMap object containing weights' values of this layer and its
        children layers.
      *args: Input args
​
    Returns:
      A list contains one tensor of [batch_size, feature_height, feature_width,
        channel].
    """
    num_layers = len(self.params.cell_tpl)
    cluster = self.cluster
​
    # 遍历 _before_layers，对于其中每层调用其FProp  
    for (name, l) in self._before_layers:
      l_theta = theta[name]
      args = _ToTuple(args)
      args = l.FProp(l_theta, *args)
    # 遍历 _cells，对于其中每层，在`cluster.WorkerDeviceInModelSplit(i)`之上调用其FProp  
    for i in range(num_layers):
      with tf.device(cluster.WorkerDeviceInModelSplit(i)):
        cell_name, cell = self._cells[i]
        args = _ToTuple(args)
        args = cell.FProp(theta[cell_name], *args)
​
    return args

6.4.1.4 Конкретная реализация

Весь код SeqLayer выглядит следующим образом:

class SeqLayer(base_layer.BaseLayer):
  """Round-robin every children cells in cell_tpl among worker devices."""
​
  @classmethod
  def Params(cls):
    p = super().Params()
    p.Define('before_tpl', [],
             'Config for the CNN layers that runs before pipelining.')
    p.Define('cell_tpl', [], 'A list of FeatureExtractionLayer layers.')
    return p
​
  def __init__(self, params):
    super().__init__(params)
    p = self.params
    self._before_layers = []
    self._cells = []
    for l in p.before_tpl:
      self.CreateChild(l.name, l)
      self._before_layers.append((l.name, self.children[l.name]))
    for l in p.cell_tpl:
      self.CreateChild(l.name, l)
      self._cells.append((l.name, self.children[l.name]))
​
  def _CreateChildrenVariables(self):
    p = self.params
​
    num_cells = len(p.cell_tpl)
    before_tpl_device = ''
    cell_devices = [''] * num_cells
    if py_utils.use_tpu():
      cluster = self.cluster
      before_tpl_device = cluster.WorkerDeviceInModelSplit(0)
      cell_devices = [
          cluster.WorkerDeviceInModelSplit(i) for i in range(num_cells)
      ]
​
    for unused_name, l in self._before_layers:
      with tf.device(before_tpl_device):
        l.InstantiateVariables()
​
    for i, (unused_name, l) in enumerate(self._cells):
      with tf.device(cell_devices[i]):
        l.InstantiateVariables()
​
    super()._CreateChildrenVariables()
​
  def FProp(self, theta, *args):
    """Round-robin every children cells in cell_tpl among worker devices.
​
    Args:
      theta: A NestedMap object containing weights' values of this layer and its
        children layers.
      *args: Input args
​
    Returns:
      A list contains one tensor of [batch_size, feature_height, feature_width,
        channel].
    """
    num_layers = len(self.params.cell_tpl)
    cluster = self.cluster
​
    for (name, l) in self._before_layers:
      l_theta = theta[name]
      args = _ToTuple(args)
      args = l.FProp(l_theta, *args)
    for i in range(num_layers):
      with tf.device(cluster.WorkerDeviceInModelSplit(i)):
        cell_name, cell = self._cells[i]
        args = _ToTuple(args)
        args = cell.FProp(theta[cell_name], *args)
​
    return args

6.4.2 Конкретное распределение PipeliningLayer

PipeliningLayer — производный класс от SeqLayer.

• В верхней части конвейера находится устройство [0], которое обрабатывает предварительные условия.
• В середине конвейера находится ряд устройств, отвечающих за обработку определенных микропакетов.
• Концом конвейера является устройство [-1], которое отвечает за сортировку формы и, наконец, выводит окончательный тензор.

6.4.2.1 Получение выходной формы промежуточного слоя

_CalculateOutputShapes вычисляет выходную форму среднего слоя. Конкретная логика выглядит следующим образом:

• Пройдите _before_layers, вызовите его FPropMeta для каждого слоя, получите выходные фигуры и вставьте их в массив state_shapes;
• Обходим _cells, вызываем его FPropMeta для каждого слоя, получаем выходные фигуры и вставляем их в массив state_shapes;

  def _CalculateOutputShapes(self, input_shapes):
    """Calcuate the output shape of intermediate layers.
​
    Given the FPropMeta function in each FeatureExtractionLayer, calcuates
    the shapes of outputs of that layer. This is used to recover the shape
    information in StackedRecurrent.
​
    Args:
      input_shapes: NestedMap or tuple of input TensorShapes.
​
    Returns:
      Return a list of K + 1 NestedMaps or lists of tShape where K is
      the number of partitions.
    """
    p = self.params
    shapes = []
​
    # Converts TensorShape to tshape.Shape.
    def _ToTShape(x):
      if x is None:
        return None
      return tshape.Shape(x.as_list())
​
    shapes = py_utils.Transform(_ToTShape, input_shapes)
    shapes = _ToTuple(shapes)
​
    state_shapes = []
    # 遍历_before_layers，对其中每层调用其FPropMeta，得到 output shapes，插入 state_shapes 数组之中
    for (_, cell) in self._before_layers:
      shapes = cell.FPropMeta(cell.params, *shapes).out_shapes
​
    state_shapes.append(shapes[0] if p.nested_map_fprop else shapes)
​
    # 遍历 _cells，对其中每层调用其FPropMeta，得到 output shapes，插入 state_shapes 数组之中
    for (_, cell) in self._cells:
      shapes = cell.FPropMeta(cell.params, *shapes).out_shapes
      state_shapes.append(shapes[0] if p.nested_map_fprop else shapes)
​
    return state_shapes

6.4.2.2 Получить тип данных

Роль _get_state_dtype — получить тип данных.

  def _get_state_dtype(self, *args):
    if self.params.state_dtype:
      return self.params.state_dtype
    if self.params.nested_map_fprop:
      inputs = args[0].Filter(lambda x: x is not None)
      return py_utils.Flatten(inputs)[0].dtype
    return args[0].dtype

6.4.2.3 Получение формы ввода

Gpipe сначала объединит мини-партию обучающих выборок (mini-batch) на более мелкие мини-партии (micro-batches), затем направьте каждый набор мини-пакетов на устройство.

Функция _get_input_shapes состоит в том, чтобы получить входные формы.Конкретная логика выглядит следующим образом:

• получить входные данные input_tensors из args;

• Обходим input_tensors, находим первый непустой тензор, получаем размер пакета этого тензора и присваиваем его mini_batch_size;

• Получите micro_batch_size из параметра и установите для него значение micro_batch_size;

• Если micro_batch_size не имеет смысла, то:

  • Если p.num_micro_batches больше, чем mini_batch_size, то p.num_micro_batches равно mini_batch_size;
  • установить для micro_batch_size значение mini_batch_size // p.num_micro_batches;

• Создайте набор input_shapes, пройдите input_tensors для каждого тензора, получите его список форм input_shape и установите для batch_dim input_shape значение micro_batch_size;

• Если установлен p.nested_map_fprop, то сконструируйте input_shapes в рекурсивно вложенную структуру;

• вернуть input_shapes;

  def _get_input_shapes(self, *args):
    p = self.params
    if p.nested_map_fprop:
      assert len(args) == 1
      assert isinstance(args[0], py_utils.NestedMap)
      input_tensors = py_utils.Flatten(args[0])
    else:
      input_tensors = _ToTuple(args)
    
    # 遍历 input_tensors，找出第一个不为空的张量，获取这个张量的 batch size，赋给 mini_batch_size
    # Get batch size from the first tensor which is not None.
    mini_batch_size = None
    for input_tensor in input_tensors:
      if input_tensor is not None:
        mini_batch_size = input_tensor.get_shape().as_list()[p.batch_dim]
    assert mini_batch_size is not None
    micro_batch_size = p.micro_batch_size
    
    if not micro_batch_size: # 如果 micro_batch_size 没有意义
      # 如果 p.num_micro_batches 大于 mini_batch_size，则 p.num_micro_batches 为 mini_batch_size
      if p.num_micro_batches > mini_batch_size:
        p.num_micro_batches = mini_batch_size
      # 把 micro_batch_size 设置为 mini_batch_size // p.num_micro_batches  
      micro_batch_size = mini_batch_size // p.num_micro_batches
    if mini_batch_size is not None:
      if micro_batch_size * p.num_micro_batches != mini_batch_size:
        raise ValueError('micro_batch_size * num_micro_batches != batch_size.')
​
    # 遍历 input_tensors，对于每个张量，得到其shapes列表 input_shape，并且设置 input_shape 的 batch_dim 为 micro_batch_size
    input_shapes = ()
    for input_tensor in input_tensors:
      if input_tensor is not None:
        input_shape = input_tensor.get_shape().as_list()
        input_shape[p.batch_dim] = micro_batch_size
        input_shapes += (tf.TensorShape(input_shape),)
      else:
        input_shapes += (None,)
​
    # 如果设置了 p.nested_map_fprop，则把 input_shapes 构建成一个递归嵌套的结构    
    if p.nested_map_fprop:
      input_shapes = py_utils.Pack(args[0], input_shapes)
    return input_shapes

6.4.2.4 FProp

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

• Выполните некоторые подготовительные работы, такие как:

  • Получить количество дочерних ячеек;
  • получить кластер;
  • получить входные формы, dtypes;
  • Используйте входные формы для расчета выходных форм;

• Проходим по среднему слою обработки:

  • Для конкретной ячейки добавляем ячейку в слой накопления, и для каждой ячейки получаем соответствующую функцию;
  • Устанавливает его начальное состояние для последующих запусков StackedRecurrent;
  • Добавьте устройство, соответствующее cell_idx, в список устройств;

• Установите некоторые переменные для каждого устройства в середине конвейера (уберите голову и хвост);

• На первом устройстве выполните следующие действия:

  • Пройдите _before_layers, запустите FProp каждого слоя и, наконец, получите предыдущий;
  • Для предыдущей операции создайте входные данные, то есть используйте _StackAndSplit для разделения тензора;
  • Установить его вход для последующего устройства трубопровода;

• Выполнить операцию recurrent.StackedRecurrent на промежуточном устройстве;

• Агрегируйте форму micro_batches на последнем устройстве, чтобы получить окончательный выходной тензор:

  • Если вложено, вернуть последнюю форму;

  • В противном случае пройдитесь по выходу и агрегируйте форму каждого выхода;

  def FProp(self, theta, *args):
    """Run multiple cells in different devices in a pipelining manner.

    Args:
      theta: A NestedMap object containing weights' values of this layer and its
        children layers.
      *args: Non-keyworded variable length argument list of input tensors.

    Returns:
      A list of output tensors
    """
    # TODO(huangyp): handle optional None inputs.
    p = self.params
    if self.do_eval and self.cluster.num_devices_per_split == 1: # 如果设置了 do_eval 并且集群的 num_devices_per_split 为 1
      outputs = copy.copy(args)
      # 就直接串行执行
      for (name, l) in self._before_layers + self._cells:
        outputs = _ToTuple(outputs)
        outputs = l.FProp(theta[name], *outputs)
      return outputs

    num_cells = len(p.cell_tpl) # 得到 children cell个数
    cluster = self.cluster # 得到集群

    # Compute shapes of input and output tensors.
    # 得到 输入shapes，dtypes
    input_shapes = self._get_input_shapes(*args)
    state_dtype = self._get_state_dtype(*args)
    # 利用 输入shapes 计算出 输出shapes
    state_shapes = self._CalculateOutputShapes(input_shapes)
    tf.logging.info('state_shapes={}'.format(state_shapes))

    def GetCellFn(i): # 对于第 i 个层，返回一个对应的函数，这个函数将在 StackedRecurrent 内部执行
      """Get the ith feature extraction layer."""

      def CellFn(theta, state0, inputs):
        """A cell fn is exectued inside of StackedRecurrent."""
        # 没有深入研究StackedRecurrent，只从此函数看，作用是利用cell.FProp计算输出，并且得到一个state，其中包括输出和micro batch tensor
        del state0

        def _FPropInputSetShape(name, t_shape): # 给输入设置shape
          if t_shape is None:
            return None
          inputs[name].set_shape(t_shape.ToTensorShape().as_list())
          return inputs[name]

        if p.nested_map_fprop:
          # pylint: disable=protected-access
          fprop_inputs = state_shapes[i]._RecursiveMap(_FPropInputSetShape)
          # pylint: enable=protected-access
        else:
          fprop_inputs = []
          for input_idx, input_shape in enumerate(state_shapes[i]):
            name = 's{}'.format(input_idx)
            fprop_inputs.append(_FPropInputSetShape(name, input_shape))

        with py_utils.RemoveAssertContext(remove=True):
          with CellFnFPropOpReplacementWrapper():
            tf.logging.info('cell {} input {}'.format(i, fprop_inputs))
            mb_tensor = inputs[_MICRO_BATCH_STATE_NAME] # 得到输入的 micro batch tensor
            SetOverWriteGlobalStep(mb_tensor)
            _, cell = self._cells[i]
            fprop_inputs = _ToTuple(fprop_inputs)
            outputs = cell.FProp(theta, *fprop_inputs) # 计算输出

        if p.nested_map_fprop:
          assert py_utils.IsCompatible(outputs, state_shapes[i + 1])
          state1 = outputs.Filter(lambda x: x is not None)
        else:
          state1 = py_utils.NestedMap()
          outputs = _ToTuple(outputs)
          assert len(outputs) == len(state_shapes[i + 1])
          for output_idx in range(len(outputs)):
            if outputs[output_idx] is not None:
              name = 's{}'.format(output_idx)
              state1[name] = outputs[output_idx]
        state1[_MICRO_BATCH_STATE_NAME] = mb_tensor
        return state1, py_utils.NestedMap()

      return CellFn

    cell_fns = []
    accumulator_layers = [] # 为了梯度累积
    thetas = []
    init_states = []
    devices = []
    # 遍历，把cell_idx对应的设备加入到devices列表
    for cell_idx in range(num_cells): # 遍历 children cell
      cell_name, cell = self._cells[cell_idx] # 得到具体一个 cell
      accumulator_layers.append(cell) # 把cell加入到累积层中
      cell_fns.append(GetCellFn(cell_idx)) # 对于每个cell，得到对应的function
      thetas.append(theta[cell_name]) # 添加 theta

      # 返回一个带有形状t_shape的，类型为state_dtype的张量,并且所有元素都设为零.
      def _TfZeros(t_shape):
        if t_shape is None:
          return None
        return tf.zeros(t_shape.ToTensorShape().as_list(), dtype=state_dtype)

      # 为后续的 StackedRecurrent 运行设置其初始状态
      if p.nested_map_fprop:
        init_state = py_utils.Transform(_TfZeros, state_shapes[cell_idx + 1])
        init_state = init_state.Filter(lambda x: x is not None)
      else:
        init_state = py_utils.NestedMap()
        for output_idx, state in enumerate(state_shapes[cell_idx + 1]):
          state = _TfZeros(state)
          if state is not None:
            name = 's{}'.format(output_idx)
            init_state[name] = state
      init_state[_MICRO_BATCH_STATE_NAME] = tf.cast(0, dtype=state_dtype)
      init_states.append(init_state)

      # 把cell_idx对应的设备加入到devices列表
      devices.append(cluster.WorkerDeviceInModelSplit(cell_idx))

    # 为流水线中间（去除头尾）的各个设备设定一些变量
    cell_grads = [None] * num_cells
    cell_outs = [lambda x: x] * num_cells
    cell_out_grads = [lambda x: x] * num_cells

    # 在第一个设备上执行如下操作
    with tf.device(devices[0]): 
      previous = _ToTuple(args)
      for (name, l) in self._before_layers: # 遍历_before_layers，运行每层的FProp，最终得到 previous
        previous = l.FProp(theta[name], *previous)
        previous = _ToTuple(previous)

      def _StackAndSplit(x): # 把张量分割成
        # Split tensors into microbatches.
        if x is None:
          return None
        # tf.split按照行或者列分割一个矩阵
        return tf.stack(tf.split(x, p.num_micro_batches, axis=p.batch_dim))

      # 对于 previous 继续操作，构建出 inputs，即利用_StackAndSplit分割张量
      if p.nested_map_fprop: # 嵌套情况，只选取previous[0]做处理
        inputs = py_utils.Transform(_StackAndSplit, previous[0]) #利用_StackAndSplit分割张量
        inputs = inputs.Filter(lambda x: x is not None)
      else: # 非嵌套
        inputs = py_utils.NestedMap()
        for output_idx, output_tensor in enumerate(previous): # 遍历第一层的输出
          output_tensor = _StackAndSplit(output_tensor) # 利用_StackAndSplit分割张量
          if output_tensor is not None:
            name = 's{}'.format(output_idx)
            inputs[name] = output_tensor
      gs_tensor = py_utils.GetGlobalStep()
      # 为流水线后续设备设置其输入
      inputs[_MICRO_BATCH_STATE_NAME] = tf.stack([
          tf.cast(gs_tensor * p.num_micro_batches + t, dtype=state_dtype)
          for t in range(p.num_micro_batches)
      ])
      
    # 在中间设备上执行操作    
    tf.logging.info('pipeline input = {}'.format(inputs))
    output_state, _ = recurrent.StackedRecurrent( 
        devices=devices,
        cell_fns=cell_fns,
        cell_grads=cell_grads,
        cell_outs=cell_outs,
        cell_out_grads=cell_out_grads,
        thetas=thetas,
        init_states=init_states,
        inputs=inputs,
        accumulator_layers=accumulator_layers,
        unused_acc_state=True)

    # 在最后一个设备上执行如下操作，最终得到输出张量
    with tf.device(devices[-1]):
      def _ReshapeRetVal(name, t_shape): # 把micro_batches的形状聚合，得到最终输出
        """Restore shape for tensors in microbatches."""
        if t_shape is None:
          return None
        output_tensor = output_state[name]
        if p.batch_dim != 0:
          perm = list(range(1, p.batch_dim + 1)) + [0]
          perm += list(range(p.batch_dim + 1, t_shape.rank + 1))
          output_tensor = tf.transpose(output_tensor, perm=perm)
        output_shape = t_shape.ToTensorShape().as_list()
        output_shape[p.batch_dim] *= p.num_micro_batches
        output_tensor = tf.reshape(output_tensor, output_shape)
        return output_tensor

      # Construct the final return values from output_state.
      if p.nested_map_fprop: # 如果嵌套，则返回最后一个形状
        # pylint: disable=protected-access
        output_tensors = state_shapes[-1]._RecursiveMap(_ReshapeRetVal) # 聚合形状
        # pylint: enable=protected-access
      else:
        output_tensors = []
        # 遍历输出，聚合各个输出的形状
        for output_idx, state_shape in enumerate(state_shapes[-1]): 
          output_name = 's{}'.format(output_idx)
          output_tensor = _ReshapeRetVal(output_name, state_shape) # 聚合形状
          output_tensors.append(output_tensor)
        if len(output_tensors) == 1:
          output_tensors = output_tensors[0]
        else:
          output_tensors = tuple(output_tensors)
        
      tf.logging.info('pipeline output = {}'.format(output_tensors))
      return output_tensors

6.4.2.5 Определение класса

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

class PipeliningLayer(SeqLayer):
  """Pipelining a sequence of layers on multiple devices."""
​
  @classmethod
  def Params(cls):
    p = super().Params()
    p.Define('num_micro_batches', 1, 'Number of micro batches.')
    p.Define('micro_batch_size', None, 'Size of a micro batch.')
    p.Define('batch_dim', 0, 'The batch dimension.')
    p.Define('state_dtype', None, 'Externally specify dtype for states.')
    p.Define(
        'nested_map_fprop', False, 'Whether arguments and returns of '
        'cell fprop functions are nested maps')
    return p

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

+--------------------------------------------------------------+
| FProp             _CalculateOutputShapes                     |
|                             +                                |
|                             |                                |
|                             |                                |
|                             v                                |
|                        state_shapes                          |
|                             +                                |
|                             |                                |
|                             |                                |
|                             |                                |
|                             v                                |
|                for cell_idx in range(num_cells):             |
|                             +                                |
|                             |                                |
|                             |                                |
|                             v                                |
|       devices.append(WorkerDeviceInModelSplit(cell_idx))     |
|                             +                                |
|                             |                                |
|                             |                                |
|                             v                                |
|                  with tf.device(devices[0])                  |
|                             +                                |
|                             |                                |
|                             |                                |
|                             v                                |
|             recurrent.StackedRecurrent(cell_outs)            |
|                             +                                |
|                             |                                |
|                             |                                |
|                             v                                |
|                 with tf.device(devices[-1])                  |
|                             +                                |
|                             |                                |
|                             |                                |
|                             v                                |
|                       output_tensors                         |
|                                                              |
+--------------------------------------------------------------+

Логика конвейера устройства следующая:

                   devices[0]
                       +
                       |
                       |
                       |
                       v
+----------------------+-------------------------+
|Pipeline                                        |
|                         devices[1]             |
|                             +                  |
|                             |                  |
|                             |                  |
|                             v                  |
|  cell_grads[1~n]        devices[2]             |
|                             +                  |
|  cell_outs[1~n]             |                  |
|                             |                  |
|  cell_out_grads[1~n]        v                  |
|                         devices[3]             |
|                             +                  |
|                             |                  |
|                             |                  |
|                             v                  |
|                         devices[4]             |
|                                                |
+----------------------+-------------------------+
                       |
                       |
                       |
                       v
                   devices[-1]

6.4.2.6 Использование

Пример, приведенный в исходном коде, — GPipeBatchMajorTransformerStack, В настоящее время кажется, что достаточно наследовать PipeliningLayer.

class GPipeBatchMajorTransformerStack(PipeliningLayer):
  """Stacked self- multi-head attention and fully connected layers.
​
  With optional layer normalization applied to the final output.
​
  See 'Attention Is All You Need' https://arxiv.org/abs/1706.03762
  for details. 
​
  Implements a gipe stack for the batch major transformer variant.
  """

FProp из GPipeBatchMajorTransformerStack возвращает список выходных тензоров, где следующий код вызывает функцию PipeliningLayer.

logits = super().FProp(theta, source_input, source_paddings, target_input,
                       target_paddings, encoder_self_atten_segment_mask,
                       decoder_self_atten_segment_mask,
                       decoder_cross_atten_segment_mask, source_segment_pos,
                       target_segment_pos)

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

  def FProp(self,
            theta,
            source_input,
            source_paddings,
            target_input=None,
            target_paddings=None,
            source_segment_id=None,
            target_segment_id=None,
            labels=None,
            label_weights=None,
            source_segment_pos=None,
            target_segment_pos=None):
​
    p = self.params
    if p.num_decoder_layers > 0:
      assert target_input is not None
      assert target_paddings is not None
      target_time = tf.shape(target_input)[1]
      batch = tf.shape(target_input)[0]
    encoder_self_atten_segment_mask = None
    decoder_self_atten_segment_mask = None
    decoder_cross_atten_segment_mask = None
​
    # Prepare segment masks from segment ids.
    if p.packed_input:
      dtype = py_utils.FPropDtype(p)
      assert source_segment_id is not None, (
          'Need to specify src_segment_id if packed input is supported.')
      assert source_segment_pos is not None, (
          'Need to specify src_segment_pos for packed input and embeddings.')
      encoder_self_atten_segment_mask = batch_major_attention.SegmentMask(
          source_segment_id, source_segment_id, dtype, False)
      if target_segment_id is not None:
        decoder_self_atten_segment_mask = batch_major_attention.SegmentMask(
            target_segment_id, target_segment_id, dtype, False)
        causal_padding = tf.expand_dims(
            tf.tile(
                tf.expand_dims(
                    batch_major_attention.CausalPadding(
                        target_time, dtype=dtype), 0), [batch, 1, 1]), 1)
        decoder_self_atten_segment_mask = tf.math.maximum(
            causal_padding, decoder_self_atten_segment_mask)
        decoder_cross_atten_segment_mask = batch_major_attention.SegmentMask(
            target_segment_id, source_segment_id, dtype, False)
​
    # FProp through the gpipe pipeline.
    # 这里调用了基类的PipeliningLayer，完成流水线操作。
    logits = super().FProp(theta, source_input, source_paddings, target_input,
                           target_paddings, encoder_self_atten_segment_mask,
                           decoder_self_atten_segment_mask,
                           decoder_cross_atten_segment_mask, source_segment_pos,
                           target_segment_pos)
            
    label_weights = tf.reshape(label_weights, [-1])
    target_probs = None
    if p.label_smoothing:
      target_probs = self.smoother.FProp(
          theta.smoother, target_paddings, labels, target_ids=None)
      target_probs = tf.reshape(target_probs, [-1, p.softmax_tpl.num_classes])
    reshaped_logits = tf.reshape(logits, [-1, p.softmax_tpl.num_classes])
    tgt_labels = tf.reshape(labels, [-1])
    num_splits = len(p.splits)
    softmax = self.children['cell_{}'.format(num_splits - 1)].softmax
    softmax_theta = theta['cell_{}'.format(num_splits - 1)].softmax
    per_example_xent, _ = softmax.XentLossFromLogits(
        softmax_theta,
        reshaped_logits,
        class_weights=tf.reshape(label_weights, [-1]),
        class_ids=tgt_labels,
        class_probabilities=target_probs)
    xent_shape = tf.shape(logits)[:2]
    per_example_xent = tf.reshape(per_example_xent, xent_shape)
    return per_example_xent, logits

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

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

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

ссылка 0xFF

Одна только эта статья позволит вам освоить системный дизайн платформы OneFlow (Часть 1)

DeepSpeed: Extreme-scale model training for everyone

DeepSpeed: инструмент обучения гипермасштабируемой модели для всех

Почему исходная структура глубокого обучения не может быть перегружена обучением GPT-3?

Почему модель GPT-3 трудно воспроизвести? Это может быть оптимальным дизайном распределенной инфраструктуры ИИ.

FLOPs и скорость вывода модели

Количество параметров и вычисление FLOPS в глубоком обучении (на примере классической сетевой структуры AlexNet в CNN)

Какие флопы вычислительной мощности требуются для модели CNN? Как рассчитать?

Расчет вычислительных FLOP в CNN

Определение и расчет FLOPS

Тринадцатая часть серии интерпретаций статей: ZeRO — метод обучения модели для триллионов параметров

Лучшие практики параллелизма моделей (PyTorch)

Tensorflow: модель параллелизма, модель параллельных вычислений

py torch.org/tutorials/i…

модель pytorch параллельная модель параллельная

АР Вест V.org/PDF/1802.09…

Каков принцип модельного параллелизма в глубоком обучении?

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

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

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

Как оценить библиотеку параллельного ускорения Google с открытым исходным кодом GPipe?

Paper Interpretation Series 4: Google GPipe Training Очень крупномасштабные нейронные сети

Как уменьшить память нейросети?

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

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

Что делать, если памяти графического процессора недостаточно?

Обучение модели слишком медленное? Недостаточно памяти? Этот метод позволяет вашему графическому процессору объединиться с центральным процессором.

Сравнение трех библиотек TF-Replicator, GPipe, Mesh-Tensorflow

Резюме параллелизма данных в обучении глубокой нейронной сети

[Новое] Система рекомендаций глубокого обучения Facebook

py torch.org/tutorials/i…

GitHub.com/py факел/тотем…

Введение в распределенный TensorFlow

Зачем вручную обнулять градиенты перед обратным распространением в PyTorch?

Model.zero_grad() or optimizer.zero_grad()?

A trick to use bigger batches for training: gradient accumulation

Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups

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

Распределенное обучение от входа до отказа