PyTorch : Parallélisme de modèle multi-GPU

La méthodologie présentée dans cette page montre comment adapter, sur Jean Zay, un modèle trop volumineux pour tenir sur un seul GPU avec Pytorch. C'est une illustration des concepts présentés dans la page mère : Jean Zay : Distribution multi-GPU et multi-nœuds pour l'apprentissage d'un modèle TensorFlow ou PyTorch. Nous ne verrons que la version optimisée de modèle parallélisme (Pipeline Parallelism) car la version naïve n'est vraiment pas conseillée.

La documentation pytorch est très complète pour ce sujet, il ne faudra faire que quelques légères modifications pour faire marcher le modèle parallélisme avec segmentation de batch (Pipeline Parallelism) sur Jean Zay. Si vous débutez avec le modèle parallélisme, nous vous conseillons de suivre d'abord le tutoriel pytorch à ce sujet.

Ce document présente uniquement les changements à faire pour paralléliser un modèle sur Jean Zay et présente un benchmark pour donner une idée de l'optimisation des paramètres de la méthode. Une démonstration complète (qui a servi au benchmark) se trouve dans un code python téléchargeable ici ou récupérable sur Jean-Zay dans le répertoire DSDIR : /examples_IA/Torch_parallel/bench_PP_transformer.py. Pour le copier dans votre espace $WORK :

$ cp $DSDIR/examples_IA/Torch_parallel/bench_PP_transformer.py $WORK

D'autres méthodes d'implémentation du modèle parallélisme avec pytorch existent et sont expliquées dans le lien suivant : Single-Machine Model Parallel Best Practices. Néanmoins, elles semblent moins efficaces que la méthode présentée dans cette page.

Le seul changement à réaliser pour faire marcher la méthode de la documentation pytorch sur Jean Zay se trouve au niveau de l'initialisation du Framework RPC. En effet, les options par défaut sont adaptées aux nœuds de calculs dotés d'InfiniBand et Jean Zay est équipé avec des communications Omni-Path. Pour cela, il faut changer l'option _transports=[“ibv”, “uv”] par _transports=[“shm”, “uv”] comme cela :

rpc.init_rpc(
    name="worker",
    rank=0,
    world_size=1,
    rpc_backend_options=rpc.TensorPipeRpcBackendOptions(
        init_method="file://{}".format(tmpfile.name),
        # Specifying _transports and _channels is a workaround and we no longer
        # will have to specify _transports and _channels for PyTorch
        # versions >= 1.8.1
        # With Jean Zay, _transports must be equal to ["shm", "uv"] and not ["ibv", "uv"] (like in pytorch doc)
        _transports=["shm", "uv"],
        _channels=["cuda_ipc", "cuda_basic"],
    )
)

Une seule tâche doit être déclarée pour un modèle (ou ensemble de modèles) et un nombre adéquat de GPU, qui doivent tous être sur le même nœud. Si l'on prend l'exemple d'un modèle que l'on voudrait répartir sur 4 gpus, les options suivantes sont exigées :

#SBATCH --gres=gpu:4
#SBATCH --ntasks-per-node=1

La méthodologie présentée, qui ne dépend que de la bibliothèque PyTorch, se limite à un parallélisme mono-nœud multi-GPU (de 2 GPU, 4 GPU ou 8 GPU) et ne peut s'appliquer à un cas multi-nœuds. Il est vivement conseillé d'associer cette technique avec du parallélisme de données, comme décrit dans la page Parallélisme des données et de modèle avec PyTorch, pour accélérer efficacement les entraînements.

Si votre modèle nécessite un parallélisme de modèle sur plusieurs nœuds, nous vous invitons à explorer les solutions documentées sur la page Distributed Pipeline Parallelism Using RPC ou à utiliser d'autres méthodes comme Deepspeed ou Megatron.

Un benchmark a été réalisé pour comparer le Pipeline Parallelism (avec différents paramètres) avec le modèle parallélisme naïf. Les tests ont été réalisés sur 2 et 4 gpus, donc avec un modèle coupé en deux ou en quatre parties. Le modèle utilisé est un Transformer qui est entraîné sur le dataset IMDB Review (le but est de mesurer les performances de l’entraînement, utiliser un Transformer pour une tâche si peu complexe est déconseillé). Différents batch sizes ont aussi été testés, de 32 à 256, combinés avec différents chunks, de 1 (parallélisme naïf) à 16. Les résultats sont les suivants (lorsqu'une donnée est manquante c'est qu'un OOM a eu lieu) :

On remarque que le Pipeline Parallelism est toujours plus performant (en temps d’exécution et en mémoire) que le modèle parallélisme lorsque le bon chunk est trouvé. Il faudra donc faire attention à bien accorder le nombre de chunks avec la taille du batch. Attention, le gain en mémoire est du au gradient checkpointing qui est appliqué par défaut sur les modèles Pipe.

⇐ Revenir à la page mère sur l'apprentissage distribué