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

オープン 初心者向け
  • 何か問題がありますか？
#257 コメント 0 件 リアクション 0 件 担当者 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 を使った再現方法が記載されています。既存の出力ディレクトリで、単一ファイルへのエクスポートとシャード分割されたエクスポートがどのように完了処理されるかを確認してください。2 回目のエクスポート後にそのエクスポート自身のファイルだけが残り、再読み込みで新たに書き込まれたテンソルだけが正確に返されれば完了です。

索引モデルが 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時間
マージ済み PR（30日）
18

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

modelscope/mcore-bridge のほかの issue

modelscope/mcore-bridge の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。