pytorch compile

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)
+```

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"]