Pourquoi mlx-lm lora manque-t-il de mémoire sur un Mac de 16 Go ?

La perte par défaut de mlx-lm calcule les logits de tout le vocabulaire à chaque position de l’exemple, puis masque le prompt. Avec Gemma 4 (262 144 mots) et des exemples de ~5 000 jetons dont ~80 de réponse, ces logits pèsent plusieurs Go pour une poignée de positions utiles. Les calculer sur la seule réponse donne la même perte pour une fraction de la mémoire.

Comment augmenter la mémoire GPU disponible pour MLX sur macOS ?

macOS ne laisse au GPU qu’environ deux tiers de la mémoire unifiée, soit ~10,7 Go sur un Mac de 16 Go. La limite se relève avec sudo sysctl iogpu.wired_limit_mb=13312 (13 Go), réglage perdu au redémarrage : un LaunchDaemon qui lance cette commande avec RunAtLoad le réapplique à chaque démarrage. Cela ne suffit pas si la perte elle-même déborde.

Comment remplacer la fonction de perte de mlx_lm.lora sans modifier mlx-lm ?

mlx_lm.lora appelle train() et evaluate() du module mlx_lm.tuner.trainer avec la perte par défaut. Un script peut remplacer lora.train et lora.evaluate par des functools.partial qui fixent loss= (et au besoin iterate_batches=), puis appeler lora.main() : le fichier YAML et les options de ligne de commande restent les mêmes.

La perte calculée sur la réponse seule change-t-elle l’entraînement ?

Non, si elle garde exactement les positions que le masque de mlx-lm conserve. Le modèle lit toujours toute la conversation ; seuls les logits des positions masquées, qui ne comptaient pas dans la perte, ne sont plus calculés. Sur un petit Gemma 4 aux poids aléatoires, la perte et le gradient sont identiques à ceux de la perte par défaut.

LoRA MLX sur Mac 16 Go : la perte qui sature la mémoire

Un entraînement mlx-lm LoRA en manque de mémoire sur un Mac 16 Go ne venait pas du modèle, mais de la perte : la perte par défaut de mlx-lm calcule les logits des 262 144 mots du vocabulaire de Gemma 4 à chaque position de l’exemple (~5 000 jetons), puis en masque tout sauf la réponse (~80 jetons). Relever la limite mémoire du GPU à 13 Go n’a pas suffi. Le correctif : un script qui laisse le modèle lire toute la conversation, mais ne calcule les logits que sur les positions de la réponse. Même perte, même gradient.

Le contexte : gemma-4-E4B sur un Mac mini M2 Pro

J’affine gemma-4-E4B pour l’assistant du Journal des Nuisances, un produit que j’ai co-fondé, avec mlx-lm 0.32.0 (Python 3.12 installé par uv, macOS ne fournissant que le 3.9). La base est mlx-community/gemma-4-e4b-it-4bit, 5,1 Go. La machine : un Mac mini M2 Pro 16 Go sous macOS 26.3, prêté par un associé et joint par le mesh WireGuard dès le démarrage.

Le jeu compte 314 exemples de 2 800 à 5 200 jetons, un par réponse, pour la raison expliquée dans l’article sur mask_prompt et le dernier message. Les réglages sont déjà prudents : un exemple par lot, accumulation sur 4, activations recalculées.

L’essai du 12 octobre, lancé sur les plus longs exemples, a pourtant manqué de mémoire GPU au premier pas.

Fausse piste : la limite mémoire GPU de macOS

Sur un Mac à puce Apple, la mémoire est unifiée, mais macOS ne laisse au GPU qu’environ deux tiers : ~10,7 Go sur 16. La limite se relève avec sysctl iogpu.wired_limit_mb, qui ne survit pas au redémarrage. Je l’ai donc posée par un LaunchDaemon, écrit avec PlistBuddy (collé dans zsh, le EOF d’un heredoc n’avait pas été reconnu) :

sudo /usr/libexec/PlistBuddy -c "Add :Label string fr.l4w.gpu-memoire" \
  -c "Add :ProgramArguments array" \
  -c "Add :ProgramArguments:0 string /usr/sbin/sysctl" \
  -c "Add :ProgramArguments:1 string iogpu.wired_limit_mb=13312" \
  -c "Add :RunAtLoad bool true" /Library/LaunchDaemons/fr.l4w.gpu-memoire.plist
sudo launchctl bootstrap system /Library/LaunchDaemons/fr.l4w.gpu-memoire.plist
sysctl iogpu.wired_limit_mb     # 13312

Nouvel essai : toujours en manque de mémoire, même à 13 Go. Le réglage reste en place, mais le problème était ailleurs.

Pourquoi mlx-lm lora manque de mémoire sur un Mac 16 Go ?

Les logits sont les scores bruts que le modèle donne à chaque mot du vocabulaire pour prédire le jeton suivant. La perte par défaut de mlx-lm les calcule pour toutes les positions de l’exemple, puis applique le masque de mask_prompt pour ne garder que la réponse.

Sur mes exemples, c’est un rapport défavorable : ~5 000 positions × 262 144 mots, pour ~80 positions qui comptent. Ces scores pèsent plusieurs Go, et il faut aussi les garder pour le calcul du gradient.

Avec un grand vocabulaire et de longs prompts masqués, la mémoire d’un affinage ne part pas dans le modèle mais dans les logits de positions qui ne comptent pas.

La perte calculée sur la réponse seule

lancer_lora.py fait lire toute la conversation au modèle, mais ne projette sur le vocabulaire que les états des positions de la réponse :

def perte_reponse(model, lot, positions):
    texte = model.language_model
    etats = mx.take(texte.model(lot[:, :-1]), positions, axis=1)
    cibles = mx.take(lot[:, 1:], positions, axis=1)
    if texte.tie_word_embeddings:
        logits = texte.model.embed_tokens.as_linear(etats)
    else:
        logits = texte.lm_head(etats)
    if texte.final_logit_softcapping is not None:
        plafond = texte.final_logit_softcapping
        logits = mx.tanh(logits / plafond) * plafond

    ce = nn.losses.cross_entropy(logits, cibles).astype(mx.float32)
    return ce.sum() / ce.size, mx.array(ce.size)

Deux détails de Gemma 4 sont à reproduire à la main, puisqu’on court-circuite la tête du modèle. Avec tie_word_embeddings, la matrice de sortie est celle des embeddings, d’où embed_tokens.as_linear au lieu de lm_head. Et final_logit_softcapping plafonne les logits par une tangente hyperbolique.

Les positions sont exactement celles que le masque de mlx-lm garde ; elles sont calculées par un générateur de lots, lots_reponse, qui suppose batch_size: 1. Le branchement ne touche pas à mlx-lm :

lora.train = partial(trainer.train, loss=perte_reponse, iterate_batches=lots_reponse)
lora.evaluate = partial(trainer.evaluate, loss=perte_reponse, iterate_batches=lots_reponse)
sys.argv[0] = "mlx_lm.lora"
lora.main()

entrainer.sh lance désormais python lancer_lora.py -c entrainement.yaml au lieu de mlx_lm.lora, avec le même YAML.

Même perte, une fraction de la mémoire

Vérification sur un petit Gemma 4 aux poids aléatoires : perte et gradient identiques à ceux de la perte de mlx-lm. Côté mémoire, en plus du modèle, mesurée sur le processeur :

Perte	Longueur d’exemple	Mémoire en plus du modèle
Par défaut (mlx-lm)	1 000 jetons	~6 Go
Sur la réponse seule	5 000 jetons	~1 Go

Cinq fois plus de jetons pour six fois moins de mémoire. L’entraînement complet sur le Mac n’a pas encore de résultat : courbe de perte et passage au banc d’essai sont à venir.

Lessons learned

La perte par défaut de mlx-lm calcule les logits à toutes les positions avant de masquer. Avec 262 144 mots de vocabulaire et 5 000 jetons de prompt, c’est elle qui sature la mémoire.
Relever iogpu.wired_limit_mb ne corrige pas une perte qui déborde. Passer de ~10,7 à 13 Go pour le GPU n’a rien changé au premier pas.
Ne projeter sur le vocabulaire que les positions de la réponse donne la même perte. Le modèle lit toujours tout le contexte ; seuls les scores inutiles disparaissent.
Court-circuiter la tête d’un modèle oblige à reproduire ses détails. Pour Gemma 4 : embeddings liés et final_logit_softcapping.
Une perte réécrite se vérifie contre l’originale sur un petit modèle aléatoire. Perte et gradient identiques avant de lancer des heures d’entraînement.