Lokale SGD met Accelerator

Je hebt gradient accumulation en gradient checkpointing geïmplementeerd om het geheugengebruik van je vertaalmodel te stroomlijnen. De training is nog steeds wat traag, dus je besluit local SGD toe te voegen aan je trainingslus om de communicatie tussen apparaten efficiënter te maken. Bouw de trainingslus met local SGD!

Het model, de train_dataloader en de accelerator zijn al gedefinieerd, en LocalSGD is geïmporteerd.

Deze oefening maakt deel uit van de cursus

Efficiënt AI-modellen trainen met PyTorch

Bekijk cursus

Oefeninstructies

Stel local_sgd_steps in om de gradients elke acht stappen te synchroniseren.
Voer de stap uit met de local SGD-contextmanager.

Interactieve oefening met praktijkervaring

Probeer deze oefening door deze voorbeeldcode aan te vullen.

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

Code bewerken en uitvoeren