Sincronização de gradientes

O módulo distribuído do PyTorch opera se comunicando de ida e volta entre todas as GPUs do seu sistema. Essa comunicação leva tempo, e garantir que todos os processos conheçam os estados uns dos outros ocorre em pontos de gatilho específicos ao usar o módulo ddp.

Esses pontos de gatilho são adicionados ao modelo do PyTorch, especificamente aos seus métodos forward() e backward(). Isso ocorre quando o modelo é encapsulado com DistributedDataParallel:

import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel

model = nn.Linear(10, 10)
ddp_model = DistributedDataParallel(model)

No Accelerate, essa conversão ocorre automaticamente ao chamar prepare() e passar seu modelo.

+ from accelerate import Accelerator
+ accelerator = Accelerator()
  import torch.nn as nn
- from torch.nn.parallel import DistributedDataParallel

  model = nn.Linear(10,10)
+ model = accelerator.prepare(model)

O atraso na acumulação de gradientes

Agora você entende que o PyTorch adiciona hooks ao método forward e ao método backward do seu modelo PyTorch durante o treinamento em um ambiente distribuído. Mas como isso pode arriscar deixar seu código mais lento?

No DDP (paralelismo de dados distribuído), a ordem específica em que os processos são executados é esperada em pontos específicos e esses processos também devem ocorrer aproximadamente ao mesmo tempo antes de prosseguir.

O exemplo mais direto é quando você atualiza os parâmetros do modelo por meio de optimizer.step(). Sem acúmulo de gradientes, todas as instâncias do modelo precisam ter seus gradientes calculados, agregados e atualizados antes de avançar para o próximo lote de dados. Ao realizar o acúmulo de gradientes, você acumula n gradientes de perda e pula optimizer.step() até que n lotes tenham sido atingidos. Como todos os processos de treinamento precisam se sincronizar apenas no momento em que optimizer.step() é chamado, sem qualquer modificação na sua etapa de treinamento, essa comunicação desnecessária entre processos pode causar um atraso significativo.

Como você pode evitar esse overhead?

Resolvendo o problema de lentidão

Como você está pulando as atualizações dos parâmetros do modelo ao treinar com esses lotes, seus gradientes não precisam ser sincronizados até o ponto em que optimizer.step() é realmente chamado. O PyTorch não consegue determinar automaticamente quando você precisa fazer isso, mas eles fornecem uma ferramenta para ajudar por meio do gerenciador de contexto no_sync que é adicionado ao seu modelo após convertê-lo para DDP.

Neste gerenciador de contexto, o PyTorch ignorará a sincronização dos gradientes quando .backward() for chamado, e a primeira chamada a .backward() fora deste gerenciador de contexto acionará a sincronização. Veja um exemplo abaixo:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for index, batch in enumerate(dataloader):
    inputs, targets = batch
    # Trigger gradient synchronization on the last batch
    if index != (len(dataloader) - 1):
        with ddp_model.no_sync():
            # Gradients only accumulate
            outputs = ddp_model(inputs)
            loss = loss_func(outputs)
            accelerator.backward(loss)
    else:
        # Gradients finally sync
        outputs = ddp_model(inputs)
        loss = loss_func(outputs)
        accelerator.backward(loss)
        optimizer.step()

No Accelerate, para tornar isso uma API que pode ser chamada independentemente do dispositivo de treinamento (embora ela possa não fazer nada se você não estiver em um sistema distribuído!), ddp_model.no_sync é substituído por no_sync() e opera da mesma forma:

  ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

  for index, batch in enumerate(dataloader):
      inputs, targets = batch
      # Trigger gradient synchronization on the last batch
      if index != (len(dataloader)-1):
-         with ddp_model.no_sync():
+         with accelerator.no_sync(model):
              # Gradients only accumulate
              outputs = ddp_model(inputs)
              loss = loss_func(outputs, targets)
              accelerator.backward(loss)
      else:
          # Gradients finally sync
          outputs = ddp_model(inputs)
          loss = loss_func(outputs)
          accelerator.backward(loss)
          optimizer.step()
          optimizer.zero_grad()

Como você pode esperar, a função accumulate() encapsula essa verificação condicional, mantendo o controle do número atual do lote, resultando na API final de acúmulo de gradientes:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for batch in dataloader:
    with accelerator.accumulate(model):
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

Como resultado, você deve usar accelerator.accumulate ou accelerator.no_sync quando se trata da escolha da API.

Quão grande é o atraso e os erros fáceis que você pode cometer

Para configurar um exemplo realista, considere o seguinte cenário:

• Dois nós de GPU única T4 e um nó com duas GPUs
• Cada GPU é uma T4 e está hospedada no GCP
• O script utilizado é uma modificação do script NLP Example
• O tamanho do lote por GPU é 16, e os gradientes são acumulados a cada 4 etapas

Todos os scripts estão disponíveis em este repositório.

Se não houver cuidado com a sincronização de gradientes e a comunicação entre GPUs, uma grande quantidade de tempo pode ser desperdiçada quando essas GPUs se comunicam entre si durante períodos desnecessários.

Em quanto?

Referência:

• Baseline: não utiliza nenhuma das práticas de sincronização discutidas aqui
• no_sync incorretamente: no_sync apenas ao redor da chamada backward, e não da forward
• no_sync: uso adequado do padrão no_sync
• accumulate: usando accumulate() corretamente

Abaixo estão os segundos médios por lote, iterando sobre 29 lotes de dados para cada configuração, tanto em um nó único quanto na configuração de dois nós:

Como você pode ver, se não tiver cuidado com a forma como configura a sincronização de gradientes, pode acabar com uma lentidão de mais de 2x durante o treinamento!

Se você está preocupado em garantir que tudo seja feito corretamente, recomendamos fortemente utilizar a função accumulate() e passar gradient_accumulation_steps ou gradient_accumulation_plugin para o objeto Accelerator para que o Accelerate possa lidar com isso para você.

no_sync requer memória de GPU adicional ao usar FSDP

Tenha em mente que não sincronizar os gradientes pode ter efeitos adversos durante o treinamento com FSDP. Como foi alertado no torch, o gerenciador de contexto no_sync para FSDP exigirá memória adicional.

Portanto, em situações intensivas de memória ao usar FSDP, recomendamos definir sync_each_batch como True no GradientAccumulationPlugin para desativar no_sync.

Veja o exemplo abaixo, no qual fazemos fine-tuning do Mixtral (47B parâmetros) em 8 GPUs A100-80GB. Observamos que, mesmo para um modesto gradient_accumulation_steps=2, rapidamente esgotamos a memória (OOM) se no_sync estiver habilitado. Novamente, isso se deve a sobrecargas adicionais de memória causadas pelo no_sync do FSDP. No entanto, se no_sync for desabilitado via sync_each_batch=True, o consumo de memória para gradient_accumulation_steps=16 retorna ao nível de gradient_accumulation_steps=1.

Desativar no_sync significa que haverá uma redução de desempenho devido às sincronizações de dados extras, conforme explicado nas seções anteriores deste guia.