Trenowanie rozproszone za pomocą narzędzia TorchDistributor

W tym artykule opisano sposób wykonywania trenowania rozproszonego na modelach uczenia maszynowego PyTorch przy użyciu narzędzia TorchDistributor.

TorchDistributor to moduł typu open source w programie PySpark, który ułatwia użytkownikom trenowanie rozproszone za pomocą rozwiązania PyTorch w swoich klastrach Spark, dzięki czemu umożliwia uruchamianie zadań szkoleniowych PyTorch jako zadań platformy Spark. Under-the-hood inicjuje środowisko i kanały komunikacji między procesami roboczymi i wykorzystuje polecenie interfejsu wiersza polecenia torch.distributed.run do uruchamiania rozproszonego szkolenia między węzłami procesu roboczego.

Interfejs API TorchDistributor obsługuje metody przedstawione w poniższej tabeli.

Metoda i podpis opis
init(self, num_processes, local_mode, use_gpu) Utwórz wystąpienie narzędzia TorchDistributor.
run(self, main, *args) Uruchamia trenowanie rozproszone przez wywołanie main(**kwargs) , jeśli main jest funkcją i uruchamia polecenie interfejsu wiersza polecenia torchrun main *args , jeśli głównym jest ścieżka pliku.

Wymagania

  • Spark 3.4
  • Databricks Runtime 13.0 ML lub nowsza wersja

Przepływ pracy tworzenia notesów

Jeśli proces tworzenia i trenowania modelu odbywa się całkowicie z notesu na komputerze lokalnym lub notesie usługi Databricks, wystarczy wprowadzić drobne zmiany, aby przygotować kod do trenowania rozproszonego.

  1. Przygotuj kod pojedynczego węzła: przygotuj i przetestuj kod pojedynczego węzła za pomocą biblioteki PyTorch, PyTorch Lightning lub innych struktur opartych na interfejsie API PyTorch/PyTorch Lightning, takim jak Interfejs API trenera HuggingFace.

  2. Przygotowanie kodu do standardowego trenowania rozproszonego: należy przekonwertować trenowanie pojedynczego procesu na trenowanie rozproszone. Aby ten kod rozproszony obejmował wszystkie funkcje trenowania, których można używać z .TorchDistributor

  3. Przenieś importy w funkcji trenowania: dodaj niezbędne importy, takie jak import torch, w funkcji trenowania. Dzięki temu można uniknąć typowych błędów pikowania. Ponadto device_id te modele i dane są powiązane z nimi, są określane przez:

    device_id = int(os.environ["LOCAL_RANK"])
    
  4. Uruchom trenowanie rozproszone: utwórz wystąpienie TorchDistributor z żądanymi parametrami i wywołaj polecenie .run(*args) w celu uruchomienia szkolenia.

Poniżej przedstawiono przykład kodu szkoleniowego:

from pyspark.ml.torch.distributor import TorchDistributor

def train(learning_rate, use_gpu):
  import torch
  import torch.distributed as dist
  import torch.nn.parallel.DistributedDataParallel as DDP
  from torch.utils.data import DistributedSampler, DataLoader

  backend = "nccl" if use_gpu else "gloo"
  dist.init_process_group(backend)
  device = int(os.environ["LOCAL_RANK"]) if use_gpu  else "cpu"
  model = DDP(createModel(), **kwargs)
  sampler = DistributedSampler(dataset)
  loader = DataLoader(dataset, sampler=sampler)

  output = train(model, loader, learning_rate)
  dist.cleanup()
  return output

distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train, 1e-3, True)

Migrowanie szkolenia z repozytoriów zewnętrznych

Jeśli masz istniejącą procedurę trenowania rozproszonego przechowywaną w repozytorium zewnętrznym, możesz łatwo przeprowadzić migrację do usługi Azure Databricks, wykonując następujące czynności:

  1. Zaimportuj repozytorium: zaimportuj repozytorium zewnętrzne jako folder Git usługi Databricks.
  2. Utwórz nowy notes Inicjuj nowy notes usługi Azure Databricks w repozytorium.
  3. Uruchom trenowanie rozproszone W komórce notesu wywołaj metodę TorchDistributor podobną do następującej:
from pyspark.ml.torch.distributor import TorchDistributor

train_file = "/path/to/train.py"
args = ["--learning_rate=0.001", "--batch_size=16"]
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train_file, *args)

Rozwiązywanie problemów

Typowy błąd przepływu pracy notesu polega na tym, że nie można odnaleźć ani zebrać obiektów podczas uruchamiania trenowania rozproszonego. Może się tak zdarzyć, gdy instrukcje importowania biblioteki nie są dystrybuowane do innych funkcji wykonawczych.

Aby uniknąć tego problemu, uwzględnij wszystkie instrukcje importu (na przykład import torch) zarówno w górnej części funkcji trenowania, która jest wywoływana z funkcją TorchDistributor(...).run(<func>) i wewnątrz innych funkcji zdefiniowanych przez użytkownika wywoływanych w metodzie trenowania.

Przykładowe notesy

W poniższych przykładach notesu pokazano, jak przeprowadzić trenowanie rozproszone za pomocą rozwiązania PyTorch.

Kompleksowe trenowanie rozproszone w notesie usługi Databricks

Pobierz notes

Rozproszone dostrajanie notesu modelu rozpoznawania twarzy

Pobierz notes

Trenowanie rozproszone w notesie plików PyTorch

Pobierz notes

Trenowanie rozproszone przy użyciu notesu PyTorch Lightning

Pobierz notes

Ładowanie rozproszonych danych przy użyciu notesu Petastorm

Pobierz notes