#!/usr/bin/env python """Iterative auto-tuner for AMD MI300X / ROCm 7.0 workloads. Three modes, picked with `--mode`: hardcoded (default) Walks through a curated list of MI300X-specific tuning changes one at a time. Deterministic, no LLM required — experiments are derived from the rules in kb/rocm_rules.yaml. llm On each iteration, asks the LLM backend (qwen-hf via HF_TOKEN, or qwen-vllm via GOBLIN_QWEN_VLLM_URL) for ONE next experiment given the live waste_budget, history, and KB rules. Greedy coordinate descent — accept changes that beat the current best by the improvement threshold, otherwise revert. llm-explore On each iteration, asks the LLM for K candidate experiments at once (--candidates-per-iteration, default 3). Benchmarks all K, picks the one with the highest tokens/sec, and accepts only if it beats the current best. Higher GPU cost (~Kx benchmarks per iteration) but better at finding interaction effects that greedy one-at-a-time can miss. After each change, runs a real benchmark via goblin_runner.sh and keeps the change only if tokens/sec improved meaningfully (>1% by default — the threshold cuts measurement noise). Stops when N consecutive experiments produce no improvement, or when the source of experiments is exhausted. Usage: # hardcoded mode (default): python scripts/auto_tune.py workloads/train_qwen_lora.py --steps 20 # LLM-driven greedy mode: python scripts/auto_tune.py workloads/train_qwen_lora.py \\ --mode llm --steps 20 # LLM-driven multi-candidate exploration: python scripts/auto_tune.py workloads/train_qwen_lora.py \\ --mode llm-explore --candidates-per-iteration 3 --steps 20 Output: - A row-by-row log of each experiment attempted, accepted or rejected - A final summary with cumulative speedup - A pointer to a temp file containing the best workload script for diff-against-baseline inspection Extending hardcoded mode: add an Experiment to EXPERIMENTS. The substitutions field is a list of (regex_pattern, replacement) tuples applied with re.subn against the workload source. env_vars are exported into the goblin_runner.sh subprocess and persist on every accepted iteration. """ from __future__ import annotations import argparse import asyncio import json import os import re import subprocess import sys import tempfile from dataclasses import dataclass, field from pathlib import Path REPO_ROOT = Path(__file__).resolve().parent.parent GOBLIN_RUNNER = REPO_ROOT / "runner" / "goblin_runner.sh" sys.path.insert(0, str(REPO_ROOT)) # Optional structured-events output. When `--events FILE` is passed, the # script appends one JSON object per line at key milestones (baseline, # iteration start, candidate done, iteration done, summary). Used by the # Streamlit UI to render progress live; CLI users typically don't need it. _EVENTS_PATH: Path | None = None def _emit(event: dict) -> None: """Append one NDJSON event to the events file if one was configured.""" if _EVENTS_PATH is None: return try: with _EVENTS_PATH.open("a") as f: f.write(json.dumps(event, default=str) + "\n") f.flush() # so a UI tailing the file sees events promptly except OSError: pass # never crash the run on an event-write failure # Default workload template — used when the user passes --model instead # of an explicit workload path. We just substitute MODEL_ID and reuse all # the other defaults (fp16, batch=4, eager attention, LoRA r=16, …). _DEFAULT_WORKLOAD_TEMPLATE = REPO_ROOT / "workloads" / "train_qwen_lora.py" def _generate_workload_from_model(model_id: str, dest: Path) -> Path: """Build a baseline workload by substituting MODEL_ID into the demo template (`workloads/train_qwen_lora.py`). Writes to `dest`, returns the path. Caveats: - Uses the demo's LoRA target_modules (`q_proj`, `v_proj`) which work for the major decoder-only LLM families (Qwen, Llama, Mistral, Gemma). MoE / GPT-2-style architectures will need a custom workload. - The template overwrites HF_TOKEN with a redactable fake. Public models load fine; gated models (Llama, etc.) need the user to edit the generated workload or use a custom one. """ if not _DEFAULT_WORKLOAD_TEMPLATE.exists(): raise SystemExit( f"--model needs the template at {_DEFAULT_WORKLOAD_TEMPLATE}, but it's missing" ) template_src = _DEFAULT_WORKLOAD_TEMPLATE.read_text() new_src, n = re.subn( r'MODEL_ID = "[^"]*"', f'MODEL_ID = "{model_id}"', template_src, ) if n == 0: raise SystemExit( f"Couldn't find `MODEL_ID = \"...\"` in {_DEFAULT_WORKLOAD_TEMPLATE} " "to substitute. Has the template format changed?" ) dest.write_text(new_src) return dest # POSIX env var name: leading letter or underscore, then alnum/underscore. # subprocess.run() raises ValueError if any key in the env dict violates # this. We validate up-front rather than letting the subprocess crash. _VALID_ENV_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") def _sanitize_env_vars(envs: dict, context: str = "") -> dict[str, str]: """Clean an env_vars dict from the LLM: 1. Strip dotted prefixes (`env_vars.X` → `X`) the LLM mimics from the KB transform notation. 2. Drop any key that still isn't a valid POSIX env var name. Warns instead of crashing — the LLM occasionally embeds shell syntax (e.g. `'NUMACTL_INTERLEAVE=1'` as a key) which would make subprocess.run raise ValueError. """ cleaned: dict[str, str] = {} for k, v in envs.items(): key = str(k) if "." in key: stripped = key.rsplit(".", 1)[-1] tag = f" [{context}]" if context else "" print(f" [warn]{tag} dotted env key {key!r}; using {stripped!r}") key = stripped if not _VALID_ENV_NAME.match(key): tag = f" [{context}]" if context else "" print( f" [warn]{tag} dropping invalid env var name {key!r} " "(must match [A-Za-z_][A-Za-z0-9_]*)" ) continue cleaned[key] = str(v) return cleaned @dataclass class Experiment: name: str description: str rationale: str substitutions: list[tuple[str, str]] = field(default_factory=list) env_vars: dict[str, str] = field(default_factory=dict) # Curated for ROCm 7.0 + MI300X (CDNA3, 192 GB HBM3). Ordered roughly by # typical impact on Qwen-shaped LoRA fine-tuning workloads. Each # experiment stacks on top of any previously accepted ones. EXPERIMENTS: list[Experiment] = [ Experiment( name="bf16_over_fp16", description="Switch precision from fp16 to bf16", rationale=( "MI300X (CDNA3) prefers bf16: same throughput, larger numeric " "range, no loss-scaler needed. fp16 underutilizes the matrix " "engine on this arch." ), substitutions=[ (r"torch_dtype=torch\.float16", "torch_dtype=torch.bfloat16"), (r"\bfp16=True\b", "bf16=True"), ], ), Experiment( name="batch_size_8", description="Increase per_device_train_batch_size 4 → 8", rationale="MI300X has 192 GB HBM; batch=4 leaves it on the floor.", substitutions=[ (r"per_device_train_batch_size=4\b", "per_device_train_batch_size=8"), ], ), Experiment( name="batch_size_16", description="Further increase per_device_train_batch_size to 16", rationale="If batch=8 fit and improved, try doubling again.", substitutions=[ (r"per_device_train_batch_size=\d+", "per_device_train_batch_size=16"), ], ), Experiment( name="batch_size_32", description="Push per_device_train_batch_size to 32", rationale=( "MI300X has 192 GB HBM3 — batch 16 typically peaks ~130 GB. " "If 16 fit, 32 likely fits too and reduces step overhead per " "token. Reverts cleanly via OOM-as-crash if not." ), substitutions=[ (r"per_device_train_batch_size=\d+", "per_device_train_batch_size=32"), ], ), Experiment( name="sdpa_attention", description="Switch attention from eager to SDPA", rationale=( "Eager attention is the slowest path. SDPA dispatches to the " "best available kernel (flash on ROCm 7.x where supported, " "memory-efficient elsewhere)." ), substitutions=[ (r'attn_implementation="eager"', 'attn_implementation="sdpa"'), ], ), Experiment( name="dataloader_workers_4", description="Bump dataloader_num_workers 0 → 4", rationale=( "0 workers means the GPU sits idle while the host loads the " "next batch. 4 is a safe value across most CPU configs." ), substitutions=[ (r"dataloader_num_workers=0", "dataloader_num_workers=4"), (r"num_workers=0", "num_workers=4"), ], ), Experiment( name="pin_memory", description="Enable dataloader_pin_memory", rationale=( "Pinned host buffers make H2D copies async and overlap with " "the GPU. Worth it once you have >0 dataloader workers." ), substitutions=[ (r"dataloader_pin_memory=False", "dataloader_pin_memory=True"), (r"\bpin_memory=False\b", "pin_memory=True"), ], ), Experiment( name="env_hipblaslt", description="Set TORCH_BLAS_PREFER_HIPBLASLT=1", rationale=( "hipBLASLt is significantly faster than rocBLAS for the GEMM " "shapes Qwen produces (LoRA-projected attention)." ), env_vars={"TORCH_BLAS_PREFER_HIPBLASLT": "1"}, ), Experiment( name="env_tunable_op", description="Set PYTORCH_TUNABLEOP_ENABLED=1", rationale=( "Enables runtime kernel auto-tuning. Pays a first-run " "warmup cost in exchange for a steady-state win on every " "subsequent step." ), env_vars={"PYTORCH_TUNABLEOP_ENABLED": "1"}, ), Experiment( name="env_miopen_find", description="Set MIOPEN_FIND_MODE=3", rationale=( "MIOpen FAST mode picks already-tuned kernels without on-the-" "fly search. Reduces per-step variance." ), env_vars={"MIOPEN_FIND_MODE": "3"}, ), ] # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def apply_substitutions(source: str, subs: list[tuple[str, str]]) -> str | None: """Apply each (pattern, replacement) in order. Returns the new source, or None if any pattern matched zero times (already applied or N/A for this workload).""" out = source for pattern, replacement in subs: new, n = re.subn(pattern, replacement, out) if n == 0: return None out = new return out def benchmark( workload_path: Path, steps: int, env_overrides: dict[str, str], timeout: int = 600, ) -> dict | None: """Run goblin_runner.sh on the workload, return parsed RunMetrics dict or None on failure.""" with tempfile.TemporaryDirectory(prefix="auto_tune_run_") as out_dir_str: out_dir = Path(out_dir_str) env = os.environ.copy() env["USER_SCRIPT"] = str(workload_path) env["OUT_DIR"] = str(out_dir) env["STEPS"] = str(steps) # Candidate workload lives in /tmp, so its self-bootstrap line # `sys.path.insert(0, dirname(dirname(__file__)))` resolves to /tmp # — which has no `workloads/` package. Inject the real repo root via # PYTHONPATH so `from workloads._runtime import ...` succeeds. existing_pp = env.get("PYTHONPATH", "") env["PYTHONPATH"] = ( str(REPO_ROOT) + (os.pathsep + existing_pp if existing_pp else "") ) env.update(env_overrides) try: proc = subprocess.run( [str(GOBLIN_RUNNER)], env=env, capture_output=True, text=True, timeout=timeout, ) except subprocess.TimeoutExpired: print(f" TIMEOUT after {timeout}s") return None except ValueError as exc: # subprocess.run validates env var names and raises ValueError # for malformed keys (e.g. names containing '=' or spaces). The # LLM has occasionally emitted those; we sanitize earlier but # this is the last-resort backstop so a single bad candidate # doesn't crash the whole tuning run. print(f" REJECTED — illegal env var name(s): {exc}") print(f" env keys offered: {list(env_overrides.keys())}") return None except OSError as exc: print(f" REJECTED — could not spawn goblin_runner.sh: {exc}") return None if proc.returncode != 0: print(f" goblin_runner.sh failed (exit {proc.returncode})") tail = (proc.stderr or "").strip().splitlines()[-8:] for line in tail: print(f" | {line}") return None try: from runner import profile_parser metrics = profile_parser.parse(out_dir, steps=steps) return metrics.model_dump() except Exception as exc: # parser is defensive but be safe print(f" profile_parser raised: {type(exc).__name__}: {exc}") return None def _delta_pct(new: float, baseline: float) -> float: if baseline <= 0: return 0.0 return (new - baseline) / baseline * 100.0 # --------------------------------------------------------------------------- # LLM-driven experiment generator # --------------------------------------------------------------------------- _LLM_SYSTEM_PROMPT = """\ You are an expert at tuning AMD MI300X (ROCm 7.0, CDNA3 arch, 192 GB HBM3) training workloads. The user is iteratively benchmarking changes to a transformers/peft fine-tuning script. On each turn you suggest ONE specific parameter change to try next, targeting the largest non-useful waste bucket in the most recent benchmark. Your output MUST be a single JSON object with this exact shape (no prose, no markdown fences, just the object): { "name": "short_snake_case_name", "rationale": "1-3 sentences on why this change addresses the worst waste bucket", "substitutions": [["regex_pattern", "replacement"]], "env_vars": {"VAR_NAME": "value"} } CRITICAL output rules — read carefully: 1. env_vars keys are LITERAL POSIX shell environment variable names. They MUST match the regex [A-Za-z_][A-Za-z0-9_]* — letters, digits, underscores only, starting with a letter or underscore. - NEVER prefix them with "env_vars." or any other dotted path. - NEVER include "=" or shell syntax in the key — env var names are identifiers, NOT assignments and NOT commands. - If you want to invoke a command-line tool like `numactl` or `taskset`, that CANNOT be expressed as an env_var. Don't try. Either propose a `substitutions` change to the script, or skip. Wrong: {"env_vars.MIOPEN_FIND_MODE": "3"} Wrong: {"NUMACTL_INTERLEAVE=1": "numactl --interleave=all"} Wrong: {"export FOO": "bar"} Right: {"MIOPEN_FIND_MODE": "3"} Right: {"TORCH_BLAS_PREFER_HIPBLASLT": "1"} 2. substitutions are (regex_pattern, replacement) pairs applied with re.subn against the current workload source. Patterns must match at least one occurrence in the source — if zero matches, the experiment is auto-skipped (counted as no improvement). 3. When the previous change for a parameter improved tokens/sec, push that parameter further in the same direction next time. E.g. if batch_size 4 → 8 won, try 8 → 16. If 16 won and HBM is still under ~150 GB, try 32. Don't be timid — MI300X has 192 GB HBM3. 4. Don't repeat any (name OR substitution OR env_var combo) from history. If a change was rejected, don't propose the same numerical value again — try a different one. 5. If you cannot think of a productive next change, output: {"name": "STOP", "rationale": "", "substitutions": [], "env_vars": {}} CONCRETE OUTPUT EXAMPLES — match this shape exactly: Switch fp16 → bf16 (precision_path bucket): {"name": "bf16_over_fp16", "rationale": "MI300X CDNA3 matrix cores prefer bf16: same throughput, larger numeric range, no loss-scaler.", "substitutions": [["fp16=True", "bf16=True"], ["torch_dtype=torch\\\\.float16", "torch_dtype=torch.bfloat16"]], "env_vars": {}} Increase batch size to 16 (memory_headroom bucket): {"name": "batch_size_16", "rationale": "Current HBM peak is well under 192 GB; bigger batch saturates the GPU.", "substitutions": [["per_device_train_batch_size=\\\\d+", "per_device_train_batch_size=16"]], "env_vars": {}} Switch attention to SDPA (kernel_shape bucket): {"name": "sdpa_attention", "rationale": "Eager attention is the slowest path; SDPA dispatches to a tuned kernel.", "substitutions": [["attn_implementation=\\"eager\\"", "attn_implementation=\\"sdpa\\""]], "env_vars": {}} Bump dataloader workers (data_wait bucket): {"name": "dataloader_workers_4", "rationale": "0 workers starves the GPU between batches.", "substitutions": [["dataloader_num_workers=0", "dataloader_num_workers=4"]], "env_vars": {}} Set MIOpen FAST mode (kernel_shape bucket, env-only): {"name": "miopen_find_fast", "rationale": "FAST mode picks already-tuned kernels without on-the-fly search.", "substitutions": [], "env_vars": {"MIOPEN_FIND_MODE": "3"}} Prefer hipBLASLt (kernel_shape bucket, env-only): {"name": "prefer_hipblaslt", "rationale": "hipBLASLt is faster than rocBLAS for Qwen GEMM shapes on MI300X.", "substitutions": [], "env_vars": {"TORCH_BLAS_PREFER_HIPBLASLT": "1"}} """ _LLM_USER_TEMPLATE = """\ Hardware facts (use these — do not contradict): - AMD MI300X, CDNA3 architecture, 192 GB HBM3 - bf16 throughput on CDNA3 ≈ same as fp16, > fp32 (matrix engine is fp16/bf16/fp8 native) - fp32 is the SLOWEST option on this arch — never suggest it as an improvement Known incompatibilities for THIS workload (peft + LoRA on transformers Trainer): {incompatibilities} KB rules (one-liner per rule, for grounding): {kb_summary} Current accepted workload state — these are the literal values in the script after every change accepted so far. The next change you propose should mutate one of these (or set an env var). DO NOT propose a value that's already present here. {tunables} Latest benchmark (this is the result of the most recent ACCEPTED state): - tokens_per_sec: {tps:.1f} - mfu_pct: {mfu:.2f} (% of MI300X dense bf16 peak; healthy LoRA ranges 30-50%) - gpu_util_pct: {util:.1f} - hbm_peak_gb: {hbm:.2f} - waste_budget (seconds/step): {waste_lines} Sorted recoverable waste (largest first — go after these): {recoverable_sorted} History of changes already tried this run (newest first; outcomes are "accepted" / "rejected" / "crashed" / "skipped"): {history_lines} If the latest entry is "crashed", the change you propose next must be STRUCTURALLY different (different parameter, not just a different value of the same one). Suggest ONE next change targeting the largest recoverable bucket. JSON only. """ # Workload-specific incompatibilities the LLM otherwise wastes iterations on. # Keep this list short and concrete — it goes into every prompt. _KNOWN_INCOMPATIBILITIES = [ "gradient_checkpointing=True requires `model.enable_input_require_grads()`" " before peft wrapping for LoRA models. Setting it via a single substitution" " WILL CRASH the workload. Don't propose it.", "bitsandbytes-based optimizers (`adamw_8bit`, `paged_adamw_8bit`) and" " `load_in_8bit=True` are NOT supported on ROCm 7.x. Don't propose them.", "torch_compile=True with peft/LoRA on ROCm 7.x triggers compile-time" " errors with the current PyTorch nightly (2.9.x). Don't propose it" " unless you have specific evidence it works on this version.", "flash_attention_2 may not be installed (try `attn_implementation=\"sdpa\"`" " before `\"flash_attention_2\"`).", "persistent_workers=True requires num_workers > 0. PyTorch raises" " `ValueError: persistent_workers option needs num_workers > 0` if you" " enable it while num_workers=0. If the current workload has" " dataloader_num_workers=0, do NOT propose persistent_workers=True" " alone — pair it with `dataloader_num_workers=4` (or higher) in the" " SAME experiment via two substitutions, or wait until a previous" " experiment has bumped num_workers above 0.", "dataloader_prefetch_factor only works when num_workers > 0 (same" " constraint as persistent_workers). Same rule: bump num_workers in" " the same experiment, or skip.", ] def _kb_summary(rules_yaml_path: Path, max_chars: int = 6000) -> str: """Return a compact one-line-per-rule summary of kb/rocm_rules.yaml. Notably we DO NOT show the raw `transform` field — earlier versions did and the LLM ended up copying its dotted-path notation literally (`env_vars.MIOPEN_FIND_MODE` as the env var name, not as a dict accessor). The system prompt's CONCRETE EXAMPLES section is the canonical source of truth for output shape; this summary just grounds the LLM's reasoning in the catalog of known issues. """ if not rules_yaml_path.exists(): return "(KB rules file not found)" try: import yaml rules = yaml.safe_load(rules_yaml_path.read_text()) or [] except Exception as exc: return f"(failed to parse KB: {exc})" lines = [] for r in rules: if not isinstance(r, dict): continue rid = r.get("id", "?") bucket = r.get("targets_bucket", "?") sym = (r.get("symptom") or "").strip().replace("\n", " ") if len(sym) > 110: sym = sym[:107] + "..." lines.append(f"- {rid:55s} [{bucket}] {sym}") text = "\n".join(lines) if len(text) > max_chars: text = text[:max_chars] + "\n... (truncated)" return text # Map of (substring-in-source) → (parameter description, example regex # pattern, example replacement template). Each entry is a hint shown to # the LLM so it has a concrete target to point its substitutions at — # instead of guessing what the workload's literal config text looks like. _TUNABLE_HINTS: list[tuple[str, str, str, str]] = [ # (token to detect, description, regex_for_substitution, replacement_template) ("torch_dtype=torch.float16", "model precision (matches `torch_dtype=torch.float16`)", r"torch_dtype=torch\.float16", "torch_dtype=torch.bfloat16"), ("torch_dtype=torch.bfloat16", "model precision (already bf16)", r"torch_dtype=torch\.bfloat16", "torch_dtype=torch.float16"), ("fp16=True", "TrainingArguments fp16 (matches `fp16=True`)", r"\bfp16=True\b", "bf16=True"), ("bf16=True", "TrainingArguments bf16 (already bf16)", r"\bbf16=True\b", "fp16=True"), ("attn_implementation=\"eager\"", "attention impl (matches `attn_implementation=\"eager\"`)", r'attn_implementation="eager"', 'attn_implementation="sdpa"'), ("attn_implementation=\"sdpa\"", "attention impl (currently sdpa; could try flash_attention_2)", r'attn_implementation="sdpa"', 'attn_implementation="flash_attention_2"'), ("per_device_train_batch_size=", "per-device batch size (matches `per_device_train_batch_size=`)", r"per_device_train_batch_size=\d+", "per_device_train_batch_size="), ("dataloader_num_workers=", "dataloader workers (matches `dataloader_num_workers=`)", r"dataloader_num_workers=\d+", "dataloader_num_workers="), ("dataloader_pin_memory=", "dataloader pin_memory (matches `dataloader_pin_memory=`)", r"dataloader_pin_memory=(True|False)", "dataloader_pin_memory=True"), ("gradient_checkpointing=", "gradient checkpointing toggle", r"gradient_checkpointing=(True|False)", "gradient_checkpointing=True"), ("torch_compile=", "torch.compile toggle (use cautiously on ROCm 7.x)", r"torch_compile=(True|False)", "torch_compile=True"), ("optim=\"adamw_torch\"", "optimizer choice (currently adamw_torch)", r'optim="adamw_torch"', 'optim="adamw_torch_fused"'), ] def _tunables_summary(source: str) -> str: """Detect which tunable parameters are present in the workload source and surface their current literal values + ready-to-use regex patterns so the LLM has concrete substitution targets. Skips comment lines when reporting the "current" value — many workloads document expected findings in a top-of-file comment block, and we want the LLM to see the live config line, not the doc string. """ lines: list[str] = [] source_lines = source.splitlines() for token, desc, pattern, replacement in _TUNABLE_HINTS: live_line: str | None = None for raw in source_lines: stripped = raw.lstrip() if stripped.startswith("#"): continue if token in raw: live_line = raw.strip() break if live_line is None: continue lines.append( f" • {desc}\n" f" current: {live_line}\n" f" pattern: {pattern!r} replacement template: {replacement!r}" ) if not lines: return " (no recognized tunables — substitutions will need to match other text)" return "\n".join(lines) def _recoverable_sorted(waste: dict) -> str: """List the non-useful_gpu buckets sorted by size, so the LLM can explicitly target the biggest one first.""" if not waste: return " (no waste_budget available)" items = [ (name, value) for name, value in waste.items() if name != "useful_gpu" and isinstance(value, (int, float)) ] items.sort(key=lambda kv: kv[1], reverse=True) if not items: return " (no recoverable buckets)" return "\n".join(f" {i + 1}. {name:18s} = {value:.4f}" for i, (name, value) in enumerate(items)) def _config_snippet(source: str, max_lines: int = 80) -> str: """Return the lines around `TrainingArguments(` and `from_pretrained(` so the LLM sees the actual config it's modifying without us shipping the whole script. Gives ~max_lines of context. """ lines = source.splitlines() keep: list[tuple[int, str]] = [] for i, line in enumerate(lines): lower = line.lower() if any( tok in lower for tok in ( "trainingarguments(", "from_pretrained(", "loraconfig(", "dataloader(", "torch_dtype", "attn_implementation", "fp16=", "bf16=", "per_device_train_batch_size", "dataloader_num_workers", "dataloader_pin_memory", "gradient_checkpointing", "torch_compile", "optim=", ) ): keep.append((i, line)) if not keep: return source[:2000] # Coalesce nearby line indices into windows for readability windows: list[list[str]] = [] last_idx = -10 cur: list[str] = [] for i, line in keep: if i - last_idx > 3: if cur: windows.append(cur) cur = [] cur.append(f"{i + 1:4d}: {line}") last_idx = i if cur: windows.append(cur) out = "\n\n".join("\n".join(w) for w in windows) if out.count("\n") > max_lines: out_lines = out.splitlines()[:max_lines] out = "\n".join(out_lines) + "\n... (truncated)" return out def _format_history(history: list[dict]) -> str: if not history: return "(none yet — this is the first iteration)" lines = [] for h in reversed(history[-12:]): # last 12 newest-first outcome = h.get("outcome", "?") delta = h.get("delta_pct") delta_s = f"{delta:+.2f}%" if delta is not None else "n/a" subs = h.get("substitutions") or [] envs = h.get("env_vars") or {} change_repr = f"subs={subs} env={envs}" lines.append(f"- {h['name']:25s} {outcome:9s} Δ {delta_s:8s} {change_repr}") return "\n".join(lines) def _format_waste(waste: dict) -> str: keys = ( "useful_gpu", "data_wait", "host_gap", "comm_excess", "memory_headroom", "precision_path", "kernel_shape", ) return "\n".join(f" {k:18s} = {waste.get(k, 0.0):.4f}" for k in keys) def _build_llm_backend(system_prompt: str = _LLM_SYSTEM_PROMPT, max_tokens: int = 1024): """Construct the same backend the agent loop uses. Surfaces a clear message if neither HF_TOKEN nor a vLLM URL is configured.""" has_hf = bool(os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN")) has_vllm = bool(os.environ.get("GOBLIN_QWEN_VLLM_URL")) backend_kind = os.environ.get("GOBLIN_AGENT_BACKEND", "qwen-hf").lower() if backend_kind in ("qwen-hf", "qwen", "hf", "") and not has_hf: raise SystemExit( "LLM mode requires HF_TOKEN (qwen-hf backend) or " "GOBLIN_AGENT_BACKEND=qwen-vllm + GOBLIN_QWEN_VLLM_URL." ) if backend_kind in ("qwen-vllm", "qwen_vllm", "vllm", "local") and not has_vllm: raise SystemExit( "LLM mode with qwen-vllm backend requires GOBLIN_QWEN_VLLM_URL." ) from agent.backends import make_backend return make_backend(system_prompt=system_prompt, max_tokens=max_tokens) async def _ask_llm_for_experiment( backend, *, kb_summary: str, source: str, metrics: dict, history: list[dict], ) -> Experiment | None: """One LLM turn → one Experiment (or None for STOP / parse failure).""" waste = metrics.get("waste_budget") or {} prompt = _LLM_USER_TEMPLATE.format( incompatibilities="\n".join(f"- {line}" for line in _KNOWN_INCOMPATIBILITIES), kb_summary=kb_summary, tunables=_tunables_summary(source), tps=metrics.get("tokens_per_sec", 0.0), mfu=metrics.get("mfu_pct", 0.0), util=metrics.get("gpu_util_pct", 0.0), hbm=metrics.get("hbm_peak_gb", 0.0), waste_lines=_format_waste(waste), recoverable_sorted=_recoverable_sorted(waste), history_lines=_format_history(history), ) backend.add_user_message(prompt) turn = await backend.next_turn(tool_schemas=[]) raw = " ".join(turn.text_blocks).strip() obj = _extract_json_object(raw) if obj is None: print(f" LLM response was not parseable JSON. Raw: {raw[:300]!r}") return None name = (obj.get("name") or "").strip() if not name or name.upper() == "STOP": print(f" LLM signaled STOP: {obj.get('rationale', '(no rationale)')}") return None subs_raw = obj.get("substitutions") or [] envs = obj.get("env_vars") or {} if not subs_raw and not envs: print(f" LLM returned an empty experiment ({name}); skipping") return None subs: list[tuple[str, str]] = [] for entry in subs_raw: if isinstance(entry, list) and len(entry) == 2: subs.append((str(entry[0]), str(entry[1]))) elif isinstance(entry, dict) and "pattern" in entry and "replacement" in entry: subs.append((str(entry["pattern"]), str(entry["replacement"]))) cleaned_envs = _sanitize_env_vars(envs, context=name) if not subs and not cleaned_envs: # Everything got dropped during sanitization (bad env names + no # valid substitutions). Treat as a no-op rather than benchmarking # an unchanged workload. print(f" LLM experiment {name!r} had nothing valid after sanitization; skipping") return None return Experiment( name=name, description=obj.get("description") or name, rationale=str(obj.get("rationale") or ""), substitutions=subs, env_vars=cleaned_envs, ) # --------------------------------------------------------------------------- # llm-explore mode: ask for K candidates per iteration # --------------------------------------------------------------------------- _LLM_EXPLORE_SYSTEM_PROMPT = """\ You are an expert at tuning AMD MI300X (ROCm 7.0, CDNA3 arch, 192 GB HBM3) training workloads. The user is running a multi-candidate exploration: on every iteration you suggest K STRUCTURALLY-DIFFERENT candidate changes, the user benchmarks all of them, and the best one is accepted (if it beats the current best by the threshold). Your output MUST be a JSON ARRAY of K objects, no prose, no markdown fences, just the array: [ {"name": "...", "rationale": "...", "substitutions": [["regex", "repl"]], "env_vars": {"VAR": "value"}}, {"name": "...", "rationale": "...", "substitutions": [["regex", "repl"]], "env_vars": {"VAR": "value"}}, {"name": "...", "rationale": "...", "substitutions": [["regex", "repl"]], "env_vars": {"VAR": "value"}} ] CRITICAL output rules: 1. Each candidate must target a DIFFERENT waste bucket or parameter category than the others. Diversity beats redundancy — don't propose three batch-size bumps; propose one batch bump, one env var, one precision/attention/dataloader change. 2. env_vars keys are LITERAL POSIX shell environment variable names — they MUST match the regex [A-Za-z_][A-Za-z0-9_]*. NEVER prefix them with "env_vars." or any other dotted path. NEVER include "=" or shell syntax in the key. If you want to invoke a CLI tool like `numactl`, that's NOT an env var — skip the candidate entirely. Wrong: {"env_vars.MIOPEN_FIND_MODE": "3"} Wrong: {"NUMACTL_INTERLEAVE=1": "numactl --interleave=all"} Right: {"MIOPEN_FIND_MODE": "3"} 3. substitutions are (regex_pattern, replacement) pairs applied with re.subn. Patterns must match at least one occurrence — if zero matches, that candidate is skipped. 4. NEVER propose a (substitutions, env_vars) combination that already appears in history with outcome rejected/crashed. Diversify within the array AND across the run. 5. If you genuinely cannot find K productive candidates, output fewer (e.g. 2 if K=3). The user will benchmark whatever you provide. If you have zero productive candidates, output: [{"name": "STOP", "rationale": "", "substitutions": [], "env_vars": {}}] CONCRETE OUTPUT EXAMPLES (for K=3): [ {"name": "bf16_over_fp16", "rationale": "Largest recoverable bucket is precision_path; CDNA3 prefers bf16.", "substitutions": [["fp16=True", "bf16=True"], ["torch_dtype=torch\\\\.float16", "torch_dtype=torch.bfloat16"]], "env_vars": {}}, {"name": "batch_size_16", "rationale": "HBM peak well under 192 GB; bigger batch saturates the GPU.", "substitutions": [["per_device_train_batch_size=\\\\d+", "per_device_train_batch_size=16"]], "env_vars": {}}, {"name": "prefer_hipblaslt", "rationale": "hipBLASLt outperforms rocBLAS on Qwen GEMM shapes.", "substitutions": [], "env_vars": {"TORCH_BLAS_PREFER_HIPBLASLT": "1"}} ] """ _LLM_EXPLORE_USER_TEMPLATE = """\ Hardware facts (use these — do not contradict): - AMD MI300X, CDNA3 architecture, 192 GB HBM3 - bf16 throughput on CDNA3 ≈ same as fp16, > fp32 (matrix engine is fp16/bf16/fp8 native) - fp32 is the SLOWEST option on this arch — never suggest it as an improvement Known incompatibilities for THIS workload (peft + LoRA on transformers Trainer): {incompatibilities} KB rules (one-liner per rule, for grounding): {kb_summary} Current accepted workload state — the literal values in the script after every change accepted so far. Each candidate you propose should mutate one of these (or set an env var). DO NOT propose a value that's already present here. {tunables} Latest benchmark (this is the result of the most recent ACCEPTED state): - tokens_per_sec: {tps:.1f} - mfu_pct: {mfu:.2f} (% of MI300X dense bf16 peak; healthy LoRA ranges 30-50%) - gpu_util_pct: {util:.1f} - hbm_peak_gb: {hbm:.2f} - waste_budget (seconds/step): {waste_lines} Sorted recoverable waste (largest first — go after these): {recoverable_sorted} Previously rejected (full fingerprint — DO NOT repropose any of these): {rejected_fingerprints} History of changes already tried this run (newest first; outcomes are "accepted" / "rejected" / "crashed" / "skipped"): {history_lines} Suggest {num_candidates} STRUCTURALLY-DIFFERENT candidate changes. Each must target a different waste bucket or parameter category. JSON array only. """ async def _ask_llm_for_experiments( backend, *, kb_summary: str, source: str, metrics: dict, history: list[dict], num_candidates: int, ) -> list[Experiment]: """One LLM turn → up to `num_candidates` Experiments. Returns an empty list on parse failure or STOP signal. """ waste = metrics.get("waste_budget") or {} prompt = _LLM_EXPLORE_USER_TEMPLATE.format( num_candidates=num_candidates, incompatibilities="\n".join(f"- {line}" for line in _KNOWN_INCOMPATIBILITIES), kb_summary=kb_summary, tunables=_tunables_summary(source), tps=metrics.get("tokens_per_sec", 0.0), mfu=metrics.get("mfu_pct", 0.0), util=metrics.get("gpu_util_pct", 0.0), hbm=metrics.get("hbm_peak_gb", 0.0), waste_lines=_format_waste(waste), recoverable_sorted=_recoverable_sorted(waste), rejected_fingerprints=_format_rejected_fingerprints(history), history_lines=_format_history(history), ) backend.add_user_message(prompt) turn = await backend.next_turn(tool_schemas=[]) raw = " ".join(turn.text_blocks).strip() arr = _extract_json_array(raw) if not arr: print(f" LLM response was not parseable JSON array. Raw: {raw[:300]!r}") return [] experiments: list[Experiment] = [] for obj in arr: if not isinstance(obj, dict): continue name = (obj.get("name") or "").strip() if not name: continue if name.upper() == "STOP": print(f" LLM signaled STOP: {obj.get('rationale', '(no rationale)')}") return [] subs_raw = obj.get("substitutions") or [] envs_raw = obj.get("env_vars") or {} if not subs_raw and not envs_raw: continue subs = [] for entry in subs_raw: if isinstance(entry, list) and len(entry) == 2: subs.append((str(entry[0]), str(entry[1]))) elif isinstance(entry, dict) and "pattern" in entry and "replacement" in entry: subs.append((str(entry["pattern"]), str(entry["replacement"]))) cleaned_envs = _sanitize_env_vars(envs_raw, context=name) if not subs and not cleaned_envs: print(f" candidate {name!r} had nothing valid after sanitization; dropping") continue experiments.append( Experiment( name=name, description=obj.get("description") or name, rationale=str(obj.get("rationale") or ""), substitutions=subs, env_vars=cleaned_envs, ) ) return experiments def _extract_json_array(text: str) -> list | None: """Pull the first JSON array out of an LLM response, tolerating markdown fences and leading prose. Returns None if nothing parseable.""" if not text: return None fence_match = re.search(r"```(?:json)?\s*(\[.*?\])\s*```", text, re.DOTALL) if fence_match: try: obj = json.loads(fence_match.group(1)) if isinstance(obj, list): return obj except json.JSONDecodeError: pass depth = 0 start = -1 for i, ch in enumerate(text): if ch == "[": if depth == 0: start = i depth += 1 elif ch == "]": depth -= 1 if depth == 0 and start >= 0: blob = text[start : i + 1] try: obj = json.loads(blob) if isinstance(obj, list): return obj except json.JSONDecodeError: start = -1 continue return None # --------------------------------------------------------------------------- # Dedup + history utilities (used by all LLM modes) # --------------------------------------------------------------------------- def _experiment_fingerprint(exp: Experiment) -> tuple: """Hashable identity for an experiment — substitutions + env_vars, NOT name (the LLM tends to give the same change different names).""" subs = tuple(sorted(tuple(s) for s in exp.substitutions)) envs = tuple(sorted(exp.env_vars.items())) return (subs, envs) def _build_merged_experiment( exps: list[Experiment], base_source: str ) -> tuple[Experiment | None, str]: """Try to combine 2+ experiments into one. The merged experiment applies all of their substitutions in sequence and unions their env_vars. Returns (merged, "") on success, (None, reason) when the merge is structurally unsafe — caller should fall back to using just the individual winner. Conflict detection: - A later substitution's pattern must still match after earlier substitutions have been applied (zero matches → conflict, e.g. cand A rewrote `fp16=True` and cand B was also targeting it). - Env var keys with conflicting values (same name, different value) → conflict. - Bad regex anywhere → conflict. """ if len(exps) < 2: return None, "need at least 2 experiments" merged_subs: list[tuple[str, str]] = [] merged_envs: dict[str, str] = {} test_source = base_source for exp in exps: for pattern, replacement in exp.substitutions: try: new_source, n = re.subn(pattern, replacement, test_source) except re.error as e: return None, f"bad regex in '{exp.name}': {e}" if n == 0: return None, ( f"'{exp.name}' substitution {pattern!r} no longer matches " "after prior merges (likely overwrites an earlier change)" ) test_source = new_source merged_subs.append((pattern, replacement)) for k, v in exp.env_vars.items(): if k in merged_envs and merged_envs[k] != v: return None, ( f"env var conflict on {k!r}: {merged_envs[k]!r} vs {v!r}" ) merged_envs[k] = v short_names = "+".join(e.name[:14] for e in exps) full_names = " + ".join(e.name for e in exps) return ( Experiment( name=f"merge[{short_names}]"[:60], description=f"Merged: {full_names}", rationale=( f"Combined {len(exps)} candidates that each had positive delta " "against the current best this iteration. Tests the compound " "effect; falls back to the individual winner if it doesn't help." ), substitutions=merged_subs, env_vars=merged_envs, ), "", ) def _is_duplicate_of_history(exp: Experiment, history: list[dict]) -> dict | None: """If `exp` matches a prior history entry by fingerprint, return that entry. Otherwise None.""" fp = _experiment_fingerprint(exp) for h in history: h_subs = tuple( sorted( (str(s[0]), str(s[1])) for s in (h.get("substitutions") or []) if isinstance(s, (list, tuple)) and len(s) == 2 ) ) h_envs = tuple(sorted((h.get("env_vars") or {}).items())) if fp == (h_subs, h_envs): return h return None def _format_rejected_fingerprints(history: list[dict]) -> str: """Compact list of every (substitutions, env_vars) the LLM has already tried with outcome rejected/crashed/skipped — so it can't propose them again under a different name.""" seen: set[tuple] = set() lines: list[str] = [] for h in history: outcome = h.get("outcome", "") if outcome not in ("rejected", "crashed", "skipped"): continue subs = tuple( sorted( (str(s[0]), str(s[1])) for s in (h.get("substitutions") or []) if isinstance(s, (list, tuple)) and len(s) == 2 ) ) envs = tuple(sorted((h.get("env_vars") or {}).items())) fp = (subs, envs) if fp in seen: continue seen.add(fp) lines.append(f" - {outcome:9s} subs={list(subs)} env={dict(envs)}") if not lines: return " (none yet)" return "\n".join(lines) def _print_waste(metrics: dict, prefix: str = " waste: ") -> None: """Print a one-line summary of waste_budget — useful is highlighted first, then non-zero recoverable buckets sorted by size.""" wb = metrics.get("waste_budget") or {} if not wb: return parts = [f"useful_gpu={wb.get('useful_gpu', 0.0):.3f}"] others = [(k, v) for k, v in wb.items() if k != "useful_gpu" and isinstance(v, (int, float)) and v > 0] others.sort(key=lambda kv: kv[1], reverse=True) parts.extend(f"{k}={v:.3f}" for k, v in others) print(prefix + ", ".join(parts)) # --------------------------------------------------------------------------- # JSON object extractor (used by single-experiment llm mode) # --------------------------------------------------------------------------- def _extract_json_object(text: str) -> dict | None: """Pull the first JSON object out of an LLM response, tolerating markdown fences / leading prose.""" if not text: return None # strip ```json ... ``` fences if present fence_match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL) if fence_match: try: return json.loads(fence_match.group(1)) except json.JSONDecodeError: pass # otherwise grab the first balanced { ... } depth = 0 start = -1 for i, ch in enumerate(text): if ch == "{": if depth == 0: start = i depth += 1 elif ch == "}": depth -= 1 if depth == 0 and start >= 0: blob = text[start : i + 1] try: return json.loads(blob) except json.JSONDecodeError: start = -1 continue return None # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> int: p = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) p.add_argument( "workload", type=Path, nargs="?", default=None, help=( "Path to a workload script (omit if using --model). When given, " "the script is used as-is for the baseline benchmark." ), ) p.add_argument( "--model", type=str, default=None, help=( "HuggingFace model id (e.g. Qwen/Qwen2.5-7B-Instruct, " "meta-llama/Llama-3.2-3B). Generates a baseline workload from " "workloads/train_qwen_lora.py with this MODEL_ID substituted in. " "Use this OR a workload path, not both. For gated models, " "ensure HF_TOKEN is set in your shell." ), ) p.add_argument( "--mode", choices=("hardcoded", "llm", "llm-explore"), default="hardcoded", help=( "hardcoded (default): walk through the priority-ordered EXPERIMENTS list. " "llm: ask the LLM for ONE next experiment per iteration (greedy). " "llm-explore: ask for K candidates per iteration, benchmark all, keep " "the best (slower but better at finding interaction effects)." ), ) p.add_argument( "--candidates-per-iteration", type=int, default=3, help="Only used when --mode llm-explore. Default 3.", ) p.add_argument("--steps", type=int, default=20, help="Steps per benchmark") p.add_argument( "--max-iterations", type=int, default=0, help=( "Cap on experiments to try. Default: len(EXPERIMENTS) for hardcoded mode, " "10 for llm mode." ), ) p.add_argument( "--early-stop-after", type=int, default=3, help=( "Stop after N consecutive non-improvements. Crashes do NOT count " "toward this — crashes mean the change was structurally bad, not " "that we've exhausted ideas." ), ) p.add_argument( "--max-crashes", type=int, default=4, help=( "Stop after N total subprocess crashes (separate from " "--early-stop-after). Default 4 leaves room for the LLM to try " "structurally different changes after a bad one." ), ) p.add_argument( "--improvement-threshold", type=float, default=0.0, help=( "Min %% improvement over current best to accept. Default 0.0 " "(any positive delta wins). Bump to 1.0 if your benchmarks are " "noisy and you want to ignore sub-1%% deltas." ), ) p.add_argument( "--events", type=Path, default=None, help=( "Optional NDJSON event stream output. If set, the script appends " "one JSON event per line at baseline / iter / candidate / summary " "milestones. Used by the Streamlit UI; CLI users don't need this." ), ) args = p.parse_args() if args.events is not None: global _EVENTS_PATH _EVENTS_PATH = args.events try: args.events.write_text("") # truncate any prior contents except OSError as exc: sys.stderr.write(f"--events: cannot open {args.events} for writing ({exc})\n") return 1 if args.max_iterations <= 0: if args.mode == "hardcoded": args.max_iterations = len(EXPERIMENTS) elif args.mode == "llm-explore": args.max_iterations = 5 # K candidates per iter so 5 iters = 5K benchmarks else: args.max_iterations = 10 # Validate that exactly one workload source was provided if args.workload is None and args.model is None: sys.stderr.write( "Pass either a workload path or --model MODEL_ID. " "Examples:\n" " python scripts/auto_tune.py workloads/train_qwen_lora.py\n" " python scripts/auto_tune.py --model Qwen/Qwen2.5-7B-Instruct\n" ) return 1 if args.workload is not None and args.model is not None: sys.stderr.write( "Pass EITHER a workload path OR --model, not both.\n" ) return 1 if not GOBLIN_RUNNER.exists(): sys.stderr.write(f"goblin_runner.sh not found at {GOBLIN_RUNNER}\n") return 1 workspace = Path(tempfile.mkdtemp(prefix="auto_tune_workloads_")) if args.workload is not None: workload = args.workload.resolve() if not workload.exists(): sys.stderr.write(f"workload not found: {workload}\n") return 1 workload_label = str(workload) else: # Generate baseline workload from --model generated = workspace / "_generated_baseline.py" workload = _generate_workload_from_model(args.model, generated) workload_label = f"(generated from --model {args.model})\n " workload_label += f" {workload}\n " workload_label += f" template: {_DEFAULT_WORKLOAD_TEMPLATE}" _emit({ "type": "started", "mode": args.mode, "workload": str(workload), "model": args.model, "steps": args.steps, "max_iterations": args.max_iterations, "early_stop_after": args.early_stop_after, "max_crashes": args.max_crashes, "improvement_threshold": args.improvement_threshold, "candidates_per_iteration": ( args.candidates_per_iteration if args.mode == "llm-explore" else 1 ), "workspace": str(workspace), }) print(f"Auto-tune workspace: {workspace}") print(f"Mode: {args.mode}") print(f"Workload: {workload_label}") print(f"Steps per benchmark: {args.steps}") print(f"Max iterations: {args.max_iterations}") print(f"Early stop after: {args.early_stop_after} non-improvements") print(f"Max crashes: {args.max_crashes} total") print(f"Accept threshold: {args.improvement_threshold:.1f}%\n") # LLM mode setup happens before the baseline so we fail fast on missing # credentials rather than after burning a baseline benchmark. Each LLM # mode gets its own system prompt — the explore mode needs a much # larger token budget to emit K JSON objects. llm_backend = None kb_summary = "" if args.mode == "llm": llm_backend = _build_llm_backend(_LLM_SYSTEM_PROMPT, max_tokens=1024) kb_summary = _kb_summary(REPO_ROOT / "kb" / "rocm_rules.yaml") print("LLM backend ready (single-candidate). KB summary loaded.\n") elif args.mode == "llm-explore": llm_backend = _build_llm_backend(_LLM_EXPLORE_SYSTEM_PROMPT, max_tokens=2048) kb_summary = _kb_summary(REPO_ROOT / "kb" / "rocm_rules.yaml") print( f"LLM backend ready (multi-candidate, K={args.candidates_per_iteration}). " "KB summary loaded.\n" ) baseline_source = workload.read_text() baseline_path = workspace / "00_baseline.py" baseline_path.write_text(baseline_source) print("=" * 60) print("Baseline benchmark") print("=" * 60) baseline = benchmark(baseline_path, args.steps, {}) if baseline is None: sys.stderr.write("Baseline benchmark failed; cannot continue.\n") return 1 baseline_tps = baseline["tokens_per_sec"] print(f" tokens/sec: {baseline_tps:.1f}") print(f" mfu_pct: {baseline.get('mfu_pct', 0.0):.2f}") print(f" hbm_peak_gb: {baseline['hbm_peak_gb']:.2f}") print(f" gpu_util_pct: {baseline['gpu_util_pct']:.1f}") print( " waste_budget: " + ", ".join(f"{k}={v:.3f}" for k, v in baseline["waste_budget"].items() if v > 0) ) _emit({"type": "baseline", "metrics": baseline}) best_source = baseline_source best_tps = baseline_tps best_env: dict[str, str] = {} last_metrics = baseline accepted: list[tuple[str, float, float]] = [] # (name, tps, delta_pct) rejected: list[tuple[str, str]] = [] # (name, reason) history: list[dict] = [] # for LLM context consecutive_no_improvement = 0 total_crashes = 0 file_counter = 0 # monotonically increases across all candidates for i in range(args.max_iterations): # ---- Get candidates list (1 for hardcoded/llm, K for llm-explore) ---- if args.mode == "hardcoded": if i >= len(EXPERIMENTS): print("\nReached end of EXPERIMENTS list.") break candidates = [EXPERIMENTS[i]] elif args.mode == "llm": print(f"\n[asking LLM for next experiment, iteration {i + 1}...]") try: exp = asyncio.run( _ask_llm_for_experiment( llm_backend, kb_summary=kb_summary, source=best_source, metrics=last_metrics, history=history, ) ) except Exception as exc: print(f" LLM call failed: {type(exc).__name__}: {exc}") exp = None if exp is None: print("LLM produced no experiment — stopping.") break candidates = [exp] else: # llm-explore K = args.candidates_per_iteration print(f"\n[asking LLM for {K} candidates, iteration {i + 1}...]") try: candidates = asyncio.run( _ask_llm_for_experiments( llm_backend, kb_summary=kb_summary, source=best_source, metrics=last_metrics, history=history, num_candidates=K, ) ) except Exception as exc: print(f" LLM call failed: {type(exc).__name__}: {exc}") candidates = [] if not candidates: print("LLM produced no candidates — stopping.") break print(f" LLM proposed {len(candidates)} candidate(s): " + ", ".join(c.name for c in candidates)) print() print("=" * 60) n_label = f" ({len(candidates)} candidates)" if len(candidates) > 1 else "" print(f"Iteration {i + 1}{n_label}") print("=" * 60) _emit({ "type": "iter_start", "iteration": i + 1, "candidates": [ { "name": c.name, "rationale": c.rationale, "substitutions": c.substitutions, "env_vars": c.env_vars, } for c in candidates ], }) # ---- Evaluate each candidate against the CURRENT best ---- # Crucial for llm-explore: every candidate is benchmarked against # the same best_source / best_env baseline, so the comparison is # apples-to-apples. State updates only happen after the iteration's # winner is chosen. eval_results: list[dict] = [] # candidates that produced metrics seen_this_iter: set[tuple] = set() # within-batch dedup crashed_this_iter = False max_crashes_hit = False for j, exp in enumerate(candidates): cand_label = f" Candidate {j + 1}/{len(candidates)}" if len(candidates) > 1 else " Candidate" print(f"\n{cand_label}: {exp.name}") print(f" description: {exp.description}") print(f" rationale: {exp.rationale}") # Helper to emit a per-candidate event with the consistent shape # the UI expects. Called at every terminus below. def _cand_event(outcome: str, metrics: dict | None = None, delta_vs_best: float | None = None, reason: str = "") -> None: _emit({ "type": "candidate", "iteration": i + 1, "candidate_index": j + 1, "n_candidates": len(candidates), "name": exp.name, "rationale": exp.rationale, "substitutions": exp.substitutions, "env_vars": exp.env_vars, "outcome": outcome, "metrics": metrics, "delta_vs_best": delta_vs_best, "reason": reason, }) # Dedup: against prior iterations' history dup = _is_duplicate_of_history(exp, history) if dup is not None: print(f" SKIPPED — already tried as '{dup.get('name', '?')}' " f"(outcome '{dup.get('outcome', '?')}')") history.append({ "name": exp.name, "outcome": "skipped", "delta_pct": None, "substitutions": exp.substitutions, "env_vars": exp.env_vars, }) _cand_event("skipped", reason=f"duplicate of '{dup.get('name', '?')}'") continue # Dedup: within the current batch (llm-explore can collide) fp = _experiment_fingerprint(exp) if fp in seen_this_iter: print(" SKIPPED — duplicate of an earlier candidate in this iteration") history.append({ "name": exp.name, "outcome": "skipped", "delta_pct": None, "substitutions": exp.substitutions, "env_vars": exp.env_vars, }) _cand_event("skipped", reason="duplicate of an earlier candidate this iteration") continue seen_this_iter.add(fp) # Apply substitutions if exp.substitutions: try: candidate_source = apply_substitutions(best_source, exp.substitutions) except re.error as exc: print(f" SKIPPED — invalid regex from LLM: {exc}") rejected.append((exp.name, f"bad regex: {exc}")) history.append({ "name": exp.name, "outcome": "rejected", "delta_pct": None, "substitutions": exp.substitutions, "env_vars": exp.env_vars, }) _cand_event("rejected", reason=f"bad regex: {exc}") continue if candidate_source is None: print(" SKIPPED — substitution patterns didn't match") rejected.append((exp.name, "patterns didn't match")) history.append({ "name": exp.name, "outcome": "skipped", "delta_pct": None, "substitutions": exp.substitutions, "env_vars": exp.env_vars, }) _cand_event("skipped", reason="substitution patterns didn't match") continue else: candidate_source = best_source file_counter += 1 safe_name = re.sub(r"[^A-Za-z0-9_]+", "_", exp.name)[:40] or "exp" candidate_path = workspace / f"{file_counter:03d}_iter{i + 1:02d}_{safe_name}.py" candidate_path.write_text(candidate_source) candidate_env = {**best_env, **exp.env_vars} if exp.env_vars: print(f" env vars: {exp.env_vars}") m = benchmark(candidate_path, args.steps, candidate_env) if m is None: rejected.append((exp.name, "benchmark crashed")) history.append({ "name": exp.name, "outcome": "crashed", "delta_pct": None, "substitutions": exp.substitutions, "env_vars": exp.env_vars, }) total_crashes += 1 crashed_this_iter = True print( f" CRASHED — counted toward max-crashes " f"({total_crashes}/{args.max_crashes})" ) _cand_event("crashed", reason="benchmark subprocess failed") if total_crashes >= args.max_crashes: max_crashes_hit = True break continue tps = m["tokens_per_sec"] delta_vs_best = _delta_pct(tps, best_tps) print(f" tokens/sec: {tps:.1f} (Δ {delta_vs_best:+.2f}% vs current best)") print(f" mfu_pct: {m.get('mfu_pct', 0.0):.2f}") print(f" hbm_peak_gb: {m['hbm_peak_gb']:.2f}") print(f" gpu_util_pct:{m['gpu_util_pct']:.1f}") _print_waste(m, prefix=" waste: ") # Emit "evaluated" — outcome (accepted/rejected) is decided # later when the iteration's winner is picked across all # candidates. For UI display purposes the per-candidate metrics # are already useful. _cand_event("evaluated", metrics=m, delta_vs_best=delta_vs_best) eval_results.append({ "exp": exp, "candidate_source": candidate_source, "candidate_env": candidate_env, "metrics": m, "delta_vs_best": delta_vs_best, }) if max_crashes_hit: print( f"\nReached max-crashes ({args.max_crashes}) — stopping to " "avoid burning more GPU on structurally bad changes." ) break # ---- Pick the iteration's winner from eval_results ---- if not eval_results: # Every candidate was skipped or crashed if crashed_this_iter: print("\n All candidates crashed or were skipped this iteration.") else: print("\n All candidates were skipped this iteration.") consecutive_no_improvement += 1 else: winner = max(eval_results, key=lambda r: r["metrics"]["tokens_per_sec"]) winner_delta = winner["delta_vs_best"] # ---- Optional merge step (llm-explore only) ---- # If 2+ candidates this iteration each beat the baseline, try # combining them into one experiment and benchmark the merge. # The merge replaces `winner` only if it strictly exceeds the # individual winner's tokens/sec. if args.mode == "llm-explore": positives = [r for r in eval_results if r["delta_vs_best"] > 0] if len(positives) >= 2: merged_exp, merge_reason = _build_merged_experiment( [r["exp"] for r in positives], best_source ) if merged_exp is None: print(f"\n MERGE SKIPPED — {merge_reason}") _emit({ "type": "merge_attempt", "iteration": i + 1, "outcome": "skipped", "reason": merge_reason, "candidate_names": [r["exp"].name for r in positives], }) else: print( f"\n Merging {len(positives)} positive candidates: " f"{merged_exp.description}" ) # Apply substitutions to get the merged source merged_source = best_source for pattern, replacement in merged_exp.substitutions: merged_source = re.sub(pattern, replacement, merged_source) merged_env = {**best_env, **merged_exp.env_vars} file_counter += 1 merged_path = workspace / f"{file_counter:03d}_iter{i + 1:02d}_merge.py" merged_path.write_text(merged_source) if merged_exp.env_vars: print(f" env vars: {merged_exp.env_vars}") m = benchmark(merged_path, args.steps, merged_env) if m is None: total_crashes += 1 crashed_this_iter = True print( f" MERGE CRASHED — counted toward max-crashes " f"({total_crashes}/{args.max_crashes})" ) _emit({ "type": "merge_attempt", "iteration": i + 1, "outcome": "crashed", "candidate_names": [r["exp"].name for r in positives], "merged_name": merged_exp.name, }) if total_crashes >= args.max_crashes: max_crashes_hit = True else: tps = m["tokens_per_sec"] delta_vs_best = _delta_pct(tps, best_tps) print( f" Merged tokens/sec: {tps:.1f} " f"(Δ {delta_vs_best:+.2f}% vs baseline)" ) print(f" mfu_pct: {m.get('mfu_pct', 0.0):.2f}") print(f" hbm_peak_gb: {m['hbm_peak_gb']:.2f}") individual_best_tps = winner["metrics"]["tokens_per_sec"] if tps > individual_best_tps: print( f" MERGE WINS — exceeds individual winner " f"'{winner['exp'].name}' " f"({tps:.1f} > {individual_best_tps:.1f})" ) _emit({ "type": "merge_attempt", "iteration": i + 1, "outcome": "wins", "candidate_names": [r["exp"].name for r in positives], "merged_name": merged_exp.name, "metrics": m, "delta_vs_best": delta_vs_best, "individual_best_name": winner["exp"].name, "individual_best_tps": individual_best_tps, }) # Promote merged to be the new winner winner = { "exp": merged_exp, "candidate_source": merged_source, "candidate_env": merged_env, "metrics": m, "delta_vs_best": delta_vs_best, } winner_delta = delta_vs_best else: print( f" Merge didn't beat individual winner; " f"keeping '{winner['exp'].name}'" ) _emit({ "type": "merge_attempt", "iteration": i + 1, "outcome": "lost", "candidate_names": [r["exp"].name for r in positives], "merged_name": merged_exp.name, "metrics": m, "delta_vs_best": delta_vs_best, "individual_best_name": winner["exp"].name, "individual_best_tps": individual_best_tps, }) if winner_delta >= args.improvement_threshold: print( f"\n ACCEPTED — '{winner['exp'].name}' wins " f"(Δ {winner_delta:+.2f}% vs current best)" ) best_source = winner["candidate_source"] best_tps = winner["metrics"]["tokens_per_sec"] best_env = winner["candidate_env"] last_metrics = winner["metrics"] accepted.append((winner["exp"].name, best_tps, winner_delta)) history.append({ "name": winner["exp"].name, "outcome": "accepted", "delta_pct": winner_delta, "substitutions": winner["exp"].substitutions, "env_vars": winner["exp"].env_vars, }) # Other candidates of this iteration get marked rejected for r in eval_results: if r is winner: continue rejected.append((r["exp"].name, f"{r['delta_vs_best']:+.2f}%")) history.append({ "name": r["exp"].name, "outcome": "rejected", "delta_pct": r["delta_vs_best"], "substitutions": r["exp"].substitutions, "env_vars": r["exp"].env_vars, }) consecutive_no_improvement = 0 _emit({ "type": "iter_done", "iteration": i + 1, "outcome": "accepted", "winner_name": winner["exp"].name, "winner_delta": winner_delta, "best_tps": best_tps, "best_metrics": winner["metrics"], "best_env_vars": best_env, }) else: print( f"\n ALL REJECTED — best candidate '{winner['exp'].name}' " f"only Δ {winner_delta:+.2f}% (threshold {args.improvement_threshold:.1f}%)" ) for r in eval_results: rejected.append((r["exp"].name, f"{r['delta_vs_best']:+.2f}%")) history.append({ "name": r["exp"].name, "outcome": "rejected", "delta_pct": r["delta_vs_best"], "substitutions": r["exp"].substitutions, "env_vars": r["exp"].env_vars, }) # Update last_metrics with the winner anyway so the LLM sees # the latest waste_budget on the next turn. if args.mode in ("llm", "llm-explore"): last_metrics = winner["metrics"] consecutive_no_improvement += 1 _emit({ "type": "iter_done", "iteration": i + 1, "outcome": "all_rejected", "winner_name": winner["exp"].name, "winner_delta": winner_delta, "best_tps": best_tps, }) if consecutive_no_improvement >= args.early_stop_after: print( f"\nNo improvement for {args.early_stop_after} consecutive iterations — early stopping." ) break # Save best best_path = workspace / "best.py" best_path.write_text(best_source) # Summary print() print("=" * 60) print("AUTO-TUNE SUMMARY") print("=" * 60) print(f"Baseline tokens/sec: {baseline_tps:.1f}") print( f"Best tokens/sec: {best_tps:.1f} " f"({_delta_pct(best_tps, baseline_tps):+.2f}% vs baseline)" ) print() print(f"Accepted ({len(accepted)}):") for name, tps, delta in accepted: print(f" + {name:25s} {tps:8.1f} tok/s (Δ {delta:+.2f}%)") print() print(f"Rejected ({len(rejected)}):") for name, reason in rejected: print(f" - {name:25s} {reason}") print() if best_env: print("Required env vars for best config:") for k, v in best_env.items(): print(f" export {k}={v}") print() print(f"Best workload script: {best_path}") print(f"Diff vs baseline: diff {workload} {best_path}") _emit({ "type": "summary", "baseline_metrics": baseline, "best_metrics": last_metrics, "baseline_tps": baseline_tps, "best_tps": best_tps, "improvement_pct": _delta_pct(best_tps, baseline_tps), "accepted": [ {"name": name, "tps": tps, "delta_pct": delta} for name, tps, delta in accepted ], "rejected": [ {"name": name, "reason": reason} for name, reason in rejected ], "best_env_vars": best_env, "best_workload_path": str(best_path), "baseline_workload_path": str(workload), }) return 0 if __name__ == "__main__": raise SystemExit(main())