Distribuerad träning med TorchDistributor

Den här artikeln beskriver hur du utför distribuerad träning på PyTorch ML-modeller med Hjälp av TorchDistributor.

TorchDistributor är en modul med öppen källkod i PySpark som hjälper användare att utföra distribuerad utbildning med PyTorch i sina Spark-kluster, så att du kan starta PyTorch-träningsjobb som Spark-jobb. Under huven initierar den miljön och kommunikationskanalerna mellan arbetarna och använder CLI-kommandot torch.distributed.run för att köra distribuerad utbildning över arbetsnoderna.

TorchDistributor-API:et stöder de metoder som visas i följande tabell.

Metod och signatur	beskrivning
init(self, num_processes, local_mode, use_gpu)	Skapa en instans av TorchDistributor.
run(self, main, *args)	Kör distribuerad träning genom att main(**kwargs) anropa om main är en funktion och kör CLI-kommandot torchrun main *args om main är en filsökväg.

Krav

• Spark 3.4
• Databricks Runtime 13.0 ML eller senare

Utvecklingsarbetsflöde för notebook-filer

Om modellskapandet och träningsprocessen sker helt från en notebook-fil på din lokala dator eller en Databricks Notebook behöver du bara göra mindre ändringar för att förbereda koden för distribuerad träning.

1. Förbered kod med en nod: Förbered och testa ennodskoden med PyTorch, PyTorch Lightning eller andra ramverk som baseras på PyTorch/PyTorch Lightning som HuggingFace Trainer-API:et.

2. Förbered kod för standard distribuerad träning: Du måste konvertera din enskild processträning till distribuerad träning. Låt den här distribuerade koden ingå i en träningsfunktion som du kan använda med TorchDistributor.

3. Flytta importer inom träningsfunktionen: Lägg till nödvändiga importer, till exempel import torch, i träningsfunktionen. På så sätt kan du undvika vanliga hämtningsfel. Dessutom bestäms de device_id modeller och data som är knutna till av:

   device_id = int(os.environ["LOCAL_RANK"])
   

4. Starta distribuerad träning: Instansiera TorchDistributor med önskade parametrar och anropa .run(*args) för att starta träningen.

Följande är ett exempel på träningskod:

from pyspark.ml.torch.distributor import TorchDistributor

def train(learning_rate, use_gpu):
  import torch
  import torch.distributed as dist
  import torch.nn.parallel.DistributedDataParallel as DDP
  from torch.utils.data import DistributedSampler, DataLoader

  backend = "nccl" if use_gpu else "gloo"
  dist.init_process_group(backend)
  device = int(os.environ["LOCAL_RANK"]) if use_gpu  else "cpu"
  model = DDP(createModel(), **kwargs)
  sampler = DistributedSampler(dataset)
  loader = DataLoader(dataset, sampler=sampler)

  output = train(model, loader, learning_rate)
  dist.cleanup()
  return output

distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train, 1e-3, True)

Migrera träning från externa lagringsplatser

Om du har en befintlig distribuerad träningsprocedur lagrad på en extern lagringsplats kan du enkelt migrera till Azure Databricks genom att göra följande:

1. Importera lagringsplatsen: Importera den externa lagringsplatsen som en Databricks Git-mapp.
2. Skapa en ny notebook-fil Initiera en ny Azure Databricks Notebook på lagringsplatsen.
3. Starta distribuerad träning I en notebook-cell anropar du TorchDistributor på följande sätt:

from pyspark.ml.torch.distributor import TorchDistributor

train_file = "/path/to/train.py"
args = ["--learning_rate=0.001", "--batch_size=16"]
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train_file, *args)

Felsökning

Ett vanligt fel för notebook-arbetsflödet är att objekt inte kan hittas eller picklas in när distribuerad träning körs. Detta kan inträffa när biblioteksimportinstruktionerna inte distribueras till andra köre.

Undvik det här problemet genom att ta med alla importinstruktioner (till exempel import torch) överst i träningsfunktionen som anropas med TorchDistributor(...).run(<func>) och i andra användardefinierade funktioner som anropas i träningsmetoden.

Exempelnotebook-filer

Följande notebook-exempel visar hur du utför distribuerad träning med PyTorch.

Distribuerad utbildning från slutpunkt till slutpunkt på Databricks Notebook

Hämta notebook-fil

Distribuerad finjustering av en notebook-fil för huggande ansiktsmodell

Hämta notebook-fil

Distribuerad träning på en PyTorch-filanteckningsbok

Hämta notebook-fil

Distribuerad träning med PyTorch Lightning Notebook

Hämta notebook-fil

Distribuerad datainläsning med Petastorm Notebook

Hämta notebook-fil