mlx-lm `mask_prompt` : seul le dernier message est appris

Avec mask_prompt: true, l’entraînement LoRA de mlx-lm sur une conversation multi-tours n’apprend que le dernier message de chaque exemple : une conversation de dix tours passée telle quelle n’enseigne qu’une seule réponse, la dixième. Pour l’assistant du Journal des Nuisances, j’ai donc fait de chaque réponse retenue un exemple à part, arrêté à cette réponse : 314 exemples au lieu de 112 conversations. Contexte : mlx-lm 0.32.0, base mlx-community/gemma-4-e4b-it-4bit, sur un Mac mini M2 Pro 16 Go.

Le contexte : affiner gemma-4-E4B pour l’assistant

L’assistant du Journal des Nuisances, un produit que j’ai co-fondé, tourne sur pc-llm, un PC Snapdragon, en CPU. La base retenue pour le remplacer est gemma-4-E4B, choisie après des mesures llama.cpp sur le CPU et le GPU Adreno. Elle n’a pas encore le comportement voulu : il faut l’affiner.

L’affinage LoRA (Low-Rank Adaptation) consiste à geler les poids du modèle et à n’entraîner que de petites matrices ajoutées à certaines couches. MLX, le framework d’Apple, sait le faire sur le GPU d’un Mac à puce Apple, avec mlx-lm. Le Mac est prêté par un associé et piloté à distance (mac-train).

Le jeu d’entraînement vient de conversations réelles de recette, rejouées tour par tour contre le modèle actuel ; une réponse n’entre que si elle passe les contrôles et le LLM-juge du banc d’essai. Résultat : 314 tours retenus sur 356 (88 %), dans 112 conversations.

Que fait mask_prompt dans une conversation multi-tours ?

mlx-lm lit des fichiers JSONL au format chat, une ligne par exemple :

{"messages": [{"role": "system", "content": "…"}, {"role": "user", "content": "…"}, {"role": "assistant", "content": "…"}, …]}

Sans masque, la perte porte sur tous les jetons, prompt système compris : le modèle apprendrait à réciter le socle. Avec mask_prompt, mlx-lm ne garde que le dernier message : tout ce qui précède est du contexte. Sur une conversation multi-tours, les réponses intermédiaires de l’assistant font partie de ce contexte. Elles ne sont jamais apprises.

Avec mask_prompt, un exemple mlx-lm enseigne une seule réponse : la dernière. Une conversation de dix tours doit devenir dix exemples.

Un exemple par réponse, arrêté à sa réponse

preparer.py découpe chaque conversation à chaque message de l’assistant :

def exemples(conversation: dict) -> list[dict]:
    messages = conversation["messages"]
    return [
        {"messages": messages[: i + 1]}
        for i, m in enumerate(messages)
        if m["role"] == "assistant"
    ]

Chaque réponse retenue est ainsi apprise une fois, avec la conversation qui la précède. Le coût se voit dans les chiffres : 314 exemples de 2 800 à 5 200 jetons, pour 26 000 jetons de réponses à apprendre. Le socle, c’est-à-dire le prompt système de l’assistant (~2 700 jetons à lui seul), est relu dans chaque exemple.

Le même script écrit un sous-dossier essai/ avec les huit plus longs exemples d’entraînement et les deux plus longs de validation : si la mémoire tient sur eux, elle tiendra sur tout le jeu.

La validation par conversation entière, tirée par CRC32

Une conversation sur dix part entière en validation. Si on répartissait les exemples un par un, deux exemples d’une même conversation, qui partagent leur début, pourraient tomber de part et d’autre : la validation mesurerait du déjà-vu.

def en_validation(identifiant: str) -> bool:
    # Stable d'une machine et d'un passage à l'autre, contrairement à hash().
    return zlib.crc32(identifiant.encode()) % 10 == 0

Pourquoi pas hash() ? Par connaissance générale de Python : depuis Python 3.3, le hash() d’une chaîne est randomisé à chaque lancement de l’interpréteur, sauf si la variable PYTHONHASHSEED est fixée. La répartition changerait donc entre mon PC et le Mac, et d’un lancement à l’autre. CRC32 dépend des seuls octets de l’identifiant. Bilan : 32 réponses en validation, 282 en entraînement.

Les réglages LoRA, dimensionnés sur le jeu

entrainement.yaml, lu par mlx_lm.lora -c :

⚠️ max_seq_length doit couvrir le plus long exemple : un exemple plus long serait tronqué, et c’est la fin, donc la réponse, qui disparaîtrait.

C’est un premier réglage. L’entraînement complet n’a pas encore de résultat : la courbe de perte et le passage au banc d’essai sont à venir.

Lessons learned

• mask_prompt de mlx-lm n’apprend que le dernier message de l’exemple. Les réponses intermédiaires d’une conversation multi-tours ne sont que du contexte.
• Une conversation multi-tours se découpe en un exemple par réponse. Chaque exemple s’arrête à la réponse qu’il enseigne, au prix d’un préfixe relu à chaque fois.
• La validation se tire par conversation entière. Des exemples issus d’une même conversation partagent leur début et fausseraient la mesure.
• zlib.crc32 remplace hash() pour une répartition reproductible. Le hash() d’une chaîne Python change d’un lancement à l’autre.
• max_seq_length se règle sur le plus long exemple réel. 5 632 pour 5 193 jetons, sinon la réponse est tronquée.