From 762657afd24673323dd78ef7fc842260406b9793 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 15 May 2026 03:51:55 -0700 Subject: [PATCH] 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. --- studio/backend/core/training/worker.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 6b3b3b6609..4434436ca3 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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,