Repository navigation
Conversation
`StreamingSafetensorSaver` writes every shard under a temporary `model-NNNNN-of-?????.safetensors` name and only renames them in `finalize()`. Nothing removed the files an earlier export left in the same directory, so a second export that needs fewer shards inherited the old ones: * a 2-shard export followed by a single-file export kept `model.safetensors.index.json`, and `SafetensorLazyLoader` prefers the index over `model.safetensors`, so reloading the directory returned the previous weights under the same keys; * a 4-shard export followed by a 2-shard export kept the old shard files; * `os.rename` also raises `FileExistsError` when the target shard of the previous export is still there, which aborted the save. Record the model files present when the saver starts and delete the ones the new export does not produce, before the renames so the targets are free. `os.replace` keeps the rename working if a file appears in between. PEFT exports (`adapter_model.safetensors`) keep their files untouched. Tests cover the single-file case, the fewer-shards case, an export repeated twice and the PEFT path; the first three fail on the previous code. Signed-off-by: li-lizhe <147392333@qq.com>
This file contains hidden or 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
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
StreamingSafetensorSavernow removes the files of an earlier export in the same directory when they are not part of the new one, and renames shards withos.replace.Fixes #257.
Why
finalize()writes the shards it just produced and never looks at what a previous export left behind, so re-saving into the same output directory kept stale files:model.safetensors.index.json, andSafetensorLazyLoaderreads the index in preference tomodel.safetensors- reloading the directory returned the old tensors under the same keys, with no error;os.renameraisedFileExistsErrorwhen the target name of the previous export still existed, which aborted the save.How
__init__records the model files present when the saver starts (model*.safetensorsplusmodel.safetensors.index.json).finalize()knows the exact set the new export produces, deletes the difference before the renames so the rename targets are free, and usesos.replacefor the rename itself. PEFT exports are untouched: they writeadapter_model.safetensorsand leave the model files alone.Tests
tests/test_safetensors_shard_cleanup.pycovers four cases:Run against the module source with a stub
mcore_bridge.utils.env(the package__init__needs megatron and peft), so the sharding and cleanup logic under test is the real code.flake8with the repo'ssetup.cfgselection,isortandyapfare clean on both files.One caveat on the environment: the shard filename contains a literal
?????placeholder, which cannot be created on Windows, so the multi-shard cases were exercised with that placeholder rewritten to a legal token in the module copy - the rename and cleanup paths themselves are unchanged. Fullpyteston Linux was not run;pre-commitand CI were not run either.