Technical Tutorials

L'exécution locale de grands modèles de langage se heurte souvent à une limite stricte : les limitations de VRAM. Même avec des GPU discrets puissants, les modèles dépassant votre capacité de mémoire vidéo obligent les développeurs à choisir entre la quantification (perte de précision) ou une inférence CPU lente. Sur Apple Silicon et les plateformes AMD récentes, l'Architecture à Mémoire Unifiée (UMA) offre une troisième voie. En traitant la RAM système et la mémoire GPU comme un seul pool, nous pouvons stratégiquement décharger des composants spécifiques, tels que les têtes d'attention, pour utiliser la bande passante mémoire complète tout en gardant les calculs critiques sur le GPU.

Pourquoi les têtes d'attention sont des candidats de choix

Toutes les parties d'un modèle transformer ne sont pas créées égales. La multiplication Matrice-Matrice dans les réseaux Feed-Forward (FFN) est intensif en calculs et bénéficie le plus de l'accélération GPU. Cependant, les opérations d'Attention (spécifiquement les projections Q, K, V et le softmax subséquent) peuvent être plus limitées par la bande passante mémoire, surtout à de plus longues longueurs de séquence.

Dans une configuration UMA, le "GPU" est essentiellement un co-processeur qui peut accéder à la RAM système à haute vitesse (souvent dépassant 100 Go/s sur les puces M2/M3/M4 Max). Par conséquent, déplacer les matrices de poids des têtes d'attention vers la RAM système n'entraîne pas le même pénalité que ce serait le cas sur un GPU discret connecté via PCIe. La récupération des données est assez rapide pour que le gain de calcul en gardant le modèle résident en mémoire à haute bande passante dépasse souvent la légère latence d'accès depuis la RAM système "lente".

Stratégie d'implémentation en Python

En utilisant des bibliothèques comme PyTorch ou des moteurs d'inférence spécialisés, nous pouvons contrôler manuellement le placement des dispositifs. La clé est d'identifier les indices de couches et de déplacer des tenseurs de paramètres spécifiques vers le dispositif 'cpu' tout en gardant le reste sur le dispositif 'cuda' ou 'mps'.

import torch

def hybrid_load_model(model, offload_ratio=0.5, layer_type='attention'):
    """
    Décharge une fraction des couches d'attention vers le CPU.
    Suppose une structure Transformer standard.
    """
    layers = list(model.named_parameters())
    total_layers = len([l for l in layers if 'self_attn' in l[0]])
    
    # Déterminer quelles couches décharger
    offload_count = int(total_layers * offload_ratio)
    
    # Stratégie : Décharger les couches les plus anciennes (souvent moins critiques pour les logits finaux)
    # ou les alterner. Ici, nous déchargeons les premiers N blocs d'attention.
    target_layers = [l[0] for l in layers if 'self_attn' in l[0]][:offload_count]

    for name, param in model.named_parameters():
        if name in target_layers:
            param.data = param.data.to('cpu')
            param.data = param.data.float() # Garder une précision plus élevée en RAM
        else:
            # S'assurer que les couches critiques restent sur le GPU
            if param.device.type != 'cuda' and param.device.type != 'mps':
                 param.data = param.data.to('cuda' if torch.cuda.is_available() else 'mps')

    model.eval()
    return model

# Exemple d'utilisation
# model = load_llama_model(path="llama-7b-gguf")
# hybrid_model = hybrid_load_model(model, offload_ratio=0.3, layer_type='attention')

Gérer la surcharge du mouvement des données

Bien que l'UMA minimise les coûts de transfert, elle ne les élimine pas. Chaque passe avant nécessite de déplacer les poids d'attention vers le GPU pour le calcul (ou de les calculer sur le CPU si le backend le prend en charge). Pour minimiser cela, envisagez de regrouper les couches déchargées par lots. Au lieu de déplacer les poids couche par couche, regroupez plusieurs têtes d'attention et transférez-les dans une opération par lot unique.

De plus, exploitez torch.compile avec des gardes de forme dynamiques. Le compilateur de PyTorch peut souvent fusionner les opérations de mouvement de données avec les lancements de noyaux, masquant la latence derrière des opérations limitées par le calcul. Sur Apple Silicon, assurez-vous d'utiliser le backend Metal Performance Shaders (MPS), qui prend en charge nativement l'allocation de mémoire unifiée.

Exemple pratique : Exécuter un modèle de 13 milliards de paramètres

Considérez un modèle de 13 milliards de paramètres en virgule flottante 16 bits. Cela nécessite environ 26 Go de mémoire. Un GPU standard de 16 Go ne peut pas l'accommoder. Cependant, un MacBook Pro avec 32 Go de mémoire unifiée peut le faire.

Embeddings de base : Garder sur le GPU (critique pour la vitesse).
Couches FFN : Garder sur le GPU (intensif en calculs).
Têtes d'attention (Couches 1-16) : Décharger vers le CPU/RAM.

En déchargeant 50 % des têtes d'attention, vous libérez environ 6,5 Go de VRAM du GPU. Cela permet au modèle de se charger. Pendant l'inférence, le GPU gère les calculs FFN lourds, tandis que le CPU gère la projection d'attention. Sur une puce M3 Max, cette approche hybride peut atteindre 25 à 35 jetons par seconde, par rapport à moins de 5 jetons par seconde sur une configuration CPU pure.

Conclusion

L'inférence hybride CPU-GPU sur les Architectures à Mémoire Unifiée n'est pas seulement un plan de secours ; c'est une stratégie principale pour le développement local de l'IA. En déchargeant intelligemment les têtes d'attention, vous débloquez la capacité d'exécuter des modèles plus grands sans compromettre la vitesse de génération des jetons. Alors que le matériel continue d'évoluer, la maîtrise de ces stratégies de placement sera essentielle pour repousser les limites des performances des LLM locaux. Commencez par profiler l'utilisation de la mémoire de votre modèle, identifiez les composants limités par la bande passante et expérimentez avec le déchargement partiel pour trouver le point optimal pour votre configuration matérielle spécifique.