Gedistribueerde training met TorchDistributor

In dit artikel wordt beschreven hoe u gedistribueerde training uitvoert op PyTorch ML-modellen met behulp van TorchDistributor.

TorchDistributor is een opensource-module in PySpark waarmee gebruikers gedistribueerde training kunnen uitvoeren met PyTorch in hun Spark-clusters, zodat u hiermee PyTorch-trainingstaken kunt starten als Spark-taken. Onder de schermen wordt de omgeving en de communicatiekanalen tussen de werkrollen geïnitialiseerd en wordt de CLI-opdracht torch.distributed.run gebruikt om gedistribueerde training uit te voeren op de werkknooppunten.

De TorchDistributor-API ondersteunt de methoden die worden weergegeven in de volgende tabel.

Methode en handtekening	Beschrijving
init(self, num_processes, local_mode, use_gpu)	Maak een exemplaar van TorchDistributor.
run(self, main, *args)	Voert gedistribueerde training uit door aan te main(**kwargs) roepen of de hoofdfunctie een functie is en voert de CLI-opdracht torchrun main *args uit als main een bestandspad is.

Vereisten

• Spark 3.4
• Databricks Runtime 13.0 ML of hoger

Ontwikkelwerkstroom voor notebooks

Als het proces voor het maken en trainen van modellen volledig plaatsvindt vanuit een notebook op uw lokale computer of een Databricks Notebook, hoeft u alleen kleine wijzigingen aan te brengen om uw code gereed te maken voor gedistribueerde training.

1. Code voor één knooppunt voorbereiden: bereid de code van één knooppunt voor met PyTorch, PyTorch Lightning of andere frameworks die zijn gebaseerd op PyTorch/PyTorch Lightning, zoals de HuggingFace Trainer API.

2. Code voorbereiden voor standaard gedistribueerde training: u moet uw training voor één proces converteren naar gedistribueerde training. Zorg ervoor dat deze gedistribueerde code is allemaal in één trainingsfunctie die u kunt gebruiken met de TorchDistributor.

3. Importbewerkingen verplaatsen binnen de trainingsfunctie: Voeg de benodigde importbewerkingen toe, zoals import torch, in de trainingsfunctie. Hierdoor kunt u veelvoorkomende kiezerfouten voorkomen. Bovendien worden de device_id modellen en gegevens waaraan modellen en gegevens zijn gekoppeld, bepaald door:

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

4. Gedistribueerde training starten: instantieer de TorchDistributor gewenste parameters en roep .run(*args) aan om de training te starten.

Hier volgt een voorbeeld van trainingscode:

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)

Training migreren vanuit externe opslagplaatsen

Als u een bestaande gedistribueerde trainingsprocedure hebt opgeslagen in een externe opslagplaats, kunt u eenvoudig migreren naar Azure Databricks door het volgende te doen:

1. Importeer de opslagplaats: importeer de externe opslagplaats als een Databricks Git-map.
2. Maak een nieuw notebook initialiseer een nieuw Azure Databricks Notebook in de opslagplaats.
3. Start gedistribueerde training in een notebookcel en roep TorchDistributor deze als volgt aan:

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)

Probleemoplossing

Een veelvoorkomende fout voor de notebookwerkstroom is dat objecten niet kunnen worden gevonden of geselecteerd bij het uitvoeren van gedistribueerde training. Dit kan gebeuren wanneer de importinstructies van de bibliotheek niet worden gedistribueerd naar andere uitvoerders.

Als u dit probleem wilt voorkomen, moet u alle importinstructies (bijvoorbeeld import torch)opnemen boven aan de trainingsfunctie die wordt aangeroepen met TorchDistributor(...).run(<func>) en binnen alle andere door de gebruiker gedefinieerde functies die in de trainingsmethode worden aangeroepen.

NCCL-fout: ncclInternalError: Internal check failed.

Wanneer u deze fout tegenkomt tijdens het trainen met meerdere knooppunten, geeft dit meestal een probleem aan met netwerkcommunicatie tussen GPU's. Dit probleem treedt op wanneer NCCL (NVIDIA Collective Communications Library) bepaalde netwerkinterfaces voor GPU-communicatie niet kan gebruiken.

Als u deze fout wilt oplossen, voegt u het volgende codefragment toe aan uw trainingscode om de primaire netwerkinterface te gebruiken.

import os
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"

Voorbeeldnotebooks

In de volgende notebookvoorbeelden ziet u hoe u gedistribueerde training kunt uitvoeren met PyTorch.

End-to-end gedistribueerde training op Databricks-notebook

Notebook downloaden

Gedistribueerde afstemming van een notitieblok voor het Face-model hugging

Notebook downloaden

Gedistribueerde training op een PyTorch-bestandsnotitieblok

Notebook downloaden

Gedistribueerde training met PyTorch Lightning-notebook

Notebook downloaden