From c145e6296e8d25fba3cfb2b9f4fb6adf30e34204 Mon Sep 17 00:00:00 2001 From: steveoni Date: Tue, 20 Jan 2026 16:25:59 +0100 Subject: [PATCH] pytorch compile initial work --- .../pytorch_compile/README.md | 162 +++++++++++ .../pytorch_compile/__init__.py | 6 + .../pytorch_compile/benchmark.py | 258 ++++++++++++++++++ .../pytorch_compile/compiled_clip.py | 108 ++++++++ .../verify_torch_compile_gpu.py | 155 +++++++++++ 5 files changed, 689 insertions(+) create mode 100644 3_cuda_triton_mini_clip/pytorch_compile/README.md create mode 100644 3_cuda_triton_mini_clip/pytorch_compile/__init__.py create mode 100644 3_cuda_triton_mini_clip/pytorch_compile/benchmark.py create mode 100644 3_cuda_triton_mini_clip/pytorch_compile/compiled_clip.py create mode 100644 3_cuda_triton_mini_clip/pytorch_compile/verify_torch_compile_gpu.py diff --git a/3_cuda_triton_mini_clip/pytorch_compile/README.md b/3_cuda_triton_mini_clip/pytorch_compile/README.md new file mode 100644 index 0000000..d45b99c --- /dev/null +++ b/3_cuda_triton_mini_clip/pytorch_compile/README.md @@ -0,0 +1,162 @@ +# PyTorch Compile for MiniCLIP + +This module provides `torch.compile` wrappers for MiniCLIP, enabling JIT compilation and benchmarking of different compilation modes. + +## Overview + +`torch.compile` optimizes PyTorch models by tracing and compiling them into optimized kernels. This module provides: + +1. **CompiledMiniCLIP**: A wrapper that applies `torch.compile` to the vision and text encoders +2. **Benchmarking**: Compare different compile modes (default, reduce-overhead, max-autotune) + +## Graph Break Analysis + +The MiniCLIP model was analyzed for potential graph breaks: + +- **No `torch.cond` required**: The model has no data-dependent control flow +- The `if self.eos_id is not None` in `TextTransformer` is config-based (not tensor-based), so Dynamo handles it via guard specialization +- All tensor operations are standard PyTorch ops that compile cleanly + +## Compile Modes + +| Mode | Description | Use Case | +|------|-------------|----------| +| `default` | Balanced performance and compile overhead | General purpose | +| `reduce-overhead` | Uses CUDA graphs for lower latency | Small batches, inference | +| `max-autotune` | Profiles Triton kernels for best performance | Maximum throughput | + +### fullgraph=True + +Setting `fullgraph=True` requires the entire function to compile into a single graph. This: +- Raises an error if there are graph breaks +- Can provide better optimization opportunities +- Recommended for production deployment after verification + +## Usage + +### Basic Usage + +```python +from mini_clip.model.clip import build_vit_b16_clip +from pytorch_compile import compile_clip + +# Build model +model = build_vit_b16_clip().cuda().eval() + +# Compile with default mode +compiled_model = compile_clip(model, mode="default") + +# Use as normal +images = torch.randn(32, 3, 224, 224).cuda() +text = torch.randint(0, 49408, (32, 77)).cuda() +image_feats, text_feats = compiled_model(images, text) +``` + +### With reduce-overhead for Low Latency + +```python +compiled_model = compile_clip( + model, + mode="reduce-overhead", # Uses CUDA graphs + fullgraph=True, # Require full graph capture +) + +# Warmup (important for CUDA graphs) +for _ in range(3): + _ = compiled_model(images, text) + +# Now runs with minimal Python overhead +image_feats, text_feats = compiled_model(images, text) +``` + +### Max Autotune for Throughput + +```python +compiled_model = compile_clip( + model, + mode="max-autotune", # Profiles kernel configurations +) +``` + +## Benchmarking + +Run the benchmark to compare modes: + +```bash +# From project root +python -m pytorch_compile.benchmark --batch-size 32 --device cuda + +# Options +python -m pytorch_compile.benchmark \ + --batch-size 64 \ + --image-size 224 \ + --warmup 10 \ + --iters 100 \ + --device cuda +``` + +### Expected Output + +``` +================================================================================ +Mode Mean (ms) Std (ms) Speedup Throughput +================================================================================ +eager 15.234 0.421 1.00x 2100.5/s +default 8.123 0.312 1.88x 3940.2/s +reduce-overhead 5.456 0.089 2.79x 5866.3/s +max-autotune 6.234 0.156 2.44x 5134.7/s +default (fullgraph) 7.890 0.298 1.93x 4056.4/s +reduce-overhead (fullgraph) 5.123 0.076 2.97x 6247.8/s +================================================================================ +``` + +## API Reference + +### compile_clip + +```python +def compile_clip( + model: nn.Module, + mode: str = "default", + fullgraph: bool = False, + dynamic: bool | None = None, +) -> CompiledMiniCLIP +``` + +**Parameters:** +- `model`: MiniCLIP model instance +- `mode`: Compile mode ("default", "reduce-overhead", "max-autotune") +- `fullgraph`: Require full graph capture without breaks +- `dynamic`: Use dynamic shapes (None = auto-detect) + +### CompiledMiniCLIP + +Wrapper class with the same interface as MiniCLIP: +- `forward(images, text)` → `(image_feats, text_feats)` +- `encode_image(images)` → `image_feats` +- `encode_text(text)` → `text_feats` + +## Best Practices + +1. **Warmup**: Always run a few warmup iterations before benchmarking +2. **CUDA Graphs**: `reduce-overhead` mode works best with fixed input shapes +3. **Dynamic Shapes**: Use `dynamic=True` if batch sizes vary significantly +4. **Profiling**: Use `TORCH_LOGS=guards` to debug guard failures +5. **Memory**: `reduce-overhead` caches workspace memory, increasing memory usage + +## Debugging + +Enable compile logging: + +```python +import torch._logging +torch._logging.set_logs(graph_code=True) # See traced graphs +torch._logging.set_logs(graph_breaks=True) # See graph break locations +``` + +Check for graph breaks: + +```python +# This will error if there are graph breaks +compiled = compile_clip(model, fullgraph=True) +``` diff --git a/3_cuda_triton_mini_clip/pytorch_compile/__init__.py b/3_cuda_triton_mini_clip/pytorch_compile/__init__.py new file mode 100644 index 0000000..e4a1c1e --- /dev/null +++ b/3_cuda_triton_mini_clip/pytorch_compile/__init__.py @@ -0,0 +1,6 @@ +"""PyTorch compile wrappers and benchmarks for MiniCLIP.""" + +from .compiled_clip import CompiledMiniCLIP, compile_clip +from .benchmark import run_benchmark + +__all__ = ["CompiledMiniCLIP", "compile_clip", "run_benchmark"] diff --git a/3_cuda_triton_mini_clip/pytorch_compile/benchmark.py b/3_cuda_triton_mini_clip/pytorch_compile/benchmark.py new file mode 100644 index 0000000..18472ec --- /dev/null +++ b/3_cuda_triton_mini_clip/pytorch_compile/benchmark.py @@ -0,0 +1,258 @@ +""" +Benchmark torch.compile modes for MiniCLIP. + +Compares: +- Eager (baseline) +- default mode +- reduce-overhead mode (CUDA graphs) +- max-autotune mode +- fullgraph=True variants + +Usage: + python -m pytorch_compile.benchmark +""" + +import argparse +import gc +import sys +import time +from pathlib import Path + +# Add mini_clip package to path +# benchmark.py -> pytorch_compile -> 3_cuda_triton_mini_clip -> mini-clip (project root) +_project_root = Path(__file__).resolve().parents[2] +_mini_clip_src = _project_root / "1_mini_clip" / "src" +if _mini_clip_src.exists() and str(_mini_clip_src) not in sys.path: + sys.path.insert(0, str(_mini_clip_src)) +from dataclasses import dataclass +from typing import Callable, Optional + +import torch +import torch.nn as nn + + +@dataclass +class BenchmarkResult: + """Stores benchmark timing results.""" + + mode: str + warmup_time_ms: float + mean_time_ms: float + std_time_ms: float + throughput: float # samples/sec + + +def sync_cuda(): + """Synchronize CUDA if available.""" + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def timed_run(fn: Callable, use_cuda_events: bool = True) -> float: + """ + Time a function execution. + + Args: + fn: Function to time + use_cuda_events: Use CUDA events for GPU timing (more accurate) + + Returns: + Elapsed time in milliseconds + """ + if use_cuda_events and torch.cuda.is_available(): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) + else: + sync_cuda() + start = time.perf_counter() + fn() + sync_cuda() + return (time.perf_counter() - start) * 1000 + + +def benchmark_model( + model: nn.Module, + images: torch.Tensor, + text: torch.Tensor, + warmup_iters: int = 10, + bench_iters: int = 100, + mode_name: str = "unknown", +) -> BenchmarkResult: + """ + Benchmark a model's forward pass. + + Args: + model: Model to benchmark + images: Image input tensor + text: Text input tensor + warmup_iters: Number of warmup iterations + bench_iters: Number of benchmark iterations + mode_name: Name for result reporting + + Returns: + BenchmarkResult with timing statistics + """ + # Warmup + warmup_start = time.perf_counter() + for _ in range(warmup_iters): + with torch.no_grad(): + _ = model(images, text) + sync_cuda() + warmup_time = (time.perf_counter() - warmup_start) * 1000 + + # Benchmark + times = [] + for _ in range(bench_iters): + t = timed_run(lambda: model(images, text)) + times.append(t) + + times_tensor = torch.tensor(times) + mean_time = times_tensor.mean().item() + std_time = times_tensor.std().item() + batch_size = images.shape[0] + throughput = (batch_size / mean_time) * 1000 # samples/sec + + return BenchmarkResult( + mode=mode_name, + warmup_time_ms=warmup_time, + mean_time_ms=mean_time, + std_time_ms=std_time, + throughput=throughput, + ) + + +def print_results(results: list[BenchmarkResult], baseline_name: str = "eager"): + """Print benchmark results in a table format.""" + baseline = next((r for r in results if r.mode == baseline_name), results[0]) + + print("\n" + "=" * 80) + print(f"{'Mode':<30} {'Mean (ms)':<12} {'Std (ms)':<12} {'Speedup':<10} {'Throughput':<15}") + print("=" * 80) + + for r in results: + speedup = baseline.mean_time_ms / r.mean_time_ms + print( + f"{r.mode:<30} {r.mean_time_ms:>10.3f} {r.std_time_ms:>10.3f} " + f"{speedup:>8.2f}x {r.throughput:>12.1f}/s" + ) + + print("=" * 80) + + +def run_benchmark( + batch_size: int = 32, + image_size: int = 224, + warmup_iters: int = 10, + bench_iters: int = 100, + device: str = "cuda", + include_fullgraph: bool = True, +) -> list[BenchmarkResult]: + """ + Run benchmarks comparing torch.compile modes. + + Args: + batch_size: Batch size for inputs + image_size: Image resolution + warmup_iters: Warmup iterations per mode + bench_iters: Benchmark iterations per mode + device: Device to run on + include_fullgraph: Include fullgraph=True variants + + Returns: + List of BenchmarkResult for each mode + """ + # Lazy import to avoid circular dependency + from mini_clip.model.clip import build_vit_b16_clip + + from .compiled_clip import compile_clip + + if device == "cuda" and not torch.cuda.is_available(): + print("CUDA not available, falling back to CPU") + device = "cpu" + + print(f"Device: {device}") + print(f"Batch size: {batch_size}") + print(f"Image size: {image_size}x{image_size}") + print(f"Warmup: {warmup_iters}, Bench: {bench_iters}") + + # Create inputs + images = torch.randn(batch_size, 3, image_size, image_size, device=device) + text = torch.randint(0, 49408, (batch_size, 77), device=device) + + results = [] + + # Eager baseline + print("\nBenchmarking: eager...") + model = build_vit_b16_clip().to(device).eval() + result = benchmark_model(model, images, text, warmup_iters, bench_iters, "eager") + results.append(result) + del model + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + + # Compile modes to test + modes = [ + ("default", False), + ("reduce-overhead", False), + ("max-autotune", False), + ] + + if include_fullgraph: + modes.extend([ + ("default", True), + ("reduce-overhead", True), + ]) + + for mode, fullgraph in modes: + mode_name = f"{mode}" + (" (fullgraph)" if fullgraph else "") + print(f"\nBenchmarking: {mode_name}...") + + torch._dynamo.reset() + + base_model = build_vit_b16_clip().to(device).eval() + compiled_model = compile_clip(base_model, mode=mode, fullgraph=fullgraph) + + result = benchmark_model( + compiled_model, images, text, warmup_iters, bench_iters, mode_name + ) + results.append(result) + + del base_model, compiled_model + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + + return results + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark torch.compile modes for MiniCLIP") + parser.add_argument("--batch-size", type=int, default=32, help="Batch size") + parser.add_argument("--image-size", type=int, default=224, help="Image size") + parser.add_argument("--warmup", type=int, default=10, help="Warmup iterations") + parser.add_argument("--iters", type=int, default=100, help="Benchmark iterations") + parser.add_argument("--device", type=str, default="cuda", help="Device (cuda/cpu)") + parser.add_argument("--no-fullgraph", action="store_true", help="Skip fullgraph variants") + + args = parser.parse_args() + + results = run_benchmark( + batch_size=args.batch_size, + image_size=args.image_size, + warmup_iters=args.warmup, + bench_iters=args.iters, + device=args.device, + include_fullgraph=not args.no_fullgraph, + ) + + print_results(results) + + +if __name__ == "__main__": + main() diff --git a/3_cuda_triton_mini_clip/pytorch_compile/compiled_clip.py b/3_cuda_triton_mini_clip/pytorch_compile/compiled_clip.py new file mode 100644 index 0000000..d5305d3 --- /dev/null +++ b/3_cuda_triton_mini_clip/pytorch_compile/compiled_clip.py @@ -0,0 +1,108 @@ +""" +torch.compile wrappers for MiniCLIP. + +Provides compiled versions of the CLIP model with configurable compile modes: +- default: balanced performance/overhead +- reduce-overhead: uses CUDA graphs for lower latency (good for small batches) +- max-autotune: uses Triton autotuning for maximum throughput + +The MiniCLIP model has no data-dependent control flow requiring torch.cond. +The eos_id conditional in TextTransformer is config-based, not tensor-based, +so Dynamo handles it via guard specialization without graph breaks. +""" + +from typing import Literal, Optional + +import torch +import torch.nn as nn + +CompileMode = Literal["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"] + + +class CompiledMiniCLIP(nn.Module): + """ + Wrapper that applies torch.compile to MiniCLIP encoders. + + Args: + model: MiniCLIP model instance + mode: Compile mode - "default", "reduce-overhead", "max-autotune" + fullgraph: If True, require full graph capture (no graph breaks) + dynamic: If True, use dynamic shapes to reduce recompilations + + Note: + The visual and textual encoders are compiled separately to allow + independent optimization. The projection heads and logit_scale + remain in eager mode as they are simple linear ops. + """ + + def __init__( + self, + model: nn.Module, + mode: CompileMode = "default", + fullgraph: bool = False, + dynamic: Optional[bool] = None, + ): + super().__init__() + + self.projection_dim = model.projection_dim + self.logit_scale = model.logit_scale + self.image_proj = model.image_proj + self.text_proj = model.text_proj + + # Compile encoders with specified mode + self.visual = torch.compile( + model.visual, + mode=mode, + fullgraph=fullgraph, + dynamic=dynamic, + ) + self.textual = torch.compile( + model.textual, + mode=mode, + fullgraph=fullgraph, + dynamic=dynamic, + ) + + def encode_image(self, images: torch.Tensor) -> torch.Tensor: + feats = self.visual(images) + feats = self.image_proj(feats) + return torch.nn.functional.normalize(feats, dim=-1) + + def encode_text(self, text: torch.Tensor) -> torch.Tensor: + feats = self.textual(text) + feats = self.text_proj(feats) + return torch.nn.functional.normalize(feats, dim=-1) + + def forward( + self, images: torch.Tensor, text: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.encode_image(images), self.encode_text(text) + + +def compile_clip( + model: nn.Module, + mode: CompileMode = "default", + fullgraph: bool = False, + dynamic: Optional[bool] = None, +) -> CompiledMiniCLIP: + """ + Compile a MiniCLIP model with torch.compile. + + Args: + model: MiniCLIP model instance + mode: One of: + - "default": balanced performance and compile overhead + - "reduce-overhead": uses CUDA graphs, good for small batches + - "max-autotune": profiles kernels for best performance + fullgraph: Require full graph capture without breaks + dynamic: Use dynamic shapes (None = auto-detect) + + Returns: + CompiledMiniCLIP wrapper with compiled encoders + """ + return CompiledMiniCLIP( + model=model, + mode=mode, + fullgraph=fullgraph, + dynamic=dynamic, + ) diff --git a/3_cuda_triton_mini_clip/pytorch_compile/verify_torch_compile_gpu.py b/3_cuda_triton_mini_clip/pytorch_compile/verify_torch_compile_gpu.py new file mode 100644 index 0000000..9e3003f --- /dev/null +++ b/3_cuda_triton_mini_clip/pytorch_compile/verify_torch_compile_gpu.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +""" +verify_torch_compile_gpu.py + +Simple script to check whether torch.compile provides a measurable speedup +on the current NVIDIA GPU for both evaluation and training paths. + +Reference: +https://docs.pytorch.org/tutorials/intermediate/torch_compile_full_example.html +""" +import argparse +import sys + +import numpy as np +import torch +from torchvision.models import densenet121 + + +def timed(fn): + """Run fn() and return (result, elapsed_seconds). + + Uses CUDA events and synchronization for accurate GPU timing. + """ + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = fn() + end.record() + torch.cuda.synchronize() + return result, start.elapsed_time(end) / 1000.0 + + +def generate_data(b): + """Generate random inputs and integer targets on CUDA. + + b: batch size + """ + return ( + torch.randn(b, 3, 128, 128, device="cuda"), + torch.randint(0, 1000, (b,), device="cuda"), + ) + + +N_ITERS = 10 + + +def init_model(): + model = densenet121().cuda() + return model + + +def run_eval_test(n_iters=10, batch=16): + model = init_model() + model_opt = init_model() + + model.eval() + model_opt.eval() + # compile the module (recommended approach) + model_opt.compile(mode="reduce-overhead") + + eager_times = [] + compile_times = [] + + for i in range(n_iters): + inp = generate_data(batch)[0] + with torch.no_grad(): + _, eager_time = timed(lambda: model(inp)) + _, compile_time = timed(lambda: model_opt(inp)) + + eager_times.append(eager_time) + compile_times.append(compile_time) + print(f"eval iter {i}: eager {eager_time:.6f}s, compile {compile_time:.6f}s") + + eager_med = float(np.median(eager_times)) + compile_med = float(np.median(compile_times)) + speedup = eager_med / compile_med if compile_med > 0 else float("inf") + return eager_med, compile_med, speedup + + +def train_step(mod, data, opt): + opt.zero_grad(set_to_none=True) + pred = mod(data[0]) + loss = torch.nn.CrossEntropyLoss()(pred, data[1]) + loss.backward() + opt.step() + + +def run_train_test(n_iters=10, batch=16): + model = init_model() + opt = torch.optim.Adam(model.parameters()) + + eager_times = [] + for i in range(n_iters): + inp = generate_data(batch) + _, t = timed(lambda: train_step(model, inp, opt)) + eager_times.append(t) + print(f"train eager iter {i}: {t:.6f}s") + + # compile the training function (function-level compile) + model2 = init_model() + opt2 = torch.optim.Adam(model2.parameters()) + train_opt = torch.compile(train_step, mode="reduce-overhead") + + compile_times = [] + for i in range(n_iters): + inp = generate_data(batch) + _, t = timed(lambda: train_opt(model2, inp, opt2)) + compile_times.append(t) + print(f"train compile iter {i}: {t:.6f}s") + + eager_med = float(np.median(eager_times)) + compile_med = float(np.median(compile_times)) + speedup = eager_med / compile_med if compile_med > 0 else float("inf") + return eager_med, compile_med, speedup + + +def main(): + parser = argparse.ArgumentParser(description="Verify torch.compile speedup on current GPU") + parser.add_argument("--iters", type=int, default=N_ITERS, help="Number of iterations per test") + parser.add_argument("--batch", type=int, default=16, help="Batch size to use for tests") + args = parser.parse_args() + + if not torch.cuda.is_available(): + print("CUDA is not available on this machine. Exiting.") + sys.exit(2) + + prop = torch.cuda.get_device_properties(0) + try: + cc = f"{prop.major}.{prop.minor}" + except Exception: + cc = "unknown" + print(f"Device: {prop.name}, total_memory={prop.total_memory/1024**3:.2f} GB, compute_capability={cc}") + + print("\nRunning evaluation benchmark...") + eager_med, compile_med, speedup = run_eval_test(n_iters=args.iters, batch=args.batch) + print(f"(eval) eager median: {eager_med:.6f}s, compile median: {compile_med:.6f}s, speedup: {speedup:.2f}x") + if speedup <= 1.0: + print("Warning: torch.compile did not speed up eval (speedup <= 1).") + + print("\nRunning training benchmark...") + eager_med_t, compile_med_t, speedup_t = run_train_test(n_iters=args.iters, batch=args.batch) + print(f"(train) eager median: {eager_med_t:.6f}s, compile median: {compile_med_t:.6f}s, speedup: {speedup_t:.2f}x") + if speedup_t <= 1.0: + print("Warning: torch.compile did not speed up training (speedup <= 1).") + + if speedup <= 1.0 and speedup_t <= 1.0: + print("\nNo speedup observed for both eval and train. Exiting with non-zero code.") + sys.exit(1) + + print("\nAt least one of eval/train showed speedup with torch.compile. Success.") + sys.exit(0) + + +if __name__ == "__main__": + main()