studio/mlx: lower per-element grad clip default from 5.0 to 1.0 (#5440)

Studio's MLX training worker explicitly pinned ``max_grad_value=5.0``
into the ``MLXTrainingConfig`` so it would override the zoo default
regardless. The 5.0 threshold was effectively no protection -- per-
element transformer gradients in steady state are 1e-3..1e-1, so
|g_i| > 5 basically never fires even on spike batches, mixed-precision
overflow, or RL gradient bursts.

Switch to 1.0:
  - matches the universal LLM clip_grad_norm=1.0 baseline (HF Trainer
    / TRL / PEFT / AutoTrain) while staying on MLX's fast per-element
    ``tree_map(mx.clip)`` path (no global reduction)
  - actually catches outliers without distorting Adam's normalised
    updates (typical post-warmup |g_i| << 1.0)
  - lines up with the new MLXTrainingConfig default in
    unslothai/unsloth-zoo so Studio doesn't silently disagree with
    what zoo ships

No UI change; the TODO to expose grad clipping in Studio settings
remains. Existing trained runs are unaffected: only newly-spawned
training workers pick up the tighter clip.
This commit is contained in:
Daniel Han 2026-05-15 03:51:55 -07:00 committed by GitHub
commit 762657afd2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -773,9 +773,11 @@ def _run_mlx_training(event_queue, stop_queue, config):
else:
eval_steps_val = int(eval_steps_val)
# MLX: value-clip grads to [-5, 5]; norm clipping disabled for compile-friendliness.
# MLX: per-element clip to [-1, 1]; norm clip disabled (it needs a
# global reduction that breaks MLX's eager pipeline). 1.0 (not 5.0):
# |g_i| > 5 rarely fires, so the historical 5.0 was effectively no-op.
max_grad_norm = 0.0
max_grad_value = 5.0 # TODO: expose MLX grad-clip in Studio UI for power users
max_grad_value = 1.0 # TODO: expose MLX grad-clip in Studio UI for power users
trainer = MLXTrainer(
model = model,