Úloha s více GPU

Důležité

Tato funkce je v beta verzi. Správci pracovního prostoru můžou řídit přístup k této funkci ze stránky Previews . Podívejte se na Správa náhledů Azure Databricks.

Distribuované úlohy můžete spouštět napříč několika GPU na jednom uzlu pomocí bezserverového rozhraní GPU Python API. Rozhraní API poskytuje jednoduché sjednocené rozhraní, které abstrahuje podrobnosti o zřizování GPU, nastavení prostředí a distribuci úloh. S minimálními změnami kódu můžete bezproblémově přejít z trénování s jedním GPU na distribuované spuštění na více GPU ze stejného poznámkového bloku.

Podporované architektury

Rozhraní @distributed API se integruje s hlavními distribuovanými trénovacími knihovnami:

PyTorch Distributed Data Parallel (DDP) – standardní datová paralelnost pro více GPU.
Plně dělená datová paralelnost (FSDP) – paměťově efektivní trénink pro velké modely.
DeepSpeed – optimalizační knihovna Microsoftu pro trénování velkých modelů

serverless_gpu API vs. TorchDistributor

Následující tabulka porovnává serverless_gpu@distributed rozhraní API s TorchDistributorem:

funkce	`serverless_gpu` `@distributed` Rozhraní api	Distributor pochodní
Infrastruktura	Plně bezserverová, bez správy clusteru	Vyžaduje cluster Spark s pracovními procesy GPU.
Setup	Jeden dekorátor, minimální konfigurace	Vyžaduje cluster Spark a nastavení TorchDistributor.
Podpora rámce	PyTorch DDP, FSDP, DeepSpeed	Primárně PyTorch DDP
Načítání dat	Uvnitř dekorátoru se používají svazky Unity Catalog.	Přes Spark nebo systém souborů

Rozhraní serverless_gpu API je doporučeným přístupem pro nové úlohy hlubokého učení v Databricks. TorchDistributor zůstává k dispozici pro úlohy úzce propojené s clustery Spark.

Rychlý start

Bezserverové rozhraní API GPU pro distribuované trénování je předinstalováno, když jste připojeni k bezserverovému GPU v poznámkových blocích a úlohách Databricks. Doporučujeme prostředí GPU 4 a vyšší. Pokud ho chcete použít pro distribuované učení, importujte a použijte distributed dekorátor k distribuci funkce učení.

Vložte kód pro trénování modelu do funkce a ozdobte funkci dekorátorem @distributed. Zdobená funkce se stane vstupním bodem pro distribuované spouštění – všechny trénovací logiky, načítání dat a inicializace modelu by měly být definovány uvnitř této funkce.

Výstraha

Parametr gpu_type musí odpovídat typu akcelerátoru, ke kterému je připojen váš notebook. Například @distributed(gpus=8, gpu_type='H100') vyžaduje, aby byl váš poznámkový blok připojen k akcelerátoru H100. Použití neshodného typu akcelerátoru (například připojení k A10 při zadávání H100) způsobí selhání úlohy.

Následující fragment kódu ukazuje základní použití @distributed:

# Import the distributed decorator
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs and GPU type
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

Níže je úplný příklad, který trénuje model MLP (multilayer perceptron) na gpu 8 H100 z poznámkového bloku:

Nastavte model a definujte funkce nástroje.


# Define the model
import os
import torch
import torch.distributed as dist
import torch.nn as nn

def setup():
    dist.init_process_group("nccl")
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

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)

Importujte knihovnu serverless_gpu a distribuovaný modul.

import serverless_gpu
from serverless_gpu import distributed

Vložte kód pro trénování modelu do funkce a ozdobte funkci dekorátorem @distributed.

@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()

Spusťte distribuované trénování voláním distribuované funkce s uživatelsky definovanými argumenty.
```
run_train.distributed(num_epochs=3, batch_size=1)
```
Po spuštění bude ve výstupu buňky poznámkového bloku vygenerován odkaz na běh MLflow. Pokud chcete zobrazit výsledky spuštění, klikněte na odkaz spuštění MLflow nebo ho najděte na panelu Experiment .

Podrobnosti distribuovaného spuštění

Bezserverové rozhraní GPU API se skládá z několika klíčových komponent:

Správce výpočetních prostředků: Zpracovává přidělování a správu prostředků
Běhové prostředí: Spravuje prostředí a závislosti Pythonu.
Spouštěč: Orchestrace provádění a monitorování úloh

Při spuštění v distribuovaném režimu:

Funkce se serializuje a distribuuje napříč zadaným počtem GRAFICKÝch procesorů.
Každý GPU spouští kopii funkce se stejnými parametry.
Prostředí se synchronizuje napříč všemi grafickými procesory.
Výsledky se shromažďují a vrací ze všech grafických procesorů.

Rozhraní API podporuje oblíbené paralelní trénovací knihovny, jako je DDP (Distributed Data Parallel ), Fully Sharded Data Parallel (FSDP), DeepSpeed.

V příkladech poznámkových bloků najdete více skutečných distribuovaných trénovacích scénářů pomocí různých knihoven.

FAQs

Kam se má kód načítání dat umístit?

Při použití Serverless GPU API pro distribuované trénování přesuňte kód načítání dat do dekorátoru @distributed. Velikost datové sady může překročit maximální velikost povolenou funkcí pickle, takže je doporučeno vygenerovat datovou sadu uvnitř dekorátoru, jak je znázorněno níže:

from serverless_gpu import distributed

# this may cause pickle error
dataset = get_dataset(file_path)
@distributed(gpus=8, gpu_type='H100')
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

Další informace

Referenční informace k rozhraní API najdete v dokumentaci k bezserverové verzi rozhraní GPU Python API .

Váš názor

Byla tato stránka užitečná?

Last updated on 2026-03-21