Skip to content

security: prefer safetensors over torch dumps#204

Closed
coldwaterq wants to merge 2 commits intonikopueringer:mainfrom
coldwaterq:safetensors-as-alternative
Closed

security: prefer safetensors over torch dumps#204
coldwaterq wants to merge 2 commits intonikopueringer:mainfrom
coldwaterq:safetensors-as-alternative

Conversation

@coldwaterq
Copy link
Copy Markdown

What does this change?

Add a conversion script and update the loading to make it easy to swap to safetensors without breaking backwards compatability. Safetesnors ensures that pickles aren't used to save the models at all, since pytorch keeps having security advisories this ensures that you don't have to worry if you are running the latest version of pytorch and know that even if the model is corrupted, an attacker can't get a shell on your system.

How was it tested?

The conversion script was tested on the model in huggingface and then that converted file was loaded to ensure the safetensor load calls are setup correctly and nothing the model requires was lost.

ran uv run pytest

Checklist

  • [ X ] uv run pytest passes
  • [ X ] uv run ruff check passes
  • [ X ] uv run ruff format --check passes

@coldwaterq
Copy link
Copy Markdown
Author

It's not related to this change so I didn't include it in the PR, but I also believe that

HF_CHECKPOINT_FILENAME = "CorridorKey.pth"

is supposed to be

HF_CHECKPOINT_FILENAME = "CorridorKey_v1.0.pth"

@nikopueringer
Copy link
Copy Markdown
Owner

Seems like a prudent change. @HYP3R00T What are your thoughts? There are a fair amount of new lines of code. Do you think we should merge it?

@nikopueringer
Copy link
Copy Markdown
Owner

Discussed this, and it seems like it might be better for me to convert the official file to safetensors and avoid .pth altogether! Going to look into that.

@coldwaterq
Copy link
Copy Markdown
Author

Yeah, if that's easier for you it would also be better. The conversion complicates this merge but it was just an attempt to make it easier for you.

@nikopueringer
Copy link
Copy Markdown
Owner

@coldwaterq Hey there! Wanted to let you know that I am now hosting a .safetensors version of CorridorKey on HF, alongside the old .pth for now. I also merged a PR to import the .safetensor file by default, along with some old legacy fallback to .pth for the people still running the old model, and a conversion script as well. Much of that was directly informed by what you did with this PR. You laid out the groundwork for this, so thank you for doing so! Hopefully this will be a big step in getting CK to be more secure.

Since I believe we have effectively duplicated the actions of this PR, I'm going to go ahead and close it, but please let me know if I am doing this in error! Thank you again for your help!

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.

2 participants