Lokalne SGD z Acceleratorem
Zaimplementowano już akumulację gradientów i gradient checkpointing, aby zoptymalizować zużycie pamięci w modelu do tłumaczenia języków. Trening nadal jest trochę wolny, dlatego postanawiasz dodać lokalne SGD do pętli treningowej, aby zwiększyć efektywność komunikacji między urządzeniami. Zbuduj pętlę treningową z lokalnym SGD!
Zmienne model, train_dataloader i accelerator są już wstępnie zdefiniowane, a LocalSGD został zaimportowany.
To ćwiczenie jest częścią kursu
Efektywne trenowanie modeli AI z PyTorch
Instrukcje do ćwiczenia
- Ustaw
local_sgd_stepstak, aby synchronizować gradienty co osiem kroków. - Wykonaj krok menedżera kontekstu lokalnego SGD.
Interaktywne ćwiczenie praktyczne
Spróbuj tego ćwiczenia, uzupełniając ten przykładowy kod.
# 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.____()