As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.
Etapa 1: modifique seu próprio script de treinamento usando a biblioteca paralela de modelos distribuídos do SageMaker
Use esta seção para aprender a personalizar o seu script de treinamento para utilizar os recursos principais da biblioteca de paralelismo de modelos do Amazon SageMaker AI. Para usar as funções e os parâmetros de API específicos da biblioteca, recomendamos que você use essa documentação junto com as APIs da biblioteca paralela de modelos do SageMaker
Os exemplos de scripts de treinamento fornecidos nessas seções são simplificados e projetados para destacar as alterações necessárias que você deve fazer para usar a biblioteca. Para exemplos de notebook executáveis de ponta a ponta que demonstram como usar um script de treinamento TensorFlow ou PyTorch com a biblioteca de paralelismo de modelos do SageMaker, consulte Exemplos da biblioteca de paralelismo de modelos do Amazon SageMaker AI v2.
Tópicos
Divida o modelo do seu script de treinamento usando a biblioteca de paralelismo de modelos do SageMaker
Há duas maneiras de modificar seu script de treinamento para configurar a divisão do modelo: divisão automática ou divisão manual.
Divisão automatizada de modelos
Ao usar a biblioteca de paralelismo de modelos do SageMaker, você pode aproveitar a divisão automática de modelos, também conhecida como particionamento automatizado de modelos. A biblioteca utiliza um algoritmo de particionamento que equilibra a memória, minimiza a comunicação entre dispositivos e otimiza o desempenho. Você pode configurar o algoritmo de particionamento automático para otimizar a velocidade ou a memória.
Também é possível usar a divisão manual do modelo. Recomendamos a divisão automatizada do modelo, a menos que você esteja muito familiarizado com a arquitetura do modelo e tenha uma boa ideia de como particionar seu modelo de forma eficiente.
Como funciona
O particionamento automático ocorre durante a primeira etapa de treinamento, quando a função decorada com smp.step é chamada pela primeira vez. Nesta chamada, a biblioteca primeiro constrói uma versão do modelo na RAM da CPU (para evitar limitações de memória da GPU) e, em seguida, analisa o grafo do modelo e toma uma decisão de particionamento. Com base nessa decisão, cada partição do modelo é carregada em uma GPU e somente então a primeira etapa é executada. Devido a essas etapas de análise e particionamento, a primeira etapa de treinamento pode levar mais tempo.
Em qualquer framework, a biblioteca gerencia a comunicação entre dispositivos por meio de seu próprio back-end, que é otimizado para infraestrutura AWS.
O design de autopartição se adapta às características do framework, e a biblioteca realiza o particionamento no nível de granularidade que é mais natural em cada framework. Por exemplo, no TensorFlow, cada operação específica pode ser atribuída a um dispositivo diferente, enquanto no PyTorch, a atribuição é feita no nível do módulo, onde cada módulo consiste em várias operações. A seção a seguir analisa as especificidades do design em cada framework.
Durante a primeira etapa de treinamento, a biblioteca de paralelismo de modelo internamente executa uma etapa de rastreamento destinada a construir o grafo do modelo e determinar as formas dos tensores e parâmetros. Após essa etapa de rastreamento, a biblioteca constrói uma árvore, que consiste nos objetos nn.Module aninhados no modelo, bem como nos dados adicionais coletados do rastreamento, como a quantidade de objetos armazenados nn.Parameters e o tempo de execução de cada nn.Module.
Em seguida, a biblioteca percorre essa árvore a partir da raiz e executa um algoritmo de particionamento que atribui cada nn.Module a um dispositivo, equilibrando a carga computacional (medida pelo tempo de execução do módulo) e o uso de memória (medido pelo tamanho total armazenado de nn.Parameter e ativações). Se vários nn.Modules compartilham o mesmo nn.Parameter, então esses módulos são colocados no mesmo dispositivo para evitar manter várias versões do mesmo parâmetro. Assim que a decisão de particionamento é tomada, os módulos e pesos atribuídos são carregados em seus dispositivos.
Para obter instruções sobre como registrar o decorador smp.step em seu script de treinamento do PyTorch, consulte Divisão automatizada com PyTorch.
A biblioteca de paralelismo de modelos analisa os tamanhos das variáveis treináveis e a estrutura do gráfico e usa internamente um algoritmo de particionamento gráfico. Este algoritmo determina uma atribuição de dispositivo para cada operação, com o objetivo de minimizar a quantidade de comunicação necessária entre dispositivos, sujeito a duas restrições:
-
Equilibrando o número de variáveis armazenadas em cada dispositivo.
-
Equilibrando o número de operações executadas em cada dispositivo
Se você especificar speed para optimize (nos parâmetros de paralelismo do modelo no Python SDK), a biblioteca tentará equilibrar o número de operações e objetos tf.Variable em cada dispositivo. Caso contrário, ele tenta equilibrar o tamanho total de tf.Variables.
Uma vez tomada a decisão de particionamento, a biblioteca cria uma representação serializada do subgráfico que cada dispositivo precisa executar e os importa para cada dispositivo. Durante o particionamento, a biblioteca coloca as operações que consomem o mesmo tf.Variable e as operações que fazem parte da mesma camada Keras no mesmo dispositivo. Ele também respeita as restrições de colocalização impostas pelo TensorFlow. Isso significa que, por exemplo, se houver duas camadas Keras que compartilham um tf.Variable, todas as operações que fazem parte dessas camadas são colocadas em um dispositivo único.
Para obter instruções sobre como registrar o decorador smp.step em seu script de treinamento do PyTorch, consulte Divisão automatizada com o TensorFlow.
Comparação da divisão automatizada de modelos entre frameworks
No TensorFlow, a unidade fundamental de computação é um tf.Operation, e o TensorFlow representa o modelo como um grafo acíclico direcionado (DAG) de tf.Operations. Portanto, a biblioteca de paralelismo de modelo particiona esse DAG de modo que cada nó vai para um dispositivo. Crucialmente, os objetos tf.Operation são suficientemente ricos em atributos personalizáveis e são universais no sentido de que cada modelo tem a garantia de consistir em um gráfico desses objetos.
O PyTorch, por outro lado, não tem uma noção equivalente de operação que seja suficientemente rica e universal. A unidade mais próxima de computação no PyTorch que possui essas características é um nn.Module, que está em um nível de granularidade muito mais alto, e é por isso que a biblioteca faz a partição nesse nível no PyTorch.
Divisão manual de modelos
Se você quiser especificar manualmente como particionar seu modelo entre dispositivos, use o gerenciador de contexto do smp.partition. Para instruções sobre como configurar o gerenciador de contexto para particionamento manual, consulte as seguintes páginas.
Para utilizar essa opção após fazer modificações, na Etapa 2, será necessário definir auto_partition como False e criar um default_partition na classe de estimador do framework no SageMaker Python SDK. Qualquer operação que não seja colocada explicitamente em uma partição por meio do gerenciador de contexto do smp.partition é executada no default_partition. Nesse caso, a lógica de divisão automatizada é ignorada e cada operação é colocada com base na sua especificação. Com base na estrutura gráfica resultante, a biblioteca de paralelismo de modelos cria automaticamente um cronograma de execução em pipeline.
