Elosztott betanítás jegyzetfüzetekben

Important

Ez a funkció bétaverzióban érhető el.

A Databricks AI Runtime SDK (databricks-sdk-air) @distributed dekorátora a legkényelmesebb módja az elosztott betanítás futtatásának egy Databricks notebookról. Díszítsd be a képzési funkciódat, hívd meg, és az AI Runtime futtatja az összes GPU-n azon a csomóponton, amelyhez a notebookod csatlakozik. Ugyanaz a kód egyetlen GPU-tól több GPU-ig skálázható, anélkül hogy klasztert kellene kiépíteni vagy elosztott indítót konfigurálni.

Főbb pontok ezen az oldalon:

• A(z) @distributed dekorátor egy notebookból futtatja a betanító függvényt a csomópont összes GPU-ján.
• Támogatja a PyTorch DDP-t, FSDP-t és DeepSpeed-et, és minimális változtatásokkal áthelyezi az egy-GPU-s kódot több-GPU-ra.
• Csatlakoztasd a notebookodat egy 8xH100-as gyorsítóhoz a teljes több-GPU-s betanításhoz. Alapértelmezés szerint a @distributed minden GPU-t használ a csomóponton.

Note

Ez az oldal a Databricks-jegyzetfüzetekben végzett elosztott tanítást fedi le a Databricks AI Runtime SDK-val. Ha a helyi gépedről szeretnél elosztott tanítási feladatokat elküldeni, használd az AI Runtime-hoz készült Databricks CLI-parancsokat. Lásd : Használd a Databricks CLI-t AI Runtime-val.

gyorskonfigurálás

Az ezen oldalon látható példák a Databricks AI Runtime SDK-t használják, amely előre telepítve van a GPU környezet 6-os verziójában vagy annál magasabb. Az 5-ös és korábbi verziókban a(z) serverless_gpu helyett a csomagot használd. Lásd: Áttelepítés innen: serverless_gpu. A notebook gyorsítóhoz való csatlakoztatásához lásd a 1. lépést: Notebook csatlakoztatása GPU-hoz.

Díszítsd meg a tréningi funkciódat @distributed, majd nevezd el a következőket .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()

Az elosztott betanításhoz 8xH100-es gyorsító szükséges, amely egyetlen csomópontot helyez üzembe 8 GPU-val. Alapértelmezés @distributed szerint mind a nyolc GPU-t használja. Állítsa a(z) num_accelerators értékét alacsonyabbra. Ez a accelerator_type paraméter opcionális, és automatikusan észlelhető abból a gyorsítóból, amelyhez a notebookod csatlakozik.

Minden .distributed() hívás létrehoz egy MLflow futást (vagy egy beágyazott gyermekfutást, ha az már aktív), és kinyomtat egy futás-linket a cella kimenetében. A teljes, futtatható útmutatóért lásd a Teljes példát.

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.
• Teljesen horizontálisan skálázott adatok párhuzamos használata (FSDP): Memória-hatékony betanítás nagy modellekhez.
• DeepSpeed: Microsoft optimalizálási kódtára a nagy modell betanításához.

A minden könyvtárat használó valós betanítási forgatókönyvekhez lásd az elosztott betanítási példákat.

Hogyan működik a @distributed dekorátor

Ha egy dekorált függvényt .distributed() meghívsz, az AI Runtime kezeli azokat a működési részleteket, amelyeket egyébként kézzel konfigurálnál egy elosztott indítóprogrammal:

• Szerializálás és fan-out: A funkció szerializálva van, és minden kiválasztott GPU-n elindul. Minden GPU futtat egy példányt a függvényből ugyanazokkal az argumentumokkal.
• Környezeti szinkronizáció: A Python környezet és függőségei minden rangon replikálódnak, így minden folyamat ugyanazt a kódot futtatja.
• Rangsorhoz tartozó környezeti változók: A szabványos környezeti változók, mint például a(z) LOCAL_RANK, minden folyamathoz be vannak állítva. Olvasd ki őket a függvényedben, hogy a modellt és az adatokat a megfelelő eszközre helyezd.
• Eredménygyűjtés: A visszaküldési értékeket minden rangról gyűjtik össze, és visszaküldik a hívónak.
• MLflow-nyomon követés: Minden .distributed() hívás létrehoz egy MLflow-futást, vagy egy beágyazott alárendelt futást, ha már van aktív futás, így a függvény által naplózott metrikák ugyanahhoz a futáshoz kerülnek naplózásra.
• Életciklus és időtúllépés: Az elosztott végrehajtás a notebook életciklusán belül fut. A notebook leállítása megszakítja a futást. A dekorátornak alapértelmezett 3 órás időkorlátja van. A módosításához add meg timeout másodpercben, a letiltásához pedig timeout=None. A testreszabott időtúllépésekhez 5-ös vagy újabb GPU-környezet szükséges.

Az API a szabványos PyTorch könyvtárakra épül: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) és DeepSpeed.

Migrálás innen: serverless_gpu

A serverless_gpu csomag elavult, és már nem kap új funkciókat. A GPU környezet 6. verziójától a Databricks AI Runtime SDK (databricks-sdk-air) váltja fel, és ekvivalens API-kat biztosít átnevezett argumentumokkal a databricks.air névtér alatt. A használt csomag a GPU környezeti verziótól függ:

• 5-ös és korábbi verziók: Csak a(z) serverless_gpu elemet foglalja bele. Továbbra is használja a(z) serverless_gpu elemet ezeken a verziókon.
• 6. verzió: Mindkét csomagot tartalmazza, de csak az AI Runtime SDK kap frissítéseket. A 6-os verzió az utolsó GPU környezeti verzió, amely tartalmazza serverless_gpu.

Az AI Runtime SDK a Databricks SDK for Pythontól különálló csomag (databricks-sdk).

	serverless_gpu (elavult)	AI Runtime SDK (databricks.air)
GPU környezeti verziók	6 és annál alacsonyabb	6 és annál magasabb
Import	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
GPU-k száma argumentum	gpus	num_accelerators
GPU-típus argumentum	gpu_type (GPUType)	accelerator_type (AcceleratorType)
GPU-típus értékei	\, \, \	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Adatbetöltési és ellenőrzőpont-kezelő modul	serverless_gpu.data	databricks.air.data. Lásd : Adatok betöltése az AI-futtatókörnyezetben.

Jegyzetfüzet migrálása:

1. Csatlakoztasd a notebookot a 6-os vagy annál magasabb verziójú GPU-környezethez, amely tartalmazza az AI Runtime SDK-t.
2. Importokat változtasd serverless_gpu-ról databricks.air-re.
3. Nevezze át a(z) gpus elemet erre: gpu_type, és a(z) accelerator_type elemet erre: num_accelerators, és használja az új gyorsítótípus-értékeket. A régi argumentusnevek és GPU-típusértékek továbbra is működnek a(z) databricks.air elemben, de FutureWarning figyelmeztetést adnak ki. A támogatás a GPU környezet 7-es verziójával megszűnik.

Tipp

Mondja meg a Genie Code-nak , hogy tegye ezt Önért:

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

Azoknál a Databricks CLI-munkaterheléseknél, amelyek a(z) databricks.air.data elemet használják, adja hozzá a(z) databricks-sdk-air[data] elemet a munkaterhelés environment.dependencies eleméhez.

A TorchDistributor-tól érkezve

Ha TorchDistributor segítségével futtatsz elosztott képzést Spark klaszterekben, és a munkaterhelés egyetlen csomópontra illeszkedik, akkor az @distributed API az AI Runtime ajánlott helyettesítője. A DDP, FSDP vagy DeepSpeed modellkódod nem változik.

Feature	@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	A dekorátoron belül a Unity Catalog-köteteket használja (UCVolumeDataset a fájladatok adatfolyamként való kezeléséhez)	Sparkon vagy fájlrendszeren keresztül

Egy egycsomópontos munkaterhelés migrálásához:

• Cseréld le a TorchDistributor(...).run(train_fn, ...) hívást a(z) train_fn elemnél a @distributed dekorátorra, majd indítsd el a train_fn.distributed(...) paranccsal.
• Távolítsd el a Spark klasztert és a GPU munkagép konfigurációját. Inkább csatlakoztatd a notebookot egy 8xH100-as gyorsítóhoz. @distributed alapértelmezés szerint mind a nyolc GPU-t használja.
• Mozgasd át az adatbetöltést a díszített funkción belül. Lásd: Adatbetöltés.
• Tartsd meg a meglévő DDP, FSDP vagy DeepSpeed modellkódodat. A dekorátor mindhármat támogatja.

A @distributed dekorátor béta állapotban van, és egyetlen csomóponton fut (lásd Korlátozásokat), így nem pótolja az összes TorchDistributor feladatot. Tartsd meg azokat a munkaterheléseket, amelyek a Spark integrációtól függenek, a TorchDistributor-on. Ha elosztott betanítást futtatsz a helyi gépedről vagy több csomópont között, használd a Databricks CLI parancsokat az AI Runtime-hoz. Lásd : Használd a Databricks CLI-t AI Runtime-val.

Teljes példa

A következő példa egy többrétegű perceptron (MLP) modellt tanít be 8 H100 GPU-n notebookból.

1. Á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():
       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. Importáld a distributed dekorátort az AI Runtime SDK-ból.

   from databricks.air import distributed
   

3. 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 díszített függvény az elosztott végrehajtás belépési pontja, így határozd meg benne az összes betanítási logikát, adatbetöltést és modell inicializálást.

   @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. A distributed traininget úgy futtatjuk, hogy a distributed függvényt hívjuk meg felhasználó által definiált argumentumokkal.

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

5. A végrehajtáskor egy MLflow-futtatási hivatkozás jön létre a jegyzetfüzet cellakimeneté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. A kísérletnevek testreszabásáról, a metrikák nyomon követéséről és a futtatások folytatásáról további információt a Kísérletkövetés és a megfigyelhetőség című témakörben talál.

Adatbetöltés

Helyezd az adatbetöltési kódot a @distributed függvénybe. Egy adathalmaz meghaladhatja a pickle által engedélyezett maximális méretet, ezért a dekorátoron belüli előállítása vagy betöltése segít elkerülni a szerializálási hibákat:

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

A Unity Catalog-kötetekben tárolt fájlalapú adatokhoz használja a(z) UCVolumeDataset elemet a(z) databricks.air.data helyről, amely helyi gyorsítótárazással streameli a fájlokat, és automatikusan particionálja őket a rangok és a feldolgozók között. Az elosztott betanítás kötetre történő ellenőrzőpont-mentéséhez használja a(z) UCVolumeWriter és UCVolumeReader elemet. Lásd: Adatok betöltése az AI Runtime környezetben és Ellenőrzőpont a Distributed Checkpoint (DCP) használatával.

Limitations

• @distributed csak egyetlen csomópontot támogat: azon a csomóponton lévő GPU-kon fut, amelyhez a jegyzetfüzeted csatlakozik. Több csomóponton való betanításhoz használja az AI Runtime Databricks CLI-t. Lásd : Használd a Databricks CLI-t AI Runtime-val.
• Ha beállítod accelerator_type, annak egyeznie kell azzal a gyorsítóval, amelyhez a Notebookod csatlakoztatva van (például "GPU_8xH100"). Az eltérés miatt a munkaterhelés sikertelen lesz. A paraméter opcionális, és automatikusan észlelhető, ha kihagyják.
• Az AI Runtime támogatja a GPU környezet v4 vagy annál magasabb verziókat. Az egyéni időtúllépésekhez (a timeout paraméterhez) 5-ös vagy újabb verziójú GPU-környezet szükséges.
• A dekorátor alapértelmezetten 3 óra után időlejár. A módosításához add meg timeout másodpercben, a letiltásához pedig timeout=None.
• A végrehajtás a jegyzetfüzet életciklusán belül fut. A notebook leállítása megszakítja a futást.

További erőforrások

• A @distributed dekorátorhoz, a AcceleratorType elemhez és a Ray API-khoz lásd a Databricks AI Runtime SDK API-referenciát.
• Azokért mintákért, amelyek hatékonyabbá és ellenállóbbá teszik az edzési folyamatodat, nézd meg a teljesítmény- és ellenállóképességi útmutatót.
• A teljes körű betanítási forgatókönyvekről lásd az elosztott betanítási példákat.