load_pretrained_model deep-copies the whole state dict, doubling peak load memory

[yydhYYDH]

🐛 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

1. Install with: pip install funasr modelscope kaldi-native-fbank
2. Run: load any large checkpoint through AutoModel and sample peak RSS
3. 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_state is produced inside the function by torch.load, so it has a single owner and is not shared with the caller.
• ori_state is not referenced again after the deepcopy line.
• src_state is only read from that point on. The loop below reads src_state[k_src].shape and rebinds references in dst_state (dst_state[k] = src_state[k_src]); the data actually reaches the model through obj.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; main at 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.

[LauraGPT]

Fixed by #3729 and verified on exact merged main 9e7f6af. Focused loader tests pass 13/13, independent checkpoint-format probes pass 6/6, and the 128 MiB separate-process probe reduced peak RSS by about 143 MiB while preserving loaded values.