StreamingSafetensorSaver keeps the files of a previous export, so reloading returns the old weights

未关闭 适合新手
  • 哪里不对？
#257 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 1 天内回复

@li-lizhe 已经在做这个了。

开始于 2026年11月8日。

  • #258 来自 @li-lizhe —— 未关闭

评估

难度
2/5
预计耗时
1-3 小时
新手友好度
76/100
Issue 类型
缺陷
描述清晰度
描述清楚
活跃度
活跃
技术栈
python
领域
backend

调研方向

从 safetensors 工具中的 StreamingSafetensorSaver 及其 finalize() 方法入手；issue 指出了失败的保存行为，并给出了使用 SafetensorLazyLoader 的复现方法。检查在已有输出目录中如何完成单文件导出和分片导出。完成标准是：第二次导出后只留下它自己的文件，重新加载时返回的正好是新写入的张量。

由索引模型根据 Issue 内容生成。

描述

Bug Description / Bug 描述

StreamingSafetensorSaver writes each shard as model-NNNNN-of-?????.safetensors and renames them in finalize(). Nothing removes the files that an earlier export left in the same directory, so a second export inherits them:

  1. A sharded export followed by a single-file export keeps the stale index. finalize() writes model.safetensors but the model.safetensors.index.json from the previous run stays. SafetensorLazyLoader reads the index when it exists, so reloading the directory returns the previous tensors under the same keys, silently.
  2. A 4-shard export followed by a 2-shard export keeps the extra shard files. The index is rewritten, but model-0000{3,4}-of-00004.safetensors stay on disk.
  3. The rename can abort the save. finalize() uses os.rename, which raises FileExistsError when the target name of the previous export is still present, so re-saving into the same directory fails instead of replacing the checkpoint.

GPTBridge.save_weights(..., output_dir=...) takes an arbitrary directory, and re-running training into the same output directory is the normal way to refresh a merged checkpoint, so this is reachable from the documented API.

How to Reproduce / 如何复现
import tempfile
from pathlib import Path
import torch
from mcore_bridge.utils.safetensors import SafetensorLazyLoader, StreamingSafetensorSaver

d = Path(tempfile.mkdtemp())

# An earlier 3-shard export, as left behind by a previous run.
from safetensors.torch import save_file
import json
save_file({'w0': torch.zeros(3), 'w1': torch.zeros(3)}, str(d / 'model-00001-of-00003.safetensors'))
save_file({'w2': torch.zeros(3)}, str(d / 'model-00002-of-00003.safetensors'))
save_file({'w3': torch.zeros(3)}, str(d / 'model-00003-of-00003.safetensors'))
(d / 'model.safetensors.index.json').write_text(json.dumps({'metadata': {'total_size': 48}, 'weight_map': {
    'w0': 'model-00001-of-00003.safetensors', 'w1': 'model-00001-of-00003.safetensors',
    'w2': 'model-00002-of-00003.safetensors', 'w3': 'model-00003-of-00003.safetensors'}}))

# Re-export the model now that it fits in one file.
saver = StreamingSafetensorSaver(save_dir=str(d), max_shard_size='5GB')
saver.add_tensor('w0', torch.arange(4, dtype=torch.float32))
saver.finalize()

print(sorted(p.name for p in d.iterdir()))
# ['model-00001-of-00003.safetensors', 'model-00002-of-00003.safetensors',
#  'model-00003-of-00003.safetensors', 'model.safetensors', 'model.safetensors.index.json']

loaded = SafetensorLazyLoader(str(d)).get_state_dict()
print(sorted(loaded))                 # ['w0', 'w1', 'w2', 'w3'] - three tensors that are no longer saved
print(loaded['w0'].load().tolist())   # [0.0, 0.0, 0.0] - the old values, not [0, 1, 2, 3]
Expected behavior / 预期行为

The directory holds exactly the files of the export that just ran: model.safetensors here, no stale shards, no stale index. Reloading it returns the tensors that were written.

Environment info / 环境信息

  • mcore_bridge main at 6160abe

  • torch 2.8.0, safetensors 0.8.0

  • Reproduced on both Linux and Windows

  • I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

主要语言
Python
星标
106
派生
53
平均合并
2 天 6 小时
30 天内合并 PR
18

环境准备

这个项目没有提供开发容器、Dockerfile 或贡献指南，环境需要你自己搭建：先看它的 README，通用步骤见我们的新手贡献指南。

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

modelscope/mcore-bridge 的其他 Issue

查看 modelscope/mcore-bridge 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。