facebookresearch/fairseq

Python API Transformer.from_pretrained support directly to load on GPU

Open

#2 480 ouverte le 13 août 2020

Voir sur GitHub
 (6 commentaires) (0 réactions) (0 assignés)Python (6 224 forks)batch import
enhancementhelp wanted

Métriques du dépôt

Stars
 (29 107 stars)
Métriques de merge PR
 (Aucune PR mergée en 30 j)

Description

❓ Questions and Help

Hi,

I'm using the API command to load my finetuned mBART model, below is the command:

en2ar = TransformerModel.from_pretrained(
  '/home/ubuntu/mayub/models',
  checkpoint_file='checkpoint_best_en-ar_ft.pt',
  data_name_or_path='/home/ubuntu/mayub/models',
  bpe='sentencepiece',
  sentencepiece_vocab='sentencepiece.bpe.model'
)

Looks like this loads the model first on cpu and then I do en2ar.cuda() to replicate again on GPU device.

image

Is there a way I can avoid the loading on both devices and just directly load on GPU. Reason is I have 5 other models which I need to load and save on RAM.

Appreciate any help. Thanks !

Guide contributeur

Direction de recherche
Examinez la méthode TransformerModel.from pretrained et explorez les options pour charger directement le modèle sur le GPU, comme l'utilisation du paramètre `device map` ou l'appel à `.to('cuda')` pendant le chargement.
Stack technique
python
Domaine
backend
Type d'issue
Recherche
DifficultéDifficulté estimée pour un nouveau contributeur, de 1 pour un très petit changement à 5 pour un travail expert.
2
Temps estiméFourchette de temps approximative pour investiguer, implémenter, tester et préparer une pull request.
Moins d'une heure
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
PythonPyTorch
Accessibilité débutantScore de 1 à 100 estimant l'accessibilité de cette issue pour un premier contributeur.
70