Guard MLX distributed eval against zero-token batches (#888)

* Guard MLX distributed eval against zero-token batches

The baseline loss labels=None fast path divides by raw ntoks for mlx-lm
parity, so a zero-token eval batch (a distributed_pad_mode="empty" padding
row on a rank, lengths [0, 0]) yields loss = 0/0 = NaN. _evaluate_batch_totals
then did all_losses += loss * ntoks = NaN * 0 = NaN, which the distributed
all-sum spreads to every rank even when other ranks have real tokens.

Mask the zero-token contribution at the aggregation site with mx.where so the
loss function stays byte-for-byte equivalent to mlx-lm default_loss (the parity
guard test is unaffected) while distributed eval no longer poisons its all-sum.

* Use scalar 0.0 fallback in the zero-token eval guard

mx.where selects the fallback only when ntoks == 0, so a Python 0.0 is
equivalent to mx.zeros_like(all_losses) and avoids allocating a scalar
array each eval iteration.

Author: danielhanchen

unsloth_zoo/mlx/trainer.py

@@ -1500,7 +1500,10 @@ def _evaluate_batch_totals(self, eval_batches, loss_fn, is_vlm=False):
                     else:
                         batch, lengths, labels = batch_data
                         loss, ntoks = loss_fn(self.model, batch, lengths, labels)
-                    all_losses += loss * ntoks
+                    # Zero-token eval batches (distributed_pad_mode="empty" padding
+                    # rows) make loss NaN; mask them so NaN * 0 does not poison the
+                    # distributed all-sum. mx.where never selects the NaN branch.
+                    all_losses += mx.where(ntoks > 0, loss * ntoks, 0.0)
                     ntokens += ntoks
                     mx.eval(all_losses, ntokens)
                 except Exception as exc: