Trenowanie rozproszone w notesach

Ważna

Ta funkcja jest dostępna w wersji beta. Administratorzy obszaru roboczego mogą kontrolować dostęp do tej funkcji ze strony Podglądy . Zobacz Zarządzanie wersjami zapoznawczami usługi Azure Databricks.

Dekorator @distributed z Serverless GPU Python API to najwygodniejszy sposób na uruchomienie 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.

Tip

  • 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 i przygotuj gpus=8 pełny trening multi-GPU.

Note

Ta strona opisuje trening rozproszony z poziomu notesów Databricks przy użyciu bezserwerowego interfejsu API Python dla GPU. Aby przesłać rozproszone obciążenia treningowe z lokalnego komputera, użyj poleceń Databricks CLI dla AI Runtime, które są dostępne w trybie Public Preview. Zobacz Korzystanie z interfejsu wiersza polecenia Databricks z AI Runtime.

Szybki start

Pakiet serverless_gpu jest wstępnie zainstalowany, gdy notatnik jest połączony z bezserwerowym procesorem graficznym (GPU). Oznacz swoją funkcję treningową za pomocą @distributed, a następnie wywołaj ją za pomocą .distributed():

from serverless_gpu import distributed

# gpus is the number of GPUs on the node. gpu_type is optional and
# auto-detected from the accelerator your notebook is connected to.
@distributed(gpus=8, gpu_type="H100")
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. W przypadku korzystania z dekoratora @distributed ustaw wartość gpus=8. Parametr gpu_type jest opcjonalny i jest automatycznie wykrywany na podstawie akceleratora, z którym połączony jest notatnik.

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 platformy

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.

Aby zapoznać się z rzeczywistymi scenariuszami treningowymi z użyciem każdej biblioteki, zobacz przykładowe notatniki.

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 rozprzestrzenianie się: Funkcja jest serializowana i uruchamiana na każde z gpus Twoich żądań. 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.

Pochodzi z TorchDistributor

Jeśli dziś uruchamiasz rozproszony PyTorch na Sparku z TorchDistributor i Twoje obciążenie mieści się na jednym węźle, API serverless_gpu@distributed jest zalecanym zamiennikiem dla nowych obciążeń deep learning. Eliminuje klaster Spark i zapewnia tę samą ścieżkę kodu od pojedynczego GPU po wiele GPU.

Funkcja serverless_gpu @distributed 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 do akceleratora 8xH100 i ustaw gpus=8 zamiast tego.
  • Przenieś ładowanie danych do udekorowanej funkcji. Zobacz ładowanie danych.
  • Zachowaj istniejący kod modelu DDP, FSDP lub DeepSpeed. Dekorator wspiera wszystkie trzy.

@distributed działa na jednym węźle (patrz Ograniczenia), więc nie zastępuje wszystkich obciążeń TorchDistributor. Zachowaj zadania, które zależą od integracji ze Spark, w TorchDistributor. Aby uruchomić rozproszone szkolenie z lokalnego komputera lub na wielu węzłach, użyj CLI AI Runtime, które jest dostępne w trybie Public Preview. 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 bibliotekę serverless_gpudistributed i moduł.

    import serverless_gpu
from serverless_gpu 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(gpus=8, gpu_type='H100')
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-parellel 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 serverless_gpu import distributed

# This may cause a pickle error because the dataset is captured by the function.
dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type='H100')
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

  • Rozproszone treningi przebiegają na wszystkich GPU na jednym węźle, do którego podłączony jest Twój notebook. Aby trenować w pełni na wielu GPU, połącz się z akceleratorem 8x H100, który udostępnia jeden węzeł z 8 procesorami GPU, i ustaw gpus=8.
  • Typ akceleratora musi pasować. Jeśli ustawisz gpu_type w @distributed, musi on odpowiadać akceleratorowi, do którego jest podłączony Twój notebook ("H100" lub "A10"). Niedopasowanie powoduje awarię obciążenia. Parametr jest opcjonalny i automatycznie wykrywany po pominięciu parametru.
  • AI Runtime zaleca środowisko GPU v4 i wyższe. 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.

Dowiedz się więcej

  • Last updated on