Dicas e armadilhas de configuração da SageMaker Distributed Model Parallelism Library

A biblioteca é mais eficiente quando o tamanho do lote é aumentado. Para casos de uso em que o modelo cabe em um dispositivo único, mas pode ser treinado apenas com um tamanho de lote pequeno, o tamanho do lote pode e deve ser aumentado após a integração da biblioteca. O paralelismo de modelos economiza memória para modelos grandes, habilitando o treinamento usando tamanhos do lote que antes não cabiam na memória.

Escolher um número de microlotes muito pequenos ou muito grandes pode reduzir a performance. A biblioteca executa cada microlote sequencialmente em cada dispositivo, portanto, o tamanho do microlote (tamanho do lote dividido pelo número de microlotes) deve ser grande o suficiente para utilizar totalmente cada GPU. Ao mesmo tempo, a eficiência do pipeline aumenta com o número de microlotes, portanto, é importante encontrar o equilíbrio certo. Normalmente, um bom ponto de partida é experimentar 2 ou 4 microlotes, aumentando o tamanho do lote até o limite de memória e, em seguida, experimentar tamanhos do lote e números de microlotes maiores. À medida que o número de microlotes aumenta, tamanhos do lote maiores podem se tornar viáveis se um pipeline intercalado for usado.

O tamanho do lote deve ser sempre divisível pelo número de microlotes. Observe que, dependendo do tamanho do conjunto de dados, às vezes, o último lote de cada epoch pode ser menor que o resto e esse lote menor também precisa ser divisível pelo número de microlotes. Se não estiver, você pode definir drop_remainder=True a tf.Dataset.batch() chamada (in TensorFlow) ou definir drop_last=True in DataLoader (in PyTorch), para que esse último lote pequeno não seja usado. Se você estiver usando uma API diferente para o pipeline de dados, talvez seja necessário ignorar manualmente o último lote sempre que ele não for divisível pelo número de microlotes.

Se você usa o particionamento manual, esteja atento aos parâmetros que são consumidos por várias operações e módulos em seu modelo, como a tabela de incorporação nas arquiteturas do transformador. Módulos que compartilham o mesmo parâmetro devem ser colocados no mesmo dispositivo para que estejam corretos. Quando o particionamento automático é usado, a biblioteca aplica automaticamente essa restrição.

Se o modelo usar várias entradas, certifique-se de semear as operações aleatórias em seu data pipeline (por exemplo, embaralhar) com smp.dp_rank(). Se o conjunto de dados estiver sendo fragmentado de forma determinística em dispositivos em paralelo de dados, certifique-se de que o fragmento seja indexado por smp.dp_rank(). Isso irá garantir que a ordem dos dados vistos em todas as classificações que formam uma partição de modelo seja consistente.

Qualquer tensor retornado da função smp.DistributedModel.call (for TensorFlow) ou smp.DistributedModel.forward (for) é transmitido para PyTorch todas as outras classificações, a partir da classificação que calculou esse tensor específico. Como resultado, qualquer tensor que não é necessário fora dos métodos de chamada e encaminhamento (ativações intermediárias, por exemplo) não deve ser retornado, pois isso causa comunicação desnecessária e sobrecarga da memória e prejudica a performance.

Se uma função decorada por smp.step tiver um argumento de tensor que não tenha uma dimensão de lote, o nome do argumento deverá ser fornecido na lista non_split_inputs durante a chamada de smp.step. Isso evita que a biblioteca tente dividir o tensor em microlotes. Para obter mais informações, consulte smp.step na documentação de API.

Se você estiver usando uma semente para resultados determinísticos, defina a semente baseada em smp.dp_rank() (por exemplo, torch.manual_seed(42 + smp.dp_rank())). Se você não fizer isso, partições diferentes de um nn.Parameter serão inicializadas da mesma forma, afetando a convergência.

SageMakerA biblioteca de paralelismo de modelos usa NCCL para implementar os coletivos necessários para a distribuição dos módulos. Especialmente para modelos menores, se muitas chamadas de NCCL forem programadas na GPU ao mesmo tempo, o uso de memória poderá aumentar devido ao espaço adicional usado pela NCCL. Para neutralizar isso, smp controla as chamadas de NCCL para que, a qualquer momento, o número de operações contínuas de NCCL seja menor ou igual a um determinado limite. O limite padrão é 8, mas isso pode ser ajustado usando a variável de ambiente SMP_NCCL_THROTTLE_LIMIT. Se você observar o uso de memória maior do que o esperado ao usar o paralelismo de tensores, tente reduzir esse limite. No entanto, escolher um limite muito pequeno pode causar perda de taxa de transferência. Para desativar completamente o controle de utilização, você pode definir SMP_NCCL_THROTTLE_LIMIT=-1.

A seguinte identidade, que é válida quando o grau de paralelismo de tensores é 1, não é válida quando o grau de paralelismo de tensores é maior que 1: smp.mp_size() * smp.dp_size() == smp.size(). Isso ocorre porque o grupo em paralelo de tensores faz parte do grupo de paralelismo do modelo e do grupo de paralelismo de dados. Se seu código tiver referências existentes a mp_rank, mp_size, MP_GROUP e assim por diante, e se você quiser trabalhar apenas com o grupo em paralelo do pipeline, talvez seja necessário substituir as referências por smp.pp_size(). As seguintes identidades são sempre verdadeiras:

Uma vez que o wrapper do smp.DistributedModel modifica os parâmetros do modelo quando o paralelismo de tensores está ativado, o otimizador deve ser criado após a chamada de smp.DistributedModel, com os parâmetros distribuídos. Por exemplo, o seguinte não funciona:

## WRONG
model = MyModel()
optimizer = SomeOptimizer(model.parameters())
model = smp.DistributedModel(model)  # optimizer now has outdated parameters! 

Em vez disso, o otimizador deve ser criado com os seguintes parâmetros de smp.DistributedModel:

## CORRECT
model = smp.DistributedModel(MyModel())
optimizer = SomeOptimizer(model.optimizers())

Quando um módulo é substituído por sua contraparte distribuída por meio de paralelismo de tensores, o módulo distribuído não herda seus pesos do módulo original e inicializa novos pesos. Isso significa que, por exemplo, se os pesos precisarem ser inicializados em uma chamada específica (por exemplo, por meio de uma chamada de load_state_dict), isso precisará acontecer após a chamada de smp.DistributedModel, quando a distribuição do módulo ocorrer.

Ao acessar diretamente os parâmetros dos módulos distribuídos, observe que o peso não tem o mesmo formato do módulo original. Por exemplo: 

with smp.tensor_parallelism():
    linear = nn.Linear(60, 60)

# will pass
assert tuple(linear.weight.shape) == (60, 60)

distributed_linear = smp.DistributedModel(linear)

# will fail. the number of input channels will have been divided by smp.tp_size()
assert tuple(distributed_linear.module.weight.shape) == (60, 60)

O uso de torch.utils.data.distributed.DistributedSampler é altamente recomendado para paralelismo de tensores. Isso garante que cada classificação em paralelo de dados receba o mesmo número de amostras de dados, o que evita interrupções que possam resultar de diferentes dp_ranks realizando um número de etapas diferentes.

Se você usar a join API da DistributedDataParallel classe PyTorch's para lidar com casos em que diferentes classificações paralelas de dados têm números diferentes de lotes, você ainda precisa garantir que as classificações que estão na mesma TP_GROUP tenham o mesmo número de lotes; caso contrário, os coletivos de comunicação usados na execução distribuída de módulos podem travar. Classificações que estão em diferentes TP_GROUPs podem ter diferentes números de lotes, desde que a API do join seja usada.

Se você quiser que seu modelo tenha um ponto de verificação e usar o paralelismo de tensores, considere o seguinte:

Para obter mais informações sobre o ponto de verificação de um modelo com paralelismo de tensores, consulte Pontos de verificação de um modelo distribuído.