Lokal SGD med Accelerator

Du har implementerat gradientackumulering och gradientlagring för att effektivisera minnesanvändningen i din språköversättningsmodell. Träningen är fortfarande lite långsam, så du bestämmer dig för att lägga till lokal SGD i din träningsloop för att förbättra kommunikationseffektiviteten mellan enheterna. Bygg träningsloopen med lokal SGD!

model, train_dataloader och accelerator är fördefinierade, och LocalSGD har importerats.

Den här övningen är en del av kursen

Effektiv AI-modellträning med PyTorch

Visa kurs

Övningsinstruktioner

Sätt local_sgd_steps så att gradienter synkroniseras var åttonde steg.
Anropa .step() på kontexthanteraren för lokal SGD.

Interaktiv övning med praktiskt arbete

Testa den här övningen genom att slutföra den här exempelkoden.

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

Redigera och kör kod