From 0104c31dd266c79dd797a080e6ec802bd159460d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 7 May 2026 04:18:24 +0000 Subject: [PATCH] ci(mlx): expand LoRA targets to MLP + bump generation budget With batch_size=2 / gradient_accumulation_steps=3 (effective batch of 6) the q/k/v/o-only LoRA collapsed in 7 steps -- training loss kept dropping (0.55 vs the previous 1.02 with grad_accum=1) but inference output the structural skeleton ("My name") without recovering the specific "Unsloth" token. Switching to the standard unsloth target set (q/k/v/o + gate/up/down) gives the LoRA enough capacity to memorize the training row at the larger effective batch. Also bump max_tokens 24 -> 48 for the in-memory + reload generation calls so the model has more room to spew the memorized sequence; we still assert "Unsloth" appears anywhere in the completion. --- tests/studio/run_real_mlx_smoke.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 4992e7da06..7b36bb99ab 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -223,17 +223,27 @@ def cmd_train(args) -> int: mx.random.seed(SEED) with Phase("apply_lora", metrics): + # Standard unsloth LoRA target set (q/k/v/o + gate/up/down). + # With bs=2 grad_accum=3 (effective batch 6) the q/k/v/o-only + # LoRA collapsed in 7 steps -- training loss kept dropping but + # inference output the structural skeleton ("My name") without + # recovering the specific "Unsloth" token. Including the MLP + # projections gives the LoRA enough capacity to memorize the + # training row at the larger effective batch. model = FastMLXModel.get_peft_model( model, r = 8, lora_alpha = 16, lora_dropout = 0.0, - target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"], + target_modules = [ + "q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj", + ], use_gradient_checkpointing = False, random_state = SEED, finetune_language_layers = True, finetune_attention_modules = True, - finetune_mlp_modules = False, + finetune_mlp_modules = True, ) with Phase("pre_train_grad_probe", metrics): @@ -317,7 +327,7 @@ def cmd_train(args) -> int: model, tokenizer, prompt = PROMPT, - max_tokens = 24, + max_tokens = 48, verbose = False, ) metrics["in_memory_generation"] = in_mem_out @@ -435,7 +445,7 @@ def cmd_reload(args) -> int: m.eval() with Phase(f"generate_{args.format}", metrics): - out = generate(m, t, prompt = PROMPT, max_tokens = 24, verbose = False) + out = generate(m, t, prompt = PROMPT, max_tokens = 48, verbose = False) metrics["generation"] = out print(f" [reload:{args.format}] output: {out!r}", flush = True) assert (