diff --git a/CHANGELOG.md b/CHANGELOG.md index d3587a5..01272c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,8 @@ This project follows semantic versioning once published to PyPI. - Added a Colab-only verified MUSHRA bundle export for the licensed LibriSpeech sinc/LavaSR matrix. - Added a bounded Colab-only MossFormer2 feasibility workflow covering minimal checkpoints, safe conversion, offline CPU inference, channel behavior, alignment, and parity-fixture constraints. +- Added a Colab T4 LavaSR eager/`torch.compile` benchmark and recorded the evidence-based decision + to keep eager as the only selectable provider. ## 0.6.0 - 2026-07-12 diff --git a/ROADMAP.md b/ROADMAP.md index 880e6b4..48b0564 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -102,13 +102,21 @@ and native 48 kHz inputs ran successfully; stereo channels were preserved. The s exposed a consistent 224-sample output shortening and cross-process floating-point drift, so any future backend needs explicit alignment plus tolerance-based parity rather than hash equality. -## Priority 4: Optimize Only After Measurement +## Completed: Optimize Only After Measurement Keep `lavasr-compat` on `torch-eager` until another provider has real-weight benchmarks and parity evidence. Evaluate `torch.compile` first because it changes the packaging surface less than ONNX, TensorRT, or OpenVINO. An optimized provider must retain deterministic CPU fallback and the same verified weight-store boundary. +Two fresh T4 measurements rejected promoting `torch.compile`: warm results contradicted each other +(`0.96x` and `1.82x`), cold inference consistently paid roughly 32 seconds of additional +compilation, peak CUDA allocation increased, and compiled repeats were not bit-exact. Numerical +eager/compile drift was small, but Inductor could not generate optimized code for the graph's +complex operators. `torch-eager` therefore remains the only selectable package-owned provider; +ONNX and hardware-specific providers remain deferred until a different graph strategy has parity +evidence and a credible, repeatable measured benefit. + ## Deferred - package-owned FlowHigh or Resemble Enhance implementations; diff --git a/docs/ACCELERATORS.md b/docs/ACCELERATORS.md index 64775d6..b85efc5 100644 --- a/docs/ACCELERATORS.md +++ b/docs/ACCELERATORS.md @@ -130,14 +130,20 @@ evidence, not a hard release gate, until stable baselines exist per device/provi ## LavaSR Optimization Recommendation -Current recommendation: keep `lavasr-compat` on PyTorch eager. +Current recommendation: keep `lavasr-compat` on PyTorch eager. A measured T4 `torch.compile` trial +did not justify a selectable provider. Evidence and risk: - The PyTorch eager path has real-weight download verification, torch smoke coverage, upstream parity coverage, and a fresh Colab T4 validation record. -- `torch.compile` may help some shapes, but it needs per-device measurement and output drift checks - before becoming a selectable provider. +- Two fresh-session runs on the same one-second 16 kHz fixture produced contradictory warm results: + `0.96x` (compile slightly slower) and `1.82x` (compile faster). This is not a stable speedup claim. + First inference consistently increased from about `2.3-2.4 s` to `34-35 s`, and peak CUDA + allocation consistently increased from 72.2 MB to 80.8 MB. +- Compile/eager drift was numerically small (maximum absolute error `5.22e-8`, RMS error + `1.06e-8`), but compiled repeats were not bit-exact. TorchInductor also warned that complex + operators are not supported for code generation and performance may be worse than eager. - ONNX export is not yet accepted because the current LavaSR graph includes STFT/ISTFT, mel projection, complex tensors, and `FastLRMerge` behavior that need parity fixtures before export. - TensorRT and OpenVINO should be separate graph-provider experiments after ONNX parity exists. @@ -151,6 +157,22 @@ Future acceptance for replacing or adding an optimized LavaSR provider requires: 3. A documented dependency and packaging cost. 4. A CPU fallback path that remains deterministic and offline. +Reproduce the rejected compile experiment in a fresh T4 session: + +```sh +colab new -s asr-lavasr-compile --gpu T4 +colab exec -s asr-lavasr-compile \ + -f examples/colab_lavasr_compile_benchmark.py --timeout 3600 +colab download -s asr-lavasr-compile \ + /content/asr-lavasr-compile-benchmark.tar.gz \ + runs/asr-lavasr-compile-benchmark.tar.gz +colab stop -s asr-lavasr-compile +``` + +The experiment compiles the already verified package-owned model internally; it deliberately does +not register `torch-compile` as a public provider. Repeat it across fresh sessions and representative +audio lengths before treating warm throughput as stable evidence. + ## Fresh GPU Validation Workflow Use a repository checkout when validating unreleased code or gated tests: diff --git a/examples/colab_lavasr_compile_benchmark.py b/examples/colab_lavasr_compile_benchmark.py new file mode 100644 index 0000000..610ad62 --- /dev/null +++ b/examples/colab_lavasr_compile_benchmark.py @@ -0,0 +1,229 @@ +"""Benchmark LavaSR torch eager versus torch.compile in a Colab T4 runtime.""" + +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import subprocess +import sys +import tarfile +import time +from datetime import datetime, timezone +from pathlib import Path + +REPO_URL = os.environ.get("ASR_REPO_URL", "https://github.com/Tinnci/python-audio-super-resolution.git") +GIT_REF = os.environ.get("ASR_GIT_REF", "agent/lavasr-torch-compile-evidence") +WORK_ROOT = Path(os.environ.get("ASR_COLAB_WORK_ROOT", "/content/lavasr-compile-benchmark")) +REPO_DIR = WORK_ROOT / "repo" +CACHE_DIR = WORK_ROOT / "models" +EVIDENCE_DIR = WORK_ROOT / "evidence" +ARCHIVE_PATH = Path(os.environ.get("ASR_COLAB_ARCHIVE", "/content/asr-lavasr-compile-benchmark.tar.gz")) + + +def run(name: str, command: list[str], *, cwd: Path | None = None) -> None: + log_path = EVIDENCE_DIR / "logs" / f"{name}.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + print(f"[compile-benchmark] {name}: {' '.join(command)}", flush=True) + with log_path.open("w", encoding="utf-8") as log: + subprocess.run(command, cwd=cwd, stdout=log, stderr=subprocess.STDOUT, text=True, check=True) + + +def sha256_array(array) -> str: + return hashlib.sha256(array.astype(" None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def benchmark_provider(provider: str, fixture, *, repeats: int = 3) -> tuple[dict[str, object], object]: + import numpy as np + import torch + + from audio_super_resolution import InferenceConfig + from audio_super_resolution.backends.lavasr_compat import LavaSRCompatBackend + + torch.manual_seed(0) + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + config = InferenceConfig( + device="cuda", + runtime_provider="torch-eager", + model_cache_dir=CACHE_DIR, + weights_manifest=CACHE_DIR / "lavasr-v2-bwe" / "manifest.json", + ) + started = time.perf_counter() + backend = LavaSRCompatBackend(config) + if provider == "torch-compile": + eager_loader = backend._load_model + compiled = False + + def load_compiled(*args, **kwargs): + nonlocal compiled + model = eager_loader(*args, **kwargs) + if not compiled: + model = torch.compile(model, fullgraph=False, dynamic=True) + backend._model = model + compiled = True + return model + + backend._load_model = load_compiled # type: ignore[method-assign] + synchronize(torch) + construction_seconds = time.perf_counter() - started + + started = time.perf_counter() + first = backend.enhance(fixture, 16_000, 48_000) + synchronize(torch) + first_seconds = time.perf_counter() - started + + warm_times = [] + warm_outputs = [] + for _ in range(repeats): + started = time.perf_counter() + output = backend.enhance(fixture, 16_000, 48_000) + synchronize(torch) + warm_times.append(time.perf_counter() - started) + warm_outputs.append(output) + + first_array = np.asarray(first, dtype=np.float32) + exact_repeat = all(np.array_equal(first_array, np.asarray(output)) for output in warm_outputs) + record = { + "provider": provider, + "construction_seconds": construction_seconds, + "first_seconds": first_seconds, + "warm_seconds": warm_times, + "warm_mean_seconds": float(np.mean(warm_times)), + "warm_median_seconds": float(np.median(warm_times)), + "output_shape": list(first_array.shape), + "output_sha256": sha256_array(first_array), + "same_provider_exact_repeat": exact_repeat, + "peak_cuda_memory_bytes": int(torch.cuda.max_memory_allocated()), + } + return record, first_array + + +def worker_main() -> int: + import numpy as np + import torch + + timeline = np.arange(16_000, dtype=np.float32) / 16_000 + fixture = (0.08 * np.sin(2 * np.pi * 440 * timeline) + 0.02 * np.sin(2 * np.pi * 3000 * timeline)).astype( + np.float32 + ) + providers = {} + outputs = {} + errors = {} + for provider in ("torch-eager", "torch-compile"): + try: + providers[provider], outputs[provider] = benchmark_provider(provider, fixture) + except Exception as exc: + errors[provider] = f"{type(exc).__name__}: {exc}" + + comparison = None + if len(outputs) == 2: + eager = outputs["torch-eager"] + compiled = outputs["torch-compile"] + delta = compiled.astype(np.float64) - eager.astype(np.float64) + comparison = { + "shape_equal": eager.shape == compiled.shape, + "exact_equal": bool(np.array_equal(eager, compiled)), + "max_abs_error": float(np.max(np.abs(delta))), + "rms_error": float(np.sqrt(np.mean(delta**2))), + "eager_rms": float(np.sqrt(np.mean(eager.astype(np.float64) ** 2))), + "warm_speedup": providers["torch-eager"]["warm_median_seconds"] + / providers["torch-compile"]["warm_median_seconds"], + } + result = { + "schema_version": 1, + "created_at": datetime.now(timezone.utc).isoformat(), + "device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, + "torch": torch.__version__, + "cuda": torch.version.cuda, + "fixture_sample_rate": 16_000, + "fixture_samples": fixture.size, + "fixture_sha256": sha256_array(fixture), + "providers": providers, + "comparison": comparison, + "errors": errors, + "passed": len(outputs) == 2, + } + (EVIDENCE_DIR / "benchmark.json").write_text(json.dumps(result, indent=2) + "\n") + return 0 if result["passed"] else 1 + + +def write_summary(status: str, error: str | None) -> None: + commit = None + if REPO_DIR.is_dir(): + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], cwd=REPO_DIR, capture_output=True, text=True, check=False + ).stdout.strip() + summary = { + "schema_version": 1, + "created_at": datetime.now(timezone.utc).isoformat(), + "status": status, + "repository": REPO_URL, + "requested_ref": GIT_REF, + "commit": commit or None, + "error": error, + } + (EVIDENCE_DIR / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + + +def archive() -> None: + with tarfile.open(ARCHIVE_PATH, "w:gz") as bundle: + bundle.add(EVIDENCE_DIR, arcname="evidence") + print(f"[compile-benchmark] evidence archive: {ARCHIVE_PATH}", flush=True) + + +def main() -> int: + if "--worker" in sys.argv: + return worker_main() + if Path.cwd() != Path("/content"): + raise RuntimeError("This benchmark must run inside a Colab runtime rooted at /content") + if WORK_ROOT.exists(): + shutil.rmtree(WORK_ROOT) + EVIDENCE_DIR.mkdir(parents=True) + try: + runner_source = Path(__file__).read_text(encoding="utf-8") + except NameError: + runner_source = get_ipython().history_manager.input_hist_raw[-1] # noqa: F821 + runner_path = WORK_ROOT / "colab_lavasr_compile_benchmark.py" + runner_path.write_text(runner_source, encoding="utf-8") + + error = None + try: + run("clone", ["git", "clone", "--filter=blob:none", REPO_URL, str(REPO_DIR)]) + run("checkout", ["git", "checkout", GIT_REF], cwd=REPO_DIR) + run("install-uv", [sys.executable, "-m", "pip", "install", "uv"]) + run("install-project", ["uv", "pip", "install", "--system", "-e", ".[lavasr,download]"], cwd=REPO_DIR) + run( + "prepare-model-cache", + [ + "audio-super-res", + "--backend", + "lavasr-compat", + "--model-cache-dir", + str(CACHE_DIR), + "--download-weights", + "--prepare-model-cache", + ], + cwd=REPO_DIR, + ) + run("benchmark", [sys.executable, str(runner_path), "--worker"], cwd=REPO_DIR) + write_summary("passed", None) + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + write_summary("failed", error) + finally: + archive() + if error: + print(f"[compile-benchmark] failed: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_lavasr_compat.py b/tests/test_lavasr_compat.py index a9bb108..80edfe3 100644 --- a/tests/test_lavasr_compat.py +++ b/tests/test_lavasr_compat.py @@ -45,7 +45,10 @@ def test_lavasr_backend_preserves_exact_digital_silence(monkeypatch, tmp_path) - "audio_super_resolution.backends.lavasr_validation.validate_lavasr_v2_weight_bundle", lambda resolved: bundle_info, ) - monkeypatch.setattr("audio_super_resolution.backends.lavasr_compat.resolve_runtime_provider", lambda *args: None) + monkeypatch.setattr( + "audio_super_resolution.backends.lavasr_compat.resolve_runtime_provider", + lambda *args: None, + ) monkeypatch.setattr("audio_super_resolution.backends.lavasr_compat.resolve_device", lambda *args, **kwargs: "cuda") monkeypatch.setattr("audio_super_resolution.backends.lavasr_compat._require_torch_runtime", lambda: None)