Llama-Factory peut-il être utilisé pour la distillation des connaissances ?

La course aux performances des grands modèles de langage (LLM) a engendré des monstres de plusieurs centaines de milliards de paramètres. Leur puissance est incontestable, mais leur déploiement se heurte à un mur de coûts et de contraintes pratiques, les rendant inaccessibles pour de nombreux scénarios en périphérie ou sur mobiles.

Un besoin pressant apparaît alors : **comment miniaturiser l'intelligence d'un modèle géant ?** La distillation des connaissances (Knowledge Distillation) est une solution clé. Elle consiste à entraîner un modèle léger (l'étudiant) à imiter le comportement d'un modèle massif (l'enseignant), non seulement sur les réponses correctes, mais aussi sur sa manière de raisonner. L'étudiant gagne ainsi en vitesse d'inférence tout en conservant une large part des capacités de l'enseignant.

La mise en œuvre traditionnelle de cette technique est complexe. Elle requiert la gestion des étiquettes souples (soft labels) de l'enseignant, la conception de fonctions de perte composées et le pilotage d'un entraînement à double modèle. C'est ici que des cadres de travail comme **Llama-Factory** entrent en jeu. Peut-il simplifier ce processus ?

Llama-Factory est conçu principalement pour l'affinage (fine-tuning) unifié des LLM, supportant une multitude d'architectures et fournissant une chaîne complète, du prétraitement des données au déploiement. Sa commodité pour l'affinage via LoRA ou QLoRA est bien connue.

Peut-on étendre cette commodité à la distillation ? La réponse est : **Oui, moyennant des adaptations techniques.** Llama-Factory ne propose pas de bouton pour activer la distillation. Cependant, son moteur d'entraînement est extensible. Il est possible de transformer son pipeline d'affinage supervisé standard en un processus de distillation.

Le cœur de la distillation réside dans la fonction de perte. En plus de l'entropie croisée classique, on introduit une divergence KL pour aligner la distribution de sortie de l'étudiant sur celle de l'enseignant. Voici un exemple de module de perte personnalisé :

# perte_distillation.py
import torch
import torch.nn as nn
import torch.nn.functional as F

class PerteDistillation(nn.Module):
    def __init__(self, coeff_alpha=0.5, coeff_temp=2.0):
        super().__init__()
        self.coeff_alpha = coeff_alpha
        self.coeff_temp = coeff_temp
        self.perte_ce = nn.CrossEntropyLoss()

    def forward(self, logits_etudiant, logits_enseignant, etiquettes_reelles):
        # Perte pour les cibles souples : divergence KL entre distributions adoucies
        perte_souple = F.kl_div(
            F.log_softmax(logits_etudiant / self.coeff_temp, dim=-1),
            F.softmax(logits_enseignant / self.coeff_temp, dim=-1),
            reduction='batchmean'
        ) * (self.coeff_temp ** 2)

        # Perte pour les cibles dures : entropie croisée standard
        perte_dure = self.perte_ce(logits_etudiant, etiquettes_reelles)

        # Perte totale combinée
        perte_totale = (self.coeff_alpha * perte_souple) + ((1 - self.coeff_alpha) * perte_dure)
        return perte_totale

Ce code met en œuvre la stratégie classique combinant cibles souples et dures. Le paramètre coeff_temp lisse les distributions de probabilité, tandis que coeff_alpha équilibre le transfert de connaissances de l'enseignant et la supervision par les étiquettes réelles.

La version actuelle de Llama-Factory ne fournit pas d'interface native pour injecter une perte personnalisée. Intégrer cette PerteDistillation nécessite donc des modifications du code source, par exemple en étendant la classe Trainer par défaut ou en modifiant les scripts d'entraînement pour qu'ils acceptent et utilisent les logits_enseignant depuis le jeu de données.

Pourquoi malgré tout considérer Llama-Factory comme une base solide ?

  1. Support natif du QLoRA : C'est crucial, car le modèle étudiant est souvent contraint en ressources. Le QLoRA permet d'effectuer un affinage efficace sur un seul GPU grand public avec 16 Go de mémoire, réduisant considérablement les coûts d'expérimentation.
  2. Conception modulaire : Le cadre gère les tâches répétitives (chargement du modèle, tokenisation, construction des batches, planification du taux d'apprentissage), permettant au développeur de se concentrer sur la logique spécifique de la distillation.
  3. Interface WebUI pour le suivi : La surveillance en temps réel des courbes de perte, du taux d'apprentissage et de l'utilisation du GPU est essentielle pour déboguer le processus délicat de l'équilibre entre perte souple et perte dure.

Un flux de travail typique de distillation pourrait se décomposer ainsi :

1. Génération hors-ligne des soft labels : Utiliser le modèle enseignant (ex: Qwen-7B) pour effectuer une inférence sur le jeu d'entraînement et sauvegarder les logits pour chaque échantillon. Une température élevée (T=5) est recommandée pour extraire une information sémantique plus riche. Ces logits sont ensuite fusionnés avec les données originales pour créer un nouveau jeu de données contenant les champs texte_entree, etiquette et logits_enseignant.

2. Intégration et lancement de l'entraînement : Après avoir adapté le code de Llama-Factory pour supporter la perte personnalisée et le chargement des logits_enseignant, on peut lancer l'entraînement. Un choix pertinent pour le modèle étudiant est TinyLlama-1.1B, avec environ un sixième des paramètres de l'enseignant. L'utilisation de LoRA pour l'affinage réduit encore la charge.

CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
    --model_name_or_path TinyLlama/TinyLlama-1.1B \
    --do_train \
    --dataset dataset_distille \
    --finetuning_type lora \
    --loss_function distillation \
    --temp_distillation 5.0 \
    --alpha_distillation 0.7 \
    --output_dir ./sortie-distillation \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 4 \
    --learning_rate 3e-4 \
    --num_train_epochs 5

3. Évaluation et déploiement : On évalue le modèle étudiant sur un jeu de validation en comparant ses métriques (BLEU, ROUGE, Accuracy) à celles de l'enseignant. L'objectif est d'atteindre 90-95% des performances de l'enseignant avec un gain de vitesse d'inférence significatif. Les poids LoRA fusionnés peuvent ensuite être exportés et, si nécessaire, quantifiés (GGUF) pour un déploiement sur du matériel très limité.

Parmi les considérations clés du processus : la compatibilité d'architecture entre l'enseignant et l'étudiant (privilégier la même famille de modèles), la qualité des données qui conditionne le plafond de performance, et la nécessité de stocker les logits_enseignant pour accélérer les itérations d'entraînement.

En somme, bien que Llama-Factory ne soit pas un outil de distillation dédié, il offre une plateforme extrêmement adaptée pour sa mise en œuvre. Sa compatibilité multi-modèles, son support du QLoRA et sa modularité en font un point de départ pragmatique et efficace pour les développeurs souhaitant explorer la compression des LLM. C'est une voie viable pour rendre opérationnels les modèles de pointe là où les ressources sont limitées.

Étiquettes: Llama-Factory Knowledge Distillation QLoRA LLM Compression Model Distillation

Publié le 16 juillet à 10h30