Apache MXNet - Treinamento Distribuído

Este capítulo é sobre o treinamento distribuído no Apache MXNet. Vamos começar entendendo quais são os modos de computação no MXNet.

Modos de Computação

MXNet, uma biblioteca de ML multilíngue, oferece aos seus usuários os seguintes dois modos de computação -

Modo imperativo

Este modo de computação expõe uma interface como a API NumPy. Por exemplo, no MXNet, use o seguinte código imperativo para construir um tensor de zeros na CPU e também na GPU -

import mxnet as mx
tensor_cpu = mx.nd.zeros((100,), ctx=mx.cpu())
tensor_gpu= mx.nd.zeros((100,), ctx=mx.gpu(0))

Como vemos no código acima, MXNets especifica o local onde manter o tensor, seja na CPU ou no dispositivo GPU. No exemplo acima, ele está no local 0. O MXNet consegue uma utilização incrível do dispositivo, porque todos os cálculos acontecem preguiçosamente em vez de instantaneamente.

Modo simbólico

Embora o modo imperativo seja bastante útil, uma das desvantagens deste modo é sua rigidez, ou seja, todos os cálculos precisam ser conhecidos de antemão junto com as estruturas de dados predefinidas.

Por outro lado, o modo Simbólico expõe um gráfico de computação como o TensorFlow. Ele remove a desvantagem da API imperativa, permitindo que o MXNet trabalhe com símbolos ou variáveis ​​em vez de estruturas de dados fixas / predefinidas. Posteriormente, os símbolos podem ser interpretados como um conjunto de operações da seguinte forma -

import mxnet as mx
x = mx.sym.Variable(“X”)
y = mx.sym.Variable(“Y”)
z = (x+y)
m = z/100

Tipos de paralelismo

Apache MXNet oferece suporte a treinamento distribuído. Isso nos permite alavancar várias máquinas para um treinamento mais rápido e eficaz.

A seguir estão as duas maneiras pelas quais podemos distribuir a carga de trabalho de treinamento de um NN em vários dispositivos, CPU ou GPU -

Paralelismo de dados

Nesse tipo de paralelismo, cada dispositivo armazena uma cópia completa do modelo e trabalha com uma parte diferente do conjunto de dados. Os dispositivos também atualizam um modelo compartilhado coletivamente. Podemos localizar todos os dispositivos em uma única máquina ou em várias máquinas.

Paralelismo de modelo

É outro tipo de paralelismo, que é útil quando os modelos são tão grandes que não cabem na memória do dispositivo. No paralelismo do modelo, diferentes dispositivos são atribuídos a tarefa de aprender diferentes partes do modelo. O ponto importante a ser observado aqui é que atualmente o Apache MXNet oferece suporte ao paralelismo de modelo em uma única máquina apenas.

Trabalho de treinamento distribuído

Os conceitos fornecidos abaixo são a chave para entender o funcionamento do treinamento distribuído no Apache MXNet -

Tipos de processos

Os processos se comunicam entre si para realizar o treinamento de um modelo. Apache MXNet tem os três processos a seguir -

Trabalhador

O trabalho do nó trabalhador é realizar o treinamento em um lote de amostras de treinamento. Os nós de trabalho extrairão pesos do servidor antes de processar cada lote. Os nós de trabalho enviarão gradientes ao servidor, assim que o lote for processado.

Servidor

O MXNet pode ter vários servidores para armazenar os parâmetros do modelo e se comunicar com os nós de trabalho.

Agendador

A função do planejador é configurar o cluster, o que inclui esperar por mensagens de que cada nó foi ativado e em qual porta o nó está atendendo. Depois de configurar o cluster, o planejador permite que todos os processos saibam sobre todos os outros nós do cluster. É porque os processos podem se comunicar uns com os outros. Existe apenas um agendador.

Loja KV

Lojas KV significa Key-Valueloja. É um componente crítico usado para treinamento em vários dispositivos. É importante porque a comunicação de parâmetros entre dispositivos em uma única máquina e também em várias máquinas é transmitida por meio de um ou mais servidores com um KVStore para os parâmetros. Vamos entender o funcionamento do KVStore com a ajuda dos seguintes pontos -

  • Cada valor em KVStore é representado por um key e um value.

  • Cada array de parâmetros na rede é atribuído a um key e os pesos dessa matriz de parâmetros são referidos por value.

  • Depois disso, os nós de trabalho pushgradientes após o processamento de um lote. Eles tambémpull pesos atualizados antes de processar um novo lote.

A noção de servidor KVStore existe apenas durante o treinamento distribuído e o modo distribuído dele é habilitado chamando mxnet.kvstore.create função com um argumento string contendo a palavra dist -

kv = mxnet.kvstore.create(‘dist_sync’)

Distribuição de Chaves

Não é necessário que todos os servidores armazenem todas as matrizes ou chaves de parâmetros, mas eles são distribuídos em diferentes servidores. Essa distribuição de chaves em diferentes servidores é tratada de forma transparente pelo KVStore e a decisão de qual servidor armazena uma chave específica é feita aleatoriamente.

KVStore, como discutido acima, garante que sempre que a chave for puxada, sua solicitação seja enviada para aquele servidor, que possui o valor correspondente. E se o valor de alguma chave for grande? Nesse caso, ele pode ser compartilhado entre diferentes servidores.

Divida os dados de treinamento

Como sendo os usuários, queremos que cada máquina trabalhe em diferentes partes do conjunto de dados, especialmente, ao executar o treinamento distribuído no modo paralelo de dados. Sabemos que, para dividir um lote de amostras fornecidas pelo iterador de dados para treinamento paralelo de dados em um único trabalhador, podemos usarmxnet.gluon.utils.split_and_load e, em seguida, carregue cada parte do lote no dispositivo que irá processá-lo posteriormente.

Por outro lado, no caso de treinamento distribuído, no início precisamos dividir o conjunto de dados em npartes diferentes para que cada trabalhador receba uma parte diferente. Uma vez obtido, cada trabalhador pode usarsplit_and_loadpara dividir novamente essa parte do conjunto de dados entre diferentes dispositivos em uma única máquina. Tudo isso acontece por meio do iterador de dados.mxnet.io.MNISTIterator e mxnet.io.ImageRecordIter são dois desses iteradores no MXNet que oferecem suporte a esse recurso.

Atualização de pesos

Para atualizar os pesos, KVStore suporta os seguintes dois modos -

  • O primeiro método agrega os gradientes e atualiza os pesos usando esses gradientes.

  • No segundo método, o servidor agrega apenas gradientes.

Se você estiver usando o Gluon, há uma opção de escolher entre os métodos indicados acima, passando update_on_kvstorevariável. Vamos entender isso criando otrainer objeto da seguinte forma -

trainer = gluon.Trainer(net.collect_params(), optimizer='sgd',
   optimizer_params={'learning_rate': opt.lr,
      'wd': opt.wd,
      'momentum': opt.momentum,
      'multi_precision': True},
      kvstore=kv,
   update_on_kvstore=True)

Modos de treinamento distribuído

Se a string de criação do KVStore contiver a palavra dist, significa que o treinamento distribuído está ativado. A seguir estão diferentes modos de treinamento distribuído que podem ser ativados usando diferentes tipos de KVStore -

dist_sync

Como o nome indica, denota treinamento distribuído síncrono. Nesse caso, todos os trabalhadores usam o mesmo conjunto sincronizado de parâmetros do modelo no início de cada lote.

A desvantagem desse modo é que, após cada lote, o servidor deve esperar para receber gradientes de cada trabalhador antes de atualizar os parâmetros do modelo. Isso significa que se um trabalhador travar, o progresso de todos os trabalhadores será interrompido.

dist_async

Como o nome indica, denota treinamento distribuído síncrono. Neste, o servidor recebe gradientes de um trabalhador e atualiza imediatamente seu armazenamento. O servidor usa o armazenamento atualizado para responder a quaisquer solicitações adicionais.

A vantagem, em comparação com dist_sync mode, é que um trabalhador que conclui o processamento de um lote pode obter os parâmetros atuais do servidor e iniciar o próximo lote. O trabalhador pode fazer isso, mesmo que o outro trabalhador ainda não tenha concluído o processamento do lote anterior. Também é mais rápido que o modo dist_sync porque pode levar mais épocas para convergir sem nenhum custo de sincronização.

dist_sync_device

Este modo é o mesmo que dist_syncmodo. A única diferença é que, quando há várias GPUs sendo usadas em cada nódist_sync_device agrega gradientes e atualiza pesos na GPU, enquanto, dist_sync agrega gradientes e atualiza pesos na memória da CPU.

Ele reduz a comunicação cara entre GPU e CPU. É por isso que é mais rápido do quedist_sync. A desvantagem é que aumenta o uso de memória na GPU.

dist_async_device

Este modo funciona da mesma forma que dist_sync_device modo, mas no modo assíncrono.