Trenowanie rozproszone w notesach

Important

Ta funkcja jest dostępna w wersji beta.

Dekorator @distributed z Databricks AI Runtime SDK (databricks-sdk-air) to najwygodniejszy sposób uruchamiania rozproszonego treningu z notebooka Databricks. Zdekoruj funkcję treningową, wywołaj ją, a AI Runtime uruchamia ją na wszystkich GPU w węźle, do którego podłączony jest Twój notebook. Ten sam kod skaluje się od pojedynczego GPU do wielu GPU bez konieczności konfigurowania klastra i bez konieczności konfigurowania programu uruchamiającego dla środowiska rozproszonego.

Kluczowe punkty na tej stronie:

• Dekorator @distributed uruchamia funkcję treningową na każdym procesorze GPU w Twoim węźle bezpośrednio z poziomu notatnika.
• Obsługuje PyTorch DDP, FSDP i DeepSpeed oraz przenosi kod z pojedynczego GPU na multi-GPU z minimalnymi zmianami.
• Podłącz swój notebook do akceleratora 8xH100, aby wykonywać pełny trening na wielu GPU. Domyślnie @distributed używa wszystkich GPU na węźle.

Note

Ta strona opisuje trenowanie rozproszone z notebooków Databricks z pomocą Databricks AI Runtime SDK. Aby przesłać rozproszone obciążenia treningowe z lokalnego komputera, użyj poleceń Databricks CLI dla AI Runtime. Zobacz Korzystanie z interfejsu wiersza polecenia Databricks z AI Runtime.

Szybki start

Przykłady na tej stronie wykorzystują Databricks AI Runtime SDK, które jest preinstalowane w środowisku GPU w wersji 6 lub nowszej. W wersji 5 i niższej zamiast tego używaj pakietu serverless_gpu. Zobacz Migruj z serverless_gpu. Aby podłączyć notebook do akceleratora, zobacz Krok 1: Podłącz notebook do GPU.

Oznacz swoją funkcję treningową za pomocą @distributed, a następnie wywołaj ją za pomocą .distributed():

from databricks.air import distributed

# By default, @distributed uses every GPU on the node your notebook is connected
# to. To use fewer, pass num_accelerators (for example, num_accelerators=2).
@distributed()
def train():
    import os
    import torch
    import torch.distributed as dist

    # Bind this process to its own GPU before training.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group("nccl")
    # ... build the model and data on `device`, then run your training loop ...
    dist.destroy_process_group()

train.distributed()

Trenowanie rozproszone wymaga układu 8xH100, który zapewnia jeden węzeł z 8 GPU. @distributed domyślnie używa wszystkich ośmiu kart graficznych. Ustaw num_accelerators tak, aby używać mniej. Parametr accelerator_type jest również opcjonalny i automatycznie wykrywany przez akcelerator, do którego podłączony jest Twój notebook.

Każde wywołanie .distributed() tworzy uruchomienie MLflow (lub zagnieżdżone uruchomienie podrzędne, jeśli jedno jest już aktywne) i wyświetla link do uruchomienia w danych wyjściowych komórki. Kompletny, możliwy do przejścia przewodnik znajdziesz w pełnym przykładzie.

Obsługiwane struktury

Interfejs @distributed API integruje się z głównymi rozproszonymi bibliotekami szkoleniowymi:

• PyTorch Distributed Data Parallel (DDP): standardowa równoległość danych na wielu procesorach GPU.
• Fully Sharded Data Parallel (FSDP): oszczędne pamięciowo trenowanie dużych modeli.
• DeepSpeed: biblioteka optymalizacji Microsoft na potrzeby trenowania dużych modeli.

Dla rzeczywistych scenariuszy treningowych wykorzystujących każdą bibliotekę, zobacz przykłady treningu rozproszonego.

Jak działa @distributed dekorator

Gdy wywołujesz funkcję dekorowaną za pomocą .distributed(), AI Runtime obsługuje mechanikę, którą normalnie konfigurowałbyś ręcznie za pomocą rozproszonego launchera:

• Serializacja i rozproszenie: Funkcja jest serializowana i uruchamiana na każdym wybranym GPU. Każdy GPU uruchamia kopię funkcji z tymi samymi argumentami.
• Synchronizacja środowiska: Środowisko Python i zależności są replikowane na wszystkich poziomach, więc każdy proces uruchamia ten sam kod.
• Zmienne środowiskowe rank: Standardowe zmienne, takie jak LOCAL_RANK, są ustawiane dla każdego procesu. Przeczytaj je w swojej funkcji, aby umieścić model i dane na właściwym urządzeniu.
• Zbieranie wyników: Wartości zwrotne są zbierane ze wszystkich rang i zwracane do dzwoniącego.
• Śledzenie MLflow: Każde .distributed() wywołanie tworzy przebieg MLflow lub zagnieżdżony przebieg podrzędny, jeśli jeden jest już aktywny, więc metryki rejestrowane przez funkcję trafiają do tego samego przebiegu.
• Cykl życia i limit czasu: Rozproszone wykonywanie działa w ramach cyklu życia notatnika. Zakończenie notesu kończy uruchomienie. Dekorator ma domyślny limit czasu wynoszący 3 godziny. Przekaż timeout w ciągu kilku sekund, aby to zmienić, lub timeout=None, aby to wyłączyć. Niestandardowe limity czasu wymagają środowiska GPU w wersji v5 lub nowszej.

API opiera się na standardowych bibliotekach PyTorch: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) oraz DeepSpeed.

Migruj z serverless_gpu

Pakiet serverless_gpu jest wycofany i nie otrzymuje już nowych funkcji. Począwszy od wersji 6 środowiska GPU, pakiet SDK Databricks AI Runtime (databricks-sdk-air) zastępuje go i udostępnia równoważne interfejsy API z przemianowanymi argumentami w przestrzeni nazw databricks.air. Pakiet, którego używasz, zależy od wersji środowiska GPU:

• Wersja 5 i niższa: Zawieraj tylko serverless_gpu. Kontynuuj używanie serverless_gpu w przypadku tych wersji.
• Wersja 6: Zawiera oba pakiety, ale aktualizacje otrzymuje tylko AI Runtime SDK. Wersja 6 to ostatnia wersja środowiska GPU, która zawiera serverless_gpu.

AI Runtime SDK to osobny pakiet od Databricks SDK dla Python (databricks-sdk).

	serverless_gpu (przestarzałe)	SDK środowiska uruchomieniowego AI (databricks.air)
Wersje środowiska GPU	6 i mniej	6 i więcej
Import	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argument liczby GPU	gpus	num_accelerators
Argument typu GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Wartości typu GPU	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Moduł ładowania danych i checkpointingu	serverless_gpu.data	databricks.air.data. Zobacz Ładowanie danych w środowisku uruchomieniowym sztucznej inteligencji.

Aby przenieść notatnik:

1. Podłącz notebook do środowiska GPU w wersji 6 lub nowszej, które obejmuje pakiet SDK AI Runtime.
2. Zmień importy z serverless_gpu na databricks.air.
3. Przemianuj gpus na num_accelerators i gpu_type na accelerator_type, i użyj nowych wartości typu akceleratora. Stare nazwy argumentów i wartości typów GPU nadal działają w databricks.air, ale generują ostrzeżenie FutureWarning. Wsparcie dla środowiska GPU w wersji 7 zakończy się.

Wskazówka

Powiedz Genie Code, żeby zrobił to za Ciebie:

Migrate this notebook from the serverless_gpu SDK to the AI Runtime SDK.

W przypadku obciążeń databricks.air.data interfejsu Databricks CLI, które używają databricks-sdk-air[data], dodaj databricks-sdk-air[data] do elementu environment.dependencies obciążenia.

Pochodzi z TorchDistributor

Jeśli uruchamiasz trening rozproszony za pomocą TorchDistributor na klastrach Spark i Twoje obciążenie mieści się na jednym węźle, @distributed API jest zalecanym zastępctwem w AI Runtime. Twój kod modelu DDP, FSDP czy DeepSpeed się nie zmienia.

Feature	@distributed Interfejs Programowania Aplikacji (API)	TorchDistributor
Infrastruktura	W pełni bezserwerowe, bez zarządzania klastrem	Wymaga klastra Spark z pracownikami GPU
Konfiguracja	Pojedynczy dekorator, minimalna konfiguracja	Wymaga konfiguracji klastra Spark i narzędzia TorchDistributor
Obsługa komponentu Framework	PyTorch DDP, FSDP, DeepSpeed	Przede wszystkim PyTorch DDP
Ładowanie danych	W dekoratorze są używane woluminy Unity Catalog (UCVolumeDataset do przesyłania strumieniowych danych plikowych)	Za pośrednictwem platformy Spark lub systemu plików

Aby zmigrować obciążenie robocze z jednym węzłem:

• Zamień wywołanie TorchDistributor(...).run(train_fn, ...) na dekorator @distributed w elemencie train_fn, a następnie uruchom za pomocą train_fn.distributed(...).
• Usuń klaster Spark i konfigurację procesora GPU. Podłącz swój notebook zamiast tego do akceleratora 8xH100. @distributed domyślnie używa wszystkich ośmiu kart graficznych.
• Przenieś ładowanie danych do udekorowanej funkcji. Zobacz ładowanie danych.
• Zachowaj istniejący kod modelu DDP, FSDP lub DeepSpeed. Dekorator wspiera wszystkie trzy.

Dekorator @distributed jest w fazie beta i działa na jednym węźle (patrz Ograniczenia), więc nie zastępuje wszystkich zadań TorchDistributor. Zachowaj zadania, które zależą od integracji ze Spark, w TorchDistributor. Aby uruchomić trenowanie rozproszone z lokalnej maszyny lub na wielu węzłach, użyj poleceń Databricks CLI dla AI Runtime. Zobacz Korzystanie z interfejsu wiersza polecenia Databricks z AI Runtime.

Pełny przykład

Poniższy przykład trenuje wielowarstwowy model perceptron (MLP) na 8 kartach graficznych H100 z notebooka.

1. Skonfiguruj model i zdefiniuj funkcje narzędzi.

   
   # Define the model
   import os
   import torch
   import torch.distributed as dist
   import torch.nn as nn
   
   def setup():
       torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
       dist.init_process_group("nccl")
   
   def cleanup():
       dist.destroy_process_group()
   
   class SimpleMLP(nn.Module):
       def __init__(self, input_dim=10, hidden_dim=64, output_dim=1):
           super().__init__()
           self.net = nn.Sequential(
               nn.Linear(input_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, output_dim)
           )
   
       def forward(self, x):
           return self.net(x)
   

2. Zaimportuj dekorator distributed z AI Runtime SDK.

   from databricks.air import distributed
   

3. Umieść kod trenowania modelu w funkcji i udekoruj funkcję dekoratorem @distributed. Funkcja dekorowana jest punktem wejścia do rozproszonego wykonywania, więc zdefiniuj całą logikę treningową, ładowanie danych i inicjalizację modelu w jej wnętrzu.

   @distributed()
   def run_train(num_epochs: int, batch_size: int) -> None:
       import mlflow
       import torch.optim as optim
       from torch.nn.parallel import DistributedDataParallel as DDP
       from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
   
       # 1. Set up multi-GPU environment
       setup()
       device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
   
       # 2. Apply the Torch distributed data parallel (DDP) library for data-parallel training.
       model = SimpleMLP().to(device)
       model = DDP(model, device_ids=[device])
   
       # 3. Create and load dataset.
       x = torch.randn(5000, 10)
       y = torch.randn(5000, 1)
   
       dataset = TensorDataset(x, y)
       sampler = DistributedSampler(dataset)
       dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)
   
       # 4. Define the training loop.
       optimizer = optim.Adam(model.parameters(), lr=0.001)
       loss_fn = nn.MSELoss()
   
       for epoch in range(num_epochs):
           sampler.set_epoch(epoch)
           model.train()
           total_loss = 0.0
           for step, (xb, yb) in enumerate(dataloader):
               xb, yb = xb.to(device), yb.to(device)
               optimizer.zero_grad()
               loss = loss_fn(model(xb), yb)
               # Log loss to MLflow metric
               mlflow.log_metric("loss", loss.item(), step=step)
   
               loss.backward()
               optimizer.step()
               total_loss += loss.item() * xb.size(0)
   
           mlflow.log_metric("total_loss", total_loss)
           print(f"Total loss for epoch {epoch}: {total_loss}")
   
       cleanup()
   

4. Uruchom rozproszone szkolenie, wywołując funkcję rozproszoną za pomocą argumentów zdefiniowanych przez użytkownika.

   run_train.distributed(num_epochs=3, batch_size=1)
   

5. Po wykonaniu w obszarze wyjściowym komórki notesnika zostanie wygenerowany link do uruchomienia MLflow. Kliknij link przebiegu MLflow lub znajdź go w panelu Eksperyment, aby zobaczyć wyniki przebiegu. Aby uzyskać szczegółowe informacje na temat dostosowywania nazw eksperymentów, śledzenia metryk i wznawiania przebiegów, zobacz Śledzenie eksperymentów i obserwowanie.

Ładowanie danych

Umieść kod ładowania @distributed danych wewnątrz funkcji. Zbiór danych może przekraczać maksymalny rozmiar dozwolony przez pickle, więc generowanie lub ładowanie go wewnątrz dekoratora pozwala uniknąć błędów serializacji:

from databricks.air import distributed

# Don't do this: loading the dataset outside the function means it's captured
# and pickled with the function, which can raise a pickle error.
# dataset = get_dataset(file_path)

@distributed()
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

W przypadku danych przechowywanych w plikach w woluminach Unity Catalog użyj UCVolumeDataset z databricks.air.data, które strumieniuje pliki z lokalnym buforowaniem i automatycznie dzieli je między rangi i procesy robocze. Aby zapisać punkt kontrolny treningu rozproszonego do woluminu, użyj poleceń UCVolumeWriter i UCVolumeReader. Zobacz Ładowanie danych w AI Runtime oraz Tworzenie punktu kontrolnego za pomocą Distributed Checkpoint (DCP).

Limitations

• @distributed Obsługuje tylko jeden węzeł: działa na procesorach GPU w węźle, do którego podłączony jest Twój notatnik. Aby trenować na wielu węzłach, użyj AI Runtime Databricks CLI. Zobacz Korzystanie z interfejsu wiersza polecenia Databricks z AI Runtime.
• Jeśli ustawisz accelerator_type, musi on odpowiadać akceleratorowi, do którego podłączony jest Twój notebook (na przykład "GPU_8xH100"). Niedopasowanie powoduje awarię obciążenia. Parametr jest opcjonalny i automatycznie wykrywany po pominięciu parametru.
• AI Runtime obsługuje środowisko GPU v4 i nowsze. Niestandardowe limity czasowe (parametr) timeout wymagają środowiska GPU v5 i wyższego.
• Domyślnie limit czasu dekoratora wynosi 3 godziny. Przekaż timeout w ciągu kilku sekund, aby to zmienić, lub timeout=None, aby to wyłączyć.
• Wykonywanie odbywa się w cyklu życia notatnika. Zakończenie notesu kończy uruchomienie.

Dodatkowe zasoby

• Informacje o API dekoratora @distributed, AcceleratorType oraz Ray znajdziesz w dokumentacji referencyjnej interfejsu API Databricks AI Runtime SDK.
• Aby poznać wzorce, które czynią Twój proces szkoleniowy bardziej wydajnym i odpornym, zobacz przewodnik po wydajności i odporności.
• W przypadku scenariuszy kompleksowego szkolenia zobacz przykłady trenowania rozproszonego.