Distribuované trénování v noteboocích

Důležité

Tato funkce je v beta verzi.

Dekorátor @distributed z Databricks AI Runtime SDK (databricks-sdk-air) je nejpohodlnější způsob, jak spustit distribuované trénování z notebooku Databricks. Opatřete svou trénovací funkci dekorátorem, zavolejte ji a AI Runtime ji spustí na všech GPU na uzlu, ke kterému je váš notebook připojen. Stejný kód lze škálovat od jednoho GPU až po více GPU, aniž je nutné zřizovat cluster nebo konfigurovat distribuovaný spouštěč.

Klíčové body na této stránce:

• Dekorátor @distributed spustí trénovací funkci na všech GPU ve vašem uzlu přímo z notebooku.
• Podporuje PyTorch DDP, FSDP a DeepSpeed a přesouvá kód s jednou GPU na multi-GPU s minimálními změnami.
• Připojte prostředí notebooku k akcelerátoru 8xH100 pro kompletní trénink na více GPU. Ve výchozím nastavení @distributed používá všechna GPU na uzlu.

Note

Tato stránka se zabývá distribuovaným tréninkem z notebooků Databricks s Databricks AI Runtime SDK. Pro odesílání distribuovaných tréninkových zátěží z vašeho lokálního stroje použijte příkazy Databricks CLI pro AI Runtime, které jsou ve veřejném náhledu. Viz Use the Databricks CLI with AI Runtime.

Rychlý start

Příklady na této stránce používají Databricks AI Runtime SDK, které je předinstalováno v GPU prostředí verze 6 a vyšší. Ve verzi 5 a nižší použijte místo toho balíček serverless_gpu. Podívejte se na Migrace z serverless_gpu. Pro připojení notebooku k akcelerátoru viz Připojit se k notebooku.

Ozdobte svou tréninkovou funkci , @distributeda pak ji .distributed()nazvěte:

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

Distribuované trénování vyžaduje akcelerátor 8xH100, který zřídí jeden uzel s 8 GPU. Ve výchozím nastavení @distributed používá všech osm GPU. Nastavte num_accelerators tak, aby se používalo méně. Parametr accelerator_type je také volitelný a automaticky detekován podle akcelerátoru, ke kterému je notebook připojen.

Každé volání .distributed() vytvoří run MLflow (nebo vnořený podřízený run, pokud už je jeden aktivní) a vypíše odkaz na run ve výstupu buňky. Pro kompletní, průchodný návod viz Plný příklad.

Podporované architektury

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

• PyTorch Distributed Data Parallel (DDP): Standardní více-GPU datový paralelismus.
• Plně shardovaný datový paralelismus (FSDP): Paměťově efektivní trénování velkých modelů.
• DeepSpeed: optimalizační knihovna společnosti Microsoft pro trénování velkých modelů.

Pro reálné tréninkové scénáře, které využívají každou knihovnu, viz příklady distribuovaného tréninku.

Jak dekoratér @distributed funguje

Když zavoláte dekorovanou funkci pomocí .distributed(), AI Runtime spravuje mechaniky, které byste jinak konfigurovali ručně pomocí distribuovaného launcheru:

• Serializace a fan-out: Funkce je serializována a spuštěna na každém vybraném GPU. Každá GPU spustí kopii funkce se stejnými argumenty.
• Synchronizace prostředí: Prostředí Python a závislosti jsou replikovány napříč všemi úrovněmi, takže každý proces spouští stejný kód.
• Řada proměnných prostředí: Standardní proměnné, jako jsou , LOCAL_RANK jsou vyplňovány pro každý proces. Přečtěte si je ve své funkci, abyste umístili model a data na správné zařízení.
• Sběr výsledků: Hodnoty návratu jsou shromažďovány ze všech pořadí a vráceny volajícímu.
• Sledování v MLflow: Každé volání .distributed() vytvoří běh v MLflow, nebo vnořený podřízený běh, pokud je již nějaký běh aktivní, takže metriky zaznamenané vaší funkcí se uloží do stejného běhu.
• Životní cyklus a časový limit: Distribuované provádění probíhá v rámci životního cyklu notebooku. Ukončení poznámkového bloku ukončí běh. Dekoratér má výchozí časovou pauzu 3 hodiny. Zadejte timeout v sekundách, chcete-li to změnit, nebo timeout=None, chcete-li to vypnout. Vlastní časové limity vyžadují GPU prostředí verze 5 a vyšší.

Rozhraní API vychází ze standardních knihoven PyTorch: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) a DeepSpeed.

Migrovat z serverless_gpu

Balíček serverless_gpu je již zastaralý a už nedostává nové funkce. Počínaje verzí 6 prostředí GPU jej nahrazuje Databricks AI Runtime SDK (databricks-sdk-air) a poskytuje ekvivalentní API s přejmenovanými argumenty v oboru názvů databricks.air. Balíček, který používáte, závisí na verzi prostředí GPU:

• Verze 5 a nižší: Zahrňte pouze serverless_gpu. Pokračujte v používání serverless_gpu v těchto verzích.
• Verze 6: Obsahuje oba balíčky, ale aktualizace dostává pouze AI Runtime SDK. Verze 6 je poslední verzí GPU prostředí, která obsahuje serverless_gpu.

AI Runtime SDK je samostatný balíček od Databricks SDK pro Python (databricks-sdk).

	serverless_gpu (zastaralé)	AI runtime SDK (databricks.air)
Verze prostředí GPU	6 a méně	6 a výše
Import	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argument pro počet GPU	gpus	num_accelerators
Argument typu GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Hodnoty typů GPU	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", , "GPU_8xH100""GPU_8xB300"
Modul pro načítání dat a ukládání kontrolních bodů	serverless_gpu.data	databricks.air.data. Viz též Načtení dat v AI Runtime.

Pro migraci zápisníku:

1. Připojte notebook k GPU prostředí verze 6 nebo vyšší, které zahrnuje AI Runtime SDK.
2. Změňte importy z serverless_gpu na databricks.air.
3. Přejmenujte gpus na num_accelerators a gpu_type na accelerator_type, a použijte nové hodnoty typu akcelerátoru. Staré názvy argumentů a hodnoty typu GPU stále fungují v databricks.air, ale vyvolají FutureWarning. Podpora pro verzi 7 prostředí GPU skončí.

Tip

Řekněte Genie Code , aby to udělal za vás:

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

Pro zátěže Databricks CLI, které používají databricks.air.data, přidejte databricks-sdk-air[data] do environment.dependencies zátěže.

Přichází z TorchDistributor

Pokud spustíte distribuované trénování pomocí TorchDistributor na Spark clusterech a vaše pracovní zátěž se vejde na jeden uzel, rozhraní API @distributed je doporučenou náhradou v AI Runtime. Váš kód modelu DDP, FSDP nebo DeepSpeed se nemění.

funkce	@distributed 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í objemy Unity Catalog (UCVolumeDataset pro streamování souborových dat).	Přes Spark nebo systém souborů

Chcete-li migrovat jednouchlovou zátěž:

• Nahraďte volání @distributed dekorátorem train_fn na TorchDistributor(...).run(train_fn, ...) a potom spusťte pomocí train_fn.distributed(...).
• Odstraňte konfiguraci clusteru Spark a GPU worker. Místo toho připojte notebook k akcelerátoru 8xH100. @distributed ve výchozím nastavení používá všech osm GPU.
• Přesuňte načítání dat dovnitř dekorované funkce. Viz Načítání dat.
• Uchovejte si stávající kód modelu DDP, FSDP nebo DeepSpeed. Dekoratér podporuje všechny tři.

@distributed běží na jednom uzlu (viz Omezení), takže nenahrazuje všechny pracovní zátěže TorchDistributor. Pracovní zátěže, které závisí na integraci se Sparkem, si nechte na TorchDistributor. Pro spuštění distribuovaného tréninku z vašeho lokálního stroje nebo přes více uzlů použijte místo toho AI Runtime CLI, které je ve veřejném náhledu. Viz Use the Databricks CLI with AI Runtime.

Úplný příklad

Následující příklad trénuje model vícevrstvého perceptronu (MLP) na 8 GPU H100 v notebooku.

1. 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():
       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. Importujte distributed dekorátor z AI Runtime SDK.

   from databricks.air import distributed
   

3. Vložte kód pro trénování modelu do funkce a ozdobte funkci dekorátorem @distributed. Dekorovaná funkce je vstupním bodem pro distribuované vykonání, proto v ní definujte veškerou trénovací logiku, načítání dat a inicializaci modelu.

   @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. Spusť distribuované trénování voláním distribuované funkce s uživatelsky definovanými argumenty.

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

5. Při spuštění se ve výstupu buňky notebooku vygeneruje 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 o přizpůsobení názvů experimentů, sledování metrik a obnovení spuštění najdete v tématu Sledování experimentů a pozorovatelnost.

Načítání dat

Vložte kód pro načítání dat přímo do @distributed funkce. Datová sada může překročit maximální velikost povolenou funkcí pickle, takže její generování nebo načítání v dekorátoru zabrání chybám serializace:

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

Pro data založená na souborech uložená ve volumech Unity Catalog použijte UCVolumeDataset z databricks.air.data, který streamuje soubory s lokálním ukládáním do mezipaměti a automaticky je rozděluje mezi ranky a pracovní procesy. Pro uložení kontrolního bodu distribuovaného trénování do svazku použijte UCVolumeWriter a UCVolumeReader. Viz Načtení dat v AI Runtime a Checkpoint s distribuovaným kontrolním bodem (DCP).

Limitations

• @distributed podporuje pouze jeden uzel: běží napříč GPU na uzlu, ke kterému je váš notebook připojen. Pro trénování napříč více uzly použijte AI Runtime Databricks CLI. Viz Use the Databricks CLI with AI Runtime.
• Pokud nastavíte accelerator_type, musí odpovídat akcelerátoru, ke kterému je váš notebook připojen (například "GPU_8xH100"). Nesoulad způsobuje selhání pracovní zátěže. Tento parametr je volitelný a automaticky detekuje, pokud je vynechán.
• AI Runtime podporuje GPU prostředí v4 a vyšší. Vlastní časové limity (parametr timeout) vyžadují prostředí GPU verze 5 a vyšší.
• Dekorátor má ve výchozím nastavení časový limit 3 hodiny. Zadejte timeout v sekundách, chcete-li to změnit, nebo timeout=None, chcete-li to vypnout.
• Provádění probíhá v rámci životního cyklu notebooku. Ukončení poznámkového bloku ukončí běh.

Dodatečné zdroje

• Pro @distributed dekorátory, AcceleratorType a Ray API viz referenci k API Databricks AI Runtime SDK.
• Pro vzory, které zefektivní a odolnější váš tréninkový proces, si přečtěte Průvodce výkonem a odolností.
• Pro scénáře trénování end-to-end viz příklady distribuovaného trénování.