Formation distribuée avec TorchDistributor
Cet article décrit comment effectuer l'entraînement distribué sur des modèles de ML PyTorch à l'aide de TorchDistributor.
TorchDistributor est un module open source dans PySpark qui aide les utilisateurs à effectuer un entraînement distribué avec PyTorch sur leurs clusters Spark, il vous permet donc de lancer des Jobs d'entraînement PyTorch en tant que Jobs Spark. En coulisses, il initialise l'environnement et les canaux de communication entre les workers et utilise la commande CLI torch.distributed.run pour exécuter l'entraînement distribué sur les nœuds worker.
L'API TorchDistributor prend en charge les méthodes présentées dans la table suivante.
Méthode et signature. | Description |
|---|---|
| Créer une instance de TorchDistributor. |
| Exécute l'entraînement distribué en invoquant |
Exigences
- Spark 3,4
- Databricks Runtime 13.0 ML ou version ultérieure
Workflow de développement pour les Notebooks
Si le processus de création et d'entraînement du modèle se déroule entièrement à partir d'un notebook sur votre machine locale ou d'un Notebook Databricks, vous n'avez qu'à apporter des modifications mineures pour préparer votre code à l'entraînement distribué.
-
**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éparer le code pour l'entraînement distribué standard : Vous devez convertir votre entraînement à processus unique en entraînement distribué. Intégrez ce code distribué dans une seule fonction d’entraînement que vous pouvez utiliser avec le
TorchDistributor. -
Déplacez les importations au sein de la fonction d'entraînement : ajoutez les importations nécessaires, telles que
import torch, au sein de la fonction d'entraînement. Cela vous permet d'éviter les erreurs courantes de « pickling ». De plus, ledevice_idauquel les modèles et les données sont liés est déterminé par :Pythondevice_id = int(os.environ["LOCAL_RANK"]) -
Lancez la formation distribuée : instanciez le
TorchDistributoravec les paramètres souhaités et appelez le.run(*args)pour lancer la formation.
Voici un exemple de code d'entraînement :
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 l'entraînement depuis des repository externes
Si vous avez une procédure de formation distribuée existante stockée dans un repository externe, vous pouvez facilement migrer vers Databricks en effectuant les opérations suivantes :
- Importer le repository : importez le repository externe en tant que dossier Git Databricks.
- Créer un nouveau Notebook Initialisez un nouveau Notebook Databricks dans le repository.
- Lancer la formation distribuée Dans une cellule de Notebook, appelez
TorchDistributorcomme 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 pour le workflow de notebook est que les objets ne peuvent pas être trouvés ou sérialisés lors de l'exécution d'un entraînement distribué. Cela peut se produire lorsque les instructions d'importation de bibliothèque ne sont pas distribuées à d'autres exécuteurs.
Pour éviter ce problème, incluez toutes les déclarations d'importation (par exemple,)import torch *à la fois* en haut de la fonction d'entraînement 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 d'entraînement.
Échec CUDA : peer access is not supported between these two devices
Il s'agit d'une erreur potentielle sur la suite G5 de GPU sur AWS. Pour résoudre cette erreur, ajoutez l'extrait de code suivant dans votre code d'entraînement :
import os
os.environ["NCCL_P2P_DISABLE"] = "1"
Échec de NCCL : ncclInternalError: Internal check failed.
Lorsque vous rencontrez cette erreur pendant l'entraînement multinœud, elle indique généralement un problème de communication réseau entre les GPU. Ce problème survient lorsque NCCL (NVIDIA Collective Communications Library) ne peut pas utiliser certaines interfaces réseau pour la communication GPU.
Pour résoudre cette erreur, ajoutez l'extrait suivant dans votre code d'entraînement afin d'utiliser l'interface réseau principale.
import os
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
Défaillance de Gloo : RuntimeError: Connection refused
Vous pourriez potentiellement rencontrer cette erreur en utilisant Gloo pour l'entraînement distribué sur des instances de CPU. Pour résoudre cette erreur, ajoutez l'extrait de code suivant dans votre code d'entraînement :
import os
os.environ["GLOO_SOCKET_IFNAME"] = "eth0"
Exemples de Notebooks
Les exemples de notebook suivants montrent comment effectuer une formation distribuée avec PyTorch.