Adicione o Accelerate ao seu código

Cada framework de treinamento distribuído tem sua própria maneira de funcionar, o que pode exigir a escrita de muito código personalizado para adaptá-lo ao seu código de treinamento do PyTorch e ao seu ambiente de treinamento. O Accelerate oferece uma maneira amigável de interagir com esses frameworks de treinamento distribuído sem precisar aprender os detalhes específicos de cada um. O Accelerate cuida desses detalhes para você, permitindo que você se concentre no código de treinamento e o escale para qualquer ambiente de treinamento distribuído.

Neste tutorial, você aprenderá a adaptar seu código existente de PyTorch com o Accelerate e começar a treinar em sistemas distribuídos com facilidade! Você começará com um loop de treinamento básico de PyTorch (assumindo que todos os objetos de treinamento, como model e optimizer, já estejam configurados) e integrará o Accelerate a ele progressivamente.

device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()

Acelerador

O Accelerator é a classe principal para adaptar seu código para funcionar com o Accelerate. Ele conhece a configuração distribuída que você está usando, como o número de processos diferentes e o tipo de hardware. Esta classe também fornece acesso a muitos dos métodos necessários para permitir que seu código PyTorch funcione em qualquer ambiente de treinamento distribuído e para gerenciar e executar processos entre dispositivos.

É por isso que você deve sempre começar importando e criando uma instância de Accelerator no seu script.

from accelerate import Accelerator

accelerator = Accelerator()

O Accelerator também sabe para qual dispositivo mover seus objetos do PyTorch, portanto, é recomendado deixar que o Accelerate trate disso para você.

- device = "cuda"
+ device = accelerator.device
  model.to(device)

Preparar objetos do PyTorch

Em seguida, você precisa preparar seus objetos do PyTorch (modelo, otimizador, agendador, etc.) para o treinamento distribuído. O método prepare() cuida de posicionar seu modelo no contêiner apropriado (como GPU única ou multi-GPU) para a sua configuração de treinamento, adapta o otimizador e o agendador para usar o AcceleratedOptimizer e o AcceleratedScheduler do Accelerate, e cria um novo dataloader que pode ser fragmentado entre processos.

O Accelerate prepara apenas objetos que herdam de suas respectivas classes do PyTorch, como torch.optim.Optimizer.

Os objetos do PyTorch são retornados na mesma ordem em que são enviados.

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

Loop de treinamento

Por fim, remova as chamadas to(device) para os inputs e targets no loop de treinamento, pois as classes de DataLoader do Accelerate as posicionam automaticamente no dispositivo correto. Você também deve substituir a passagem usual de backward() pelo método backward() do Accelerate, que escala os gradientes para você e usa o método backward() apropriado dependendo da sua configuração distribuída (por exemplo, DeepSpeed ou Megatron).

-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)

Junte tudo e seu novo loop de treinamento do Accelerate deve ficar assim!

from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

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

Recursos de treinamento

O Accelerate oferece recursos adicionais - como acúmulo de gradiente, clipping de gradiente, treinamento com precisão mista e mais - que você pode adicionar ao seu script para melhorar sua execução de treinamento. Vamos explorar esses três recursos.

Acúmulo de gradiente

A acumulação de gradientes permite treinar com tamanhos de lote maiores, acumulando os gradientes em vários lotes antes de atualizar os pesos. Isso pode ser útil para contornar limitações de memória. Para habilitar esse recurso no Accelerate, especifique o parâmetro gradient_accumulation_steps na classe Accelerator e adicione o gerenciador de contexto accumulate() ao seu script.

+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()

Corte de gradiente

O recorte de gradiente é uma técnica para evitar “gradientes explosivos”, e o Accelerate oferece:

• clipgrad_value() para limitar os gradientes a um valor mínimo e máximo
• clipgrad_norm() para normalizar os gradientes para um determinado valor

Precisão mista

A precisão mista acelera o treinamento ao usar um tipo de dados de menor precisão, como fp16 (precisão dupla), para calcular os gradientes. Para o melhor desempenho com Accelerate, a perda deve ser calculada dentro do seu modelo (como nos modelos Transformers), pois os cálculos fora do modelo são executados em precisão total.

Defina o tipo de precisão mista a ser usado no Accelerator e, em seguida, use o gerenciador de contexto autocast() para converter automaticamente os valores para o tipo de dados especificado.

O Accelerate habilita a precisão mista automática, portanto autocast() só é necessário se houver outras operações de precisão mista além daquelas executadas na perda por backward(), que já trata o escalonamento.

+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target)

Salvar e carregar

O Accelerate também pode salvar e carregar um modelo uma vez que o treinamento esteja concluído, ou você também pode salvar o estado do modelo e do otimizador, o que pode ser útil para retomar o treinamento.

Model

Quando todos os processos estiverem concluídos, desfaça o encapsulamento do modelo com o método unwrap_model() antes de salvá-lo, pois o método prepare() encapsulou seu modelo na interface adequada para treinamento distribuído. Se você não desfizer o encapsulamento do modelo, salvar o dicionário de estado do modelo também salvará quaisquer camadas extras potenciais do modelo maior e você não conseguirá carregar os pesos de volta para o seu modelo base.

Você deve usar o método save_model() para desempacotar e salvar o dicionário de estado do modelo. Este método também pode salvar um modelo em checkpoints fragmentados ou no formato safetensors.

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)

load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})

Estado

Durante o treinamento, você pode querer salvar o estado atual do modelo, do otimizador, dos geradores aleatórios e, potencialmente, dos agendadores de taxa de aprendizado, para que possam ser restaurados no mesmo script. Você deve adicionar os métodos save_state() e load_state() ao seu script para salvar e carregar estados.

Para personalizar ainda mais onde e como os estados são salvos por meio de save_state(), use a classe ProjectConfiguration. Por exemplo, se automatic_checkpoint_naming estiver habilitado, cada checkpoint salvo é armazenado em Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}.

Quaisquer outros itens com estado a serem armazenados devem ser registrados com o método register_for_checkpointing() para que possam ser salvos e carregados. Todo objeto passado a este método para armazenamento deve ter uma função load_state_dict e uma função state_dict.

Se você tiver torchdata>=0.8.0 instalado, você pode adicionalmente passar use_stateful_dataloader=True para o seu DataLoaderConfiguration. Isso estende as classes de DataLoader do Accelerate com uma função load_state_dict e uma função state_dict, e faz com que Accelerator.save_state e Accelerator.load_state também rastreiem até onde o conjunto de dados de treinamento foi lido ao persistir o modelo.