Entraînement distribué avec TorchDistributor
Cet article décrit comment effectuer une formation distribuée sur les modèles PyTorch ML à l'aide de TorchDistributor.
TorchDistributor est un module open source de PySpark qui aide les utilisateurs à effectuer une formation distribuée avec PyTorch sur leurs clusters Spark, il vous permet donc de lancer des tâches de formation PyTorch en tant que tâches Spark. Sous le capot, il initialise l’environnement et les canaux de communication entre les workers et utilise la commande torch.distributed.run CLI pour exécuter une formation distribuée sur les nœuds Worker.
L'API TorchDistributor prend en charge les méthodes présentées dans le tableau suivant.
| Méthode et signature | Description |
|---|---|
init(self, num_processes, local_mode, use_gpu) |
Créez une instance de TorchDistributor. |
run(self, main, *args) |
Exécute la formation distribuée en appelant main(**kwargs) si main est une fonction et exécute la commande CLI torchrun main *args si main est un chemin de fichier. |
Spécifications
- Spark 3.4
- Databricks Runtime 13.0 ML ou supérieur
Workflow de développement pour les notebooks
Si le processus de création et de formation du modèle s'effectue entièrement à partir d'un bloc-notes sur votre ordinateur local ou d'un Notebook Databricks, vous n'avez qu'à apporter des modifications mineures pour préparer votre code pour la formation distribuée.
Préparer le code de nœud unique : préparez et testez le code de nœud unique avec PyTorch, PyTorch Lightning ou d'autres frameworks basés sur PyTorch/PyTorch Lightning comme l'API HuggingFace Trainer.
Préparez le code pour la formation distribuée standard : vous devez convertir votre formation à processus unique en formation distribuée. Intégrez ce code distribué dans une seule fonction de formation que vous pouvez utiliser avec le
TorchDistributor.Déplacer les importations dans la fonction de formation : ajoutez les importations nécessaires, telles que
import torch, dans la fonction de formation. Cela vous permet d’éviter les erreurs courantes de décapage. De plus, la manièredevice_iddont les modèles et les données sont liés est déterminée par :device_id = int(os.environ["LOCAL_RANK"])Lancer une formation distribuée : instanciez le
TorchDistributoravec les paramètres souhaités et appelez.run(*args)pour lancer la formation.
Voici un exemple de code de formation :
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)
Migrer la formation à partir de référentiels externes
Si vous disposez d’une procédure de formation distribuée existante stockée dans un référentiel externe, vous pouvez facilement migrer vers Azure Databricks en procédant comme suit :
- Importer le référentiel : importez le référentiel externe en tant que Dossier Databricks GIT.
- Créer un nouveau notebook Initialisez un nouveau Notebook Azure Databricks dans le référentiel.
- Lancer une formation distribuée Dans une cellule de notebook
TorchDistributor, appelez comme suit :
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)
Dépannage
Une erreur courante dans le flux de travail des notebooks est que les objets ne peuvent pas être trouvés ou récupérés lors de l'exécution d'une formation distribuée. Cela peut se produire lorsque les instructions d'importation de bibliothèque ne sont pas distribuées aux autres exécuteurs.
Pour éviter ce problème, incluez toutes les instructions d'importation (par exemple, import torch) à la fois en haut de la fonction de formation appelée avec TorchDistributor(...).run(<func>) et à l'intérieur de toute autre fonction définie par l'utilisateur appelée dans la méthode de formation.
Exemples de notebooks
Les exemples de notebook suivants montrent comment effectuer une formation distribuée avec PyTorch.