MLX `Attempting to eval an array` : `.item()` en compilé

Attempting to eval an array during function transformations : MLX refuse ce .item() parce que le pas d’entraînement de mlx-lm est compilé par mx.compile, et qu’une fonction compilée ne peut pas lire la valeur d’un tableau pendant qu’on la transforme. Ma perte personnalisée lisait deux longueurs par .item() pour découper la réponse. Le correctif : calculer ces positions avant la compilation, dans le générateur de lots, et ne plus les lire dans la perte, seulement les utiliser par mx.take. Contexte : mlx-lm 0.32.0 sur un Mac mini M2 Pro 16 Go, le 12 octobre.

Le contexte : une perte calculée sur la réponse seule

Pour tenir dans 16 Go, j’avais remplacé la perte de mlx-lm par une perte qui ne calcule les logits que sur les positions de la réponse, comme expliqué dans l’article sur la perte qui sature la mémoire d’un Mac 16 Go. La perte reçoit, comme celle de mlx-lm, un tableau lengths qui contient pour chaque exemple le début de la réponse (offset) et la longueur totale.

La première version découpait avec des entiers Python, tirés de ce tableau.

Pourquoi .item() provoque « Attempting to eval an array » ?

mx.compile trace la fonction : MLX l’exécute une fois avec des tableaux symboliques, enregistre le graphe de calcul, puis rejoue ce graphe sans repasser par Python. Pendant ce traçage, un tableau n’a pas encore de valeur. Or .item() réclame une valeur concrète, donc une évaluation : MLX la refuse avec le message cité plus haut.

Le piège vient de ce qu’on ne voit pas la compilation : elle est dans trainer.train de mlx-lm, qui compile tout le pas (perte, gradient, mise à jour). Une perte testée seule, hors de trainer.train, passe sans erreur.

Dans une fonction compilée par MLX, on peut indexer avec un tableau, pas avec une valeur lue par .item() : tout ce qui dépend des données se calcule avant.

Le correctif : les positions viennent du générateur de lots

Les lots de mlx-lm sortent de trainer.iterate_batches, qui tourne en Python ordinaire, hors compilation. lots_reponse l’enveloppe et remplace les longueurs par la liste des positions de la réponse. La perte n’a plus qu’à les prendre avec mx.take, une opération sur tableaux que la compilation sait tracer. Le diff, réduit à l’essentiel :

-def perte_reponse(model, batch, lengths):
-    if batch.shape[0] != 1:
-        raise ValueError("perte_reponse suppose batch_size: 1")
-    debut = max(int(lengths[0, 0].item()) - 1, 0)
-    fin = int(lengths[0, 1].item())
-
-    texte = model.language_model
-    etats = texte.model(batch[:, :-1])[:, debut:fin]
-    cibles = batch[:, 1:][:, debut:fin]
+def lots_reponse(*args, **kwargs):
+    for lot, longueurs in trainer.iterate_batches(*args, **kwargs):
+        if lot.shape[0] != 1:
+            raise ValueError("lots_reponse suppose batch_size: 1")
+        offset, longueur = longueurs[0].tolist()
+        # Les mêmes positions que le masque de default_loss : la cible i
+        # compte si offset <= i + 1 <= longueur.
+        yield lot, mx.arange(max(offset - 1, 0), longueur)
+
+
+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)

Le .tolist() de lots_reponse lit bien une valeur, mais dans le générateur, avant que le lot n’entre dans le pas compilé : là, c’est permis. Le reste de la perte (projection sur le vocabulaire, cross-entropie) ne change pas.

Le branchement passe le générateur à mlx-lm en même temps que la perte, pour l’entraînement comme pour l’évaluation :

lora.train = partial(trainer.train, loss=perte_reponse, iterate_batches=lots_reponse)
lora.evaluate = partial(trainer.evaluate, loss=perte_reponse, iterate_batches=lots_reponse)

La vérification : trainer.train complet

Une perte vérifiée seule ne prouve rien ici, puisque l’erreur n’apparaît qu’une fois compilée. J’ai donc lancé trainer.train complet, validation, accumulation de gradient et activations recalculées comprises, sur un petit Gemma 4 : courbes de perte identiques à celles de la perte par défaut de mlx-lm, à ordre des lots égal. L’entraînement complet sur le Mac reste à faire : pas encore de résultat à montrer.

Lessons learned

• Le pas d’entraînement de mlx-lm est compilé par mx.compile. Une perte personnalisée y est tracée, et ne peut lire aucune valeur de tableau.
• .item() dans une fonction compilée déclenche « Attempting to eval an array during function transformations ». Il réclame une évaluation pendant le traçage.
• Ce qui dépend des données se calcule dans le générateur de lots. iterate_batches tourne hors compilation et peut fournir des indices tout prêts.
• Une perte personnalisée se teste dans trainer.train, pas seule. C’est là que la compilation s’applique.