Lokales SGD mit Accelerator
Du hast Gradient Accumulation und Gradient Checkpointing implementiert, um den Speicherverbrauch deines Übersetzungsmodells zu optimieren. Das Training ist noch etwas langsam, daher entscheidest du dich, lokales SGD in deine Trainingsschleife aufzunehmen, um die Kommunikation zwischen Geräten effizienter zu machen. Baue die Trainingsschleife mit lokalem SGD!
Das model, train_dataloader und der accelerator sind vordefiniert, und LocalSGD wurde importiert.
Diese Übung ist Teil des Kurses
Effizientes KI-Modelltraining mit PyTorch
Anleitung zur Übung
- Richte einen Context-Manager für lokales SGD ein und synchronisiere die Gradienten alle acht Schritte.
- Führe den Schritt des Local-SGD-Context-Managers aus.
Interaktive Übung
Vervollständige den Beispielcode, um diese Übung erfolgreich abzuschließen.
# Set up a context manager to synchronize gradients every eight steps
with ____(accelerator=accelerator, model=model, ____=____, 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
____.____()