MLX Training updates

studio/backend/core/training/training.py

@@ -218,6 +218,7 @@ def start_training(self, job_id: str, **kwargs) -> bool:
             "save_steps": kwargs.get("save_steps", 0),
             "weight_decay": kwargs.get("weight_decay", 0.001),
             "max_grad_norm": kwargs.get("max_grad_norm", 0.0),
+            "max_grad_value": kwargs.get("max_grad_value"),
             "random_seed": kwargs.get("random_seed", 3407),
             "packing": kwargs.get("packing", False),
             "optim": kwargs.get("optim", "adamw_8bit"),

studio/backend/core/training/worker.py

@@ -1156,14 +1156,16 @@ def _send(event_type, **kwargs):
     is_dataset_image = bool(config.get("is_dataset_image", False))
     training_type = config.get("training_type", "LoRA/QLoRA")
     use_lora = training_type == "LoRA/QLoRA"
+    model_random_state = config.get("model_random_state", 3407)
+    lora_random_state = config.get("lora_random_state", 3407)
     model, tokenizer = FastMLXModel.from_pretrained(
         model_name,
         load_in_4bit = config.get("load_in_4bit", True),
         full_finetuning = not use_lora,
         text_only = None if is_dataset_image else True,
         token = hf_token,
         trust_remote_code = bool(config.get("trust_remote_code", False)),
-        random_state = config.get("random_seed", 3407),
+        random_state = model_random_state,
     )
 
     is_vlm = bool(is_dataset_image and getattr(model, "_is_vlm_model", False))
@@ -1188,7 +1190,7 @@ def _send(event_type, **kwargs):
             lora_dropout = config.get("lora_dropout", 0.0),
             use_rslora = config.get("use_rslora", False),
             init_lora_weights = config.get("init_lora_weights", True),
-            random_state = config.get("random_seed", 3407),
+            random_state = lora_random_state,
             target_modules = config.get("target_modules")
             or [
                 "q_proj",
@@ -1390,7 +1392,8 @@ def _fmt_progress(status_message = "", **_kw):
     # 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 = 1.0  # TODO: expose MLX grad-clip in Studio UI for power users
+    max_grad_value = config.get("max_grad_value")
+    max_grad_value = 1.0 if max_grad_value is None else float(max_grad_value)
 
     trainer = MLXTrainer(
         model = model,

studio/backend/models/training.py

@@ -267,6 +267,14 @@ def _check_lora_dropout(cls, v: float) -> float:
         ge = 0,
         description = "Global gradient norm clipping threshold. Set 0 to disable.",
     )
+    max_grad_value: Optional[float] = Field(
+        None,
+        ge = 0,
+        description = (
+            "MLX-only elementwise gradient value clipping threshold. "
+            "If unset, MLX uses its runtime default."
+        ),
+    )
     random_seed: int = Field(42, description = "Random seed")
     packing: bool = Field(False, description = "Enable sequence packing")
     optim: str = Field("adamw_8bit", description = "Optimizer")

studio/backend/routes/training.py

@@ -216,6 +216,7 @@ async def start_training(
             "save_steps": request.save_steps,
             "weight_decay": request.weight_decay,
             "max_grad_norm": request.max_grad_norm,
+            "max_grad_value": request.max_grad_value,
             "random_seed": request.random_seed,
             "packing": request.packing,
             "optim": request.optim,

studio/backend/tests/test_mlx_training_worker_config.py

@@ -82,3 +82,12 @@ def test_mlx_studio_rejects_unknown_optimizer():
 def test_mlx_studio_rejects_unknown_scheduler():
     with pytest.raises(ValueError, match = "Unsupported LR scheduler for MLX training"):
         _normalize_mlx_studio_scheduler("linear_typo")
+
+
+def test_mlx_studio_keeps_hf_style_tokenizer_dual_purpose():
+    source = (
+        Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py"
+    ).read_text()
+
+    assert "tokenizer = tokenizer" in source
+    assert "processor = tokenizer if is_vlm else None" not in source

studio/backend/tests/test_training_raw_support.py

@@ -107,10 +107,25 @@ def start(self):
                 model_name = "unsloth/test",
                 training_type = "LoRA/QLoRA",
                 max_grad_norm = 0.7,
+                max_grad_value = 3.0,
             )
 
         config = mock_process.call_args.kwargs["kwargs"]["config"]
         self.assertEqual(config["max_grad_norm"], 0.7)
+        self.assertEqual(config["max_grad_value"], 3.0)
+
+    def test_mlx_worker_uses_cuda_style_model_and_lora_init_seed(self):
+        source = (_BACKEND_ROOT / "core" / "training" / "worker.py").read_text()
+
+        self.assertIn(
+            'model_random_state = config.get("model_random_state", 3407)', source
+        )
+        self.assertIn(
+            'lora_random_state = config.get("lora_random_state", 3407)', source
+        )
+        self.assertIn("random_state = model_random_state", source)
+        self.assertIn("random_state = lora_random_state", source)
+        self.assertIn('seed = config.get("random_seed", 3407)', source)
 
     def test_training_route_forwards_embedding_learning_rate(self):
         training_route = _load_route_module(

studio/frontend/src/features/training/api/mappers.ts

@@ -84,6 +84,7 @@ export function buildTrainingStartPayload(
     eval_steps: config.evalSteps,
     weight_decay: config.weightDecay,
     max_grad_norm: 0.0,
+    max_grad_value: null,
     random_seed: config.randomSeed,
     packing: isEmbedding ? false : config.packing,
     optim: config.optimizerType,

studio/frontend/src/features/training/types/api.ts

@@ -32,6 +32,7 @@ export interface TrainingStartRequest {
   eval_steps: number;
   weight_decay: number;
   max_grad_norm: number;
+  max_grad_value?: number | null;
   random_seed: number;
   packing: boolean;
   optim: string;