[FEATURE] Support for a calibration dataset in PrunaAlgorithmBase

Is your feature request related to a problem? Please describe.

Currently, the PrunaAlgorithmBase has a property called dataset_required. Its intent is to "request a dataset for the algorithm".

The caveat is that when a model is saved, its smash config is saved as JSON, but the dataset is skipped. This is intentional and documented in code here. While this is not a problem for algorithms that use pickled, it is for algorithms with save_before_apply.

This is because when the model is loaded, the dataset_required property would be set to true, but there would be no data available. This will cause a crash and failure in loading the model.

Describe the solution you'd like

It would be great to have a flag that is compatible with save_before_apply, as pickled is not always a good choice, as it requires custom parsing.

Extra

Additional Context

The problem was discovered while designing the static-fp8 quantization algorithm, introduced in https://github.com/PrunaAI/pruna/pull/733.

Such an algorithm requires a dataset on smash for pre-compile calibration.

Here is the exact invocation chain that causes a problem.

1. This is the method which loads a model.
2. This is the default method used for "re-smashing"
3. Here the "re-smashing" invokes the same smash the user can use.
4. Here the smash method invokes the function which check that validates a dataset exists when dataset_required is truthy.

To circumvent this, the dataset_required is set to false. However, this necessitates documentation in code, and thus, it is not sustainable over time.

Furthermore, for convenience, as well as considering the fact that the dataset is not saved, calibration is not re-done on load. Instead, statistics necessary for quantization (e.g., the maximum magnitude observed) are stored as artifacts during save and loaded directly.

Proposed Solution

Considering this, I possible long-term solution might use two boolean flags.

• one indicating that a dataset is specifically for calibration, and
• another indicating that we are in "loading" phase, necessary to ensure runtime can make a difference between first time smash and post-save load. Alternatively, the validation might check whether the load function is set.

Keep in mind

With any solution, there are multiple caveats.

• We must ensure that algorithms relying on calibration (i.e., a calibration dataset) will save an artifact that is enough to restore the calibrated state on its own. Else, we either (1) should not support loading smashed models, (2) should find a way to store the dataset, or (3) should alert that user that the model will run in its original variant in case no artifacts are found.

Start with src/pruna/algorithms/base/pruna_base.py and trace the load and re-smash path through src/pruna/engine/load.py, src/pruna/smash.py, and src/pruna/config/pre_smash_routines.py. Clarify how calibration datasets differ from ordinary dataset requirements during initial smashing versus loading. Done means a calibration-dependent algorithm can save and reload using stored artifacts without requiring the unavailable dataset or rerunning calibration.