Ú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 . Viz Manage Azure Databricks preview.

Distribuované úlohy můžete spouštět napříč několika GPU na jednom uzlu pomocí Serverless 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.

Poznámka:

Distribuované trénování vyžaduje akcelerátor 8xH100, který zřídí jeden uzel s 8 GPU. Při použití dekorátoru @distributed nastavte gpus=8. Parametr gpu_type je volitelný a automaticky zjištěný z akcelerátoru, ke kterému je poznámkový blok připojený.

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

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	V dekorátoru se používají svazky služby Unity Catalog (UCVolumeDataset pro streamování souborových dat)	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é spuštění, takže v této funkci by měla být definována veškerá logika trénování, načítání dat a inicializace modelu.

Chcete-li spustit distribuované provádění, zavolejte dekorovanou funkci pomocí train_function.distributed(). Každé volání automaticky vytvoří běh experimentu v MLflow nebo vnořený podřízený běh, pokud už je některý běh aktivní.

Výstraha

Pokud nastavujete gpu_type v @distributed, ujistěte se, že odpovídá typu akcelerátoru, ke kterému je notebook připojen ("H100" nebo "A10"). Zadání nesprávného typu akcelerátoru způsobí selhání úlohy.

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

from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs.
# gpu_type='H100' is optional and will be auto-detected if not set.
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

run_train.distributed()

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

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 knihovnu serverless_gpudistributed a modul.

   import serverless_gpu
   from serverless_gpu import distributed
   

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

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

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.

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í: Řídí Python prostředí a závislosti.
• 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ů.
• Správa životního cyklu: Distribuované provádění běží v rámci životního cyklu poznámkového bloku. Když se notebook ukončí, ukončí se i provádění. Dekorátor @distributed má výchozí časový limit 3 hodiny. Chcete-li nastavit vlastní časový limit, zadejte hodnotu timeout v sekundách nebo hodnotu timeout=None, pokud ho chcete zakázat. Nastavení časového limitu je k dispozici v prostředí GPU verze 5 a vyšší.

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

Pro data založená na souborech uložená ve volumech Unity Catalog použijte UCVolumeDataset z serverless_gpu.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čítání dat v prostředí AI Runtime a ukládání kontrolních bodů modelu.

Další informace

Referenční informace k rozhraní API najdete v dokumentaci k rozhraní Serverless GPU Python API.