Local SGD mit Accelerator

Du hast Gradient Accumulation und Gradient Checkpointing implementiert, um den Speicherverbrauch für dein Übersetzungsmodell zu optimieren. Das Training ist noch etwas langsam, daher fügst du Local SGD in deine Trainingsschleife ein, um die Kommunikation zwischen Geräten effizienter zu machen. Baue die Trainingsschleife mit Local SGD!

model, train_dataloader und accelerator sind bereits definiert, und LocalSGD wurde importiert.

Diese Übung ist Teil des Kurses

<Kurs>Effizientes KI-Modelltraining mit PyTorch</Kurs>

Kurs ansehen

Übungsanweisungen

Setze local_sgd_steps, um die Gradienten alle acht Schritte zu synchronisieren.
Führe den Schritt des Local-SGD-Context-Managers aus.

Interaktive praktische Übung

Versuche dich an dieser Übung, indem du diesen Beispielcode vervollständigst.

# 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 bearbeiten und ausführen