Executando e adiando jobs

Quando você executa seu script habitual, as instruções são executadas em ordem. Usar o Accelerate para implantar seu script em várias GPUs ao mesmo tempo introduz uma complicação: embora cada processo execute todas as instruções em ordem, alguns podem ser mais rápidos que outros.

Você pode precisar esperar que todos os processos tenham atingido um determinado ponto antes de executar uma instrução específica. Por exemplo, você não deve salvar um modelo antes de ter certeza de que todos os processos concluíram o treinamento, e não gostaria de continuar o treinamento antes que todos os pesos do modelo tenham sido carregados. Para fazer isso, basta escrever a seguinte linha no seu código:

accelerator.wait_for_everyone()

Esta instrução bloqueará todos os processos que chegarem primeiro até que todos os outros processos tenham alcançado esse ponto (se você executar seu script em apenas uma GPU ou CPU, isso não terá nenhum efeito).

Alguns exemplos de casos em que usar esta utilidade estão listados abaixo:

Alguns deles são utilizados com o gerenciador de contexto main_process_first(), que utiliza wait_for_everyone() para executar um conjunto específico de código no processo principal antes de acionar e iniciar os outros processos

Baixando um Conjunto de Dados

Ao baixar um conjunto de dados, você deve baixá-lo primeiro no processo principal e, em seguida, carregar o conjunto de dados em cache

load_dataset executa um lock internamente para impedir que múltiplos downloads ocorram simultaneamente, mas, se você estiver baixando algo sem usar esta biblioteca, deve usar este método.

with accelerator.main_process_first():
    datasets = load_dataset("glue", "mrpc")

Por baixo dos panos, isso é o mesmo que chamar:

# First do something on the main process
if accelerator.is_main_process:
    datasets = load_dataset("glue", "mrpc")
else:
    accelerator.wait_for_everyone()

# And then send it to the rest of them
if not accelerator.is_main_process:
    datasets = load_dataset("glue", "mrpc")
else:
    accelerator.wait_for_everyone()

Salvando o state_dict

Ao salvar o state_dict do modelo, como normalmente você salvaria um arquivo apenas no processo principal, você deve especificar isso:

if accelerator.is_main_process:
    model = accelerator.unwrap_model(model)
    torch.save(model.state_dict(), "weights.pth")

Carregando o state_dict

Ao carregar o state_dict em um modelo, otimizador ou agendador, você deve aguardar que todos os workers tenham carregado os pesos antes de prosseguir para o treinamento

with accelerator.main_process_first():
    state = torch.load("weights.pth")
    model.load_state_dict(state)

Aplicando uma operação de CPU com múltiplos workers

Aplicar uma operação map() em vários workers, como a tokenização, deve ser feito primeiro no processo principal e, em seguida, propagado para cada um.

datasets = load_dataset("glue", "mrpc")

with accelerator.main_process_first():
    tokenized_datasets = datasets.map(
        tokenize_function,
        batched=True,
        remove_columns=["idx", "sentence1", "sentence2"],
    )

Aplicando verificações como Early Stopping

Para ter uma verificação que funcione com uma flag definida por um processo específico, a API set_trigger e check_trigger deve ser usada. Exemplos úteis para fazer isso podem incluir situações como o uso de early stopping e o monitoramento da perda (já que cada perda difere ligeiramente em cada processo).

Chame Accelerator.set_trigger() quando sua condição tiver sido atendida e Accelerator.check_trigger() ao verificar se essa condição foi atendida em qualquer processo:

for (x,y) in data_loader:
    logits = model(x)
    loss = loss_func(logits, y)
    # Assume `should_do_early_stopping` is a custom defined function that returns a conditional
    if should_do_early_stopping(loss):
        accelerator.set_trigger()

    # Later in the training script when we need to check for the breakpoint
    if accelerator.check_trigger():
        break