Allier vitesse et Python : comment utiliser les graphes CUDA pour exécuter plus rapidement du code Python en deep learning

Ces dernières années, la vitesse des GPU a explosé, et les méthodes d’optimisation des charges de travail de deep learning évoluent elles aussi. PyTorch ajoute également des fonctionnalités d’optimisation comme torch.compile(), mais pour certaines charges de travail, notamment les LLM, les améliorations sont encore en cours.

(En attendant que torch.compile() progresse), j’ai trouvé un article présentant CUDA Graph comme une méthode d’optimisation applicable immédiatement, ainsi qu’un retour d’expérience sur son utilisation, et je l’ai traduit. (⚠️ Attention : la fin de l’article contient un peu de promotion pour Fireworks.ai, une entreprise qui développe et propose une plateforme d’inférence LLM.)

Cet article présente CUDA Graph dans l’ordre suivant :

• présentation de la méthode d’optimisation existante qu’est le chevauchement CPU/GPU (overlap)

• zones où surviennent des surcoûts CPU

• techniques d’optimisation des surcoûts CPU et CUDA Graph

• présentation d’un cas d’application de CUDA Graph au modèle LLaMA2-7B

• présentation des gains de performance obtenus grâce à CUDA Graph

• annexe : problèmes rencontrés lors de l’utilisation de torch.compile() à l’heure actuelle (PyTorch 2.0.1) et solutions