Local SGD con Accelerator

Hai già implementato l'accumulo dei gradienti e il gradient checkpointing per ottimizzare l'uso della memoria nel tuo modello di traduzione automatica. L'addestramento è ancora un po' lento, quindi decidi di aggiungere la local SGD al ciclo di training per migliorare l'efficienza della comunicazione tra dispositivi. Costruisci il ciclo di training con la local SGD!

model, train_dataloader e accelerator sono già definiti e LocalSGD è stato importato.

Questo esercizio fa parte del corso

Efficient AI Model Training with PyTorch

Visualizza corso

Istruzioni dell'esercizio

Imposta local_sgd_steps per sincronizzare i gradienti ogni otto step.
Esegui lo step del context manager di local SGD.

esercizio interattivo pratico

Prova questo esercizio completando questo codice di esempio.

# 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.____()

Modifica ed esegui il codice