Skip to content

Commit

Permalink
[Mosaic GPU] Use events as the default profiling method
Browse files Browse the repository at this point in the history
JAX still supports old CUDA versions such as 12.0, where CUPTI leaks memory.

PiperOrigin-RevId: 705459909
  • Loading branch information
apaszke authored and Google-ML-Automation committed Dec 12, 2024
1 parent dda6b88 commit 3630756
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion jax/experimental/mosaic/gpu/profiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ def wrapper(*args, **kwargs):
return wrapper


def measure(f: Callable, *, mode: str = "cupti", aggregate: bool = True
def measure(f: Callable, *, mode: str = "events", aggregate: bool = True
) -> Callable:
"""Sets up a function ``f`` for profiling on GPU.
Expand Down

0 comments on commit 3630756

Please sign in to comment.