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. Achter de schermen worden de omgeving en de communicatiekanalen tussen de werkers 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.

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 alle gedistribueerde code binnen één trainingsfunctie valt 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 pickle-fouten voorkomen. Bovendien worden de device_id waaraan modellen en gegevens zijn gekoppeld, bepaald door:

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

4. Gedistribueerde training starten: Instantieer TorchDistributor met de 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 executors.

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"

Gloo-fout: RuntimeError: Connection refused

U kunt deze fout tegenkomen bij het gebruik van Gloo voor gedistribueerde training op CPU-exemplaren. Als u deze fout wilt oplossen, voegt u het volgende codefragment toe aan uw trainingscode:

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

Voorbeeld notebooks

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

Van begin tot eind gedistribueerde training op Databricks-notebooks

Notebook ophalen

Gedistribueerde fijnafstemming van een Hugging Face-model notitieblok

Notebook ophalen

Gedistribueerde training op een PyTorch-bestandsnotitieblok

Notebook ophalen

Gedistribueerde training met PyTorch Lightning-notebook

Notebook ophalen