Több GPU-s számítási feladat

Fontos

Ez a funkció bétaverzióban érhető el. A munkaterület rendszergazdái az Előnézetek lapon szabályozhatják a funkcióhoz való hozzáférést. Lásd: Az Azure Databricks előzetes verziójának kezelése.

A kiszolgáló nélküli GPU Python API-val egyetlen csomóponton több GPU-n is elindíthat elosztott számítási feladatokat. Az API egy egyszerű, egységes felületet biztosít, amely elvonja a GPU kiépítésének, a környezet beállításának és a számítási feladatok elosztásának részleteit. Minimális kódmódosításokkal zökkenőmentesen áttérhet az egy GPU-betanításról a több GPU-s elosztott végrehajtásra ugyanabból a jegyzetfüzetből.

Támogatott keretrendszerek

Az @distributed API integrálható a főbb elosztott betanítási kódtárakkal:

PyTorch Distributed Data Parallel (DDP) – Standard több GPU-s adat-párhuzamosság.
Teljes mértékben skálázott adatok párhuzamos használata (FSDP) – Memória-hatékony betanítás nagy modellekhez.
DeepSpeed – A Microsoft optimalizálási kódtára a nagy modell betanításához.

serverless_gpu API és a TorchDistributor

Az alábbi táblázat összehasonlítja az API-t a serverless_gpu@distributedTorchDistributorral:

Funkció	`serverless_gpu` `@distributed` API	Fáklyaelosztó
Infrastruktúra	Teljesen szerver nélküli, klaszterkezelés nélkül	Gpu-feldolgozókkal rendelkező Spark-fürtöt igényel
Setup	Szimpla dekorátor, minimális konfiguráció	Spark-fürt és TorchDistributor beállítás szükséges
Keretrendszer támogatása	PyTorch DDP, FSDP, DeepSpeed	Elsősorban PyTorch DDP
Adatbetöltés	Dekoratőrön belül a Unity Catalog Volumes-t használja	Sparkon vagy fájlrendszeren keresztül

Az serverless_gpu API a Databricks új mélytanulási számítási feladatainak ajánlott megközelítése. A TorchDistributor továbbra is elérhető marad a Spark-fürtökkel szorosan összekapcsolt számítási feladatokhoz.

Gyors kezdés

Az elosztott betanítás kiszolgáló nélküli GPU API-ja előre telepítve van, ha kiszolgáló nélküli GPU-hoz csatlakozik a Databricks-jegyzetfüzetekben és -feladatokban. A 4- és újabb GPU-környezetet javasoljuk. Ha elosztott betanításhoz szeretné használni, importálja és használja a distributed dekorátort a betanítási függvény terjesztéséhez.

Csomagolja be a modell betanítási kódját egy függvénybe, és díszítse a függvényt a @distributed dekorátorsal. A kitüntetett függvény lesz az elosztott végrehajtás belépésipontja – ebben a függvényben minden betanítási logikát, adatbetöltést és modell inicializálást meg kell határozni.

Figyelmeztetés

A gpu_type paraméternek meg kell egyeznie a jegyzetfüzethez csatlakoztatott gyorsító típusával. Szükség van például arra, @distributed(gpus=8, gpu_type='H100') hogy a jegyzetfüzet egy H100-gyorsítóhoz legyen csatlakoztatva. A nem egyező gyorsítótípus használata (például az A10-hez való csatlakozás a H100 megadásakor) a számítási feladat meghiúsulását okozza.

Az alábbi kódrészlet bemutatja a @distributed alapszintű használatát.

# 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():
    ...

Az alábbiakban egy teljes példa látható, amely egy többrétegű perceptron (MLP) modellt képez be 8 H100 GPU-n egy jegyzetfüzetből:

Állítsa be a modellt, és definiálja a segédprogramfüggvényeket.


# 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)

Importálja a serverless_gpu kódtárat és az elosztott modult .

import serverless_gpu
from serverless_gpu import distributed

Csomagolja be a modell betanítási kódját egy függvénybe, és díszítse a függvényt a @distributed dekorátorsal.

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

Az elosztott betanítás végrehajtásához hívja meg az elosztott függvényt felhasználó által megadott argumentumokkal.
```
run_train.distributed(num_epochs=3, batch_size=1)
```
A végrehajtáskor egy MLflow-futtatási hivatkozás jön létre a jegyzetfüzetcella kimenetében. Kattintson az MLflow futtatási hivatkozására, vagy keresse meg a Kísérlet panelen a futtatási eredmények megtekintéséhez.

Elosztott végrehajtás részletei

A kiszolgáló nélküli GPU API több fő összetevőből áll:

Compute Manager: Kezeli az erőforrások elosztását és kezelését
Futtatókörnyezet: Python-környezetek és -függőségek kezelése
Launcher: A feladatok végrehajtásának és figyelésének vezénylése

Elosztott módban való futtatás esetén:

A függvény szerializálva van és elosztva a megadott számú GPU-ban
Minden GPU ugyanazokkal a paraméterekkel futtatja a függvény egy példányát
A környezet szinkronizálva van az összes GPU-ban
A rendszer összegyűjti és visszaadja az eredményeket az összes GPU-ból

Az API támogatja a népszerű párhuzamos betanítási kódtárakat, például az Elosztott Adatpárhuzamosítást (DDP), a Teljesen Töredezett Adatpárhuzamosítást (FSDP), a DeepSpeedet.

A jegyzetfüzet-példákban szereplő különböző kódtárak használatával valósabb elosztott betanítási forgatókönyveket találhat.

FAQs

Hol kell elhelyezni az adatbetöltési kódot?

Ha kiszolgáló nélküli GPU API-t használ az elosztott betanításhoz, helyezze át az adatbetöltési kódot a @distributed dekorátorba. Az adathalmaz mérete meghaladhatja a pickle által megengedett maximális méretet, ezért javasoljuk, hogy az adathalmazt a dekorátoron belül hozza létre, az alábbiak szerint:

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)
  ....

Tudj meg többet

Az API-referenciaért tekintse meg a Kiszolgáló nélküli GPU Python API dokumentációját.

Visszajelzés

Hasznosnak találta ezt az oldalt?

Last updated on 2026-03-21