Deploy auto-tune UI + scripts (work-from-91d0cf0)
Browse filesBranch: work-from-91d0cf0
Adds scripts/auto_tune.py (with hardcoded / llm / llm-explore modes, --model flag, merge-and-test step, --events stream), ui/auto_tune_ui.py (Streamlit frontend), and the /auto-tune FastAPI endpoint for the UI to call against a remote MI300X server. The Space's app_file now points at the auto-tune UI.
- .gitignore +26 -0
- README.md +39 -92
- agent/server.py +183 -10
- brainstorming/architecture.md +359 -0
- brainstorming/goals.md +233 -0
- brainstorming/idea.md +105 -0
- pyproject.toml +43 -0
- runner/goblin_runner.sh +169 -0
- runner/profile_parser.py +693 -0
- runner/protocol.py +7 -51
- scripts/auto_tune.py +1918 -0
- tests/fixtures/sample_train.json +27 -0
- tests/fixtures/sample_train.py +101 -0
- tests/fixtures/sample_train.yaml +26 -0
- tests/test_compare_runs_normalize.py +155 -0
- tests/test_loop.py +322 -0
- tests/test_misnested_args.py +95 -0
- tests/test_parse_config.py +377 -0
- tests/test_query_rocm_kb.py +348 -0
- tests/test_qwen_vllm_backend.py +280 -0
- tests/test_runner.py +311 -0
- ui/auto_tune_ui.py +799 -0
- workloads/_runtime.py +32 -35
- workloads/scenarios/README.md +30 -0
- workloads/scenarios/train_qwen_bnb.py +110 -0
- workloads/scenarios/train_qwen_distributed_bad.py +122 -0
- workloads/scenarios/train_qwen_full_ft.py +103 -0
- workloads/scenarios/train_qwen_long_context.py +111 -0
- workloads/scenarios/train_qwen_well_tuned.py +143 -0
- workloads/train_qwen_lora.py +8 -6
.gitignore
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
.claude/
|
| 2 |
+
|
| 3 |
+
# Python
|
| 4 |
+
__pycache__/
|
| 5 |
+
*.py[cod]
|
| 6 |
+
*$py.class
|
| 7 |
+
*.egg-info/
|
| 8 |
+
.venv/
|
| 9 |
+
venv/
|
| 10 |
+
.pytest_cache/
|
| 11 |
+
.ruff_cache/
|
| 12 |
+
|
| 13 |
+
# Goblin-specific
|
| 14 |
+
bench_cache/
|
| 15 |
+
# kb/.embeddings_cache_*.npy is intentionally TRACKED — see README "Deploying
|
| 16 |
+
# to Hugging Face Spaces". Shipping the cache keyed on the current YAML's
|
| 17 |
+
# sha256 keeps the Space cold-start fast (no sentence-transformers download).
|
| 18 |
+
*.trace.csv
|
| 19 |
+
*.trace.json
|
| 20 |
+
.env
|
| 21 |
+
.anthropic-key
|
| 22 |
+
|
| 23 |
+
# Editor / OS
|
| 24 |
+
.vscode/
|
| 25 |
+
.idea/
|
| 26 |
+
.DS_Store
|
README.md
CHANGED
|
@@ -3,18 +3,17 @@ title: GPU Goblin
|
|
| 3 |
emoji: 🧌
|
| 4 |
colorFrom: red
|
| 5 |
colorTo: red
|
| 6 |
-
sdk:
|
| 7 |
-
|
|
|
|
| 8 |
pinned: false
|
| 9 |
license: mit
|
| 10 |
-
short_description:
|
| 11 |
tags:
|
| 12 |
- amd
|
| 13 |
-
- amd-hackathon-2026
|
| 14 |
- mi300x
|
| 15 |
- rocm
|
| 16 |
- qwen
|
| 17 |
-
- vllm
|
| 18 |
- huggingface
|
| 19 |
- agent
|
| 20 |
- fine-tuning
|
|
@@ -382,53 +381,36 @@ Application Platform + Application URL" submission fields.
|
|
| 382 |
### One-time setup
|
| 383 |
|
| 384 |
1. Create a Hugging Face account at [huggingface.co](https://huggingface.co/)
|
| 385 |
-
and
|
| 386 |
-
([
|
| 387 |
-
|
| 388 |
-
for the lablab submission to validate.
|
| 389 |
2. Create a token at [Settings → Access Tokens](https://huggingface.co/settings/tokens)
|
| 390 |
-
with **`write`** scope
|
| 391 |
-
``
|
| 392 |
-
|
| 393 |
-
|
| 394 |
-
|
| 395 |
-
-
|
| 396 |
-
- **
|
| 397 |
-
- **
|
| 398 |
-
-
|
| 399 |
-
- **Hardware:** **CPU basic** (free; the Space's offline-replay default
|
| 400 |
-
loads no GPU code path)
|
| 401 |
-
- **Visibility:** **Public** (required for the hackathon prize)
|
| 402 |
4. Don't initialize the Space with anything — leave it empty so the first
|
| 403 |
push lands cleanly.
|
| 404 |
|
| 405 |
-
### Deploy
|
| 406 |
-
|
| 407 |
-
```bash
|
| 408 |
-
export HF_TOKEN=hf_...
|
| 409 |
-
python scripts/deploy_to_hf_space.py --space-name gpu-goblin
|
| 410 |
-
```
|
| 411 |
-
|
| 412 |
-
The script (`scripts/deploy_to_hf_space.py`) uses `HfApi.upload_file` to
|
| 413 |
-
push exactly the files the Space needs (README, requirements.txt, the
|
| 414 |
-
`agent/` package, `kb/`, `ui/`, the cached audit fixture, etc.) and
|
| 415 |
-
deliberately omits build artifacts like `bench_cache/` and
|
| 416 |
-
`__pycache__/` that would bloat the Space repo.
|
| 417 |
-
|
| 418 |
-
### Deploy — Option B: git push (works but uploads everything tracked)
|
| 419 |
|
| 420 |
-
From the project root
|
|
|
|
| 421 |
|
| 422 |
```bash
|
| 423 |
-
#
|
| 424 |
-
git remote add space https://huggingface.co/spaces/
|
| 425 |
|
| 426 |
-
# HF Spaces use 'main' as the default branch:
|
| 427 |
-
git push space main
|
| 428 |
```
|
| 429 |
|
| 430 |
-
You'll see a build log at
|
| 431 |
-
`https://huggingface.co/spaces/lablab-ai-amd-developer-hackathon/gpu-goblin`.
|
| 432 |
Cold-start takes 30-60 seconds (Streamlit + the pure-pydantic deps); once
|
| 433 |
up, the canonical demo trajectory replays in ~10 seconds when a judge
|
| 434 |
clicks **"Use sample workload"**.
|
|
@@ -451,64 +433,30 @@ When a judge opens the Space URL:
|
|
| 451 |
|
| 452 |
### Updating the Space
|
| 453 |
|
| 454 |
-
After any change to the repo, redeploy
|
| 455 |
|
| 456 |
```bash
|
| 457 |
-
|
| 458 |
-
python scripts/deploy_to_hf_space.py --space-name gpu-goblin
|
| 459 |
-
|
| 460 |
-
# Option B:
|
| 461 |
-
git push space main
|
| 462 |
```
|
| 463 |
|
| 464 |
-
HF rebuilds the Space automatically on
|
| 465 |
|
| 466 |
### (Stretch) Live agent in the Space
|
| 467 |
|
| 468 |
-
The shipped Space
|
| 469 |
-
|
| 470 |
-
|
| 471 |
-
**Option 1 — In-process Qwen via HF Inference Providers (already wired)**
|
| 472 |
-
|
| 473 |
-
Set these as Space **Settings → Variables and secrets**:
|
| 474 |
-
|
| 475 |
-
| Secret | Value |
|
| 476 |
-
|---|---|
|
| 477 |
-
| `HF_TOKEN` | Your HF token with `inference` scope |
|
| 478 |
-
| `GOBLIN_AGENT_BACKEND` | `qwen-hf` (default) |
|
| 479 |
-
| `GOBLIN_QWEN_MODEL` | `Qwen/Qwen2.5-7B-Instruct` (or any HF Qwen model id) |
|
| 480 |
-
|
| 481 |
-
After redeploying, the lane caption flips to `🟢 Live: agent runs Qwen
|
| 482 |
-
in-process via Hugging Face Inference Providers` and judges audit live.
|
| 483 |
-
|
| 484 |
-
**Option 2 — Connect to a self-hosted vLLM on your AMD droplet**
|
| 485 |
-
|
| 486 |
-
If you've followed the AMD Developer Cloud tutorial and have vLLM serving
|
| 487 |
-
Qwen on your MI300X at `http://YOUR_DROPLET_IP:8000/v1`, point the Space
|
| 488 |
-
at it via Space secrets:
|
| 489 |
-
|
| 490 |
-
| Secret | Value |
|
| 491 |
-
|---|---|
|
| 492 |
-
| `GOBLIN_AGENT_BACKEND` | `qwen-vllm` |
|
| 493 |
-
| `GOBLIN_QWEN_VLLM_URL` | `http://YOUR_DROPLET_IP:8000/v1` |
|
| 494 |
-
| `GOBLIN_QWEN_VLLM_MODEL` | The model ID vLLM advertises at `/v1/models` |
|
| 495 |
-
|
| 496 |
-
**Important:** the AMD Developer Cloud droplet blocks port 8000 by default.
|
| 497 |
-
SSH into the droplet and run `ufw allow 8000` (per the [lablab tutorial](https://lablab.ai/ai-tutorials/amd-huggingface-deployment-for-ai-hackathons))
|
| 498 |
-
before the Space can reach the endpoint. Verify from outside:
|
| 499 |
-
```bash
|
| 500 |
-
curl -s http://YOUR_DROPLET_IP:8000/v1/models
|
| 501 |
-
```
|
| 502 |
-
|
| 503 |
-
**Option 3 — Connect to a separate FastAPI backend**
|
| 504 |
|
| 505 |
-
|
| 506 |
-
|
| 507 |
-
`GOBLIN_BACKEND_URL`
|
| 508 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 509 |
|
| 510 |
-
|
| 511 |
-
|
| 512 |
|
| 513 |
## Configuration Reference
|
| 514 |
|
|
@@ -524,4 +472,3 @@ satisfies the submission requirement.
|
|
| 524 |
| `GOBLIN_BACKEND_URL` | `http://localhost:8000/audit` | UI's backend endpoint. |
|
| 525 |
| `ROCM_IMAGE_TAG` | `unknown` | Container tag mixed into the benchmark cache key. |
|
| 526 |
| `GOBLIN_GPU_ID` | `0` | Which `/dev/dri/renderD*` to bind in `goblin_runner.sh`. |
|
| 527 |
-
| `GOBLIN_RUNNER_TIMEOUT_SECONDS` | `1800` | LiveRunner subprocess timeout. Bump if cold-cache model downloads or kernel JIT push past 30 min; LiveRunner falls back to FakeRunner once exceeded. |
|
|
|
|
| 3 |
emoji: 🧌
|
| 4 |
colorFrom: red
|
| 5 |
colorTo: red
|
| 6 |
+
sdk: streamlit
|
| 7 |
+
sdk_version: "1.32.0"
|
| 8 |
+
app_file: ui/auto_tune_ui.py
|
| 9 |
pinned: false
|
| 10 |
license: mit
|
| 11 |
+
short_description: AI auto-tuner for MI300X fine-tuning workloads.
|
| 12 |
tags:
|
| 13 |
- amd
|
|
|
|
| 14 |
- mi300x
|
| 15 |
- rocm
|
| 16 |
- qwen
|
|
|
|
| 17 |
- huggingface
|
| 18 |
- agent
|
| 19 |
- fine-tuning
|
|
|
|
| 381 |
### One-time setup
|
| 382 |
|
| 383 |
1. Create a Hugging Face account at [huggingface.co](https://huggingface.co/)
|
| 384 |
+
and accept the invite to the **AMD Developer Hackathon HF Organization**
|
| 385 |
+
(link is on the [hackathon page](https://lablab.ai/ai-hackathons/amd-developer)
|
| 386 |
+
under the Hugging Face section).
|
|
|
|
| 387 |
2. Create a token at [Settings → Access Tokens](https://huggingface.co/settings/tokens)
|
| 388 |
+
with **`write`** scope (you need write access to push to the Space repo).
|
| 389 |
+
Save it as `HF_PUSH_TOKEN`.
|
| 390 |
+
3. On the HF organization's page, click **"New Space"**:
|
| 391 |
+
- Owner: AMD Developer Hackathon org
|
| 392 |
+
- Space name: `gpu-goblin` (or your preferred slug)
|
| 393 |
+
- License: MIT
|
| 394 |
+
- SDK: **Streamlit**
|
| 395 |
+
- Hardware: **CPU basic** (free; the Space loads no GPU code path)
|
| 396 |
+
- Visibility: Public
|
|
|
|
|
|
|
|
|
|
| 397 |
4. Don't initialize the Space with anything — leave it empty so the first
|
| 398 |
push lands cleanly.
|
| 399 |
|
| 400 |
+
### Deploy
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 401 |
|
| 402 |
+
From the project root, push the existing `feat/scaffold` branch to the
|
| 403 |
+
Space's git remote:
|
| 404 |
|
| 405 |
```bash
|
| 406 |
+
# Add the Space remote (use HTTPS with your username + HF_PUSH_TOKEN as password):
|
| 407 |
+
git remote add space https://huggingface.co/spaces/<org-slug>/gpu-goblin
|
| 408 |
|
| 409 |
+
# Push (HF Spaces use 'main' as the default branch):
|
| 410 |
+
git push space feat/scaffold:main
|
| 411 |
```
|
| 412 |
|
| 413 |
+
You'll see a build log at `https://huggingface.co/spaces/<org-slug>/gpu-goblin`.
|
|
|
|
| 414 |
Cold-start takes 30-60 seconds (Streamlit + the pure-pydantic deps); once
|
| 415 |
up, the canonical demo trajectory replays in ~10 seconds when a judge
|
| 416 |
clicks **"Use sample workload"**.
|
|
|
|
| 433 |
|
| 434 |
### Updating the Space
|
| 435 |
|
| 436 |
+
After any change to the main repo, redeploy:
|
| 437 |
|
| 438 |
```bash
|
| 439 |
+
git push space feat/scaffold:main
|
|
|
|
|
|
|
|
|
|
|
|
|
| 440 |
```
|
| 441 |
|
| 442 |
+
HF rebuilds the Space automatically on push.
|
| 443 |
|
| 444 |
### (Stretch) Live agent in the Space
|
| 445 |
|
| 446 |
+
The shipped Space is read-only — it doesn't reach a real LLM. If you want
|
| 447 |
+
judges to drive the agent live, two paths:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 448 |
|
| 449 |
+
1. **Stand up the FastAPI backend somewhere reachable** (an MI300X on AMD
|
| 450 |
+
Developer Cloud, an HF Inference Endpoint, a small CPU box) and set the
|
| 451 |
+
Space's `GOBLIN_BACKEND_URL` secret to that URL. The Streamlit app will
|
| 452 |
+
stream real SSE from your backend instead of the cached replay.
|
| 453 |
+
2. **Embed the agent loop in-process** (refactor `ui/app.py` to call
|
| 454 |
+
`agent.loop.run_audit` directly via `asyncio.run`). This adds
|
| 455 |
+
`huggingface_hub` to `requirements.txt` and requires `HF_TOKEN` as a
|
| 456 |
+
Space secret. Larger cold-start, fully self-contained.
|
| 457 |
|
| 458 |
+
Both are post-MVP; the offline-replay Space is what satisfies the
|
| 459 |
+
submission requirement.
|
| 460 |
|
| 461 |
## Configuration Reference
|
| 462 |
|
|
|
|
| 472 |
| `GOBLIN_BACKEND_URL` | `http://localhost:8000/audit` | UI's backend endpoint. |
|
| 473 |
| `ROCM_IMAGE_TAG` | `unknown` | Container tag mixed into the benchmark cache key. |
|
| 474 |
| `GOBLIN_GPU_ID` | `0` | Which `/dev/dri/renderD*` to bind in `goblin_runner.sh`. |
|
|
|
agent/server.py
CHANGED
|
@@ -12,14 +12,19 @@ We never crash on missing keys.
|
|
| 12 |
|
| 13 |
from __future__ import annotations
|
| 14 |
|
|
|
|
| 15 |
import json
|
| 16 |
import os
|
|
|
|
|
|
|
| 17 |
import tempfile
|
| 18 |
from collections.abc import AsyncIterator
|
| 19 |
from pathlib import Path
|
|
|
|
| 20 |
|
| 21 |
-
from fastapi import FastAPI, File, UploadFile
|
| 22 |
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
| 23 |
from sse_starlette.sse import EventSourceResponse
|
| 24 |
|
| 25 |
from agent.backends import active_backend_name
|
|
@@ -27,6 +32,9 @@ from agent.loop import run_audit
|
|
| 27 |
from agent.schemas import SSEEvent
|
| 28 |
from agent.tools import ALL_TOOLS
|
| 29 |
|
|
|
|
|
|
|
|
|
|
| 30 |
app = FastAPI(title="GPU Goblin Agent", version="0.1.0")
|
| 31 |
|
| 32 |
app.add_middleware(
|
|
@@ -82,20 +90,16 @@ async def _stream_audit(file_path: str) -> AsyncIterator[dict]:
|
|
| 82 |
sse-starlette expects. Each yielded dict becomes one `data: ...\\n\\n`
|
| 83 |
SSE message.
|
| 84 |
"""
|
| 85 |
-
|
| 86 |
-
|
| 87 |
-
# tokens by default. Without this guard the server unconditionally errored
|
| 88 |
-
# when HF_TOKEN was unset, even though the active backend didn't need it.
|
| 89 |
-
if active_backend_name() == "qwen-hf" and not _has_hf_token():
|
| 90 |
yield {
|
| 91 |
"data": SSEEvent(
|
| 92 |
type="error",
|
| 93 |
data={
|
| 94 |
"message": (
|
| 95 |
-
"HF_TOKEN not set on the server — Qwen
|
| 96 |
-
"unavailable. Set HF_TOKEN (or HUGGINGFACEHUB_API_TOKEN)
|
| 97 |
-
"
|
| 98 |
-
"offline-replay UI lane."
|
| 99 |
)
|
| 100 |
},
|
| 101 |
).model_dump_json()
|
|
@@ -143,6 +147,175 @@ async def audit(file: UploadFile = File(...)) -> EventSourceResponse:
|
|
| 143 |
return EventSourceResponse(_stream_audit(tmp_path))
|
| 144 |
|
| 145 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 146 |
# Convenience: support `python -m uvicorn agent.server:app --reload`.
|
| 147 |
__all__ = ["app"]
|
| 148 |
|
|
|
|
| 12 |
|
| 13 |
from __future__ import annotations
|
| 14 |
|
| 15 |
+
import asyncio
|
| 16 |
import json
|
| 17 |
import os
|
| 18 |
+
import subprocess
|
| 19 |
+
import sys
|
| 20 |
import tempfile
|
| 21 |
from collections.abc import AsyncIterator
|
| 22 |
from pathlib import Path
|
| 23 |
+
from typing import Any
|
| 24 |
|
| 25 |
+
from fastapi import FastAPI, File, HTTPException, UploadFile
|
| 26 |
from fastapi.middleware.cors import CORSMiddleware
|
| 27 |
+
from pydantic import BaseModel, Field
|
| 28 |
from sse_starlette.sse import EventSourceResponse
|
| 29 |
|
| 30 |
from agent.backends import active_backend_name
|
|
|
|
| 32 |
from agent.schemas import SSEEvent
|
| 33 |
from agent.tools import ALL_TOOLS
|
| 34 |
|
| 35 |
+
_REPO_ROOT = Path(__file__).resolve().parent.parent
|
| 36 |
+
_AUTO_TUNE_SCRIPT = _REPO_ROOT / "scripts" / "auto_tune.py"
|
| 37 |
+
|
| 38 |
app = FastAPI(title="GPU Goblin Agent", version="0.1.0")
|
| 39 |
|
| 40 |
app.add_middleware(
|
|
|
|
| 90 |
sse-starlette expects. Each yielded dict becomes one `data: ...\\n\\n`
|
| 91 |
SSE message.
|
| 92 |
"""
|
| 93 |
+
if not _has_hf_token():
|
| 94 |
+
# Surface a clean error instead of letting the loop crash on missing key.
|
|
|
|
|
|
|
|
|
|
| 95 |
yield {
|
| 96 |
"data": SSEEvent(
|
| 97 |
type="error",
|
| 98 |
data={
|
| 99 |
"message": (
|
| 100 |
+
"HF_TOKEN not set on the server — Qwen agent loop is "
|
| 101 |
+
"unavailable. Set HF_TOKEN (or HUGGINGFACEHUB_API_TOKEN) "
|
| 102 |
+
"or use the offline-replay UI lane."
|
|
|
|
| 103 |
)
|
| 104 |
},
|
| 105 |
).model_dump_json()
|
|
|
|
| 147 |
return EventSourceResponse(_stream_audit(tmp_path))
|
| 148 |
|
| 149 |
|
| 150 |
+
# ---------------------------------------------------------------------------
|
| 151 |
+
# Auto-tune endpoint — lets a UI on a CPU-only host (e.g. an HF Space) drive
|
| 152 |
+
# scripts/auto_tune.py running on a remote MI300X server. The endpoint
|
| 153 |
+
# spawns the CLI, tails its --events NDJSON stream, and re-emits each line
|
| 154 |
+
# as an SSE message. Subprocess output is discarded; everything the UI
|
| 155 |
+
# needs is in the structured events.
|
| 156 |
+
# ---------------------------------------------------------------------------
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
class AutoTuneRequest(BaseModel):
|
| 160 |
+
"""JSON shape the /auto-tune endpoint accepts. Mirrors the auto_tune.py
|
| 161 |
+
CLI surface so the UI just sends what the user picked in the form."""
|
| 162 |
+
|
| 163 |
+
model: str | None = Field(
|
| 164 |
+
default=None,
|
| 165 |
+
description="HuggingFace model id (e.g. Qwen/Qwen2.5-7B-Instruct). "
|
| 166 |
+
"Mutually exclusive with `workload`.",
|
| 167 |
+
)
|
| 168 |
+
workload: str | None = Field(
|
| 169 |
+
default=None,
|
| 170 |
+
description="Path to a workload script ON THE SERVER's filesystem. "
|
| 171 |
+
"Mutually exclusive with `model`.",
|
| 172 |
+
)
|
| 173 |
+
mode: str = Field(default="hardcoded", pattern="^(hardcoded|llm|llm-explore)$")
|
| 174 |
+
candidates_per_iteration: int = Field(default=3, ge=2, le=10)
|
| 175 |
+
steps: int = Field(default=20, ge=1, le=500)
|
| 176 |
+
max_iterations: int = Field(default=10, ge=1, le=50)
|
| 177 |
+
early_stop_after: int = Field(default=3, ge=1, le=20)
|
| 178 |
+
max_crashes: int = Field(default=4, ge=1, le=20)
|
| 179 |
+
improvement_threshold: float = Field(default=0.0, ge=0.0, le=20.0)
|
| 180 |
+
|
| 181 |
+
|
| 182 |
+
def _build_auto_tune_cmd(req: AutoTuneRequest, events_file: Path) -> list[str]:
|
| 183 |
+
cmd: list[str] = [sys.executable, "-u", str(_AUTO_TUNE_SCRIPT)]
|
| 184 |
+
if req.model:
|
| 185 |
+
cmd.extend(["--model", req.model])
|
| 186 |
+
elif req.workload:
|
| 187 |
+
cmd.append(req.workload)
|
| 188 |
+
cmd.extend([
|
| 189 |
+
"--mode", req.mode,
|
| 190 |
+
"--steps", str(req.steps),
|
| 191 |
+
"--max-iterations", str(req.max_iterations),
|
| 192 |
+
"--early-stop-after", str(req.early_stop_after),
|
| 193 |
+
"--max-crashes", str(req.max_crashes),
|
| 194 |
+
"--improvement-threshold", str(req.improvement_threshold),
|
| 195 |
+
"--events", str(events_file),
|
| 196 |
+
])
|
| 197 |
+
if req.mode == "llm-explore":
|
| 198 |
+
cmd.extend(["--candidates-per-iteration", str(req.candidates_per_iteration)])
|
| 199 |
+
return cmd
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
async def _stream_auto_tune(req: AutoTuneRequest) -> AsyncIterator[dict]:
|
| 203 |
+
"""Spawn auto_tune.py and forward its NDJSON --events stream as SSE.
|
| 204 |
+
|
| 205 |
+
Each event is forwarded verbatim — the UI gets the same structured
|
| 206 |
+
payload it would see when running auto_tune.py locally. We discard
|
| 207 |
+
the subprocess's stdout/stderr; any errors are surfaced via the
|
| 208 |
+
`summary` event's absence at process exit.
|
| 209 |
+
"""
|
| 210 |
+
events_file = Path(tempfile.mktemp(prefix="auto_tune_events_", suffix=".ndjson"))
|
| 211 |
+
events_file.write_text("")
|
| 212 |
+
|
| 213 |
+
cmd = _build_auto_tune_cmd(req, events_file)
|
| 214 |
+
|
| 215 |
+
# Validate at least one of model/workload was provided. (Pydantic
|
| 216 |
+
# can't express "exactly one of A or B" cleanly, so we check here.)
|
| 217 |
+
if not req.model and not req.workload:
|
| 218 |
+
yield {"data": json.dumps({
|
| 219 |
+
"type": "error",
|
| 220 |
+
"message": "Pass either `model` or `workload`, not neither."
|
| 221 |
+
})}
|
| 222 |
+
return
|
| 223 |
+
if req.model and req.workload:
|
| 224 |
+
yield {"data": json.dumps({
|
| 225 |
+
"type": "error",
|
| 226 |
+
"message": "Pass either `model` or `workload`, not both."
|
| 227 |
+
})}
|
| 228 |
+
return
|
| 229 |
+
|
| 230 |
+
proc = subprocess.Popen(
|
| 231 |
+
cmd,
|
| 232 |
+
cwd=str(_REPO_ROOT),
|
| 233 |
+
stdout=subprocess.DEVNULL,
|
| 234 |
+
stderr=subprocess.DEVNULL,
|
| 235 |
+
env={**os.environ},
|
| 236 |
+
)
|
| 237 |
+
|
| 238 |
+
seen_bytes = 0
|
| 239 |
+
try:
|
| 240 |
+
while True:
|
| 241 |
+
# Poll the events file for new lines
|
| 242 |
+
try:
|
| 243 |
+
with events_file.open("r") as f:
|
| 244 |
+
f.seek(seen_bytes)
|
| 245 |
+
chunk = f.read()
|
| 246 |
+
new_seen = f.tell()
|
| 247 |
+
except OSError:
|
| 248 |
+
chunk = ""
|
| 249 |
+
new_seen = seen_bytes
|
| 250 |
+
|
| 251 |
+
if chunk:
|
| 252 |
+
# Drop a trailing partial line — re-read it next tick once
|
| 253 |
+
# the writer has flushed the rest.
|
| 254 |
+
lines = chunk.splitlines(keepends=True)
|
| 255 |
+
if lines and not lines[-1].endswith("\n"):
|
| 256 |
+
partial = lines.pop()
|
| 257 |
+
new_seen -= len(partial.encode("utf-8"))
|
| 258 |
+
for line in lines:
|
| 259 |
+
line = line.strip()
|
| 260 |
+
if line:
|
| 261 |
+
yield {"data": line}
|
| 262 |
+
seen_bytes = new_seen
|
| 263 |
+
|
| 264 |
+
if proc.poll() is not None:
|
| 265 |
+
# Subprocess exited. Drain whatever's left on disk.
|
| 266 |
+
try:
|
| 267 |
+
with events_file.open("r") as f:
|
| 268 |
+
f.seek(seen_bytes)
|
| 269 |
+
tail = f.read()
|
| 270 |
+
except OSError:
|
| 271 |
+
tail = ""
|
| 272 |
+
for line in tail.splitlines():
|
| 273 |
+
line = line.strip()
|
| 274 |
+
if line:
|
| 275 |
+
yield {"data": line}
|
| 276 |
+
if proc.returncode != 0:
|
| 277 |
+
yield {"data": json.dumps({
|
| 278 |
+
"type": "process_exit",
|
| 279 |
+
"returncode": proc.returncode,
|
| 280 |
+
"message": (
|
| 281 |
+
f"auto_tune.py exited with code {proc.returncode}. "
|
| 282 |
+
"Check the server's stderr or check `last_runner_failure_*` "
|
| 283 |
+
"in `bench_cache/` for goblin_runner.sh failure logs."
|
| 284 |
+
),
|
| 285 |
+
})}
|
| 286 |
+
break
|
| 287 |
+
|
| 288 |
+
await asyncio.sleep(0.5)
|
| 289 |
+
finally:
|
| 290 |
+
if proc.poll() is None:
|
| 291 |
+
proc.terminate()
|
| 292 |
+
try:
|
| 293 |
+
proc.wait(timeout=3)
|
| 294 |
+
except subprocess.TimeoutExpired:
|
| 295 |
+
proc.kill()
|
| 296 |
+
try:
|
| 297 |
+
events_file.unlink()
|
| 298 |
+
except OSError:
|
| 299 |
+
pass
|
| 300 |
+
|
| 301 |
+
|
| 302 |
+
@app.post("/auto-tune")
|
| 303 |
+
async def auto_tune_endpoint(req: AutoTuneRequest) -> EventSourceResponse:
|
| 304 |
+
"""Stream auto_tune.py events back to the caller as SSE.
|
| 305 |
+
|
| 306 |
+
Run a UI on any host (HF Spaces, local laptop), point it at this
|
| 307 |
+
endpoint, and the actual GPU work happens on the server hosting the
|
| 308 |
+
FastAPI app. Subprocess output is discarded — only the --events
|
| 309 |
+
NDJSON stream crosses the wire, one structured event per SSE message.
|
| 310 |
+
"""
|
| 311 |
+
if not _AUTO_TUNE_SCRIPT.exists():
|
| 312 |
+
raise HTTPException(
|
| 313 |
+
status_code=500,
|
| 314 |
+
detail=f"auto_tune.py not found at {_AUTO_TUNE_SCRIPT}",
|
| 315 |
+
)
|
| 316 |
+
return EventSourceResponse(_stream_auto_tune(req))
|
| 317 |
+
|
| 318 |
+
|
| 319 |
# Convenience: support `python -m uvicorn agent.server:app --reload`.
|
| 320 |
__all__ = ["app"]
|
| 321 |
|
brainstorming/architecture.md
ADDED
|
@@ -0,0 +1,359 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# GPU Goblin — Architecture
|
| 2 |
+
|
| 3 |
+
## System Topology
|
| 4 |
+
|
| 5 |
+
```
|
| 6 |
+
┌──────────────────────────────────────────────────────────────────┐
|
| 7 |
+
│ Streamlit Chat UI (browser) │
|
| 8 |
+
│ upload script · live tool-call stream · final report │
|
| 9 |
+
└────────────────────────────┬─────────────────────────────────────┘
|
| 10 |
+
│ HTTP + Server-Sent Events
|
| 11 |
+
▼
|
| 12 |
+
┌──────────────────────────────────────────────────────────────────┐
|
| 13 |
+
│ Goblin Agent — FastAPI (port 8000) │
|
| 14 |
+
│ Qwen2.5-7B (HF Inference Providers) · agent loop · session │
|
| 15 |
+
└─┬────────┬─────────┬──────────┬──────────┬──────────┬────────────┘
|
| 16 |
+
│ │ │ │ │ │
|
| 17 |
+
▼ ▼ ▼ ▼ ▼ ▼
|
| 18 |
+
┌─────┐ ┌────────┐ ┌──────┐ ┌────────┐ ┌─────────┐ ┌────────┐
|
| 19 |
+
│parse│ │profile │ │rocm │ │propose │ │benchmark│ │compare │
|
| 20 |
+
│cfg │ │_run │ │_kb │ │_patch │ │ │ │_runs │
|
| 21 |
+
└──┬──┘ └────┬───┘ └──┬───┘ └────┬───┘ └────┬────┘ └────┬───┘
|
| 22 |
+
│ │ │ │ │ │
|
| 23 |
+
▼ ▼ ▼ ▼ ▼ ▼
|
| 24 |
+
┌─────┐ ┌────────────────┐ ┌──────────┐ ┌──────────────────┐
|
| 25 |
+
│AST +│ │torch.profiler +│ │YAML KB + │ │MI300X cloud │
|
| 26 |
+
│regex│ │rocprofv3 + amd- │ │sentence- │ │job runner │
|
| 27 |
+
│ │ │smi │ │transform │ │(subprocess + │
|
| 28 |
+
│ │ │ │ │ers │ │ result cache) │
|
| 29 |
+
└─────┘ └────────────────┘ └──────────┘ └──────────────────┘
|
| 30 |
+
```
|
| 31 |
+
|
| 32 |
+
## Components
|
| 33 |
+
|
| 34 |
+
### 1. Streamlit Chat UI
|
| 35 |
+
|
| 36 |
+
- Single-page app: file upload + chat panel + report panel
|
| 37 |
+
- Streams tool calls as cards: *"calling profile_run…"* with live status
|
| 38 |
+
- Final report: side-by-side metrics, diff viewer, kernel waterfall chart
|
| 39 |
+
- Backup mode: replay a cached "golden run" if MI300X unreachable
|
| 40 |
+
|
| 41 |
+
### 2. Goblin Agent (FastAPI)
|
| 42 |
+
|
| 43 |
+
- Single endpoint: `POST /audit` — takes uploaded script, returns SSE stream of agent steps
|
| 44 |
+
- Session state in-memory only (hackathon, not production)
|
| 45 |
+
- Hard cap: **max 8 tool calls per audit** (prevents loops)
|
| 46 |
+
- LLM: **Qwen2.5-7B-Instruct** via Hugging Face Inference Providers (`huggingface_hub.AsyncInferenceClient.chat_completion`, OpenAI-shape tool calls). Pluggable via `agent/backends/`; a future `LiveQwenBackend` swaps to a self-hosted vLLM-on-MI300X endpoint with no other code changes.
|
| 47 |
+
- System prompt establishes ROCm expert persona + tool-call etiquette + report format
|
| 48 |
+
|
| 49 |
+
### 3. The Six Tools
|
| 50 |
+
|
| 51 |
+
#### `parse_config(file_path: str) -> ConfigDict`
|
| 52 |
+
Inputs: HF `TrainingArguments` (Python or JSON), raw PyTorch training script, YAML config.
|
| 53 |
+
Implementation: AST parsing for `.py` (extract `TrainingArguments(...)` kwargs), `json.load` / `yaml.safe_load` for configs, regex fallback.
|
| 54 |
+
Output schema:
|
| 55 |
+
```python
|
| 56 |
+
{
|
| 57 |
+
"model_name": str,
|
| 58 |
+
"batch_size": int,
|
| 59 |
+
"grad_accum_steps": int,
|
| 60 |
+
"seq_len": int,
|
| 61 |
+
"precision": "fp16" | "bf16" | "fp32",
|
| 62 |
+
"optimizer": str,
|
| 63 |
+
"attention_impl": "sdpa" | "flash" | "eager" | "unknown",
|
| 64 |
+
"gradient_checkpointing": bool,
|
| 65 |
+
"lora_rank": int | None,
|
| 66 |
+
"dataloader_workers": int,
|
| 67 |
+
"lr": float,
|
| 68 |
+
"warmup_steps": int,
|
| 69 |
+
"raw_source": str,
|
| 70 |
+
}
|
| 71 |
+
```
|
| 72 |
+
|
| 73 |
+
#### `profile_run(config: ConfigDict, steps: int = 10) -> RunMetrics`
|
| 74 |
+
Wraps the user's training command in `torch.profiler` + `rocprofv3`. Runs for N steps after a 2-step warmup. Captures and **summarizes** (never returns raw traces — too big for LLM context).
|
| 75 |
+
Output is a `RunMetrics` (the same type returned by `benchmark` — see schemas below):
|
| 76 |
+
```python
|
| 77 |
+
RunMetrics(
|
| 78 |
+
steps=10,
|
| 79 |
+
tokens_per_sec=...,
|
| 80 |
+
mfu_pct=...,
|
| 81 |
+
hbm_peak_gb=..., hbm_avg_gb=...,
|
| 82 |
+
gpu_util_pct=...,
|
| 83 |
+
top_kernels=[KernelEntry(name=..., pct_time=...), ...], # top 5
|
| 84 |
+
attention_kernel_loaded=...,
|
| 85 |
+
waste_budget=WasteBudget(...), # see below
|
| 86 |
+
warnings=[...],
|
| 87 |
+
)
|
| 88 |
+
```
|
| 89 |
+
|
| 90 |
+
##### Waste Budget Decomposition
|
| 91 |
+
|
| 92 |
+
Inside `RunMetrics.waste_budget`, total step time is decomposed into interpretable buckets:
|
| 93 |
+
|
| 94 |
+
```
|
| 95 |
+
T_total = T_useful_gpu
|
| 96 |
+
+ T_data_wait # dataloader / host→device stalls
|
| 97 |
+
+ T_host_gap # CPU→GPU launch latency, eager-mode kernel gaps
|
| 98 |
+
+ T_comm_excess # collectives, all-reduce, RCCL overhead
|
| 99 |
+
+ T_memory_headroom # HBM left idle relative to model needs
|
| 100 |
+
+ T_precision_path # throughput lost vs ideal precision (fp32 where bf16 OK)
|
| 101 |
+
+ T_kernel_shape # GEMM tile mismatch, untuned hipBLASLt/MIOpen
|
| 102 |
+
```
|
| 103 |
+
|
| 104 |
+
This decomposition is the basis for the "where time was lost" chart in the final report and for ranking which rules to apply first. Each rule in the KB declares which bucket it targets (`targets_bucket: data_wait`), so `propose_patch` can avoid stacking redundant fixes for the same bucket.
|
| 105 |
+
|
| 106 |
+
#### `query_rocm_kb(symptom: str, top_k: int = 5) -> List[Rule]`
|
| 107 |
+
Semantic search over the YAML KB. Embeds the symptom string with `sentence-transformers/all-MiniLM-L6-v2`, cosine-similarity against pre-embedded rules.
|
| 108 |
+
Output: list of rules sorted by relevance.
|
| 109 |
+
|
| 110 |
+
#### `propose_patch(config: ConfigDict, rules: List[Rule], metrics: RunMetrics) -> Patch`
|
| 111 |
+
Deterministic — no LLM call. Applies rule-to-config transforms (each rule has a `transform` field describing the diff) and computes a per-rule and aggregate uplift estimate from the waste budget.
|
| 112 |
+
Output:
|
| 113 |
+
```python
|
| 114 |
+
Patch(
|
| 115 |
+
new_config=ConfigDict(...),
|
| 116 |
+
diff=str, # unified diff
|
| 117 |
+
rationale=[RuleApplication(...)], # one entry per applied rule
|
| 118 |
+
expected_speedup_low=float, # conservative end of range
|
| 119 |
+
expected_speedup_high=float, # optimistic end of range
|
| 120 |
+
confidence=float, # 0..1, see formula below
|
| 121 |
+
)
|
| 122 |
+
```
|
| 123 |
+
|
| 124 |
+
##### Uplift Estimate
|
| 125 |
+
|
| 126 |
+
For each applied rule with `targets_bucket=B`, predicted recovery is `recovery_fraction × T_B / T_total`. Aggregate predicted recoverable waste = sum across applied rules (capped per bucket). Predicted speedup = `1 / (1 - recoverable_fraction)`. Reported as a range, not a point estimate.
|
| 127 |
+
|
| 128 |
+
##### Confidence Score
|
| 129 |
+
|
| 130 |
+
```
|
| 131 |
+
confidence = evidence_coverage × rule_consistency
|
| 132 |
+
```
|
| 133 |
+
|
| 134 |
+
- `evidence_coverage` — fraction of waste-budget buckets that had real measurement (vs. defaulted-to-zero because the profiler couldn't observe it). Lower if `profile_run` produced partial traces.
|
| 135 |
+
- `rule_consistency` — 1.0 if all applied rules target distinct buckets and don't have conflicting `transform` fields; lower if rules overlap or conflict.
|
| 136 |
+
|
| 137 |
+
(The report's `historical_calibration` term is intentionally dropped — we have no historical data in a hackathon timeframe. Document this honestly in the system prompt.)
|
| 138 |
+
|
| 139 |
+
#### `benchmark(config: ConfigDict, steps: int = 50) -> RunMetrics`
|
| 140 |
+
Same pipeline as `profile_run` but runs longer and at full quality. Returns the same `RunMetrics` type. Result is cached by `sha256((canonical_config_json, workload_script_sha, rocm_image_tag, runner_script_sha))` — re-running the same config is free, and the version-tagged hash prevents stale cache hits when the container or runner changes.
|
| 141 |
+
|
| 142 |
+
#### `compare_runs(before: RunMetrics, after: RunMetrics) -> Report`
|
| 143 |
+
Pure function. Builds the side-by-side report dict the UI renders. Includes the side-by-side waste-budget bar chart that visually shows which buckets shrank.
|
| 144 |
+
|
| 145 |
+
### 4. ROCm Knowledge Base
|
| 146 |
+
|
| 147 |
+
Single file: `kb/rocm_rules.yaml`. ~25 rules at MVP. Schema:
|
| 148 |
+
|
| 149 |
+
```yaml
|
| 150 |
+
- id: precision.bf16_over_fp16_on_mi300x
|
| 151 |
+
category: precision
|
| 152 |
+
symptom: "fp16 used on MI300X"
|
| 153 |
+
detect:
|
| 154 |
+
config.precision: fp16
|
| 155 |
+
fix:
|
| 156 |
+
config.precision: bf16
|
| 157 |
+
expected_impact: "Same throughput, +numerical stability. Reduces NaN risk."
|
| 158 |
+
rocm_version_min: "6.0"
|
| 159 |
+
citation: "AMD ROCm Best Practices Guide §3.2"
|
| 160 |
+
```
|
| 161 |
+
|
| 162 |
+
Categories:
|
| 163 |
+
- `precision` (bf16 over fp16 on CDNA3 matrix cores; FP8 for stretch inference scenarios)
|
| 164 |
+
- `attention` (flash-attn ROCm fork via Optimum-AMD, PyTorch SDPA, packed sequences)
|
| 165 |
+
- `memory` (batch-size sweetspot for 192 GB HBM3, gradient checkpoint thresholds, activation offloading)
|
| 166 |
+
- `kernels` (hipBLASLt hint logging + offline tuning files; MIOpen `MIOPEN_FIND_*` autotune)
|
| 167 |
+
- `env_vars` (`HSA_FORCE_FINE_GRAIN_PCIE`, `MIOPEN_FIND_MODE`, `NCCL_MIN_NCHANNELS=112`, NUMA auto-balancing disable)
|
| 168 |
+
- `optimizer` (8-bit Adam on ROCm — **warn that bitsandbytes is not officially supported on ROCm**; recommend Optimum-AMD-validated alternatives or CPU-offload optimizers)
|
| 169 |
+
- `data` (`num_workers`, `pin_memory=True`, `prefetch_factor`, `persistent_workers=True`, IterableDataset sharding correctness)
|
| 170 |
+
- `compile` (`torch.compile` when graph-break evidence is low; `torch_compile=True` in `TrainingArguments`)
|
| 171 |
+
- `collectives` (RCCL channel count, one-process-per-GPU vs one-process-many-GPUs)
|
| 172 |
+
- `topology` (tensor parallelism within a single XGMI island; prefer TP over EP on single-node)
|
| 173 |
+
|
| 174 |
+
Rules are hand-curated on Day 1 from ROCm docs + AMD blog posts. **This is the moat.** No LLM-generated rules. Each rule entry includes a `citation` field linking back to a ROCm doc page or AMD blog post — every recommendation in the final report carries this citation.
|
| 175 |
+
|
| 176 |
+
**Footgun guardrail — `ROCPROFSYS_*` are NOT tuning vars.** ROCm Systems Profiler env vars (`ROCPROFSYS_MODE`, `ROCPROFSYS_USE_SAMPLING`, etc.) configure how the *profiler* observes a run; they do not affect how work is dispatched to the GPU. The KB **excludes** all `ROCPROFSYS_*` vars from optimization rules. Additionally, `parse_config` flags any user config that sets `ROCPROFSYS_*` as if it were a perf knob and surfaces a warning in the report ("These configure the profiler, not the workload — they will not change throughput").
|
| 177 |
+
|
| 178 |
+
**Workload-validity disclaimer.** Every recommendation is valid for the specific tuple `(workload script, model, GPU=MI300X, ROCm version, framework version, batch/seq pattern)` observed during profiling. The system prompt instructs the agent to state this explicitly when uncertainty is high, and the final report includes a footer line: *"Recommendations validated against MI300X with ROCm <version> and PyTorch <version>. Re-run audit if you change model, hardware, or framework version."*
|
| 179 |
+
|
| 180 |
+
### 5. Profiling Pipeline
|
| 181 |
+
|
| 182 |
+
A wrapper script `goblin_runner.sh` invokes the user's training command inside the ROCm container, with:
|
| 183 |
+
|
| 184 |
+
```bash
|
| 185 |
+
# isolate to a single MI300X to keep concurrent benchmark runs sane
|
| 186 |
+
export ROCR_VISIBLE_DEVICES=${GOBLIN_GPU_ID:-0}
|
| 187 |
+
|
| 188 |
+
rocprofv3 --hsa-trace --kernel-trace -o trace.csv -- \
|
| 189 |
+
python -m torch.profiler.scheduler ... \
|
| 190 |
+
python <user_script> --max_steps=10
|
| 191 |
+
```
|
| 192 |
+
|
| 193 |
+
Post-processing (`profile_parser.py`):
|
| 194 |
+
- parse `trace.csv` for kernel breakdown
|
| 195 |
+
- parse `torch.profiler` JSON for tokens/sec + MFU
|
| 196 |
+
- parse `amd-smi` polling output for HBM + util
|
| 197 |
+
- merge into single `RunMetrics` with populated `WasteBudget`
|
| 198 |
+
|
| 199 |
+
### 6. Benchmark Result Cache
|
| 200 |
+
|
| 201 |
+
`bench_cache/` directory. Key: `sha256((canonical_config_json, workload_script_sha, rocm_image_tag, runner_script_sha))`. Value: full `RunMetrics` plus the unhashed key tuple for debuggability. Lets us re-run the demo without burning cloud time. The Streamlit UI exposes a `--no-cache` toggle for the Day-3 dry-run pass that confirms cached results haven't gone stale.
|
| 202 |
+
|
| 203 |
+
### 7. Synthetic Corpus (`workloads/synthetic/`)
|
| 204 |
+
|
| 205 |
+
A directory of pre-generated, deliberately misconfigured fine-tuning runs. Each entry is `{config.py, expected_findings.yaml, cached_metrics.json}`. Generated on Day 1 by ROCm Lead by perturbing the canonical Qwen2.5-7B LoRA workload along single dimensions: `num_workers=0`, FP32 instead of BF16, naive attention, `torch.compile=False`, untuned hipBLASLt, etc.
|
| 206 |
+
|
| 207 |
+
Two purposes:
|
| 208 |
+
1. **Backend Lead can develop on a laptop** — test the agent loop, KB queries, and patch generation against `cached_metrics.json` without touching the GPU.
|
| 209 |
+
2. **Demo lane 1 (offline replay)** — judges can audit any synthetic scenario in <30 seconds without live MI300X. If cloud is unreachable on demo day, this is the safety net.
|
| 210 |
+
|
| 211 |
+
Each scenario's `expected_findings.yaml` lists which rules *should* fire — gives us a regression test for KB recall.
|
| 212 |
+
|
| 213 |
+
### 8. Demo Lanes
|
| 214 |
+
|
| 215 |
+
The system runs in one of two modes:
|
| 216 |
+
|
| 217 |
+
| Lane | Source of metrics | Use |
|
| 218 |
+
|---|---|---|
|
| 219 |
+
| **Offline replay** | Synthetic corpus + cached `RunMetrics` JSON | Dev, regression testing, demo backup if MI300X unreachable |
|
| 220 |
+
| **Live MI300X** | `goblin_runner.sh` → real GPU | Canonical demo, "before/after with real numbers" pitch moment |
|
| 221 |
+
|
| 222 |
+
The agent loop is identical in both lanes — only the `RunnerProtocol` implementation differs. This is the seam introduced to address the audit's testability finding.
|
| 223 |
+
|
| 224 |
+
### 9. Secrets & Privacy
|
| 225 |
+
|
| 226 |
+
Uploaded scripts and logs may contain API tokens, dataset paths, internal URLs. Before any artefact is persisted to disk or sent to the LLM:
|
| 227 |
+
|
| 228 |
+
1. **Regex redaction pass** in `parse_config` for common patterns: `Bearer [A-Za-z0-9._-]+`, `sk-[A-Za-z0-9]{20,}`, `hf_[A-Za-z0-9]{20,}`, `/home/<user>/`, `s3://`, `wss?://`.
|
| 229 |
+
2. **Ephemeral storage** — `bench_cache/` is gitignored; uploaded scripts are deleted at session end.
|
| 230 |
+
3. **No model weights or training data required** — we audit configs and metrics only. The system prompt explicitly tells the agent it cannot ask for weights or datasets.
|
| 231 |
+
|
| 232 |
+
This is a one-screen redaction pass, not a full PII pipeline. Sufficient for hackathon; would need expansion for production.
|
| 233 |
+
|
| 234 |
+
## Data Flow — One Audit
|
| 235 |
+
|
| 236 |
+
```
|
| 237 |
+
[user uploads script] ──► [redaction pass]
|
| 238 |
+
│
|
| 239 |
+
▼
|
| 240 |
+
parse_config ──► ConfigDict ────────────────────────────┐
|
| 241 |
+
│ │
|
| 242 |
+
▼ │
|
| 243 |
+
profile_run ──► RunMetrics + WasteBudget ──► query_rocm_kb │
|
| 244 |
+
│ │ │
|
| 245 |
+
│ ▼ │
|
| 246 |
+
│ List[Rule] ──────►│
|
| 247 |
+
│ ▼
|
| 248 |
+
│ propose_patch (uplift + confidence)
|
| 249 |
+
│ │
|
| 250 |
+
│ ▼
|
| 251 |
+
│ Patch (new ConfigDict + diff)
|
| 252 |
+
│ │
|
| 253 |
+
▼ ▼
|
| 254 |
+
benchmark(original) ◄──── cache ─────► benchmark(patched)
|
| 255 |
+
│ │
|
| 256 |
+
└────────► compare_runs ◄────────────────────────┘
|
| 257 |
+
│
|
| 258 |
+
▼
|
| 259 |
+
Report ──► UI
|
| 260 |
+
(waste-budget bar chart + metrics + diff)
|
| 261 |
+
```
|
| 262 |
+
|
| 263 |
+
## Tech Stack (frozen)
|
| 264 |
+
|
| 265 |
+
| Layer | Choice | Why |
|
| 266 |
+
|---|---|---|
|
| 267 |
+
| Agent LLM | Qwen2.5-7B-Instruct via HF Inference Providers | Tool calling via OpenAI-compatible API; routes through Together / Fireworks-AI / Nebius (auto). Stretch: self-host on MI300X via vLLM. |
|
| 268 |
+
| Backend | Python 3.11 + FastAPI + anthropic SDK | Standard, async-friendly |
|
| 269 |
+
| Frontend | Streamlit | Ships in hours, not days |
|
| 270 |
+
| Container | `rocm/pytorch:rocm6.1_ubuntu22.04_py3.10_pytorch_2.3` | Official AMD image |
|
| 271 |
+
| Profiling | torch.profiler + rocprofv3 + amd-smi | Native ROCm tools |
|
| 272 |
+
| KB index | sentence-transformers + numpy cosine | No vector DB needed for 25 rules |
|
| 273 |
+
| Workload | Qwen2.5-7B + LoRA + alpaca-cleaned | Canonical, reproducible |
|
| 274 |
+
| Hardware | MI300X cloud (single GPU, 192 GB HBM) | Hackathon constraint |
|
| 275 |
+
|
| 276 |
+
## Repo Layout
|
| 277 |
+
|
| 278 |
+
```
|
| 279 |
+
amd-hackathon/
|
| 280 |
+
├── brainstorming/ # this folder
|
| 281 |
+
├── kb/
|
| 282 |
+
│ └── rocm_rules.yaml # the moat
|
| 283 |
+
├── agent/
|
| 284 |
+
│ ├── server.py # FastAPI app
|
| 285 |
+
│ ├── loop.py # agent loop driver
|
| 286 |
+
│ ├── prompts.py # system prompt
|
| 287 |
+
│ └── tools/
|
| 288 |
+
│ ├── parse_config.py
|
| 289 |
+
│ ├── profile_run.py
|
| 290 |
+
│ ├── query_rocm_kb.py
|
| 291 |
+
│ ├── propose_patch.py
|
| 292 |
+
│ ├── benchmark.py
|
| 293 |
+
│ └── compare_runs.py
|
| 294 |
+
├── runner/
|
| 295 |
+
│ ├── goblin_runner.sh # rocprofv3 wrapper
|
| 296 |
+
│ └── profile_parser.py
|
| 297 |
+
├── ui/
|
| 298 |
+
│ └── app.py # Streamlit
|
| 299 |
+
├── workloads/
|
| 300 |
+
│ ├── train_qwen_lora.py # canonical demo workload
|
| 301 |
+
│ └── synthetic/ # pre-generated misconfigured runs + cached metrics
|
| 302 |
+
│ ├── 01_no_workers/
|
| 303 |
+
│ ├── 02_fp32_default/
|
| 304 |
+
│ ├── 03_naive_attention/
|
| 305 |
+
│ └── ...
|
| 306 |
+
├── bench_cache/ # gitignored
|
| 307 |
+
└── docker/
|
| 308 |
+
└── Dockerfile # rocm/pytorch + our deps
|
| 309 |
+
```
|
| 310 |
+
|
| 311 |
+
## Agent Loop — Pseudocode
|
| 312 |
+
|
| 313 |
+
The loop is provider-agnostic. It talks to a `Backend` (see `agent/backends/`); today's only concrete is `QwenHFBackend`.
|
| 314 |
+
|
| 315 |
+
```python
|
| 316 |
+
def run_audit(user_script_path: str) -> SSEStream:
|
| 317 |
+
backend = make_backend(system_prompt=SYSTEM_PROMPT) # QwenHFBackend
|
| 318 |
+
backend.add_user_message(f"Audit this fine-tuning workload: {user_script_path}")
|
| 319 |
+
|
| 320 |
+
for step in range(MAX_STEPS): # MAX_STEPS = 8
|
| 321 |
+
turn = await backend.next_turn(tool_schemas())
|
| 322 |
+
for text in turn.text_blocks:
|
| 323 |
+
yield {"type": "thought", "text": text}
|
| 324 |
+
for tc in turn.tool_calls:
|
| 325 |
+
yield {"type": "tool_call", "name": tc.name, "input": tc.input}
|
| 326 |
+
result = call_tool(tc.name, **tc.input)
|
| 327 |
+
yield {"type": "tool_result", "name": tc.name, "result": result}
|
| 328 |
+
backend.add_tool_result(tc.id, tc.name, result.content, is_error=not result.ok)
|
| 329 |
+
if turn.stop_reason == "end_turn":
|
| 330 |
+
break
|
| 331 |
+
|
| 332 |
+
yield {"type": "final_report", "report": extract_final_report(tool_results)}
|
| 333 |
+
```
|
| 334 |
+
|
| 335 |
+
The Backend protocol (`agent/backends/base.py`) is a 3-method contract:
|
| 336 |
+
`add_user_message`, `next_turn`, `add_tool_result`. `QwenHFBackend` translates
|
| 337 |
+
between this neutral shape and OpenAI-compatible chat-completion calls
|
| 338 |
+
through HF Inference Providers.
|
| 339 |
+
|
| 340 |
+
## Boundaries & Interfaces
|
| 341 |
+
|
| 342 |
+
- **Agent ↔ Tools:** Each tool is a pure function with typed input/output. No shared state. Tools never call other tools — only the agent orchestrates.
|
| 343 |
+
- **Tools ↔ MI300X:** Only `profile_run` and `benchmark` touch the GPU. Both go through `goblin_runner.sh` for consistency.
|
| 344 |
+
- **Backend ↔ UI:** SSE stream of typed events (`thought`, `tool_call`, `tool_result`, `final_report`). UI just renders events; no agent logic on frontend.
|
| 345 |
+
- **KB ↔ Code:** YAML is the single source of truth. Code consumes; never generates rules at runtime.
|
| 346 |
+
|
| 347 |
+
## What's Deliberately Excluded (YAGNI)
|
| 348 |
+
|
| 349 |
+
- ❌ Multi-GPU / distributed training analysis (single MI300X is enough; collective rules in KB still apply when user uploads multi-GPU configs, we just don't benchmark them)
|
| 350 |
+
- ❌ Pretraining workloads (fine-tuning only)
|
| 351 |
+
- ❌ vLLM / inference serving as a primary mode (Day-4 stretch only; vLLM-specific rules NOT in MVP KB)
|
| 352 |
+
- ❌ Live job watching with intervention (offline audit only — much simpler)
|
| 353 |
+
- ❌ User accounts / persistence / DB (in-memory session, hackathon)
|
| 354 |
+
- ❌ Vector DB (25 rules + numpy is fine)
|
| 355 |
+
- ❌ Polars/DuckDB for run storage (Pydantic + JSON is enough at this scale)
|
| 356 |
+
- ❌ ML model for uplift prediction (deterministic waste-budget calculus is enough)
|
| 357 |
+
- ❌ Auto-applying patches to the user's repo (we output a diff, user applies)
|
| 358 |
+
- ❌ Sliders/what-if interactive panel (chat layer answers what-if conversationally; sliders are stretch only)
|
| 359 |
+
- ❌ Cost calculator beyond a static "$ per training run" line (stretch goal only)
|
brainstorming/goals.md
ADDED
|
@@ -0,0 +1,233 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# GPU Goblin — Goals & Implementation Plan
|
| 2 |
+
|
| 3 |
+
## North Star
|
| 4 |
+
|
| 5 |
+
Win the AMD hackathon (**Track 1: AI Agents & Agentic Workflows**) by demonstrating a **real, reproducible 2×+ throughput improvement on a Qwen2.5-7B LoRA fine-tune on MI300X**, driven end-to-end by a Qwen-powered tool-using agent. Hits the Qwen Technology Partner challenge with end-to-end Qwen-on-AMD; deploys as a Hugging Face Space within the hackathon HF Organization.
|
| 6 |
+
|
| 7 |
+
Everything else is in service of that demo.
|
| 8 |
+
|
| 9 |
+
## Success Criteria
|
| 10 |
+
|
| 11 |
+
| # | Criterion | Measurable As |
|
| 12 |
+
|---|---|---|
|
| 13 |
+
| 1 | Real MI300X speedup | ≥ 2.0× tokens/sec, before vs. after, on the canonical workload |
|
| 14 |
+
| 2 | Agent is actually agentic | ≥ 4 tool calls per audit, visible in UI |
|
| 15 |
+
| 3 | AMD differentiation | 100% of recommendations cite a ROCm-specific KB rule |
|
| 16 |
+
| 4 | Uplift estimate honesty | Predicted speedup range (low–high) brackets the measured speedup on ≥ 8 of 10 synthetic corpus scenarios |
|
| 17 |
+
| 5 | Offline parse coverage | ≥ 90% of synthetic corpus parsed without manual edits |
|
| 18 |
+
| 6 | End-to-end reliability | Demo runs cleanly 5× in a row without manual intervention |
|
| 19 |
+
| 7 | Pitch quality | 3-min demo + 1-min architecture + 1-min impact, rehearsed |
|
| 20 |
+
| 8 | Backup safety | Offline-replay lane works fully when MI300X disconnected |
|
| 21 |
+
|
| 22 |
+
## Mapping to Judging Criteria
|
| 23 |
+
|
| 24 |
+
| Hackathon criterion | Where we score |
|
| 25 |
+
|---|---|
|
| 26 |
+
| **Application of Technology** | Tool-using agent loop on MI300X; rocprofv3 + torch.profiler integration; deterministic uplift via waste-budget calculus; ROCm-specific KB |
|
| 27 |
+
| **Presentation** | Live before/after benchmark on real MI300X; waste-budget bar chart; chat-based "why?" interaction; 90-sec backup video |
|
| 28 |
+
| **Business Value** | Each audit reports `$ saved per training run` and `time saved per epoch`, framed against $1.99/hr public MI300X pricing reference |
|
| 29 |
+
| **Originality** | "AI for AI builders" — meta-positioning. Real benchmarks, not just LLM advice. Waste-budget decomposition is novel framing |
|
| 30 |
+
|
| 31 |
+
## Team Roles (3 people × 4 days)
|
| 32 |
+
|
| 33 |
+
### ROCm / ML Lead
|
| 34 |
+
- Day 1: MI300X cloud env, ROCm container, baseline workload (Qwen2.5-7B LoRA on alpaca)
|
| 35 |
+
- Day 1-2: Hand-curate 20-25 KB rules from ROCm docs + AMD blog
|
| 36 |
+
- Day 2: `profile_run` + `benchmark` tools (rocprofv3 wrapper, parser)
|
| 37 |
+
- Day 3: Validate end-to-end speedup on canonical demo, generate cached results
|
| 38 |
+
- **Owns:** anything that touches the GPU
|
| 39 |
+
|
| 40 |
+
### Agent / Backend Lead
|
| 41 |
+
- Day 1: FastAPI skeleton, tool schemas, Qwen-via-HF tool-use plumbing
|
| 42 |
+
- Day 1-2: `parse_config`, `propose_patch`, `query_rocm_kb`, `compare_runs` tools
|
| 43 |
+
- Day 2: Full agent loop, system prompt, SSE streaming
|
| 44 |
+
- Day 3: Hardening, error handling, max-steps cap, fallback behaviors
|
| 45 |
+
- **Owns:** agent reasoning quality + tool wiring
|
| 46 |
+
|
| 47 |
+
### Frontend / Demo Lead
|
| 48 |
+
- Day 1: Streamlit skeleton, file upload, message log
|
| 49 |
+
- Day 2: Tool-call cards (live status), report renderer, diff viewer
|
| 50 |
+
- Day 3: Demo polish, charts, golden-run replay mode
|
| 51 |
+
- Day 4: Pitch deck, 90-second backup video, dry runs
|
| 52 |
+
- **Owns:** what the judges actually see
|
| 53 |
+
|
| 54 |
+
Roles overlap on integration days — pair-program when blocked.
|
| 55 |
+
|
| 56 |
+
## Day-by-Day Plan
|
| 57 |
+
|
| 58 |
+
### Day 1 (May 5–6) — Foundations
|
| 59 |
+
|
| 60 |
+
**ROCm Lead**
|
| 61 |
+
- [ ] Provision MI300X cloud instance via AMD Developer Cloud ($100 credits), SSH access, persistent storage
|
| 62 |
+
- [ ] Pull `rocm/pytorch:rocm6.1_*` image, verify GPU visible inside container
|
| 63 |
+
- [ ] Run baseline `train_qwen_lora.py` (batch=4, fp16, naive attention, alpaca, 100 steps)
|
| 64 |
+
- [ ] Capture baseline tokens/sec, MFU, HBM peak — *this is our "before"*
|
| 65 |
+
- [ ] Generate **synthetic corpus** — 5-8 misconfigured variants of the canonical workload (FP32, num_workers=0, naive attention, etc.) with cached `RunMetrics` JSON for each
|
| 66 |
+
- [ ] Start drafting KB rules (target 10 rules by EOD), each tagged with `targets_bucket` matching the waste-budget decomposition
|
| 67 |
+
|
| 68 |
+
**Backend Lead**
|
| 69 |
+
- [ ] Repo scaffold per architecture.md layout
|
| 70 |
+
- [ ] `pip install fastapi anthropic sentence-transformers pyyaml pydantic`
|
| 71 |
+
- [ ] Define `agent/schemas.py` with `RunMetrics`, `WasteBudget`, `ConfigDict`, `Patch`, `Rule`, `Report` as pydantic models — **Day-1 priority** (blocks all tools)
|
| 72 |
+
- [ ] Define `RunnerProtocol` interface; build `FakeRunner` that loads cached metrics from `workloads/synthetic/` (lets backend dev without MI300X)
|
| 73 |
+
- [ ] FastAPI `POST /audit` skeleton with SSE
|
| 74 |
+
- [ ] `parse_config` tool — handle HF `TrainingArguments` first; include regex redaction pass for tokens/paths
|
| 75 |
+
- [ ] Qwen tool-use hello-world (one tool, one round-trip via HF Inference Providers)
|
| 76 |
+
|
| 77 |
+
**Frontend Lead**
|
| 78 |
+
- [ ] Streamlit skeleton with file upload + chat panel
|
| 79 |
+
- [ ] Hardcoded "fake audit" — render canned tool calls + report
|
| 80 |
+
- [ ] Pick chart library (Altair recommended for Streamlit)
|
| 81 |
+
|
| 82 |
+
**Day 1 Exit Criteria**
|
| 83 |
+
- Baseline benchmark numbers in hand
|
| 84 |
+
- Synthetic corpus has ≥ 3 cached scenarios (Backend Lead can now dev without GPU)
|
| 85 |
+
- Schemas (`RunMetrics`, `WasteBudget`, etc.) frozen
|
| 86 |
+
- `RunnerProtocol` + `FakeRunner` working end-to-end
|
| 87 |
+
- Backend can call Qwen with one tool via HF Inference Providers
|
| 88 |
+
- UI renders a fake audit
|
| 89 |
+
- 10 KB rules drafted
|
| 90 |
+
|
| 91 |
+
### Day 2 (May 6–7) — Core Build
|
| 92 |
+
|
| 93 |
+
**ROCm Lead**
|
| 94 |
+
- [ ] Finish KB to 20-25 rules; hand-tag categories + `targets_bucket`; pre-embed with sentence-transformers
|
| 95 |
+
- [ ] Include the high-impact MI300X rules: BF16-over-FP16, AITER-flash-attn-via-Optimum-AMD, `NCCL_MIN_NCHANNELS=112`, NUMA disable, one-process-per-GPU, hipBLASLt hint logging, MIOpen `MIOPEN_FIND_*`, **bitsandbytes-not-supported-on-ROCm warning**, `num_workers`/`pin_memory`/`prefetch_factor`/`persistent_workers`
|
| 96 |
+
- [ ] `profile_run` tool: rocprofv3 wrapper + torch.profiler + amd-smi → `RunMetrics` with `WasteBudget`
|
| 97 |
+
- [ ] `benchmark` tool: same pipeline, longer run, with version-tagged cache
|
| 98 |
+
- [ ] Validate on baseline workload: profile output makes sense, MFU is plausible, waste budget sums to ~T_total
|
| 99 |
+
|
| 100 |
+
**Backend Lead**
|
| 101 |
+
- [ ] All 6 tools wired and individually tested with fixtures from synthetic corpus
|
| 102 |
+
- [ ] Full agent loop with max-steps cap, SSE event types finalized, error envelope per tool (`{ok, result, error}`)
|
| 103 |
+
- [ ] System prompt iterated against test workloads — include MI300X hardware specs (304 CUs, 192 GB HBM3, ~5.3 TB/s, FP8 native) so the agent reasons quantitatively
|
| 104 |
+
- [ ] `propose_patch` deterministic transformer + uplift estimator (waste-budget × bucket recovery) + confidence formula
|
| 105 |
+
|
| 106 |
+
**Frontend Lead**
|
| 107 |
+
- [ ] Live tool-call cards consuming real SSE stream
|
| 108 |
+
- [ ] Final report layout: side-by-side metrics + diff + kernel chart + **waste-budget bar chart** (where time was lost, before vs after)
|
| 109 |
+
- [ ] Lane toggle: `Live MI300X` vs `Offline replay (synthetic corpus)` — judges can pick either
|
| 110 |
+
- [ ] First end-to-end run through the actual backend
|
| 111 |
+
|
| 112 |
+
**Day 2 Exit Criteria**
|
| 113 |
+
- Real audit runs end-to-end on a toy workload
|
| 114 |
+
- Profile + benchmark return real MI300X numbers
|
| 115 |
+
- KB has 20+ rules, embeddings pre-computed
|
| 116 |
+
- UI streams real agent activity
|
| 117 |
+
|
| 118 |
+
### Day 3 (May 7–8) — Demo Day Prep
|
| 119 |
+
|
| 120 |
+
**All hands**
|
| 121 |
+
- [ ] Run canonical demo (Qwen2.5-7B LoRA) end-to-end → confirm ≥ 2× speedup
|
| 122 |
+
- [ ] Cache the demo benchmark results — don't burn cloud time on every rehearsal
|
| 123 |
+
- [ ] Build "golden run" replay mode (read cached SSE events, replay timing)
|
| 124 |
+
- [ ] Validate uplift accuracy on synthetic corpus: predicted range should bracket measured speedup on ≥ 8 of 10 scenarios
|
| 125 |
+
- [ ] Polish system prompt for demo-friendly narration ("I'll start by…")
|
| 126 |
+
- [ ] Tighten error handling — agent should never panic in front of judges; tool failures degrade gracefully with `{ok:false}` envelopes
|
| 127 |
+
- [ ] Run with `--no-cache` once to verify cached results aren't masking real bugs
|
| 128 |
+
- [ ] 5× clean dry runs, fix anything flaky
|
| 129 |
+
|
| 130 |
+
**Day 3 Exit Criteria**
|
| 131 |
+
- Canonical demo: cleanly runs in ≤ 4 minutes, ≥ 2× speedup, no manual fixes
|
| 132 |
+
- Cached results enable offline replay
|
| 133 |
+
- Offline-replay lane works fully when MI300X is disconnected (proven by unplugging cloud)
|
| 134 |
+
- Golden-run video recorded as backup
|
| 135 |
+
|
| 136 |
+
### Day 4 (May 8–9) — Pitch & Stretch
|
| 137 |
+
|
| 138 |
+
**Frontend Lead**
|
| 139 |
+
- [ ] Pitch deck (5 slides): problem, agent loop diagram, demo (live), KB rules sample, impact + ask
|
| 140 |
+
- [ ] Cover image for the submission listing
|
| 141 |
+
- [ ] Final 90-second backup video
|
| 142 |
+
- [ ] Submission form filled (title, short + long description, tags), repo public, README crisp
|
| 143 |
+
- [ ] **Build-in-Public bonus track:** at least one public update post (X/LinkedIn/Discord) showing the agent's first audit; one ROCm/Optimum-AMD feedback note based on what was rough during build
|
| 144 |
+
|
| 145 |
+
**ROCm + Backend Leads (in parallel, optional stretch)**
|
| 146 |
+
- [ ] vLLM inference workload as second demo (only if rock-solid on Day 3)
|
| 147 |
+
- [ ] Cost calculator: `$ saved per training run` line, anchored on $1.99/hr public MI300X reference
|
| 148 |
+
- [ ] What-if slider panel for batch / precision / attention (chat already does this conversationally, sliders are visual icing)
|
| 149 |
+
- [ ] Stretch dream: self-host Qwen on MI300X via vLLM (replacing the HF-Inference-Providers path) — closes the loop entirely on AMD silicon. Mention in pitch even if not running live.
|
| 150 |
+
|
| 151 |
+
**Day 4 Exit Criteria**
|
| 152 |
+
- Submission complete by deadline
|
| 153 |
+
- 5+ rehearsed dry runs of the pitch
|
| 154 |
+
- Cover image, video, slides, repo all linked from submission form
|
| 155 |
+
- Backup video on USB stick, in cloud, on phone
|
| 156 |
+
|
| 157 |
+
## Scope Discipline (YAGNI)
|
| 158 |
+
|
| 159 |
+
If we're behind schedule, **cut in this order**:
|
| 160 |
+
|
| 161 |
+
1. ❌ All stretch goals (vLLM, cost calc, self-hosted agent)
|
| 162 |
+
2. ❌ Live tool-call UI animations — static cards work
|
| 163 |
+
3. ❌ Some KB categories (keep precision/attention/memory; drop env_vars/collectives)
|
| 164 |
+
4. ❌ Multi-file script parsing — single-file only
|
| 165 |
+
5. ❌ Dynamic batch-size search — hardcode the recommended batch for demo workload
|
| 166 |
+
|
| 167 |
+
**Never cut:**
|
| 168 |
+
- Real MI300X benchmark (that's the entire pitch)
|
| 169 |
+
- Tool-using agent loop (that's the track fit)
|
| 170 |
+
- ROCm-specific KB citations (that's the differentiation)
|
| 171 |
+
|
| 172 |
+
## Risk Register
|
| 173 |
+
|
| 174 |
+
| Risk | Likelihood | Impact | Mitigation |
|
| 175 |
+
|---|---|---|---|
|
| 176 |
+
| MI300X cloud quota burns out | Medium | High | Cache every benchmark, dev with 10-step traces, full benchmarks only for canonical demo |
|
| 177 |
+
| `rocprofv3` flaky / version mismatch | Medium | High | Always run `torch.profiler` in parallel as fallback. Pre-record golden trace. Use `rocprofv3` not deprecated `rocprof`/`rocprofv2` |
|
| 178 |
+
| MI300X cloud unreachable on demo day | Low | Critical | **Offline-replay lane** (synthetic corpus + cached metrics) provides full demo without cloud — judges can't tell the difference for the agent-reasoning portion |
|
| 179 |
+
| Uploaded scripts contain user secrets / tokens | Medium | Medium | Regex redaction pass in `parse_config` before any persistence or LLM call; ephemeral storage; explicit "we never store weights or datasets" line in UI |
|
| 180 |
+
| Agent loops infinitely | Low | Medium | Hard cap of 8 tool calls, fall back to "best effort" report after cap |
|
| 181 |
+
| Recommendations are generic, not ROCm-specific | Medium | Critical | Hand-curate KB on Day 1 *before* wiring agent — KB is the moat |
|
| 182 |
+
| Demo crashes during pitch | Low | Critical | Pre-recorded video backup. Golden-run replay mode. USB + cloud + phone copies |
|
| 183 |
+
| Qwen2.5-7B doesn't fit a 12-batch on MI300X | Low | Medium | Have a fallback config in hand (batch=8 + grad_accum=2) |
|
| 184 |
+
| LoRA on alpaca too easy — speedup looks staged | Low | High | Measure on a non-trivial seq_len (1024+), include MFU not just tokens/sec, show kernel breakdown to prove it's real |
|
| 185 |
+
| HF Inference Provider rate limit / outage during demo | Low | Medium | Offline-replay UI lane plays cached_audit.json without any backend; pre-cache a full recorded session; have a backup HF token; `provider="auto"` already routes around individual provider outages |
|
| 186 |
+
| Team member unavailable (illness, etc.) | Low | High | Pair on critical path (agent loop, KB) so no single point of failure |
|
| 187 |
+
|
| 188 |
+
## Definition of Done — MVP
|
| 189 |
+
|
| 190 |
+
GPU Goblin is "done enough to ship" when **all** of these are true:
|
| 191 |
+
|
| 192 |
+
- ✅ A judge can upload `train_qwen_lora.py` (we provide it) and get a real audit
|
| 193 |
+
- ✅ Agent makes ≥ 4 visible tool calls
|
| 194 |
+
- ✅ Final report shows ≥ 2× tokens/sec, real numbers from MI300X
|
| 195 |
+
- ✅ Every recommendation in the report cites a ROCm KB rule by ID
|
| 196 |
+
- ✅ User can ask follow-up questions in chat ("why bf16?") and get cited answers
|
| 197 |
+
- ✅ Full audit completes in ≤ 4 minutes
|
| 198 |
+
- ✅ 5 consecutive dry runs succeed without manual intervention
|
| 199 |
+
- ✅ Backup video and cached replay both work
|
| 200 |
+
|
| 201 |
+
## Stretch Definition of Done
|
| 202 |
+
|
| 203 |
+
- 🎯 Second canonical workload (vLLM inference) audited end-to-end
|
| 204 |
+
- 🎯 Cost calculator: "you save $X per training run, $Y per epoch"
|
| 205 |
+
- 🎯 Agent backed by self-hosted Qwen via vLLM on the same MI300X (the ultimate AMD story — replaces today's HF Inference Providers path with on-cluster serving)
|
| 206 |
+
|
| 207 |
+
## Compute Budget — AMD Developer Cloud Credits
|
| 208 |
+
|
| 209 |
+
Eligible participants get **$100 in AMD Developer Cloud credits**. Public reference price for MI300X is around $1.99/hr (single VM) or $3.39/hr (8× bare metal). Plan accordingly:
|
| 210 |
+
|
| 211 |
+
| Activity | GPU-hours | Notes |
|
| 212 |
+
|---|---|---|
|
| 213 |
+
| Day-1 baseline + synthetic corpus generation | 4–6 | One-shot, results cached |
|
| 214 |
+
| Day-2 KB validation runs | 2–3 | Sanity-check the rules fire on synthetic scenarios |
|
| 215 |
+
| Day-3 canonical demo dry runs (cached) | 2–4 | Cache hits after the first run |
|
| 216 |
+
| Day-3 `--no-cache` cold validation | 1 | Confirms nothing's stale |
|
| 217 |
+
| Day-4 final dry runs + record video | 2–3 | Lock the demo |
|
| 218 |
+
| **Total estimate** | **~12–17 hrs** | Well within $100 even at bare-metal rates |
|
| 219 |
+
|
| 220 |
+
Backend Lead spends zero MI300X time after Day 1 — develops against synthetic corpus + `FakeRunner`.
|
| 221 |
+
|
| 222 |
+
## Submission Checklist
|
| 223 |
+
|
| 224 |
+
- [ ] Public GitHub repo with clear README + architecture diagram + setup instructions
|
| 225 |
+
- [ ] 90-second demo video (live agent run, real MI300X numbers)
|
| 226 |
+
- [ ] Pitch deck (PDF or slides URL)
|
| 227 |
+
- [ ] Cover image (project listing visual)
|
| 228 |
+
- [ ] Architecture diagram (PNG)
|
| 229 |
+
- [ ] Sample audit report PDF (one canonical run, before/after, including waste-budget chart)
|
| 230 |
+
- [ ] Short description (1-2 sentences) + long description (paragraph) + technology/category tags
|
| 231 |
+
- [ ] At least one Build-in-Public post + one ROCm/Optimum-AMD feedback note (bonus track)
|
| 232 |
+
- [ ] Hackathon submission form filled
|
| 233 |
+
- [ ] Team member credits + contact
|
brainstorming/idea.md
ADDED
|
@@ -0,0 +1,105 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# GPU Goblin — Idea
|
| 2 |
+
|
| 3 |
+
> **An AI agent that hunts wasted compute on AMD MI300X.**
|
| 4 |
+
> Upload a fine-tuning config; the agent profiles a real run, diagnoses inefficiency, recommends ROCm-specific fixes, then re-runs and proves the speedup with hard numbers.
|
| 5 |
+
|
| 6 |
+
---
|
| 7 |
+
|
| 8 |
+
## The Problem
|
| 9 |
+
|
| 10 |
+
Most teams fine-tuning LLMs on AMD hardware leave **2–3× of throughput on the floor** without realizing it. The waste hides in plain sight:
|
| 11 |
+
|
| 12 |
+
- **HBM under-fed** — batch sizes copied from NVIDIA tutorials don't exploit MI300X's 192 GB.
|
| 13 |
+
- **Wrong precision** — `fp16` is the reflexive default; `bf16` matches throughput on CDNA3 with better numerics.
|
| 14 |
+
- **Naive attention** — teams ship without flash-attn ROCm fork or PyTorch SDPA enabled.
|
| 15 |
+
- **Generic kernels** — hipBLASLt tuning, MIOpen autotune, RCCL collective tweaks all left at defaults.
|
| 16 |
+
- **Wrong libraries** — `bitsandbytes` is the reflexive choice for 8-bit Adam, but it's **not officially supported on ROCm**. Optimum-AMD validates Flash Attention 2 + GPTQ/AWQ paths instead.
|
| 17 |
+
- **Distributed footguns** — one-process-driving-many-GPUs serializes kernel launches on ROCm; NUMA auto-balancing causes variance; `NCCL_MIN_NCHANNELS` left at default.
|
| 18 |
+
- **Knobs nobody knows about** — `HSA_FORCE_FINE_GRAIN_PCIE`, `MIOPEN_FIND_MODE`, hipBLASLt cache paths.
|
| 19 |
+
|
| 20 |
+
Engineers don't fix this not because they're lazy — it's because the knowledge is **scattered across ROCm docs, AMD blog posts, GitHub issues, and Discord threads**. Nobody has time to read it all before launching a run.
|
| 21 |
+
|
| 22 |
+
## The Solution — GPU Goblin
|
| 23 |
+
|
| 24 |
+
A tool-using agent that does what an experienced AMD performance engineer would do, in one minute:
|
| 25 |
+
|
| 26 |
+
1. **Read** the user's fine-tuning script or HF `TrainingArguments`.
|
| 27 |
+
2. **Profile** a 10-step warm run on MI300X (torch.profiler + `rocprofv3` + amd-smi).
|
| 28 |
+
3. **Diagnose** against a curated knowledge base of ROCm-specific optimizations.
|
| 29 |
+
4. **Patch** the config — concrete diff, not vague advice.
|
| 30 |
+
5. **Benchmark** the patched config on the same MI300X.
|
| 31 |
+
6. **Report** before/after side-by-side with real numbers and citations.
|
| 32 |
+
|
| 33 |
+
The user keeps interacting in chat: *"why bf16?"*, *"what if I can't change the optimizer?"*, *"how much $ does this save per epoch?"*
|
| 34 |
+
|
| 35 |
+
## Why This Wins
|
| 36 |
+
|
| 37 |
+
**Most teams build user-facing AI. We build AI for AI builders.** That alone is judge-bait.
|
| 38 |
+
|
| 39 |
+
But the deeper hook is that GPU Goblin is **provably useful** in a way most hackathon projects aren't. We don't have to convince judges the recommendations are good — we *show* them, on the same MI300X, with the same model, in the same demo. Tokens/sec: `142 → 318`. End of debate.
|
| 40 |
+
|
| 41 |
+
### Track Fit (Track 1: AI Agents & Agentic Workflows)
|
| 42 |
+
|
| 43 |
+
- **Primary track:** AI Agents & Agentic Workflows. Real tool-using loop. The agent observes (profile), hypothesizes (KB query), tests (benchmark), refines (patch). Every step visible in chat. Not a one-shot LLM call dressed up as an agent.
|
| 44 |
+
- **Fine-tuning is what we *audit*, not the track we enter.** Canonical workload is **Qwen2.5-7B-Instruct + LoRA on MI300X**. Recommendations are *fine-tuning specific* — batch sizing for LoRA, gradient checkpointing thresholds, optimizer choice (8-bit Adam on ROCm), seq-length packing.
|
| 45 |
+
|
| 46 |
+
### Qwen Technology Partner Challenge
|
| 47 |
+
|
| 48 |
+
GPU Goblin satisfies the Qwen partner challenge two ways:
|
| 49 |
+
|
| 50 |
+
1. **Qwen as the agent brain.** The audit loop runs on `Qwen/Qwen2.5-7B-Instruct` served via Hugging Face Inference Providers. Every tool call, every recommendation, every chat answer is generated by Qwen.
|
| 51 |
+
2. **Qwen as the audit target.** The canonical demo workload is a Qwen2.5-7B-Instruct LoRA fine-tune. The agent audits *the same model family it runs on*.
|
| 52 |
+
|
| 53 |
+
Result: end-to-end Qwen on AMD silicon. Directly answers the judging criterion *"How effectively the chosen model(s) are integrated into the solution."*
|
| 54 |
+
|
| 55 |
+
### Hugging Face Integration
|
| 56 |
+
|
| 57 |
+
Hugging Face is the named Technology Partner — both **model hub** and **deployment layer**:
|
| 58 |
+
|
| 59 |
+
- **Model hub:** Qwen models pulled from HF Hub at audit time + at agent-spin-up time.
|
| 60 |
+
- **Optimum-AMD:** several KB rules cite Optimum-AMD validated paths (Flash Attention 2 on MI300, GPTQ/AWQ ROCm path).
|
| 61 |
+
- **Inference Providers:** the agent's Qwen calls go through HF's Inference Providers router (`provider="auto"` selects Together / Fireworks-AI / Nebius based on availability).
|
| 62 |
+
- **Deployment as HF Space:** the Streamlit UI ships as a Space within the AMD Developer Hackathon HF Organization, satisfying the "Demo Application Platform" + "Application URL" submission fields. The Space runs in offline-replay mode by default so it works without any live backend.
|
| 63 |
+
|
| 64 |
+
### AMD Differentiation
|
| 65 |
+
|
| 66 |
+
Every single recommendation cites a **ROCm-specific rule**, not generic PyTorch advice. The knowledge base is the moat:
|
| 67 |
+
|
| 68 |
+
- ROCm env-var tuning (`HSA_*`, `MIOPEN_*`, `NCCL_MIN_NCHANNELS=112`)
|
| 69 |
+
- hipBLASLt hint logging + offline tuning files; MIOpen `MIOPEN_FIND_*` autotune
|
| 70 |
+
- Flash-attn ROCm fork (validated on MI300 via Optimum-AMD)
|
| 71 |
+
- RCCL topology + tensor-parallel placement within a single XGMI island
|
| 72 |
+
- bitsandbytes-on-ROCm gotchas (not officially supported — recommend alternatives)
|
| 73 |
+
- bf16 vs fp16 on CDNA3 matrix cores
|
| 74 |
+
- One-process-per-GPU vs one-process-many-GPUs (ROCm serializes launches in the latter)
|
| 75 |
+
- NUMA auto-balancing disable for stable benchmarks
|
| 76 |
+
- MI300X-specific: 304 CUs, 192 GB HBM3, ~5.3 TB/s peak bandwidth, native FP8
|
| 77 |
+
|
| 78 |
+
This is the kind of insight you only get from someone who has actually shipped on MI300X. We bottle it.
|
| 79 |
+
|
| 80 |
+
## Demo Narrative (3 minutes)
|
| 81 |
+
|
| 82 |
+
**Setup (15s):** "Most teams waste 50%+ of their MI300X. Watch GPU Goblin find that waste live."
|
| 83 |
+
|
| 84 |
+
**Live demo (2 min):**
|
| 85 |
+
|
| 86 |
+
1. Drop in `train_qwen_lora.py` — batch=4, fp16, naive attention. (Looks normal.)
|
| 87 |
+
2. Agent: *"Parsing config…"* → extracts hyperparams.
|
| 88 |
+
3. Agent: *"Running 10-step profile on MI300X…"* → shows real metrics: HBM 38%, MFU 24%.
|
| 89 |
+
4. Agent: *"Querying ROCm playbook…"* → 4 issues found, each with citation.
|
| 90 |
+
5. Agent: *"Generating optimized config…"* → diff appears: bf16, batch=12, flash-attn ROCm, hipBLASLt env, packed sequences.
|
| 91 |
+
6. Agent: *"Benchmarking new config — 50 steps on MI300X…"* → live progress.
|
| 92 |
+
7. Final report: **142 → 318 tokens/sec (2.24×). MFU 24% → 51%. $X saved per epoch.**
|
| 93 |
+
|
| 94 |
+
**Why it works (45s):** Show the agent loop. Show the KB. Land the line: *"the agent runs Qwen on AMD silicon, the audit target is Qwen on AMD silicon, the optimizations are AMD-specific. End-to-end AMD."*
|
| 95 |
+
|
| 96 |
+
## What This Is *Not*
|
| 97 |
+
|
| 98 |
+
- Not a generic PyTorch profiler GUI (`rocprofv3` + tensorboard already exist)
|
| 99 |
+
- Not a chatbot wrapping a static FAQ (the agent runs real benchmarks)
|
| 100 |
+
- Not a multi-cloud cost tool (focused on MI300X compute waste, not infra cost)
|
| 101 |
+
- Not a fine-tuning trainer itself (we *audit* others' runs)
|
| 102 |
+
|
| 103 |
+
## One-Line Pitch
|
| 104 |
+
|
| 105 |
+
> **GPU Goblin is the AMD performance engineer you wish was on your team — except it costs five minutes and audits any fine-tuning run on MI300X.**
|
pyproject.toml
ADDED
|
@@ -0,0 +1,43 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[project]
|
| 2 |
+
name = "gpu-goblin"
|
| 3 |
+
version = "0.1.0"
|
| 4 |
+
description = "An AI agent that hunts wasted compute on AMD MI300X"
|
| 5 |
+
requires-python = ">=3.10"
|
| 6 |
+
dependencies = [
|
| 7 |
+
"fastapi>=0.110",
|
| 8 |
+
"uvicorn[standard]>=0.27",
|
| 9 |
+
"python-multipart>=0.0.9",
|
| 10 |
+
"huggingface_hub>=0.28",
|
| 11 |
+
"pydantic>=2.6",
|
| 12 |
+
"pyyaml>=6.0",
|
| 13 |
+
"sentence-transformers>=2.7",
|
| 14 |
+
"numpy>=1.26",
|
| 15 |
+
"sse-starlette>=2.0",
|
| 16 |
+
"streamlit>=1.32",
|
| 17 |
+
"altair>=5.2",
|
| 18 |
+
"pandas>=2.2",
|
| 19 |
+
"requests>=2.31",
|
| 20 |
+
"openai>=1.30",
|
| 21 |
+
]
|
| 22 |
+
|
| 23 |
+
[project.optional-dependencies]
|
| 24 |
+
dev = [
|
| 25 |
+
"pytest>=8.0",
|
| 26 |
+
"pytest-asyncio>=0.23",
|
| 27 |
+
"ruff>=0.4",
|
| 28 |
+
]
|
| 29 |
+
|
| 30 |
+
[build-system]
|
| 31 |
+
requires = ["setuptools>=68"]
|
| 32 |
+
build-backend = "setuptools.build_meta"
|
| 33 |
+
|
| 34 |
+
[tool.setuptools.packages.find]
|
| 35 |
+
include = ["agent*", "runner*", "ui*"]
|
| 36 |
+
|
| 37 |
+
[tool.ruff]
|
| 38 |
+
line-length = 100
|
| 39 |
+
target-version = "py310"
|
| 40 |
+
|
| 41 |
+
[tool.ruff.lint]
|
| 42 |
+
select = ["E", "F", "W", "I", "B", "UP"]
|
| 43 |
+
ignore = ["E501"]
|
runner/goblin_runner.sh
ADDED
|
@@ -0,0 +1,169 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
# goblin_runner.sh — wrap the user's training command with rocprofv3 + amd-smi.
|
| 3 |
+
#
|
| 4 |
+
# Architecture: brainstorming/architecture.md §5 (Profiling Pipeline).
|
| 5 |
+
#
|
| 6 |
+
# Inputs (env vars):
|
| 7 |
+
# USER_SCRIPT Path to the workload python file. Required.
|
| 8 |
+
# OUT_DIR Directory to write trace.csv / torch_profile.json /
|
| 9 |
+
# amd_smi.csv. Required. Created if missing.
|
| 10 |
+
# STEPS --max_steps argument forwarded to the user script. Default 10.
|
| 11 |
+
# GOBLIN_GPU_ID ROCR_VISIBLE_DEVICES value. Default 0.
|
| 12 |
+
#
|
| 13 |
+
# Outputs (in $OUT_DIR):
|
| 14 |
+
# trace.csv rocprofv3 kernel trace
|
| 15 |
+
# torch_profile.json torch.profiler chrome trace (the user script writes this)
|
| 16 |
+
# amd_smi.csv amd-smi telemetry sampled at 200 ms
|
| 17 |
+
# stdout.log user-script stdout
|
| 18 |
+
# stderr.log user-script stderr
|
| 19 |
+
#
|
| 20 |
+
# Failure mode: any non-zero rocprofv3 exit short-circuits the script. On
|
| 21 |
+
# failure we dump the captured stdout/stderr logs to THIS script's own stderr
|
| 22 |
+
# so `subprocess.run(capture_output=True)` in LiveRunner sees the real error
|
| 23 |
+
# (not just an empty `[]` tail). LiveRunner then archives the whole OUT_DIR
|
| 24 |
+
# under bench_cache/last_runner_failure_<ts>/ so you can inspect after-the-fact.
|
| 25 |
+
|
| 26 |
+
set -uo pipefail
|
| 27 |
+
|
| 28 |
+
: "${USER_SCRIPT:?USER_SCRIPT env var is required}"
|
| 29 |
+
: "${OUT_DIR:?OUT_DIR env var is required}"
|
| 30 |
+
STEPS="${STEPS:-10}"
|
| 31 |
+
|
| 32 |
+
mkdir -p "$OUT_DIR"
|
| 33 |
+
|
| 34 |
+
# Pin to a single MI300X so concurrent benchmark runs don't fight.
|
| 35 |
+
export ROCR_VISIBLE_DEVICES="${GOBLIN_GPU_ID:-0}"
|
| 36 |
+
|
| 37 |
+
# Background HBM/power telemetry. Different amd-smi versions ship slightly
|
| 38 |
+
# different flag names (--interval / --watch / no flag at all in older
|
| 39 |
+
# builds), so try a few variants and gracefully degrade. Telemetry is
|
| 40 |
+
# optional — profile_parser tolerates a missing amd_smi.csv.
|
| 41 |
+
start_amd_smi() {
|
| 42 |
+
if ! command -v amd-smi >/dev/null 2>&1; then
|
| 43 |
+
echo "amd-smi not on PATH; skipping telemetry sidecar" \
|
| 44 |
+
> "$OUT_DIR/amd_smi.err"
|
| 45 |
+
return 1
|
| 46 |
+
fi
|
| 47 |
+
# amd-smi telemetry surface drifted hard across rocm versions. We try
|
| 48 |
+
# subcommands in this order, newest-first; first one that survives 0.5s
|
| 49 |
+
# is kept:
|
| 50 |
+
#
|
| 51 |
+
# 1. `amd-smi metric --watch <s> --mem-usage --usage --csv` (ROCm 7.x).
|
| 52 |
+
# Different code path from `monitor`, so it dodges the known
|
| 53 |
+
# `AttributeError: 'Namespace' object has no attribute 'violation'`
|
| 54 |
+
# crash that some ROCm 7.x point releases ship inside the `monitor`
|
| 55 |
+
# subcommand. Produces VRAM_USED_MB / GFX_ACTIVITY / MEM_ACTIVITY
|
| 56 |
+
# columns that profile_parser._pick_column already recognises.
|
| 57 |
+
#
|
| 58 |
+
# 2. `amd-smi monitor --watch <s> ...` (ROCm 6.x and the 7.x builds
|
| 59 |
+
# where `monitor` actually works). Kept as a fallback for older
|
| 60 |
+
# installs that may not have the `metric --watch` form.
|
| 61 |
+
#
|
| 62 |
+
# 3. `amd-smi monitor --interval <s>` (pre-6.0) and finally bare
|
| 63 |
+
# `amd-smi monitor` (very old / implicit-all).
|
| 64 |
+
local variants=(
|
| 65 |
+
# ROCm 7.x: prefer the metric subcommand — bypasses the monitor bug.
|
| 66 |
+
"amd-smi metric --watch 1 --mem-usage --usage --csv"
|
| 67 |
+
"amd-smi metric -w 1 -m -u --csv"
|
| 68 |
+
# ROCm 7.x monitor (works on builds without the violation-attribute bug)
|
| 69 |
+
"amd-smi monitor --watch 1 --power-usage --gfx --mem --vram-usage --csv"
|
| 70 |
+
"amd-smi monitor -w 1 -p -u -m -v --csv"
|
| 71 |
+
# ROCm 6.x intermediate forms
|
| 72 |
+
"amd-smi monitor --watch 1 --csv"
|
| 73 |
+
"amd-smi monitor --watch 1"
|
| 74 |
+
# Older / fallback
|
| 75 |
+
"amd-smi monitor --interval 1 --csv"
|
| 76 |
+
"amd-smi monitor --interval 1"
|
| 77 |
+
"amd-smi monitor --csv"
|
| 78 |
+
"amd-smi monitor"
|
| 79 |
+
)
|
| 80 |
+
for cmd in "${variants[@]}"; do
|
| 81 |
+
# shellcheck disable=SC2086
|
| 82 |
+
$cmd > "$OUT_DIR/amd_smi.csv" 2> "$OUT_DIR/amd_smi.err" &
|
| 83 |
+
local pid=$!
|
| 84 |
+
sleep 0.5
|
| 85 |
+
if kill -0 "$pid" 2>/dev/null; then
|
| 86 |
+
AMD_SMI_PID=$pid
|
| 87 |
+
return 0
|
| 88 |
+
fi
|
| 89 |
+
done
|
| 90 |
+
# All variants failed. Telemetry is optional; mark it as deliberately
|
| 91 |
+
# skipped and let the main run proceed — profile_parser tolerates a
|
| 92 |
+
# missing/empty amd_smi.csv.
|
| 93 |
+
echo "amd-smi monitor: no compatible flag set in this build; telemetry skipped" \
|
| 94 |
+
> "$OUT_DIR/amd_smi.skipped"
|
| 95 |
+
rm -f "$OUT_DIR/amd_smi.csv"
|
| 96 |
+
return 1
|
| 97 |
+
}
|
| 98 |
+
|
| 99 |
+
AMD_SMI_PID=
|
| 100 |
+
start_amd_smi || true # never block the main run on telemetry
|
| 101 |
+
|
| 102 |
+
cleanup() {
|
| 103 |
+
if [[ -n "${AMD_SMI_PID:-}" ]] && kill -0 "$AMD_SMI_PID" 2>/dev/null; then
|
| 104 |
+
kill "$AMD_SMI_PID" 2>/dev/null || true
|
| 105 |
+
wait "$AMD_SMI_PID" 2>/dev/null || true
|
| 106 |
+
fi
|
| 107 |
+
}
|
| 108 |
+
trap cleanup EXIT
|
| 109 |
+
|
| 110 |
+
# Dump the captured logs to *our* stderr on a non-zero rocprofv3 exit so the
|
| 111 |
+
# Python subprocess that spawned us actually sees the real error message.
|
| 112 |
+
# Without this the redirected stdout.log / stderr.log live inside the tempdir
|
| 113 |
+
# only and LiveRunner's stderr-tail check sees nothing.
|
| 114 |
+
dump_failure_logs() {
|
| 115 |
+
local code=$?
|
| 116 |
+
if [[ $code -ne 0 ]]; then
|
| 117 |
+
{
|
| 118 |
+
echo "=== goblin_runner.sh failed with exit code $code ==="
|
| 119 |
+
echo "=== USER_SCRIPT: $USER_SCRIPT ==="
|
| 120 |
+
echo "=== OUT_DIR: $OUT_DIR ==="
|
| 121 |
+
echo "=== ROCR_VISIBLE_DEVICES: ${ROCR_VISIBLE_DEVICES:-unset} ==="
|
| 122 |
+
echo
|
| 123 |
+
echo "=== last 50 lines of $OUT_DIR/stdout.log ==="
|
| 124 |
+
tail -n 50 "$OUT_DIR/stdout.log" 2>/dev/null || echo "(stdout.log missing)"
|
| 125 |
+
echo
|
| 126 |
+
echo "=== last 50 lines of $OUT_DIR/stderr.log ==="
|
| 127 |
+
tail -n 50 "$OUT_DIR/stderr.log" 2>/dev/null || echo "(stderr.log missing)"
|
| 128 |
+
echo
|
| 129 |
+
echo "=== last 20 lines of $OUT_DIR/amd_smi.err ==="
|
| 130 |
+
tail -n 20 "$OUT_DIR/amd_smi.err" 2>/dev/null || echo "(amd_smi.err missing)"
|
| 131 |
+
} 1>&2
|
| 132 |
+
fi
|
| 133 |
+
return $code
|
| 134 |
+
}
|
| 135 |
+
|
| 136 |
+
# rocprofv3 collects HSA + kernel traces. The user script is responsible for
|
| 137 |
+
# writing torch_profile.json (the agent injects torch.profiler around the
|
| 138 |
+
# training loop in Phase 3). --output-format csv keeps parsing simple.
|
| 139 |
+
set +e
|
| 140 |
+
rocprofv3 \
|
| 141 |
+
--hsa-trace --kernel-trace \
|
| 142 |
+
--output-directory "$OUT_DIR" \
|
| 143 |
+
--output-file trace \
|
| 144 |
+
--output-format csv \
|
| 145 |
+
-- \
|
| 146 |
+
python "$USER_SCRIPT" \
|
| 147 |
+
--max_steps="$STEPS" \
|
| 148 |
+
--torch_profile_out="$OUT_DIR/torch_profile.json" \
|
| 149 |
+
> "$OUT_DIR/stdout.log" 2> "$OUT_DIR/stderr.log"
|
| 150 |
+
ROCPROF_EXIT=$?
|
| 151 |
+
set -e
|
| 152 |
+
|
| 153 |
+
if [[ $ROCPROF_EXIT -ne 0 ]]; then
|
| 154 |
+
(exit $ROCPROF_EXIT) || dump_failure_logs
|
| 155 |
+
exit $ROCPROF_EXIT
|
| 156 |
+
fi
|
| 157 |
+
|
| 158 |
+
# rocprofv3 may write trace_kernel_trace.csv etc. — normalize to trace.csv so
|
| 159 |
+
# profile_parser has one stable filename to look for.
|
| 160 |
+
if [[ ! -f "$OUT_DIR/trace.csv" ]]; then
|
| 161 |
+
for candidate in "$OUT_DIR"/trace*kernel*.csv "$OUT_DIR"/*kernel_trace.csv; do
|
| 162 |
+
if [[ -f "$candidate" ]]; then
|
| 163 |
+
cp "$candidate" "$OUT_DIR/trace.csv"
|
| 164 |
+
break
|
| 165 |
+
fi
|
| 166 |
+
done
|
| 167 |
+
fi
|
| 168 |
+
|
| 169 |
+
exit 0
|
runner/profile_parser.py
ADDED
|
@@ -0,0 +1,693 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""profile_parser — turn goblin_runner.sh artefacts into RunMetrics.
|
| 2 |
+
|
| 3 |
+
Reads three files written by goblin_runner.sh (architecture.md §5):
|
| 4 |
+
|
| 5 |
+
<out_dir>/trace.csv rocprofv3 kernel trace
|
| 6 |
+
<out_dir>/torch_profile.json torch.profiler chrome trace
|
| 7 |
+
<out_dir>/amd_smi.csv amd-smi telemetry, ~200 ms cadence
|
| 8 |
+
|
| 9 |
+
and produces one `RunMetrics` with a populated `WasteBudget`.
|
| 10 |
+
|
| 11 |
+
Each waste-budget bucket is computed with a deliberately simple, documented
|
| 12 |
+
heuristic — these are best-effort signals for the agent, NOT measured
|
| 13 |
+
ground-truth. See the docstring on `_waste_budget()` for the per-bucket
|
| 14 |
+
formula and its known failure modes.
|
| 15 |
+
|
| 16 |
+
This module is import-tolerant on machines without the rocprofv3 stack —
|
| 17 |
+
it only reads files. Missing or unparseable files degrade individual
|
| 18 |
+
metrics to zero and append a warning to RunMetrics.warnings rather than
|
| 19 |
+
raising. LiveRunner ultimately decides what to do with parse failures.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
from __future__ import annotations
|
| 23 |
+
|
| 24 |
+
import csv
|
| 25 |
+
import json
|
| 26 |
+
import logging
|
| 27 |
+
import re
|
| 28 |
+
from dataclasses import dataclass
|
| 29 |
+
from pathlib import Path
|
| 30 |
+
from typing import Any
|
| 31 |
+
|
| 32 |
+
from agent.schemas import KernelEntry, RunMetrics, WasteBudget, WorkloadConfig
|
| 33 |
+
|
| 34 |
+
_LOG = logging.getLogger(__name__)
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
# MI300X has 192 GB HBM3. We treat sustained <70% utilisation as headroom.
|
| 38 |
+
_HBM_TOTAL_GB = 192.0
|
| 39 |
+
_HBM_HEALTHY_TARGET = 0.70
|
| 40 |
+
|
| 41 |
+
# Kernel-name patterns used by the heuristics.
|
| 42 |
+
#
|
| 43 |
+
# We use `(?<![A-Za-z0-9])` / `(?![A-Za-z0-9])` instead of `\b` so that
|
| 44 |
+
# underscores act as token separators — kernel names like
|
| 45 |
+
# `rccl_AllReduce` and `hipBLASLt_generic_gemm` should match.
|
| 46 |
+
_BOUND_L = r"(?<![A-Za-z0-9])"
|
| 47 |
+
_BOUND_R = r"(?![A-Za-z0-9])"
|
| 48 |
+
_RCCL_PATTERN = re.compile(
|
| 49 |
+
_BOUND_L + r"(rccl|nccl|all[_-]?reduce|broadcast|reduce[_-]?scatter)" + _BOUND_R, re.I
|
| 50 |
+
)
|
| 51 |
+
_GEMM_PATTERN = re.compile(_BOUND_L + r"(gemm|matmul|hgemm|sgemm|hipblaslt)" + _BOUND_R, re.I)
|
| 52 |
+
_GENERIC_GEMM_PATTERN = re.compile(
|
| 53 |
+
_BOUND_L + r"(generic|fallback|naive|reference)" + _BOUND_R, re.I
|
| 54 |
+
)
|
| 55 |
+
_FP16_PATTERN = re.compile(_BOUND_L + r"(fp16|half|f16)" + _BOUND_R, re.I)
|
| 56 |
+
_BF16_PATTERN = re.compile(_BOUND_L + r"(bf16|bfloat16)" + _BOUND_R, re.I)
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
# ---------------------------------------------------------------------------
|
| 60 |
+
# Public entry point
|
| 61 |
+
# ---------------------------------------------------------------------------
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def parse(
|
| 65 |
+
out_dir: Path,
|
| 66 |
+
config: WorkloadConfig | None = None,
|
| 67 |
+
steps: int = 10,
|
| 68 |
+
) -> RunMetrics:
|
| 69 |
+
"""Build a `RunMetrics` from goblin_runner.sh artefacts in `out_dir`.
|
| 70 |
+
|
| 71 |
+
`config` is used by some heuristics (e.g. precision_path skips the
|
| 72 |
+
estimate when the config is already bf16). It can be None at the cost
|
| 73 |
+
of a more conservative waste-budget estimate.
|
| 74 |
+
|
| 75 |
+
Always returns a RunMetrics object — individual metrics degrade to
|
| 76 |
+
zero with a warning rather than raising, because LiveRunner ultimately
|
| 77 |
+
decides whether parse failures should trigger fallback.
|
| 78 |
+
"""
|
| 79 |
+
warnings: list[str] = []
|
| 80 |
+
|
| 81 |
+
kernels = _read_kernels(out_dir / "trace.csv", warnings)
|
| 82 |
+
torch_summary = _read_torch_profile(out_dir / "torch_profile.json", warnings)
|
| 83 |
+
smi = _read_amd_smi(out_dir / "amd_smi.csv", warnings)
|
| 84 |
+
|
| 85 |
+
top_kernels = _top_kernels(kernels, top_n=5)
|
| 86 |
+
gpu_util_pct = smi.gpu_util_pct if smi.gpu_util_pct is not None else _gpu_util_from_kernels(
|
| 87 |
+
kernels, torch_summary
|
| 88 |
+
)
|
| 89 |
+
# Write the resolved gpu_util back into smi so _waste_budget sees the
|
| 90 |
+
# same number RunMetrics reports. Without this, a build where amd-smi's
|
| 91 |
+
# `usage` column returns N/A leaves smi.gpu_util_pct=None — and
|
| 92 |
+
# _waste_budget reads it as 0%, which forces host_gap to consume the
|
| 93 |
+
# full step time and collapses kernel_shape / precision_path to zero
|
| 94 |
+
# (both multiplied by gpu_util).
|
| 95 |
+
smi.gpu_util_pct = gpu_util_pct
|
| 96 |
+
|
| 97 |
+
waste_budget = _waste_budget(
|
| 98 |
+
kernels=kernels,
|
| 99 |
+
torch_summary=torch_summary,
|
| 100 |
+
smi=smi,
|
| 101 |
+
config=config,
|
| 102 |
+
)
|
| 103 |
+
|
| 104 |
+
attention_kernel_loaded = _detect_attention_kernel(kernels)
|
| 105 |
+
|
| 106 |
+
return RunMetrics(
|
| 107 |
+
steps=steps,
|
| 108 |
+
tokens_per_sec=torch_summary.tokens_per_sec or 0.0,
|
| 109 |
+
mfu_pct=torch_summary.mfu_pct or 0.0,
|
| 110 |
+
hbm_peak_gb=smi.hbm_peak_gb or 0.0,
|
| 111 |
+
hbm_avg_gb=smi.hbm_avg_gb or 0.0,
|
| 112 |
+
gpu_util_pct=gpu_util_pct,
|
| 113 |
+
top_kernels=top_kernels,
|
| 114 |
+
attention_kernel_loaded=attention_kernel_loaded,
|
| 115 |
+
waste_budget=waste_budget,
|
| 116 |
+
warnings=warnings,
|
| 117 |
+
rocm_version=smi.rocm_version or "unknown",
|
| 118 |
+
pytorch_version=torch_summary.pytorch_version or "unknown",
|
| 119 |
+
runner_kind="live",
|
| 120 |
+
)
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
# ---------------------------------------------------------------------------
|
| 124 |
+
# Internal data carriers
|
| 125 |
+
# ---------------------------------------------------------------------------
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
@dataclass
|
| 129 |
+
class _Kernel:
|
| 130 |
+
name: str
|
| 131 |
+
duration_ns: int
|
| 132 |
+
"""Kernel duration in nanoseconds."""
|
| 133 |
+
|
| 134 |
+
is_collective: bool = False
|
| 135 |
+
is_gemm: bool = False
|
| 136 |
+
is_generic_gemm: bool = False
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
@dataclass
|
| 140 |
+
class _TorchSummary:
|
| 141 |
+
tokens_per_sec: float | None = None
|
| 142 |
+
mfu_pct: float | None = None
|
| 143 |
+
pytorch_version: str | None = None
|
| 144 |
+
step_time_seconds: float | None = None
|
| 145 |
+
"""Average wall-clock seconds per training step (used by waste budget)."""
|
| 146 |
+
|
| 147 |
+
host_busy_fraction: float | None = None
|
| 148 |
+
"""Fraction of step time the host (CPU) spent doing non-launch work.
|
| 149 |
+
Heuristic: `cpu_self_time / step_time` reported by torch.profiler."""
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
@dataclass
|
| 153 |
+
class _SmiSummary:
|
| 154 |
+
hbm_peak_gb: float | None = None
|
| 155 |
+
hbm_avg_gb: float | None = None
|
| 156 |
+
gpu_util_pct: float | None = None
|
| 157 |
+
rocm_version: str | None = None
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
# ---------------------------------------------------------------------------
|
| 161 |
+
# trace.csv (rocprofv3) → list[_Kernel]
|
| 162 |
+
# ---------------------------------------------------------------------------
|
| 163 |
+
|
| 164 |
+
|
| 165 |
+
def _read_kernels(path: Path, warnings: list[str]) -> list[_Kernel]:
|
| 166 |
+
"""Parse rocprofv3 kernel-trace CSV into a list of `_Kernel` records.
|
| 167 |
+
|
| 168 |
+
rocprofv3 column names vary slightly by version. We look up by header and
|
| 169 |
+
accept the common aliases; if the file is missing or unparseable we
|
| 170 |
+
return an empty list and append a warning so the caller can decide.
|
| 171 |
+
"""
|
| 172 |
+
if not path.exists():
|
| 173 |
+
warnings.append(f"profile_parser: kernel trace not found at {path}")
|
| 174 |
+
return []
|
| 175 |
+
|
| 176 |
+
try:
|
| 177 |
+
with path.open(newline="") as f:
|
| 178 |
+
reader = csv.DictReader(f)
|
| 179 |
+
if reader.fieldnames is None:
|
| 180 |
+
warnings.append(f"profile_parser: empty kernel trace at {path}")
|
| 181 |
+
return []
|
| 182 |
+
name_col = _pick_column(reader.fieldnames, ["KernelName", "Kernel_Name", "kernel_name", "Name"])
|
| 183 |
+
start_col = _pick_column(reader.fieldnames, ["BeginNs", "start_ns", "BeginNS", "Start"])
|
| 184 |
+
end_col = _pick_column(reader.fieldnames, ["EndNs", "end_ns", "EndNS", "End"])
|
| 185 |
+
duration_col = _pick_column(reader.fieldnames, ["DurationNs", "duration_ns", "Duration"])
|
| 186 |
+
|
| 187 |
+
kernels: list[_Kernel] = []
|
| 188 |
+
for row in reader:
|
| 189 |
+
name = (row.get(name_col) or "").strip() if name_col else ""
|
| 190 |
+
if not name:
|
| 191 |
+
continue
|
| 192 |
+
duration = _row_duration_ns(row, duration_col, start_col, end_col)
|
| 193 |
+
if duration <= 0:
|
| 194 |
+
continue
|
| 195 |
+
kernels.append(
|
| 196 |
+
_Kernel(
|
| 197 |
+
name=name,
|
| 198 |
+
duration_ns=duration,
|
| 199 |
+
is_collective=bool(_RCCL_PATTERN.search(name)),
|
| 200 |
+
is_gemm=bool(_GEMM_PATTERN.search(name)),
|
| 201 |
+
is_generic_gemm=bool(
|
| 202 |
+
_GEMM_PATTERN.search(name) and _GENERIC_GEMM_PATTERN.search(name)
|
| 203 |
+
),
|
| 204 |
+
)
|
| 205 |
+
)
|
| 206 |
+
return kernels
|
| 207 |
+
except (OSError, csv.Error) as exc:
|
| 208 |
+
warnings.append(f"profile_parser: failed to read kernel trace ({exc})")
|
| 209 |
+
return []
|
| 210 |
+
|
| 211 |
+
|
| 212 |
+
def _pick_column(fieldnames: list[str], candidates: list[str]) -> str | None:
|
| 213 |
+
"""Pick the first matching column name, with three fallback tiers:
|
| 214 |
+
1. Exact match.
|
| 215 |
+
2. Case-insensitive exact match.
|
| 216 |
+
3. Substring match (case-insensitive) — every token of any candidate
|
| 217 |
+
must appear somewhere in the field name. Tolerates the column-name
|
| 218 |
+
drift between rocprofv3 / amd-smi versions (e.g. `VRAM_USED` vs
|
| 219 |
+
`vram_used_mb` vs `VRAM USED MB`).
|
| 220 |
+
"""
|
| 221 |
+
for c in candidates:
|
| 222 |
+
if c in fieldnames:
|
| 223 |
+
return c
|
| 224 |
+
lower = {f.lower(): f for f in fieldnames}
|
| 225 |
+
for c in candidates:
|
| 226 |
+
if c.lower() in lower:
|
| 227 |
+
return lower[c.lower()]
|
| 228 |
+
# Substring tier: split each candidate on _/space, require all tokens
|
| 229 |
+
# appear in the (lowercased) field name. Avoids matching too eagerly
|
| 230 |
+
# by requiring every token of the candidate.
|
| 231 |
+
for c in candidates:
|
| 232 |
+
tokens = [t for t in c.lower().replace("_", " ").split() if t]
|
| 233 |
+
if not tokens:
|
| 234 |
+
continue
|
| 235 |
+
for fname in fieldnames:
|
| 236 |
+
fl = fname.lower()
|
| 237 |
+
if all(t in fl for t in tokens):
|
| 238 |
+
return fname
|
| 239 |
+
return None
|
| 240 |
+
|
| 241 |
+
|
| 242 |
+
def _row_duration_ns(
|
| 243 |
+
row: dict[str, str], duration_col: str | None, start_col: str | None, end_col: str | None
|
| 244 |
+
) -> int:
|
| 245 |
+
if duration_col and row.get(duration_col):
|
| 246 |
+
try:
|
| 247 |
+
return int(float(row[duration_col]))
|
| 248 |
+
except ValueError:
|
| 249 |
+
return 0
|
| 250 |
+
if start_col and end_col and row.get(start_col) and row.get(end_col):
|
| 251 |
+
try:
|
| 252 |
+
return int(float(row[end_col])) - int(float(row[start_col]))
|
| 253 |
+
except ValueError:
|
| 254 |
+
return 0
|
| 255 |
+
return 0
|
| 256 |
+
|
| 257 |
+
|
| 258 |
+
def _top_kernels(kernels: list[_Kernel], top_n: int) -> list[KernelEntry]:
|
| 259 |
+
if not kernels:
|
| 260 |
+
return []
|
| 261 |
+
total = sum(k.duration_ns for k in kernels) or 1
|
| 262 |
+
by_name: dict[str, int] = {}
|
| 263 |
+
for k in kernels:
|
| 264 |
+
by_name[k.name] = by_name.get(k.name, 0) + k.duration_ns
|
| 265 |
+
ranked = sorted(by_name.items(), key=lambda kv: kv[1], reverse=True)[:top_n]
|
| 266 |
+
return [KernelEntry(name=name, pct_time=ns / total * 100.0) for name, ns in ranked]
|
| 267 |
+
|
| 268 |
+
|
| 269 |
+
def _detect_attention_kernel(kernels: list[_Kernel]) -> str:
|
| 270 |
+
for k in kernels:
|
| 271 |
+
n = k.name.lower()
|
| 272 |
+
if "flash_attn" in n and "rocm" in n:
|
| 273 |
+
return "flash_rocm"
|
| 274 |
+
if "flash" in n and "attn" in n:
|
| 275 |
+
return "flash"
|
| 276 |
+
if "scaled_dot_product_attention" in n:
|
| 277 |
+
return "sdpa"
|
| 278 |
+
if kernels:
|
| 279 |
+
return "eager" # nothing flash-shaped; default conservative label
|
| 280 |
+
return "unknown"
|
| 281 |
+
|
| 282 |
+
|
| 283 |
+
def _gpu_util_from_kernels(kernels: list[_Kernel], torch_summary: _TorchSummary) -> float:
|
| 284 |
+
"""Fallback GPU util when amd-smi is missing.
|
| 285 |
+
|
| 286 |
+
util ≈ sum(kernel duration) / total wall-clock step time.
|
| 287 |
+
"""
|
| 288 |
+
if not kernels or torch_summary.step_time_seconds in (None, 0):
|
| 289 |
+
return 0.0
|
| 290 |
+
total_kernel_ns = sum(k.duration_ns for k in kernels)
|
| 291 |
+
wall_ns = torch_summary.step_time_seconds * 1e9
|
| 292 |
+
if wall_ns <= 0:
|
| 293 |
+
return 0.0
|
| 294 |
+
return min(100.0, total_kernel_ns / wall_ns * 100.0)
|
| 295 |
+
|
| 296 |
+
|
| 297 |
+
# ---------------------------------------------------------------------------
|
| 298 |
+
# torch_profile.json (torch.profiler chrome trace) → _TorchSummary
|
| 299 |
+
# ---------------------------------------------------------------------------
|
| 300 |
+
|
| 301 |
+
|
| 302 |
+
def _read_torch_profile(path: Path, warnings: list[str]) -> _TorchSummary:
|
| 303 |
+
"""Pull tokens/sec, MFU, and step timing from a torch.profiler artefact.
|
| 304 |
+
|
| 305 |
+
The user script (workloads/train_qwen_lora.py in Phase 3) is
|
| 306 |
+
responsible for embedding `tokens_per_sec`, `mfu_pct`, `pytorch_version`
|
| 307 |
+
and `step_time_seconds` in the trace as `metadata` events. If those are
|
| 308 |
+
missing, we estimate `step_time_seconds` from the total trace duration.
|
| 309 |
+
"""
|
| 310 |
+
summary = _TorchSummary()
|
| 311 |
+
if not path.exists():
|
| 312 |
+
warnings.append(f"profile_parser: torch profile not found at {path}")
|
| 313 |
+
return summary
|
| 314 |
+
try:
|
| 315 |
+
data = json.loads(path.read_text())
|
| 316 |
+
except (OSError, json.JSONDecodeError) as exc:
|
| 317 |
+
warnings.append(f"profile_parser: failed to read torch profile ({exc})")
|
| 318 |
+
return summary
|
| 319 |
+
|
| 320 |
+
# torch.profiler chrome trace is `{"traceEvents": [...], "metadata": {...}}`
|
| 321 |
+
metadata = data.get("metadata") if isinstance(data, dict) else None
|
| 322 |
+
if isinstance(metadata, dict):
|
| 323 |
+
summary.tokens_per_sec = _coerce_float(metadata.get("tokens_per_sec"))
|
| 324 |
+
summary.mfu_pct = _coerce_float(metadata.get("mfu_pct"))
|
| 325 |
+
summary.pytorch_version = metadata.get("pytorch_version") or metadata.get("torch_version")
|
| 326 |
+
summary.step_time_seconds = _coerce_float(metadata.get("step_time_seconds"))
|
| 327 |
+
summary.host_busy_fraction = _coerce_float(metadata.get("host_busy_fraction"))
|
| 328 |
+
|
| 329 |
+
events = data.get("traceEvents") if isinstance(data, dict) else None
|
| 330 |
+
if isinstance(events, list):
|
| 331 |
+
if summary.step_time_seconds is None:
|
| 332 |
+
summary.step_time_seconds = _step_time_from_events(events)
|
| 333 |
+
if summary.host_busy_fraction is None:
|
| 334 |
+
summary.host_busy_fraction = _host_busy_from_events(events)
|
| 335 |
+
|
| 336 |
+
return summary
|
| 337 |
+
|
| 338 |
+
|
| 339 |
+
def _coerce_float(v: Any) -> float | None:
|
| 340 |
+
if v is None:
|
| 341 |
+
return None
|
| 342 |
+
try:
|
| 343 |
+
return float(v)
|
| 344 |
+
except (TypeError, ValueError):
|
| 345 |
+
return None
|
| 346 |
+
|
| 347 |
+
|
| 348 |
+
def _step_time_from_events(events: list[dict]) -> float | None:
|
| 349 |
+
"""Estimate per-step wall-clock seconds from chrome-trace duration events.
|
| 350 |
+
|
| 351 |
+
Looks for `name == "ProfilerStep#*"` complete events; falls back to the
|
| 352 |
+
overall trace span if those aren't present.
|
| 353 |
+
"""
|
| 354 |
+
durations: list[float] = []
|
| 355 |
+
overall_start: float | None = None
|
| 356 |
+
overall_end: float | None = None
|
| 357 |
+
for ev in events:
|
| 358 |
+
if not isinstance(ev, dict):
|
| 359 |
+
continue
|
| 360 |
+
name = ev.get("name", "")
|
| 361 |
+
ts = ev.get("ts")
|
| 362 |
+
dur = ev.get("dur")
|
| 363 |
+
if isinstance(name, str) and name.startswith("ProfilerStep") and isinstance(dur, (int, float)):
|
| 364 |
+
durations.append(float(dur) / 1e6) # us → seconds
|
| 365 |
+
if isinstance(ts, (int, float)) and isinstance(dur, (int, float)):
|
| 366 |
+
start = float(ts)
|
| 367 |
+
end = start + float(dur)
|
| 368 |
+
overall_start = start if overall_start is None else min(overall_start, start)
|
| 369 |
+
overall_end = end if overall_end is None else max(overall_end, end)
|
| 370 |
+
if durations:
|
| 371 |
+
return sum(durations) / len(durations)
|
| 372 |
+
if overall_start is not None and overall_end is not None and overall_end > overall_start:
|
| 373 |
+
return (overall_end - overall_start) / 1e6
|
| 374 |
+
return None
|
| 375 |
+
|
| 376 |
+
|
| 377 |
+
def _host_busy_from_events(events: list[dict]) -> float | None:
|
| 378 |
+
"""Heuristic: cpu_op event time / total event span.
|
| 379 |
+
|
| 380 |
+
Used by the data_wait waste bucket to disambiguate "GPU idle because the
|
| 381 |
+
host is busy preparing the next batch" from "GPU idle because nothing
|
| 382 |
+
is running anywhere".
|
| 383 |
+
"""
|
| 384 |
+
cpu_op_us = 0.0
|
| 385 |
+
span_min: float | None = None
|
| 386 |
+
span_max: float | None = None
|
| 387 |
+
for ev in events:
|
| 388 |
+
if not isinstance(ev, dict):
|
| 389 |
+
continue
|
| 390 |
+
cat = ev.get("cat", "")
|
| 391 |
+
ts = ev.get("ts")
|
| 392 |
+
dur = ev.get("dur")
|
| 393 |
+
if not (isinstance(ts, (int, float)) and isinstance(dur, (int, float))):
|
| 394 |
+
continue
|
| 395 |
+
ts_f = float(ts)
|
| 396 |
+
dur_f = float(dur)
|
| 397 |
+
end = ts_f + dur_f
|
| 398 |
+
span_min = ts_f if span_min is None else min(span_min, ts_f)
|
| 399 |
+
span_max = end if span_max is None else max(span_max, end)
|
| 400 |
+
if isinstance(cat, str) and "cpu_op" in cat.lower():
|
| 401 |
+
cpu_op_us += dur_f
|
| 402 |
+
if span_min is None or span_max is None or span_max <= span_min:
|
| 403 |
+
return None
|
| 404 |
+
span = span_max - span_min
|
| 405 |
+
if span <= 0:
|
| 406 |
+
return None
|
| 407 |
+
return min(1.0, cpu_op_us / span)
|
| 408 |
+
|
| 409 |
+
|
| 410 |
+
# ---------------------------------------------------------------------------
|
| 411 |
+
# amd_smi.csv → _SmiSummary
|
| 412 |
+
# ---------------------------------------------------------------------------
|
| 413 |
+
|
| 414 |
+
|
| 415 |
+
def _read_amd_smi(path: Path, warnings: list[str]) -> _SmiSummary:
|
| 416 |
+
"""Aggregate amd-smi polling output into HBM peak/avg + GPU util."""
|
| 417 |
+
summary = _SmiSummary()
|
| 418 |
+
if not path.exists():
|
| 419 |
+
warnings.append(f"profile_parser: amd-smi telemetry not found at {path}")
|
| 420 |
+
return summary
|
| 421 |
+
|
| 422 |
+
try:
|
| 423 |
+
raw = path.read_text()
|
| 424 |
+
except OSError as exc:
|
| 425 |
+
warnings.append(f"profile_parser: failed to read amd-smi telemetry ({exc})")
|
| 426 |
+
return summary
|
| 427 |
+
|
| 428 |
+
csv_text = _strip_amd_smi_preamble(raw)
|
| 429 |
+
if csv_text is None:
|
| 430 |
+
warnings.append(f"profile_parser: no parseable header in amd-smi telemetry at {path}")
|
| 431 |
+
return summary
|
| 432 |
+
|
| 433 |
+
try:
|
| 434 |
+
import io as _io
|
| 435 |
+
|
| 436 |
+
reader = csv.DictReader(_io.StringIO(csv_text))
|
| 437 |
+
if reader.fieldnames is None:
|
| 438 |
+
warnings.append(f"profile_parser: empty amd-smi telemetry at {path}")
|
| 439 |
+
return summary
|
| 440 |
+
hbm_col = _pick_column(
|
| 441 |
+
reader.fieldnames,
|
| 442 |
+
[
|
| 443 |
+
# `amd-smi metric --mem-usage --csv` (ROCm 7.x) emits
|
| 444 |
+
# "used_vram" in MB. Older `amd-smi monitor --vram-usage`
|
| 445 |
+
# emits "VRAM_USED" / "vram_used_mb" depending on minor
|
| 446 |
+
# version.
|
| 447 |
+
"used_vram",
|
| 448 |
+
"USED_VRAM",
|
| 449 |
+
"VRAM_USED_MB",
|
| 450 |
+
"vram_used_mb",
|
| 451 |
+
"VRAM_USED",
|
| 452 |
+
"VRAM_USED_GB",
|
| 453 |
+
"vram_used",
|
| 454 |
+
"VRAM Used",
|
| 455 |
+
# Older rocm 6.x naming
|
| 456 |
+
"MEM_USED",
|
| 457 |
+
"mem_used",
|
| 458 |
+
],
|
| 459 |
+
)
|
| 460 |
+
util_col = _pick_column(
|
| 461 |
+
reader.fieldnames,
|
| 462 |
+
[
|
| 463 |
+
# `amd-smi metric --usage --csv` (ROCm 7.x) emits a single
|
| 464 |
+
# consolidated "usage" column. Some builds report N/A here
|
| 465 |
+
# — _coerce_float drops those silently and we fall back to
|
| 466 |
+
# the kernel-trace gpu_util estimate downstream.
|
| 467 |
+
"usage",
|
| 468 |
+
"USAGE",
|
| 469 |
+
# `amd-smi monitor --gfx` (ROCm 7.x) → "gfx_util"
|
| 470 |
+
"GFX_UTIL",
|
| 471 |
+
"gfx_util",
|
| 472 |
+
"GFX_UTILIZATION",
|
| 473 |
+
"gfx_utilization",
|
| 474 |
+
# rocm 6.x and older
|
| 475 |
+
"GFX_ACTIVITY",
|
| 476 |
+
"gfx_activity",
|
| 477 |
+
"GPU_USE",
|
| 478 |
+
"GFX %",
|
| 479 |
+
"Util",
|
| 480 |
+
],
|
| 481 |
+
)
|
| 482 |
+
rocm_col = _pick_column(reader.fieldnames, ["ROCM_VERSION", "rocm_version"])
|
| 483 |
+
|
| 484 |
+
hbm_samples: list[float] = []
|
| 485 |
+
util_samples: list[float] = []
|
| 486 |
+
for row in reader:
|
| 487 |
+
if hbm_col:
|
| 488 |
+
hbm_gb = _hbm_to_gb(row.get(hbm_col), hbm_col)
|
| 489 |
+
if hbm_gb is not None:
|
| 490 |
+
hbm_samples.append(hbm_gb)
|
| 491 |
+
if util_col:
|
| 492 |
+
util = _coerce_float(row.get(util_col))
|
| 493 |
+
if util is not None:
|
| 494 |
+
util_samples.append(min(100.0, util))
|
| 495 |
+
if rocm_col and summary.rocm_version is None:
|
| 496 |
+
val = (row.get(rocm_col) or "").strip()
|
| 497 |
+
if val:
|
| 498 |
+
summary.rocm_version = val
|
| 499 |
+
if hbm_samples:
|
| 500 |
+
summary.hbm_peak_gb = max(hbm_samples)
|
| 501 |
+
summary.hbm_avg_gb = sum(hbm_samples) / len(hbm_samples)
|
| 502 |
+
if util_samples:
|
| 503 |
+
summary.gpu_util_pct = sum(util_samples) / len(util_samples)
|
| 504 |
+
return summary
|
| 505 |
+
except csv.Error as exc:
|
| 506 |
+
warnings.append(f"profile_parser: failed to parse amd-smi telemetry ({exc})")
|
| 507 |
+
return summary
|
| 508 |
+
|
| 509 |
+
|
| 510 |
+
def _strip_amd_smi_preamble(raw: str) -> str | None:
|
| 511 |
+
"""Drop everything before the first real CSV header and dedupe repeated
|
| 512 |
+
header lines — both noise produced by `amd-smi <subcmd> --watch`.
|
| 513 |
+
|
| 514 |
+
--watch prints a "'CTRL' + 'C' to stop watching output:" banner once at
|
| 515 |
+
the top, then re-emits the CSV header on every iteration. csv.DictReader
|
| 516 |
+
naively reads the banner as fieldnames and treats every subsequent
|
| 517 |
+
header as a misshapen data row. Pre-strip both before handing it off.
|
| 518 |
+
|
| 519 |
+
Returns a CSV string ready for DictReader, or None if no header line
|
| 520 |
+
is recognisable.
|
| 521 |
+
"""
|
| 522 |
+
lines = raw.splitlines()
|
| 523 |
+
header_idx: int | None = None
|
| 524 |
+
for i, line in enumerate(lines):
|
| 525 |
+
if "," not in line:
|
| 526 |
+
continue
|
| 527 |
+
lower = line.lower()
|
| 528 |
+
# Recognised tokens come from the columns amd-smi metric / monitor
|
| 529 |
+
# actually emit. Match conservatively — banners or other noise can
|
| 530 |
+
# contain commas too.
|
| 531 |
+
if any(tok in lower for tok in ("vram", "gfx_", "timestamp,", "gpu_use", "gpu,")):
|
| 532 |
+
header_idx = i
|
| 533 |
+
break
|
| 534 |
+
if header_idx is None:
|
| 535 |
+
return None
|
| 536 |
+
|
| 537 |
+
header = lines[header_idx]
|
| 538 |
+
data = [
|
| 539 |
+
line
|
| 540 |
+
for line in lines[header_idx + 1 :]
|
| 541 |
+
if line.strip() and line != header
|
| 542 |
+
]
|
| 543 |
+
return header + "\n" + "\n".join(data) + ("\n" if data else "")
|
| 544 |
+
|
| 545 |
+
|
| 546 |
+
def _hbm_to_gb(raw: str | None, column_name: str | None = None) -> float | None:
|
| 547 |
+
"""amd-smi sometimes reports VRAM in MB, sometimes in GB.
|
| 548 |
+
|
| 549 |
+
First check the column name — `*vram*` / `*_mb` columns are MB-typed
|
| 550 |
+
in every amd-smi build we've seen; `*_gb` is GB. Without a column-name
|
| 551 |
+
hint, fall back to a value heuristic. The old "v > 1024 ⇒ MB" heuristic
|
| 552 |
+
misclassified small idle samples (e.g. 285 MB at GPU idle) as GB and
|
| 553 |
+
inflated the peak across the run, which then forced memory_headroom to
|
| 554 |
+
zero downstream.
|
| 555 |
+
"""
|
| 556 |
+
if not raw:
|
| 557 |
+
return None
|
| 558 |
+
try:
|
| 559 |
+
v = float(str(raw).strip().replace("MB", "").replace("GB", "").replace(",", ""))
|
| 560 |
+
except ValueError:
|
| 561 |
+
return None
|
| 562 |
+
if column_name:
|
| 563 |
+
lower = column_name.lower()
|
| 564 |
+
if "_mb" in lower or lower.endswith("mb"):
|
| 565 |
+
return v / 1024.0
|
| 566 |
+
if "_gb" in lower or lower.endswith("gb"):
|
| 567 |
+
return v
|
| 568 |
+
if "vram" in lower or "mem" in lower:
|
| 569 |
+
# amd-smi metric / monitor default to MB for the bare
|
| 570 |
+
# `used_vram` / `mem_used` columns on every ROCm 6.x+ build.
|
| 571 |
+
return v / 1024.0
|
| 572 |
+
# No column hint — fall back to value heuristic.
|
| 573 |
+
if v > 1024.0:
|
| 574 |
+
return v / 1024.0
|
| 575 |
+
return v
|
| 576 |
+
|
| 577 |
+
|
| 578 |
+
# ---------------------------------------------------------------------------
|
| 579 |
+
# Waste-budget heuristics
|
| 580 |
+
# ---------------------------------------------------------------------------
|
| 581 |
+
|
| 582 |
+
|
| 583 |
+
def _waste_budget(
|
| 584 |
+
*,
|
| 585 |
+
kernels: list[_Kernel],
|
| 586 |
+
torch_summary: _TorchSummary,
|
| 587 |
+
smi: _SmiSummary,
|
| 588 |
+
config: WorkloadConfig | None,
|
| 589 |
+
) -> WasteBudget:
|
| 590 |
+
"""Decompose step time into the seven WasteBudget buckets (architecture.md §3).
|
| 591 |
+
|
| 592 |
+
These are HEURISTICS, not measurements. Each bucket is in seconds-per-step
|
| 593 |
+
so they can be summed against `step_time_seconds`. If we can't observe a
|
| 594 |
+
bucket we set it to 0 — `evidence_coverage` in `propose_patch` then
|
| 595 |
+
discounts confidence accordingly.
|
| 596 |
+
|
| 597 |
+
Per-bucket logic:
|
| 598 |
+
|
| 599 |
+
data_wait
|
| 600 |
+
Fraction of step time where GPU util was below 30% AND the host was
|
| 601 |
+
busy (host_busy_fraction > 0.5). Maps to dataloader / H2D copy stalls.
|
| 602 |
+
|
| 603 |
+
precision_path
|
| 604 |
+
If the user is already on bf16/fp8 we skip this bucket (recovery is 0).
|
| 605 |
+
Otherwise we estimate from kernel names: time spent in fp16-tagged
|
| 606 |
+
GEMMs is the recoverable surface; bf16-tagged kernels are not.
|
| 607 |
+
|
| 608 |
+
kernel_shape
|
| 609 |
+
Fraction of total GEMM kernel time spent on kernels whose names
|
| 610 |
+
match `generic|fallback|naive|reference` — i.e. hipBLASLt/MIOpen
|
| 611 |
+
couldn't pick a tuned tile size and fell back to a slow path.
|
| 612 |
+
|
| 613 |
+
host_gap
|
| 614 |
+
Time the GPU was idle while the host was NOT busy either — pure
|
| 615 |
+
launch latency / eager-mode kernel gaps. We approximate as
|
| 616 |
+
`(1 - gpu_util) * (1 - host_busy)` × step_time.
|
| 617 |
+
|
| 618 |
+
comm_excess
|
| 619 |
+
Sum of all collective-kernel duration (anything matching the rccl
|
| 620 |
+
pattern). Treated as 100% recoverable elsewhere — the rule decides
|
| 621 |
+
recovery_fraction.
|
| 622 |
+
|
| 623 |
+
memory_headroom
|
| 624 |
+
`(192 - hbm_peak) / 192 × small_constant`. Only counts the headroom
|
| 625 |
+
that exceeded a healthy 70% target — running at 60% of HBM costs us
|
| 626 |
+
roughly 0.07 of step time worth of optimisation surface.
|
| 627 |
+
|
| 628 |
+
useful_gpu
|
| 629 |
+
Whatever step time is left. Will be roughly the busy GPU time minus
|
| 630 |
+
comm_excess and the kernel-shape penalty.
|
| 631 |
+
"""
|
| 632 |
+
step_t = torch_summary.step_time_seconds or 0.0
|
| 633 |
+
gpu_util = (smi.gpu_util_pct or 0.0) / 100.0 # 0..1
|
| 634 |
+
host_busy = torch_summary.host_busy_fraction or 0.0
|
| 635 |
+
if step_t <= 0:
|
| 636 |
+
return WasteBudget()
|
| 637 |
+
|
| 638 |
+
# data_wait: GPU under-utilised AND host busy → dataloader bottleneck.
|
| 639 |
+
if gpu_util < 0.30 and host_busy > 0.5:
|
| 640 |
+
data_wait = step_t * (1.0 - gpu_util) * host_busy
|
| 641 |
+
else:
|
| 642 |
+
data_wait = 0.0
|
| 643 |
+
|
| 644 |
+
# host_gap: GPU idle while host idle too → launch latency / kernel gaps.
|
| 645 |
+
host_gap = step_t * (1.0 - gpu_util) * (1.0 - host_busy)
|
| 646 |
+
|
| 647 |
+
# comm_excess: total time in collective kernels (in seconds).
|
| 648 |
+
comm_excess_ns = sum(k.duration_ns for k in kernels if k.is_collective)
|
| 649 |
+
comm_excess = comm_excess_ns / 1e9
|
| 650 |
+
|
| 651 |
+
# kernel_shape: fraction of GEMM time spent on un-tuned/generic kernels.
|
| 652 |
+
gemm_total_ns = sum(k.duration_ns for k in kernels if k.is_gemm) or 1
|
| 653 |
+
generic_gemm_ns = sum(k.duration_ns for k in kernels if k.is_generic_gemm)
|
| 654 |
+
kernel_shape = (generic_gemm_ns / gemm_total_ns) * step_t * gpu_util # cap at "real" GPU time
|
| 655 |
+
|
| 656 |
+
# precision_path: only meaningful if config is fp16/fp32. Estimate from
|
| 657 |
+
# kernel name tags. If config is bf16+ we leave as 0 — already optimal.
|
| 658 |
+
precision_path = 0.0
|
| 659 |
+
if config is None or config.precision in {"fp16", "fp32"}:
|
| 660 |
+
fp16_ns = sum(k.duration_ns for k in kernels if _FP16_PATTERN.search(k.name))
|
| 661 |
+
bf16_ns = sum(k.duration_ns for k in kernels if _BF16_PATTERN.search(k.name))
|
| 662 |
+
denom = fp16_ns + bf16_ns
|
| 663 |
+
if denom > 0:
|
| 664 |
+
# Fraction of compute still on fp16. On MI300X bf16 is faster +
|
| 665 |
+
# more numerically stable, so this fraction is the recoverable
|
| 666 |
+
# precision_path surface.
|
| 667 |
+
precision_path = (fp16_ns / denom) * step_t * gpu_util * 0.10
|
| 668 |
+
|
| 669 |
+
# memory_headroom: headroom past the 70% healthy target × small constant.
|
| 670 |
+
memory_headroom = 0.0
|
| 671 |
+
hbm_peak = smi.hbm_peak_gb
|
| 672 |
+
if hbm_peak is not None and hbm_peak > 0:
|
| 673 |
+
utilisation = hbm_peak / _HBM_TOTAL_GB
|
| 674 |
+
if utilisation < _HBM_HEALTHY_TARGET:
|
| 675 |
+
slack = (_HBM_HEALTHY_TARGET - utilisation) / _HBM_HEALTHY_TARGET
|
| 676 |
+
# Small constant: HBM slack is real but only enables a fraction of
|
| 677 |
+
# potential gain (you still need a larger batch to use it). 0.05
|
| 678 |
+
# of step time per "unit of slack" is a conservative anchor.
|
| 679 |
+
memory_headroom = slack * step_t * 0.05
|
| 680 |
+
|
| 681 |
+
# useful_gpu: everything else. Clamp to >= 0.
|
| 682 |
+
spent = data_wait + host_gap + comm_excess + kernel_shape + precision_path + memory_headroom
|
| 683 |
+
useful_gpu = max(0.0, step_t - spent)
|
| 684 |
+
|
| 685 |
+
return WasteBudget(
|
| 686 |
+
useful_gpu=useful_gpu,
|
| 687 |
+
data_wait=data_wait,
|
| 688 |
+
host_gap=host_gap,
|
| 689 |
+
comm_excess=comm_excess,
|
| 690 |
+
memory_headroom=memory_headroom,
|
| 691 |
+
precision_path=precision_path,
|
| 692 |
+
kernel_shape=kernel_shape,
|
| 693 |
+
)
|
runner/protocol.py
CHANGED
|
@@ -222,41 +222,6 @@ def _archive_failure(out_dir: Path, proc: subprocess.CompletedProcess) -> Path:
|
|
| 222 |
return dest
|
| 223 |
|
| 224 |
|
| 225 |
-
def _default_runner_timeout_seconds() -> int:
|
| 226 |
-
"""Resolve the LiveRunner subprocess timeout from env, with safe defaults.
|
| 227 |
-
|
| 228 |
-
Reads ``GOBLIN_RUNNER_TIMEOUT_SECONDS`` and validates it as a positive
|
| 229 |
-
int. Anything missing, empty, non-numeric, or non-positive falls back to
|
| 230 |
-
1800 seconds (30 minutes).
|
| 231 |
-
|
| 232 |
-
Why this helper exists: live MI300X audits sometimes overshoot the old
|
| 233 |
-
600s default — model download on a cold cache, ROCm kernel JIT on the
|
| 234 |
-
first step, or torch silently running on CPU after a botched pip
|
| 235 |
-
install. Operators need a knob to extend the budget without editing
|
| 236 |
-
code; this turns it into a single env var.
|
| 237 |
-
"""
|
| 238 |
-
raw = os.environ.get("GOBLIN_RUNNER_TIMEOUT_SECONDS", "").strip()
|
| 239 |
-
if not raw:
|
| 240 |
-
return 1800
|
| 241 |
-
try:
|
| 242 |
-
val = int(raw)
|
| 243 |
-
except ValueError:
|
| 244 |
-
_LOG.warning(
|
| 245 |
-
"LiveRunner: GOBLIN_RUNNER_TIMEOUT_SECONDS=%r is not an int; "
|
| 246 |
-
"falling back to 1800s.",
|
| 247 |
-
raw,
|
| 248 |
-
)
|
| 249 |
-
return 1800
|
| 250 |
-
if val <= 0:
|
| 251 |
-
_LOG.warning(
|
| 252 |
-
"LiveRunner: GOBLIN_RUNNER_TIMEOUT_SECONDS=%d is not positive; "
|
| 253 |
-
"falling back to 1800s.",
|
| 254 |
-
val,
|
| 255 |
-
)
|
| 256 |
-
return 1800
|
| 257 |
-
return val
|
| 258 |
-
|
| 259 |
-
|
| 260 |
class LiveRunner:
|
| 261 |
"""Real-MI300X path: shells out to goblin_runner.sh and parses artefacts.
|
| 262 |
|
|
@@ -271,26 +236,17 @@ class LiveRunner:
|
|
| 271 |
self,
|
| 272 |
runner_script: Path | str = _DEFAULT_RUNNER_SCRIPT,
|
| 273 |
user_script: Path | str = _DEFAULT_USER_SCRIPT,
|
| 274 |
-
timeout_seconds: int
|
| 275 |
fake_fallback: FakeRunner | None = None,
|
| 276 |
) -> None:
|
| 277 |
-
#
|
| 278 |
-
#
|
| 279 |
-
#
|
| 280 |
-
#
|
| 281 |
-
#
|
| 282 |
-
# Qwen2.5-7B LoRA at bs=1/seq=512 finishes in 2–5 min when torch is
|
| 283 |
-
# correctly using ROCm. The extra headroom absorbs cold model
|
| 284 |
-
# downloads, kernel JIT on first run, and slow CPU-fallback hiccups
|
| 285 |
-
# if torch ends up CPU-only. If your workload needs longer, bump
|
| 286 |
-
# the env var:
|
| 287 |
-
# export GOBLIN_RUNNER_TIMEOUT_SECONDS=3600
|
| 288 |
self.runner_script = Path(runner_script)
|
| 289 |
self.user_script = Path(user_script)
|
| 290 |
-
self.timeout_seconds =
|
| 291 |
-
timeout_seconds if timeout_seconds is not None
|
| 292 |
-
else _default_runner_timeout_seconds()
|
| 293 |
-
)
|
| 294 |
self._fake = fake_fallback or FakeRunner()
|
| 295 |
|
| 296 |
# ------------------------------------------------------------------
|
|
|
|
| 222 |
return dest
|
| 223 |
|
| 224 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 225 |
class LiveRunner:
|
| 226 |
"""Real-MI300X path: shells out to goblin_runner.sh and parses artefacts.
|
| 227 |
|
|
|
|
| 236 |
self,
|
| 237 |
runner_script: Path | str = _DEFAULT_RUNNER_SCRIPT,
|
| 238 |
user_script: Path | str = _DEFAULT_USER_SCRIPT,
|
| 239 |
+
timeout_seconds: int = 600,
|
| 240 |
fake_fallback: FakeRunner | None = None,
|
| 241 |
) -> None:
|
| 242 |
+
# Default 600s (10 min). Profile runs (10 steps) finish in seconds
|
| 243 |
+
# on a healthy MI300X; benchmarks (50 steps) in a couple of minutes.
|
| 244 |
+
# 30 minutes was a leftover from a workload that wasn't honoring
|
| 245 |
+
# --max_steps and silently trained for hours. With max_steps wired
|
| 246 |
+
# correctly, 600s is generous.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 247 |
self.runner_script = Path(runner_script)
|
| 248 |
self.user_script = Path(user_script)
|
| 249 |
+
self.timeout_seconds = timeout_seconds
|
|
|
|
|
|
|
|
|
|
| 250 |
self._fake = fake_fallback or FakeRunner()
|
| 251 |
|
| 252 |
# ------------------------------------------------------------------
|
scripts/auto_tune.py
ADDED
|
@@ -0,0 +1,1918 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python
|
| 2 |
+
"""Iterative auto-tuner for AMD MI300X / ROCm 7.0 workloads.
|
| 3 |
+
|
| 4 |
+
Three modes, picked with `--mode`:
|
| 5 |
+
|
| 6 |
+
hardcoded (default)
|
| 7 |
+
Walks through a curated list of MI300X-specific tuning changes one
|
| 8 |
+
at a time. Deterministic, no LLM required — experiments are
|
| 9 |
+
derived from the rules in kb/rocm_rules.yaml.
|
| 10 |
+
|
| 11 |
+
llm
|
| 12 |
+
On each iteration, asks the LLM backend (qwen-hf via HF_TOKEN, or
|
| 13 |
+
qwen-vllm via GOBLIN_QWEN_VLLM_URL) for ONE next experiment given
|
| 14 |
+
the live waste_budget, history, and KB rules. Greedy coordinate
|
| 15 |
+
descent — accept changes that beat the current best by the
|
| 16 |
+
improvement threshold, otherwise revert.
|
| 17 |
+
|
| 18 |
+
llm-explore
|
| 19 |
+
On each iteration, asks the LLM for K candidate experiments at
|
| 20 |
+
once (--candidates-per-iteration, default 3). Benchmarks all K,
|
| 21 |
+
picks the one with the highest tokens/sec, and accepts only if it
|
| 22 |
+
beats the current best. Higher GPU cost (~Kx benchmarks per
|
| 23 |
+
iteration) but better at finding interaction effects that greedy
|
| 24 |
+
one-at-a-time can miss.
|
| 25 |
+
|
| 26 |
+
After each change, runs a real benchmark via goblin_runner.sh and keeps
|
| 27 |
+
the change only if tokens/sec improved meaningfully (>1% by default —
|
| 28 |
+
the threshold cuts measurement noise). Stops when N consecutive
|
| 29 |
+
experiments produce no improvement, or when the source of experiments
|
| 30 |
+
is exhausted.
|
| 31 |
+
|
| 32 |
+
Usage:
|
| 33 |
+
# hardcoded mode (default):
|
| 34 |
+
python scripts/auto_tune.py workloads/train_qwen_lora.py --steps 20
|
| 35 |
+
|
| 36 |
+
# LLM-driven greedy mode:
|
| 37 |
+
python scripts/auto_tune.py workloads/train_qwen_lora.py \\
|
| 38 |
+
--mode llm --steps 20
|
| 39 |
+
|
| 40 |
+
# LLM-driven multi-candidate exploration:
|
| 41 |
+
python scripts/auto_tune.py workloads/train_qwen_lora.py \\
|
| 42 |
+
--mode llm-explore --candidates-per-iteration 3 --steps 20
|
| 43 |
+
|
| 44 |
+
Output:
|
| 45 |
+
- A row-by-row log of each experiment attempted, accepted or rejected
|
| 46 |
+
- A final summary with cumulative speedup
|
| 47 |
+
- A pointer to a temp file containing the best workload script for
|
| 48 |
+
diff-against-baseline inspection
|
| 49 |
+
|
| 50 |
+
Extending hardcoded mode: add an Experiment to EXPERIMENTS. The
|
| 51 |
+
substitutions field is a list of (regex_pattern, replacement) tuples
|
| 52 |
+
applied with re.subn against the workload source. env_vars are exported
|
| 53 |
+
into the goblin_runner.sh subprocess and persist on every accepted
|
| 54 |
+
iteration.
|
| 55 |
+
"""
|
| 56 |
+
|
| 57 |
+
from __future__ import annotations
|
| 58 |
+
|
| 59 |
+
import argparse
|
| 60 |
+
import asyncio
|
| 61 |
+
import json
|
| 62 |
+
import os
|
| 63 |
+
import re
|
| 64 |
+
import subprocess
|
| 65 |
+
import sys
|
| 66 |
+
import tempfile
|
| 67 |
+
from dataclasses import dataclass, field
|
| 68 |
+
from pathlib import Path
|
| 69 |
+
|
| 70 |
+
REPO_ROOT = Path(__file__).resolve().parent.parent
|
| 71 |
+
GOBLIN_RUNNER = REPO_ROOT / "runner" / "goblin_runner.sh"
|
| 72 |
+
sys.path.insert(0, str(REPO_ROOT))
|
| 73 |
+
|
| 74 |
+
# Optional structured-events output. When `--events FILE` is passed, the
|
| 75 |
+
# script appends one JSON object per line at key milestones (baseline,
|
| 76 |
+
# iteration start, candidate done, iteration done, summary). Used by the
|
| 77 |
+
# Streamlit UI to render progress live; CLI users typically don't need it.
|
| 78 |
+
_EVENTS_PATH: Path | None = None
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
def _emit(event: dict) -> None:
|
| 82 |
+
"""Append one NDJSON event to the events file if one was configured."""
|
| 83 |
+
if _EVENTS_PATH is None:
|
| 84 |
+
return
|
| 85 |
+
try:
|
| 86 |
+
with _EVENTS_PATH.open("a") as f:
|
| 87 |
+
f.write(json.dumps(event, default=str) + "\n")
|
| 88 |
+
f.flush() # so a UI tailing the file sees events promptly
|
| 89 |
+
except OSError:
|
| 90 |
+
pass # never crash the run on an event-write failure
|
| 91 |
+
|
| 92 |
+
# Default workload template — used when the user passes --model instead
|
| 93 |
+
# of an explicit workload path. We just substitute MODEL_ID and reuse all
|
| 94 |
+
# the other defaults (fp16, batch=4, eager attention, LoRA r=16, …).
|
| 95 |
+
_DEFAULT_WORKLOAD_TEMPLATE = REPO_ROOT / "workloads" / "train_qwen_lora.py"
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def _generate_workload_from_model(model_id: str, dest: Path) -> Path:
|
| 99 |
+
"""Build a baseline workload by substituting MODEL_ID into the demo
|
| 100 |
+
template (`workloads/train_qwen_lora.py`). Writes to `dest`, returns
|
| 101 |
+
the path.
|
| 102 |
+
|
| 103 |
+
Caveats:
|
| 104 |
+
- Uses the demo's LoRA target_modules (`q_proj`, `v_proj`) which work
|
| 105 |
+
for the major decoder-only LLM families (Qwen, Llama, Mistral,
|
| 106 |
+
Gemma). MoE / GPT-2-style architectures will need a custom workload.
|
| 107 |
+
- The template overwrites HF_TOKEN with a redactable fake. Public
|
| 108 |
+
models load fine; gated models (Llama, etc.) need the user to edit
|
| 109 |
+
the generated workload or use a custom one.
|
| 110 |
+
"""
|
| 111 |
+
if not _DEFAULT_WORKLOAD_TEMPLATE.exists():
|
| 112 |
+
raise SystemExit(
|
| 113 |
+
f"--model needs the template at {_DEFAULT_WORKLOAD_TEMPLATE}, but it's missing"
|
| 114 |
+
)
|
| 115 |
+
template_src = _DEFAULT_WORKLOAD_TEMPLATE.read_text()
|
| 116 |
+
new_src, n = re.subn(
|
| 117 |
+
r'MODEL_ID = "[^"]*"',
|
| 118 |
+
f'MODEL_ID = "{model_id}"',
|
| 119 |
+
template_src,
|
| 120 |
+
)
|
| 121 |
+
if n == 0:
|
| 122 |
+
raise SystemExit(
|
| 123 |
+
f"Couldn't find `MODEL_ID = \"...\"` in {_DEFAULT_WORKLOAD_TEMPLATE} "
|
| 124 |
+
"to substitute. Has the template format changed?"
|
| 125 |
+
)
|
| 126 |
+
dest.write_text(new_src)
|
| 127 |
+
return dest
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
# POSIX env var name: leading letter or underscore, then alnum/underscore.
|
| 131 |
+
# subprocess.run() raises ValueError if any key in the env dict violates
|
| 132 |
+
# this. We validate up-front rather than letting the subprocess crash.
|
| 133 |
+
_VALID_ENV_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
def _sanitize_env_vars(envs: dict, context: str = "") -> dict[str, str]:
|
| 137 |
+
"""Clean an env_vars dict from the LLM:
|
| 138 |
+
1. Strip dotted prefixes (`env_vars.X` → `X`) the LLM mimics from the
|
| 139 |
+
KB transform notation.
|
| 140 |
+
2. Drop any key that still isn't a valid POSIX env var name. Warns
|
| 141 |
+
instead of crashing — the LLM occasionally embeds shell syntax
|
| 142 |
+
(e.g. `'NUMACTL_INTERLEAVE=1'` as a key) which would make
|
| 143 |
+
subprocess.run raise ValueError.
|
| 144 |
+
"""
|
| 145 |
+
cleaned: dict[str, str] = {}
|
| 146 |
+
for k, v in envs.items():
|
| 147 |
+
key = str(k)
|
| 148 |
+
if "." in key:
|
| 149 |
+
stripped = key.rsplit(".", 1)[-1]
|
| 150 |
+
tag = f" [{context}]" if context else ""
|
| 151 |
+
print(f" [warn]{tag} dotted env key {key!r}; using {stripped!r}")
|
| 152 |
+
key = stripped
|
| 153 |
+
if not _VALID_ENV_NAME.match(key):
|
| 154 |
+
tag = f" [{context}]" if context else ""
|
| 155 |
+
print(
|
| 156 |
+
f" [warn]{tag} dropping invalid env var name {key!r} "
|
| 157 |
+
"(must match [A-Za-z_][A-Za-z0-9_]*)"
|
| 158 |
+
)
|
| 159 |
+
continue
|
| 160 |
+
cleaned[key] = str(v)
|
| 161 |
+
return cleaned
|
| 162 |
+
|
| 163 |
+
|
| 164 |
+
@dataclass
|
| 165 |
+
class Experiment:
|
| 166 |
+
name: str
|
| 167 |
+
description: str
|
| 168 |
+
rationale: str
|
| 169 |
+
substitutions: list[tuple[str, str]] = field(default_factory=list)
|
| 170 |
+
env_vars: dict[str, str] = field(default_factory=dict)
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
# Curated for ROCm 7.0 + MI300X (CDNA3, 192 GB HBM3). Ordered roughly by
|
| 174 |
+
# typical impact on Qwen-shaped LoRA fine-tuning workloads. Each
|
| 175 |
+
# experiment stacks on top of any previously accepted ones.
|
| 176 |
+
EXPERIMENTS: list[Experiment] = [
|
| 177 |
+
Experiment(
|
| 178 |
+
name="bf16_over_fp16",
|
| 179 |
+
description="Switch precision from fp16 to bf16",
|
| 180 |
+
rationale=(
|
| 181 |
+
"MI300X (CDNA3) prefers bf16: same throughput, larger numeric "
|
| 182 |
+
"range, no loss-scaler needed. fp16 underutilizes the matrix "
|
| 183 |
+
"engine on this arch."
|
| 184 |
+
),
|
| 185 |
+
substitutions=[
|
| 186 |
+
(r"torch_dtype=torch\.float16", "torch_dtype=torch.bfloat16"),
|
| 187 |
+
(r"\bfp16=True\b", "bf16=True"),
|
| 188 |
+
],
|
| 189 |
+
),
|
| 190 |
+
Experiment(
|
| 191 |
+
name="batch_size_8",
|
| 192 |
+
description="Increase per_device_train_batch_size 4 → 8",
|
| 193 |
+
rationale="MI300X has 192 GB HBM; batch=4 leaves it on the floor.",
|
| 194 |
+
substitutions=[
|
| 195 |
+
(r"per_device_train_batch_size=4\b", "per_device_train_batch_size=8"),
|
| 196 |
+
],
|
| 197 |
+
),
|
| 198 |
+
Experiment(
|
| 199 |
+
name="batch_size_16",
|
| 200 |
+
description="Further increase per_device_train_batch_size to 16",
|
| 201 |
+
rationale="If batch=8 fit and improved, try doubling again.",
|
| 202 |
+
substitutions=[
|
| 203 |
+
(r"per_device_train_batch_size=\d+", "per_device_train_batch_size=16"),
|
| 204 |
+
],
|
| 205 |
+
),
|
| 206 |
+
Experiment(
|
| 207 |
+
name="batch_size_32",
|
| 208 |
+
description="Push per_device_train_batch_size to 32",
|
| 209 |
+
rationale=(
|
| 210 |
+
"MI300X has 192 GB HBM3 — batch 16 typically peaks ~130 GB. "
|
| 211 |
+
"If 16 fit, 32 likely fits too and reduces step overhead per "
|
| 212 |
+
"token. Reverts cleanly via OOM-as-crash if not."
|
| 213 |
+
),
|
| 214 |
+
substitutions=[
|
| 215 |
+
(r"per_device_train_batch_size=\d+", "per_device_train_batch_size=32"),
|
| 216 |
+
],
|
| 217 |
+
),
|
| 218 |
+
Experiment(
|
| 219 |
+
name="sdpa_attention",
|
| 220 |
+
description="Switch attention from eager to SDPA",
|
| 221 |
+
rationale=(
|
| 222 |
+
"Eager attention is the slowest path. SDPA dispatches to the "
|
| 223 |
+
"best available kernel (flash on ROCm 7.x where supported, "
|
| 224 |
+
"memory-efficient elsewhere)."
|
| 225 |
+
),
|
| 226 |
+
substitutions=[
|
| 227 |
+
(r'attn_implementation="eager"', 'attn_implementation="sdpa"'),
|
| 228 |
+
],
|
| 229 |
+
),
|
| 230 |
+
Experiment(
|
| 231 |
+
name="dataloader_workers_4",
|
| 232 |
+
description="Bump dataloader_num_workers 0 → 4",
|
| 233 |
+
rationale=(
|
| 234 |
+
"0 workers means the GPU sits idle while the host loads the "
|
| 235 |
+
"next batch. 4 is a safe value across most CPU configs."
|
| 236 |
+
),
|
| 237 |
+
substitutions=[
|
| 238 |
+
(r"dataloader_num_workers=0", "dataloader_num_workers=4"),
|
| 239 |
+
(r"num_workers=0", "num_workers=4"),
|
| 240 |
+
],
|
| 241 |
+
),
|
| 242 |
+
Experiment(
|
| 243 |
+
name="pin_memory",
|
| 244 |
+
description="Enable dataloader_pin_memory",
|
| 245 |
+
rationale=(
|
| 246 |
+
"Pinned host buffers make H2D copies async and overlap with "
|
| 247 |
+
"the GPU. Worth it once you have >0 dataloader workers."
|
| 248 |
+
),
|
| 249 |
+
substitutions=[
|
| 250 |
+
(r"dataloader_pin_memory=False", "dataloader_pin_memory=True"),
|
| 251 |
+
(r"\bpin_memory=False\b", "pin_memory=True"),
|
| 252 |
+
],
|
| 253 |
+
),
|
| 254 |
+
Experiment(
|
| 255 |
+
name="env_hipblaslt",
|
| 256 |
+
description="Set TORCH_BLAS_PREFER_HIPBLASLT=1",
|
| 257 |
+
rationale=(
|
| 258 |
+
"hipBLASLt is significantly faster than rocBLAS for the GEMM "
|
| 259 |
+
"shapes Qwen produces (LoRA-projected attention)."
|
| 260 |
+
),
|
| 261 |
+
env_vars={"TORCH_BLAS_PREFER_HIPBLASLT": "1"},
|
| 262 |
+
),
|
| 263 |
+
Experiment(
|
| 264 |
+
name="env_tunable_op",
|
| 265 |
+
description="Set PYTORCH_TUNABLEOP_ENABLED=1",
|
| 266 |
+
rationale=(
|
| 267 |
+
"Enables runtime kernel auto-tuning. Pays a first-run "
|
| 268 |
+
"warmup cost in exchange for a steady-state win on every "
|
| 269 |
+
"subsequent step."
|
| 270 |
+
),
|
| 271 |
+
env_vars={"PYTORCH_TUNABLEOP_ENABLED": "1"},
|
| 272 |
+
),
|
| 273 |
+
Experiment(
|
| 274 |
+
name="env_miopen_find",
|
| 275 |
+
description="Set MIOPEN_FIND_MODE=3",
|
| 276 |
+
rationale=(
|
| 277 |
+
"MIOpen FAST mode picks already-tuned kernels without on-the-"
|
| 278 |
+
"fly search. Reduces per-step variance."
|
| 279 |
+
),
|
| 280 |
+
env_vars={"MIOPEN_FIND_MODE": "3"},
|
| 281 |
+
),
|
| 282 |
+
]
|
| 283 |
+
|
| 284 |
+
|
| 285 |
+
# ---------------------------------------------------------------------------
|
| 286 |
+
# Helpers
|
| 287 |
+
# ---------------------------------------------------------------------------
|
| 288 |
+
|
| 289 |
+
|
| 290 |
+
def apply_substitutions(source: str, subs: list[tuple[str, str]]) -> str | None:
|
| 291 |
+
"""Apply each (pattern, replacement) in order. Returns the new source,
|
| 292 |
+
or None if any pattern matched zero times (already applied or N/A for
|
| 293 |
+
this workload)."""
|
| 294 |
+
out = source
|
| 295 |
+
for pattern, replacement in subs:
|
| 296 |
+
new, n = re.subn(pattern, replacement, out)
|
| 297 |
+
if n == 0:
|
| 298 |
+
return None
|
| 299 |
+
out = new
|
| 300 |
+
return out
|
| 301 |
+
|
| 302 |
+
|
| 303 |
+
def benchmark(
|
| 304 |
+
workload_path: Path,
|
| 305 |
+
steps: int,
|
| 306 |
+
env_overrides: dict[str, str],
|
| 307 |
+
timeout: int = 600,
|
| 308 |
+
) -> dict | None:
|
| 309 |
+
"""Run goblin_runner.sh on the workload, return parsed RunMetrics dict
|
| 310 |
+
or None on failure."""
|
| 311 |
+
with tempfile.TemporaryDirectory(prefix="auto_tune_run_") as out_dir_str:
|
| 312 |
+
out_dir = Path(out_dir_str)
|
| 313 |
+
env = os.environ.copy()
|
| 314 |
+
env["USER_SCRIPT"] = str(workload_path)
|
| 315 |
+
env["OUT_DIR"] = str(out_dir)
|
| 316 |
+
env["STEPS"] = str(steps)
|
| 317 |
+
# Candidate workload lives in /tmp, so its self-bootstrap line
|
| 318 |
+
# `sys.path.insert(0, dirname(dirname(__file__)))` resolves to /tmp
|
| 319 |
+
# — which has no `workloads/` package. Inject the real repo root via
|
| 320 |
+
# PYTHONPATH so `from workloads._runtime import ...` succeeds.
|
| 321 |
+
existing_pp = env.get("PYTHONPATH", "")
|
| 322 |
+
env["PYTHONPATH"] = (
|
| 323 |
+
str(REPO_ROOT) + (os.pathsep + existing_pp if existing_pp else "")
|
| 324 |
+
)
|
| 325 |
+
env.update(env_overrides)
|
| 326 |
+
|
| 327 |
+
try:
|
| 328 |
+
proc = subprocess.run(
|
| 329 |
+
[str(GOBLIN_RUNNER)],
|
| 330 |
+
env=env,
|
| 331 |
+
capture_output=True,
|
| 332 |
+
text=True,
|
| 333 |
+
timeout=timeout,
|
| 334 |
+
)
|
| 335 |
+
except subprocess.TimeoutExpired:
|
| 336 |
+
print(f" TIMEOUT after {timeout}s")
|
| 337 |
+
return None
|
| 338 |
+
except ValueError as exc:
|
| 339 |
+
# subprocess.run validates env var names and raises ValueError
|
| 340 |
+
# for malformed keys (e.g. names containing '=' or spaces). The
|
| 341 |
+
# LLM has occasionally emitted those; we sanitize earlier but
|
| 342 |
+
# this is the last-resort backstop so a single bad candidate
|
| 343 |
+
# doesn't crash the whole tuning run.
|
| 344 |
+
print(f" REJECTED — illegal env var name(s): {exc}")
|
| 345 |
+
print(f" env keys offered: {list(env_overrides.keys())}")
|
| 346 |
+
return None
|
| 347 |
+
except OSError as exc:
|
| 348 |
+
print(f" REJECTED — could not spawn goblin_runner.sh: {exc}")
|
| 349 |
+
return None
|
| 350 |
+
|
| 351 |
+
if proc.returncode != 0:
|
| 352 |
+
print(f" goblin_runner.sh failed (exit {proc.returncode})")
|
| 353 |
+
tail = (proc.stderr or "").strip().splitlines()[-8:]
|
| 354 |
+
for line in tail:
|
| 355 |
+
print(f" | {line}")
|
| 356 |
+
return None
|
| 357 |
+
|
| 358 |
+
try:
|
| 359 |
+
from runner import profile_parser
|
| 360 |
+
|
| 361 |
+
metrics = profile_parser.parse(out_dir, steps=steps)
|
| 362 |
+
return metrics.model_dump()
|
| 363 |
+
except Exception as exc: # parser is defensive but be safe
|
| 364 |
+
print(f" profile_parser raised: {type(exc).__name__}: {exc}")
|
| 365 |
+
return None
|
| 366 |
+
|
| 367 |
+
|
| 368 |
+
def _delta_pct(new: float, baseline: float) -> float:
|
| 369 |
+
if baseline <= 0:
|
| 370 |
+
return 0.0
|
| 371 |
+
return (new - baseline) / baseline * 100.0
|
| 372 |
+
|
| 373 |
+
|
| 374 |
+
# ---------------------------------------------------------------------------
|
| 375 |
+
# LLM-driven experiment generator
|
| 376 |
+
# ---------------------------------------------------------------------------
|
| 377 |
+
|
| 378 |
+
|
| 379 |
+
_LLM_SYSTEM_PROMPT = """\
|
| 380 |
+
You are an expert at tuning AMD MI300X (ROCm 7.0, CDNA3 arch, 192 GB
|
| 381 |
+
HBM3) training workloads. The user is iteratively benchmarking changes
|
| 382 |
+
to a transformers/peft fine-tuning script. On each turn you suggest ONE
|
| 383 |
+
specific parameter change to try next, targeting the largest non-useful
|
| 384 |
+
waste bucket in the most recent benchmark.
|
| 385 |
+
|
| 386 |
+
Your output MUST be a single JSON object with this exact shape (no
|
| 387 |
+
prose, no markdown fences, just the object):
|
| 388 |
+
|
| 389 |
+
{
|
| 390 |
+
"name": "short_snake_case_name",
|
| 391 |
+
"rationale": "1-3 sentences on why this change addresses the worst waste bucket",
|
| 392 |
+
"substitutions": [["regex_pattern", "replacement"]],
|
| 393 |
+
"env_vars": {"VAR_NAME": "value"}
|
| 394 |
+
}
|
| 395 |
+
|
| 396 |
+
CRITICAL output rules — read carefully:
|
| 397 |
+
|
| 398 |
+
1. env_vars keys are LITERAL POSIX shell environment variable names.
|
| 399 |
+
They MUST match the regex [A-Za-z_][A-Za-z0-9_]* — letters, digits,
|
| 400 |
+
underscores only, starting with a letter or underscore.
|
| 401 |
+
- NEVER prefix them with "env_vars." or any other dotted path.
|
| 402 |
+
- NEVER include "=" or shell syntax in the key — env var names are
|
| 403 |
+
identifiers, NOT assignments and NOT commands.
|
| 404 |
+
- If you want to invoke a command-line tool like `numactl` or
|
| 405 |
+
`taskset`, that CANNOT be expressed as an env_var. Don't try.
|
| 406 |
+
Either propose a `substitutions` change to the script, or skip.
|
| 407 |
+
Wrong: {"env_vars.MIOPEN_FIND_MODE": "3"}
|
| 408 |
+
Wrong: {"NUMACTL_INTERLEAVE=1": "numactl --interleave=all"}
|
| 409 |
+
Wrong: {"export FOO": "bar"}
|
| 410 |
+
Right: {"MIOPEN_FIND_MODE": "3"}
|
| 411 |
+
Right: {"TORCH_BLAS_PREFER_HIPBLASLT": "1"}
|
| 412 |
+
|
| 413 |
+
2. substitutions are (regex_pattern, replacement) pairs applied with
|
| 414 |
+
re.subn against the current workload source. Patterns must match at
|
| 415 |
+
least one occurrence in the source — if zero matches, the experiment
|
| 416 |
+
is auto-skipped (counted as no improvement).
|
| 417 |
+
|
| 418 |
+
3. When the previous change for a parameter improved tokens/sec, push
|
| 419 |
+
that parameter further in the same direction next time. E.g. if
|
| 420 |
+
batch_size 4 → 8 won, try 8 → 16. If 16 won and HBM is still under
|
| 421 |
+
~150 GB, try 32. Don't be timid — MI300X has 192 GB HBM3.
|
| 422 |
+
|
| 423 |
+
4. Don't repeat any (name OR substitution OR env_var combo) from
|
| 424 |
+
history. If a change was rejected, don't propose the same numerical
|
| 425 |
+
value again — try a different one.
|
| 426 |
+
|
| 427 |
+
5. If you cannot think of a productive next change, output:
|
| 428 |
+
{"name": "STOP", "rationale": "<why>", "substitutions": [], "env_vars": {}}
|
| 429 |
+
|
| 430 |
+
CONCRETE OUTPUT EXAMPLES — match this shape exactly:
|
| 431 |
+
|
| 432 |
+
Switch fp16 → bf16 (precision_path bucket):
|
| 433 |
+
{"name": "bf16_over_fp16",
|
| 434 |
+
"rationale": "MI300X CDNA3 matrix cores prefer bf16: same throughput, larger numeric range, no loss-scaler.",
|
| 435 |
+
"substitutions": [["fp16=True", "bf16=True"], ["torch_dtype=torch\\\\.float16", "torch_dtype=torch.bfloat16"]],
|
| 436 |
+
"env_vars": {}}
|
| 437 |
+
|
| 438 |
+
Increase batch size to 16 (memory_headroom bucket):
|
| 439 |
+
{"name": "batch_size_16",
|
| 440 |
+
"rationale": "Current HBM peak is well under 192 GB; bigger batch saturates the GPU.",
|
| 441 |
+
"substitutions": [["per_device_train_batch_size=\\\\d+", "per_device_train_batch_size=16"]],
|
| 442 |
+
"env_vars": {}}
|
| 443 |
+
|
| 444 |
+
Switch attention to SDPA (kernel_shape bucket):
|
| 445 |
+
{"name": "sdpa_attention",
|
| 446 |
+
"rationale": "Eager attention is the slowest path; SDPA dispatches to a tuned kernel.",
|
| 447 |
+
"substitutions": [["attn_implementation=\\"eager\\"", "attn_implementation=\\"sdpa\\""]],
|
| 448 |
+
"env_vars": {}}
|
| 449 |
+
|
| 450 |
+
Bump dataloader workers (data_wait bucket):
|
| 451 |
+
{"name": "dataloader_workers_4",
|
| 452 |
+
"rationale": "0 workers starves the GPU between batches.",
|
| 453 |
+
"substitutions": [["dataloader_num_workers=0", "dataloader_num_workers=4"]],
|
| 454 |
+
"env_vars": {}}
|
| 455 |
+
|
| 456 |
+
Set MIOpen FAST mode (kernel_shape bucket, env-only):
|
| 457 |
+
{"name": "miopen_find_fast",
|
| 458 |
+
"rationale": "FAST mode picks already-tuned kernels without on-the-fly search.",
|
| 459 |
+
"substitutions": [],
|
| 460 |
+
"env_vars": {"MIOPEN_FIND_MODE": "3"}}
|
| 461 |
+
|
| 462 |
+
Prefer hipBLASLt (kernel_shape bucket, env-only):
|
| 463 |
+
{"name": "prefer_hipblaslt",
|
| 464 |
+
"rationale": "hipBLASLt is faster than rocBLAS for Qwen GEMM shapes on MI300X.",
|
| 465 |
+
"substitutions": [],
|
| 466 |
+
"env_vars": {"TORCH_BLAS_PREFER_HIPBLASLT": "1"}}
|
| 467 |
+
"""
|
| 468 |
+
|
| 469 |
+
|
| 470 |
+
_LLM_USER_TEMPLATE = """\
|
| 471 |
+
Hardware facts (use these — do not contradict):
|
| 472 |
+
- AMD MI300X, CDNA3 architecture, 192 GB HBM3
|
| 473 |
+
- bf16 throughput on CDNA3 ≈ same as fp16, > fp32 (matrix engine is fp16/bf16/fp8 native)
|
| 474 |
+
- fp32 is the SLOWEST option on this arch — never suggest it as an improvement
|
| 475 |
+
|
| 476 |
+
Known incompatibilities for THIS workload (peft + LoRA on transformers Trainer):
|
| 477 |
+
{incompatibilities}
|
| 478 |
+
|
| 479 |
+
KB rules (one-liner per rule, for grounding):
|
| 480 |
+
{kb_summary}
|
| 481 |
+
|
| 482 |
+
Current accepted workload state — these are the literal values in the
|
| 483 |
+
script after every change accepted so far. The next change you propose
|
| 484 |
+
should mutate one of these (or set an env var). DO NOT propose a value
|
| 485 |
+
that's already present here.
|
| 486 |
+
{tunables}
|
| 487 |
+
|
| 488 |
+
Latest benchmark (this is the result of the most recent ACCEPTED state):
|
| 489 |
+
- tokens_per_sec: {tps:.1f}
|
| 490 |
+
- mfu_pct: {mfu:.2f} (% of MI300X dense bf16 peak; healthy LoRA ranges 30-50%)
|
| 491 |
+
- gpu_util_pct: {util:.1f}
|
| 492 |
+
- hbm_peak_gb: {hbm:.2f}
|
| 493 |
+
- waste_budget (seconds/step):
|
| 494 |
+
{waste_lines}
|
| 495 |
+
|
| 496 |
+
Sorted recoverable waste (largest first — go after these):
|
| 497 |
+
{recoverable_sorted}
|
| 498 |
+
|
| 499 |
+
History of changes already tried this run (newest first; outcomes are
|
| 500 |
+
"accepted" / "rejected" / "crashed" / "skipped"):
|
| 501 |
+
{history_lines}
|
| 502 |
+
|
| 503 |
+
If the latest entry is "crashed", the change you propose next must be
|
| 504 |
+
STRUCTURALLY different (different parameter, not just a different value
|
| 505 |
+
of the same one).
|
| 506 |
+
|
| 507 |
+
Suggest ONE next change targeting the largest recoverable bucket. JSON only.
|
| 508 |
+
"""
|
| 509 |
+
|
| 510 |
+
|
| 511 |
+
# Workload-specific incompatibilities the LLM otherwise wastes iterations on.
|
| 512 |
+
# Keep this list short and concrete — it goes into every prompt.
|
| 513 |
+
_KNOWN_INCOMPATIBILITIES = [
|
| 514 |
+
"gradient_checkpointing=True requires `model.enable_input_require_grads()`"
|
| 515 |
+
" before peft wrapping for LoRA models. Setting it via a single substitution"
|
| 516 |
+
" WILL CRASH the workload. Don't propose it.",
|
| 517 |
+
"bitsandbytes-based optimizers (`adamw_8bit`, `paged_adamw_8bit`) and"
|
| 518 |
+
" `load_in_8bit=True` are NOT supported on ROCm 7.x. Don't propose them.",
|
| 519 |
+
"torch_compile=True with peft/LoRA on ROCm 7.x triggers compile-time"
|
| 520 |
+
" errors with the current PyTorch nightly (2.9.x). Don't propose it"
|
| 521 |
+
" unless you have specific evidence it works on this version.",
|
| 522 |
+
"flash_attention_2 may not be installed (try `attn_implementation=\"sdpa\"`"
|
| 523 |
+
" before `\"flash_attention_2\"`).",
|
| 524 |
+
"persistent_workers=True requires num_workers > 0. PyTorch raises"
|
| 525 |
+
" `ValueError: persistent_workers option needs num_workers > 0` if you"
|
| 526 |
+
" enable it while num_workers=0. If the current workload has"
|
| 527 |
+
" dataloader_num_workers=0, do NOT propose persistent_workers=True"
|
| 528 |
+
" alone — pair it with `dataloader_num_workers=4` (or higher) in the"
|
| 529 |
+
" SAME experiment via two substitutions, or wait until a previous"
|
| 530 |
+
" experiment has bumped num_workers above 0.",
|
| 531 |
+
"dataloader_prefetch_factor only works when num_workers > 0 (same"
|
| 532 |
+
" constraint as persistent_workers). Same rule: bump num_workers in"
|
| 533 |
+
" the same experiment, or skip.",
|
| 534 |
+
]
|
| 535 |
+
|
| 536 |
+
|
| 537 |
+
def _kb_summary(rules_yaml_path: Path, max_chars: int = 6000) -> str:
|
| 538 |
+
"""Return a compact one-line-per-rule summary of kb/rocm_rules.yaml.
|
| 539 |
+
|
| 540 |
+
Notably we DO NOT show the raw `transform` field — earlier versions
|
| 541 |
+
did and the LLM ended up copying its dotted-path notation literally
|
| 542 |
+
(`env_vars.MIOPEN_FIND_MODE` as the env var name, not as a dict
|
| 543 |
+
accessor). The system prompt's CONCRETE EXAMPLES section is the
|
| 544 |
+
canonical source of truth for output shape; this summary just
|
| 545 |
+
grounds the LLM's reasoning in the catalog of known issues.
|
| 546 |
+
"""
|
| 547 |
+
if not rules_yaml_path.exists():
|
| 548 |
+
return "(KB rules file not found)"
|
| 549 |
+
try:
|
| 550 |
+
import yaml
|
| 551 |
+
|
| 552 |
+
rules = yaml.safe_load(rules_yaml_path.read_text()) or []
|
| 553 |
+
except Exception as exc:
|
| 554 |
+
return f"(failed to parse KB: {exc})"
|
| 555 |
+
|
| 556 |
+
lines = []
|
| 557 |
+
for r in rules:
|
| 558 |
+
if not isinstance(r, dict):
|
| 559 |
+
continue
|
| 560 |
+
rid = r.get("id", "?")
|
| 561 |
+
bucket = r.get("targets_bucket", "?")
|
| 562 |
+
sym = (r.get("symptom") or "").strip().replace("\n", " ")
|
| 563 |
+
if len(sym) > 110:
|
| 564 |
+
sym = sym[:107] + "..."
|
| 565 |
+
lines.append(f"- {rid:55s} [{bucket}] {sym}")
|
| 566 |
+
text = "\n".join(lines)
|
| 567 |
+
if len(text) > max_chars:
|
| 568 |
+
text = text[:max_chars] + "\n... (truncated)"
|
| 569 |
+
return text
|
| 570 |
+
|
| 571 |
+
|
| 572 |
+
# Map of (substring-in-source) → (parameter description, example regex
|
| 573 |
+
# pattern, example replacement template). Each entry is a hint shown to
|
| 574 |
+
# the LLM so it has a concrete target to point its substitutions at —
|
| 575 |
+
# instead of guessing what the workload's literal config text looks like.
|
| 576 |
+
_TUNABLE_HINTS: list[tuple[str, str, str, str]] = [
|
| 577 |
+
# (token to detect, description, regex_for_substitution, replacement_template)
|
| 578 |
+
("torch_dtype=torch.float16",
|
| 579 |
+
"model precision (matches `torch_dtype=torch.float16`)",
|
| 580 |
+
r"torch_dtype=torch\.float16",
|
| 581 |
+
"torch_dtype=torch.bfloat16"),
|
| 582 |
+
("torch_dtype=torch.bfloat16",
|
| 583 |
+
"model precision (already bf16)",
|
| 584 |
+
r"torch_dtype=torch\.bfloat16",
|
| 585 |
+
"torch_dtype=torch.float16"),
|
| 586 |
+
("fp16=True",
|
| 587 |
+
"TrainingArguments fp16 (matches `fp16=True`)",
|
| 588 |
+
r"\bfp16=True\b",
|
| 589 |
+
"bf16=True"),
|
| 590 |
+
("bf16=True",
|
| 591 |
+
"TrainingArguments bf16 (already bf16)",
|
| 592 |
+
r"\bbf16=True\b",
|
| 593 |
+
"fp16=True"),
|
| 594 |
+
("attn_implementation=\"eager\"",
|
| 595 |
+
"attention impl (matches `attn_implementation=\"eager\"`)",
|
| 596 |
+
r'attn_implementation="eager"',
|
| 597 |
+
'attn_implementation="sdpa"'),
|
| 598 |
+
("attn_implementation=\"sdpa\"",
|
| 599 |
+
"attention impl (currently sdpa; could try flash_attention_2)",
|
| 600 |
+
r'attn_implementation="sdpa"',
|
| 601 |
+
'attn_implementation="flash_attention_2"'),
|
| 602 |
+
("per_device_train_batch_size=",
|
| 603 |
+
"per-device batch size (matches `per_device_train_batch_size=<N>`)",
|
| 604 |
+
r"per_device_train_batch_size=\d+",
|
| 605 |
+
"per_device_train_batch_size=<NEW_VALUE>"),
|
| 606 |
+
("dataloader_num_workers=",
|
| 607 |
+
"dataloader workers (matches `dataloader_num_workers=<N>`)",
|
| 608 |
+
r"dataloader_num_workers=\d+",
|
| 609 |
+
"dataloader_num_workers=<NEW_VALUE>"),
|
| 610 |
+
("dataloader_pin_memory=",
|
| 611 |
+
"dataloader pin_memory (matches `dataloader_pin_memory=<bool>`)",
|
| 612 |
+
r"dataloader_pin_memory=(True|False)",
|
| 613 |
+
"dataloader_pin_memory=True"),
|
| 614 |
+
("gradient_checkpointing=",
|
| 615 |
+
"gradient checkpointing toggle",
|
| 616 |
+
r"gradient_checkpointing=(True|False)",
|
| 617 |
+
"gradient_checkpointing=True"),
|
| 618 |
+
("torch_compile=",
|
| 619 |
+
"torch.compile toggle (use cautiously on ROCm 7.x)",
|
| 620 |
+
r"torch_compile=(True|False)",
|
| 621 |
+
"torch_compile=True"),
|
| 622 |
+
("optim=\"adamw_torch\"",
|
| 623 |
+
"optimizer choice (currently adamw_torch)",
|
| 624 |
+
r'optim="adamw_torch"',
|
| 625 |
+
'optim="adamw_torch_fused"'),
|
| 626 |
+
]
|
| 627 |
+
|
| 628 |
+
|
| 629 |
+
def _tunables_summary(source: str) -> str:
|
| 630 |
+
"""Detect which tunable parameters are present in the workload source
|
| 631 |
+
and surface their current literal values + ready-to-use regex patterns
|
| 632 |
+
so the LLM has concrete substitution targets.
|
| 633 |
+
|
| 634 |
+
Skips comment lines when reporting the "current" value — many workloads
|
| 635 |
+
document expected findings in a top-of-file comment block, and we want
|
| 636 |
+
the LLM to see the live config line, not the doc string.
|
| 637 |
+
"""
|
| 638 |
+
lines: list[str] = []
|
| 639 |
+
source_lines = source.splitlines()
|
| 640 |
+
for token, desc, pattern, replacement in _TUNABLE_HINTS:
|
| 641 |
+
live_line: str | None = None
|
| 642 |
+
for raw in source_lines:
|
| 643 |
+
stripped = raw.lstrip()
|
| 644 |
+
if stripped.startswith("#"):
|
| 645 |
+
continue
|
| 646 |
+
if token in raw:
|
| 647 |
+
live_line = raw.strip()
|
| 648 |
+
break
|
| 649 |
+
if live_line is None:
|
| 650 |
+
continue
|
| 651 |
+
lines.append(
|
| 652 |
+
f" • {desc}\n"
|
| 653 |
+
f" current: {live_line}\n"
|
| 654 |
+
f" pattern: {pattern!r} replacement template: {replacement!r}"
|
| 655 |
+
)
|
| 656 |
+
if not lines:
|
| 657 |
+
return " (no recognized tunables — substitutions will need to match other text)"
|
| 658 |
+
return "\n".join(lines)
|
| 659 |
+
|
| 660 |
+
|
| 661 |
+
def _recoverable_sorted(waste: dict) -> str:
|
| 662 |
+
"""List the non-useful_gpu buckets sorted by size, so the LLM can
|
| 663 |
+
explicitly target the biggest one first."""
|
| 664 |
+
if not waste:
|
| 665 |
+
return " (no waste_budget available)"
|
| 666 |
+
items = [
|
| 667 |
+
(name, value)
|
| 668 |
+
for name, value in waste.items()
|
| 669 |
+
if name != "useful_gpu" and isinstance(value, (int, float))
|
| 670 |
+
]
|
| 671 |
+
items.sort(key=lambda kv: kv[1], reverse=True)
|
| 672 |
+
if not items:
|
| 673 |
+
return " (no recoverable buckets)"
|
| 674 |
+
return "\n".join(f" {i + 1}. {name:18s} = {value:.4f}" for i, (name, value) in enumerate(items))
|
| 675 |
+
|
| 676 |
+
|
| 677 |
+
def _config_snippet(source: str, max_lines: int = 80) -> str:
|
| 678 |
+
"""Return the lines around `TrainingArguments(` and `from_pretrained(` so
|
| 679 |
+
the LLM sees the actual config it's modifying without us shipping the
|
| 680 |
+
whole script. Gives ~max_lines of context.
|
| 681 |
+
"""
|
| 682 |
+
lines = source.splitlines()
|
| 683 |
+
keep: list[tuple[int, str]] = []
|
| 684 |
+
for i, line in enumerate(lines):
|
| 685 |
+
lower = line.lower()
|
| 686 |
+
if any(
|
| 687 |
+
tok in lower
|
| 688 |
+
for tok in (
|
| 689 |
+
"trainingarguments(",
|
| 690 |
+
"from_pretrained(",
|
| 691 |
+
"loraconfig(",
|
| 692 |
+
"dataloader(",
|
| 693 |
+
"torch_dtype",
|
| 694 |
+
"attn_implementation",
|
| 695 |
+
"fp16=",
|
| 696 |
+
"bf16=",
|
| 697 |
+
"per_device_train_batch_size",
|
| 698 |
+
"dataloader_num_workers",
|
| 699 |
+
"dataloader_pin_memory",
|
| 700 |
+
"gradient_checkpointing",
|
| 701 |
+
"torch_compile",
|
| 702 |
+
"optim=",
|
| 703 |
+
)
|
| 704 |
+
):
|
| 705 |
+
keep.append((i, line))
|
| 706 |
+
if not keep:
|
| 707 |
+
return source[:2000]
|
| 708 |
+
# Coalesce nearby line indices into windows for readability
|
| 709 |
+
windows: list[list[str]] = []
|
| 710 |
+
last_idx = -10
|
| 711 |
+
cur: list[str] = []
|
| 712 |
+
for i, line in keep:
|
| 713 |
+
if i - last_idx > 3:
|
| 714 |
+
if cur:
|
| 715 |
+
windows.append(cur)
|
| 716 |
+
cur = []
|
| 717 |
+
cur.append(f"{i + 1:4d}: {line}")
|
| 718 |
+
last_idx = i
|
| 719 |
+
if cur:
|
| 720 |
+
windows.append(cur)
|
| 721 |
+
out = "\n\n".join("\n".join(w) for w in windows)
|
| 722 |
+
if out.count("\n") > max_lines:
|
| 723 |
+
out_lines = out.splitlines()[:max_lines]
|
| 724 |
+
out = "\n".join(out_lines) + "\n... (truncated)"
|
| 725 |
+
return out
|
| 726 |
+
|
| 727 |
+
|
| 728 |
+
def _format_history(history: list[dict]) -> str:
|
| 729 |
+
if not history:
|
| 730 |
+
return "(none yet — this is the first iteration)"
|
| 731 |
+
lines = []
|
| 732 |
+
for h in reversed(history[-12:]): # last 12 newest-first
|
| 733 |
+
outcome = h.get("outcome", "?")
|
| 734 |
+
delta = h.get("delta_pct")
|
| 735 |
+
delta_s = f"{delta:+.2f}%" if delta is not None else "n/a"
|
| 736 |
+
subs = h.get("substitutions") or []
|
| 737 |
+
envs = h.get("env_vars") or {}
|
| 738 |
+
change_repr = f"subs={subs} env={envs}"
|
| 739 |
+
lines.append(f"- {h['name']:25s} {outcome:9s} Δ {delta_s:8s} {change_repr}")
|
| 740 |
+
return "\n".join(lines)
|
| 741 |
+
|
| 742 |
+
|
| 743 |
+
def _format_waste(waste: dict) -> str:
|
| 744 |
+
keys = (
|
| 745 |
+
"useful_gpu",
|
| 746 |
+
"data_wait",
|
| 747 |
+
"host_gap",
|
| 748 |
+
"comm_excess",
|
| 749 |
+
"memory_headroom",
|
| 750 |
+
"precision_path",
|
| 751 |
+
"kernel_shape",
|
| 752 |
+
)
|
| 753 |
+
return "\n".join(f" {k:18s} = {waste.get(k, 0.0):.4f}" for k in keys)
|
| 754 |
+
|
| 755 |
+
|
| 756 |
+
def _build_llm_backend(system_prompt: str = _LLM_SYSTEM_PROMPT, max_tokens: int = 1024):
|
| 757 |
+
"""Construct the same backend the agent loop uses. Surfaces a clear
|
| 758 |
+
message if neither HF_TOKEN nor a vLLM URL is configured."""
|
| 759 |
+
has_hf = bool(os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN"))
|
| 760 |
+
has_vllm = bool(os.environ.get("GOBLIN_QWEN_VLLM_URL"))
|
| 761 |
+
backend_kind = os.environ.get("GOBLIN_AGENT_BACKEND", "qwen-hf").lower()
|
| 762 |
+
if backend_kind in ("qwen-hf", "qwen", "hf", "") and not has_hf:
|
| 763 |
+
raise SystemExit(
|
| 764 |
+
"LLM mode requires HF_TOKEN (qwen-hf backend) or "
|
| 765 |
+
"GOBLIN_AGENT_BACKEND=qwen-vllm + GOBLIN_QWEN_VLLM_URL."
|
| 766 |
+
)
|
| 767 |
+
if backend_kind in ("qwen-vllm", "qwen_vllm", "vllm", "local") and not has_vllm:
|
| 768 |
+
raise SystemExit(
|
| 769 |
+
"LLM mode with qwen-vllm backend requires GOBLIN_QWEN_VLLM_URL."
|
| 770 |
+
)
|
| 771 |
+
from agent.backends import make_backend
|
| 772 |
+
|
| 773 |
+
return make_backend(system_prompt=system_prompt, max_tokens=max_tokens)
|
| 774 |
+
|
| 775 |
+
|
| 776 |
+
async def _ask_llm_for_experiment(
|
| 777 |
+
backend,
|
| 778 |
+
*,
|
| 779 |
+
kb_summary: str,
|
| 780 |
+
source: str,
|
| 781 |
+
metrics: dict,
|
| 782 |
+
history: list[dict],
|
| 783 |
+
) -> Experiment | None:
|
| 784 |
+
"""One LLM turn → one Experiment (or None for STOP / parse failure)."""
|
| 785 |
+
waste = metrics.get("waste_budget") or {}
|
| 786 |
+
prompt = _LLM_USER_TEMPLATE.format(
|
| 787 |
+
incompatibilities="\n".join(f"- {line}" for line in _KNOWN_INCOMPATIBILITIES),
|
| 788 |
+
kb_summary=kb_summary,
|
| 789 |
+
tunables=_tunables_summary(source),
|
| 790 |
+
tps=metrics.get("tokens_per_sec", 0.0),
|
| 791 |
+
mfu=metrics.get("mfu_pct", 0.0),
|
| 792 |
+
util=metrics.get("gpu_util_pct", 0.0),
|
| 793 |
+
hbm=metrics.get("hbm_peak_gb", 0.0),
|
| 794 |
+
waste_lines=_format_waste(waste),
|
| 795 |
+
recoverable_sorted=_recoverable_sorted(waste),
|
| 796 |
+
history_lines=_format_history(history),
|
| 797 |
+
)
|
| 798 |
+
backend.add_user_message(prompt)
|
| 799 |
+
turn = await backend.next_turn(tool_schemas=[])
|
| 800 |
+
raw = " ".join(turn.text_blocks).strip()
|
| 801 |
+
|
| 802 |
+
obj = _extract_json_object(raw)
|
| 803 |
+
if obj is None:
|
| 804 |
+
print(f" LLM response was not parseable JSON. Raw: {raw[:300]!r}")
|
| 805 |
+
return None
|
| 806 |
+
|
| 807 |
+
name = (obj.get("name") or "").strip()
|
| 808 |
+
if not name or name.upper() == "STOP":
|
| 809 |
+
print(f" LLM signaled STOP: {obj.get('rationale', '(no rationale)')}")
|
| 810 |
+
return None
|
| 811 |
+
|
| 812 |
+
subs_raw = obj.get("substitutions") or []
|
| 813 |
+
envs = obj.get("env_vars") or {}
|
| 814 |
+
if not subs_raw and not envs:
|
| 815 |
+
print(f" LLM returned an empty experiment ({name}); skipping")
|
| 816 |
+
return None
|
| 817 |
+
|
| 818 |
+
subs: list[tuple[str, str]] = []
|
| 819 |
+
for entry in subs_raw:
|
| 820 |
+
if isinstance(entry, list) and len(entry) == 2:
|
| 821 |
+
subs.append((str(entry[0]), str(entry[1])))
|
| 822 |
+
elif isinstance(entry, dict) and "pattern" in entry and "replacement" in entry:
|
| 823 |
+
subs.append((str(entry["pattern"]), str(entry["replacement"])))
|
| 824 |
+
|
| 825 |
+
cleaned_envs = _sanitize_env_vars(envs, context=name)
|
| 826 |
+
if not subs and not cleaned_envs:
|
| 827 |
+
# Everything got dropped during sanitization (bad env names + no
|
| 828 |
+
# valid substitutions). Treat as a no-op rather than benchmarking
|
| 829 |
+
# an unchanged workload.
|
| 830 |
+
print(f" LLM experiment {name!r} had nothing valid after sanitization; skipping")
|
| 831 |
+
return None
|
| 832 |
+
|
| 833 |
+
return Experiment(
|
| 834 |
+
name=name,
|
| 835 |
+
description=obj.get("description") or name,
|
| 836 |
+
rationale=str(obj.get("rationale") or ""),
|
| 837 |
+
substitutions=subs,
|
| 838 |
+
env_vars=cleaned_envs,
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
|
| 842 |
+
# ---------------------------------------------------------------------------
|
| 843 |
+
# llm-explore mode: ask for K candidates per iteration
|
| 844 |
+
# ---------------------------------------------------------------------------
|
| 845 |
+
|
| 846 |
+
|
| 847 |
+
_LLM_EXPLORE_SYSTEM_PROMPT = """\
|
| 848 |
+
You are an expert at tuning AMD MI300X (ROCm 7.0, CDNA3 arch, 192 GB
|
| 849 |
+
HBM3) training workloads. The user is running a multi-candidate
|
| 850 |
+
exploration: on every iteration you suggest K STRUCTURALLY-DIFFERENT
|
| 851 |
+
candidate changes, the user benchmarks all of them, and the best one
|
| 852 |
+
is accepted (if it beats the current best by the threshold).
|
| 853 |
+
|
| 854 |
+
Your output MUST be a JSON ARRAY of K objects, no prose, no markdown
|
| 855 |
+
fences, just the array:
|
| 856 |
+
|
| 857 |
+
[
|
| 858 |
+
{"name": "...", "rationale": "...", "substitutions": [["regex", "repl"]], "env_vars": {"VAR": "value"}},
|
| 859 |
+
{"name": "...", "rationale": "...", "substitutions": [["regex", "repl"]], "env_vars": {"VAR": "value"}},
|
| 860 |
+
{"name": "...", "rationale": "...", "substitutions": [["regex", "repl"]], "env_vars": {"VAR": "value"}}
|
| 861 |
+
]
|
| 862 |
+
|
| 863 |
+
CRITICAL output rules:
|
| 864 |
+
|
| 865 |
+
1. Each candidate must target a DIFFERENT waste bucket or parameter
|
| 866 |
+
category than the others. Diversity beats redundancy — don't propose
|
| 867 |
+
three batch-size bumps; propose one batch bump, one env var, one
|
| 868 |
+
precision/attention/dataloader change.
|
| 869 |
+
|
| 870 |
+
2. env_vars keys are LITERAL POSIX shell environment variable names —
|
| 871 |
+
they MUST match the regex [A-Za-z_][A-Za-z0-9_]*. NEVER prefix them
|
| 872 |
+
with "env_vars." or any other dotted path. NEVER include "=" or
|
| 873 |
+
shell syntax in the key. If you want to invoke a CLI tool like
|
| 874 |
+
`numactl`, that's NOT an env var — skip the candidate entirely.
|
| 875 |
+
Wrong: {"env_vars.MIOPEN_FIND_MODE": "3"}
|
| 876 |
+
Wrong: {"NUMACTL_INTERLEAVE=1": "numactl --interleave=all"}
|
| 877 |
+
Right: {"MIOPEN_FIND_MODE": "3"}
|
| 878 |
+
|
| 879 |
+
3. substitutions are (regex_pattern, replacement) pairs applied with
|
| 880 |
+
re.subn. Patterns must match at least one occurrence — if zero
|
| 881 |
+
matches, that candidate is skipped.
|
| 882 |
+
|
| 883 |
+
4. NEVER propose a (substitutions, env_vars) combination that already
|
| 884 |
+
appears in history with outcome rejected/crashed. Diversify within
|
| 885 |
+
the array AND across the run.
|
| 886 |
+
|
| 887 |
+
5. If you genuinely cannot find K productive candidates, output fewer
|
| 888 |
+
(e.g. 2 if K=3). The user will benchmark whatever you provide. If
|
| 889 |
+
you have zero productive candidates, output:
|
| 890 |
+
[{"name": "STOP", "rationale": "<why>", "substitutions": [], "env_vars": {}}]
|
| 891 |
+
|
| 892 |
+
CONCRETE OUTPUT EXAMPLES (for K=3):
|
| 893 |
+
|
| 894 |
+
[
|
| 895 |
+
{"name": "bf16_over_fp16",
|
| 896 |
+
"rationale": "Largest recoverable bucket is precision_path; CDNA3 prefers bf16.",
|
| 897 |
+
"substitutions": [["fp16=True", "bf16=True"], ["torch_dtype=torch\\\\.float16", "torch_dtype=torch.bfloat16"]],
|
| 898 |
+
"env_vars": {}},
|
| 899 |
+
{"name": "batch_size_16",
|
| 900 |
+
"rationale": "HBM peak well under 192 GB; bigger batch saturates the GPU.",
|
| 901 |
+
"substitutions": [["per_device_train_batch_size=\\\\d+", "per_device_train_batch_size=16"]],
|
| 902 |
+
"env_vars": {}},
|
| 903 |
+
{"name": "prefer_hipblaslt",
|
| 904 |
+
"rationale": "hipBLASLt outperforms rocBLAS on Qwen GEMM shapes.",
|
| 905 |
+
"substitutions": [],
|
| 906 |
+
"env_vars": {"TORCH_BLAS_PREFER_HIPBLASLT": "1"}}
|
| 907 |
+
]
|
| 908 |
+
"""
|
| 909 |
+
|
| 910 |
+
|
| 911 |
+
_LLM_EXPLORE_USER_TEMPLATE = """\
|
| 912 |
+
Hardware facts (use these — do not contradict):
|
| 913 |
+
- AMD MI300X, CDNA3 architecture, 192 GB HBM3
|
| 914 |
+
- bf16 throughput on CDNA3 ≈ same as fp16, > fp32 (matrix engine is fp16/bf16/fp8 native)
|
| 915 |
+
- fp32 is the SLOWEST option on this arch — never suggest it as an improvement
|
| 916 |
+
|
| 917 |
+
Known incompatibilities for THIS workload (peft + LoRA on transformers Trainer):
|
| 918 |
+
{incompatibilities}
|
| 919 |
+
|
| 920 |
+
KB rules (one-liner per rule, for grounding):
|
| 921 |
+
{kb_summary}
|
| 922 |
+
|
| 923 |
+
Current accepted workload state — the literal values in the script
|
| 924 |
+
after every change accepted so far. Each candidate you propose should
|
| 925 |
+
mutate one of these (or set an env var). DO NOT propose a value that's
|
| 926 |
+
already present here.
|
| 927 |
+
{tunables}
|
| 928 |
+
|
| 929 |
+
Latest benchmark (this is the result of the most recent ACCEPTED state):
|
| 930 |
+
- tokens_per_sec: {tps:.1f}
|
| 931 |
+
- mfu_pct: {mfu:.2f} (% of MI300X dense bf16 peak; healthy LoRA ranges 30-50%)
|
| 932 |
+
- gpu_util_pct: {util:.1f}
|
| 933 |
+
- hbm_peak_gb: {hbm:.2f}
|
| 934 |
+
- waste_budget (seconds/step):
|
| 935 |
+
{waste_lines}
|
| 936 |
+
|
| 937 |
+
Sorted recoverable waste (largest first — go after these):
|
| 938 |
+
{recoverable_sorted}
|
| 939 |
+
|
| 940 |
+
Previously rejected (full fingerprint — DO NOT repropose any of these):
|
| 941 |
+
{rejected_fingerprints}
|
| 942 |
+
|
| 943 |
+
History of changes already tried this run (newest first; outcomes are
|
| 944 |
+
"accepted" / "rejected" / "crashed" / "skipped"):
|
| 945 |
+
{history_lines}
|
| 946 |
+
|
| 947 |
+
Suggest {num_candidates} STRUCTURALLY-DIFFERENT candidate changes.
|
| 948 |
+
Each must target a different waste bucket or parameter category. JSON
|
| 949 |
+
array only.
|
| 950 |
+
"""
|
| 951 |
+
|
| 952 |
+
|
| 953 |
+
async def _ask_llm_for_experiments(
|
| 954 |
+
backend,
|
| 955 |
+
*,
|
| 956 |
+
kb_summary: str,
|
| 957 |
+
source: str,
|
| 958 |
+
metrics: dict,
|
| 959 |
+
history: list[dict],
|
| 960 |
+
num_candidates: int,
|
| 961 |
+
) -> list[Experiment]:
|
| 962 |
+
"""One LLM turn → up to `num_candidates` Experiments.
|
| 963 |
+
|
| 964 |
+
Returns an empty list on parse failure or STOP signal.
|
| 965 |
+
"""
|
| 966 |
+
waste = metrics.get("waste_budget") or {}
|
| 967 |
+
prompt = _LLM_EXPLORE_USER_TEMPLATE.format(
|
| 968 |
+
num_candidates=num_candidates,
|
| 969 |
+
incompatibilities="\n".join(f"- {line}" for line in _KNOWN_INCOMPATIBILITIES),
|
| 970 |
+
kb_summary=kb_summary,
|
| 971 |
+
tunables=_tunables_summary(source),
|
| 972 |
+
tps=metrics.get("tokens_per_sec", 0.0),
|
| 973 |
+
mfu=metrics.get("mfu_pct", 0.0),
|
| 974 |
+
util=metrics.get("gpu_util_pct", 0.0),
|
| 975 |
+
hbm=metrics.get("hbm_peak_gb", 0.0),
|
| 976 |
+
waste_lines=_format_waste(waste),
|
| 977 |
+
recoverable_sorted=_recoverable_sorted(waste),
|
| 978 |
+
rejected_fingerprints=_format_rejected_fingerprints(history),
|
| 979 |
+
history_lines=_format_history(history),
|
| 980 |
+
)
|
| 981 |
+
backend.add_user_message(prompt)
|
| 982 |
+
turn = await backend.next_turn(tool_schemas=[])
|
| 983 |
+
raw = " ".join(turn.text_blocks).strip()
|
| 984 |
+
|
| 985 |
+
arr = _extract_json_array(raw)
|
| 986 |
+
if not arr:
|
| 987 |
+
print(f" LLM response was not parseable JSON array. Raw: {raw[:300]!r}")
|
| 988 |
+
return []
|
| 989 |
+
|
| 990 |
+
experiments: list[Experiment] = []
|
| 991 |
+
for obj in arr:
|
| 992 |
+
if not isinstance(obj, dict):
|
| 993 |
+
continue
|
| 994 |
+
name = (obj.get("name") or "").strip()
|
| 995 |
+
if not name:
|
| 996 |
+
continue
|
| 997 |
+
if name.upper() == "STOP":
|
| 998 |
+
print(f" LLM signaled STOP: {obj.get('rationale', '(no rationale)')}")
|
| 999 |
+
return []
|
| 1000 |
+
subs_raw = obj.get("substitutions") or []
|
| 1001 |
+
envs_raw = obj.get("env_vars") or {}
|
| 1002 |
+
if not subs_raw and not envs_raw:
|
| 1003 |
+
continue
|
| 1004 |
+
subs = []
|
| 1005 |
+
for entry in subs_raw:
|
| 1006 |
+
if isinstance(entry, list) and len(entry) == 2:
|
| 1007 |
+
subs.append((str(entry[0]), str(entry[1])))
|
| 1008 |
+
elif isinstance(entry, dict) and "pattern" in entry and "replacement" in entry:
|
| 1009 |
+
subs.append((str(entry["pattern"]), str(entry["replacement"])))
|
| 1010 |
+
cleaned_envs = _sanitize_env_vars(envs_raw, context=name)
|
| 1011 |
+
if not subs and not cleaned_envs:
|
| 1012 |
+
print(f" candidate {name!r} had nothing valid after sanitization; dropping")
|
| 1013 |
+
continue
|
| 1014 |
+
experiments.append(
|
| 1015 |
+
Experiment(
|
| 1016 |
+
name=name,
|
| 1017 |
+
description=obj.get("description") or name,
|
| 1018 |
+
rationale=str(obj.get("rationale") or ""),
|
| 1019 |
+
substitutions=subs,
|
| 1020 |
+
env_vars=cleaned_envs,
|
| 1021 |
+
)
|
| 1022 |
+
)
|
| 1023 |
+
return experiments
|
| 1024 |
+
|
| 1025 |
+
|
| 1026 |
+
def _extract_json_array(text: str) -> list | None:
|
| 1027 |
+
"""Pull the first JSON array out of an LLM response, tolerating
|
| 1028 |
+
markdown fences and leading prose. Returns None if nothing parseable."""
|
| 1029 |
+
if not text:
|
| 1030 |
+
return None
|
| 1031 |
+
fence_match = re.search(r"```(?:json)?\s*(\[.*?\])\s*```", text, re.DOTALL)
|
| 1032 |
+
if fence_match:
|
| 1033 |
+
try:
|
| 1034 |
+
obj = json.loads(fence_match.group(1))
|
| 1035 |
+
if isinstance(obj, list):
|
| 1036 |
+
return obj
|
| 1037 |
+
except json.JSONDecodeError:
|
| 1038 |
+
pass
|
| 1039 |
+
depth = 0
|
| 1040 |
+
start = -1
|
| 1041 |
+
for i, ch in enumerate(text):
|
| 1042 |
+
if ch == "[":
|
| 1043 |
+
if depth == 0:
|
| 1044 |
+
start = i
|
| 1045 |
+
depth += 1
|
| 1046 |
+
elif ch == "]":
|
| 1047 |
+
depth -= 1
|
| 1048 |
+
if depth == 0 and start >= 0:
|
| 1049 |
+
blob = text[start : i + 1]
|
| 1050 |
+
try:
|
| 1051 |
+
obj = json.loads(blob)
|
| 1052 |
+
if isinstance(obj, list):
|
| 1053 |
+
return obj
|
| 1054 |
+
except json.JSONDecodeError:
|
| 1055 |
+
start = -1
|
| 1056 |
+
continue
|
| 1057 |
+
return None
|
| 1058 |
+
|
| 1059 |
+
|
| 1060 |
+
# ---------------------------------------------------------------------------
|
| 1061 |
+
# Dedup + history utilities (used by all LLM modes)
|
| 1062 |
+
# ---------------------------------------------------------------------------
|
| 1063 |
+
|
| 1064 |
+
|
| 1065 |
+
def _experiment_fingerprint(exp: Experiment) -> tuple:
|
| 1066 |
+
"""Hashable identity for an experiment — substitutions + env_vars,
|
| 1067 |
+
NOT name (the LLM tends to give the same change different names)."""
|
| 1068 |
+
subs = tuple(sorted(tuple(s) for s in exp.substitutions))
|
| 1069 |
+
envs = tuple(sorted(exp.env_vars.items()))
|
| 1070 |
+
return (subs, envs)
|
| 1071 |
+
|
| 1072 |
+
|
| 1073 |
+
def _build_merged_experiment(
|
| 1074 |
+
exps: list[Experiment], base_source: str
|
| 1075 |
+
) -> tuple[Experiment | None, str]:
|
| 1076 |
+
"""Try to combine 2+ experiments into one. The merged experiment
|
| 1077 |
+
applies all of their substitutions in sequence and unions their
|
| 1078 |
+
env_vars. Returns (merged, "") on success, (None, reason) when the
|
| 1079 |
+
merge is structurally unsafe — caller should fall back to using just
|
| 1080 |
+
the individual winner.
|
| 1081 |
+
|
| 1082 |
+
Conflict detection:
|
| 1083 |
+
- A later substitution's pattern must still match after earlier
|
| 1084 |
+
substitutions have been applied (zero matches → conflict, e.g.
|
| 1085 |
+
cand A rewrote `fp16=True` and cand B was also targeting it).
|
| 1086 |
+
- Env var keys with conflicting values (same name, different value)
|
| 1087 |
+
→ conflict.
|
| 1088 |
+
- Bad regex anywhere → conflict.
|
| 1089 |
+
"""
|
| 1090 |
+
if len(exps) < 2:
|
| 1091 |
+
return None, "need at least 2 experiments"
|
| 1092 |
+
|
| 1093 |
+
merged_subs: list[tuple[str, str]] = []
|
| 1094 |
+
merged_envs: dict[str, str] = {}
|
| 1095 |
+
test_source = base_source
|
| 1096 |
+
|
| 1097 |
+
for exp in exps:
|
| 1098 |
+
for pattern, replacement in exp.substitutions:
|
| 1099 |
+
try:
|
| 1100 |
+
new_source, n = re.subn(pattern, replacement, test_source)
|
| 1101 |
+
except re.error as e:
|
| 1102 |
+
return None, f"bad regex in '{exp.name}': {e}"
|
| 1103 |
+
if n == 0:
|
| 1104 |
+
return None, (
|
| 1105 |
+
f"'{exp.name}' substitution {pattern!r} no longer matches "
|
| 1106 |
+
"after prior merges (likely overwrites an earlier change)"
|
| 1107 |
+
)
|
| 1108 |
+
test_source = new_source
|
| 1109 |
+
merged_subs.append((pattern, replacement))
|
| 1110 |
+
for k, v in exp.env_vars.items():
|
| 1111 |
+
if k in merged_envs and merged_envs[k] != v:
|
| 1112 |
+
return None, (
|
| 1113 |
+
f"env var conflict on {k!r}: {merged_envs[k]!r} vs {v!r}"
|
| 1114 |
+
)
|
| 1115 |
+
merged_envs[k] = v
|
| 1116 |
+
|
| 1117 |
+
short_names = "+".join(e.name[:14] for e in exps)
|
| 1118 |
+
full_names = " + ".join(e.name for e in exps)
|
| 1119 |
+
return (
|
| 1120 |
+
Experiment(
|
| 1121 |
+
name=f"merge[{short_names}]"[:60],
|
| 1122 |
+
description=f"Merged: {full_names}",
|
| 1123 |
+
rationale=(
|
| 1124 |
+
f"Combined {len(exps)} candidates that each had positive delta "
|
| 1125 |
+
"against the current best this iteration. Tests the compound "
|
| 1126 |
+
"effect; falls back to the individual winner if it doesn't help."
|
| 1127 |
+
),
|
| 1128 |
+
substitutions=merged_subs,
|
| 1129 |
+
env_vars=merged_envs,
|
| 1130 |
+
),
|
| 1131 |
+
"",
|
| 1132 |
+
)
|
| 1133 |
+
|
| 1134 |
+
|
| 1135 |
+
def _is_duplicate_of_history(exp: Experiment, history: list[dict]) -> dict | None:
|
| 1136 |
+
"""If `exp` matches a prior history entry by fingerprint, return that
|
| 1137 |
+
entry. Otherwise None."""
|
| 1138 |
+
fp = _experiment_fingerprint(exp)
|
| 1139 |
+
for h in history:
|
| 1140 |
+
h_subs = tuple(
|
| 1141 |
+
sorted(
|
| 1142 |
+
(str(s[0]), str(s[1]))
|
| 1143 |
+
for s in (h.get("substitutions") or [])
|
| 1144 |
+
if isinstance(s, (list, tuple)) and len(s) == 2
|
| 1145 |
+
)
|
| 1146 |
+
)
|
| 1147 |
+
h_envs = tuple(sorted((h.get("env_vars") or {}).items()))
|
| 1148 |
+
if fp == (h_subs, h_envs):
|
| 1149 |
+
return h
|
| 1150 |
+
return None
|
| 1151 |
+
|
| 1152 |
+
|
| 1153 |
+
def _format_rejected_fingerprints(history: list[dict]) -> str:
|
| 1154 |
+
"""Compact list of every (substitutions, env_vars) the LLM has already
|
| 1155 |
+
tried with outcome rejected/crashed/skipped — so it can't propose them
|
| 1156 |
+
again under a different name."""
|
| 1157 |
+
seen: set[tuple] = set()
|
| 1158 |
+
lines: list[str] = []
|
| 1159 |
+
for h in history:
|
| 1160 |
+
outcome = h.get("outcome", "")
|
| 1161 |
+
if outcome not in ("rejected", "crashed", "skipped"):
|
| 1162 |
+
continue
|
| 1163 |
+
subs = tuple(
|
| 1164 |
+
sorted(
|
| 1165 |
+
(str(s[0]), str(s[1]))
|
| 1166 |
+
for s in (h.get("substitutions") or [])
|
| 1167 |
+
if isinstance(s, (list, tuple)) and len(s) == 2
|
| 1168 |
+
)
|
| 1169 |
+
)
|
| 1170 |
+
envs = tuple(sorted((h.get("env_vars") or {}).items()))
|
| 1171 |
+
fp = (subs, envs)
|
| 1172 |
+
if fp in seen:
|
| 1173 |
+
continue
|
| 1174 |
+
seen.add(fp)
|
| 1175 |
+
lines.append(f" - {outcome:9s} subs={list(subs)} env={dict(envs)}")
|
| 1176 |
+
if not lines:
|
| 1177 |
+
return " (none yet)"
|
| 1178 |
+
return "\n".join(lines)
|
| 1179 |
+
|
| 1180 |
+
|
| 1181 |
+
def _print_waste(metrics: dict, prefix: str = " waste: ") -> None:
|
| 1182 |
+
"""Print a one-line summary of waste_budget — useful is highlighted
|
| 1183 |
+
first, then non-zero recoverable buckets sorted by size."""
|
| 1184 |
+
wb = metrics.get("waste_budget") or {}
|
| 1185 |
+
if not wb:
|
| 1186 |
+
return
|
| 1187 |
+
parts = [f"useful_gpu={wb.get('useful_gpu', 0.0):.3f}"]
|
| 1188 |
+
others = [(k, v) for k, v in wb.items() if k != "useful_gpu" and isinstance(v, (int, float)) and v > 0]
|
| 1189 |
+
others.sort(key=lambda kv: kv[1], reverse=True)
|
| 1190 |
+
parts.extend(f"{k}={v:.3f}" for k, v in others)
|
| 1191 |
+
print(prefix + ", ".join(parts))
|
| 1192 |
+
|
| 1193 |
+
|
| 1194 |
+
# ---------------------------------------------------------------------------
|
| 1195 |
+
# JSON object extractor (used by single-experiment llm mode)
|
| 1196 |
+
# ---------------------------------------------------------------------------
|
| 1197 |
+
|
| 1198 |
+
|
| 1199 |
+
def _extract_json_object(text: str) -> dict | None:
|
| 1200 |
+
"""Pull the first JSON object out of an LLM response, tolerating
|
| 1201 |
+
markdown fences / leading prose."""
|
| 1202 |
+
if not text:
|
| 1203 |
+
return None
|
| 1204 |
+
# strip ```json ... ``` fences if present
|
| 1205 |
+
fence_match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
|
| 1206 |
+
if fence_match:
|
| 1207 |
+
try:
|
| 1208 |
+
return json.loads(fence_match.group(1))
|
| 1209 |
+
except json.JSONDecodeError:
|
| 1210 |
+
pass
|
| 1211 |
+
# otherwise grab the first balanced { ... }
|
| 1212 |
+
depth = 0
|
| 1213 |
+
start = -1
|
| 1214 |
+
for i, ch in enumerate(text):
|
| 1215 |
+
if ch == "{":
|
| 1216 |
+
if depth == 0:
|
| 1217 |
+
start = i
|
| 1218 |
+
depth += 1
|
| 1219 |
+
elif ch == "}":
|
| 1220 |
+
depth -= 1
|
| 1221 |
+
if depth == 0 and start >= 0:
|
| 1222 |
+
blob = text[start : i + 1]
|
| 1223 |
+
try:
|
| 1224 |
+
return json.loads(blob)
|
| 1225 |
+
except json.JSONDecodeError:
|
| 1226 |
+
start = -1
|
| 1227 |
+
continue
|
| 1228 |
+
return None
|
| 1229 |
+
|
| 1230 |
+
|
| 1231 |
+
# ---------------------------------------------------------------------------
|
| 1232 |
+
# Main
|
| 1233 |
+
# ---------------------------------------------------------------------------
|
| 1234 |
+
|
| 1235 |
+
|
| 1236 |
+
def main() -> int:
|
| 1237 |
+
p = argparse.ArgumentParser(
|
| 1238 |
+
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
| 1239 |
+
)
|
| 1240 |
+
p.add_argument(
|
| 1241 |
+
"workload",
|
| 1242 |
+
type=Path,
|
| 1243 |
+
nargs="?",
|
| 1244 |
+
default=None,
|
| 1245 |
+
help=(
|
| 1246 |
+
"Path to a workload script (omit if using --model). When given, "
|
| 1247 |
+
"the script is used as-is for the baseline benchmark."
|
| 1248 |
+
),
|
| 1249 |
+
)
|
| 1250 |
+
p.add_argument(
|
| 1251 |
+
"--model",
|
| 1252 |
+
type=str,
|
| 1253 |
+
default=None,
|
| 1254 |
+
help=(
|
| 1255 |
+
"HuggingFace model id (e.g. Qwen/Qwen2.5-7B-Instruct, "
|
| 1256 |
+
"meta-llama/Llama-3.2-3B). Generates a baseline workload from "
|
| 1257 |
+
"workloads/train_qwen_lora.py with this MODEL_ID substituted in. "
|
| 1258 |
+
"Use this OR a workload path, not both. For gated models, "
|
| 1259 |
+
"ensure HF_TOKEN is set in your shell."
|
| 1260 |
+
),
|
| 1261 |
+
)
|
| 1262 |
+
p.add_argument(
|
| 1263 |
+
"--mode",
|
| 1264 |
+
choices=("hardcoded", "llm", "llm-explore"),
|
| 1265 |
+
default="hardcoded",
|
| 1266 |
+
help=(
|
| 1267 |
+
"hardcoded (default): walk through the priority-ordered EXPERIMENTS list. "
|
| 1268 |
+
"llm: ask the LLM for ONE next experiment per iteration (greedy). "
|
| 1269 |
+
"llm-explore: ask for K candidates per iteration, benchmark all, keep "
|
| 1270 |
+
"the best (slower but better at finding interaction effects)."
|
| 1271 |
+
),
|
| 1272 |
+
)
|
| 1273 |
+
p.add_argument(
|
| 1274 |
+
"--candidates-per-iteration",
|
| 1275 |
+
type=int,
|
| 1276 |
+
default=3,
|
| 1277 |
+
help="Only used when --mode llm-explore. Default 3.",
|
| 1278 |
+
)
|
| 1279 |
+
p.add_argument("--steps", type=int, default=20, help="Steps per benchmark")
|
| 1280 |
+
p.add_argument(
|
| 1281 |
+
"--max-iterations",
|
| 1282 |
+
type=int,
|
| 1283 |
+
default=0,
|
| 1284 |
+
help=(
|
| 1285 |
+
"Cap on experiments to try. Default: len(EXPERIMENTS) for hardcoded mode, "
|
| 1286 |
+
"10 for llm mode."
|
| 1287 |
+
),
|
| 1288 |
+
)
|
| 1289 |
+
p.add_argument(
|
| 1290 |
+
"--early-stop-after",
|
| 1291 |
+
type=int,
|
| 1292 |
+
default=3,
|
| 1293 |
+
help=(
|
| 1294 |
+
"Stop after N consecutive non-improvements. Crashes do NOT count "
|
| 1295 |
+
"toward this — crashes mean the change was structurally bad, not "
|
| 1296 |
+
"that we've exhausted ideas."
|
| 1297 |
+
),
|
| 1298 |
+
)
|
| 1299 |
+
p.add_argument(
|
| 1300 |
+
"--max-crashes",
|
| 1301 |
+
type=int,
|
| 1302 |
+
default=4,
|
| 1303 |
+
help=(
|
| 1304 |
+
"Stop after N total subprocess crashes (separate from "
|
| 1305 |
+
"--early-stop-after). Default 4 leaves room for the LLM to try "
|
| 1306 |
+
"structurally different changes after a bad one."
|
| 1307 |
+
),
|
| 1308 |
+
)
|
| 1309 |
+
p.add_argument(
|
| 1310 |
+
"--improvement-threshold",
|
| 1311 |
+
type=float,
|
| 1312 |
+
default=0.0,
|
| 1313 |
+
help=(
|
| 1314 |
+
"Min %% improvement over current best to accept. Default 0.0 "
|
| 1315 |
+
"(any positive delta wins). Bump to 1.0 if your benchmarks are "
|
| 1316 |
+
"noisy and you want to ignore sub-1%% deltas."
|
| 1317 |
+
),
|
| 1318 |
+
)
|
| 1319 |
+
p.add_argument(
|
| 1320 |
+
"--events",
|
| 1321 |
+
type=Path,
|
| 1322 |
+
default=None,
|
| 1323 |
+
help=(
|
| 1324 |
+
"Optional NDJSON event stream output. If set, the script appends "
|
| 1325 |
+
"one JSON event per line at baseline / iter / candidate / summary "
|
| 1326 |
+
"milestones. Used by the Streamlit UI; CLI users don't need this."
|
| 1327 |
+
),
|
| 1328 |
+
)
|
| 1329 |
+
args = p.parse_args()
|
| 1330 |
+
if args.events is not None:
|
| 1331 |
+
global _EVENTS_PATH
|
| 1332 |
+
_EVENTS_PATH = args.events
|
| 1333 |
+
try:
|
| 1334 |
+
args.events.write_text("") # truncate any prior contents
|
| 1335 |
+
except OSError as exc:
|
| 1336 |
+
sys.stderr.write(f"--events: cannot open {args.events} for writing ({exc})\n")
|
| 1337 |
+
return 1
|
| 1338 |
+
if args.max_iterations <= 0:
|
| 1339 |
+
if args.mode == "hardcoded":
|
| 1340 |
+
args.max_iterations = len(EXPERIMENTS)
|
| 1341 |
+
elif args.mode == "llm-explore":
|
| 1342 |
+
args.max_iterations = 5 # K candidates per iter so 5 iters = 5K benchmarks
|
| 1343 |
+
else:
|
| 1344 |
+
args.max_iterations = 10
|
| 1345 |
+
|
| 1346 |
+
# Validate that exactly one workload source was provided
|
| 1347 |
+
if args.workload is None and args.model is None:
|
| 1348 |
+
sys.stderr.write(
|
| 1349 |
+
"Pass either a workload path or --model MODEL_ID. "
|
| 1350 |
+
"Examples:\n"
|
| 1351 |
+
" python scripts/auto_tune.py workloads/train_qwen_lora.py\n"
|
| 1352 |
+
" python scripts/auto_tune.py --model Qwen/Qwen2.5-7B-Instruct\n"
|
| 1353 |
+
)
|
| 1354 |
+
return 1
|
| 1355 |
+
if args.workload is not None and args.model is not None:
|
| 1356 |
+
sys.stderr.write(
|
| 1357 |
+
"Pass EITHER a workload path OR --model, not both.\n"
|
| 1358 |
+
)
|
| 1359 |
+
return 1
|
| 1360 |
+
if not GOBLIN_RUNNER.exists():
|
| 1361 |
+
sys.stderr.write(f"goblin_runner.sh not found at {GOBLIN_RUNNER}\n")
|
| 1362 |
+
return 1
|
| 1363 |
+
|
| 1364 |
+
workspace = Path(tempfile.mkdtemp(prefix="auto_tune_workloads_"))
|
| 1365 |
+
|
| 1366 |
+
if args.workload is not None:
|
| 1367 |
+
workload = args.workload.resolve()
|
| 1368 |
+
if not workload.exists():
|
| 1369 |
+
sys.stderr.write(f"workload not found: {workload}\n")
|
| 1370 |
+
return 1
|
| 1371 |
+
workload_label = str(workload)
|
| 1372 |
+
else:
|
| 1373 |
+
# Generate baseline workload from --model
|
| 1374 |
+
generated = workspace / "_generated_baseline.py"
|
| 1375 |
+
workload = _generate_workload_from_model(args.model, generated)
|
| 1376 |
+
workload_label = f"(generated from --model {args.model})\n "
|
| 1377 |
+
workload_label += f" {workload}\n "
|
| 1378 |
+
workload_label += f" template: {_DEFAULT_WORKLOAD_TEMPLATE}"
|
| 1379 |
+
|
| 1380 |
+
_emit({
|
| 1381 |
+
"type": "started",
|
| 1382 |
+
"mode": args.mode,
|
| 1383 |
+
"workload": str(workload),
|
| 1384 |
+
"model": args.model,
|
| 1385 |
+
"steps": args.steps,
|
| 1386 |
+
"max_iterations": args.max_iterations,
|
| 1387 |
+
"early_stop_after": args.early_stop_after,
|
| 1388 |
+
"max_crashes": args.max_crashes,
|
| 1389 |
+
"improvement_threshold": args.improvement_threshold,
|
| 1390 |
+
"candidates_per_iteration": (
|
| 1391 |
+
args.candidates_per_iteration if args.mode == "llm-explore" else 1
|
| 1392 |
+
),
|
| 1393 |
+
"workspace": str(workspace),
|
| 1394 |
+
})
|
| 1395 |
+
print(f"Auto-tune workspace: {workspace}")
|
| 1396 |
+
print(f"Mode: {args.mode}")
|
| 1397 |
+
print(f"Workload: {workload_label}")
|
| 1398 |
+
print(f"Steps per benchmark: {args.steps}")
|
| 1399 |
+
print(f"Max iterations: {args.max_iterations}")
|
| 1400 |
+
print(f"Early stop after: {args.early_stop_after} non-improvements")
|
| 1401 |
+
print(f"Max crashes: {args.max_crashes} total")
|
| 1402 |
+
print(f"Accept threshold: {args.improvement_threshold:.1f}%\n")
|
| 1403 |
+
|
| 1404 |
+
# LLM mode setup happens before the baseline so we fail fast on missing
|
| 1405 |
+
# credentials rather than after burning a baseline benchmark. Each LLM
|
| 1406 |
+
# mode gets its own system prompt — the explore mode needs a much
|
| 1407 |
+
# larger token budget to emit K JSON objects.
|
| 1408 |
+
llm_backend = None
|
| 1409 |
+
kb_summary = ""
|
| 1410 |
+
if args.mode == "llm":
|
| 1411 |
+
llm_backend = _build_llm_backend(_LLM_SYSTEM_PROMPT, max_tokens=1024)
|
| 1412 |
+
kb_summary = _kb_summary(REPO_ROOT / "kb" / "rocm_rules.yaml")
|
| 1413 |
+
print("LLM backend ready (single-candidate). KB summary loaded.\n")
|
| 1414 |
+
elif args.mode == "llm-explore":
|
| 1415 |
+
llm_backend = _build_llm_backend(_LLM_EXPLORE_SYSTEM_PROMPT, max_tokens=2048)
|
| 1416 |
+
kb_summary = _kb_summary(REPO_ROOT / "kb" / "rocm_rules.yaml")
|
| 1417 |
+
print(
|
| 1418 |
+
f"LLM backend ready (multi-candidate, K={args.candidates_per_iteration}). "
|
| 1419 |
+
"KB summary loaded.\n"
|
| 1420 |
+
)
|
| 1421 |
+
|
| 1422 |
+
baseline_source = workload.read_text()
|
| 1423 |
+
baseline_path = workspace / "00_baseline.py"
|
| 1424 |
+
baseline_path.write_text(baseline_source)
|
| 1425 |
+
|
| 1426 |
+
print("=" * 60)
|
| 1427 |
+
print("Baseline benchmark")
|
| 1428 |
+
print("=" * 60)
|
| 1429 |
+
baseline = benchmark(baseline_path, args.steps, {})
|
| 1430 |
+
if baseline is None:
|
| 1431 |
+
sys.stderr.write("Baseline benchmark failed; cannot continue.\n")
|
| 1432 |
+
return 1
|
| 1433 |
+
|
| 1434 |
+
baseline_tps = baseline["tokens_per_sec"]
|
| 1435 |
+
print(f" tokens/sec: {baseline_tps:.1f}")
|
| 1436 |
+
print(f" mfu_pct: {baseline.get('mfu_pct', 0.0):.2f}")
|
| 1437 |
+
print(f" hbm_peak_gb: {baseline['hbm_peak_gb']:.2f}")
|
| 1438 |
+
print(f" gpu_util_pct: {baseline['gpu_util_pct']:.1f}")
|
| 1439 |
+
print(
|
| 1440 |
+
" waste_budget: "
|
| 1441 |
+
+ ", ".join(f"{k}={v:.3f}" for k, v in baseline["waste_budget"].items() if v > 0)
|
| 1442 |
+
)
|
| 1443 |
+
_emit({"type": "baseline", "metrics": baseline})
|
| 1444 |
+
|
| 1445 |
+
best_source = baseline_source
|
| 1446 |
+
best_tps = baseline_tps
|
| 1447 |
+
best_env: dict[str, str] = {}
|
| 1448 |
+
last_metrics = baseline
|
| 1449 |
+
accepted: list[tuple[str, float, float]] = [] # (name, tps, delta_pct)
|
| 1450 |
+
rejected: list[tuple[str, str]] = [] # (name, reason)
|
| 1451 |
+
history: list[dict] = [] # for LLM context
|
| 1452 |
+
consecutive_no_improvement = 0
|
| 1453 |
+
total_crashes = 0
|
| 1454 |
+
file_counter = 0 # monotonically increases across all candidates
|
| 1455 |
+
|
| 1456 |
+
for i in range(args.max_iterations):
|
| 1457 |
+
# ---- Get candidates list (1 for hardcoded/llm, K for llm-explore) ----
|
| 1458 |
+
if args.mode == "hardcoded":
|
| 1459 |
+
if i >= len(EXPERIMENTS):
|
| 1460 |
+
print("\nReached end of EXPERIMENTS list.")
|
| 1461 |
+
break
|
| 1462 |
+
candidates = [EXPERIMENTS[i]]
|
| 1463 |
+
elif args.mode == "llm":
|
| 1464 |
+
print(f"\n[asking LLM for next experiment, iteration {i + 1}...]")
|
| 1465 |
+
try:
|
| 1466 |
+
exp = asyncio.run(
|
| 1467 |
+
_ask_llm_for_experiment(
|
| 1468 |
+
llm_backend,
|
| 1469 |
+
kb_summary=kb_summary,
|
| 1470 |
+
source=best_source,
|
| 1471 |
+
metrics=last_metrics,
|
| 1472 |
+
history=history,
|
| 1473 |
+
)
|
| 1474 |
+
)
|
| 1475 |
+
except Exception as exc:
|
| 1476 |
+
print(f" LLM call failed: {type(exc).__name__}: {exc}")
|
| 1477 |
+
exp = None
|
| 1478 |
+
if exp is None:
|
| 1479 |
+
print("LLM produced no experiment — stopping.")
|
| 1480 |
+
break
|
| 1481 |
+
candidates = [exp]
|
| 1482 |
+
else: # llm-explore
|
| 1483 |
+
K = args.candidates_per_iteration
|
| 1484 |
+
print(f"\n[asking LLM for {K} candidates, iteration {i + 1}...]")
|
| 1485 |
+
try:
|
| 1486 |
+
candidates = asyncio.run(
|
| 1487 |
+
_ask_llm_for_experiments(
|
| 1488 |
+
llm_backend,
|
| 1489 |
+
kb_summary=kb_summary,
|
| 1490 |
+
source=best_source,
|
| 1491 |
+
metrics=last_metrics,
|
| 1492 |
+
history=history,
|
| 1493 |
+
num_candidates=K,
|
| 1494 |
+
)
|
| 1495 |
+
)
|
| 1496 |
+
except Exception as exc:
|
| 1497 |
+
print(f" LLM call failed: {type(exc).__name__}: {exc}")
|
| 1498 |
+
candidates = []
|
| 1499 |
+
if not candidates:
|
| 1500 |
+
print("LLM produced no candidates — stopping.")
|
| 1501 |
+
break
|
| 1502 |
+
print(f" LLM proposed {len(candidates)} candidate(s): "
|
| 1503 |
+
+ ", ".join(c.name for c in candidates))
|
| 1504 |
+
|
| 1505 |
+
print()
|
| 1506 |
+
print("=" * 60)
|
| 1507 |
+
n_label = f" ({len(candidates)} candidates)" if len(candidates) > 1 else ""
|
| 1508 |
+
print(f"Iteration {i + 1}{n_label}")
|
| 1509 |
+
print("=" * 60)
|
| 1510 |
+
_emit({
|
| 1511 |
+
"type": "iter_start",
|
| 1512 |
+
"iteration": i + 1,
|
| 1513 |
+
"candidates": [
|
| 1514 |
+
{
|
| 1515 |
+
"name": c.name,
|
| 1516 |
+
"rationale": c.rationale,
|
| 1517 |
+
"substitutions": c.substitutions,
|
| 1518 |
+
"env_vars": c.env_vars,
|
| 1519 |
+
}
|
| 1520 |
+
for c in candidates
|
| 1521 |
+
],
|
| 1522 |
+
})
|
| 1523 |
+
|
| 1524 |
+
# ---- Evaluate each candidate against the CURRENT best ----
|
| 1525 |
+
# Crucial for llm-explore: every candidate is benchmarked against
|
| 1526 |
+
# the same best_source / best_env baseline, so the comparison is
|
| 1527 |
+
# apples-to-apples. State updates only happen after the iteration's
|
| 1528 |
+
# winner is chosen.
|
| 1529 |
+
eval_results: list[dict] = [] # candidates that produced metrics
|
| 1530 |
+
seen_this_iter: set[tuple] = set() # within-batch dedup
|
| 1531 |
+
crashed_this_iter = False
|
| 1532 |
+
max_crashes_hit = False
|
| 1533 |
+
|
| 1534 |
+
for j, exp in enumerate(candidates):
|
| 1535 |
+
cand_label = f" Candidate {j + 1}/{len(candidates)}" if len(candidates) > 1 else " Candidate"
|
| 1536 |
+
print(f"\n{cand_label}: {exp.name}")
|
| 1537 |
+
print(f" description: {exp.description}")
|
| 1538 |
+
print(f" rationale: {exp.rationale}")
|
| 1539 |
+
|
| 1540 |
+
# Helper to emit a per-candidate event with the consistent shape
|
| 1541 |
+
# the UI expects. Called at every terminus below.
|
| 1542 |
+
def _cand_event(outcome: str, metrics: dict | None = None,
|
| 1543 |
+
delta_vs_best: float | None = None,
|
| 1544 |
+
reason: str = "") -> None:
|
| 1545 |
+
_emit({
|
| 1546 |
+
"type": "candidate",
|
| 1547 |
+
"iteration": i + 1,
|
| 1548 |
+
"candidate_index": j + 1,
|
| 1549 |
+
"n_candidates": len(candidates),
|
| 1550 |
+
"name": exp.name,
|
| 1551 |
+
"rationale": exp.rationale,
|
| 1552 |
+
"substitutions": exp.substitutions,
|
| 1553 |
+
"env_vars": exp.env_vars,
|
| 1554 |
+
"outcome": outcome,
|
| 1555 |
+
"metrics": metrics,
|
| 1556 |
+
"delta_vs_best": delta_vs_best,
|
| 1557 |
+
"reason": reason,
|
| 1558 |
+
})
|
| 1559 |
+
|
| 1560 |
+
# Dedup: against prior iterations' history
|
| 1561 |
+
dup = _is_duplicate_of_history(exp, history)
|
| 1562 |
+
if dup is not None:
|
| 1563 |
+
print(f" SKIPPED — already tried as '{dup.get('name', '?')}' "
|
| 1564 |
+
f"(outcome '{dup.get('outcome', '?')}')")
|
| 1565 |
+
history.append({
|
| 1566 |
+
"name": exp.name, "outcome": "skipped",
|
| 1567 |
+
"delta_pct": None,
|
| 1568 |
+
"substitutions": exp.substitutions, "env_vars": exp.env_vars,
|
| 1569 |
+
})
|
| 1570 |
+
_cand_event("skipped", reason=f"duplicate of '{dup.get('name', '?')}'")
|
| 1571 |
+
continue
|
| 1572 |
+
|
| 1573 |
+
# Dedup: within the current batch (llm-explore can collide)
|
| 1574 |
+
fp = _experiment_fingerprint(exp)
|
| 1575 |
+
if fp in seen_this_iter:
|
| 1576 |
+
print(" SKIPPED — duplicate of an earlier candidate in this iteration")
|
| 1577 |
+
history.append({
|
| 1578 |
+
"name": exp.name, "outcome": "skipped",
|
| 1579 |
+
"delta_pct": None,
|
| 1580 |
+
"substitutions": exp.substitutions, "env_vars": exp.env_vars,
|
| 1581 |
+
})
|
| 1582 |
+
_cand_event("skipped", reason="duplicate of an earlier candidate this iteration")
|
| 1583 |
+
continue
|
| 1584 |
+
seen_this_iter.add(fp)
|
| 1585 |
+
|
| 1586 |
+
# Apply substitutions
|
| 1587 |
+
if exp.substitutions:
|
| 1588 |
+
try:
|
| 1589 |
+
candidate_source = apply_substitutions(best_source, exp.substitutions)
|
| 1590 |
+
except re.error as exc:
|
| 1591 |
+
print(f" SKIPPED — invalid regex from LLM: {exc}")
|
| 1592 |
+
rejected.append((exp.name, f"bad regex: {exc}"))
|
| 1593 |
+
history.append({
|
| 1594 |
+
"name": exp.name, "outcome": "rejected",
|
| 1595 |
+
"delta_pct": None,
|
| 1596 |
+
"substitutions": exp.substitutions, "env_vars": exp.env_vars,
|
| 1597 |
+
})
|
| 1598 |
+
_cand_event("rejected", reason=f"bad regex: {exc}")
|
| 1599 |
+
continue
|
| 1600 |
+
if candidate_source is None:
|
| 1601 |
+
print(" SKIPPED — substitution patterns didn't match")
|
| 1602 |
+
rejected.append((exp.name, "patterns didn't match"))
|
| 1603 |
+
history.append({
|
| 1604 |
+
"name": exp.name, "outcome": "skipped",
|
| 1605 |
+
"delta_pct": None,
|
| 1606 |
+
"substitutions": exp.substitutions, "env_vars": exp.env_vars,
|
| 1607 |
+
})
|
| 1608 |
+
_cand_event("skipped", reason="substitution patterns didn't match")
|
| 1609 |
+
continue
|
| 1610 |
+
else:
|
| 1611 |
+
candidate_source = best_source
|
| 1612 |
+
|
| 1613 |
+
file_counter += 1
|
| 1614 |
+
safe_name = re.sub(r"[^A-Za-z0-9_]+", "_", exp.name)[:40] or "exp"
|
| 1615 |
+
candidate_path = workspace / f"{file_counter:03d}_iter{i + 1:02d}_{safe_name}.py"
|
| 1616 |
+
candidate_path.write_text(candidate_source)
|
| 1617 |
+
|
| 1618 |
+
candidate_env = {**best_env, **exp.env_vars}
|
| 1619 |
+
if exp.env_vars:
|
| 1620 |
+
print(f" env vars: {exp.env_vars}")
|
| 1621 |
+
|
| 1622 |
+
m = benchmark(candidate_path, args.steps, candidate_env)
|
| 1623 |
+
if m is None:
|
| 1624 |
+
rejected.append((exp.name, "benchmark crashed"))
|
| 1625 |
+
history.append({
|
| 1626 |
+
"name": exp.name, "outcome": "crashed",
|
| 1627 |
+
"delta_pct": None,
|
| 1628 |
+
"substitutions": exp.substitutions, "env_vars": exp.env_vars,
|
| 1629 |
+
})
|
| 1630 |
+
total_crashes += 1
|
| 1631 |
+
crashed_this_iter = True
|
| 1632 |
+
print(
|
| 1633 |
+
f" CRASHED — counted toward max-crashes "
|
| 1634 |
+
f"({total_crashes}/{args.max_crashes})"
|
| 1635 |
+
)
|
| 1636 |
+
_cand_event("crashed", reason="benchmark subprocess failed")
|
| 1637 |
+
if total_crashes >= args.max_crashes:
|
| 1638 |
+
max_crashes_hit = True
|
| 1639 |
+
break
|
| 1640 |
+
continue
|
| 1641 |
+
|
| 1642 |
+
tps = m["tokens_per_sec"]
|
| 1643 |
+
delta_vs_best = _delta_pct(tps, best_tps)
|
| 1644 |
+
print(f" tokens/sec: {tps:.1f} (Δ {delta_vs_best:+.2f}% vs current best)")
|
| 1645 |
+
print(f" mfu_pct: {m.get('mfu_pct', 0.0):.2f}")
|
| 1646 |
+
print(f" hbm_peak_gb: {m['hbm_peak_gb']:.2f}")
|
| 1647 |
+
print(f" gpu_util_pct:{m['gpu_util_pct']:.1f}")
|
| 1648 |
+
_print_waste(m, prefix=" waste: ")
|
| 1649 |
+
|
| 1650 |
+
# Emit "evaluated" — outcome (accepted/rejected) is decided
|
| 1651 |
+
# later when the iteration's winner is picked across all
|
| 1652 |
+
# candidates. For UI display purposes the per-candidate metrics
|
| 1653 |
+
# are already useful.
|
| 1654 |
+
_cand_event("evaluated", metrics=m, delta_vs_best=delta_vs_best)
|
| 1655 |
+
|
| 1656 |
+
eval_results.append({
|
| 1657 |
+
"exp": exp,
|
| 1658 |
+
"candidate_source": candidate_source,
|
| 1659 |
+
"candidate_env": candidate_env,
|
| 1660 |
+
"metrics": m,
|
| 1661 |
+
"delta_vs_best": delta_vs_best,
|
| 1662 |
+
})
|
| 1663 |
+
|
| 1664 |
+
if max_crashes_hit:
|
| 1665 |
+
print(
|
| 1666 |
+
f"\nReached max-crashes ({args.max_crashes}) — stopping to "
|
| 1667 |
+
"avoid burning more GPU on structurally bad changes."
|
| 1668 |
+
)
|
| 1669 |
+
break
|
| 1670 |
+
|
| 1671 |
+
# ---- Pick the iteration's winner from eval_results ----
|
| 1672 |
+
if not eval_results:
|
| 1673 |
+
# Every candidate was skipped or crashed
|
| 1674 |
+
if crashed_this_iter:
|
| 1675 |
+
print("\n All candidates crashed or were skipped this iteration.")
|
| 1676 |
+
else:
|
| 1677 |
+
print("\n All candidates were skipped this iteration.")
|
| 1678 |
+
consecutive_no_improvement += 1
|
| 1679 |
+
else:
|
| 1680 |
+
winner = max(eval_results, key=lambda r: r["metrics"]["tokens_per_sec"])
|
| 1681 |
+
winner_delta = winner["delta_vs_best"]
|
| 1682 |
+
|
| 1683 |
+
# ---- Optional merge step (llm-explore only) ----
|
| 1684 |
+
# If 2+ candidates this iteration each beat the baseline, try
|
| 1685 |
+
# combining them into one experiment and benchmark the merge.
|
| 1686 |
+
# The merge replaces `winner` only if it strictly exceeds the
|
| 1687 |
+
# individual winner's tokens/sec.
|
| 1688 |
+
if args.mode == "llm-explore":
|
| 1689 |
+
positives = [r for r in eval_results if r["delta_vs_best"] > 0]
|
| 1690 |
+
if len(positives) >= 2:
|
| 1691 |
+
merged_exp, merge_reason = _build_merged_experiment(
|
| 1692 |
+
[r["exp"] for r in positives], best_source
|
| 1693 |
+
)
|
| 1694 |
+
if merged_exp is None:
|
| 1695 |
+
print(f"\n MERGE SKIPPED — {merge_reason}")
|
| 1696 |
+
_emit({
|
| 1697 |
+
"type": "merge_attempt",
|
| 1698 |
+
"iteration": i + 1,
|
| 1699 |
+
"outcome": "skipped",
|
| 1700 |
+
"reason": merge_reason,
|
| 1701 |
+
"candidate_names": [r["exp"].name for r in positives],
|
| 1702 |
+
})
|
| 1703 |
+
else:
|
| 1704 |
+
print(
|
| 1705 |
+
f"\n Merging {len(positives)} positive candidates: "
|
| 1706 |
+
f"{merged_exp.description}"
|
| 1707 |
+
)
|
| 1708 |
+
# Apply substitutions to get the merged source
|
| 1709 |
+
merged_source = best_source
|
| 1710 |
+
for pattern, replacement in merged_exp.substitutions:
|
| 1711 |
+
merged_source = re.sub(pattern, replacement, merged_source)
|
| 1712 |
+
merged_env = {**best_env, **merged_exp.env_vars}
|
| 1713 |
+
|
| 1714 |
+
file_counter += 1
|
| 1715 |
+
merged_path = workspace / f"{file_counter:03d}_iter{i + 1:02d}_merge.py"
|
| 1716 |
+
merged_path.write_text(merged_source)
|
| 1717 |
+
if merged_exp.env_vars:
|
| 1718 |
+
print(f" env vars: {merged_exp.env_vars}")
|
| 1719 |
+
|
| 1720 |
+
m = benchmark(merged_path, args.steps, merged_env)
|
| 1721 |
+
if m is None:
|
| 1722 |
+
total_crashes += 1
|
| 1723 |
+
crashed_this_iter = True
|
| 1724 |
+
print(
|
| 1725 |
+
f" MERGE CRASHED — counted toward max-crashes "
|
| 1726 |
+
f"({total_crashes}/{args.max_crashes})"
|
| 1727 |
+
)
|
| 1728 |
+
_emit({
|
| 1729 |
+
"type": "merge_attempt",
|
| 1730 |
+
"iteration": i + 1,
|
| 1731 |
+
"outcome": "crashed",
|
| 1732 |
+
"candidate_names": [r["exp"].name for r in positives],
|
| 1733 |
+
"merged_name": merged_exp.name,
|
| 1734 |
+
})
|
| 1735 |
+
if total_crashes >= args.max_crashes:
|
| 1736 |
+
max_crashes_hit = True
|
| 1737 |
+
else:
|
| 1738 |
+
tps = m["tokens_per_sec"]
|
| 1739 |
+
delta_vs_best = _delta_pct(tps, best_tps)
|
| 1740 |
+
print(
|
| 1741 |
+
f" Merged tokens/sec: {tps:.1f} "
|
| 1742 |
+
f"(Δ {delta_vs_best:+.2f}% vs baseline)"
|
| 1743 |
+
)
|
| 1744 |
+
print(f" mfu_pct: {m.get('mfu_pct', 0.0):.2f}")
|
| 1745 |
+
print(f" hbm_peak_gb: {m['hbm_peak_gb']:.2f}")
|
| 1746 |
+
|
| 1747 |
+
individual_best_tps = winner["metrics"]["tokens_per_sec"]
|
| 1748 |
+
if tps > individual_best_tps:
|
| 1749 |
+
print(
|
| 1750 |
+
f" MERGE WINS — exceeds individual winner "
|
| 1751 |
+
f"'{winner['exp'].name}' "
|
| 1752 |
+
f"({tps:.1f} > {individual_best_tps:.1f})"
|
| 1753 |
+
)
|
| 1754 |
+
_emit({
|
| 1755 |
+
"type": "merge_attempt",
|
| 1756 |
+
"iteration": i + 1,
|
| 1757 |
+
"outcome": "wins",
|
| 1758 |
+
"candidate_names": [r["exp"].name for r in positives],
|
| 1759 |
+
"merged_name": merged_exp.name,
|
| 1760 |
+
"metrics": m,
|
| 1761 |
+
"delta_vs_best": delta_vs_best,
|
| 1762 |
+
"individual_best_name": winner["exp"].name,
|
| 1763 |
+
"individual_best_tps": individual_best_tps,
|
| 1764 |
+
})
|
| 1765 |
+
# Promote merged to be the new winner
|
| 1766 |
+
winner = {
|
| 1767 |
+
"exp": merged_exp,
|
| 1768 |
+
"candidate_source": merged_source,
|
| 1769 |
+
"candidate_env": merged_env,
|
| 1770 |
+
"metrics": m,
|
| 1771 |
+
"delta_vs_best": delta_vs_best,
|
| 1772 |
+
}
|
| 1773 |
+
winner_delta = delta_vs_best
|
| 1774 |
+
else:
|
| 1775 |
+
print(
|
| 1776 |
+
f" Merge didn't beat individual winner; "
|
| 1777 |
+
f"keeping '{winner['exp'].name}'"
|
| 1778 |
+
)
|
| 1779 |
+
_emit({
|
| 1780 |
+
"type": "merge_attempt",
|
| 1781 |
+
"iteration": i + 1,
|
| 1782 |
+
"outcome": "lost",
|
| 1783 |
+
"candidate_names": [r["exp"].name for r in positives],
|
| 1784 |
+
"merged_name": merged_exp.name,
|
| 1785 |
+
"metrics": m,
|
| 1786 |
+
"delta_vs_best": delta_vs_best,
|
| 1787 |
+
"individual_best_name": winner["exp"].name,
|
| 1788 |
+
"individual_best_tps": individual_best_tps,
|
| 1789 |
+
})
|
| 1790 |
+
|
| 1791 |
+
if winner_delta >= args.improvement_threshold:
|
| 1792 |
+
print(
|
| 1793 |
+
f"\n ACCEPTED — '{winner['exp'].name}' wins "
|
| 1794 |
+
f"(Δ {winner_delta:+.2f}% vs current best)"
|
| 1795 |
+
)
|
| 1796 |
+
best_source = winner["candidate_source"]
|
| 1797 |
+
best_tps = winner["metrics"]["tokens_per_sec"]
|
| 1798 |
+
best_env = winner["candidate_env"]
|
| 1799 |
+
last_metrics = winner["metrics"]
|
| 1800 |
+
accepted.append((winner["exp"].name, best_tps, winner_delta))
|
| 1801 |
+
history.append({
|
| 1802 |
+
"name": winner["exp"].name, "outcome": "accepted",
|
| 1803 |
+
"delta_pct": winner_delta,
|
| 1804 |
+
"substitutions": winner["exp"].substitutions,
|
| 1805 |
+
"env_vars": winner["exp"].env_vars,
|
| 1806 |
+
})
|
| 1807 |
+
# Other candidates of this iteration get marked rejected
|
| 1808 |
+
for r in eval_results:
|
| 1809 |
+
if r is winner:
|
| 1810 |
+
continue
|
| 1811 |
+
rejected.append((r["exp"].name, f"{r['delta_vs_best']:+.2f}%"))
|
| 1812 |
+
history.append({
|
| 1813 |
+
"name": r["exp"].name, "outcome": "rejected",
|
| 1814 |
+
"delta_pct": r["delta_vs_best"],
|
| 1815 |
+
"substitutions": r["exp"].substitutions,
|
| 1816 |
+
"env_vars": r["exp"].env_vars,
|
| 1817 |
+
})
|
| 1818 |
+
consecutive_no_improvement = 0
|
| 1819 |
+
_emit({
|
| 1820 |
+
"type": "iter_done",
|
| 1821 |
+
"iteration": i + 1,
|
| 1822 |
+
"outcome": "accepted",
|
| 1823 |
+
"winner_name": winner["exp"].name,
|
| 1824 |
+
"winner_delta": winner_delta,
|
| 1825 |
+
"best_tps": best_tps,
|
| 1826 |
+
"best_metrics": winner["metrics"],
|
| 1827 |
+
"best_env_vars": best_env,
|
| 1828 |
+
})
|
| 1829 |
+
else:
|
| 1830 |
+
print(
|
| 1831 |
+
f"\n ALL REJECTED — best candidate '{winner['exp'].name}' "
|
| 1832 |
+
f"only Δ {winner_delta:+.2f}% (threshold {args.improvement_threshold:.1f}%)"
|
| 1833 |
+
)
|
| 1834 |
+
for r in eval_results:
|
| 1835 |
+
rejected.append((r["exp"].name, f"{r['delta_vs_best']:+.2f}%"))
|
| 1836 |
+
history.append({
|
| 1837 |
+
"name": r["exp"].name, "outcome": "rejected",
|
| 1838 |
+
"delta_pct": r["delta_vs_best"],
|
| 1839 |
+
"substitutions": r["exp"].substitutions,
|
| 1840 |
+
"env_vars": r["exp"].env_vars,
|
| 1841 |
+
})
|
| 1842 |
+
# Update last_metrics with the winner anyway so the LLM sees
|
| 1843 |
+
# the latest waste_budget on the next turn.
|
| 1844 |
+
if args.mode in ("llm", "llm-explore"):
|
| 1845 |
+
last_metrics = winner["metrics"]
|
| 1846 |
+
consecutive_no_improvement += 1
|
| 1847 |
+
_emit({
|
| 1848 |
+
"type": "iter_done",
|
| 1849 |
+
"iteration": i + 1,
|
| 1850 |
+
"outcome": "all_rejected",
|
| 1851 |
+
"winner_name": winner["exp"].name,
|
| 1852 |
+
"winner_delta": winner_delta,
|
| 1853 |
+
"best_tps": best_tps,
|
| 1854 |
+
})
|
| 1855 |
+
|
| 1856 |
+
if consecutive_no_improvement >= args.early_stop_after:
|
| 1857 |
+
print(
|
| 1858 |
+
f"\nNo improvement for {args.early_stop_after} consecutive iterations — early stopping."
|
| 1859 |
+
)
|
| 1860 |
+
break
|
| 1861 |
+
|
| 1862 |
+
# Save best
|
| 1863 |
+
best_path = workspace / "best.py"
|
| 1864 |
+
best_path.write_text(best_source)
|
| 1865 |
+
|
| 1866 |
+
# Summary
|
| 1867 |
+
print()
|
| 1868 |
+
print("=" * 60)
|
| 1869 |
+
print("AUTO-TUNE SUMMARY")
|
| 1870 |
+
print("=" * 60)
|
| 1871 |
+
print(f"Baseline tokens/sec: {baseline_tps:.1f}")
|
| 1872 |
+
print(
|
| 1873 |
+
f"Best tokens/sec: {best_tps:.1f} "
|
| 1874 |
+
f"({_delta_pct(best_tps, baseline_tps):+.2f}% vs baseline)"
|
| 1875 |
+
)
|
| 1876 |
+
print()
|
| 1877 |
+
print(f"Accepted ({len(accepted)}):")
|
| 1878 |
+
for name, tps, delta in accepted:
|
| 1879 |
+
print(f" + {name:25s} {tps:8.1f} tok/s (Δ {delta:+.2f}%)")
|
| 1880 |
+
print()
|
| 1881 |
+
print(f"Rejected ({len(rejected)}):")
|
| 1882 |
+
for name, reason in rejected:
|
| 1883 |
+
print(f" - {name:25s} {reason}")
|
| 1884 |
+
print()
|
| 1885 |
+
|
| 1886 |
+
if best_env:
|
| 1887 |
+
print("Required env vars for best config:")
|
| 1888 |
+
for k, v in best_env.items():
|
| 1889 |
+
print(f" export {k}={v}")
|
| 1890 |
+
print()
|
| 1891 |
+
|
| 1892 |
+
print(f"Best workload script: {best_path}")
|
| 1893 |
+
print(f"Diff vs baseline: diff {workload} {best_path}")
|
| 1894 |
+
|
| 1895 |
+
_emit({
|
| 1896 |
+
"type": "summary",
|
| 1897 |
+
"baseline_metrics": baseline,
|
| 1898 |
+
"best_metrics": last_metrics,
|
| 1899 |
+
"baseline_tps": baseline_tps,
|
| 1900 |
+
"best_tps": best_tps,
|
| 1901 |
+
"improvement_pct": _delta_pct(best_tps, baseline_tps),
|
| 1902 |
+
"accepted": [
|
| 1903 |
+
{"name": name, "tps": tps, "delta_pct": delta}
|
| 1904 |
+
for name, tps, delta in accepted
|
| 1905 |
+
],
|
| 1906 |
+
"rejected": [
|
| 1907 |
+
{"name": name, "reason": reason}
|
| 1908 |
+
for name, reason in rejected
|
| 1909 |
+
],
|
| 1910 |
+
"best_env_vars": best_env,
|
| 1911 |
+
"best_workload_path": str(best_path),
|
| 1912 |
+
"baseline_workload_path": str(workload),
|
| 1913 |
+
})
|
| 1914 |
+
return 0
|
| 1915 |
+
|
| 1916 |
+
|
| 1917 |
+
if __name__ == "__main__":
|
| 1918 |
+
raise SystemExit(main())
|
tests/fixtures/sample_train.json
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name_or_path": "Qwen/Qwen2.5-7B-Instruct",
|
| 3 |
+
"per_device_train_batch_size": 8,
|
| 4 |
+
"gradient_accumulation_steps": 4,
|
| 5 |
+
"max_seq_length": 4096,
|
| 6 |
+
"learning_rate": 0.0003,
|
| 7 |
+
"warmup_steps": 200,
|
| 8 |
+
"bf16": true,
|
| 9 |
+
"optim": "adamw_torch_fused",
|
| 10 |
+
"gradient_checkpointing": true,
|
| 11 |
+
"torch_compile": true,
|
| 12 |
+
"dataloader_num_workers": 4,
|
| 13 |
+
"dataloader_pin_memory": true,
|
| 14 |
+
"dataloader_prefetch_factor": 4,
|
| 15 |
+
"dataloader_persistent_workers": true,
|
| 16 |
+
"attn_implementation": "flash",
|
| 17 |
+
"num_train_epochs": 3,
|
| 18 |
+
"save_steps": 500,
|
| 19 |
+
"logging_steps": 25,
|
| 20 |
+
"output_dir": "./out",
|
| 21 |
+
"hub_token": "hf_jsonsamplehfabcdefghijklmnopqrs",
|
| 22 |
+
"checkpoint_uri": "s3://team-bucket/runs/qwen-lora-001/",
|
| 23 |
+
"env_vars": {
|
| 24 |
+
"HSA_FORCE_FINE_GRAIN_PCIE": "1",
|
| 25 |
+
"NCCL_MIN_NCHANNELS": "112"
|
| 26 |
+
}
|
| 27 |
+
}
|
tests/fixtures/sample_train.py
ADDED
|
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Realistic-looking HF Trainer fine-tuning script with secrets sprinkled in.
|
| 2 |
+
|
| 3 |
+
Used as a fixture for parse_config — exercises every code path we care about:
|
| 4 |
+
TrainingArguments kwargs, DataLoader kwargs, torch.compile, gradient
|
| 5 |
+
checkpointing, os.environ assignments, LoRA config, and from_pretrained.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import os
|
| 9 |
+
|
| 10 |
+
import torch
|
| 11 |
+
from torch.utils.data import DataLoader
|
| 12 |
+
from transformers import (
|
| 13 |
+
AutoModelForCausalLM,
|
| 14 |
+
AutoTokenizer,
|
| 15 |
+
Trainer,
|
| 16 |
+
TrainingArguments,
|
| 17 |
+
)
|
| 18 |
+
from peft import LoraConfig, get_peft_model
|
| 19 |
+
from datasets import load_dataset
|
| 20 |
+
|
| 21 |
+
# Secrets we expect parse_config to redact before storing raw_source.
|
| 22 |
+
HF_TOKEN = "hf_abcdefghijklmnopqrstuvwxyz123456"
|
| 23 |
+
OPENAI_KEY = "sk-abcdefghijklmnopqrstuvwxyz1234567890"
|
| 24 |
+
GH_TOKEN = "gho_abcdefghijklmnopqrstuvwxyz123456"
|
| 25 |
+
AUTH_HEADER = "Authorization: Bearer eyJhbGciOi.JIUzI1NiJ9.signature123"
|
| 26 |
+
DATA_ROOT = "/home/researcher/datasets/alpaca"
|
| 27 |
+
S3_BUCKET = "s3://my-team/checkpoints/qwen-lora/"
|
| 28 |
+
WS_LOG = "wss://logs.internal.example.com/stream"
|
| 29 |
+
|
| 30 |
+
# Environment variables the agent should capture into env_vars.
|
| 31 |
+
os.environ["HSA_FORCE_FINE_GRAIN_PCIE"] = "1"
|
| 32 |
+
os.environ["MIOPEN_FIND_MODE"] = "3"
|
| 33 |
+
os.environ["NCCL_MIN_NCHANNELS"] = "112"
|
| 34 |
+
|
| 35 |
+
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
|
| 36 |
+
|
| 37 |
+
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
|
| 38 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 39 |
+
MODEL_ID,
|
| 40 |
+
torch_dtype=torch.bfloat16,
|
| 41 |
+
attn_implementation="eager",
|
| 42 |
+
token=HF_TOKEN,
|
| 43 |
+
)
|
| 44 |
+
|
| 45 |
+
# LoRA — rank should land in WorkloadConfig.lora_rank.
|
| 46 |
+
lora_config = LoraConfig(
|
| 47 |
+
r=16,
|
| 48 |
+
lora_alpha=32,
|
| 49 |
+
target_modules=["q_proj", "v_proj"],
|
| 50 |
+
lora_dropout=0.05,
|
| 51 |
+
bias="none",
|
| 52 |
+
task_type="CAUSAL_LM",
|
| 53 |
+
)
|
| 54 |
+
model = get_peft_model(model, lora_config)
|
| 55 |
+
|
| 56 |
+
# Should set gradient_checkpointing=True via the explicit enable() call.
|
| 57 |
+
model.gradient_checkpointing_enable()
|
| 58 |
+
|
| 59 |
+
# Should flip torch_compile=True.
|
| 60 |
+
model = torch.compile(model, mode="reduce-overhead")
|
| 61 |
+
|
| 62 |
+
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
|
| 63 |
+
|
| 64 |
+
train_loader = DataLoader(
|
| 65 |
+
dataset,
|
| 66 |
+
batch_size=4,
|
| 67 |
+
num_workers=0,
|
| 68 |
+
pin_memory=False,
|
| 69 |
+
prefetch_factor=2,
|
| 70 |
+
persistent_workers=False,
|
| 71 |
+
)
|
| 72 |
+
|
| 73 |
+
training_args = TrainingArguments(
|
| 74 |
+
output_dir="./out",
|
| 75 |
+
per_device_train_batch_size=4,
|
| 76 |
+
gradient_accumulation_steps=8,
|
| 77 |
+
num_train_epochs=3,
|
| 78 |
+
learning_rate=2e-4,
|
| 79 |
+
warmup_steps=100,
|
| 80 |
+
fp16=True,
|
| 81 |
+
optim="adamw_torch",
|
| 82 |
+
logging_steps=10,
|
| 83 |
+
save_steps=500,
|
| 84 |
+
dataloader_num_workers=0,
|
| 85 |
+
dataloader_pin_memory=False,
|
| 86 |
+
gradient_checkpointing=True,
|
| 87 |
+
torch_compile=False,
|
| 88 |
+
report_to="none",
|
| 89 |
+
push_to_hub=False,
|
| 90 |
+
hub_token=HF_TOKEN,
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
trainer = Trainer(
|
| 94 |
+
model=model,
|
| 95 |
+
args=training_args,
|
| 96 |
+
train_dataset=dataset,
|
| 97 |
+
tokenizer=tokenizer,
|
| 98 |
+
)
|
| 99 |
+
|
| 100 |
+
if __name__ == "__main__":
|
| 101 |
+
trainer.train()
|
tests/fixtures/sample_train.yaml
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# YAML training config — same schema as JSON, exercised separately so we know
|
| 2 |
+
# yaml.safe_load + dict extraction agree with the JSON path.
|
| 3 |
+
model_name_or_path: Qwen/Qwen2.5-7B-Instruct
|
| 4 |
+
per_device_train_batch_size: 2
|
| 5 |
+
gradient_accumulation_steps: 16
|
| 6 |
+
max_seq_length: 8192
|
| 7 |
+
learning_rate: 1.0e-4
|
| 8 |
+
warmup_steps: 50
|
| 9 |
+
bf16: true
|
| 10 |
+
optim: adamw_torch
|
| 11 |
+
gradient_checkpointing: true
|
| 12 |
+
torch_compile: false
|
| 13 |
+
dataloader_num_workers: 8
|
| 14 |
+
dataloader_pin_memory: true
|
| 15 |
+
dataloader_prefetch_factor: 2
|
| 16 |
+
dataloader_persistent_workers: true
|
| 17 |
+
attn_implementation: sdpa
|
| 18 |
+
num_train_epochs: 1
|
| 19 |
+
output_dir: ./out
|
| 20 |
+
# Secrets that should be scrubbed:
|
| 21 |
+
hub_token: "hf_yamlsamplehfabcdefghijklmnopqrs1"
|
| 22 |
+
auth_header: "Bearer eyJ.payload.signaturetoken"
|
| 23 |
+
data_path: "/home/teamuser/datasets/alpaca-cleaned"
|
| 24 |
+
env_vars:
|
| 25 |
+
HSA_FORCE_FINE_GRAIN_PCIE: "1"
|
| 26 |
+
MIOPEN_FIND_MODE: "3"
|
tests/test_compare_runs_normalize.py
ADDED
|
@@ -0,0 +1,155 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for compare_runs's flat-Patch recovery path.
|
| 2 |
+
|
| 3 |
+
Live AMD-GPU lesson: Qwen models routinely forward a flat WorkloadConfig (or
|
| 4 |
+
even just the changed-fields subset) as the ``patch=`` argument to
|
| 5 |
+
compare_runs, instead of the full Patch envelope. ``_normalize_patch`` is
|
| 6 |
+
the safety net — it must:
|
| 7 |
+
|
| 8 |
+
1. Pass real Patch dicts through unchanged.
|
| 9 |
+
2. Detect any flat-config shape (full WorkloadConfig, just dataloader
|
| 10 |
+
fields, just env_vars, etc.) — NOT just dicts with model_name.
|
| 11 |
+
3. Recover by substituting the cached propose_patch result when one
|
| 12 |
+
exists, or wrapping the flat config minimally as a last resort.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import pytest
|
| 18 |
+
|
| 19 |
+
from agent.tools import compare_runs as cr_mod
|
| 20 |
+
from agent.tools import propose_patch as pp_mod
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
# ---------------------------------------------------------------------------
|
| 24 |
+
# _looks_like_flat_config detection
|
| 25 |
+
# ---------------------------------------------------------------------------
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
class TestFlatConfigDetection:
|
| 29 |
+
def test_real_patch_is_not_flat(self) -> None:
|
| 30 |
+
real = {
|
| 31 |
+
"new_config": {"model_name": "x"},
|
| 32 |
+
"diff": "(no changes)",
|
| 33 |
+
"rationale": [],
|
| 34 |
+
"expected_speedup_low": 1.0,
|
| 35 |
+
"expected_speedup_high": 1.0,
|
| 36 |
+
"confidence": 0.0,
|
| 37 |
+
}
|
| 38 |
+
assert cr_mod._looks_like_flat_config(real) is False
|
| 39 |
+
|
| 40 |
+
def test_full_workload_config_is_flat(self) -> None:
|
| 41 |
+
flat = {
|
| 42 |
+
"model_name": "Qwen/Qwen2.5-7B-Instruct",
|
| 43 |
+
"precision": "bf16",
|
| 44 |
+
"attention_impl": "flash_rocm",
|
| 45 |
+
"batch_size": 12,
|
| 46 |
+
}
|
| 47 |
+
assert cr_mod._looks_like_flat_config(flat) is True
|
| 48 |
+
|
| 49 |
+
def test_dataloader_only_diff_is_flat(self) -> None:
|
| 50 |
+
# The exact failure mode from the live MI300X audit: model only
|
| 51 |
+
# passed the *changed* dataloader fields, no model_name in sight.
|
| 52 |
+
flat = {
|
| 53 |
+
"dataloader_persistent_workers": True,
|
| 54 |
+
"dataloader_pin_memory": True,
|
| 55 |
+
"dataloader_workers": 8,
|
| 56 |
+
}
|
| 57 |
+
assert cr_mod._looks_like_flat_config(flat) is True
|
| 58 |
+
|
| 59 |
+
def test_env_vars_only_diff_is_flat(self) -> None:
|
| 60 |
+
flat = {"env_vars": {"NCCL_MIN_NCHANNELS": "112"}}
|
| 61 |
+
assert cr_mod._looks_like_flat_config(flat) is True
|
| 62 |
+
|
| 63 |
+
def test_precision_only_diff_is_flat(self) -> None:
|
| 64 |
+
flat = {"precision": "bf16"}
|
| 65 |
+
assert cr_mod._looks_like_flat_config(flat) is True
|
| 66 |
+
|
| 67 |
+
def test_unrelated_dict_not_flat(self) -> None:
|
| 68 |
+
# Garbage dict with no WorkloadConfig fields → don't claim it's flat.
|
| 69 |
+
assert cr_mod._looks_like_flat_config({"foo": 1, "bar": 2}) is False
|
| 70 |
+
|
| 71 |
+
def test_non_dict_not_flat(self) -> None:
|
| 72 |
+
assert cr_mod._looks_like_flat_config(None) is False # type: ignore[arg-type]
|
| 73 |
+
assert cr_mod._looks_like_flat_config("a string") is False # type: ignore[arg-type]
|
| 74 |
+
assert cr_mod._looks_like_flat_config([1, 2, 3]) is False # type: ignore[arg-type]
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
# ---------------------------------------------------------------------------
|
| 78 |
+
# _normalize_patch + cached-patch recovery
|
| 79 |
+
# ---------------------------------------------------------------------------
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
@pytest.fixture
|
| 83 |
+
def cached_patch(monkeypatch):
|
| 84 |
+
"""Plant a fake `latest_patch()` so the recovery path picks it up."""
|
| 85 |
+
fake = {
|
| 86 |
+
"new_config": {"model_name": "Qwen/Qwen2.5-7B-Instruct", "precision": "bf16"},
|
| 87 |
+
"diff": "- precision: fp16\n+ precision: bf16",
|
| 88 |
+
"rationale": [
|
| 89 |
+
{
|
| 90 |
+
"rule_id": "precision.bf16_over_fp16_on_mi300x",
|
| 91 |
+
"rationale": "r",
|
| 92 |
+
"citation": "c",
|
| 93 |
+
"targets_bucket": "precision_path",
|
| 94 |
+
"estimated_recovery_seconds": 0.09,
|
| 95 |
+
}
|
| 96 |
+
],
|
| 97 |
+
"expected_speedup_low": 1.05,
|
| 98 |
+
"expected_speedup_high": 1.30,
|
| 99 |
+
"confidence": 0.85,
|
| 100 |
+
}
|
| 101 |
+
monkeypatch.setattr(pp_mod, "_LAST_PATCH", fake)
|
| 102 |
+
yield fake
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
class TestNormalizePatch:
|
| 106 |
+
def test_real_patch_passes_through(self) -> None:
|
| 107 |
+
real = {
|
| 108 |
+
"new_config": {"model_name": "x"},
|
| 109 |
+
"diff": "...",
|
| 110 |
+
"rationale": [],
|
| 111 |
+
"expected_speedup_low": 1.0,
|
| 112 |
+
"expected_speedup_high": 1.0,
|
| 113 |
+
"confidence": 0.0,
|
| 114 |
+
}
|
| 115 |
+
out, notes = cr_mod._normalize_patch(real)
|
| 116 |
+
assert out is real
|
| 117 |
+
assert notes == []
|
| 118 |
+
|
| 119 |
+
def test_dataloader_only_diff_recovers_via_cached(self, cached_patch) -> None:
|
| 120 |
+
# The exact live-AMD-GPU failure: model forwarded only the changed
|
| 121 |
+
# dataloader fields. Old code's narrow sentinel set (model_name etc.)
|
| 122 |
+
# would miss this. New behavior: detected, cached patch substituted.
|
| 123 |
+
flat = {
|
| 124 |
+
"dataloader_persistent_workers": True,
|
| 125 |
+
"dataloader_pin_memory": True,
|
| 126 |
+
"dataloader_workers": 8,
|
| 127 |
+
}
|
| 128 |
+
out, notes = cr_mod._normalize_patch(flat)
|
| 129 |
+
assert out is cached_patch # full fidelity restored
|
| 130 |
+
assert any("substituted the cached" in n for n in notes)
|
| 131 |
+
|
| 132 |
+
def test_flat_config_falls_back_to_minimal_wrap_when_no_cache(
|
| 133 |
+
self, monkeypatch
|
| 134 |
+
) -> None:
|
| 135 |
+
# No cached patch — must still produce a Patch-shape dict so
|
| 136 |
+
# compare_runs doesn't crash on Pydantic validation.
|
| 137 |
+
monkeypatch.setattr(pp_mod, "_LAST_PATCH", None)
|
| 138 |
+
flat = {"precision": "bf16"}
|
| 139 |
+
out, notes = cr_mod._normalize_patch(flat)
|
| 140 |
+
assert "new_config" in out
|
| 141 |
+
assert "diff" in out
|
| 142 |
+
assert out["expected_speedup_low"] == 1.0
|
| 143 |
+
assert out["confidence"] == 0.0
|
| 144 |
+
assert any("synthesized a minimal Patch" in n for n in notes)
|
| 145 |
+
|
| 146 |
+
def test_non_flat_garbage_passes_through_for_pydantic_to_reject(
|
| 147 |
+
self, monkeypatch
|
| 148 |
+
) -> None:
|
| 149 |
+
# If it's neither a real Patch nor a recognizable flat config, let
|
| 150 |
+
# pydantic produce the clear ValidationError — don't silently mangle.
|
| 151 |
+
monkeypatch.setattr(pp_mod, "_LAST_PATCH", None)
|
| 152 |
+
garbage = {"foo": 1, "bar": [2]}
|
| 153 |
+
out, notes = cr_mod._normalize_patch(garbage)
|
| 154 |
+
assert out is garbage
|
| 155 |
+
assert notes == []
|
tests/test_loop.py
ADDED
|
@@ -0,0 +1,322 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for the agent loop driver.
|
| 2 |
+
|
| 3 |
+
We never hit a real LLM. The loop talks to a `Backend` (see
|
| 4 |
+
`agent/backends/`); each test injects a `FakeBackend` whose `next_turn`
|
| 5 |
+
returns a queued sequence of scripted `AgentTurn` objects. Tools are
|
| 6 |
+
stubbed so we can drive specific control-flow paths.
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
from __future__ import annotations
|
| 10 |
+
|
| 11 |
+
from typing import Any
|
| 12 |
+
|
| 13 |
+
import pytest
|
| 14 |
+
|
| 15 |
+
from agent import loop as loop_module
|
| 16 |
+
from agent.backends.base import AgentTurn, Backend, ToolCall
|
| 17 |
+
from agent.schemas import SSEEvent, ToolResult
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
# ---------------------------------------------------------------------------
|
| 21 |
+
# Fake backend
|
| 22 |
+
# ---------------------------------------------------------------------------
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
class FakeBackend(Backend):
|
| 26 |
+
"""A scripted Backend for testing the loop in isolation.
|
| 27 |
+
|
| 28 |
+
Each test queues a list of `AgentTurn`s; calling `next_turn` pops the
|
| 29 |
+
next one. We also record every tool result the loop hands back so tests
|
| 30 |
+
can assert that error / id / content were threaded through correctly.
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
name = "fake"
|
| 34 |
+
|
| 35 |
+
def __init__(
|
| 36 |
+
self,
|
| 37 |
+
scripted_turns: list[AgentTurn] | None = None,
|
| 38 |
+
next_turn_raises: BaseException | None = None,
|
| 39 |
+
) -> None:
|
| 40 |
+
self._scripted = list(scripted_turns or [])
|
| 41 |
+
self._raise_on_next = next_turn_raises
|
| 42 |
+
self.user_messages: list[str] = []
|
| 43 |
+
self.tool_results: list[dict[str, Any]] = []
|
| 44 |
+
self.turn_count = 0
|
| 45 |
+
|
| 46 |
+
def add_user_message(self, content: str) -> None:
|
| 47 |
+
self.user_messages.append(content)
|
| 48 |
+
|
| 49 |
+
def add_tool_result(
|
| 50 |
+
self,
|
| 51 |
+
tool_call_id: str,
|
| 52 |
+
name: str,
|
| 53 |
+
content: str,
|
| 54 |
+
is_error: bool,
|
| 55 |
+
) -> None:
|
| 56 |
+
self.tool_results.append(
|
| 57 |
+
{
|
| 58 |
+
"id": tool_call_id,
|
| 59 |
+
"name": name,
|
| 60 |
+
"content": content,
|
| 61 |
+
"is_error": is_error,
|
| 62 |
+
}
|
| 63 |
+
)
|
| 64 |
+
|
| 65 |
+
async def next_turn(self, tool_schemas: list[dict[str, Any]]) -> AgentTurn:
|
| 66 |
+
self.turn_count += 1
|
| 67 |
+
if self._raise_on_next is not None:
|
| 68 |
+
exc = self._raise_on_next
|
| 69 |
+
self._raise_on_next = None
|
| 70 |
+
raise exc
|
| 71 |
+
if not self._scripted:
|
| 72 |
+
raise AssertionError(
|
| 73 |
+
"FakeBackend exhausted — loop made more turns than expected"
|
| 74 |
+
)
|
| 75 |
+
return self._scripted.pop(0)
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
# ---------------------------------------------------------------------------
|
| 79 |
+
# Helpers
|
| 80 |
+
# ---------------------------------------------------------------------------
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
def _install_backend(monkeypatch, backend: Backend) -> Backend:
|
| 84 |
+
"""Replace `make_backend` so the loop sees our fake."""
|
| 85 |
+
monkeypatch.setattr(loop_module, "make_backend", lambda **_kwargs: backend)
|
| 86 |
+
return backend
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def _install_make_backend_raises(monkeypatch, exc: BaseException) -> None:
|
| 90 |
+
def boom(**_kwargs: Any) -> Backend:
|
| 91 |
+
raise exc
|
| 92 |
+
|
| 93 |
+
monkeypatch.setattr(loop_module, "make_backend", boom)
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
def _install_fake_tools(
|
| 97 |
+
monkeypatch, tool_responses: dict[str, ToolResult]
|
| 98 |
+
) -> list[str]:
|
| 99 |
+
"""Replace `tools_module.call` and `tool_schemas`. Returns a list
|
| 100 |
+
that records the order tools were invoked.
|
| 101 |
+
"""
|
| 102 |
+
invoked: list[str] = []
|
| 103 |
+
|
| 104 |
+
def fake_call(name: str, **_kwargs: Any) -> ToolResult:
|
| 105 |
+
invoked.append(name)
|
| 106 |
+
return tool_responses.get(
|
| 107 |
+
name, ToolResult(ok=False, error=f"no fake registered for {name}")
|
| 108 |
+
)
|
| 109 |
+
|
| 110 |
+
monkeypatch.setattr(loop_module.tools_module, "call", fake_call)
|
| 111 |
+
monkeypatch.setattr(loop_module.tools_module, "tool_schemas", lambda: [])
|
| 112 |
+
return invoked
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
async def _collect(stream) -> list[SSEEvent]:
|
| 116 |
+
out: list[SSEEvent] = []
|
| 117 |
+
async for event in stream:
|
| 118 |
+
out.append(event)
|
| 119 |
+
return out
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
# ---------------------------------------------------------------------------
|
| 123 |
+
# Tests
|
| 124 |
+
# ---------------------------------------------------------------------------
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
@pytest.mark.asyncio
|
| 128 |
+
async def test_emits_thought_then_tool_call_then_tool_result(monkeypatch) -> None:
|
| 129 |
+
backend = FakeBackend(
|
| 130 |
+
scripted_turns=[
|
| 131 |
+
AgentTurn(
|
| 132 |
+
text_blocks=["I'll start by parsing the config."],
|
| 133 |
+
tool_calls=[
|
| 134 |
+
ToolCall(id="tu_1", name="parse_config", input={"file_path": "/x.py"})
|
| 135 |
+
],
|
| 136 |
+
stop_reason="tool_use",
|
| 137 |
+
),
|
| 138 |
+
AgentTurn(text_blocks=["Done."], tool_calls=[], stop_reason="end_turn"),
|
| 139 |
+
]
|
| 140 |
+
)
|
| 141 |
+
_install_backend(monkeypatch, backend)
|
| 142 |
+
invoked = _install_fake_tools(
|
| 143 |
+
monkeypatch,
|
| 144 |
+
{"parse_config": ToolResult(ok=True, result={"model_name": "x"})},
|
| 145 |
+
)
|
| 146 |
+
|
| 147 |
+
events = await _collect(loop_module.run_audit("/x.py"))
|
| 148 |
+
types = [e.type for e in events]
|
| 149 |
+
|
| 150 |
+
assert types[0] == "thought"
|
| 151 |
+
assert types[1] == "tool_call"
|
| 152 |
+
assert types[2] == "tool_result"
|
| 153 |
+
# No compare_runs ⇒ final event is the "no final report" error.
|
| 154 |
+
assert types[-1] == "error"
|
| 155 |
+
assert "without producing a final report" in events[-1].data["message"]
|
| 156 |
+
assert invoked == ["parse_config"]
|
| 157 |
+
|
| 158 |
+
# tool_call carries id/name/input; tool_result mirrors that plus ok/result/error.
|
| 159 |
+
assert events[1].data == {
|
| 160 |
+
"id": "tu_1",
|
| 161 |
+
"name": "parse_config",
|
| 162 |
+
"input": {"file_path": "/x.py"},
|
| 163 |
+
}
|
| 164 |
+
assert events[2].data["ok"] is True
|
| 165 |
+
assert events[2].data["result"] == {"model_name": "x"}
|
| 166 |
+
assert events[2].data["error"] is None
|
| 167 |
+
|
| 168 |
+
# The user message and tool result were threaded into the backend.
|
| 169 |
+
assert backend.user_messages == ["Audit this fine-tuning workload: /x.py"]
|
| 170 |
+
assert backend.tool_results == [
|
| 171 |
+
{
|
| 172 |
+
"id": "tu_1",
|
| 173 |
+
"name": "parse_config",
|
| 174 |
+
"content": '{"model_name": "x"}',
|
| 175 |
+
"is_error": False,
|
| 176 |
+
}
|
| 177 |
+
]
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
@pytest.mark.asyncio
|
| 181 |
+
async def test_final_report_extracted_from_compare_runs(monkeypatch) -> None:
|
| 182 |
+
fake_report = {"workload_name": "test", "speedup_actual": 2.0}
|
| 183 |
+
backend = FakeBackend(
|
| 184 |
+
scripted_turns=[
|
| 185 |
+
AgentTurn(
|
| 186 |
+
text_blocks=["Wrapping up."],
|
| 187 |
+
tool_calls=[
|
| 188 |
+
ToolCall(
|
| 189 |
+
id="tu_compare",
|
| 190 |
+
name="compare_runs",
|
| 191 |
+
input={
|
| 192 |
+
"workload_name": "t",
|
| 193 |
+
"before": {},
|
| 194 |
+
"after": {},
|
| 195 |
+
"patch": {},
|
| 196 |
+
},
|
| 197 |
+
)
|
| 198 |
+
],
|
| 199 |
+
stop_reason="end_turn",
|
| 200 |
+
),
|
| 201 |
+
]
|
| 202 |
+
)
|
| 203 |
+
_install_backend(monkeypatch, backend)
|
| 204 |
+
_install_fake_tools(
|
| 205 |
+
monkeypatch, {"compare_runs": ToolResult(ok=True, result=fake_report)}
|
| 206 |
+
)
|
| 207 |
+
|
| 208 |
+
events = await _collect(loop_module.run_audit("/x.py"))
|
| 209 |
+
|
| 210 |
+
assert events[-1].type == "final_report"
|
| 211 |
+
assert events[-1].data["report"] == fake_report
|
| 212 |
+
|
| 213 |
+
|
| 214 |
+
@pytest.mark.asyncio
|
| 215 |
+
async def test_tool_error_passes_through_does_not_crash(monkeypatch) -> None:
|
| 216 |
+
backend = FakeBackend(
|
| 217 |
+
scripted_turns=[
|
| 218 |
+
AgentTurn(
|
| 219 |
+
text_blocks=["Trying parse."],
|
| 220 |
+
tool_calls=[
|
| 221 |
+
ToolCall(id="tu_1", name="parse_config", input={"file_path": "/bogus"})
|
| 222 |
+
],
|
| 223 |
+
stop_reason="tool_use",
|
| 224 |
+
),
|
| 225 |
+
AgentTurn(text_blocks=["Giving up."], tool_calls=[], stop_reason="end_turn"),
|
| 226 |
+
]
|
| 227 |
+
)
|
| 228 |
+
_install_backend(monkeypatch, backend)
|
| 229 |
+
_install_fake_tools(
|
| 230 |
+
monkeypatch,
|
| 231 |
+
{"parse_config": ToolResult(ok=False, error="file not found")},
|
| 232 |
+
)
|
| 233 |
+
|
| 234 |
+
events = await _collect(loop_module.run_audit("/bogus"))
|
| 235 |
+
tool_result_events = [e for e in events if e.type == "tool_result"]
|
| 236 |
+
assert len(tool_result_events) == 1
|
| 237 |
+
assert tool_result_events[0].data["ok"] is False
|
| 238 |
+
assert tool_result_events[0].data["error"] == "file not found"
|
| 239 |
+
# The loop kept iterating rather than bailing.
|
| 240 |
+
assert events[-1].type == "error" # no compare_runs ⇒ "no final report"
|
| 241 |
+
|
| 242 |
+
# Backend received an is_error=True tool result with the error message.
|
| 243 |
+
assert backend.tool_results[-1]["is_error"] is True
|
| 244 |
+
assert backend.tool_results[-1]["content"] == "file not found"
|
| 245 |
+
|
| 246 |
+
|
| 247 |
+
@pytest.mark.asyncio
|
| 248 |
+
async def test_backend_construction_failure_yields_error_event(monkeypatch) -> None:
|
| 249 |
+
_install_make_backend_raises(
|
| 250 |
+
monkeypatch, RuntimeError("HF_TOKEN is not set; Qwen backend cannot run.")
|
| 251 |
+
)
|
| 252 |
+
events = await _collect(loop_module.run_audit("/x.py"))
|
| 253 |
+
assert len(events) == 1
|
| 254 |
+
assert events[0].type == "error"
|
| 255 |
+
assert "HF_TOKEN" in events[0].data["message"]
|
| 256 |
+
|
| 257 |
+
|
| 258 |
+
@pytest.mark.asyncio
|
| 259 |
+
async def test_mid_loop_exception_is_caught(monkeypatch) -> None:
|
| 260 |
+
backend = FakeBackend(next_turn_raises=RuntimeError("boom"))
|
| 261 |
+
_install_backend(monkeypatch, backend)
|
| 262 |
+
monkeypatch.setattr(loop_module.tools_module, "tool_schemas", lambda: [])
|
| 263 |
+
|
| 264 |
+
events = await _collect(loop_module.run_audit("/x.py"))
|
| 265 |
+
assert events[-1].type == "error"
|
| 266 |
+
assert "boom" in events[-1].data["message"]
|
| 267 |
+
|
| 268 |
+
|
| 269 |
+
@pytest.mark.asyncio
|
| 270 |
+
async def test_loop_caps_at_max_steps(monkeypatch) -> None:
|
| 271 |
+
"""Even if the model never says end_turn, we bail after MAX_STEPS."""
|
| 272 |
+
backend = FakeBackend(
|
| 273 |
+
scripted_turns=[
|
| 274 |
+
AgentTurn(
|
| 275 |
+
text_blocks=[f"step {i}"],
|
| 276 |
+
tool_calls=[
|
| 277 |
+
ToolCall(id=f"tu_{i}", name="parse_config", input={"file_path": "/x.py"})
|
| 278 |
+
],
|
| 279 |
+
stop_reason="tool_use",
|
| 280 |
+
)
|
| 281 |
+
for i in range(loop_module.MAX_STEPS + 2) # extra so we'd overrun
|
| 282 |
+
]
|
| 283 |
+
)
|
| 284 |
+
_install_backend(monkeypatch, backend)
|
| 285 |
+
_install_fake_tools(monkeypatch, {"parse_config": ToolResult(ok=True, result={})})
|
| 286 |
+
|
| 287 |
+
events = await _collect(loop_module.run_audit("/x.py"))
|
| 288 |
+
# Backend's next_turn was called exactly MAX_STEPS times.
|
| 289 |
+
assert backend.turn_count == loop_module.MAX_STEPS
|
| 290 |
+
# Last event is the "no final report" error (not a crash).
|
| 291 |
+
assert events[-1].type == "error"
|
| 292 |
+
|
| 293 |
+
|
| 294 |
+
@pytest.mark.asyncio
|
| 295 |
+
async def test_tool_call_id_is_threaded_to_backend(monkeypatch) -> None:
|
| 296 |
+
"""The loop must hand the tool_call id back to the backend so the next
|
| 297 |
+
turn's request can correlate the tool_result with the originating call.
|
| 298 |
+
"""
|
| 299 |
+
backend = FakeBackend(
|
| 300 |
+
scripted_turns=[
|
| 301 |
+
AgentTurn(
|
| 302 |
+
text_blocks=["parse"],
|
| 303 |
+
tool_calls=[
|
| 304 |
+
ToolCall(id="tu_abc", name="parse_config", input={"file_path": "/x"})
|
| 305 |
+
],
|
| 306 |
+
stop_reason="tool_use",
|
| 307 |
+
),
|
| 308 |
+
AgentTurn(text_blocks=["done"], tool_calls=[], stop_reason="end_turn"),
|
| 309 |
+
]
|
| 310 |
+
)
|
| 311 |
+
_install_backend(monkeypatch, backend)
|
| 312 |
+
_install_fake_tools(
|
| 313 |
+
monkeypatch, {"parse_config": ToolResult(ok=True, result={"a": 1})}
|
| 314 |
+
)
|
| 315 |
+
|
| 316 |
+
await _collect(loop_module.run_audit("/x"))
|
| 317 |
+
|
| 318 |
+
# Backend got exactly one tool_result with id=tu_abc.
|
| 319 |
+
assert len(backend.tool_results) == 1
|
| 320 |
+
assert backend.tool_results[0]["id"] == "tu_abc"
|
| 321 |
+
assert backend.tool_results[0]["name"] == "parse_config"
|
| 322 |
+
assert backend.tool_results[0]["is_error"] is False
|
tests/test_misnested_args.py
ADDED
|
@@ -0,0 +1,95 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for the defensive misnested-arg extraction in benchmark + profile_run.
|
| 2 |
+
|
| 3 |
+
Live AMD-GPU lesson: Qwen2.5-7B (and probably others) occasionally JSON-nests
|
| 4 |
+
``steps`` / ``cache`` *inside* the ``config`` dict instead of at the top level
|
| 5 |
+
alongside it. WorkloadConfig strict-validates extras, so without this defense
|
| 6 |
+
the call errors out and a tool slot is wasted. The well-tuned scenario run
|
| 7 |
+
on 2026-05-07 burned two of the eight available slots on this exact mistake;
|
| 8 |
+
fixing it costs nothing and saves the audit.
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import shutil
|
| 14 |
+
|
| 15 |
+
from agent.tools import call
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def _baseline_config() -> dict:
|
| 19 |
+
return {
|
| 20 |
+
"model_name": "Qwen/Qwen2.5-7B-Instruct",
|
| 21 |
+
"batch_size": 4,
|
| 22 |
+
"precision": "fp16",
|
| 23 |
+
"attention_impl": "eager",
|
| 24 |
+
"dataloader_workers": 0,
|
| 25 |
+
}
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
class TestBenchmarkMisnestedArgs:
|
| 29 |
+
def setup_method(self) -> None:
|
| 30 |
+
# Each test starts with an empty cache so cache-hit doesn't mask the
|
| 31 |
+
# behavior under test.
|
| 32 |
+
shutil.rmtree("bench_cache", ignore_errors=True)
|
| 33 |
+
|
| 34 |
+
def test_steps_nested_in_config_is_extracted(self) -> None:
|
| 35 |
+
"""Old behavior: ``WorkloadConfig`` validation explodes with
|
| 36 |
+
'Extra inputs are not permitted [steps]'. New behavior: defensive
|
| 37 |
+
extraction pulls ``steps`` back to the top-level arg, call succeeds.
|
| 38 |
+
"""
|
| 39 |
+
cfg = {**_baseline_config(), "steps": 25}
|
| 40 |
+
result = call("benchmark", config=cfg)
|
| 41 |
+
assert result.ok, result.error
|
| 42 |
+
assert result.result["steps"] == 25
|
| 43 |
+
|
| 44 |
+
def test_cache_nested_in_config_is_extracted(self) -> None:
|
| 45 |
+
cfg = {**_baseline_config(), "cache": False}
|
| 46 |
+
result = call("benchmark", config=cfg)
|
| 47 |
+
assert result.ok, result.error
|
| 48 |
+
|
| 49 |
+
def test_force_rerun_nested_in_config_is_extracted(self) -> None:
|
| 50 |
+
cfg = {**_baseline_config(), "force_rerun": True}
|
| 51 |
+
result = call("benchmark", config=cfg)
|
| 52 |
+
assert result.ok, result.error
|
| 53 |
+
|
| 54 |
+
def test_explicit_top_level_wins_over_nested(self) -> None:
|
| 55 |
+
"""If caller passes BOTH (config has steps + top-level steps), the
|
| 56 |
+
explicit non-default top-level wins. Defensive code is for the
|
| 57 |
+
accident case, not for letting nesting silently override."""
|
| 58 |
+
cfg = {**_baseline_config(), "steps": 25}
|
| 59 |
+
result = call("benchmark", config=cfg, steps=37)
|
| 60 |
+
assert result.ok, result.error
|
| 61 |
+
assert result.result["steps"] == 37
|
| 62 |
+
|
| 63 |
+
def test_all_three_nested_at_once(self) -> None:
|
| 64 |
+
"""The exact failure mode from the live run: model nested three
|
| 65 |
+
runtime args inside config. All three should get pulled out.
|
| 66 |
+
"""
|
| 67 |
+
cfg = {
|
| 68 |
+
**_baseline_config(),
|
| 69 |
+
"steps": 30,
|
| 70 |
+
"cache": False,
|
| 71 |
+
"force_rerun": True,
|
| 72 |
+
}
|
| 73 |
+
result = call("benchmark", config=cfg)
|
| 74 |
+
assert result.ok, result.error
|
| 75 |
+
assert result.result["steps"] == 30
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
class TestProfileRunMisnestedArgs:
|
| 79 |
+
def test_steps_nested_in_config_is_extracted(self) -> None:
|
| 80 |
+
cfg = {**_baseline_config(), "steps": 7}
|
| 81 |
+
result = call("profile_run", config=cfg)
|
| 82 |
+
assert result.ok, result.error
|
| 83 |
+
assert result.result["steps"] == 7
|
| 84 |
+
|
| 85 |
+
def test_explicit_top_level_wins(self) -> None:
|
| 86 |
+
cfg = {**_baseline_config(), "steps": 7}
|
| 87 |
+
result = call("profile_run", config=cfg, steps=15)
|
| 88 |
+
assert result.ok, result.error
|
| 89 |
+
assert result.result["steps"] == 15
|
| 90 |
+
|
| 91 |
+
def test_clean_config_unaffected(self) -> None:
|
| 92 |
+
"""Sanity: when nothing is misnested, behavior is unchanged."""
|
| 93 |
+
result = call("profile_run", config=_baseline_config())
|
| 94 |
+
assert result.ok, result.error
|
| 95 |
+
assert result.result["steps"] == 10 # default
|
tests/test_parse_config.py
ADDED
|
@@ -0,0 +1,377 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for ``agent.tools.parse_config._parse_config``.
|
| 2 |
+
|
| 3 |
+
Covers all three input shapes (Python AST, JSON, YAML), redaction of every
|
| 4 |
+
secret pattern we ship, error paths for missing/malformed inputs, and the
|
| 5 |
+
WorkloadConfig field mapping for HF TrainingArguments + DataLoader kwargs.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
import json
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
|
| 13 |
+
import pytest
|
| 14 |
+
|
| 15 |
+
from agent.schemas import WorkloadConfig
|
| 16 |
+
from agent.tools.parse_config import PARSE_CONFIG, _parse_config, _parse_config_full
|
| 17 |
+
|
| 18 |
+
FIXTURES = Path(__file__).parent / "fixtures"
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
# ---------------------------------------------------------------------------
|
| 22 |
+
# Python script path
|
| 23 |
+
# ---------------------------------------------------------------------------
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
class TestPythonScript:
|
| 27 |
+
def test_returns_ok(self) -> None:
|
| 28 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 29 |
+
assert result.ok, result.error
|
| 30 |
+
assert result.result is not None
|
| 31 |
+
|
| 32 |
+
def test_extracts_training_arguments_kwargs(self) -> None:
|
| 33 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 34 |
+
cfg = result.result
|
| 35 |
+
assert cfg["batch_size"] == 4
|
| 36 |
+
assert cfg["grad_accum_steps"] == 8
|
| 37 |
+
assert cfg["lr"] == pytest.approx(2e-4)
|
| 38 |
+
assert cfg["warmup_steps"] == 100
|
| 39 |
+
assert cfg["optimizer"] == "adamw_torch"
|
| 40 |
+
# fp16=True in TrainingArguments → precision should resolve to fp16.
|
| 41 |
+
assert cfg["precision"] == "fp16"
|
| 42 |
+
|
| 43 |
+
def test_dataloader_kwargs_captured(self) -> None:
|
| 44 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 45 |
+
cfg = result.result
|
| 46 |
+
assert cfg["dataloader_workers"] == 0
|
| 47 |
+
assert cfg["dataloader_pin_memory"] is False
|
| 48 |
+
assert cfg["dataloader_prefetch_factor"] == 2
|
| 49 |
+
assert cfg["dataloader_persistent_workers"] is False
|
| 50 |
+
|
| 51 |
+
def test_torch_compile_call_flips_flag(self) -> None:
|
| 52 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 53 |
+
# The script calls torch.compile(model, ...), even though the
|
| 54 |
+
# TrainingArguments has torch_compile=False — the explicit call wins.
|
| 55 |
+
assert result.result["torch_compile"] is True
|
| 56 |
+
|
| 57 |
+
def test_gradient_checkpointing_enable_call(self) -> None:
|
| 58 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 59 |
+
assert result.result["gradient_checkpointing"] is True
|
| 60 |
+
|
| 61 |
+
def test_env_vars_captured(self) -> None:
|
| 62 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 63 |
+
env = result.result["env_vars"]
|
| 64 |
+
assert env["HSA_FORCE_FINE_GRAIN_PCIE"] == "1"
|
| 65 |
+
assert env["MIOPEN_FIND_MODE"] == "3"
|
| 66 |
+
assert env["NCCL_MIN_NCHANNELS"] == "112"
|
| 67 |
+
|
| 68 |
+
def test_lora_rank_extracted(self) -> None:
|
| 69 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 70 |
+
assert result.result["lora_rank"] == 16
|
| 71 |
+
|
| 72 |
+
def test_attention_impl_from_from_pretrained(self) -> None:
|
| 73 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 74 |
+
assert result.result["attention_impl"] == "eager"
|
| 75 |
+
|
| 76 |
+
def test_model_name_resolved(self) -> None:
|
| 77 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 78 |
+
assert result.result["model_name"] == "Qwen/Qwen2.5-7B-Instruct"
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
# ---------------------------------------------------------------------------
|
| 82 |
+
# JSON path
|
| 83 |
+
# ---------------------------------------------------------------------------
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
class TestJsonConfig:
|
| 87 |
+
def test_returns_ok(self) -> None:
|
| 88 |
+
result = _parse_config(str(FIXTURES / "sample_train.json"))
|
| 89 |
+
assert result.ok, result.error
|
| 90 |
+
|
| 91 |
+
def test_field_mapping(self) -> None:
|
| 92 |
+
cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
|
| 93 |
+
assert cfg["model_name"] == "Qwen/Qwen2.5-7B-Instruct"
|
| 94 |
+
assert cfg["batch_size"] == 8
|
| 95 |
+
assert cfg["grad_accum_steps"] == 4
|
| 96 |
+
assert cfg["seq_len"] == 4096
|
| 97 |
+
assert cfg["precision"] == "bf16"
|
| 98 |
+
assert cfg["optimizer"] == "adamw_torch_fused"
|
| 99 |
+
assert cfg["torch_compile"] is True
|
| 100 |
+
assert cfg["gradient_checkpointing"] is True
|
| 101 |
+
assert cfg["dataloader_workers"] == 4
|
| 102 |
+
assert cfg["dataloader_pin_memory"] is True
|
| 103 |
+
assert cfg["dataloader_persistent_workers"] is True
|
| 104 |
+
assert cfg["attention_impl"] == "flash"
|
| 105 |
+
|
| 106 |
+
def test_env_vars_dict(self) -> None:
|
| 107 |
+
cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
|
| 108 |
+
assert cfg["env_vars"]["HSA_FORCE_FINE_GRAIN_PCIE"] == "1"
|
| 109 |
+
|
| 110 |
+
def test_extras_collects_unmapped_fields(self) -> None:
|
| 111 |
+
cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
|
| 112 |
+
# num_train_epochs has no slot in WorkloadConfig — must land in extras.
|
| 113 |
+
assert cfg["extras"]["num_train_epochs"] == 3
|
| 114 |
+
assert cfg["extras"]["save_steps"] == 500
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
# ---------------------------------------------------------------------------
|
| 118 |
+
# YAML path
|
| 119 |
+
# ---------------------------------------------------------------------------
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
class TestYamlConfig:
|
| 123 |
+
def test_returns_ok(self) -> None:
|
| 124 |
+
result = _parse_config(str(FIXTURES / "sample_train.yaml"))
|
| 125 |
+
assert result.ok, result.error
|
| 126 |
+
|
| 127 |
+
def test_field_mapping(self) -> None:
|
| 128 |
+
cfg = _parse_config(str(FIXTURES / "sample_train.yaml")).result
|
| 129 |
+
assert cfg["batch_size"] == 2
|
| 130 |
+
assert cfg["grad_accum_steps"] == 16
|
| 131 |
+
assert cfg["seq_len"] == 8192
|
| 132 |
+
assert cfg["precision"] == "bf16"
|
| 133 |
+
assert cfg["torch_compile"] is False
|
| 134 |
+
assert cfg["gradient_checkpointing"] is True
|
| 135 |
+
assert cfg["attention_impl"] == "sdpa"
|
| 136 |
+
assert cfg["dataloader_workers"] == 8
|
| 137 |
+
assert cfg["dataloader_persistent_workers"] is True
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
# ---------------------------------------------------------------------------
|
| 141 |
+
# Redaction
|
| 142 |
+
# ---------------------------------------------------------------------------
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
class TestRedaction:
|
| 146 |
+
def test_python_script_redactions(self) -> None:
|
| 147 |
+
cfg = _parse_config(str(FIXTURES / "sample_train.py")).result
|
| 148 |
+
labels = set(cfg["redactions"])
|
| 149 |
+
# Every secret pattern in sample_train.py should fire.
|
| 150 |
+
assert "hf_token" in labels
|
| 151 |
+
assert "openai_key" in labels
|
| 152 |
+
assert "github_token" in labels
|
| 153 |
+
assert "bearer_token" in labels
|
| 154 |
+
assert "home_path" in labels
|
| 155 |
+
assert "s3_uri" in labels
|
| 156 |
+
assert "ws_uri" in labels
|
| 157 |
+
|
| 158 |
+
def test_raw_source_is_scrubbed(self) -> None:
|
| 159 |
+
# `raw_source` is intentionally stripped from the tool result envelope
|
| 160 |
+
# (keeps the LLM conversation small) — use the `_full` helper to read
|
| 161 |
+
# it. The redaction labels list still proves which patterns fired.
|
| 162 |
+
cfg = _parse_config_full(str(FIXTURES / "sample_train.py"))
|
| 163 |
+
assert isinstance(cfg, WorkloadConfig)
|
| 164 |
+
raw = cfg.raw_source
|
| 165 |
+
assert "hf_abcdefghijklmnopqrstuvwxyz123456" not in raw
|
| 166 |
+
assert "sk-abcdefghijklmnopqrstuvwxyz1234567890" not in raw
|
| 167 |
+
assert "gho_abcdefghijklmnopqrstuvwxyz123456" not in raw
|
| 168 |
+
assert "/home/researcher/datasets/alpaca" not in raw
|
| 169 |
+
assert "s3://my-team/checkpoints/qwen-lora/" not in raw
|
| 170 |
+
assert "wss://logs.internal.example.com/stream" not in raw
|
| 171 |
+
assert "<REDACTED:hf_token>" in raw
|
| 172 |
+
assert "<REDACTED:openai_key>" in raw
|
| 173 |
+
|
| 174 |
+
def test_raw_source_excluded_from_tool_result(self) -> None:
|
| 175 |
+
# The tool result MUST NOT carry raw_source — it bloated the audit
|
| 176 |
+
# conversation past 8K on Qwen2.5-7B during the live AMD GPU run.
|
| 177 |
+
cfg = _parse_config(str(FIXTURES / "sample_train.py")).result
|
| 178 |
+
assert "raw_source" not in cfg
|
| 179 |
+
|
| 180 |
+
def test_json_redactions(self) -> None:
|
| 181 |
+
cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
|
| 182 |
+
labels = set(cfg["redactions"])
|
| 183 |
+
assert "hf_token" in labels
|
| 184 |
+
assert "s3_uri" in labels
|
| 185 |
+
# raw_source is no longer in the result; verify scrubbing via the
|
| 186 |
+
# full-config helper.
|
| 187 |
+
full = _parse_config_full(str(FIXTURES / "sample_train.json"))
|
| 188 |
+
assert isinstance(full, WorkloadConfig)
|
| 189 |
+
assert "hf_jsonsamplehfabcdefghijklmnopqrs" not in full.raw_source
|
| 190 |
+
|
| 191 |
+
def test_extras_values_are_scrubbed(self) -> None:
|
| 192 |
+
# Secret-shaped values that landed in extras must also be redacted —
|
| 193 |
+
# otherwise the leak just moves from raw_source into extras.
|
| 194 |
+
cfg = _parse_config(str(FIXTURES / "sample_train.json")).result
|
| 195 |
+
extras = cfg["extras"]
|
| 196 |
+
assert "hf_jsonsamplehfabcdefghijklmnopqrs" not in extras.get("hub_token", "")
|
| 197 |
+
assert extras.get("hub_token", "").startswith("<REDACTED:")
|
| 198 |
+
assert extras.get("checkpoint_uri", "").startswith("<REDACTED:")
|
| 199 |
+
|
| 200 |
+
def test_yaml_redactions(self) -> None:
|
| 201 |
+
cfg = _parse_config(str(FIXTURES / "sample_train.yaml")).result
|
| 202 |
+
labels = set(cfg["redactions"])
|
| 203 |
+
assert "hf_token" in labels
|
| 204 |
+
assert "bearer_token" in labels
|
| 205 |
+
assert "home_path" in labels
|
| 206 |
+
|
| 207 |
+
|
| 208 |
+
# ---------------------------------------------------------------------------
|
| 209 |
+
# Failure modes
|
| 210 |
+
# ---------------------------------------------------------------------------
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
class TestErrors:
|
| 214 |
+
def test_missing_file(self) -> None:
|
| 215 |
+
result = _parse_config("/tmp/definitely-does-not-exist-xyz.py")
|
| 216 |
+
assert result.ok is False
|
| 217 |
+
assert "not found" in (result.error or "").lower()
|
| 218 |
+
|
| 219 |
+
def test_unsupported_extension(self, tmp_path: Path) -> None:
|
| 220 |
+
bad = tmp_path / "config.toml"
|
| 221 |
+
bad.write_text("model_name = 'foo'\n")
|
| 222 |
+
result = _parse_config(str(bad))
|
| 223 |
+
assert result.ok is False
|
| 224 |
+
assert "unsupported" in (result.error or "").lower()
|
| 225 |
+
|
| 226 |
+
def test_malformed_python(self, tmp_path: Path) -> None:
|
| 227 |
+
bad = tmp_path / "broken.py"
|
| 228 |
+
bad.write_text("def oops(:\n pass\n")
|
| 229 |
+
result = _parse_config(str(bad))
|
| 230 |
+
assert result.ok is False
|
| 231 |
+
assert "parse error" in (result.error or "").lower()
|
| 232 |
+
|
| 233 |
+
def test_malformed_json(self, tmp_path: Path) -> None:
|
| 234 |
+
bad = tmp_path / "broken.json"
|
| 235 |
+
bad.write_text("{not really json")
|
| 236 |
+
result = _parse_config(str(bad))
|
| 237 |
+
assert result.ok is False
|
| 238 |
+
assert "json" in (result.error or "").lower()
|
| 239 |
+
|
| 240 |
+
def test_json_top_level_must_be_dict(self, tmp_path: Path) -> None:
|
| 241 |
+
bad = tmp_path / "list.json"
|
| 242 |
+
bad.write_text(json.dumps([{"foo": 1}]))
|
| 243 |
+
result = _parse_config(str(bad))
|
| 244 |
+
assert result.ok is False
|
| 245 |
+
|
| 246 |
+
def test_yaml_top_level_must_be_mapping(self, tmp_path: Path) -> None:
|
| 247 |
+
bad = tmp_path / "scalar.yaml"
|
| 248 |
+
bad.write_text("- 1\n- 2\n")
|
| 249 |
+
result = _parse_config(str(bad))
|
| 250 |
+
assert result.ok is False
|
| 251 |
+
|
| 252 |
+
|
| 253 |
+
# ---------------------------------------------------------------------------
|
| 254 |
+
# Schema invariants
|
| 255 |
+
# ---------------------------------------------------------------------------
|
| 256 |
+
|
| 257 |
+
|
| 258 |
+
class TestSchema:
|
| 259 |
+
def test_result_round_trips_through_workload_config(self) -> None:
|
| 260 |
+
result = _parse_config(str(FIXTURES / "sample_train.py"))
|
| 261 |
+
# Must be reconstructible — guards against extras-vs-fields collisions.
|
| 262 |
+
cfg = WorkloadConfig(**result.result)
|
| 263 |
+
assert cfg.model_name == "Qwen/Qwen2.5-7B-Instruct"
|
| 264 |
+
|
| 265 |
+
def test_defaults_when_field_absent(self, tmp_path: Path) -> None:
|
| 266 |
+
# Minimal config — only model_name. Everything else should fall back to schema defaults.
|
| 267 |
+
path = tmp_path / "tiny.json"
|
| 268 |
+
path.write_text(json.dumps({"model_name": "test/tiny"}))
|
| 269 |
+
result = _parse_config(str(path))
|
| 270 |
+
assert result.ok
|
| 271 |
+
cfg = result.result
|
| 272 |
+
assert cfg["batch_size"] == 1
|
| 273 |
+
assert cfg["precision"] == "fp16"
|
| 274 |
+
assert cfg["optimizer"] == "adamw_torch"
|
| 275 |
+
assert cfg["gradient_checkpointing"] is False
|
| 276 |
+
assert cfg["redactions"] == []
|
| 277 |
+
|
| 278 |
+
def test_tool_definition_unchanged_in_shape(self) -> None:
|
| 279 |
+
# The Tool definition should still expose name/description/input_schema/fn.
|
| 280 |
+
assert PARSE_CONFIG.name == "parse_config"
|
| 281 |
+
assert PARSE_CONFIG.fn is _parse_config
|
| 282 |
+
assert "file_path" in PARSE_CONFIG.input_schema["properties"]
|
| 283 |
+
|
| 284 |
+
|
| 285 |
+
# ---------------------------------------------------------------------------
|
| 286 |
+
# Regression: canonical + scenario workloads must parse with all the right
|
| 287 |
+
# audit-relevant fields. These are what the live agent actually sees, so a
|
| 288 |
+
# regression here directly degrades audit quality (the agent reasons over
|
| 289 |
+
# HF defaults instead of the script's settings).
|
| 290 |
+
# ---------------------------------------------------------------------------
|
| 291 |
+
|
| 292 |
+
|
| 293 |
+
REPO_ROOT = Path(__file__).resolve().parent.parent
|
| 294 |
+
|
| 295 |
+
|
| 296 |
+
class TestCanonicalWorkload:
|
| 297 |
+
def test_canonical_workload_extracts_full_config(self) -> None:
|
| 298 |
+
"""The canonical demo workload must yield batch_size=4, lr=2e-4, etc.
|
| 299 |
+
— not HF defaults. Catches the `**dict_var` splat regression where
|
| 300 |
+
every TrainingArguments kwarg disappears.
|
| 301 |
+
"""
|
| 302 |
+
result = _parse_config(str(REPO_ROOT / "workloads" / "train_qwen_lora.py"))
|
| 303 |
+
assert result.ok, result.error
|
| 304 |
+
cfg = result.result
|
| 305 |
+
assert cfg["model_name"] == "Qwen/Qwen2.5-7B-Instruct"
|
| 306 |
+
assert cfg["batch_size"] == 4, (
|
| 307 |
+
"expected batch_size=4 from per_device_train_batch_size; "
|
| 308 |
+
"did `**_ta_kwargs` splat hide the kwargs?"
|
| 309 |
+
)
|
| 310 |
+
assert cfg["grad_accum_steps"] == 8
|
| 311 |
+
assert cfg["lr"] == 2e-4
|
| 312 |
+
assert cfg["warmup_steps"] == 100
|
| 313 |
+
assert cfg["precision"] == "fp16"
|
| 314 |
+
assert cfg["attention_impl"] == "eager"
|
| 315 |
+
assert cfg["dataloader_workers"] == 0
|
| 316 |
+
assert cfg["dataloader_pin_memory"] is False
|
| 317 |
+
assert cfg["lora_rank"] == 16
|
| 318 |
+
assert cfg["torch_compile"] is False
|
| 319 |
+
assert cfg["env_vars"]["HSA_FORCE_FINE_GRAIN_PCIE"] == "1"
|
| 320 |
+
|
| 321 |
+
|
| 322 |
+
class TestSplatKwargsResolution:
|
| 323 |
+
"""`_ta = dict(k=v); Foo(**_ta)` must resolve back through the dict
|
| 324 |
+
constant. Defensive — the canonical workload no longer uses this
|
| 325 |
+
pattern, but third-party scripts often do.
|
| 326 |
+
"""
|
| 327 |
+
|
| 328 |
+
def test_dict_function_call_splat(self, tmp_path) -> None:
|
| 329 |
+
src = """
|
| 330 |
+
from transformers import TrainingArguments
|
| 331 |
+
|
| 332 |
+
_ta = dict(
|
| 333 |
+
per_device_train_batch_size=8,
|
| 334 |
+
gradient_accumulation_steps=2,
|
| 335 |
+
fp16=True,
|
| 336 |
+
optim=\"adamw_torch_fused\",
|
| 337 |
+
)
|
| 338 |
+
training_args = TrainingArguments(output_dir=\"./out\", **_ta)
|
| 339 |
+
"""
|
| 340 |
+
p = tmp_path / "splat.py"
|
| 341 |
+
p.write_text(src)
|
| 342 |
+
cfg = _parse_config(str(p)).result
|
| 343 |
+
assert cfg["batch_size"] == 8
|
| 344 |
+
assert cfg["grad_accum_steps"] == 2
|
| 345 |
+
assert cfg["precision"] == "fp16"
|
| 346 |
+
assert cfg["optimizer"] == "adamw_torch_fused"
|
| 347 |
+
|
| 348 |
+
def test_dict_literal_splat(self, tmp_path) -> None:
|
| 349 |
+
src = """
|
| 350 |
+
from transformers import TrainingArguments
|
| 351 |
+
|
| 352 |
+
_ta = {
|
| 353 |
+
"per_device_train_batch_size": 16,
|
| 354 |
+
"bf16": True,
|
| 355 |
+
}
|
| 356 |
+
training_args = TrainingArguments(output_dir=\"./out\", **_ta)
|
| 357 |
+
"""
|
| 358 |
+
p = tmp_path / "splat_literal.py"
|
| 359 |
+
p.write_text(src)
|
| 360 |
+
cfg = _parse_config(str(p)).result
|
| 361 |
+
assert cfg["batch_size"] == 16
|
| 362 |
+
assert cfg["precision"] == "bf16"
|
| 363 |
+
|
| 364 |
+
def test_explicit_kwarg_overrides_splat(self, tmp_path) -> None:
|
| 365 |
+
src = """
|
| 366 |
+
from transformers import TrainingArguments
|
| 367 |
+
|
| 368 |
+
_ta = dict(per_device_train_batch_size=8)
|
| 369 |
+
training_args = TrainingArguments(per_device_train_batch_size=32, **_ta)
|
| 370 |
+
"""
|
| 371 |
+
p = tmp_path / "splat_override.py"
|
| 372 |
+
p.write_text(src)
|
| 373 |
+
cfg = _parse_config(str(p)).result
|
| 374 |
+
# Explicit kwarg wins over splat (both occur in the kwargs list,
|
| 375 |
+
# explicit comes first in the AST → setdefault keeps it).
|
| 376 |
+
assert cfg["batch_size"] == 32
|
| 377 |
+
|
tests/test_query_rocm_kb.py
ADDED
|
@@ -0,0 +1,348 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for ``agent.tools.query_rocm_kb``.
|
| 2 |
+
|
| 3 |
+
Coverage:
|
| 4 |
+
* The shipped ``kb/rocm_rules.yaml`` is loadable, validates against
|
| 5 |
+
:class:`Rule`, and contains every required-by-spec rule id.
|
| 6 |
+
* Each rule's ``targets_bucket`` is a valid :class:`WasteBucket` and its
|
| 7 |
+
``category`` is a valid :class:`RuleCategory` (caught for free by
|
| 8 |
+
pydantic, but assert here for a clearer failure when the YAML drifts).
|
| 9 |
+
* Semantic search returns the bf16 rule first for an "fp16 on MI300X"
|
| 10 |
+
query and the flash-attn / sdpa rules first for an attention query.
|
| 11 |
+
* Bad inputs (empty symptom, ``top_k <= 0``, ``top_k`` larger than the
|
| 12 |
+
rule count) are handled gracefully via ``ToolResult.ok=False`` or
|
| 13 |
+
clamped to the rule count.
|
| 14 |
+
* Embeddings cache: the cache file is created on first import, and a
|
| 15 |
+
second ``_embed_rules`` call with the same YAML bytes hits the cache
|
| 16 |
+
without re-encoding.
|
| 17 |
+
* The frozen ``Tool`` definition retains its ``name``, ``description``,
|
| 18 |
+
and ``input_schema`` shape — the agent registry depends on these.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
from __future__ import annotations
|
| 22 |
+
|
| 23 |
+
from pathlib import Path
|
| 24 |
+
from typing import get_args
|
| 25 |
+
|
| 26 |
+
import pytest
|
| 27 |
+
import yaml
|
| 28 |
+
|
| 29 |
+
from agent.schemas import Rule, RuleCategory, ToolResult, WasteBucket
|
| 30 |
+
from agent.tools.query_rocm_kb import (
|
| 31 |
+
_KB_YAML,
|
| 32 |
+
_RULES,
|
| 33 |
+
QUERY_ROCM_KB,
|
| 34 |
+
_cache_path,
|
| 35 |
+
_embed_rules,
|
| 36 |
+
_load_rules,
|
| 37 |
+
_query_rocm_kb,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
REQUIRED_RULE_IDS = {
|
| 42 |
+
"precision.bf16_over_fp16_on_mi300x",
|
| 43 |
+
"attention.flash_rocm_over_eager",
|
| 44 |
+
"attention.sdpa_over_eager",
|
| 45 |
+
"memory.batch_too_small_for_192gb",
|
| 46 |
+
"memory.gradient_checkpointing_for_long_seq",
|
| 47 |
+
"data.dataloader_workers_zero",
|
| 48 |
+
"data.pin_memory_false",
|
| 49 |
+
"data.prefetch_factor_default",
|
| 50 |
+
"data.persistent_workers_false",
|
| 51 |
+
"compile.torch_compile_off",
|
| 52 |
+
"env.nccl_min_nchannels",
|
| 53 |
+
"env.numa_auto_balancing_disable",
|
| 54 |
+
"env.hsa_force_fine_grain_pcie",
|
| 55 |
+
"kernels.hipblaslt_hint_logging",
|
| 56 |
+
"kernels.miopen_find_mode_2",
|
| 57 |
+
"optimizer.bitsandbytes_not_supported_warning",
|
| 58 |
+
"collectives.one_process_per_gpu",
|
| 59 |
+
"topology.tp_within_xgmi_island",
|
| 60 |
+
}
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
# ---------------------------------------------------------------------------
|
| 64 |
+
# YAML KB invariants
|
| 65 |
+
# ---------------------------------------------------------------------------
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
class TestYamlKB:
|
| 69 |
+
def test_kb_yaml_exists(self) -> None:
|
| 70 |
+
assert _KB_YAML.exists(), f"Missing {_KB_YAML}"
|
| 71 |
+
|
| 72 |
+
def test_kb_loads_and_validates(self) -> None:
|
| 73 |
+
rules, _raw = _load_rules(_KB_YAML)
|
| 74 |
+
# Every rule pydantic-validated.
|
| 75 |
+
for r in rules:
|
| 76 |
+
assert isinstance(r, Rule)
|
| 77 |
+
# Spec calls for 20-25 rules; allow a small wiggle room either way.
|
| 78 |
+
assert 18 <= len(rules) <= 30, f"Expected 18-30 rules, got {len(rules)}"
|
| 79 |
+
|
| 80 |
+
def test_required_rule_ids_present(self) -> None:
|
| 81 |
+
ids = {r.id for r in _RULES}
|
| 82 |
+
missing = REQUIRED_RULE_IDS - ids
|
| 83 |
+
assert not missing, f"Required rule ids missing from KB: {sorted(missing)}"
|
| 84 |
+
|
| 85 |
+
def test_rule_ids_unique(self) -> None:
|
| 86 |
+
ids = [r.id for r in _RULES]
|
| 87 |
+
dupes = {i for i in ids if ids.count(i) > 1}
|
| 88 |
+
assert not dupes, f"Duplicate rule ids: {sorted(dupes)}"
|
| 89 |
+
|
| 90 |
+
def test_categories_are_valid(self) -> None:
|
| 91 |
+
valid = set(get_args(RuleCategory))
|
| 92 |
+
for r in _RULES:
|
| 93 |
+
assert r.category in valid, f"{r.id}: invalid category {r.category!r}"
|
| 94 |
+
|
| 95 |
+
def test_targets_bucket_valid(self) -> None:
|
| 96 |
+
valid = set(get_args(WasteBucket))
|
| 97 |
+
for r in _RULES:
|
| 98 |
+
assert r.targets_bucket in valid, (
|
| 99 |
+
f"{r.id}: invalid targets_bucket {r.targets_bucket!r}"
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
def test_recovery_fraction_in_range(self) -> None:
|
| 103 |
+
for r in _RULES:
|
| 104 |
+
assert 0.0 <= r.expected_recovery_fraction <= 1.0, (
|
| 105 |
+
f"{r.id}: expected_recovery_fraction out of [0,1] "
|
| 106 |
+
f"({r.expected_recovery_fraction})"
|
| 107 |
+
)
|
| 108 |
+
|
| 109 |
+
def test_citations_non_empty(self) -> None:
|
| 110 |
+
for r in _RULES:
|
| 111 |
+
assert r.citation and r.citation.strip(), f"{r.id}: empty citation"
|
| 112 |
+
|
| 113 |
+
def test_bitsandbytes_is_warning_only(self) -> None:
|
| 114 |
+
rule = next(r for r in _RULES if r.id == "optimizer.bitsandbytes_not_supported_warning")
|
| 115 |
+
# Warning rule has empty transform — propose_patch must not auto-fix.
|
| 116 |
+
assert rule.transform == {}
|
| 117 |
+
|
| 118 |
+
def test_categories_cover_spec(self) -> None:
|
| 119 |
+
# Architecture §4 lists 10 categories. We expect rules in at least
|
| 120 |
+
# the high-impact ones the spec calls out as required.
|
| 121 |
+
cats = {r.category for r in _RULES}
|
| 122 |
+
for required in (
|
| 123 |
+
"precision",
|
| 124 |
+
"attention",
|
| 125 |
+
"memory",
|
| 126 |
+
"data",
|
| 127 |
+
"compile",
|
| 128 |
+
"env_vars",
|
| 129 |
+
"kernels",
|
| 130 |
+
"optimizer",
|
| 131 |
+
"collectives",
|
| 132 |
+
"topology",
|
| 133 |
+
):
|
| 134 |
+
assert required in cats, f"No rule for category {required!r}"
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
# ---------------------------------------------------------------------------
|
| 138 |
+
# Skip-and-warn behaviour for invalid entries
|
| 139 |
+
# ---------------------------------------------------------------------------
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
class TestLoadRulesResilience:
|
| 143 |
+
def test_invalid_entry_is_skipped_with_warning(self, tmp_path: Path) -> None:
|
| 144 |
+
bad = tmp_path / "rules.yaml"
|
| 145 |
+
bad.write_text(
|
| 146 |
+
yaml.safe_dump(
|
| 147 |
+
[
|
| 148 |
+
{
|
| 149 |
+
"id": "good.rule",
|
| 150 |
+
"category": "precision",
|
| 151 |
+
"targets_bucket": "precision_path",
|
| 152 |
+
"symptom": "fp16 used",
|
| 153 |
+
"expected_impact": "switch to bf16",
|
| 154 |
+
"citation": "ROCm guide",
|
| 155 |
+
},
|
| 156 |
+
{"id": "bad.rule", "category": "not_a_real_category"},
|
| 157 |
+
]
|
| 158 |
+
)
|
| 159 |
+
)
|
| 160 |
+
with pytest.warns(UserWarning):
|
| 161 |
+
rules, _raw = _load_rules(bad)
|
| 162 |
+
assert [r.id for r in rules] == ["good.rule"]
|
| 163 |
+
|
| 164 |
+
def test_top_level_must_be_list(self, tmp_path: Path) -> None:
|
| 165 |
+
bad = tmp_path / "rules.yaml"
|
| 166 |
+
bad.write_text("not_a_list: 1\n")
|
| 167 |
+
with pytest.raises(ValueError, match="top-level"):
|
| 168 |
+
_load_rules(bad)
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
# ---------------------------------------------------------------------------
|
| 172 |
+
# Semantic search behaviour
|
| 173 |
+
# ---------------------------------------------------------------------------
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
class TestQuery:
|
| 177 |
+
def test_returns_ok_for_real_query(self) -> None:
|
| 178 |
+
result = _query_rocm_kb("fp16 used on MI300X with eager attention", top_k=5)
|
| 179 |
+
assert isinstance(result, ToolResult)
|
| 180 |
+
assert result.ok, result.error
|
| 181 |
+
rules = result.result["rules"]
|
| 182 |
+
assert 1 <= len(rules) <= 5
|
| 183 |
+
|
| 184 |
+
def test_fp16_query_returns_bf16_rule_in_top_results(self) -> None:
|
| 185 |
+
result = _query_rocm_kb("fp16 used on MI300X / CDNA3", top_k=3)
|
| 186 |
+
ids = [r["id"] for r in result.result["rules"]]
|
| 187 |
+
assert "precision.bf16_over_fp16_on_mi300x" in ids
|
| 188 |
+
|
| 189 |
+
def test_eager_attention_query_returns_attention_rules(self) -> None:
|
| 190 |
+
result = _query_rocm_kb(
|
| 191 |
+
"eager attention with no flash kernel loaded on MI300X", top_k=3
|
| 192 |
+
)
|
| 193 |
+
ids = [r["id"] for r in result.result["rules"]]
|
| 194 |
+
attention_ids = {
|
| 195 |
+
"attention.flash_rocm_over_eager",
|
| 196 |
+
"attention.sdpa_over_eager",
|
| 197 |
+
}
|
| 198 |
+
assert attention_ids & set(ids), (
|
| 199 |
+
f"Expected at least one of {attention_ids} in top 3, got {ids}"
|
| 200 |
+
)
|
| 201 |
+
|
| 202 |
+
def test_dataloader_query_returns_data_rules(self) -> None:
|
| 203 |
+
result = _query_rocm_kb(
|
| 204 |
+
"DataLoader num_workers is zero, GPU starves waiting for batches",
|
| 205 |
+
top_k=3,
|
| 206 |
+
)
|
| 207 |
+
ids = [r["id"] for r in result.result["rules"]]
|
| 208 |
+
data_ids = {
|
| 209 |
+
"data.dataloader_workers_zero",
|
| 210 |
+
"data.pin_memory_false",
|
| 211 |
+
"data.prefetch_factor_default",
|
| 212 |
+
"data.persistent_workers_false",
|
| 213 |
+
}
|
| 214 |
+
assert data_ids & set(ids), (
|
| 215 |
+
f"Expected at least one data.* rule in top 3, got {ids}"
|
| 216 |
+
)
|
| 217 |
+
|
| 218 |
+
def test_top_k_bounds_returned(self) -> None:
|
| 219 |
+
result = _query_rocm_kb("anything", top_k=2)
|
| 220 |
+
assert len(result.result["rules"]) == 2
|
| 221 |
+
|
| 222 |
+
def test_top_k_clamped_to_rule_count(self) -> None:
|
| 223 |
+
# Asking for more than we have should not crash; we should get every
|
| 224 |
+
# rule back, ordered by score.
|
| 225 |
+
result = _query_rocm_kb("anything", top_k=100)
|
| 226 |
+
assert len(result.result["rules"]) == len(_RULES)
|
| 227 |
+
|
| 228 |
+
def test_results_sorted_by_descending_score(self) -> None:
|
| 229 |
+
# If two queries with different focus return different top rules,
|
| 230 |
+
# ordering is real — top_1 differs by query.
|
| 231 |
+
a = _query_rocm_kb("fp16 numerical instability", top_k=1).result["rules"][0]
|
| 232 |
+
b = _query_rocm_kb("eager attention slow on long sequences", top_k=1).result[
|
| 233 |
+
"rules"
|
| 234 |
+
][0]
|
| 235 |
+
assert a["id"] != b["id"], (
|
| 236 |
+
"Top rule should depend on query, but both queries returned "
|
| 237 |
+
f"{a['id']} — semantic search is degenerate."
|
| 238 |
+
)
|
| 239 |
+
|
| 240 |
+
def test_rule_payload_is_lite_shape(self) -> None:
|
| 241 |
+
# The LLM-facing rule payload is intentionally trimmed from the full
|
| 242 |
+
# Rule schema — only id / symptom / transform / expected_impact /
|
| 243 |
+
# citation make it through. This shrinks the audit conversation enough
|
| 244 |
+
# to fit Qwen2.5-7B's 8K window. Full Rule lookup happens server-side
|
| 245 |
+
# in propose_patch via the loaded KB.
|
| 246 |
+
result = _query_rocm_kb("any query", top_k=1)
|
| 247 |
+
payload = result.result["rules"][0]
|
| 248 |
+
assert set(payload.keys()) == {
|
| 249 |
+
"id",
|
| 250 |
+
"symptom",
|
| 251 |
+
"transform",
|
| 252 |
+
"expected_impact",
|
| 253 |
+
"citation",
|
| 254 |
+
}
|
| 255 |
+
# And the id resolves against the loaded KB so propose_patch can
|
| 256 |
+
# reconstruct the full Rule.
|
| 257 |
+
from agent.tools.query_rocm_kb import _RULES
|
| 258 |
+
|
| 259 |
+
kb_ids = {r.id for r in _RULES}
|
| 260 |
+
assert payload["id"] in kb_ids
|
| 261 |
+
|
| 262 |
+
|
| 263 |
+
# ---------------------------------------------------------------------------
|
| 264 |
+
# Failure modes
|
| 265 |
+
# ---------------------------------------------------------------------------
|
| 266 |
+
|
| 267 |
+
|
| 268 |
+
class TestErrors:
|
| 269 |
+
def test_empty_symptom(self) -> None:
|
| 270 |
+
result = _query_rocm_kb("", top_k=3)
|
| 271 |
+
assert result.ok is False
|
| 272 |
+
assert "symptom" in (result.error or "")
|
| 273 |
+
|
| 274 |
+
def test_whitespace_only_symptom(self) -> None:
|
| 275 |
+
result = _query_rocm_kb(" \t\n ", top_k=3)
|
| 276 |
+
assert result.ok is False
|
| 277 |
+
|
| 278 |
+
def test_top_k_zero(self) -> None:
|
| 279 |
+
result = _query_rocm_kb("fp16", top_k=0)
|
| 280 |
+
assert result.ok is False
|
| 281 |
+
assert "top_k" in (result.error or "")
|
| 282 |
+
|
| 283 |
+
def test_top_k_negative(self) -> None:
|
| 284 |
+
result = _query_rocm_kb("fp16", top_k=-3)
|
| 285 |
+
assert result.ok is False
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
# ---------------------------------------------------------------------------
|
| 289 |
+
# Embeddings cache
|
| 290 |
+
# ---------------------------------------------------------------------------
|
| 291 |
+
|
| 292 |
+
|
| 293 |
+
class TestEmbeddingsCache:
|
| 294 |
+
def test_cache_file_was_created_on_import(self) -> None:
|
| 295 |
+
raw = _KB_YAML.read_bytes()
|
| 296 |
+
cache = _cache_path(raw)
|
| 297 |
+
assert cache.exists(), (
|
| 298 |
+
f"Expected embeddings cache at {cache}; cache write failed silently."
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
def test_second_embed_call_uses_cache_without_recoding(
|
| 302 |
+
self, monkeypatch: pytest.MonkeyPatch
|
| 303 |
+
) -> None:
|
| 304 |
+
rules, raw = _load_rules(_KB_YAML)
|
| 305 |
+
# Sentinel: replace the lazy model getter so any encode() call would
|
| 306 |
+
# blow up. If the cache is hit, the model is never consulted.
|
| 307 |
+
from agent.tools import query_rocm_kb as kb_module
|
| 308 |
+
|
| 309 |
+
def explode() -> None:
|
| 310 |
+
raise AssertionError(
|
| 311 |
+
"_get_model() called on a cache-hit path; cache is not being used."
|
| 312 |
+
)
|
| 313 |
+
|
| 314 |
+
monkeypatch.setattr(kb_module, "_get_model", explode)
|
| 315 |
+
embeddings = _embed_rules(rules, raw)
|
| 316 |
+
assert embeddings.shape[0] == len(rules)
|
| 317 |
+
assert embeddings.dtype.kind == "f"
|
| 318 |
+
|
| 319 |
+
|
| 320 |
+
# ---------------------------------------------------------------------------
|
| 321 |
+
# Tool registry shape — the agent loop depends on these fields being stable.
|
| 322 |
+
# ---------------------------------------------------------------------------
|
| 323 |
+
|
| 324 |
+
|
| 325 |
+
class TestToolDefinition:
|
| 326 |
+
def test_name_is_query_rocm_kb(self) -> None:
|
| 327 |
+
assert QUERY_ROCM_KB.name == "query_rocm_kb"
|
| 328 |
+
|
| 329 |
+
def test_description_unchanged_keywords(self) -> None:
|
| 330 |
+
# Must still describe the search-by-symptom semantics; the system
|
| 331 |
+
# prompt references this language.
|
| 332 |
+
desc = QUERY_ROCM_KB.description
|
| 333 |
+
assert "ROCm" in desc and "symptom" in desc
|
| 334 |
+
|
| 335 |
+
def test_input_schema_shape(self) -> None:
|
| 336 |
+
schema = QUERY_ROCM_KB.input_schema
|
| 337 |
+
assert schema["type"] == "object"
|
| 338 |
+
# Both single-query (`symptom`) and batched (`symptoms`) shapes are
|
| 339 |
+
# advertised — either works at runtime, neither is strictly required
|
| 340 |
+
# in the schema because the impl validates "at least one" itself.
|
| 341 |
+
assert "symptom" in schema["properties"]
|
| 342 |
+
assert "symptoms" in schema["properties"]
|
| 343 |
+
assert "top_k" in schema["properties"]
|
| 344 |
+
assert schema["properties"]["symptoms"]["type"] == "array"
|
| 345 |
+
assert schema["properties"]["symptoms"]["items"] == {"type": "string"}
|
| 346 |
+
|
| 347 |
+
def test_fn_is_module_query(self) -> None:
|
| 348 |
+
assert QUERY_ROCM_KB.fn is _query_rocm_kb
|
tests/test_qwen_vllm_backend.py
ADDED
|
@@ -0,0 +1,280 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for QwenVLLMBackend.
|
| 2 |
+
|
| 3 |
+
Mocks the openai AsyncClient at the chat.completions level so no real
|
| 4 |
+
network call ever happens. Verifies the OpenAI-shape conversation
|
| 5 |
+
threading, tool-call extraction, and finish_reason mapping.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
import os
|
| 11 |
+
from types import SimpleNamespace
|
| 12 |
+
from typing import Any
|
| 13 |
+
|
| 14 |
+
import pytest
|
| 15 |
+
|
| 16 |
+
from agent.backends import active_backend_name, make_backend
|
| 17 |
+
from agent.backends.base import AgentTurn
|
| 18 |
+
from agent.backends.qwen_vllm import (
|
| 19 |
+
QwenVLLMBackend,
|
| 20 |
+
_normalize_finish_reason,
|
| 21 |
+
_to_openai_tools,
|
| 22 |
+
)
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
# ---------------------------------------------------------------------------
|
| 26 |
+
# Helpers
|
| 27 |
+
# ---------------------------------------------------------------------------
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class _FakeChatCompletions:
|
| 31 |
+
"""Stands in for ``client.chat.completions``. Records every call and
|
| 32 |
+
returns scripted responses one-by-one."""
|
| 33 |
+
|
| 34 |
+
def __init__(self, responses: list[Any]) -> None:
|
| 35 |
+
self.responses = list(responses)
|
| 36 |
+
self.calls: list[dict[str, Any]] = []
|
| 37 |
+
|
| 38 |
+
async def create(self, **kwargs: Any) -> Any:
|
| 39 |
+
self.calls.append(kwargs)
|
| 40 |
+
if not self.responses:
|
| 41 |
+
raise AssertionError(
|
| 42 |
+
"FakeChatCompletions exhausted — backend made more calls than expected"
|
| 43 |
+
)
|
| 44 |
+
return self.responses.pop(0)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
class _FakeClient:
|
| 48 |
+
def __init__(self, responses: list[Any]) -> None:
|
| 49 |
+
self.chat = SimpleNamespace(completions=_FakeChatCompletions(responses))
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def _scripted_response(
|
| 53 |
+
*,
|
| 54 |
+
content: str | None = None,
|
| 55 |
+
tool_calls: list[dict[str, Any]] | None = None,
|
| 56 |
+
finish_reason: str = "stop",
|
| 57 |
+
) -> Any:
|
| 58 |
+
"""Build the SimpleNamespace shape the openai SDK returns from
|
| 59 |
+
chat.completions.create(...).
|
| 60 |
+
"""
|
| 61 |
+
tcs = []
|
| 62 |
+
if tool_calls:
|
| 63 |
+
for tc in tool_calls:
|
| 64 |
+
tcs.append(
|
| 65 |
+
SimpleNamespace(
|
| 66 |
+
id=tc["id"],
|
| 67 |
+
function=SimpleNamespace(
|
| 68 |
+
name=tc["name"],
|
| 69 |
+
arguments=tc.get("arguments", "{}"),
|
| 70 |
+
),
|
| 71 |
+
)
|
| 72 |
+
)
|
| 73 |
+
msg = SimpleNamespace(content=content, tool_calls=tcs or None)
|
| 74 |
+
choice = SimpleNamespace(message=msg, finish_reason=finish_reason)
|
| 75 |
+
return SimpleNamespace(choices=[choice])
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def _backend_with(responses: list[Any]) -> QwenVLLMBackend:
|
| 79 |
+
"""Construct a QwenVLLMBackend with a scripted client."""
|
| 80 |
+
backend = QwenVLLMBackend.__new__(QwenVLLMBackend)
|
| 81 |
+
backend._system = "you are a test agent"
|
| 82 |
+
backend._model = "Qwen/Qwen2.5-7B-Instruct"
|
| 83 |
+
backend._base_url = "http://fake-vllm:8000/v1"
|
| 84 |
+
backend._api_key = "EMPTY"
|
| 85 |
+
backend._max_tokens = 1024
|
| 86 |
+
backend._client = _FakeClient(responses)
|
| 87 |
+
backend._conversation = [{"role": "system", "content": backend._system}]
|
| 88 |
+
return backend
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
# ---------------------------------------------------------------------------
|
| 92 |
+
# Tool-schema translation
|
| 93 |
+
# ---------------------------------------------------------------------------
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
def test_to_openai_tools_translates_neutral_shape() -> None:
|
| 97 |
+
neutral = [
|
| 98 |
+
{
|
| 99 |
+
"name": "parse_config",
|
| 100 |
+
"description": "Parse the file.",
|
| 101 |
+
"input_schema": {"type": "object", "properties": {"file_path": {"type": "string"}}},
|
| 102 |
+
},
|
| 103 |
+
]
|
| 104 |
+
out = _to_openai_tools(neutral)
|
| 105 |
+
assert out == [
|
| 106 |
+
{
|
| 107 |
+
"type": "function",
|
| 108 |
+
"function": {
|
| 109 |
+
"name": "parse_config",
|
| 110 |
+
"description": "Parse the file.",
|
| 111 |
+
"parameters": {
|
| 112 |
+
"type": "object",
|
| 113 |
+
"properties": {"file_path": {"type": "string"}},
|
| 114 |
+
},
|
| 115 |
+
},
|
| 116 |
+
}
|
| 117 |
+
]
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def test_to_openai_tools_handles_missing_input_schema() -> None:
|
| 121 |
+
out = _to_openai_tools([{"name": "x", "description": "y"}])
|
| 122 |
+
assert out[0]["function"]["parameters"] == {"type": "object", "properties": {}}
|
| 123 |
+
|
| 124 |
+
|
| 125 |
+
def test_finish_reason_normalization() -> None:
|
| 126 |
+
assert _normalize_finish_reason("stop") == "end_turn"
|
| 127 |
+
assert _normalize_finish_reason("tool_calls") == "tool_use"
|
| 128 |
+
assert _normalize_finish_reason("length") == "max_tokens"
|
| 129 |
+
assert _normalize_finish_reason(None) == "other"
|
| 130 |
+
assert _normalize_finish_reason("weird") == "weird"
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
# ---------------------------------------------------------------------------
|
| 134 |
+
# next_turn behavior
|
| 135 |
+
# ---------------------------------------------------------------------------
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
@pytest.mark.asyncio
|
| 139 |
+
async def test_next_turn_emits_text_block_and_end_turn() -> None:
|
| 140 |
+
backend = _backend_with([_scripted_response(content="hello there", finish_reason="stop")])
|
| 141 |
+
backend.add_user_message("audit /tmp/x.py")
|
| 142 |
+
turn = await backend.next_turn(tool_schemas=[])
|
| 143 |
+
assert isinstance(turn, AgentTurn)
|
| 144 |
+
assert turn.text_blocks == ["hello there"]
|
| 145 |
+
assert turn.tool_calls == []
|
| 146 |
+
assert turn.stop_reason == "end_turn"
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
@pytest.mark.asyncio
|
| 150 |
+
async def test_next_turn_emits_tool_calls_with_parsed_args() -> None:
|
| 151 |
+
backend = _backend_with(
|
| 152 |
+
[
|
| 153 |
+
_scripted_response(
|
| 154 |
+
content="calling parse_config",
|
| 155 |
+
tool_calls=[
|
| 156 |
+
{
|
| 157 |
+
"id": "tc-1",
|
| 158 |
+
"name": "parse_config",
|
| 159 |
+
"arguments": '{"file_path": "/tmp/x.py"}',
|
| 160 |
+
}
|
| 161 |
+
],
|
| 162 |
+
finish_reason="tool_calls",
|
| 163 |
+
)
|
| 164 |
+
]
|
| 165 |
+
)
|
| 166 |
+
backend.add_user_message("audit /tmp/x.py")
|
| 167 |
+
turn = await backend.next_turn(tool_schemas=[])
|
| 168 |
+
assert turn.stop_reason == "tool_use"
|
| 169 |
+
assert len(turn.tool_calls) == 1
|
| 170 |
+
tc = turn.tool_calls[0]
|
| 171 |
+
assert tc.id == "tc-1"
|
| 172 |
+
assert tc.name == "parse_config"
|
| 173 |
+
assert tc.input == {"file_path": "/tmp/x.py"}
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
@pytest.mark.asyncio
|
| 177 |
+
async def test_next_turn_handles_malformed_tool_arguments() -> None:
|
| 178 |
+
"""vLLM occasionally emits unparseable JSON in arguments — don't crash."""
|
| 179 |
+
backend = _backend_with(
|
| 180 |
+
[
|
| 181 |
+
_scripted_response(
|
| 182 |
+
tool_calls=[
|
| 183 |
+
{"id": "tc-1", "name": "parse_config", "arguments": "{not-json"}
|
| 184 |
+
],
|
| 185 |
+
finish_reason="tool_calls",
|
| 186 |
+
)
|
| 187 |
+
]
|
| 188 |
+
)
|
| 189 |
+
backend.add_user_message("x")
|
| 190 |
+
turn = await backend.next_turn(tool_schemas=[])
|
| 191 |
+
# We get the call but with empty args rather than raising.
|
| 192 |
+
assert turn.tool_calls[0].input == {}
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
@pytest.mark.asyncio
|
| 196 |
+
async def test_tool_result_is_threaded_into_next_request() -> None:
|
| 197 |
+
backend = _backend_with(
|
| 198 |
+
[
|
| 199 |
+
_scripted_response(
|
| 200 |
+
tool_calls=[
|
| 201 |
+
{"id": "tc-1", "name": "parse_config", "arguments": "{}"}
|
| 202 |
+
],
|
| 203 |
+
finish_reason="tool_calls",
|
| 204 |
+
),
|
| 205 |
+
_scripted_response(content="done", finish_reason="stop"),
|
| 206 |
+
]
|
| 207 |
+
)
|
| 208 |
+
backend.add_user_message("audit")
|
| 209 |
+
await backend.next_turn(tool_schemas=[])
|
| 210 |
+
backend.add_tool_result("tc-1", "parse_config", '{"ok": true}', is_error=False)
|
| 211 |
+
await backend.next_turn(tool_schemas=[])
|
| 212 |
+
|
| 213 |
+
# The second create() call should include role="tool" referencing tc-1.
|
| 214 |
+
second_call = backend._client.chat.completions.calls[1]
|
| 215 |
+
msgs = second_call["messages"]
|
| 216 |
+
tool_msgs = [m for m in msgs if m["role"] == "tool"]
|
| 217 |
+
assert len(tool_msgs) == 1
|
| 218 |
+
assert tool_msgs[0]["tool_call_id"] == "tc-1"
|
| 219 |
+
assert tool_msgs[0]["content"] == '{"ok": true}'
|
| 220 |
+
|
| 221 |
+
|
| 222 |
+
@pytest.mark.asyncio
|
| 223 |
+
async def test_is_error_prefix_added_to_failed_tool_results() -> None:
|
| 224 |
+
backend = _backend_with([_scripted_response(content="adapting", finish_reason="stop")])
|
| 225 |
+
backend.add_tool_result("tc-1", "parse_config", "file not found", is_error=True)
|
| 226 |
+
await backend.next_turn(tool_schemas=[])
|
| 227 |
+
msgs = backend._client.chat.completions.calls[0]["messages"]
|
| 228 |
+
tool_msg = next(m for m in msgs if m["role"] == "tool")
|
| 229 |
+
assert tool_msg["content"].startswith("ERROR:")
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# ---------------------------------------------------------------------------
|
| 233 |
+
# Factory selection via env var
|
| 234 |
+
# ---------------------------------------------------------------------------
|
| 235 |
+
|
| 236 |
+
|
| 237 |
+
def test_make_backend_picks_vllm_when_env_var_set(monkeypatch) -> None:
|
| 238 |
+
monkeypatch.setenv("GOBLIN_AGENT_BACKEND", "qwen-vllm")
|
| 239 |
+
monkeypatch.setenv("GOBLIN_QWEN_VLLM_URL", "http://test:8000/v1")
|
| 240 |
+
assert active_backend_name() == "qwen-vllm"
|
| 241 |
+
backend = make_backend(system_prompt="x")
|
| 242 |
+
assert isinstance(backend, QwenVLLMBackend)
|
| 243 |
+
assert backend._base_url == "http://test:8000/v1"
|
| 244 |
+
|
| 245 |
+
|
| 246 |
+
def test_active_backend_name_aliases(monkeypatch) -> None:
|
| 247 |
+
for alias in ("qwen-vllm", "qwen_vllm", "vllm", "local", "QWEN-VLLM", "Vllm"):
|
| 248 |
+
monkeypatch.setenv("GOBLIN_AGENT_BACKEND", alias)
|
| 249 |
+
assert active_backend_name() == "qwen-vllm", alias
|
| 250 |
+
for alias in ("qwen-hf", "qwen", "hf", ""):
|
| 251 |
+
monkeypatch.setenv("GOBLIN_AGENT_BACKEND", alias)
|
| 252 |
+
assert active_backend_name() == "qwen-hf", alias
|
| 253 |
+
|
| 254 |
+
|
| 255 |
+
def test_make_backend_default_is_hf(monkeypatch) -> None:
|
| 256 |
+
monkeypatch.delenv("GOBLIN_AGENT_BACKEND", raising=False)
|
| 257 |
+
monkeypatch.setenv("HF_TOKEN", "fake-test-token")
|
| 258 |
+
assert active_backend_name() == "qwen-hf"
|
| 259 |
+
|
| 260 |
+
|
| 261 |
+
# ---------------------------------------------------------------------------
|
| 262 |
+
# Construction-time error: openai SDK missing
|
| 263 |
+
# ---------------------------------------------------------------------------
|
| 264 |
+
|
| 265 |
+
|
| 266 |
+
def test_constructor_raises_when_openai_not_installed(monkeypatch) -> None:
|
| 267 |
+
"""If openai isn't on PYTHONPATH, the backend's _build_client raises a
|
| 268 |
+
clear RuntimeError instead of an opaque ImportError."""
|
| 269 |
+
import sys
|
| 270 |
+
|
| 271 |
+
saved = sys.modules.pop("openai", None)
|
| 272 |
+
sys.modules["openai"] = None # block re-import
|
| 273 |
+
try:
|
| 274 |
+
with pytest.raises(RuntimeError, match="openai"):
|
| 275 |
+
QwenVLLMBackend(system_prompt="x")
|
| 276 |
+
finally:
|
| 277 |
+
if saved is not None:
|
| 278 |
+
sys.modules["openai"] = saved
|
| 279 |
+
else:
|
| 280 |
+
sys.modules.pop("openai", None)
|
tests/test_runner.py
ADDED
|
@@ -0,0 +1,311 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for runner/protocol.py and runner/profile_parser.py.
|
| 2 |
+
|
| 3 |
+
Two laptop-only invariants:
|
| 4 |
+
1. FakeRunner still works exactly as before (the Phase 1 contract).
|
| 5 |
+
2. LiveRunner gracefully falls back to FakeRunner whenever GPU/profiler
|
| 6 |
+
tools are missing — this dev box has no AMD GPU, so every test here
|
| 7 |
+
should exercise the fallback path.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import csv
|
| 13 |
+
import json
|
| 14 |
+
from pathlib import Path
|
| 15 |
+
from unittest import mock
|
| 16 |
+
|
| 17 |
+
import pytest
|
| 18 |
+
|
| 19 |
+
from agent.schemas import RunMetrics, WorkloadConfig
|
| 20 |
+
from runner import profile_parser
|
| 21 |
+
from runner.protocol import FakeRunner, LiveRunner, _default_runner, gpu_available
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
# ---------------------------------------------------------------------------
|
| 25 |
+
# Helpers
|
| 26 |
+
# ---------------------------------------------------------------------------
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def _baseline_config() -> WorkloadConfig:
|
| 30 |
+
return WorkloadConfig(
|
| 31 |
+
model_name="Qwen/Qwen2.5-7B-Instruct",
|
| 32 |
+
precision="fp16",
|
| 33 |
+
batch_size=4,
|
| 34 |
+
attention_impl="eager",
|
| 35 |
+
dataloader_workers=0,
|
| 36 |
+
)
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
# ---------------------------------------------------------------------------
|
| 40 |
+
# FakeRunner — unchanged contract
|
| 41 |
+
# ---------------------------------------------------------------------------
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
class TestFakeRunner:
|
| 45 |
+
def test_matches_baseline_scenario(self):
|
| 46 |
+
runner = FakeRunner()
|
| 47 |
+
metrics = runner.run(_baseline_config(), steps=10)
|
| 48 |
+
assert isinstance(metrics, RunMetrics)
|
| 49 |
+
assert metrics.runner_kind == "fake"
|
| 50 |
+
assert metrics.steps == 10
|
| 51 |
+
# 01_baseline_bad fixture
|
| 52 |
+
assert metrics.tokens_per_sec == pytest.approx(142.0)
|
| 53 |
+
|
| 54 |
+
def test_steps_override_takes_precedence(self):
|
| 55 |
+
runner = FakeRunner()
|
| 56 |
+
metrics = runner.run(_baseline_config(), steps=99)
|
| 57 |
+
assert metrics.steps == 99
|
| 58 |
+
|
| 59 |
+
def test_default_metrics_when_no_match(self):
|
| 60 |
+
runner = FakeRunner()
|
| 61 |
+
# An unknown model_name forces the no-match path.
|
| 62 |
+
cfg = _baseline_config().model_copy(update={"model_name": "unknown/model"})
|
| 63 |
+
metrics = runner.run(cfg, steps=7)
|
| 64 |
+
assert metrics.runner_kind == "fake"
|
| 65 |
+
assert metrics.steps == 7
|
| 66 |
+
assert any("FakeRunner" in w for w in metrics.warnings)
|
| 67 |
+
|
| 68 |
+
def test_corpus_dir_missing_returns_default(self, tmp_path):
|
| 69 |
+
runner = FakeRunner(corpus_dir=tmp_path / "nope")
|
| 70 |
+
metrics = runner.run(_baseline_config(), steps=10)
|
| 71 |
+
assert metrics.runner_kind == "fake"
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
# ---------------------------------------------------------------------------
|
| 75 |
+
# gpu_available — pure detection
|
| 76 |
+
# ---------------------------------------------------------------------------
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
class TestGpuAvailable:
|
| 80 |
+
def test_no_rocprofv3(self):
|
| 81 |
+
with mock.patch("runner.protocol.shutil.which", return_value=None):
|
| 82 |
+
ok, reason = gpu_available()
|
| 83 |
+
assert ok is False
|
| 84 |
+
assert reason and "rocprofv3" in reason
|
| 85 |
+
|
| 86 |
+
def test_no_amd_smi(self):
|
| 87 |
+
def which(name):
|
| 88 |
+
return "/usr/bin/rocprofv3" if name == "rocprofv3" else None
|
| 89 |
+
|
| 90 |
+
with mock.patch("runner.protocol.shutil.which", side_effect=which):
|
| 91 |
+
ok, reason = gpu_available()
|
| 92 |
+
assert ok is False
|
| 93 |
+
assert reason and "amd-smi" in reason
|
| 94 |
+
|
| 95 |
+
def test_no_render_device(self):
|
| 96 |
+
with mock.patch(
|
| 97 |
+
"runner.protocol.shutil.which",
|
| 98 |
+
side_effect=lambda name: f"/usr/bin/{name}",
|
| 99 |
+
), mock.patch("runner.protocol._has_render_device", return_value=False):
|
| 100 |
+
ok, reason = gpu_available()
|
| 101 |
+
assert ok is False
|
| 102 |
+
assert reason and "renderD" in reason
|
| 103 |
+
|
| 104 |
+
def test_all_present(self):
|
| 105 |
+
with mock.patch(
|
| 106 |
+
"runner.protocol.shutil.which",
|
| 107 |
+
side_effect=lambda name: f"/usr/bin/{name}",
|
| 108 |
+
), mock.patch("runner.protocol._has_render_device", return_value=True):
|
| 109 |
+
ok, reason = gpu_available()
|
| 110 |
+
assert ok is True
|
| 111 |
+
assert reason is None
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
# ---------------------------------------------------------------------------
|
| 115 |
+
# LiveRunner — must fall back on this no-GPU dev machine
|
| 116 |
+
# ---------------------------------------------------------------------------
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
class TestLiveRunnerFallback:
|
| 120 |
+
def test_falls_back_when_gpu_unavailable(self):
|
| 121 |
+
runner = LiveRunner()
|
| 122 |
+
metrics = runner.run(_baseline_config(), steps=10)
|
| 123 |
+
# On a laptop, gpu_available() returns False → FakeRunner path.
|
| 124 |
+
assert metrics.runner_kind == "fake"
|
| 125 |
+
# The warning must be the FIRST entry (LiveRunner prepends it).
|
| 126 |
+
assert metrics.warnings, "LiveRunner must surface a fallback warning"
|
| 127 |
+
assert "LiveRunner" in metrics.warnings[0]
|
| 128 |
+
|
| 129 |
+
def test_falls_back_when_runner_script_missing(self, tmp_path):
|
| 130 |
+
with mock.patch("runner.protocol.gpu_available", return_value=(True, None)):
|
| 131 |
+
runner = LiveRunner(runner_script=tmp_path / "does_not_exist.sh")
|
| 132 |
+
metrics = runner.run(_baseline_config(), steps=10)
|
| 133 |
+
assert metrics.runner_kind == "fake"
|
| 134 |
+
assert any("runner script not found" in w for w in metrics.warnings)
|
| 135 |
+
|
| 136 |
+
def test_falls_back_when_runner_script_not_executable(self, tmp_path):
|
| 137 |
+
script = tmp_path / "goblin_runner.sh"
|
| 138 |
+
script.write_text("#!/bin/sh\nexit 0\n")
|
| 139 |
+
# Deliberately don't chmod +x
|
| 140 |
+
with mock.patch("runner.protocol.gpu_available", return_value=(True, None)):
|
| 141 |
+
runner = LiveRunner(runner_script=script)
|
| 142 |
+
metrics = runner.run(_baseline_config(), steps=10)
|
| 143 |
+
assert metrics.runner_kind == "fake"
|
| 144 |
+
assert any("not executable" in w for w in metrics.warnings)
|
| 145 |
+
|
| 146 |
+
def test_falls_back_when_subprocess_returns_nonzero(self, tmp_path):
|
| 147 |
+
script = tmp_path / "goblin_runner.sh"
|
| 148 |
+
script.write_text("#!/usr/bin/env bash\nexit 7\n")
|
| 149 |
+
script.chmod(0o755)
|
| 150 |
+
with mock.patch("runner.protocol.gpu_available", return_value=(True, None)):
|
| 151 |
+
runner = LiveRunner(runner_script=script)
|
| 152 |
+
metrics = runner.run(_baseline_config(), steps=10)
|
| 153 |
+
assert metrics.runner_kind == "fake"
|
| 154 |
+
assert any("exited with code 7" in w for w in metrics.warnings)
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
# ---------------------------------------------------------------------------
|
| 158 |
+
# _default_runner — module-level factory
|
| 159 |
+
# ---------------------------------------------------------------------------
|
| 160 |
+
|
| 161 |
+
|
| 162 |
+
def test_default_runner_is_live_runner():
|
| 163 |
+
runner = _default_runner()
|
| 164 |
+
assert isinstance(runner, LiveRunner)
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
# ---------------------------------------------------------------------------
|
| 168 |
+
# profile_parser — graceful degradation when artefacts missing
|
| 169 |
+
# ---------------------------------------------------------------------------
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
class TestProfileParser:
|
| 173 |
+
def test_empty_dir_returns_zero_metrics_with_warnings(self, tmp_path):
|
| 174 |
+
metrics = profile_parser.parse(tmp_path, config=_baseline_config(), steps=10)
|
| 175 |
+
assert metrics.tokens_per_sec == 0.0
|
| 176 |
+
assert metrics.mfu_pct == 0.0
|
| 177 |
+
assert metrics.gpu_util_pct == 0.0
|
| 178 |
+
assert len(metrics.warnings) >= 3 # one warning per missing artefact
|
| 179 |
+
|
| 180 |
+
def test_parses_synthetic_artefacts(self, tmp_path):
|
| 181 |
+
# Minimal rocprofv3-shaped CSV
|
| 182 |
+
trace = tmp_path / "trace.csv"
|
| 183 |
+
with trace.open("w", newline="") as f:
|
| 184 |
+
w = csv.writer(f)
|
| 185 |
+
w.writerow(["KernelName", "DurationNs"])
|
| 186 |
+
w.writerow(["aten::matmul (fp16)", 5_000_000])
|
| 187 |
+
w.writerow(["aten::scaled_dot_product_attention", 3_000_000])
|
| 188 |
+
w.writerow(["rccl_AllReduce", 1_000_000])
|
| 189 |
+
w.writerow(["hipBLASLt_generic_gemm", 2_000_000])
|
| 190 |
+
|
| 191 |
+
# Minimal torch.profiler chrome trace with embedded metadata
|
| 192 |
+
torch_profile = {
|
| 193 |
+
"metadata": {
|
| 194 |
+
"tokens_per_sec": 142.0,
|
| 195 |
+
"mfu_pct": 24.0,
|
| 196 |
+
"pytorch_version": "2.3.0+rocm6.1",
|
| 197 |
+
"step_time_seconds": 0.5,
|
| 198 |
+
"host_busy_fraction": 0.6,
|
| 199 |
+
},
|
| 200 |
+
"traceEvents": [],
|
| 201 |
+
}
|
| 202 |
+
(tmp_path / "torch_profile.json").write_text(json.dumps(torch_profile))
|
| 203 |
+
|
| 204 |
+
# Minimal amd-smi telemetry
|
| 205 |
+
smi = tmp_path / "amd_smi.csv"
|
| 206 |
+
with smi.open("w", newline="") as f:
|
| 207 |
+
w = csv.writer(f)
|
| 208 |
+
w.writerow(["VRAM_USED_GB", "GFX_ACTIVITY"])
|
| 209 |
+
w.writerow(["72.0", "20.0"]) # < 30% util → triggers data_wait
|
| 210 |
+
w.writerow(["75.0", "22.0"])
|
| 211 |
+
|
| 212 |
+
metrics = profile_parser.parse(tmp_path, config=_baseline_config(), steps=10)
|
| 213 |
+
assert metrics.tokens_per_sec == pytest.approx(142.0)
|
| 214 |
+
assert metrics.mfu_pct == pytest.approx(24.0)
|
| 215 |
+
assert metrics.hbm_peak_gb == pytest.approx(75.0)
|
| 216 |
+
assert metrics.hbm_avg_gb == pytest.approx(73.5)
|
| 217 |
+
# comm_excess detected (rccl kernel, 1 ms)
|
| 218 |
+
assert metrics.waste_budget.comm_excess == pytest.approx(0.001)
|
| 219 |
+
# data_wait triggered (gpu util < 30, host_busy > 0.5)
|
| 220 |
+
assert metrics.waste_budget.data_wait > 0.0
|
| 221 |
+
# precision_path triggered (config.precision='fp16' AND fp16 kernels present)
|
| 222 |
+
assert metrics.waste_budget.precision_path > 0.0
|
| 223 |
+
# kernel_shape: generic GEMM detected
|
| 224 |
+
assert metrics.waste_budget.kernel_shape > 0.0
|
| 225 |
+
# memory_headroom: 75 GB used << 70% × 192 GB = 134.4 GB → slack
|
| 226 |
+
assert metrics.waste_budget.memory_headroom > 0.0
|
| 227 |
+
|
| 228 |
+
def test_bf16_config_skips_precision_path(self, tmp_path):
|
| 229 |
+
# Even with fp16-tagged kernels, a bf16 config means precision_path = 0
|
| 230 |
+
# because the user is already on the optimal precision.
|
| 231 |
+
trace = tmp_path / "trace.csv"
|
| 232 |
+
with trace.open("w", newline="") as f:
|
| 233 |
+
w = csv.writer(f)
|
| 234 |
+
w.writerow(["KernelName", "DurationNs"])
|
| 235 |
+
w.writerow(["matmul_fp16_kernel", 5_000_000])
|
| 236 |
+
torch_profile = {
|
| 237 |
+
"metadata": {
|
| 238 |
+
"tokens_per_sec": 318.0,
|
| 239 |
+
"mfu_pct": 51.0,
|
| 240 |
+
"step_time_seconds": 0.3,
|
| 241 |
+
"host_busy_fraction": 0.2,
|
| 242 |
+
},
|
| 243 |
+
"traceEvents": [],
|
| 244 |
+
}
|
| 245 |
+
(tmp_path / "torch_profile.json").write_text(json.dumps(torch_profile))
|
| 246 |
+
smi = tmp_path / "amd_smi.csv"
|
| 247 |
+
smi.write_text("VRAM_USED_GB,GFX_ACTIVITY\n168.0,86.0\n")
|
| 248 |
+
|
| 249 |
+
bf16_config = _baseline_config().model_copy(update={"precision": "bf16"})
|
| 250 |
+
metrics = profile_parser.parse(tmp_path, config=bf16_config, steps=50)
|
| 251 |
+
assert metrics.waste_budget.precision_path == 0.0
|
| 252 |
+
|
| 253 |
+
|
| 254 |
+
# ---------------------------------------------------------------------------
|
| 255 |
+
# Caching — exercise the benchmark tool's cache layer
|
| 256 |
+
# ---------------------------------------------------------------------------
|
| 257 |
+
|
| 258 |
+
|
| 259 |
+
class TestBenchmarkCache:
|
| 260 |
+
"""The benchmark tool writes to the real bench_cache/ directory; isolate it."""
|
| 261 |
+
|
| 262 |
+
@pytest.fixture(autouse=True)
|
| 263 |
+
def _isolate_cache(self, tmp_path, monkeypatch):
|
| 264 |
+
monkeypatch.setattr("agent.tools.benchmark._CACHE_DIR", tmp_path / "bench_cache")
|
| 265 |
+
# Force ROCM_IMAGE_TAG to a known value so the key is reproducible.
|
| 266 |
+
monkeypatch.setenv("ROCM_IMAGE_TAG", "test-tag")
|
| 267 |
+
yield
|
| 268 |
+
|
| 269 |
+
def test_cache_hit_on_second_call(self):
|
| 270 |
+
from agent.tools.benchmark import _benchmark
|
| 271 |
+
|
| 272 |
+
cfg = _baseline_config().model_dump()
|
| 273 |
+
r1 = _benchmark(cfg, steps=50)
|
| 274 |
+
assert r1.ok
|
| 275 |
+
# Second call should HIT the cache and warn about it.
|
| 276 |
+
r2 = _benchmark(cfg, steps=50)
|
| 277 |
+
assert r2.ok
|
| 278 |
+
assert any("cache hit" in w for w in r2.result["warnings"])
|
| 279 |
+
|
| 280 |
+
def test_force_rerun_bypasses_cache(self):
|
| 281 |
+
from agent.tools.benchmark import _benchmark
|
| 282 |
+
|
| 283 |
+
cfg = _baseline_config().model_dump()
|
| 284 |
+
_benchmark(cfg, steps=50)
|
| 285 |
+
r2 = _benchmark(cfg, steps=50, force_rerun=True)
|
| 286 |
+
assert r2.ok
|
| 287 |
+
assert not any("cache hit" in w for w in r2.result["warnings"])
|
| 288 |
+
|
| 289 |
+
def test_different_steps_invalidate_cache(self):
|
| 290 |
+
from agent.tools.benchmark import _benchmark
|
| 291 |
+
|
| 292 |
+
cfg = _baseline_config().model_dump()
|
| 293 |
+
_benchmark(cfg, steps=50)
|
| 294 |
+
r2 = _benchmark(cfg, steps=100)
|
| 295 |
+
# Same config, different steps → different cache key → cold call.
|
| 296 |
+
assert not any("cache hit" in w for w in r2.result["warnings"])
|
| 297 |
+
|
| 298 |
+
def test_runner_script_change_invalidates_cache(self, tmp_path, monkeypatch):
|
| 299 |
+
from agent.tools.benchmark import _benchmark
|
| 300 |
+
|
| 301 |
+
cfg = _baseline_config().model_dump()
|
| 302 |
+
_benchmark(cfg, steps=50)
|
| 303 |
+
|
| 304 |
+
# Pretend the runner script changed by swapping the path the cache
|
| 305 |
+
# key reads from. (Simulates "container/runner version bump".)
|
| 306 |
+
fake_script = tmp_path / "different_runner.sh"
|
| 307 |
+
fake_script.write_text("# different content\n")
|
| 308 |
+
monkeypatch.setattr("agent.tools.benchmark._RUNNER_SCRIPT", fake_script)
|
| 309 |
+
|
| 310 |
+
r2 = _benchmark(cfg, steps=50)
|
| 311 |
+
assert not any("cache hit" in w for w in r2.result["warnings"])
|
ui/auto_tune_ui.py
ADDED
|
@@ -0,0 +1,799 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""GPU Goblin auto-tune UI — Streamlit frontend for scripts/auto_tune.py.
|
| 2 |
+
|
| 3 |
+
Run:
|
| 4 |
+
streamlit run ui/auto_tune_ui.py
|
| 5 |
+
|
| 6 |
+
The UI:
|
| 7 |
+
1. Form: pick model OR workload path, mode, steps, etc.
|
| 8 |
+
2. Run: launches scripts/auto_tune.py as a subprocess with --events FILE
|
| 9 |
+
3. Live progress: tails the events file, renders iteration cards as they
|
| 10 |
+
arrive, updates the best-tokens/sec metric on every accepted change
|
| 11 |
+
4. Final report: improvement vs baseline, accepted vs rejected
|
| 12 |
+
experiments, waste-budget reduction chart, and a copyable diff.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import json
|
| 18 |
+
import os
|
| 19 |
+
import subprocess
|
| 20 |
+
import sys
|
| 21 |
+
import tempfile
|
| 22 |
+
import time
|
| 23 |
+
from pathlib import Path
|
| 24 |
+
from typing import Any
|
| 25 |
+
|
| 26 |
+
# Streamlit only puts ui/ on sys.path, so add the repo root for any helper
|
| 27 |
+
# imports that might land in shared modules later.
|
| 28 |
+
_REPO_ROOT = Path(__file__).resolve().parent.parent
|
| 29 |
+
if str(_REPO_ROOT) not in sys.path:
|
| 30 |
+
sys.path.insert(0, str(_REPO_ROOT))
|
| 31 |
+
|
| 32 |
+
import altair as alt
|
| 33 |
+
import pandas as pd
|
| 34 |
+
import requests
|
| 35 |
+
import streamlit as st
|
| 36 |
+
|
| 37 |
+
REPO_ROOT = _REPO_ROOT
|
| 38 |
+
AUTO_TUNE_SCRIPT = REPO_ROOT / "scripts" / "auto_tune.py"
|
| 39 |
+
|
| 40 |
+
WASTE_BUCKETS = [
|
| 41 |
+
"useful_gpu",
|
| 42 |
+
"data_wait",
|
| 43 |
+
"host_gap",
|
| 44 |
+
"comm_excess",
|
| 45 |
+
"memory_headroom",
|
| 46 |
+
"precision_path",
|
| 47 |
+
"kernel_shape",
|
| 48 |
+
]
|
| 49 |
+
|
| 50 |
+
# ---------------------------------------------------------------------------
|
| 51 |
+
# Page setup
|
| 52 |
+
# ---------------------------------------------------------------------------
|
| 53 |
+
|
| 54 |
+
st.set_page_config(
|
| 55 |
+
page_title="GPU Goblin — Auto-Tune",
|
| 56 |
+
page_icon="🔧",
|
| 57 |
+
layout="wide",
|
| 58 |
+
)
|
| 59 |
+
|
| 60 |
+
st.title("🔧 GPU Goblin — Auto-Tune")
|
| 61 |
+
st.caption(
|
| 62 |
+
"Iteratively find the fastest configuration for a fine-tuning workload "
|
| 63 |
+
"on AMD MI300X. Pick a model, hit Run, watch tokens/sec climb."
|
| 64 |
+
)
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
# ---------------------------------------------------------------------------
|
| 68 |
+
# Form: inputs
|
| 69 |
+
# ---------------------------------------------------------------------------
|
| 70 |
+
|
| 71 |
+
with st.sidebar:
|
| 72 |
+
st.header("Run configuration")
|
| 73 |
+
|
| 74 |
+
# ---- Backend mode (Local subprocess vs remote GPU server) ----
|
| 75 |
+
default_backend = os.environ.get("GOBLIN_AUTO_TUNE_URL", "")
|
| 76 |
+
backend_mode = st.radio(
|
| 77 |
+
"Backend",
|
| 78 |
+
options=("Local subprocess", "Remote GPU server"),
|
| 79 |
+
index=1 if default_backend else 0,
|
| 80 |
+
help=(
|
| 81 |
+
"Local: launches scripts/auto_tune.py here (this host needs an MI300X). "
|
| 82 |
+
"Remote: POSTs to /auto-tune on a FastAPI server you've stood up on "
|
| 83 |
+
"the GPU host. Use Remote when running the UI on HF Spaces."
|
| 84 |
+
),
|
| 85 |
+
)
|
| 86 |
+
backend_url = ""
|
| 87 |
+
if backend_mode == "Remote GPU server":
|
| 88 |
+
backend_url = st.text_input(
|
| 89 |
+
"Backend URL",
|
| 90 |
+
value=default_backend or "http://localhost:8000",
|
| 91 |
+
help="Base URL of the FastAPI server (no trailing slash). "
|
| 92 |
+
"/auto-tune is appended automatically.",
|
| 93 |
+
)
|
| 94 |
+
|
| 95 |
+
st.divider()
|
| 96 |
+
workload_source = st.radio(
|
| 97 |
+
"Workload source",
|
| 98 |
+
options=("Model id", "Custom workload script"),
|
| 99 |
+
index=0,
|
| 100 |
+
help=(
|
| 101 |
+
"Model id: auto-generates a baseline workload from the demo "
|
| 102 |
+
"template (Qwen-style LoRA fine-tune). Custom: point at any "
|
| 103 |
+
"Python training script."
|
| 104 |
+
),
|
| 105 |
+
)
|
| 106 |
+
|
| 107 |
+
model_id = ""
|
| 108 |
+
workload_path = ""
|
| 109 |
+
if workload_source == "Model id":
|
| 110 |
+
model_id = st.text_input(
|
| 111 |
+
"Model id",
|
| 112 |
+
value="Qwen/Qwen2.5-7B-Instruct",
|
| 113 |
+
help=(
|
| 114 |
+
"Any HuggingFace causal-LM model id. For gated models "
|
| 115 |
+
"(Llama, etc.) ensure HF_TOKEN is set in the environment."
|
| 116 |
+
),
|
| 117 |
+
)
|
| 118 |
+
else:
|
| 119 |
+
workload_path = st.text_input(
|
| 120 |
+
"Path to workload script (relative to repo root)",
|
| 121 |
+
value="workloads/train_qwen_lora.py",
|
| 122 |
+
)
|
| 123 |
+
|
| 124 |
+
st.divider()
|
| 125 |
+
st.subheader("Tuning strategy")
|
| 126 |
+
|
| 127 |
+
mode = st.selectbox(
|
| 128 |
+
"Mode",
|
| 129 |
+
options=("hardcoded", "llm", "llm-explore"),
|
| 130 |
+
index=0,
|
| 131 |
+
help=(
|
| 132 |
+
"hardcoded: priority-ordered playbook (no API key). "
|
| 133 |
+
"llm: LLM picks one experiment per iteration (greedy). "
|
| 134 |
+
"llm-explore: LLM proposes K candidates per iteration; best wins."
|
| 135 |
+
),
|
| 136 |
+
)
|
| 137 |
+
candidates_per_iteration = 3
|
| 138 |
+
if mode == "llm-explore":
|
| 139 |
+
candidates_per_iteration = st.slider(
|
| 140 |
+
"Candidates per iteration (K)",
|
| 141 |
+
min_value=2,
|
| 142 |
+
max_value=6,
|
| 143 |
+
value=3,
|
| 144 |
+
help="Each iteration runs K benchmarks. Higher = broader search, more GPU time.",
|
| 145 |
+
)
|
| 146 |
+
|
| 147 |
+
steps = st.slider(
|
| 148 |
+
"Steps per benchmark",
|
| 149 |
+
min_value=10,
|
| 150 |
+
max_value=100,
|
| 151 |
+
value=20,
|
| 152 |
+
step=5,
|
| 153 |
+
help="More steps = lower variance but longer benchmarks.",
|
| 154 |
+
)
|
| 155 |
+
|
| 156 |
+
max_iterations = st.slider(
|
| 157 |
+
"Max iterations",
|
| 158 |
+
min_value=1,
|
| 159 |
+
max_value=20,
|
| 160 |
+
value=10,
|
| 161 |
+
help="Cap on tuning iterations. Default 10 for llm modes, 5 for llm-explore.",
|
| 162 |
+
)
|
| 163 |
+
|
| 164 |
+
early_stop_after = st.slider(
|
| 165 |
+
"Early stop after N non-improvements",
|
| 166 |
+
min_value=1,
|
| 167 |
+
max_value=10,
|
| 168 |
+
value=3,
|
| 169 |
+
)
|
| 170 |
+
|
| 171 |
+
max_crashes = st.slider(
|
| 172 |
+
"Max total crashes",
|
| 173 |
+
min_value=1,
|
| 174 |
+
max_value=10,
|
| 175 |
+
value=4,
|
| 176 |
+
)
|
| 177 |
+
|
| 178 |
+
improvement_threshold = st.number_input(
|
| 179 |
+
"Improvement threshold (%)",
|
| 180 |
+
min_value=0.0,
|
| 181 |
+
max_value=10.0,
|
| 182 |
+
value=0.0,
|
| 183 |
+
step=0.1,
|
| 184 |
+
help="Min % gain to accept. 0.0 = any positive delta wins.",
|
| 185 |
+
)
|
| 186 |
+
|
| 187 |
+
st.divider()
|
| 188 |
+
run_pressed = st.button("🚀 Run auto-tune", type="primary", use_container_width=True)
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
# ---------------------------------------------------------------------------
|
| 192 |
+
# Render helpers
|
| 193 |
+
# ---------------------------------------------------------------------------
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
def _format_tps(v: float) -> str:
|
| 197 |
+
return f"{v:,.0f}" if v else "—"
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
def _format_pct(v: float | None, suffix: str = "%") -> str:
|
| 201 |
+
if v is None:
|
| 202 |
+
return "—"
|
| 203 |
+
return f"{v:+.2f}{suffix}"
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
def _waste_chart(baseline: dict, current: dict | None) -> alt.Chart:
|
| 207 |
+
"""Two-row stacked bar: baseline vs current waste_budget breakdown."""
|
| 208 |
+
rows = []
|
| 209 |
+
sources = [("baseline", baseline)]
|
| 210 |
+
if current is not None and current is not baseline:
|
| 211 |
+
sources.append(("current best", current))
|
| 212 |
+
for label, m in sources:
|
| 213 |
+
wb = (m or {}).get("waste_budget") or {}
|
| 214 |
+
for bucket in WASTE_BUCKETS:
|
| 215 |
+
v = float(wb.get(bucket, 0.0))
|
| 216 |
+
if v <= 0:
|
| 217 |
+
continue
|
| 218 |
+
rows.append({"run": label, "bucket": bucket, "seconds": v})
|
| 219 |
+
if not rows:
|
| 220 |
+
return None
|
| 221 |
+
df = pd.DataFrame(rows)
|
| 222 |
+
return (
|
| 223 |
+
alt.Chart(df)
|
| 224 |
+
.mark_bar()
|
| 225 |
+
.encode(
|
| 226 |
+
x=alt.X("seconds:Q", title="seconds / step"),
|
| 227 |
+
y=alt.Y("run:N", title=None, sort=["baseline", "current best"]),
|
| 228 |
+
color=alt.Color(
|
| 229 |
+
"bucket:N",
|
| 230 |
+
scale=alt.Scale(scheme="tableau10"),
|
| 231 |
+
sort=WASTE_BUCKETS,
|
| 232 |
+
),
|
| 233 |
+
order=alt.Order("bucket:N", sort="ascending"),
|
| 234 |
+
tooltip=["run", "bucket", alt.Tooltip("seconds:Q", format=".3f")],
|
| 235 |
+
)
|
| 236 |
+
.properties(height=120)
|
| 237 |
+
)
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
def _render_metrics_row(baseline: dict | None, current: dict | None) -> None:
|
| 241 |
+
"""Top-of-page metrics: tokens/sec, mfu_pct, hbm, with deltas vs baseline."""
|
| 242 |
+
cols = st.columns(4)
|
| 243 |
+
if baseline is None:
|
| 244 |
+
for c, label in zip(cols, ["tokens/sec", "mfu_pct", "hbm_peak (GB)", "iterations"]):
|
| 245 |
+
c.metric(label, "—")
|
| 246 |
+
return
|
| 247 |
+
|
| 248 |
+
cur = current if current is not None else baseline
|
| 249 |
+
base_tps = float(baseline.get("tokens_per_sec") or 0)
|
| 250 |
+
cur_tps = float(cur.get("tokens_per_sec") or 0)
|
| 251 |
+
base_mfu = float(baseline.get("mfu_pct") or 0)
|
| 252 |
+
cur_mfu = float(cur.get("mfu_pct") or 0)
|
| 253 |
+
base_hbm = float(baseline.get("hbm_peak_gb") or 0)
|
| 254 |
+
cur_hbm = float(cur.get("hbm_peak_gb") or 0)
|
| 255 |
+
|
| 256 |
+
cols[0].metric(
|
| 257 |
+
"tokens/sec",
|
| 258 |
+
_format_tps(cur_tps),
|
| 259 |
+
delta=f"{cur_tps - base_tps:+,.0f}" if base_tps else None,
|
| 260 |
+
)
|
| 261 |
+
cols[1].metric(
|
| 262 |
+
"MFU %",
|
| 263 |
+
f"{cur_mfu:.2f}",
|
| 264 |
+
delta=f"{cur_mfu - base_mfu:+.2f} pts" if base_mfu else None,
|
| 265 |
+
)
|
| 266 |
+
cols[2].metric(
|
| 267 |
+
"HBM peak (GB)",
|
| 268 |
+
f"{cur_hbm:.1f}",
|
| 269 |
+
delta=f"{cur_hbm - base_hbm:+.1f}" if base_hbm else None,
|
| 270 |
+
)
|
| 271 |
+
|
| 272 |
+
|
| 273 |
+
def _render_candidate_card(event: dict) -> None:
|
| 274 |
+
"""One candidate's outcome inside an iteration."""
|
| 275 |
+
name = event.get("name", "?")
|
| 276 |
+
rationale = event.get("rationale", "")
|
| 277 |
+
outcome = event.get("outcome", "?")
|
| 278 |
+
metrics = event.get("metrics") or {}
|
| 279 |
+
delta = event.get("delta_vs_best")
|
| 280 |
+
reason = event.get("reason", "")
|
| 281 |
+
|
| 282 |
+
icon = {
|
| 283 |
+
"evaluated": "📊",
|
| 284 |
+
"skipped": "⏭️",
|
| 285 |
+
"rejected": "❌",
|
| 286 |
+
"crashed": "💥",
|
| 287 |
+
}.get(outcome, "•")
|
| 288 |
+
|
| 289 |
+
title = f"{icon} **{name}**"
|
| 290 |
+
if outcome == "evaluated" and delta is not None:
|
| 291 |
+
title += f" — Δ {delta:+.2f}%"
|
| 292 |
+
elif outcome != "evaluated":
|
| 293 |
+
title += f" — *{outcome}*"
|
| 294 |
+
|
| 295 |
+
with st.container(border=True):
|
| 296 |
+
st.markdown(title)
|
| 297 |
+
if rationale:
|
| 298 |
+
st.caption(rationale)
|
| 299 |
+
if metrics:
|
| 300 |
+
sub = st.columns(4)
|
| 301 |
+
sub[0].metric("tokens/sec", f"{metrics.get('tokens_per_sec', 0):,.0f}")
|
| 302 |
+
sub[1].metric("MFU %", f"{metrics.get('mfu_pct', 0):.2f}")
|
| 303 |
+
sub[2].metric("HBM peak GB", f"{metrics.get('hbm_peak_gb', 0):.1f}")
|
| 304 |
+
sub[3].metric("Δ vs best", f"{delta:+.2f}%" if delta is not None else "—")
|
| 305 |
+
wb = metrics.get("waste_budget") or {}
|
| 306 |
+
non_zero = {k: v for k, v in wb.items() if k != "useful_gpu" and v > 0}
|
| 307 |
+
if non_zero:
|
| 308 |
+
wb_text = ", ".join(
|
| 309 |
+
f"`{k}={v:.3f}`"
|
| 310 |
+
for k, v in sorted(non_zero.items(), key=lambda kv: kv[1], reverse=True)
|
| 311 |
+
)
|
| 312 |
+
st.caption(f"Waste: useful_gpu=`{wb.get('useful_gpu', 0):.3f}`, {wb_text}")
|
| 313 |
+
if reason and outcome != "evaluated":
|
| 314 |
+
st.caption(f"Reason: {reason}")
|
| 315 |
+
|
| 316 |
+
|
| 317 |
+
# ---------------------------------------------------------------------------
|
| 318 |
+
# Run
|
| 319 |
+
# ---------------------------------------------------------------------------
|
| 320 |
+
|
| 321 |
+
|
| 322 |
+
def _build_command(events_file: Path) -> list[str]:
|
| 323 |
+
cmd: list[str] = [sys.executable, "-u", str(AUTO_TUNE_SCRIPT)]
|
| 324 |
+
if workload_source == "Model id":
|
| 325 |
+
if not model_id.strip():
|
| 326 |
+
raise ValueError("Model id is required.")
|
| 327 |
+
cmd.extend(["--model", model_id.strip()])
|
| 328 |
+
else:
|
| 329 |
+
wp = workload_path.strip()
|
| 330 |
+
if not wp:
|
| 331 |
+
raise ValueError("Workload path is required.")
|
| 332 |
+
full = (REPO_ROOT / wp).resolve() if not Path(wp).is_absolute() else Path(wp)
|
| 333 |
+
if not full.exists():
|
| 334 |
+
raise ValueError(f"Workload not found: {full}")
|
| 335 |
+
cmd.append(str(full))
|
| 336 |
+
cmd.extend([
|
| 337 |
+
"--mode", mode,
|
| 338 |
+
"--steps", str(steps),
|
| 339 |
+
"--max-iterations", str(max_iterations),
|
| 340 |
+
"--early-stop-after", str(early_stop_after),
|
| 341 |
+
"--max-crashes", str(max_crashes),
|
| 342 |
+
"--improvement-threshold", str(improvement_threshold),
|
| 343 |
+
"--events", str(events_file),
|
| 344 |
+
])
|
| 345 |
+
if mode == "llm-explore":
|
| 346 |
+
cmd.extend(["--candidates-per-iteration", str(candidates_per_iteration)])
|
| 347 |
+
return cmd
|
| 348 |
+
|
| 349 |
+
|
| 350 |
+
def _build_request_body() -> dict:
|
| 351 |
+
"""Same config as _build_command, but as a JSON body for POST /auto-tune."""
|
| 352 |
+
body: dict[str, Any] = {
|
| 353 |
+
"mode": mode,
|
| 354 |
+
"steps": steps,
|
| 355 |
+
"max_iterations": max_iterations,
|
| 356 |
+
"early_stop_after": early_stop_after,
|
| 357 |
+
"max_crashes": max_crashes,
|
| 358 |
+
"improvement_threshold": float(improvement_threshold),
|
| 359 |
+
}
|
| 360 |
+
if mode == "llm-explore":
|
| 361 |
+
body["candidates_per_iteration"] = candidates_per_iteration
|
| 362 |
+
if workload_source == "Model id":
|
| 363 |
+
if not model_id.strip():
|
| 364 |
+
raise ValueError("Model id is required.")
|
| 365 |
+
body["model"] = model_id.strip()
|
| 366 |
+
else:
|
| 367 |
+
if not workload_path.strip():
|
| 368 |
+
raise ValueError("Workload path is required.")
|
| 369 |
+
body["workload"] = workload_path.strip()
|
| 370 |
+
return body
|
| 371 |
+
|
| 372 |
+
|
| 373 |
+
def _read_events(path: Path, seen: int) -> tuple[list[dict], int]:
|
| 374 |
+
"""Read events from byte position `seen` onward; return (new_events, new_seen)."""
|
| 375 |
+
if not path.exists():
|
| 376 |
+
return [], seen
|
| 377 |
+
try:
|
| 378 |
+
with path.open("r") as f:
|
| 379 |
+
f.seek(seen)
|
| 380 |
+
chunk = f.read()
|
| 381 |
+
new_seen = f.tell()
|
| 382 |
+
except OSError:
|
| 383 |
+
return [], seen
|
| 384 |
+
if not chunk:
|
| 385 |
+
return [], new_seen
|
| 386 |
+
out: list[dict] = []
|
| 387 |
+
# The last line may be partial if the writer is mid-flush; drop it
|
| 388 |
+
# and back the pointer up so we re-read it next tick.
|
| 389 |
+
pieces = chunk.splitlines(keepends=True)
|
| 390 |
+
if pieces and not pieces[-1].endswith("\n"):
|
| 391 |
+
partial = pieces.pop()
|
| 392 |
+
new_seen -= len(partial.encode("utf-8"))
|
| 393 |
+
for line in pieces:
|
| 394 |
+
line = line.strip()
|
| 395 |
+
if not line:
|
| 396 |
+
continue
|
| 397 |
+
try:
|
| 398 |
+
out.append(json.loads(line))
|
| 399 |
+
except json.JSONDecodeError:
|
| 400 |
+
continue
|
| 401 |
+
return out, new_seen
|
| 402 |
+
|
| 403 |
+
|
| 404 |
+
# ---------------------------------------------------------------------------
|
| 405 |
+
# Initial state — empty page
|
| 406 |
+
# ---------------------------------------------------------------------------
|
| 407 |
+
|
| 408 |
+
if not run_pressed:
|
| 409 |
+
st.info(
|
| 410 |
+
"👈 Configure inputs in the sidebar and press **Run auto-tune**. "
|
| 411 |
+
"The script will profile your workload, try MI300X-specific tuning "
|
| 412 |
+
"changes, and report what worked."
|
| 413 |
+
)
|
| 414 |
+
_render_metrics_row(None, None)
|
| 415 |
+
st.subheader("How it works")
|
| 416 |
+
st.markdown(
|
| 417 |
+
"""
|
| 418 |
+
1. **Baseline benchmark** — runs your workload as-is (or as generated
|
| 419 |
+
from `--model`) on the GPU and measures tokens/sec, MFU, HBM peak,
|
| 420 |
+
and a per-bucket waste budget.
|
| 421 |
+
2. **Iterative tuning** — applies one MI300X-specific change at a
|
| 422 |
+
time (e.g. `bf16`, `batch_size=16`, `TORCH_BLAS_PREFER_HIPBLASLT=1`)
|
| 423 |
+
and re-benchmarks. Keeps changes that beat the current best.
|
| 424 |
+
3. **Live progress** — every accepted/rejected/crashed candidate
|
| 425 |
+
shows up here as it happens.
|
| 426 |
+
4. **Final report** — overall improvement, which experiments won,
|
| 427 |
+
how much wastage was recovered, and a diff against the baseline
|
| 428 |
+
workload.
|
| 429 |
+
"""
|
| 430 |
+
)
|
| 431 |
+
st.stop()
|
| 432 |
+
|
| 433 |
+
|
| 434 |
+
# ---------------------------------------------------------------------------
|
| 435 |
+
# Run path: launch subprocess + tail events
|
| 436 |
+
# ---------------------------------------------------------------------------
|
| 437 |
+
|
| 438 |
+
st.subheader("Live run")
|
| 439 |
+
header = st.empty()
|
| 440 |
+
|
| 441 |
+
# Validate inputs early so the spinner doesn't show before a clear error
|
| 442 |
+
try:
|
| 443 |
+
if backend_mode == "Local subprocess":
|
| 444 |
+
events_file = Path(tempfile.NamedTemporaryFile(
|
| 445 |
+
prefix="auto_tune_events_", suffix=".ndjson", delete=False
|
| 446 |
+
).name)
|
| 447 |
+
cmd = _build_command(events_file)
|
| 448 |
+
with header.container():
|
| 449 |
+
st.code(" ".join(cmd), language="bash")
|
| 450 |
+
st.caption(f"Events stream: `{events_file}`")
|
| 451 |
+
else:
|
| 452 |
+
if not backend_url.strip():
|
| 453 |
+
raise ValueError("Backend URL is required for Remote GPU server mode.")
|
| 454 |
+
body = _build_request_body()
|
| 455 |
+
events_file = None
|
| 456 |
+
cmd = None
|
| 457 |
+
with header.container():
|
| 458 |
+
st.code(
|
| 459 |
+
f"POST {backend_url.rstrip('/')}/auto-tune\n"
|
| 460 |
+
+ json.dumps(body, indent=2),
|
| 461 |
+
language="bash",
|
| 462 |
+
)
|
| 463 |
+
st.caption(
|
| 464 |
+
"Events stream: SSE from remote server. "
|
| 465 |
+
"All GPU work happens on that host."
|
| 466 |
+
)
|
| 467 |
+
except ValueError as exc:
|
| 468 |
+
st.error(str(exc))
|
| 469 |
+
st.stop()
|
| 470 |
+
|
| 471 |
+
baseline_metrics: dict | None = None
|
| 472 |
+
best_metrics: dict | None = None # most recent accepted iteration's metrics
|
| 473 |
+
final_summary: dict | None = None
|
| 474 |
+
iter_idx_to_container: dict[int, Any] = {}
|
| 475 |
+
|
| 476 |
+
metrics_row = st.empty()
|
| 477 |
+
with metrics_row.container():
|
| 478 |
+
_render_metrics_row(None, None)
|
| 479 |
+
|
| 480 |
+
progress_bar = st.progress(0, text="Starting auto_tune…")
|
| 481 |
+
|
| 482 |
+
iters_section = st.container()
|
| 483 |
+
iters_section.subheader("Iterations")
|
| 484 |
+
|
| 485 |
+
stdout_buffer: list[str] = []
|
| 486 |
+
expected_iters = max(1, max_iterations)
|
| 487 |
+
return_code: int | None = None
|
| 488 |
+
|
| 489 |
+
|
| 490 |
+
def _handle_event(event: dict) -> None:
|
| 491 |
+
"""Render one event into the live UI. Mutates module-level state for
|
| 492 |
+
baseline/best_metrics/final_summary so the summary block below has
|
| 493 |
+
the data it needs after the run."""
|
| 494 |
+
global baseline_metrics, best_metrics, final_summary # noqa: PLW0603
|
| 495 |
+
etype = event.get("type")
|
| 496 |
+
if etype == "started":
|
| 497 |
+
st.session_state["started_event"] = event
|
| 498 |
+
elif etype == "baseline":
|
| 499 |
+
baseline_metrics = event["metrics"]
|
| 500 |
+
best_metrics = baseline_metrics
|
| 501 |
+
with metrics_row.container():
|
| 502 |
+
_render_metrics_row(baseline_metrics, best_metrics)
|
| 503 |
+
with iters_section:
|
| 504 |
+
with st.container(border=True):
|
| 505 |
+
st.markdown("**Baseline**")
|
| 506 |
+
sub = st.columns(4)
|
| 507 |
+
sub[0].metric("tokens/sec", f"{baseline_metrics.get('tokens_per_sec', 0):,.0f}")
|
| 508 |
+
sub[1].metric("MFU %", f"{baseline_metrics.get('mfu_pct', 0):.2f}")
|
| 509 |
+
sub[2].metric("HBM peak GB", f"{baseline_metrics.get('hbm_peak_gb', 0):.1f}")
|
| 510 |
+
sub[3].metric("GPU util %", f"{baseline_metrics.get('gpu_util_pct', 0):.1f}")
|
| 511 |
+
wb = baseline_metrics.get("waste_budget") or {}
|
| 512 |
+
non_zero = {k: v for k, v in wb.items() if k != "useful_gpu" and v > 0}
|
| 513 |
+
if non_zero:
|
| 514 |
+
wb_str = ", ".join(
|
| 515 |
+
f"`{k}={v:.3f}`"
|
| 516 |
+
for k, v in sorted(non_zero.items(), key=lambda kv: kv[1], reverse=True)
|
| 517 |
+
)
|
| 518 |
+
st.caption(f"Recoverable waste: {wb_str}")
|
| 519 |
+
elif etype == "iter_start":
|
| 520 |
+
i = event["iteration"]
|
| 521 |
+
with iters_section:
|
| 522 |
+
container = st.container(border=True)
|
| 523 |
+
iter_idx_to_container[i] = container
|
| 524 |
+
n_cand = len(event.get("candidates") or [])
|
| 525 |
+
cand_summary = ", ".join(
|
| 526 |
+
c.get("name", "?") for c in event.get("candidates") or []
|
| 527 |
+
)
|
| 528 |
+
container.markdown(
|
| 529 |
+
f"**Iteration {i}** · {n_cand} candidate{'s' if n_cand != 1 else ''}: {cand_summary}"
|
| 530 |
+
)
|
| 531 |
+
progress_bar.progress(
|
| 532 |
+
min(0.99, (i - 1) / expected_iters),
|
| 533 |
+
text=f"Iteration {i} of up to {expected_iters} — proposed: {cand_summary}",
|
| 534 |
+
)
|
| 535 |
+
elif etype == "candidate":
|
| 536 |
+
i = event["iteration"]
|
| 537 |
+
container = iter_idx_to_container.get(i)
|
| 538 |
+
if container is not None:
|
| 539 |
+
with container:
|
| 540 |
+
_render_candidate_card(event)
|
| 541 |
+
progress_bar.progress(
|
| 542 |
+
min(0.99, (i - 1) / expected_iters + 0.05),
|
| 543 |
+
text=f"Iter {i} · candidate {event.get('candidate_index')}/"
|
| 544 |
+
f"{event.get('n_candidates')}: {event.get('name')} "
|
| 545 |
+
f"({event.get('outcome')})",
|
| 546 |
+
)
|
| 547 |
+
elif etype == "merge_attempt":
|
| 548 |
+
i = event["iteration"]
|
| 549 |
+
container = iter_idx_to_container.get(i)
|
| 550 |
+
if container is not None:
|
| 551 |
+
with container:
|
| 552 |
+
outcome = event.get("outcome", "?")
|
| 553 |
+
names = ", ".join(event.get("candidate_names") or [])
|
| 554 |
+
if outcome == "wins":
|
| 555 |
+
delta = event.get("delta_vs_best", 0)
|
| 556 |
+
container.success(
|
| 557 |
+
f"🔗 MERGE WINS — combined ({names}) hit Δ {delta:+.2f}% "
|
| 558 |
+
f"(beats individual best `{event.get('individual_best_name')}`)"
|
| 559 |
+
)
|
| 560 |
+
elif outcome == "lost":
|
| 561 |
+
delta = event.get("delta_vs_best", 0)
|
| 562 |
+
container.info(
|
| 563 |
+
f"🔗 Merge tested ({names}) — Δ {delta:+.2f}%, didn't beat "
|
| 564 |
+
f"individual `{event.get('individual_best_name')}`"
|
| 565 |
+
)
|
| 566 |
+
elif outcome == "crashed":
|
| 567 |
+
container.warning(f"💥 Merge crashed: {names}")
|
| 568 |
+
else:
|
| 569 |
+
container.info(f"🔗 Merge skipped: {event.get('reason', '?')}")
|
| 570 |
+
elif etype == "iter_done":
|
| 571 |
+
i = event["iteration"]
|
| 572 |
+
container = iter_idx_to_container.get(i)
|
| 573 |
+
outcome = event.get("outcome")
|
| 574 |
+
if container is not None:
|
| 575 |
+
with container:
|
| 576 |
+
if outcome == "accepted":
|
| 577 |
+
container.success(
|
| 578 |
+
f"✅ ACCEPTED — `{event.get('winner_name')}` "
|
| 579 |
+
f"(Δ {event.get('winner_delta', 0):+.2f}%)"
|
| 580 |
+
)
|
| 581 |
+
else:
|
| 582 |
+
container.warning(
|
| 583 |
+
f"⏭️ ALL REJECTED — best was `{event.get('winner_name')}` "
|
| 584 |
+
f"(Δ {event.get('winner_delta', 0):+.2f}%, below threshold)"
|
| 585 |
+
)
|
| 586 |
+
if outcome == "accepted" and event.get("best_metrics"):
|
| 587 |
+
best_metrics = event["best_metrics"]
|
| 588 |
+
with metrics_row.container():
|
| 589 |
+
_render_metrics_row(baseline_metrics, best_metrics)
|
| 590 |
+
progress_bar.progress(
|
| 591 |
+
min(0.99, i / expected_iters),
|
| 592 |
+
text=f"Iter {i} done — best so far: {event.get('best_tps', 0):,.0f} tok/s",
|
| 593 |
+
)
|
| 594 |
+
elif etype == "summary":
|
| 595 |
+
final_summary = event
|
| 596 |
+
progress_bar.progress(1.0, text="Auto-tune complete")
|
| 597 |
+
elif etype == "error":
|
| 598 |
+
st.error(event.get("message", "unknown error"))
|
| 599 |
+
elif etype == "process_exit":
|
| 600 |
+
rc = event.get("returncode", "?")
|
| 601 |
+
st.error(
|
| 602 |
+
f"Backend subprocess exited (code {rc}): "
|
| 603 |
+
+ event.get("message", "")
|
| 604 |
+
)
|
| 605 |
+
|
| 606 |
+
|
| 607 |
+
# ---- Source the events from either local subprocess or remote SSE ----
|
| 608 |
+
if backend_mode == "Local subprocess":
|
| 609 |
+
proc = subprocess.Popen(
|
| 610 |
+
cmd,
|
| 611 |
+
stdout=subprocess.PIPE,
|
| 612 |
+
stderr=subprocess.STDOUT,
|
| 613 |
+
text=True,
|
| 614 |
+
bufsize=1,
|
| 615 |
+
cwd=str(REPO_ROOT),
|
| 616 |
+
env={**os.environ},
|
| 617 |
+
)
|
| 618 |
+
|
| 619 |
+
seen_bytes = 0
|
| 620 |
+
try:
|
| 621 |
+
while True:
|
| 622 |
+
# Pull stdout incrementally so the user sees the raw script log too
|
| 623 |
+
while True:
|
| 624 |
+
try:
|
| 625 |
+
line = proc.stdout.readline() if proc.stdout else ""
|
| 626 |
+
except Exception:
|
| 627 |
+
line = ""
|
| 628 |
+
if not line:
|
| 629 |
+
break
|
| 630 |
+
stdout_buffer.append(line.rstrip())
|
| 631 |
+
|
| 632 |
+
new_events, seen_bytes = _read_events(events_file, seen_bytes)
|
| 633 |
+
for event in new_events:
|
| 634 |
+
_handle_event(event)
|
| 635 |
+
|
| 636 |
+
if proc.poll() is not None and not new_events:
|
| 637 |
+
new_events, seen_bytes = _read_events(events_file, seen_bytes)
|
| 638 |
+
if not new_events:
|
| 639 |
+
break
|
| 640 |
+
time.sleep(0.4)
|
| 641 |
+
finally:
|
| 642 |
+
if proc.poll() is None:
|
| 643 |
+
proc.terminate()
|
| 644 |
+
try:
|
| 645 |
+
proc.wait(timeout=3)
|
| 646 |
+
except subprocess.TimeoutExpired:
|
| 647 |
+
proc.kill()
|
| 648 |
+
return_code = proc.returncode
|
| 649 |
+
else:
|
| 650 |
+
# Remote SSE: stream from the FastAPI server
|
| 651 |
+
url = backend_url.rstrip("/") + "/auto-tune"
|
| 652 |
+
try:
|
| 653 |
+
response = requests.post(
|
| 654 |
+
url,
|
| 655 |
+
json=body,
|
| 656 |
+
stream=True,
|
| 657 |
+
timeout=(10, None), # connect timeout 10s, read timeout indefinite
|
| 658 |
+
headers={"Accept": "text/event-stream"},
|
| 659 |
+
)
|
| 660 |
+
response.raise_for_status()
|
| 661 |
+
except requests.RequestException as exc:
|
| 662 |
+
st.error(f"Failed to reach backend at {url}: {exc}")
|
| 663 |
+
st.stop()
|
| 664 |
+
|
| 665 |
+
try:
|
| 666 |
+
for raw_line in response.iter_lines(decode_unicode=True):
|
| 667 |
+
if not raw_line:
|
| 668 |
+
continue
|
| 669 |
+
stdout_buffer.append(raw_line)
|
| 670 |
+
# SSE framing: lines starting with `data: <json>`
|
| 671 |
+
if raw_line.startswith("data:"):
|
| 672 |
+
payload = raw_line[len("data:"):].strip()
|
| 673 |
+
try:
|
| 674 |
+
event = json.loads(payload)
|
| 675 |
+
except json.JSONDecodeError:
|
| 676 |
+
continue
|
| 677 |
+
_handle_event(event)
|
| 678 |
+
if event.get("type") == "summary":
|
| 679 |
+
# Final event; the server may still send a process_exit
|
| 680 |
+
# but we can stop blocking the UI.
|
| 681 |
+
pass
|
| 682 |
+
elif raw_line.startswith("event:") or raw_line.startswith(":"):
|
| 683 |
+
# SSE event-name line or comment — ignored
|
| 684 |
+
continue
|
| 685 |
+
except requests.RequestException as exc:
|
| 686 |
+
st.error(f"Stream interrupted: {exc}")
|
| 687 |
+
return_code = 0 if final_summary is not None else 1
|
| 688 |
+
|
| 689 |
+
# ---------------------------------------------------------------------------
|
| 690 |
+
# Final summary
|
| 691 |
+
# ---------------------------------------------------------------------------
|
| 692 |
+
|
| 693 |
+
st.divider()
|
| 694 |
+
st.subheader("Summary")
|
| 695 |
+
|
| 696 |
+
if final_summary is None:
|
| 697 |
+
st.error(
|
| 698 |
+
f"Auto-tune subprocess exited with code {return_code} but emitted "
|
| 699 |
+
"no summary event. Check the raw stdout below for details."
|
| 700 |
+
)
|
| 701 |
+
else:
|
| 702 |
+
base_tps = float(final_summary.get("baseline_tps") or 0)
|
| 703 |
+
best_tps = float(final_summary.get("best_tps") or 0)
|
| 704 |
+
improvement_pct = float(final_summary.get("improvement_pct") or 0)
|
| 705 |
+
base = final_summary.get("baseline_metrics") or {}
|
| 706 |
+
best = final_summary.get("best_metrics") or {}
|
| 707 |
+
|
| 708 |
+
summary_cols = st.columns(4)
|
| 709 |
+
summary_cols[0].metric(
|
| 710 |
+
"Baseline tokens/sec",
|
| 711 |
+
f"{base_tps:,.0f}",
|
| 712 |
+
)
|
| 713 |
+
summary_cols[1].metric(
|
| 714 |
+
"Best tokens/sec",
|
| 715 |
+
f"{best_tps:,.0f}",
|
| 716 |
+
delta=f"{best_tps - base_tps:+,.0f}",
|
| 717 |
+
)
|
| 718 |
+
summary_cols[2].metric(
|
| 719 |
+
"Improvement",
|
| 720 |
+
f"{improvement_pct:+.2f}%",
|
| 721 |
+
)
|
| 722 |
+
summary_cols[3].metric(
|
| 723 |
+
"MFU baseline → best",
|
| 724 |
+
f"{base.get('mfu_pct', 0):.1f} → {best.get('mfu_pct', 0):.1f} %",
|
| 725 |
+
)
|
| 726 |
+
|
| 727 |
+
chart = _waste_chart(base, best)
|
| 728 |
+
if chart is not None:
|
| 729 |
+
st.markdown("**Waste reduction (seconds/step, by bucket)**")
|
| 730 |
+
st.altair_chart(chart, use_container_width=True)
|
| 731 |
+
# Per-bucket reduction table
|
| 732 |
+
diff_rows = []
|
| 733 |
+
bwb = base.get("waste_budget") or {}
|
| 734 |
+
cwb = best.get("waste_budget") or {}
|
| 735 |
+
for bucket in WASTE_BUCKETS:
|
| 736 |
+
bv = float(bwb.get(bucket, 0.0))
|
| 737 |
+
cv = float(cwb.get(bucket, 0.0))
|
| 738 |
+
if bv == 0 and cv == 0:
|
| 739 |
+
continue
|
| 740 |
+
diff_rows.append({
|
| 741 |
+
"bucket": bucket,
|
| 742 |
+
"baseline (s)": round(bv, 4),
|
| 743 |
+
"best (s)": round(cv, 4),
|
| 744 |
+
"Δ (s)": round(cv - bv, 4),
|
| 745 |
+
})
|
| 746 |
+
if diff_rows:
|
| 747 |
+
st.dataframe(pd.DataFrame(diff_rows), use_container_width=True)
|
| 748 |
+
|
| 749 |
+
accepted = final_summary.get("accepted") or []
|
| 750 |
+
rejected = final_summary.get("rejected") or []
|
| 751 |
+
|
| 752 |
+
col_a, col_r = st.columns(2)
|
| 753 |
+
with col_a:
|
| 754 |
+
st.markdown(f"**✅ Accepted ({len(accepted)})**")
|
| 755 |
+
if accepted:
|
| 756 |
+
st.dataframe(
|
| 757 |
+
pd.DataFrame(accepted)[["name", "tps", "delta_pct"]].rename(
|
| 758 |
+
columns={"tps": "tokens/sec", "delta_pct": "Δ %"}
|
| 759 |
+
),
|
| 760 |
+
use_container_width=True,
|
| 761 |
+
hide_index=True,
|
| 762 |
+
)
|
| 763 |
+
else:
|
| 764 |
+
st.caption("(no experiments accepted)")
|
| 765 |
+
with col_r:
|
| 766 |
+
st.markdown(f"**❌ Rejected ({len(rejected)})**")
|
| 767 |
+
if rejected:
|
| 768 |
+
st.dataframe(
|
| 769 |
+
pd.DataFrame(rejected),
|
| 770 |
+
use_container_width=True,
|
| 771 |
+
hide_index=True,
|
| 772 |
+
)
|
| 773 |
+
else:
|
| 774 |
+
st.caption("(none)")
|
| 775 |
+
|
| 776 |
+
env_vars = final_summary.get("best_env_vars") or {}
|
| 777 |
+
if env_vars:
|
| 778 |
+
st.markdown("**Required env vars for best config**")
|
| 779 |
+
st.code("\n".join(f"export {k}={v}" for k, v in env_vars.items()), language="bash")
|
| 780 |
+
|
| 781 |
+
best_path = final_summary.get("best_workload_path")
|
| 782 |
+
base_path = final_summary.get("baseline_workload_path")
|
| 783 |
+
if best_path and base_path:
|
| 784 |
+
st.markdown("**Best workload script**")
|
| 785 |
+
st.code(f"diff {base_path} {best_path}", language="bash")
|
| 786 |
+
try:
|
| 787 |
+
best_text = Path(best_path).read_text()
|
| 788 |
+
st.download_button(
|
| 789 |
+
"⬇️ Download best.py",
|
| 790 |
+
data=best_text,
|
| 791 |
+
file_name="best.py",
|
| 792 |
+
mime="text/x-python",
|
| 793 |
+
)
|
| 794 |
+
except OSError:
|
| 795 |
+
pass
|
| 796 |
+
|
| 797 |
+
# Raw stdout always at the bottom for debugging
|
| 798 |
+
with st.expander("Raw subprocess output"):
|
| 799 |
+
st.code("\n".join(stdout_buffer), language="text")
|
workloads/_runtime.py
CHANGED
|
@@ -86,6 +86,22 @@ def parse_runtime_args() -> RuntimeArgs:
|
|
| 86 |
)
|
| 87 |
|
| 88 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 89 |
def emit_torch_profile(
|
| 90 |
path: str,
|
| 91 |
*,
|
|
@@ -94,13 +110,16 @@ def emit_torch_profile(
|
|
| 94 |
per_device_batch: int,
|
| 95 |
grad_accum: int = 1,
|
| 96 |
seq_len_cap: int = 512,
|
|
|
|
| 97 |
) -> None:
|
| 98 |
"""Write the smallest torch_profile-shape JSON profile_parser will read.
|
| 99 |
|
| 100 |
profile_parser._read_torch_profile looks for these top-level fields under
|
| 101 |
``metadata``: tokens_per_sec, mfu_pct, step_time_seconds, pytorch_version.
|
| 102 |
-
|
| 103 |
-
|
|
|
|
|
|
|
| 104 |
|
| 105 |
No-ops when ``path`` is empty (script run outside goblin_runner.sh) or
|
| 106 |
when ``n_steps`` is 0 (training crashed before finishing a step).
|
|
@@ -113,43 +132,21 @@ def emit_torch_profile(
|
|
| 113 |
global_batch = max(1, per_device_batch) * max(1, grad_accum)
|
| 114 |
approx_tokens = n_steps * global_batch * seq_len_cap
|
| 115 |
tokens_per_sec = approx_tokens / elapsed if elapsed > 0 else 0.0
|
| 116 |
-
|
| 117 |
-
"
|
| 118 |
-
|
| 119 |
-
|
| 120 |
-
|
| 121 |
-
"n_steps": n_steps,
|
| 122 |
-
}
|
| 123 |
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 124 |
with open(path, "w") as f:
|
| 125 |
json.dump(payload, f)
|
| 126 |
except Exception as exc: # pragma: no cover — diagnostic only
|
| 127 |
# Don't tank the run on a profile-emit failure; the agent will
|
| 128 |
# just see "fake" metrics for this step instead of "live".
|
| 129 |
print(f"[workloads._runtime] failed to write {path}: {exc}")
|
| 130 |
-
|
| 131 |
-
|
| 132 |
-
def trainer_tokenizer_kwargs(trainer_cls, tokenizer) -> dict:
|
| 133 |
-
"""Return the right kwarg for handing a tokenizer to ``Trainer``.
|
| 134 |
-
|
| 135 |
-
transformers ≥ 4.46 renamed ``tokenizer=`` to ``processing_class=`` (the
|
| 136 |
-
old name is still accepted with a DeprecationWarning, but a future
|
| 137 |
-
release drops it). Older versions only know ``tokenizer=``. We
|
| 138 |
-
introspect ``Trainer.__init__`` so the workloads run on either
|
| 139 |
-
generation without pinning a transformers version.
|
| 140 |
-
|
| 141 |
-
Use site:
|
| 142 |
-
|
| 143 |
-
trainer = Trainer(
|
| 144 |
-
model=model,
|
| 145 |
-
args=training_args,
|
| 146 |
-
train_dataset=dataset,
|
| 147 |
-
**trainer_tokenizer_kwargs(Trainer, tokenizer),
|
| 148 |
-
data_collator=_toy_collate,
|
| 149 |
-
)
|
| 150 |
-
"""
|
| 151 |
-
import inspect
|
| 152 |
-
|
| 153 |
-
if "processing_class" in inspect.signature(trainer_cls.__init__).parameters:
|
| 154 |
-
return {"processing_class": tokenizer}
|
| 155 |
-
return {"tokenizer": tokenizer}
|
|
|
|
| 86 |
)
|
| 87 |
|
| 88 |
|
| 89 |
+
# MI300X (CDNA3) peak throughput, dense, bf16/fp16 — both arrive at the
|
| 90 |
+
# same number on this arch since the matrix engine is the same. Source:
|
| 91 |
+
# AMD Instinct MI300X datasheet. With sparsity it's ~2.6 PFLOPS, but
|
| 92 |
+
# transformers training rarely hits the sparse path so we use dense as
|
| 93 |
+
# the realistic peak.
|
| 94 |
+
_MI300X_PEAK_FLOPS_DENSE_BF16 = 1.307e15
|
| 95 |
+
|
| 96 |
+
# FLOPs per token for forward + backward. The standard 6N approximation
|
| 97 |
+
# (forward 2N + backward 4N for full fine-tuning) slightly overestimates
|
| 98 |
+
# LoRA — pure LoRA backward only computes weight gradients for the small
|
| 99 |
+
# adapter matrices, not the frozen base — so true LoRA flops/token is
|
| 100 |
+
# closer to 4N. We use 6N as the conventional choice and accept a ~30%
|
| 101 |
+
# pessimistic MFU for LoRA. Still useful as a relative metric run-to-run.
|
| 102 |
+
_FLOPS_PER_TOKEN_FACTOR = 6
|
| 103 |
+
|
| 104 |
+
|
| 105 |
def emit_torch_profile(
|
| 106 |
path: str,
|
| 107 |
*,
|
|
|
|
| 110 |
per_device_batch: int,
|
| 111 |
grad_accum: int = 1,
|
| 112 |
seq_len_cap: int = 512,
|
| 113 |
+
model_params: int = 0,
|
| 114 |
) -> None:
|
| 115 |
"""Write the smallest torch_profile-shape JSON profile_parser will read.
|
| 116 |
|
| 117 |
profile_parser._read_torch_profile looks for these top-level fields under
|
| 118 |
``metadata``: tokens_per_sec, mfu_pct, step_time_seconds, pytorch_version.
|
| 119 |
+
|
| 120 |
+
`model_params` is optional — pass `sum(p.numel() for p in
|
| 121 |
+
model.parameters())` from the workload to get a populated `mfu_pct`.
|
| 122 |
+
Without it, mfu_pct stays unset (profile_parser will default to 0).
|
| 123 |
|
| 124 |
No-ops when ``path`` is empty (script run outside goblin_runner.sh) or
|
| 125 |
when ``n_steps`` is 0 (training crashed before finishing a step).
|
|
|
|
| 132 |
global_batch = max(1, per_device_batch) * max(1, grad_accum)
|
| 133 |
approx_tokens = n_steps * global_batch * seq_len_cap
|
| 134 |
tokens_per_sec = approx_tokens / elapsed if elapsed > 0 else 0.0
|
| 135 |
+
metadata = {
|
| 136 |
+
"tokens_per_sec": round(tokens_per_sec, 2),
|
| 137 |
+
"step_time_seconds": round(elapsed / n_steps, 4),
|
| 138 |
+
"pytorch_version": torch.__version__,
|
| 139 |
+
"n_steps": n_steps,
|
|
|
|
|
|
|
| 140 |
}
|
| 141 |
+
if model_params > 0 and tokens_per_sec > 0:
|
| 142 |
+
flops_per_token = _FLOPS_PER_TOKEN_FACTOR * model_params
|
| 143 |
+
mfu_pct = (flops_per_token * tokens_per_sec) / _MI300X_PEAK_FLOPS_DENSE_BF16 * 100
|
| 144 |
+
metadata["mfu_pct"] = round(mfu_pct, 2)
|
| 145 |
+
metadata["model_params"] = model_params
|
| 146 |
+
payload = {"metadata": metadata}
|
| 147 |
with open(path, "w") as f:
|
| 148 |
json.dump(payload, f)
|
| 149 |
except Exception as exc: # pragma: no cover — diagnostic only
|
| 150 |
# Don't tank the run on a profile-emit failure; the agent will
|
| 151 |
# just see "fake" metrics for this step instead of "live".
|
| 152 |
print(f"[workloads._runtime] failed to write {path}: {exc}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
workloads/scenarios/README.md
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Test workloads
|
| 2 |
+
|
| 3 |
+
Hand-built training scripts that demo different misconfiguration patterns.
|
| 4 |
+
The agent should produce **different audit reports** for each one — useful
|
| 5 |
+
for stress-testing the KB, for live demos with variety, and for sanity-
|
| 6 |
+
checking that recommendations track the actual configuration.
|
| 7 |
+
|
| 8 |
+
| File | Headline misconfiguration | KB rules expected to fire |
|
| 9 |
+
|---|---|---|
|
| 10 |
+
| `train_qwen_full_ft.py` | Full fine-tune (no LoRA), fp32, no gradient checkpointing | `precision.fp32_default_wastes_matrix_cores`, `memory.gradient_checkpointing_for_long_seq`, `attention.sdpa_over_eager` |
|
| 11 |
+
| `train_qwen_long_context.py` | seq_len=8192, no gradient checkpointing, fp16, eager attention | `memory.gradient_checkpointing_for_long_seq`, `attention.flash_rocm_over_eager`, `precision.bf16_over_fp16_on_mi300x` |
|
| 12 |
+
| `train_qwen_distributed_bad.py` | One process driving 8 GPUs (ROCm launch serialization antipattern) | `collectives.one_process_per_gpu`, `env.nccl_min_nchannels`, `env.numa_auto_balancing_disable` |
|
| 13 |
+
| `train_qwen_bnb.py` | Uses bitsandbytes 8-bit Adam (not officially ROCm-supported) | `optimizer.bitsandbytes_not_supported_warning`, `precision.bf16_over_fp16_on_mi300x` |
|
| 14 |
+
| `train_qwen_well_tuned.py` | Already optimized (bf16, flash_rocm, prefetch=4, persistent_workers) | None or only minor suggestions — sanity check that the agent says "nothing to do" instead of inventing problems |
|
| 15 |
+
|
| 16 |
+
Run any of them through the agent:
|
| 17 |
+
|
| 18 |
+
```bash
|
| 19 |
+
export GOBLIN_AGENT_BACKEND=qwen-vllm
|
| 20 |
+
export GOBLIN_QWEN_VLLM_URL=http://localhost:8001/v1
|
| 21 |
+
export GOBLIN_QWEN_VLLM_MODEL=Qwen/Qwen3-32B
|
| 22 |
+
python -m agent workloads/scenarios/train_qwen_long_context.py
|
| 23 |
+
```
|
| 24 |
+
|
| 25 |
+
These are **AST-parseable, optionally executable**. Each one redirects on
|
| 26 |
+
the same fix — `parse_config` only walks the AST, so even if the script
|
| 27 |
+
doesn't actually train cleanly, the agent's audit still works. If you
|
| 28 |
+
want rocprofv3 to capture real numbers (vs FakeRunner fallback), the
|
| 29 |
+
scripts marked "executable: yes" below are runnable; the others rely on
|
| 30 |
+
FakeRunner replay.
|
workloads/scenarios/train_qwen_bnb.py
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Test workload: bitsandbytes 8-bit Adam — the classic "ported from NVIDIA"
|
| 2 |
+
# script that doesn't realize bitsandbytes isn't officially supported on ROCm.
|
| 3 |
+
#
|
| 4 |
+
# Headline misconfigurations the agent should catch:
|
| 5 |
+
# - optimizer.bitsandbytes_not_supported_warning (optim="paged_adamw_8bit")
|
| 6 |
+
# - precision.bf16_over_fp16_on_mi300x (fp16=True)
|
| 7 |
+
# - attention.flash_rocm_over_eager (attn_implementation="eager")
|
| 8 |
+
# - data.dataloader_workers_zero (num_workers=0)
|
| 9 |
+
#
|
| 10 |
+
# Executable: not on ROCm — bitsandbytes will fail at import or on the
|
| 11 |
+
# first 8-bit GEMM. The whole point of this fixture is that the agent
|
| 12 |
+
# warns the user BEFORE they hit that wall. AST parse always works.
|
| 13 |
+
|
| 14 |
+
import os
|
| 15 |
+
import sys
|
| 16 |
+
import time
|
| 17 |
+
|
| 18 |
+
sys.path.insert(
|
| 19 |
+
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 20 |
+
)
|
| 21 |
+
|
| 22 |
+
import torch
|
| 23 |
+
from datasets import load_dataset
|
| 24 |
+
from peft import LoraConfig, get_peft_model
|
| 25 |
+
from torch.utils.data import DataLoader
|
| 26 |
+
from transformers import (
|
| 27 |
+
AutoModelForCausalLM,
|
| 28 |
+
AutoTokenizer,
|
| 29 |
+
Trainer,
|
| 30 |
+
TrainingArguments,
|
| 31 |
+
)
|
| 32 |
+
|
| 33 |
+
from workloads._runtime import emit_torch_profile, parse_runtime_args
|
| 34 |
+
|
| 35 |
+
_runtime = parse_runtime_args()
|
| 36 |
+
|
| 37 |
+
os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
| 38 |
+
HF_TOKEN = os.environ["HF_TOKEN"]
|
| 39 |
+
|
| 40 |
+
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
|
| 41 |
+
|
| 42 |
+
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
|
| 43 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 44 |
+
MODEL_ID,
|
| 45 |
+
torch_dtype=torch.float16,
|
| 46 |
+
attn_implementation="eager",
|
| 47 |
+
token=HF_TOKEN,
|
| 48 |
+
load_in_8bit=True, # bitsandbytes load — ROCm support unofficial
|
| 49 |
+
)
|
| 50 |
+
|
| 51 |
+
lora_config = LoraConfig(
|
| 52 |
+
r=16,
|
| 53 |
+
lora_alpha=32,
|
| 54 |
+
target_modules=["q_proj", "v_proj"],
|
| 55 |
+
lora_dropout=0.05,
|
| 56 |
+
bias="none",
|
| 57 |
+
task_type="CAUSAL_LM",
|
| 58 |
+
)
|
| 59 |
+
model = get_peft_model(model, lora_config)
|
| 60 |
+
|
| 61 |
+
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
|
| 62 |
+
|
| 63 |
+
train_loader = DataLoader(
|
| 64 |
+
dataset,
|
| 65 |
+
batch_size=4,
|
| 66 |
+
num_workers=0,
|
| 67 |
+
pin_memory=False,
|
| 68 |
+
persistent_workers=False,
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
_RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
|
| 72 |
+
|
| 73 |
+
training_args = TrainingArguments(
|
| 74 |
+
output_dir="./out",
|
| 75 |
+
per_device_train_batch_size=4,
|
| 76 |
+
gradient_accumulation_steps=8,
|
| 77 |
+
num_train_epochs=3,
|
| 78 |
+
learning_rate=2e-4,
|
| 79 |
+
warmup_steps=100,
|
| 80 |
+
fp16=True,
|
| 81 |
+
optim="paged_adamw_8bit", # bitsandbytes optimizer — ROCm unofficial
|
| 82 |
+
logging_steps=10,
|
| 83 |
+
save_steps=500,
|
| 84 |
+
dataloader_num_workers=0,
|
| 85 |
+
dataloader_pin_memory=False,
|
| 86 |
+
gradient_checkpointing=False,
|
| 87 |
+
torch_compile=False,
|
| 88 |
+
report_to="none",
|
| 89 |
+
push_to_hub=False,
|
| 90 |
+
max_steps=_RUNTIME_MAX_STEPS,
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
trainer = Trainer(
|
| 94 |
+
model=model,
|
| 95 |
+
args=training_args,
|
| 96 |
+
train_dataset=dataset,
|
| 97 |
+
tokenizer=tokenizer,
|
| 98 |
+
)
|
| 99 |
+
|
| 100 |
+
if __name__ == "__main__":
|
| 101 |
+
_t0 = time.time()
|
| 102 |
+
trainer.train()
|
| 103 |
+
emit_torch_profile(
|
| 104 |
+
_runtime.torch_profile_out,
|
| 105 |
+
elapsed=time.time() - _t0,
|
| 106 |
+
n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
|
| 107 |
+
per_device_batch=training_args.per_device_train_batch_size,
|
| 108 |
+
grad_accum=training_args.gradient_accumulation_steps,
|
| 109 |
+
seq_len_cap=512,
|
| 110 |
+
)
|
workloads/scenarios/train_qwen_distributed_bad.py
ADDED
|
@@ -0,0 +1,122 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Test workload: multi-GPU LoRA with the "one-process-many-GPUs" antipattern.
|
| 2 |
+
#
|
| 3 |
+
# A single Python process (no torchrun, no accelerate launch) tries to drive
|
| 4 |
+
# multiple MI300X devices via DataParallel. ROCm serializes kernel launches
|
| 5 |
+
# across devices when one process owns multiple HIP streams, so every
|
| 6 |
+
# collective becomes a launch-queue bottleneck.
|
| 7 |
+
#
|
| 8 |
+
# Headline misconfigurations the agent should catch:
|
| 9 |
+
# - collectives.one_process_per_gpu (single launcher, multiple devices)
|
| 10 |
+
# - env.nccl_min_nchannels (NCCL_MIN_NCHANNELS not set)
|
| 11 |
+
# - env.numa_auto_balancing_disable (NUMA balancing left on by default)
|
| 12 |
+
# - precision.bf16_over_fp16_on_mi300x (fp16=True)
|
| 13 |
+
# - attention.flash_rocm_over_eager
|
| 14 |
+
#
|
| 15 |
+
# Executable: requires multiple MI300X visible to one process. Even when
|
| 16 |
+
# runnable, ROCm warns on the launch serialization. AST parse is fine.
|
| 17 |
+
|
| 18 |
+
import os
|
| 19 |
+
import sys
|
| 20 |
+
import time
|
| 21 |
+
|
| 22 |
+
sys.path.insert(
|
| 23 |
+
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 24 |
+
)
|
| 25 |
+
|
| 26 |
+
import torch
|
| 27 |
+
import torch.nn as nn
|
| 28 |
+
from datasets import load_dataset
|
| 29 |
+
from peft import LoraConfig, get_peft_model
|
| 30 |
+
from torch.utils.data import DataLoader
|
| 31 |
+
from transformers import (
|
| 32 |
+
AutoModelForCausalLM,
|
| 33 |
+
AutoTokenizer,
|
| 34 |
+
Trainer,
|
| 35 |
+
TrainingArguments,
|
| 36 |
+
)
|
| 37 |
+
|
| 38 |
+
from workloads._runtime import emit_torch_profile, parse_runtime_args
|
| 39 |
+
|
| 40 |
+
_runtime = parse_runtime_args()
|
| 41 |
+
|
| 42 |
+
os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
| 43 |
+
HF_TOKEN = os.environ["HF_TOKEN"]
|
| 44 |
+
os.environ["HSA_FORCE_FINE_GRAIN_PCIE"] = "1"
|
| 45 |
+
# Notably absent: NCCL_MIN_NCHANNELS, GOBLIN_HINT_NUMA_AUTO_BALANCING.
|
| 46 |
+
|
| 47 |
+
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
|
| 48 |
+
|
| 49 |
+
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
|
| 50 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 51 |
+
MODEL_ID,
|
| 52 |
+
torch_dtype=torch.float16,
|
| 53 |
+
attn_implementation="eager",
|
| 54 |
+
token=HF_TOKEN,
|
| 55 |
+
)
|
| 56 |
+
|
| 57 |
+
lora_config = LoraConfig(
|
| 58 |
+
r=16,
|
| 59 |
+
lora_alpha=32,
|
| 60 |
+
target_modules=["q_proj", "v_proj"],
|
| 61 |
+
lora_dropout=0.05,
|
| 62 |
+
bias="none",
|
| 63 |
+
task_type="CAUSAL_LM",
|
| 64 |
+
)
|
| 65 |
+
model = get_peft_model(model, lora_config)
|
| 66 |
+
|
| 67 |
+
# THE ANTIPATTERN: DataParallel from a single process across all visible GPUs.
|
| 68 |
+
# Production code should use torchrun --nproc_per_node=N or accelerate launch.
|
| 69 |
+
if torch.cuda.device_count() > 1:
|
| 70 |
+
model = nn.DataParallel(model)
|
| 71 |
+
|
| 72 |
+
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
|
| 73 |
+
|
| 74 |
+
train_loader = DataLoader(
|
| 75 |
+
dataset,
|
| 76 |
+
batch_size=8,
|
| 77 |
+
num_workers=4,
|
| 78 |
+
pin_memory=True,
|
| 79 |
+
persistent_workers=True,
|
| 80 |
+
)
|
| 81 |
+
|
| 82 |
+
_RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
|
| 83 |
+
|
| 84 |
+
training_args = TrainingArguments(
|
| 85 |
+
output_dir="./out",
|
| 86 |
+
per_device_train_batch_size=8,
|
| 87 |
+
gradient_accumulation_steps=4,
|
| 88 |
+
num_train_epochs=1,
|
| 89 |
+
learning_rate=2e-4,
|
| 90 |
+
warmup_steps=100,
|
| 91 |
+
fp16=True,
|
| 92 |
+
optim="adamw_torch",
|
| 93 |
+
logging_steps=10,
|
| 94 |
+
save_steps=500,
|
| 95 |
+
dataloader_num_workers=4,
|
| 96 |
+
dataloader_pin_memory=True,
|
| 97 |
+
dataloader_persistent_workers=True,
|
| 98 |
+
gradient_checkpointing=False,
|
| 99 |
+
torch_compile=False,
|
| 100 |
+
report_to="none",
|
| 101 |
+
push_to_hub=False,
|
| 102 |
+
max_steps=_RUNTIME_MAX_STEPS,
|
| 103 |
+
)
|
| 104 |
+
|
| 105 |
+
trainer = Trainer(
|
| 106 |
+
model=model,
|
| 107 |
+
args=training_args,
|
| 108 |
+
train_dataset=dataset,
|
| 109 |
+
tokenizer=tokenizer,
|
| 110 |
+
)
|
| 111 |
+
|
| 112 |
+
if __name__ == "__main__":
|
| 113 |
+
_t0 = time.time()
|
| 114 |
+
trainer.train()
|
| 115 |
+
emit_torch_profile(
|
| 116 |
+
_runtime.torch_profile_out,
|
| 117 |
+
elapsed=time.time() - _t0,
|
| 118 |
+
n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
|
| 119 |
+
per_device_batch=training_args.per_device_train_batch_size,
|
| 120 |
+
grad_accum=training_args.gradient_accumulation_steps,
|
| 121 |
+
seq_len_cap=512,
|
| 122 |
+
)
|
workloads/scenarios/train_qwen_full_ft.py
ADDED
|
@@ -0,0 +1,103 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Test workload: full fine-tune (no LoRA), fp32, no gradient checkpointing.
|
| 2 |
+
#
|
| 3 |
+
# Headline misconfigurations the agent should catch:
|
| 4 |
+
# - precision.fp32_default_wastes_matrix_cores (fp16=False, bf16=False)
|
| 5 |
+
# - memory.gradient_checkpointing_for_long_seq (gradient_checkpointing=False, seq_len=2048)
|
| 6 |
+
# - attention.sdpa_over_eager OR flash_rocm (attn_implementation="eager")
|
| 7 |
+
# - data.dataloader_workers_zero (num_workers=0)
|
| 8 |
+
# - data.pin_memory_false (pin_memory=False)
|
| 9 |
+
#
|
| 10 |
+
# Executable: not really — full FT of a 7B model OOMs on a single MI300X
|
| 11 |
+
# without LoRA. AST parse is fine; rocprofv3 will fail and FakeRunner kicks in.
|
| 12 |
+
|
| 13 |
+
import os
|
| 14 |
+
import sys
|
| 15 |
+
import time
|
| 16 |
+
|
| 17 |
+
# Bootstrap repo root so `from workloads._runtime import ...` resolves
|
| 18 |
+
# regardless of the cwd goblin_runner.sh launches us from.
|
| 19 |
+
sys.path.insert(
|
| 20 |
+
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 21 |
+
)
|
| 22 |
+
|
| 23 |
+
import torch
|
| 24 |
+
from datasets import load_dataset
|
| 25 |
+
from torch.utils.data import DataLoader
|
| 26 |
+
from transformers import (
|
| 27 |
+
AutoModelForCausalLM,
|
| 28 |
+
AutoTokenizer,
|
| 29 |
+
Trainer,
|
| 30 |
+
TrainingArguments,
|
| 31 |
+
)
|
| 32 |
+
|
| 33 |
+
from workloads._runtime import emit_torch_profile, parse_runtime_args
|
| 34 |
+
|
| 35 |
+
_runtime = parse_runtime_args()
|
| 36 |
+
|
| 37 |
+
os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
| 38 |
+
HF_TOKEN = os.environ["HF_TOKEN"]
|
| 39 |
+
os.environ["HSA_FORCE_FINE_GRAIN_PCIE"] = "1"
|
| 40 |
+
|
| 41 |
+
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
|
| 42 |
+
|
| 43 |
+
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
|
| 44 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 45 |
+
MODEL_ID,
|
| 46 |
+
torch_dtype=torch.float32, # fp32 — wastes CDNA3 matrix cores
|
| 47 |
+
attn_implementation="eager",
|
| 48 |
+
token=HF_TOKEN,
|
| 49 |
+
)
|
| 50 |
+
|
| 51 |
+
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
|
| 52 |
+
|
| 53 |
+
train_loader = DataLoader(
|
| 54 |
+
dataset,
|
| 55 |
+
batch_size=2, # tiny because fp32 + full FT eats HBM
|
| 56 |
+
num_workers=0,
|
| 57 |
+
pin_memory=False,
|
| 58 |
+
persistent_workers=False,
|
| 59 |
+
)
|
| 60 |
+
|
| 61 |
+
# Literal kwargs so parse_config picks them all up (see canonical workload).
|
| 62 |
+
_RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
|
| 63 |
+
|
| 64 |
+
training_args = TrainingArguments(
|
| 65 |
+
output_dir="./out",
|
| 66 |
+
per_device_train_batch_size=2,
|
| 67 |
+
gradient_accumulation_steps=16,
|
| 68 |
+
num_train_epochs=1,
|
| 69 |
+
learning_rate=1e-5,
|
| 70 |
+
warmup_steps=200,
|
| 71 |
+
fp16=False, # no precision flag → fp32 path
|
| 72 |
+
bf16=False,
|
| 73 |
+
optim="adamw_torch",
|
| 74 |
+
max_seq_length=2048,
|
| 75 |
+
logging_steps=10,
|
| 76 |
+
save_steps=500,
|
| 77 |
+
dataloader_num_workers=0,
|
| 78 |
+
dataloader_pin_memory=False,
|
| 79 |
+
gradient_checkpointing=False, # at seq=2048 this leaves a lot of HBM tied up
|
| 80 |
+
torch_compile=False,
|
| 81 |
+
report_to="none",
|
| 82 |
+
push_to_hub=False,
|
| 83 |
+
max_steps=_RUNTIME_MAX_STEPS,
|
| 84 |
+
)
|
| 85 |
+
|
| 86 |
+
trainer = Trainer(
|
| 87 |
+
model=model,
|
| 88 |
+
args=training_args,
|
| 89 |
+
train_dataset=dataset,
|
| 90 |
+
tokenizer=tokenizer,
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
if __name__ == "__main__":
|
| 94 |
+
_t0 = time.time()
|
| 95 |
+
trainer.train()
|
| 96 |
+
emit_torch_profile(
|
| 97 |
+
_runtime.torch_profile_out,
|
| 98 |
+
elapsed=time.time() - _t0,
|
| 99 |
+
n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
|
| 100 |
+
per_device_batch=training_args.per_device_train_batch_size,
|
| 101 |
+
grad_accum=training_args.gradient_accumulation_steps,
|
| 102 |
+
seq_len_cap=2048,
|
| 103 |
+
)
|
workloads/scenarios/train_qwen_long_context.py
ADDED
|
@@ -0,0 +1,111 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Test workload: long-context LoRA fine-tune.
|
| 2 |
+
#
|
| 3 |
+
# Headline misconfigurations the agent should catch:
|
| 4 |
+
# - memory.gradient_checkpointing_for_long_seq (seq_len=8192, gradient_checkpointing=False)
|
| 5 |
+
# - attention.flash_rocm_over_eager (attn_implementation="eager" — KILLS at seq 8K)
|
| 6 |
+
# - precision.bf16_over_fp16_on_mi300x (fp16=True)
|
| 7 |
+
# - data.persistent_workers_false (persistent_workers=False)
|
| 8 |
+
# - data.pin_memory_false (pin_memory=False)
|
| 9 |
+
#
|
| 10 |
+
# Executable: only with flash attention installed. Without it, eager attention
|
| 11 |
+
# at seq=8192 blows up HBM. AST parse always works.
|
| 12 |
+
|
| 13 |
+
import os
|
| 14 |
+
import sys
|
| 15 |
+
import time
|
| 16 |
+
|
| 17 |
+
sys.path.insert(
|
| 18 |
+
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 19 |
+
)
|
| 20 |
+
|
| 21 |
+
import torch
|
| 22 |
+
from datasets import load_dataset
|
| 23 |
+
from peft import LoraConfig, get_peft_model
|
| 24 |
+
from torch.utils.data import DataLoader
|
| 25 |
+
from transformers import (
|
| 26 |
+
AutoModelForCausalLM,
|
| 27 |
+
AutoTokenizer,
|
| 28 |
+
Trainer,
|
| 29 |
+
TrainingArguments,
|
| 30 |
+
)
|
| 31 |
+
|
| 32 |
+
from workloads._runtime import emit_torch_profile, parse_runtime_args
|
| 33 |
+
|
| 34 |
+
_runtime = parse_runtime_args()
|
| 35 |
+
|
| 36 |
+
os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
| 37 |
+
HF_TOKEN = os.environ["HF_TOKEN"]
|
| 38 |
+
os.environ["MIOPEN_FIND_MODE"] = "3"
|
| 39 |
+
|
| 40 |
+
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
|
| 41 |
+
|
| 42 |
+
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
|
| 43 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 44 |
+
MODEL_ID,
|
| 45 |
+
torch_dtype=torch.float16,
|
| 46 |
+
attn_implementation="eager", # at seq=8192 this is catastrophic
|
| 47 |
+
token=HF_TOKEN,
|
| 48 |
+
)
|
| 49 |
+
|
| 50 |
+
lora_config = LoraConfig(
|
| 51 |
+
r=32,
|
| 52 |
+
lora_alpha=64,
|
| 53 |
+
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
| 54 |
+
lora_dropout=0.05,
|
| 55 |
+
bias="none",
|
| 56 |
+
task_type="CAUSAL_LM",
|
| 57 |
+
)
|
| 58 |
+
model = get_peft_model(model, lora_config)
|
| 59 |
+
|
| 60 |
+
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
|
| 61 |
+
|
| 62 |
+
train_loader = DataLoader(
|
| 63 |
+
dataset,
|
| 64 |
+
batch_size=1, # forced down by long context
|
| 65 |
+
num_workers=4, # workers OK; problem is elsewhere
|
| 66 |
+
pin_memory=False,
|
| 67 |
+
persistent_workers=False,
|
| 68 |
+
)
|
| 69 |
+
|
| 70 |
+
_RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
|
| 71 |
+
|
| 72 |
+
training_args = TrainingArguments(
|
| 73 |
+
output_dir="./out",
|
| 74 |
+
per_device_train_batch_size=1,
|
| 75 |
+
gradient_accumulation_steps=32,
|
| 76 |
+
num_train_epochs=2,
|
| 77 |
+
learning_rate=2e-4,
|
| 78 |
+
warmup_steps=50,
|
| 79 |
+
fp16=True, # bf16 is the right call on CDNA3
|
| 80 |
+
optim="adamw_torch",
|
| 81 |
+
max_seq_length=8192, # the long-context part
|
| 82 |
+
logging_steps=10,
|
| 83 |
+
save_steps=500,
|
| 84 |
+
dataloader_num_workers=4,
|
| 85 |
+
dataloader_pin_memory=False,
|
| 86 |
+
dataloader_persistent_workers=False,
|
| 87 |
+
gradient_checkpointing=False, # missing — at seq 8K activations dominate HBM
|
| 88 |
+
torch_compile=False,
|
| 89 |
+
report_to="none",
|
| 90 |
+
push_to_hub=False,
|
| 91 |
+
max_steps=_RUNTIME_MAX_STEPS,
|
| 92 |
+
)
|
| 93 |
+
|
| 94 |
+
trainer = Trainer(
|
| 95 |
+
model=model,
|
| 96 |
+
args=training_args,
|
| 97 |
+
train_dataset=dataset,
|
| 98 |
+
tokenizer=tokenizer,
|
| 99 |
+
)
|
| 100 |
+
|
| 101 |
+
if __name__ == "__main__":
|
| 102 |
+
_t0 = time.time()
|
| 103 |
+
trainer.train()
|
| 104 |
+
emit_torch_profile(
|
| 105 |
+
_runtime.torch_profile_out,
|
| 106 |
+
elapsed=time.time() - _t0,
|
| 107 |
+
n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
|
| 108 |
+
per_device_batch=training_args.per_device_train_batch_size,
|
| 109 |
+
grad_accum=training_args.gradient_accumulation_steps,
|
| 110 |
+
seq_len_cap=8192,
|
| 111 |
+
)
|
workloads/scenarios/train_qwen_well_tuned.py
ADDED
|
@@ -0,0 +1,143 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Test workload: already well-tuned LoRA fine-tune.
|
| 2 |
+
#
|
| 3 |
+
# This is the negative-control fixture. The agent should produce a SHORT
|
| 4 |
+
# audit with few or no recommendations — the goal is to catch hallucinated
|
| 5 |
+
# rule applications. If GPU Goblin invents problems on a clean config, the
|
| 6 |
+
# audit isn't trustworthy.
|
| 7 |
+
#
|
| 8 |
+
# What's already correct here:
|
| 9 |
+
# - precision = bf16
|
| 10 |
+
# - attention_impl = flash_rocm (Optimum-AMD path)
|
| 11 |
+
# - dataloader_num_workers = 8, pin_memory=True, prefetch_factor=4,
|
| 12 |
+
# persistent_workers=True
|
| 13 |
+
# - gradient_checkpointing = True (long context)
|
| 14 |
+
# - torch_compile = True
|
| 15 |
+
# - NCCL_MIN_NCHANNELS, HSA_FORCE_FINE_GRAIN_PCIE both set
|
| 16 |
+
#
|
| 17 |
+
# Executable: yes (fastest path; meant to actually run cleanly).
|
| 18 |
+
|
| 19 |
+
import os
|
| 20 |
+
import sys
|
| 21 |
+
import time
|
| 22 |
+
|
| 23 |
+
sys.path.insert(
|
| 24 |
+
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 25 |
+
)
|
| 26 |
+
|
| 27 |
+
import torch
|
| 28 |
+
from datasets import load_dataset
|
| 29 |
+
from peft import LoraConfig, get_peft_model
|
| 30 |
+
from torch.utils.data import DataLoader
|
| 31 |
+
from transformers import (
|
| 32 |
+
AutoModelForCausalLM,
|
| 33 |
+
AutoTokenizer,
|
| 34 |
+
Trainer,
|
| 35 |
+
TrainingArguments,
|
| 36 |
+
)
|
| 37 |
+
|
| 38 |
+
from workloads._runtime import emit_torch_profile, parse_runtime_args
|
| 39 |
+
|
| 40 |
+
_runtime = parse_runtime_args()
|
| 41 |
+
|
| 42 |
+
os.environ["HF_TOKEN"] = "hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
| 43 |
+
HF_TOKEN = os.environ["HF_TOKEN"]
|
| 44 |
+
|
| 45 |
+
# All the env knobs from the KB, set correctly.
|
| 46 |
+
os.environ["HSA_FORCE_FINE_GRAIN_PCIE"] = "1"
|
| 47 |
+
os.environ["MIOPEN_FIND_MODE"] = "3"
|
| 48 |
+
os.environ["NCCL_MIN_NCHANNELS"] = "112"
|
| 49 |
+
|
| 50 |
+
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
|
| 51 |
+
|
| 52 |
+
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
|
| 53 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 54 |
+
MODEL_ID,
|
| 55 |
+
torch_dtype=torch.bfloat16, # bf16 ✓
|
| 56 |
+
attn_implementation="flash_attention_2", # parser maps to flash / flash_rocm
|
| 57 |
+
token=HF_TOKEN,
|
| 58 |
+
)
|
| 59 |
+
|
| 60 |
+
# Compile path on — Qwen is on the eager-friendly list.
|
| 61 |
+
model = torch.compile(model, mode="reduce-overhead")
|
| 62 |
+
|
| 63 |
+
lora_config = LoraConfig(
|
| 64 |
+
r=16,
|
| 65 |
+
lora_alpha=32,
|
| 66 |
+
target_modules=["q_proj", "v_proj"],
|
| 67 |
+
lora_dropout=0.05,
|
| 68 |
+
bias="none",
|
| 69 |
+
task_type="CAUSAL_LM",
|
| 70 |
+
)
|
| 71 |
+
model = get_peft_model(model, lora_config)
|
| 72 |
+
model.gradient_checkpointing_enable()
|
| 73 |
+
|
| 74 |
+
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
|
| 75 |
+
|
| 76 |
+
train_loader = DataLoader(
|
| 77 |
+
dataset,
|
| 78 |
+
batch_size=12, # comfortable on 192 GB HBM
|
| 79 |
+
num_workers=8,
|
| 80 |
+
pin_memory=True,
|
| 81 |
+
prefetch_factor=4,
|
| 82 |
+
persistent_workers=True,
|
| 83 |
+
)
|
| 84 |
+
|
| 85 |
+
_RUNTIME_MAX_STEPS = _runtime.max_steps if _runtime.max_steps > 0 else -1
|
| 86 |
+
|
| 87 |
+
training_args = TrainingArguments(
|
| 88 |
+
output_dir="./out",
|
| 89 |
+
per_device_train_batch_size=12,
|
| 90 |
+
gradient_accumulation_steps=2,
|
| 91 |
+
num_train_epochs=1,
|
| 92 |
+
learning_rate=2e-4,
|
| 93 |
+
warmup_steps=100,
|
| 94 |
+
bf16=True, # bf16 ✓
|
| 95 |
+
optim="adamw_torch",
|
| 96 |
+
max_seq_length=4096,
|
| 97 |
+
logging_steps=10,
|
| 98 |
+
save_steps=500,
|
| 99 |
+
dataloader_num_workers=8,
|
| 100 |
+
dataloader_pin_memory=True,
|
| 101 |
+
dataloader_prefetch_factor=4,
|
| 102 |
+
dataloader_persistent_workers=True,
|
| 103 |
+
gradient_checkpointing=True,
|
| 104 |
+
torch_compile=True,
|
| 105 |
+
report_to="none",
|
| 106 |
+
push_to_hub=False,
|
| 107 |
+
remove_unused_columns=False,
|
| 108 |
+
max_steps=_RUNTIME_MAX_STEPS,
|
| 109 |
+
)
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def _toy_collate(rows):
|
| 113 |
+
texts = [
|
| 114 |
+
(r.get("instruction") or "")
|
| 115 |
+
+ ("\n" + r["input"] if r.get("input") else "")
|
| 116 |
+
+ "\n"
|
| 117 |
+
+ (r.get("output") or "")
|
| 118 |
+
for r in rows
|
| 119 |
+
]
|
| 120 |
+
enc = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
|
| 121 |
+
enc["labels"] = enc["input_ids"].clone()
|
| 122 |
+
return enc
|
| 123 |
+
|
| 124 |
+
|
| 125 |
+
trainer = Trainer(
|
| 126 |
+
model=model,
|
| 127 |
+
args=training_args,
|
| 128 |
+
train_dataset=dataset,
|
| 129 |
+
tokenizer=tokenizer,
|
| 130 |
+
data_collator=_toy_collate,
|
| 131 |
+
)
|
| 132 |
+
|
| 133 |
+
if __name__ == "__main__":
|
| 134 |
+
_t0 = time.time()
|
| 135 |
+
trainer.train()
|
| 136 |
+
emit_torch_profile(
|
| 137 |
+
_runtime.torch_profile_out,
|
| 138 |
+
elapsed=time.time() - _t0,
|
| 139 |
+
n_steps=int(getattr(trainer.state, "global_step", 0) or _runtime.max_steps),
|
| 140 |
+
per_device_batch=training_args.per_device_train_batch_size,
|
| 141 |
+
grad_accum=training_args.gradient_accumulation_steps,
|
| 142 |
+
seq_len_cap=4096,
|
| 143 |
+
)
|
workloads/train_qwen_lora.py
CHANGED
|
@@ -30,11 +30,7 @@ from transformers import (
|
|
| 30 |
TrainingArguments,
|
| 31 |
)
|
| 32 |
|
| 33 |
-
from workloads._runtime import
|
| 34 |
-
emit_torch_profile,
|
| 35 |
-
parse_runtime_args,
|
| 36 |
-
trainer_tokenizer_kwargs,
|
| 37 |
-
)
|
| 38 |
|
| 39 |
# Parse the goblin_runner.sh injected flags (--max_steps, --torch_profile_out).
|
| 40 |
_runtime = parse_runtime_args()
|
|
@@ -150,12 +146,17 @@ trainer = Trainer(
|
|
| 150 |
model=model,
|
| 151 |
args=training_args,
|
| 152 |
train_dataset=dataset,
|
| 153 |
-
|
| 154 |
data_collator=_toy_collate,
|
| 155 |
)
|
| 156 |
|
| 157 |
|
| 158 |
if __name__ == "__main__":
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 159 |
_t0 = time.time()
|
| 160 |
trainer.train()
|
| 161 |
_elapsed = time.time() - _t0
|
|
@@ -170,4 +171,5 @@ if __name__ == "__main__":
|
|
| 170 |
per_device_batch=training_args.per_device_train_batch_size,
|
| 171 |
grad_accum=training_args.gradient_accumulation_steps,
|
| 172 |
seq_len_cap=512,
|
|
|
|
| 173 |
)
|
|
|
|
| 30 |
TrainingArguments,
|
| 31 |
)
|
| 32 |
|
| 33 |
+
from workloads._runtime import emit_torch_profile, parse_runtime_args
|
|
|
|
|
|
|
|
|
|
|
|
|
| 34 |
|
| 35 |
# Parse the goblin_runner.sh injected flags (--max_steps, --torch_profile_out).
|
| 36 |
_runtime = parse_runtime_args()
|
|
|
|
| 146 |
model=model,
|
| 147 |
args=training_args,
|
| 148 |
train_dataset=dataset,
|
| 149 |
+
tokenizer=tokenizer,
|
| 150 |
data_collator=_toy_collate,
|
| 151 |
)
|
| 152 |
|
| 153 |
|
| 154 |
if __name__ == "__main__":
|
| 155 |
+
# Total parameter count drives MFU. For a peft-wrapped model this
|
| 156 |
+
# includes both the frozen base (which still does forward) and the
|
| 157 |
+
# tiny LoRA adapters. The forward+backward FLOPs estimate uses the
|
| 158 |
+
# full count (the standard 6N rule lives inside emit_torch_profile).
|
| 159 |
+
_model_params = sum(p.numel() for p in model.parameters())
|
| 160 |
_t0 = time.time()
|
| 161 |
trainer.train()
|
| 162 |
_elapsed = time.time() - _t0
|
|
|
|
| 171 |
per_device_batch=training_args.per_device_train_batch_size,
|
| 172 |
grad_accum=training_args.gradient_accumulation_steps,
|
| 173 |
seq_len_cap=512,
|
| 174 |
+
model_params=_model_params,
|
| 175 |
)
|