Sdílet prostřednictvím

Distribuované trénování s využitím TorchDistributoru

Tento článek popisuje, jak provádět distribuované trénování na modelech PyTorch ML pomocí TorchDistributor.

TorchDistributor je opensourcový modul v PySparku, který uživatelům pomáhá provádět distribuované trénování pomocí PyTorchu v jejich clusterech Spark, takže umožňuje spouštět trénovací úlohy PyTorch jako úlohy Sparku. Na pozadí inicializuje prostředí a komunikační kanály mezi pracovními procesy a využívá příkaz torch.distributed.run rozhraní příkazového řádku ke spuštění distribuovaného trénování napříč pracovními uzly.

Rozhraní API TorchDistributor podporuje metody uvedené v následující tabulce.

Metoda a podpis Popis
init(self, num_processes, local_mode, use_gpu) Vytvořte instanci TorchDistributor.
run(self, main, *args) Spustí distribuované trénování vyvoláním main(**kwargs), pokud je 'main' funkcí, a spuštěním příkazu torchrun main *args, pokud je 'main' cestou k souboru.

Požadavky

  • Spark 3.4
  • Databricks Runtime 13.0 ML nebo vyšší

Pracovní postup vývoje pro notebooky

Pokud proces vytváření a trénování modelu probíhá přímo v poznámkovém sešitu na vašem místním počítači nebo v Databricks Notebook, postačí provést drobné úpravy, abyste mohli připravit kód pro distribuované trénování.

  1. Příprava kódu s jedním uzlem: Připravte a otestujte kód s jedním uzlem pomocí PyTorchu, PyTorch Lightning nebo jiných architektur založených na PyTorch/PyTorch Lightning, jako je rozhraní API HuggingFace Trainer.

  2. Příprava kódu pro standardní distribuované trénování: Potřebujete převést trénování jednoho procesu na distribuované trénování. Mít tento distribuovaný kód zcela zahrnutý v jedné trénovací funkci, kterou můžete použít s TorchDistributor.

  3. Přesun importů v rámci trénovací funkce: Přidejte potřebné importy, například import torch, v rámci trénovací funkce. Díky tomu se můžete vyhnout běžným chybám při výběru. Kromě toho je určeno, k čemu jsou modely a data device_id svázány:

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

  4. Spusťte distribuované trénování: Vytvořte instanci TorchDistributor s požadovanými parametry a volání .run(*args) pro spuštění trénování.

Následuje příklad trénovacího kódu:

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)

Migrace školení z externích úložišť

Pokud máte existující distribuovanou trénovací proceduru uloženou v externím úložišti, můžete snadno migrovat do Azure Databricks následujícím způsobem:

  1. Import úložiště: Importujte externí úložiště jako složku Databricks Git.
  2. Vytvořte nový poznámkový blok. Inicializujte nový poznámkový blok Azure Databricks v rámci úložiště.
  3. Spusťte distribuované trénování v buňce poznámkového bloku a volejte TorchDistributor takto:
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)

Řešení problému

Běžnou chybou pracovního postupu poznámkového bloku je, že při spuštění distribuovaného trénování se objekty nedají najít ani vybrat. K tomu může dojít v případě, že příkazy importu knihovny nejsou distribuovány do jiných vykonávačů.

Abyste se tomuto problému vyhnuli, zahrňte všechny příkazy importu (například import torch) jak do horní části trénovací funkce, která je volána s , tak i uvnitř každé jiné uživatelsky definované funkce, které jsou zavolány v tréninkové metodě.

Selhání NCCL: ncclInternalError: Internal check failed.

Když při trénování s více uzly narazíte na tuto chybu, obvykle to značí problém se síťovými komunikacemi mezi grafickými procesory. K tomuto problému dochází, když NCCL (NVIDIA Collective Communications Library) nemůže pro komunikaci s GPU používat určitá síťová rozhraní.

Pokud chcete tuto chybu vyřešit, přidejte do trénovacího kódu následující fragment kódu pro použití primárního síťového rozhraní.

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

Selhání Gloo: RuntimeError: Connection refused

K této chybě může dojít při použití Gloo pro distribuované trénování instancí procesoru. Pokud chcete tuto chybu vyřešit, přidejte do trénovacího kódu následující fragment kódu:

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

Příklady poznámkových bloků

Následující příklady poznámkového bloku ukazují, jak provádět distribuované trénování pomocí PyTorchu.

Kompletní distribuované trénování v poznámkovém bloku Databricks

Získat notebook

Distribuované doladění modelu Hugging Face v poznámkovém bloku

Získat notebook

Distribuované trénování v poznámkovém bloku PyTorch File

Získat notebook

Distribuované trénování s využitím poznámkového bloku PyTorch Lightning

Získat notebook