[FEATURE] Add Safetensors Support for Saving Smashed Models

[ParagEkbote]

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

The models which have been smashed with Pruna and pushed to HF Hub, they do not have explicit safetensors support. I think it would be useful to add this support to have secure model loading, prevent execution of malicious code with pickle (.pt) and faster throughput of model output. WDYT?

cc: @minettekaum

Describe the solution you'd like

Integrate Safetensors Support to save Smashed Model with Safetensors.

Additional context

I think that the support for safetensors could be added to the save.py within engine, feel free to let me know otherwise.

[begumcig]

Hi @ParagEkbote, I agree this is a valid safety concern, and we already have related discussions in #592 and PR #593.

That said, I have a few concerns about introducing this change:

• We currently use this pattern in both save and load across Pruna. Updating only one side would introduce inconsistencies in the codebase.
• Some architectures may rely entirely on pickling, so removing it outright could break those use cases.
• Finetuners and distillers (including LoRA saving) still depend on pickle-based serialization. To avoid regressions, we would likely need to refactor those parts in the same PR.

Given this, one possible approach would be to introduce the safer mechanism while keeping pickling as a fallback for compatibility.

Let me know what you think!

[ParagEkbote]

Hi @ParagEkbote, I agree this is a valid safety concern, and we already have related discussions in #592 and PR #593.

That said, I have a few concerns about introducing this change:

• We currently use this pattern in both save and load across Pruna. Updating only one side would introduce inconsistencies in the codebase.
• Some architectures may rely entirely on pickling, so removing it outright could break those use cases.
• Finetuners and distillers (including LoRA saving) still depend on pickle-based serialization. To avoid regressions, we would likely need to refactor those parts in the same PR.

Given this, one possible approach would be to introduce the safer mechanism while keeping pickling as a fallback for compatibility.

Let me know what you think!

I was not aware that pickling is also applied in loading of Pruna, thanks for letting me know 🙌 I think a safe approach can be to add the safetensors support for saving and loading model per algorithm class as introduced in v0.3.0 . So, each PR could target a specific class of algorithm like quantizers, cachers compilers, etc. Furthermore, we could use either SmolLM or Qwen 3 models to test whether the loading works without issues and to keep compute cost low. Does this approach seem feasible to you?

[begumcig]

Hi @ParagEkbote,

I think the idea of testing this per algorithm is a great way to manage risk and keep the PRs manageable!

One small technical note on the implementation: Based on our current architecture, the logic for saving doesn't actually live inside the individual algorithm classes. Instead, we use a centralized pattern where all logic resides in pruna/engine/save.py within the SAVE_FUNCTIONS registry. The algorithm classes simply declare which function to use (e.g., self.save_fn = SAVE_FUNCTIONS.pickled).

I think then the most probable way could be something like:

1 Implement a new function in save.py (perhaps something like SAVE_FUNCTIONS.safetensors_with_fallback) that attempts a safetensors save but falls back to pickling if the model architecture doesn't support it.

2 Do the same for load.

3 We can then go through the algorithm classes that use pickling one by one and switch to this new saving / loading. And As you mentioned, using SmolLM or Qwen for the initial tests would be perfect for keeping compute costs low.

Do you think this aligns with your idea?

[ParagEkbote]

Hi @ParagEkbote,

I think the idea of testing this per algorithm is a great way to manage risk and keep the PRs manageable!

One small technical note on the implementation: Based on our current architecture, the logic for saving doesn't actually live inside the individual algorithm classes. Instead, we use a centralized pattern where all logic resides in pruna/engine/save.py within the SAVE_FUNCTIONS registry. The algorithm classes simply declare which function to use (e.g., self.save_fn = SAVE_FUNCTIONS.pickled).

I think then the most probable way could be something like:

1 Implement a new function in save.py (perhaps something like SAVE_FUNCTIONS.safetensors_with_fallback) that attempts a safetensors save but falls back to pickling if the model architecture doesn't support it.

2 Do the same for load.

3 We can then go through the algorithm classes that use pickling one by one and switch to this new saving / loading. And As you mentioned, using SmolLM or Qwen for the initial tests would be perfect for keeping compute costs low.

Do you think this aligns with your idea?

Yeah, this does align with the idea I had in mind, which class of algorithm could we proceed with?