Lightning-AI/pytorch-lightning

Integrating pytorch XLA when using multiple GPUs

Open

#16 130 ouverte le 20 déc. 2022

Voir sur GitHub
 (13 commentaires) (4 réactions) (0 assignés)Python (3 233 forks)batch import
accelerator: cudafeaturehelp wantedstrategy: xla

Métriques du dépôt

Stars
 (26 687 stars)
Métriques de merge PR
 (Merge moyen 9j 15h) (3 PRs mergées en 30 j)

Description

Description & Motivation

I've experienced with pytorch XLA using multitple NVIDIA A100 GPU and I observed that in most cases training is faster. So it would be really nice to have the option to use XLA for training in pytorch lightning.

The main advantage is faster training.

Additional context

Here is a code link : https://github.com/Dhouib-med/Test-XLA/blob/17e5b6bd6c77fffa67818462856277a57877ff3b/test_xla.py to train a simple CNN on the MNIST dataset using XLA (on 2 GPUS). The main parts where taken from https://github.com/pytorch/xla. This wheel needs to be installed along with adequate pytorch and torchvision versions (1.11 and 0.14) https://storage.googleapis.com/tpu-pytorch/wheels/cuda/112/torch_xla-1.13-cp37-cp37m-linux_x86_64.whl @justusschock

cc @borda @justusschock @awaelchli @carmocca @JackCaoG @steventk-g @Liyang90

Guide contributeur

Direction de recherche
Étudiez l'implémentation de la stratégie multi GPU de PyTorch Lightning et l'intégration du périphérique XLA. Examinez le support XLA existant dans Lightning et identifiez où ajouter la gestion multi GPU XLA.
Stack technique
pythonpytorch
Domaine
machine learningai
Type d'issue
Fonctionnalité
DifficultéDifficulté estimée pour un nouveau contributeur, de 1 pour un très petit changement à 5 pour un travail expert.
3
Temps estiméFourchette de temps approximative pour investiguer, implémenter, tester et préparer une pull request.
1-2 jours
Statut d'activitéDisponibilité apparente de l'issue : fraîche, active, ancienne, bloquée ou en attente d'un mainteneur.
Active
ClartéClarté avec laquelle l'issue explique le changement attendu, les critères d'acceptation et la prochaine étape.
Claire
Prérequis
PythonPyTorchPyTorch LightningCUDA
Accessibilité débutantScore de 1 à 100 estimant l'accessibilité de cette issue pour un premier contributeur.
45