diff --git a/benchmarks/mlx/reproduce_loss_curves.py b/benchmarks/mlx/reproduce_loss_curves.py new file mode 100644 index 0000000000..8c40c6b7aa --- /dev/null +++ b/benchmarks/mlx/reproduce_loss_curves.py @@ -0,0 +1,376 @@ +#!/usr/bin/env python3 +""" +Unsloth MLX Benchmark: Baseline+compile vs CCE+compile +======================================================= +Each (model, config) runs in its own subprocess for isolation. +Timing excludes warmup/compile steps (first 5 steps skipped). +All configs use gradient checkpointing + mx.compile. +Config order alternates per model to avoid systematic bias. +""" + +import json +import subprocess +import sys +import os +import time + +# ============================================================================= +# CONFIG +# ============================================================================= + +MODELS = [ + # (model_name, display_name, use_lora) + # --- Tiny (<1B) --- + ("mlx-community/Qwen3-0.6B-4bit", "Qwen3-0.6B-4bit", True), + # --- Small (1-2B) --- + ("mlx-community/Llama-3.2-1B-Instruct-bf16", "Llama-1B-full", False), + # --- Medium (3-4B) --- + ("mlx-community/Llama-3.2-3B-Instruct-4bit", "Llama-3B-4bit", True), + ("mlx-community/Qwen2.5-3B-Instruct-4bit", "Qwen2.5-3B-4bit", True), + ("mlx-community/Phi-3.5-mini-instruct-4bit", "Phi-3.5-mini-4bit", True), + ("mlx-community/Qwen2.5-3B-Instruct-8bit", "Qwen2.5-3B-8bit", True), + ("mlx-community/Llama-3.2-3B-Instruct-bf16", "Llama-3B-LoRA", True), + # --- Large (7-9B) --- + ("mlx-community/Mistral-7B-Instruct-v0.3-4bit", "Mistral-7B-4bit", True), +] + +# (label, use_cce) — all use compile=True, gradient_checkpointing=True +CONFIGS = [ + ("Baseline+compile", False), + ("CCE+compile", True), +] + +BATCH_SIZE = 8 +SEQ_LEN = 1024 +WARMUP_STEPS = 5 +MEASURE_STEPS = 95 +SEED = 42 +LR = 1e-5 +LORA_RANK = 8 +LORA_ALPHA = 16 + + +# ============================================================================= +# WORKER — runs a single (model, config) in isolation +# ============================================================================= + +def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None): + """Run in a subprocess. Prints training progress to stderr, JSON result to stdout.""" + import gc + import mlx.core as mx + import mlx.optimizers as mx_opt + from datasets import load_dataset + from transformers import AutoTokenizer + from unsloth.kernels.mlx.models import MLXLlamaForCausalLM + from unsloth.kernels.mlx.lora import get_peft_model, LoRAConfig as LoRAConfigLora + from unsloth.kernels.mlx.trainer import MLXTrainer, TrainingConfig + + # Load tokenizer from HuggingFace + tokenizer = AutoTokenizer.from_pretrained(model_name) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + # Load pure MLX model (dequantizes 4-bit weights to full precision MLX arrays) + print(f"Loading MLX model: {model_name}", file=sys.stderr) + model = MLXLlamaForCausalLM.from_pretrained(model_name) + + if use_lora: + lora_config = LoRAConfigLora( + r=LORA_RANK, + lora_alpha=LORA_ALPHA, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj"], + ) + model = get_peft_model(model, lora_config) + + # Store use_cce flag on the model so it can be passed during forward + model._use_cce = use_cce + + # Check if mx.fast.cce_loss is available (for mlx-cce) + # If available, we need to explicitly disable for baseline + import mlx.core as mx + has_fast_cce = hasattr(mx, "fast") and hasattr(mx.fast, "cce_loss") + + # If mx.fast.cce_loss exists, CCE is auto-enabled by default + # For baseline, we need to explicitly disable it + if has_fast_cce and not use_cce: + # Force disable CCE for baseline by passing explicit flag + use_cce_for_call = False + else: + use_cce_for_call = use_cce + + # Wrap __call__ to inject use_cce + _original_call = model.__call__ + def _call_with_cce(*args, **kwargs): + kwargs["use_cce"] = use_cce_for_call + return _original_call(*args, **kwargs) + model.__call__ = _call_with_cce + + # Local dataset loading and batching + dataset = load_dataset("emozilla/pg19-test", split="test", streaming=True) + + def create_batches(n): + batch_input_ids = [] + count = 0 + for item in dataset: + encoded = tokenizer( + item["text"], + max_length=SEQ_LEN, + padding="max_length", + truncation=True, + return_tensors="np", + ) + batch_input_ids.append(mx.array(encoded["input_ids"][0])) + if len(batch_input_ids) == BATCH_SIZE: + yield {"input_ids": mx.stack(batch_input_ids)} + batch_input_ids = [] + count += 1 + if count >= n: + break + + # Use MLX native Adafactor + optimizer = mx_opt.Adafactor(learning_rate=LR) + + config = TrainingConfig( + batch_size=BATCH_SIZE, + num_epochs=1, + logging_steps=1, + ) + + trainer = MLXTrainer( + model=model, + optimizer=optimizer, + config=config, + ) + + gc.collect() + mx.synchronize() + mx.reset_peak_memory() + + # We manually run the training loop to track per-step metrics accurately + loss_history = [] + step_times = [] + + total_steps = WARMUP_STEPS + MEASURE_STEPS + data_iter = create_batches(total_steps) + + print(f"Starting training loop (warmup={WARMUP_STEPS}, measure={MEASURE_STEPS})...", file=sys.stderr) + for i in range(total_steps): + try: + batch = next(data_iter) + except StopIteration: + break + + t0 = time.time() + step_result = trainer.training_step(batch) + mx.synchronize() + t1 = time.time() + + loss = step_result["loss"] + loss_history.append(loss) + + if i >= WARMUP_STEPS: + step_times.append(t1 - t0) + + if (i + 1) % 5 == 0 or i < 5: + print(f" Step {i+1}/{total_steps} | Loss: {loss:.4f} | Time: {(t1-t0)*1000:.0f}ms", file=sys.stderr) + + mx.synchronize() + peak_gb = mx.get_peak_memory() / 1e9 + + if not step_times: + ms_per_step = 0 + else: + ms_per_step = (sum(step_times) / len(step_times)) * 1000 + + final_loss = loss_history[-1] if loss_history else 0 + nan_count = sum(1 for l in loss_history if l != l) + + if wandb_project: + try: + import wandb + import mlx.core as mx + # Add suffix to distinguish mlx-cce vs regular mlx in W&B + has_fast_cce = hasattr(mx, "fast") and hasattr(mx.fast, "cce_loss") + cce_suffix = " (mlx-cce)" if has_fast_cce else " (mlx)" + label = ("CCE+compile" if use_cce else "Baseline+compile") + cce_suffix + run = wandb.init( + project=wandb_project, + name=f"{display_name} ({label})", + group=display_name, + reinit=True, + config={ + "model_name": model_name, + "display_name": display_name, + "use_lora": use_lora, + "use_cce": use_cce, + "batch_size": BATCH_SIZE, + "seq_len": SEQ_LEN, + "warmup_steps": WARMUP_STEPS, + "measure_steps": MEASURE_STEPS, + "lr": LR, + } + ) + for i, loss in enumerate(loss_history): + wandb.log({"loss": loss, "step": i + 1}) + + wandb.run.summary["ms_per_step"] = ms_per_step + wandb.run.summary["peak_gb"] = peak_gb + wandb.run.summary["final_loss"] = final_loss + wandb.run.summary["nan_count"] = nan_count + run.finish() + except ImportError: + print("Warning: wandb not installed, skipping logging.", file=sys.stderr) + + result = { + "ms_per_step": ms_per_step, + "peak_gb": peak_gb, + "final_loss": final_loss, + "nan_count": nan_count, + } + # Print JSON on a marker line so orchestrator can parse it + print(f"__RESULT__ {json.dumps(result)}", flush=True) + + +# ============================================================================= +# ORCHESTRATOR — spawns subprocesses and collects results +# ============================================================================= + +def run_subprocess(cmd): + """Run a subprocess, streaming stdout line by line. Returns (stdout_lines, returncode).""" + proc = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, bufsize=1, + ) + lines = [] + for line in proc.stdout: + line = line.rstrip("\n") + lines.append(line) + if not line.startswith("__RESULT__"): + print(f" {line}", flush=True) + proc.wait() + return lines, proc.returncode + + +def main(args): + print("=" * 80) + print("Unsloth MLX Benchmark: Baseline+compile vs CCE+compile") + print(" Gradient Checkpointing: ON | Compile: ON | Isolation: subprocess") + print("=" * 80) + print(f"Batch: {BATCH_SIZE}, Seq: {SEQ_LEN}, Warmup: {WARMUP_STEPS}, Measure: {MEASURE_STEPS}") + print(f"Optimizer: Adafactor (lr={LR})") + print(f"Models: {len(MODELS)}, Configs: {len(CONFIGS)}") + print("=" * 80) + + script_path = os.path.abspath(__file__) + python = sys.executable + all_results = [] + + for mi, (model_name, display_name, use_lora) in enumerate(MODELS): + print(f"\n{'='*80}") + print(f"[{mi+1}/{len(MODELS)}] {display_name}") + print(f" Repo: {model_name}, LoRA: {use_lora}") + print("=" * 80) + + # Alternate config order per model to avoid systematic bias + configs = CONFIGS if mi % 2 == 0 else list(reversed(CONFIGS)) + model_results = {} + + for label, use_cce in configs: + # Skip 8bit models for baseline since they won't load in standard transformers on MPS + if "8bit" in model_name.lower() and label == "Baseline+compile": + print(f"\n --- {label} (SKIPPED: 8bit baseline not supported on MPS) ---") + model_results[label] = None + continue + + print(f"\n --- {label} ---") + + cmd = [ + python, script_path, "--worker", + "--model", model_name, + "--display_name", display_name, + "--use_cce", str(int(use_cce)), + "--use_lora", str(int(use_lora)), + ] + if args.wandb: + cmd.extend(["--wandb_project", args.project]) + + lines, returncode = run_subprocess(cmd) + + # Parse result from stdout + result = None + for line in lines: + if line.startswith("__RESULT__"): + result = json.loads(line[len("__RESULT__"):]) + + if result and returncode == 0: + ms = result["ms_per_step"] + mem = result["peak_gb"] + loss = result["final_loss"] + nans = result["nan_count"] + model_results[label] = (ms, mem, loss, nans) + status = "OK" if nans == 0 else f"{nans} NaN" + print(f" >> {label}: {ms:.0f} ms/step | {mem:.2f} GB | loss={loss:.4f} | {status}") + else: + model_results[label] = None + print(f" >> {label}: FAILED (exit={returncode})") + + # Per-model summary + bl = model_results.get("Baseline+compile") + cce = model_results.get("CCE+compile") + if bl and cce: + speedup = bl[0] / cce[0] if cce[0] > 0 else 0 + mem_save = (bl[1] - cce[1]) / bl[1] * 100 if bl[1] > 0 else 0 + print(f"\n >> {display_name}: CCE+compile is {speedup:.2f}x speed, {mem_save:.1f}% mem saved") + + all_results.append((display_name, model_results)) + + # ========================================================================= + # FINAL SUMMARY + # ========================================================================= + print(f"\n\n{'='*80}") + print("FINAL SUMMARY — Baseline+compile vs CCE+compile (GC=ON, compile=ON)") + print("=" * 80) + print(f"{ 'Model':<22} {'BL+compile':>14} {'CCE+compile':>14} {'Speedup':>8} {'MemSave':>8}") + print("-" * 80) + + for display_name, model_results in all_results: + bl = model_results.get("Baseline+compile") + cce = model_results.get("CCE+compile") + + if not bl: + print(f"{display_name:<22} {'FAILED':>14}") + continue + if not cce: + print(f"{display_name:<22} {bl[0]:>7.0f}ms {bl[1]:>5.1f}GB {'FAILED':>14}") + continue + + speedup = bl[0] / cce[0] if cce[0] > 0 else 0 + mem_save = (bl[1] - cce[1]) / bl[1] * 100 if bl[1] > 0 else 0 + print( + f"{display_name:<22} " + f"{bl[0]:>7.0f}ms {bl[1]:>5.1f}GB " + f"{cce[0]:>7.0f}ms {cce[1]:>5.1f}GB " + f"{speedup:>7.2f}x {mem_save:>7.1f}%") + + print("=" * 80) + print("Done!") + + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser(allow_abbrev=False) + parser.add_argument("--worker", action="store_true") + parser.add_argument("--model", type=str) + parser.add_argument("--display_name", type=str) + parser.add_argument("--use_cce", type=int, default=0) + parser.add_argument("--use_lora", type=int, default=1) + parser.add_argument("--wandb_project", type=str, default=None) + parser.add_argument("--wandb", action="store_true") + parser.add_argument("--project", type=str, default="unsloth-mlx-benchmark") + args = parser.parse_args() + + if args.worker: + run_worker(args.model, args.display_name, bool(args.use_lora), bool(args.use_cce), args.wandb_project) + else: + main(args) diff --git a/test_embedding.py b/test_embedding.py new file mode 100644 index 0000000000..fc54259a21 --- /dev/null +++ b/test_embedding.py @@ -0,0 +1,23 @@ +import torch +import mlx.core as mx +from unsloth.kernels.mlx.bridge import mlx_to_torch + +arr = mx.array([[1, 2, 3], [4, 5, 6]], dtype=mx.int32) +tensor = mlx_to_torch(arr) +print("Tensor:", tensor) +print("Tensor device:", tensor.device) + +emb = torch.nn.Embedding(10, 5).to("mps") +try: + res = emb(tensor.to("mps")) + print("Success") +except Exception as e: + print("Error:", e) + +# Test with a clone +tensor2 = tensor.clone() +try: + res = emb(tensor2.to("mps")) + print("Success with clone") +except Exception as e: + print("Error with clone:", e)