Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add BenchmarkFamily iterable #199

Merged
merged 4 commits into from
Jan 20, 2025
Merged

Add BenchmarkFamily iterable #199

merged 4 commits into from
Jan 20, 2025

Conversation

nicholasjng
Copy link
Collaborator

@nicholasjng nicholasjng commented Jan 15, 2025

Designed to consume the arguments of parametrize/product, and yield benchmarks lazily.

Params needs to support lazy iterables as well, to not fall into the eager materialization + memory blowup trap.


Testing should contain these scenarios at minimum:

  1. Confirm that BenchmarkFamilys are being picked up by nnbench.collect(), similarly to existing collect checks.
  2. Add a test that would fail with an eager iterable (e.g. list[Benchmark like it is now), asserting memory consumption stays low e.g. with memray.

Proposal for 2):

Add a benchmark matmul(a: np.ndarray, b: np.ndarray) with two arrays $a, b \in \mathbb{R}^{N \times N}$, parametrize over one of them (five np.random.randns should suffice), assert that memory consumption stays near 3 * N^2 * sizeof(np.float) (a, b, and the intermediary for the result).

Designed to consume the arguments of `parametrize`/`product`, and yield
benchmarks lazily.

Params needs to support lazy iterables as well, to not fall into the eager
materialization + memory blowup trap.
@nicholasjng
Copy link
Collaborator Author

Sweet. Results are horrible:

================================================================================= MEMRAY REPORT ==================================================================================
Allocation results for tests/integration/test_benchmark_family_memory_consumption.py::test_foobar at the high watermark

         📦 Total memory allocated: 145.0MiB
         📏 Total allocations: 96
         📊 Histogram of allocation sizes: | █▄  |
         🥇 Biggest allocating functions:
                - matmul:/Users/nicholasjunge/Workspaces/python/nnbench/tests/integration/test_benchmark_family_memory_consumption.py:24 -> 68.1MiB
                - <genexpr>:/Users/nicholasjunge/Workspaces/python/nnbench/tests/integration/test_benchmark_family_memory_consumption.py:22 -> 64.0MiB
                - _call_with_frames_removed:<frozen importlib._bootstrap>:241 -> 267.5KiB

At an allowance of 25MiB, this puts us at 6x the theoretically necessary memory consumption.

This is an easy way to assert that stale parameters are dealloc'ed once benchmarks
are over.
This means lazy evaluation of inputs and construction of benchmarks.
In the case of parametrization, we get optimal memory usage.
For products, we load all the iterators into memory, since the cartesian product
of iterables cannot be evaluated without eager consumption of all inputs.
Since we don't return lists anymore, we cannot check length or use __getitem__ calls.

Also, name checks are suspended for the time being, since the user should specify their
own names without us interfering / berating them for it.
@nicholasjng
Copy link
Collaborator Author

Merging, docs updates to follow in a separate PR.

@nicholasjng nicholasjng merged commit 0b6f9e6 into main Jan 20, 2025
14 checks passed
@nicholasjng nicholasjng deleted the benchmark-iterators branch January 20, 2025 14:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant