[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
Daniel Han
parent
5bf74072e1
commit
0400928135
2 changed files with 109 additions and 63 deletions
|
|
@ -37,7 +37,7 @@ MODELS = [
|
|||
# (label, use_cce) — all use compile=True, gradient_checkpointing=True
|
||||
CONFIGS = [
|
||||
("Baseline+compile", False),
|
||||
("CCE+compile", True),
|
||||
("CCE+compile", True),
|
||||
]
|
||||
|
||||
BATCH_SIZE = 8
|
||||
|
|
@ -54,7 +54,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):
|
||||
|
||||
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
|
||||
|
|
@ -71,15 +72,22 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=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)
|
||||
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"],
|
||||
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)
|
||||
|
||||
|
|
@ -89,6 +97,7 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None):
|
|||
# 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
|
||||
|
|
@ -101,13 +110,15 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None):
|
|||
|
||||
# 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)
|
||||
dataset = load_dataset("emozilla/pg19-test", split = "test", streaming = True)
|
||||
|
||||
def create_batches(n):
|
||||
batch_input_ids = []
|
||||
|
|
@ -115,10 +126,10 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None):
|
|||
for item in dataset:
|
||||
encoded = tokenizer(
|
||||
item["text"],
|
||||
max_length=SEQ_LEN,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_tensors="np",
|
||||
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:
|
||||
|
|
@ -129,18 +140,18 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None):
|
|||
break
|
||||
|
||||
# Use MLX native Adafactor
|
||||
optimizer = mx_opt.Adafactor(learning_rate=LR)
|
||||
optimizer = mx_opt.Adafactor(learning_rate = LR)
|
||||
|
||||
config = TrainingConfig(
|
||||
batch_size=BATCH_SIZE,
|
||||
num_epochs=1,
|
||||
logging_steps=1,
|
||||
batch_size = BATCH_SIZE,
|
||||
num_epochs = 1,
|
||||
logging_steps = 1,
|
||||
)
|
||||
|
||||
trainer = MLXTrainer(
|
||||
model=model,
|
||||
optimizer=optimizer,
|
||||
config=config,
|
||||
model = model,
|
||||
optimizer = optimizer,
|
||||
config = config,
|
||||
)
|
||||
|
||||
gc.collect()
|
||||
|
|
@ -150,30 +161,36 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None):
|
|||
# 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)
|
||||
|
||||
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)
|
||||
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
|
||||
|
|
@ -182,7 +199,7 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None):
|
|||
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)
|
||||
|
||||
|
|
@ -190,16 +207,17 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None):
|
|||
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={
|
||||
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,
|
||||
|
|
@ -209,18 +227,18 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None):
|
|||
"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)
|
||||
print("Warning: wandb not installed, skipping logging.", file = sys.stderr)
|
||||
|
||||
result = {
|
||||
"ms_per_step": ms_per_step,
|
||||
|
|
@ -229,25 +247,29 @@ def run_worker(model_name, display_name, use_lora, use_cce, wandb_project=None):
|
|||
"nan_count": nan_count,
|
||||
}
|
||||
# Print JSON on a marker line so orchestrator can parse it
|
||||
print(f"__RESULT__ {json.dumps(result)}", flush=True)
|
||||
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,
|
||||
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)
|
||||
print(f" {line}", flush = True)
|
||||
proc.wait()
|
||||
return lines, proc.returncode
|
||||
|
||||
|
|
@ -257,7 +279,9 @@ def main(args):
|
|||
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"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)
|
||||
|
|
@ -279,18 +303,26 @@ def main(args):
|
|||
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) ---")
|
||||
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)),
|
||||
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])
|
||||
|
|
@ -301,7 +333,7 @@ def main(args):
|
|||
result = None
|
||||
for line in lines:
|
||||
if line.startswith("__RESULT__"):
|
||||
result = json.loads(line[len("__RESULT__"):])
|
||||
result = json.loads(line[len("__RESULT__") :])
|
||||
|
||||
if result and returncode == 0:
|
||||
ms = result["ms_per_step"]
|
||||
|
|
@ -310,7 +342,9 @@ def main(args):
|
|||
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}")
|
||||
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})")
|
||||
|
|
@ -321,7 +355,9 @@ def main(args):
|
|||
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")
|
||||
print(
|
||||
f"\n >> {display_name}: CCE+compile is {speedup:.2f}x speed, {mem_save:.1f}% mem saved"
|
||||
)
|
||||
|
||||
all_results.append((display_name, model_results))
|
||||
|
||||
|
|
@ -331,7 +367,9 @@ def main(args):
|
|||
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(
|
||||
f"{ 'Model':<22} {'BL+compile':>14} {'CCE+compile':>14} {'Speedup':>8} {'MemSave':>8}"
|
||||
)
|
||||
print("-" * 80)
|
||||
|
||||
for display_name, model_results in all_results:
|
||||
|
|
@ -351,7 +389,8 @@ def main(args):
|
|||
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}%")
|
||||
f"{speedup:>7.2f}x {mem_save:>7.1f}%"
|
||||
)
|
||||
|
||||
print("=" * 80)
|
||||
print("Done!")
|
||||
|
|
@ -359,18 +398,25 @@ def main(args):
|
|||
|
||||
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")
|
||||
|
||||
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)
|
||||
run_worker(
|
||||
args.model,
|
||||
args.display_name,
|
||||
bool(args.use_lora),
|
||||
bool(args.use_cce),
|
||||
args.wandb_project,
|
||||
)
|
||||
else:
|
||||
main(args)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ 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)
|
||||
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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue