【Bug已解决】[WebGPU] Optimize load-time memory peak by streaming external data instead of fully materia…

【Bug已解决】[WebGPU] Optimize load-time memory peak by streaming external data instead of fully materializing it 解决方案一、现象长什么样在 WebGPU 环境加载一个带外部数据external data即权重存在单独的.data文件里、ONNX 模型只存引用的大模型时加载阶段的内存峰值异常高——ORT 在 load 时把整个外部数据文件一次性全读进内存、物化materialize成完整张量导致加载期内存占用是模型实际大小的 2~3 倍在内存受限的设备手机/浏览器 tab / 嵌入式上直接 OOM 或触发系统杀进程。现象# 现象 A加载期内存峰值是模型实际大小的 2~3 倍 # 模型权重 2GB加载时内存冲到 5GB外部数据被整文件读入 解析成张量 副本 # 现象 B只在带 external data 的大模型触发 # 小模型或不拆外部数据的模型内存正常大模型一拆外部数据反而更吃内存 # 现象 C不报错但加载慢且卡 # 整文件读入 全物化加载耗时久、内存抖动大浏览器 tab 被冻结最坑的是现象 A本意是“拆外部数据省内存”结果加载时反而更吃内存因为实现把外部数据整文件读入再切分等于同时在内存里有“原始文件”和“解析后张量”两份。二、背景ONNX 的 external data 机制大模型权重放在独立的二进制.data文件模型.onnx里每个大型初始值initializer只记录external_data属性offset/length/checksum指向.data文件里的某一段。正确加载方式是按需、按段读取为每个 initializer 用pread/内存映射只取它对应的offset..offsetlength那一段物化成该张量避免把整个.data文件读进内存。ORT 的 WebGPU EP 加载路径在重构时为了简单先read()把整个.data文件读成一块大 buffer再从中切出各 initializer——这就是“完全物化”。于是加载期内存 整文件大小 所有张量总和峰值远超模型实际占用。这是资源加载/内存审查里典型的坑外部数据被整文件读入而非按需按段流式读取导致加载内存峰值翻倍。三、根因整文件读入而非按段流式加载时f.read()整个.data文件而非用offset/length按需pread内存峰值 文件 张量现象 A。未用内存映射mmap本可用mmap让 OS 按页按需调页避免一次性占用但实现没用。缺少加载内存峰值回归CI 没测“加载期峰值内存”翻倍问题长期存在。本质是外部数据加载时被整文件物化而非按需流式/分页数读取导致加载内存峰值翻倍且缺峰值回归。四、最小可运行复现下面用 Python 模拟“整文件读入 vs 按需按段读取的内存差异”import io class FakeExternalData: def __init__(self, total_size): self.data b\x00 * total_size # 模拟 2GB 外部数据文件 def load_fully_buggy(self, initializers): buggy: 整文件读入再切。 whole self.data # 占 total_size tensors {} for name, (off, length) in initializers.items(): tensors[name] whole[off:offlength] # 又占各段之和 return whole, tensors # 峰值 ≈ total_size sum(lengths) def load_streamed_fixed(self, initializers): fixed: 按需按段读取不整文件占内存。 tensors {} for name, (off, length) in initializers.items(): tensors[name] self.data[off:offlength] # 只这段可配合 mmap return None, tensors # 峰值 ≈ sum(lengths)无整文件副本 ext FakeExternalData(total_size2_000_000_000) # 2GB inits {w1: (0, 1_000_000_000), w2: (1_000_000_000, 1_000_000_000)} whole, _ ext.load_fully_buggy(inits) print(buggy peak footprint includes full file:, len(whole) 2_000_000_000) _, tens ext.load_streamed_fixed(inits) print(fixed peak footprint skips full file:, tens[w1] is not None)buggy峰值含整文件 2GBfixed只按段无整文件副本。五、解决方案第一层最小直接修复最小修复加载外部数据时按需按段读取用offset/length做pread或 mmap 视图不再整文件读入// 修正按 initializer 的 offset/length 流式读取外部数据 for (const auto init : graph.initializers()) { auto ext init.external_data(); // {offset, length, checksum} // 只读取该段而非整文件 std::vectorchar buf(ext.length); file.pread(buf.data(), ext.length, ext.offset); // 按需定位读 init.set_data(std::move(buf)); // 整文件句柄保持打开用 mmap 视图更佳零拷贝 }这一层改动最小从整文件read()改为按offset/length流式pread加载内存峰值从“文件张量”降到“张量”且可用 mmap 进一步降。但依赖“每处外部数据加载都改对”下看第二层。六、解决方案第二层结构性改进把“外部数据必须按需流式读取、禁止整文件物化”固化成单一事实来源。下面这个 dataclass 集中管理外部数据加载契约from dataclasses import dataclass, field from typing import Dict, Tuple dataclass class WebGpuStreamExternalPolicy: 单一事实来源外部数据按需流式加载契约。 _fully_materialized: bool False # 禁止整文件物化 def load_initializer(self, ext_store: Dict[str, bytes], name: str, offset: int, length: int) - bytes: if self._fully_materialized: raise AssertionError(external data must be streamed, not fully materialized) # 只取该 initializer 对应的段 return ext_store[__file__][offset:offset length] def assert_streamed(self, peak_footprint: int, model_size: int) - None: # 加载峰值不应超过模型实际大小允许小余量 if peak_footprint model_size * 1.2: raise AssertionError( fload peak {peak_footprint} model size {model_size}; fexternal data not streamed)这一层的关键收益禁止整文件物化_fully_materializedFalse作为契约违反即报错按需按段load_initializer只取offset..offsetlength段峰值校验assert_streamed验证加载峰值不超过模型大小防回归单一事实来源所有外部数据加载约定收口在WebGpuStreamExternalPolicy。七、解决方案第三层断言 / CI 守护把第二层钉成 pytest挂进 CI覆盖流式加载与内存峰值import pytest from your_package.webgpu_stream import WebGpuStreamExternalPolicy def test_streamed_reads_only_segment(): # 断言 1按需只取该 initializer 对应段不整文件 p WebGpuStreamExternalPolicy() store {__file__: bytes(range(256))} seg p.load_initializer(store, w1, 10, 5) assert seg bytes(range(10, 15)) def test_fully_materialized_forbidden(): # 断言 2开启整文件物化必须报错 p WebGpuStreamExternalPolicy(_fully_materializedTrue) with pytest.raises(AssertionError): p.load_initializer({__file__: bx}, w, 0, 1) def test_peak_within_model_size(): # 断言 3加载峰值不超过模型大小 p WebGpuStreamExternalPolicy() model_size 2_000_000_000 p.assert_streamed(peak_footprint2_100_000_000, model_sizemodel_size) def test_peak_exceeds_caught(): # 断言 4峰值翻倍必须被抓复现原 bug p WebGpuStreamExternalPolicy() with pytest.raises(AssertionError): p.assert_streamed(peak_footprint5_000_000_000, model_size2_000_000_000)四条断言从“按段读取”“禁止物化”“峰值合规”“峰值翻倍被抓”四面把内存回归钉死在 CI。八、排查清单WebGPU 加载带 external data 的大模型内存峰值时加载期内存是否是模型实际大小的 2~3 倍是就确认外部数据被整文件读入现象 A。是否用了offset/length按需pread或 mmap没用就是整文件物化。小模型正常、大模型反吃内存查外部数据加载路径是否“先整文件后切分”。用第二层WebGpuStreamExternalPolicy禁止物化 按需按段 峰值校验。加第三层 pytest断言“按段读取、禁止物化、峰值合规、翻倍被抓”。外部数据的本意是省内存加载实现必须按需流式否则适得其反。九、小结WebGPU EP 加载 external data 大模型时内存峰值翻倍本质是外部数据被整文件读入再切分完全物化而非按offset/length按需流式/pread/mmap 读取导致加载期内存 整文件 所有张量远超模型实际占用且缺加载峰值回归。修复分三层——第一层改为按需按段pread第二层用WebGpuStreamExternalPolicy这个 dataclass 把“禁止整文件物化 按需按段 峰值校验”收口成单一事实来源第三层用四条 pytest 把“按段读取、禁止物化、峰值合规、翻倍被抓”钉死在 CI。核心心法外部数据必须按需按段流式读取pread/mmap绝不能整文件物化否则加载内存峰值翻倍违背拆外部数据的初衷。