Local SGD com Accelerator

Você implementou acumulação de gradientes e gradient checkpointing para otimizar o uso de memória no seu modelo de tradução de linguagem. O treinamento ainda está um pouco lento, então você decide adicionar Local SGD ao seu loop de treinamento para melhorar a eficiência de comunicação entre dispositivos. Construa o loop de treinamento com Local SGD!

O model, train_dataloader e o accelerator já foram definidos, e LocalSGD já foi importado.

Este exercicio faz parte do curso

Treinamento Eficiente de Modelos de IA com PyTorch

Ver curso

Instruções do exercicio

Defina local_sgd_steps para sincronizar os gradientes a cada oito passos.
Execute o gerenciador de contexto de Local SGD com .step().

exercicio interativo prático

Tente este exercicio completando este código de exemplo.

# Set up a context manager to synchronize gradients every eight steps
with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=____, enabled=True) as local_sgd:
    for batch in train_dataloader:
        with accelerator.accumulate(model):
            inputs, targets = batch["input_ids"], batch["labels"]
            outputs = model(inputs, labels=targets)
            loss = outputs.loss
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            # Step the local SGD context manager
            local_sgd.____()

Editar e Executar Código