Add safetensor option when saving and restoring models #11549
+1,134
−62
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: stevehuang52 <[email protected]>
stevehuang52
commented
Dec 11, 2024
| @@ -508,6 +511,7 @@ def restore_from( | |||
| override_config_path: Optional[str] = None, | |||
| map_location: Optional[torch.device] = None, | |||
| strict: bool = False, | |||
| safe: bool = False, | |||
There was a problem hiding this comment.
@tango4j this won't work for diarization models, since the vad and speaker models are restored in the _init_vad_model and _init_speaker_model methods that don't take safe as additional param. Do you think we should fix it or maybe plan in the future?
There was a problem hiding this comment.
My opinion is to leave it for the future. Enabling safe tensors unblocks an important use case for us.
stevehuang52
commented
Dec 11, 2024
| self.msdd_model.save_to(neural_diar_model) | ||
| self.clus_diar._vad_model.save_to(vad_model, safe=safe) | ||
| self.clus_diar._speaker_model.save_to(spkr_model, safe=safe) | ||
| self.msdd_model.save_to(neural_diar_model, safe=safe) |
There was a problem hiding this comment.
@tango4j Similar issue as ClusterDiarizer, here the _init_msdd_model doesn't take safe param
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| @@ -1010,7 +1010,7 @@ | |||
| ) | |||
| super().__init__() | |||
|
|
|||
| def save_to(self, model, save_path: str): | |||
| def save_to(self, model, save_path: str, safe: bool = False): | |||
Check notice
Code scanning / CodeQL
Explicit returns mixed with implicit (fall through) returns Note
| save_path = save_path.replace(".nemo", "_XYZ.nemo") | ||
| super().save_to(model, save_path, safe=safe) | ||
|
|
||
| class MockModelV2(MockModel): |
Check notice
Code scanning / CodeQL
Unused local variable Note test
Signed-off-by: stevehuang52 <[email protected]>
Signed-off-by: stevehuang52 <[email protected]>
…nto heh/add_safetensor
Signed-off-by: stevehuang52 <[email protected]>
|
I just wanted to chime in that if there is any way to get this pull request merged into the main branch and a release branch that would be hugely helpful. I can hobble along using my clone of this branch, but that's obviously not ideal. Thank you guys for you great work and for this really helpful contribution in particular! |
|
This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. |
|
This PR was closed because it has been inactive for 7 days since being marked as stale. |
|
Reopened. It got stale over the holidays. |
galv
reviewed
Jan 9, 2025
| if safe: | ||
| raise e | ||
| else: | ||
| logging.info(e) |
There was a problem hiding this comment.
My opinion is that this is going to create a lot of confusingly noise. Basically every time that someone does an unsafe load on a model that wasn't saved with the safe=True option (which is frequent because it's the default), they will get an unhelpful log message of a file not found error, even though the application is otherwise working correctly.
galv
reviewed
Jan 9, 2025
| None | ||
| """ | ||
| try: | ||
| storch.save_file(tensors, filename + SAFE_EXTENSION) |
There was a problem hiding this comment.
Why use the safetensors format if the user does not request it? You will basically double the size of the model on the disk this way. This does not seem like a good idea to me.
There was a problem hiding this comment.
The old PR enables safetensor by default, but yes I agree that it should be disabled by default, and only enable when required.
|
I just wanted to say I fully endorse this pull request. I tested it and it works well for my use case. It will greatly enable my work if some form of this safetensor PR can make it into the main branch. Thank you!!! |
|
beep boop 🤖: 🙏 The following files have warnings. In case you are familiar with these, please try helping us to improve the code base. Your code was analyzed with PyLint. The following annotations have been identified: Mitigation guide:
By applying these rules, we reduce the occurance of this message in future. Thank you for improving NeMo's documentation! |
|
This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Re-do changes in a previously closed PR #7812 to add safetensor options when saving and restoring models. as customers requested.
CC @galv
Collection: [core,asr,nlp,common]
Changelog
Usage
Following usage is copied from #7812
When using nemo with untrusted .nemo files this will greatly reduce the potential for the user to be attacked.
PR Type: