Optimisation de la mémoire pour l’inférence et la fine-tuning de grands modèles de langage

Les grands modèles de langage (LLM) comme GPT-4, Bloom et LLaMA ont atteint des capacités remarquables en augmentant leur taille à plusieurs milliards de paramètres. Cependant, déployer ces modèles massifs pour l’inférence ou la fine-tuning est difficile en raison de leurs énormes besoins en mémoire. Dans ce blog technique, nous allons explorer les techniques pour estimer et optimiser la consommation de mémoire pendant l’inférence et la fine-tuning des LLM sur différents matériel.

Comprendre les besoins en mémoire

La mémoire requise pour charger un LLM est principalement déterminée par le nombre de paramètres et la précision numérique utilisée pour stocker les paramètres. Une règle simple est :

• Charger un modèle avec X milliards de paramètres nécessite environ 4X GB de VRAM en précision 32 bits
• Charger un modèle avec X milliards de paramètres nécessite environ 2X GB de VRAM en précision 16 bits bfloat16/float16

Par exemple, charger le modèle GPT-3 de 175 milliards de paramètres nécessiterait environ 350 GB de VRAM en précision bfloat16. Actuellement, les plus grandes cartes graphiques disponibles sur le marché, comme les NVIDIA A100 et H100, n’offrent que 80 GB de VRAM, ce qui nécessite des techniques de parallélisme de tenseur et de modèle.

Pendant l’inférence, l’empreinte mémoire est dominée par les paramètres du modèle et les tenseurs d’activation temporaires produits. Une estimation de haut niveau de l’utilisation de mémoire maximale pendant l’inférence est la somme de la mémoire requise pour charger les paramètres du modèle et la mémoire pour les activations.

Quantifier la mémoire d’inférence

Essayons de quantifier les besoins en mémoire pour l’inférence en utilisant le modèle OctoCode, qui a environ 15 milliards de paramètres au format bfloat16 (~ 31 GB). Nous allons utiliser la bibliothèque Transformers pour charger le modèle et générer du texte :

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder",
torch_dtype=torch.bfloat16,
device_map="auto",
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a Python function to convert bytes to gigabytes.\n\nAnswer:"
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]

def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024

bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Sortie :

L’utilisation maximale de mémoire GPU est d’environ 29 GB, ce qui correspond à notre estimation de 31 GB pour charger les paramètres du modèle au format bfloat16.

Optimiser la mémoire d’inférence avec la quantification

Alors que bfloat16 est la précision couramment utilisée pour l’entraînement des LLM, les chercheurs ont constaté que la quantification des poids du modèle à des types de données de précision inférieure comme des entiers 8 bits (int8) ou des entiers 4 bits peut réduire considérablement l’utilisation de la mémoire avec une perte d’exactitude minimale pour les tâches d’inférence comme la génération de texte.

Voyons les économies de mémoire provenant de la quantification 8 bits et 4 bits du modèle OctoCode :

</div>
# Quantification 8 bits
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True,
pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

Sortie :

# Quantification 4 bits
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Sortie :

Avec la quantification 8 bits, les besoins en mémoire passent de 31 GB à 15 GB, tandis que la quantification 4 bits les réduit encore à 9,5 GB ! Cela permet d’exécuter le modèle OctoCode de 15 milliards de paramètres sur des cartes graphiques grand public comme la RTX 3090 (24 GB de VRAM).

Cependant, notez que la quantification plus agressive comme 4 bits peut parfois entraîner une dégradation de l’exactitude par rapport à la précision 8 bits ou bfloat16. Il existe un compromis entre les économies de mémoire et l’exactitude que les utilisateurs doivent évaluer pour leur cas d’utilisation.

La quantification est une technique puissante qui peut permettre le déploiement de LLM sur des environnements à ressources limitées comme les instances cloud, les appareils périphériques ou même les téléphones mobiles en réduisant considérablement l’empreinte mémoire.

Estimer la mémoire pour la fine-tuning

Alors que la quantification est principalement utilisée pour l’inférence efficace, des techniques comme le parallélisme de tenseur et le parallélisme de modèle sont cruciales pour gérer les besoins en mémoire pendant l’entraînement ou la fine-tuning des grands modèles de langage.

La consommation de mémoire maximale pendant la fine-tuning est généralement 3 à 4 fois supérieure à l’inférence en raison de besoins en mémoire supplémentaires pour :

• Les gradients
• Les états de l’optimiseur
• Les activations de la passe avant stockées pour la rétropropagation

Une estimation conservatrice est que la fine-tuning d’un LLM avec X milliards de paramètres nécessite environ 4 * (2X) = 8X GB de VRAM en précision bfloat16.

Par exemple, la fine-tuning du modèle LLaMA de 7 milliards de paramètres nécessiterait environ 7 * 8 = 56 GB de VRAM par carte graphique en précision bfloat16. Cela dépasse la capacité de mémoire des cartes graphiques actuelles, ce qui nécessite des techniques de fine-tuning distribuées.

Techniques de fine-tuning distribuées

Plusieurs méthodes de fine-tuning distribuées ont été proposées pour surmonter les contraintes de mémoire GPU pour les grands modèles :

1. Parallélisme de données : L’approche classique de parallélisme de données réplique le modèle entier sur plusieurs cartes graphiques tout en divisant et en distribuant les lots de données d’entraînement. Cela réduit le temps d’entraînement de manière linéaire avec le nombre de cartes graphiques, mais ne réduit pas les besoins en mémoire maximaux sur chaque carte graphique.
2. ZeRO Stage 3 : Une forme avancée de parallélisme de données qui partitionne les paramètres du modèle, les gradients et les états de l’optimiseur sur les cartes graphiques. Cela réduit la mémoire par rapport au parallélisme de données classique en ne gardant que les données partitionnées requises sur chaque carte graphique pendant les différentes phases de l’entraînement.
3. Parallélisme de tenseur : Au lieu de répliquer le modèle, le parallélisme de tenseur divise les paramètres du modèle en lignes ou en colonnes et les distribue sur les cartes graphiques. Chaque carte graphique opère sur un ensemble partitionné de paramètres, de gradients et d’états de l’optimiseur, ce qui entraîne des économies de mémoire considérables.
4. Parallélisme de pipeline : Cette technique partitionne les couches du modèle sur différents processeurs, chaque processeur exécutant un sous-ensemble des couches. Les activations sont transmises entre les processeurs, ce qui réduit la mémoire maximale mais augmente la charge de communication.

Estimer les besoins en mémoire pour ces méthodes distribuées est non trivial, car la distribution des paramètres, des gradients, des activations et des états de l’optimiseur varie selon les techniques. De plus, les différents composants comme le corps du transformateur et la tête de modélisation de langage peuvent présenter des comportements d’allocation de mémoire différents.

La solution LLMem

Les chercheurs ont récemment proposé LLMem, une solution qui estime avec précision la consommation de mémoire GPU lors de l’application de méthodes de fine-tuning distribuées à des LLM sur plusieurs cartes graphiques.

Estimating GPU Memory Usage for Fine-Tuning Pre-Trained LLM

LLMem prend en compte des facteurs tels que la recombinaison des paramètres avant le calcul (ZeRO Stage 3), la collecte de sortie dans la passe arrière (parallélisme de tenseur) et les différentes stratégies d’allocation de mémoire pour le corps du transformateur et la tête de modélisation de langage.

Les résultats expérimentaux montrent que LLMem peut estimer l’utilisation maximale de mémoire GPU pour la fine-tuning de LLM sur une seule carte graphique avec des taux d’erreur allant jusqu’à 1,6 %, surpassant le taux d’erreur moyen de 42,6 % de DNNMem. Lors de l’application de méthodes de fine-tuning distribuées à des LLM avec plus d’un milliard de paramètres sur plusieurs cartes graphiques, LLMem atteint un taux d’erreur moyen de 3,0 %.