Parameter representations instead of parameters in benchmark records

[nicholasjng]

TL,DR: Retaining references to parameters in benchmark records prevents garbage collection and wastes memory - how can we do better?

103 introduced saving the parameters to the records. This is fine for standard Python types, but wasteful for models and datasets which have a large memory footprint. In the worst (and unfortunately common) case, garbage collection is inhibited since the reference counts of models and data that are not needed anymore never drop to zero.

There are a few ideas here:

1. Rolling without the parameters. This was the case until #103, but means that we have no straightforward way to find out what parameters were used in a run.
2. Saving a unique representation of the parameters instead of the parameters. This can a no-op for standard types, and turn large models and data into a small struct (e.g. by saving file path(s), hash(es)). This requires knowing how to translate custom types into representations, and probably filtering out parameters that are too large / not trivially representable. This could also take a parameter schema.
3. Saving some kind of stamp that uniquely represents the parameters (potentially not even bidirectional, e.g. a hash function).

I'm leaning towards 2), but if the serialization is too difficult, I prefer dropping the parameters again.

[nicholasjng]

Let's discuss object lifetimes here without parameters in the record, assuming no unknown outside references.

• Case 1: I put a model (large size, say 16GB) into my params, which I pass into runner.run(). The params are not bound to and thus not referenced by any benchmark, so the model can be reaped as soon as the last benchmark accessing model is run and the params go out of scope. This is the desired case. ✅
  • Case 2: I put a model into a parametrization (say @nnbench.parametrize(model=model)), which binds it to that benchmark instance's params. This is still zero-copy, and the references are destroyed with the benchmarks. But: Since benchmarks are module-level members, the references are not destroyed until either the benchmark function goes out of scope or the module containing the benchmarks is unloaded. ❌
• Case 3: I put the model in a parametrization, but use a memo/artifact to lazy-load it before the benchmark. This puts the model referenced by the memo into a cache managed by nnbench, and the memo is bound to the benchmarks instead of the actual model. After the model benchmarks are done, I fire a hook that clears the model from the cache, causing it to be garbage-collected (maybe even sprinkling in a gc.collect() to help). This is the desired case for at-rest parametrizations like @nnbench.parametrize/@nnbench.product. ✅

So it seems either way that the best way to get rid of parameters is to supply models by hand, and if you want to parametrize, doing it with memoization.

[nicholasjng]

124 contains a built-in transform for dealing with compressed parameter representations. For the current memoization approach, this is good enough, since params are still freed, and memos do not themselves hold references to their values.

Once the global cache + eviction API is in, we'll run a memory profiler on a multi-model benchmark and see what happens if we evict by hand after the completion of a model benchmark family.

TL,DR: Blocked by #125, revisit afterwards.

[Maciej818]

Addressed by #124 and #120 . We close the ticket.