load_pretrained_model deep-copies the whole state dict, doubling peak load memory
维护者通常 1 天内回复
还没有人认领这个 Issue。
评估
- 难度
- 2/5
- 预计耗时
- 1-3 小时
- 新手友好度
- 88/100
- Issue 类型
- 缺陷
- 描述清晰度
- 描述清楚
- 活跃度
- 活跃
- 技术栈
- python, pytorch
调研方向
从 funasr/train_utils/load_pretrained_model.py 开始,跟踪 AutoModel 在 torch.load 和多余的 deepcopy 附近的 checkpoint 加载路径。移除不必要的第二份 state-dict 副本,然后使用指定的大型 checkpoint 运行复现,并确认加载仍然报告所有键都匹配,同时峰值 RSS 不再包含第二份 checkpoint 大小的副本。
由索引模型根据 Issue 内容生成。
描述
🐛 Bug
load_pretrained_model() deep-copies the entire checkpoint state dict on every model load. The copy is redundant, and it holds a second full copy of the weights in memory for the duration of the load, so peak host memory is roughly doubled by the checkpoint size.
Loading iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online (220M params, 956 tensors, 840MB model.pt) peaks at 3461 MB RSS. With the copy removed the same load peaks at 2624 MB — a 837MB difference that matches the checkpoint size.
The practical consequence is that a model which fits in memory can still fail to load, and container memory limits have to be set to twice the checkpoint size.
To Reproduce
- Install with:
pip install funasr modelscope kaldi-native-fbank - Run: load any large checkpoint through
AutoModeland sample peak RSS - See: no exception on a roomy host — the symptom is peak RSS; on a constrained host it is an OOM kill
pip install funasr==1.4.16 modelscope kaldi-native-fbank
python - <<'PY'
import resource, time
from funasr import AutoModel
t0 = time.time()
AutoModel(
model="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online",
device="cpu",
disable_update=True,
)
peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
print(f"{time.time() - t0:.1f}s peak RSS {peak:.0f} MB")
PY
Code sample
funasr/train_utils/load_pretrained_model.py:
ori_state = torch.load(path, map_location=map_location) # fresh, local, never read again
src_state = copy.deepcopy(ori_state) # every tensor copied a second time
src_state = src_state["state_dict"] if "state_dict" in src_state else src_state
src_state = src_state["model_state_dict"] if "model_state_dict" in src_state else src_state
src_state = src_state["model"] if "model" in src_state else src_state
The copy is not needed:
ori_stateis produced inside the function bytorch.load, so it has a single owner and is not shared with the caller.ori_stateis not referenced again after thedeepcopyline.src_stateis only read from that point on. The loop below readssrc_state[k_src].shapeand rebinds references indst_state(dst_state[k] = src_state[k_src]); the data actually reaches the model throughobj.load_state_dict(dst_state, strict=True).- No tensor is mutated in place anywhere on this path, so the deep-copied version and the shared version are equivalent.
Expected behavior
Loading a checkpoint should not require memory for two copies of it. A 220M-parameter model should not need ~2x its own weight size in headroom.
Error logs
No exception. Measured with resource.getrusage(RUSAGE_SELF).ru_maxrss, same machine, same checkpoint, back to back:
funasr 1.4.16 : LOAD 13.3s peak RSS 3461 MB <All keys matched successfully>
same, deep copy removed : LOAD 26.9s peak RSS 2624 MB <All keys matched successfully>
The 837MB difference tracks the 840MB checkpoint, which is the copy being dropped. The wall-clock column is noisy and is not part of the report; the memory difference is the reproducible result. All keys matched successfully on both sides is the correctness check.
Environment
- OS: Linux 6.6.87.2-microsoft-standard-WSL2 (Ubuntu userspace)
- Python version: 3.12.14
- FunASR version: 1.4.16 (latest release;
mainat 41778c4 is identical here) - ModelScope version: 1.40.1
- PyTorch version: 2.14.0+cu126
- Install method:
pip - Device: cpu for the load measurement
- GPU model: NVIDIA GeForce RTX 4060 Laptop GPU
- CUDA version: 12.6
Audio details
Not audio-related. The model used for the measurement is iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online (220M params, 956 tensors, 840MB model.pt). The overhead scales with checkpoint size, so it should reproduce with any large FunASR checkpoint, including the LLM-ASR models.
- 主要语言
- Python
- 星标
- 20.4k
- 派生
- 2k
- 平均合并
- 9 小时 32 分钟
- 30 天内合并 PR
- 171
环境准备
从这里开始
- 先读完整个 Issue,再读项目的贡献指南。
- 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
- Fork 仓库,在一个分支上完成修改。
- 提交 Pull Request,并在描述里引用这个 Issue 编号。
modelscope/FunASR 的其他 Issue
-
难度 2/5 1-3 小时 新手友好度 86/100
modelscope/FunASR#3704 · 1 条评论 ·
维护者通常 1 天内回复
-
bug needs feedback
难度 2/5 1-3 小时 新手友好度 76/100
modelscope/FunASR#3401 · 2 条评论 ·
维护者通常 1 天内回复
-
bug needs triage
难度 4/5 3-5 天 新手友好度 30/100
modelscope/FunASR#3727 · 1 条评论 ·
维护者通常 1 天内回复
-
needs triage question
难度 4/5 3-5 天 新手友好度 30/100
modelscope/FunASR#3718 · 3 条评论 ·
维护者通常 1 天内回复
-
难度 3/5 1-2 天 新手友好度 48/100
modelscope/FunASR#3717 · 1 条评论 ·
维护者通常 1 天内回复
查看 modelscope/FunASR 的全部 Issue
相似的 Issue
-
难度 2/5 1-3 小时 新手友好度 74/100
维护者通常 1 天内回复
-
难度 2/5 1-3 小时 新手友好度 84/100
gradio-app/gradio#13895 ·
维护者通常 1 天内回复
-
build-error
难度 2/5 1-3 小时 新手友好度 76/100
spack/spack-packages#6713 ·
维护者通常 1 天内回复
-
难度 2/5 1-3 小时 新手友好度 72/100
ActivityWatch/activitywatch#1464 · 1 个 reaction ·
维护者通常 1 天内回复
-
[Bug]: The ckg tool drops the return type of every decorated Python method in class search results未关闭
难度 2/5 1-3 小时 新手友好度 78/100
bytedance/trae-agent#483 ·
维护者通常 1 天内回复