1. 先说结论init_empty_weights 不是真的“空”而是给你一张“图纸”如果你在深度学习圈子里混了一段时间大概率遇到过这种场景想看一下某个大模型的结构比如层数、每个模块的维度、dropout 位置但模型文件动辄几十个 GB显存根本塞不下。或者你只是想画个网络结构图、做参数统计、排查某个层是否加载正确这时候把整个模型完整加载进来成本高得离谱。init_empty_weights就是为这种场景设计的。它来自transformers库作用是让模型在初始化时不分配真正的数据内存只保留权重张量的“元信息”shape、dtype、device 等。换句话说你用with init_empty_weights():包裹模型实例化过程得到的模型像一张建筑图纸结构完整但是没有实际的钢筋水泥。听起来很完美对吧但实际操作中尤其在探查某些模型结构时你可能会踩到一个很有意思的异常——NotImplementedError。这个报错出现得莫名其妙有时候同一个模型换个设备就没事了有时候换个注意力实现就崩了。这篇博文就来聊聊它背后的原理、我实际踩过的场景以及怎么定位和解决。这篇文章适合谁适合那些想在低显存环境下探查模型结构、做模型迁移、自定义加载逻辑或者纯粹想弄懂transformers底层机制的开发者。不需要你有多深的框架源码功底但建议对 PyTorch 的基础概念比如torch.nn.Module、device、张量有基本了解。2. init_empty_weights 机制拆解它是怎么做到“不占显存”的2.1 meta device 是个什么东西要理解init_empty_weights必须先理解 PyTorch 的 meta device。在 PyTorch 中每个张量都有device属性常见的有cpu、cuda以及分布式场景下的xpu、npu等。而meta是一个特殊的 device它的核心特点是张量没有实际数据存储空间。你在 meta device 上创建张量时PyTorch 只记录这个张量应该有的形状和数据类型但不会分配任何内存。这里有个生活化类比你去车管所查一辆车的档案能看到它的品牌、型号、排量、颜色但档案馆里并没有停放这辆真实的车。meta device 就是这套档案系统它保存了“这辆车应该长什么样”的全部信息但没有任何真实车辆的存在。import torch # 在 meta device 上创建一个“虚拟”张量 x torch.empty(1024, 1024, devicemeta) print(x.shape) # torch.Size([1024, 1024]) print(x.dtype) # torch.float32默认 print(x.device) # meta print(x.sum()) # NotImplementedError 大概率出现当你尝试对这个张量做sum()这样的数值计算时PyTorch 不知道该怎么办——因为它根本没有任何真实数据可以加。于是它只能抛出NotImplementedError。这个机制是整个问题的起点也是init_empty_weights能工作的基础。2.2 init_empty_weights 的工作原理transformers库的init_empty_weights做的事情核心就是把模型__init__过程中创建的所有参数张量都放到 meta device 上。它通过上下文管理器临时修改 PyTorch 的参数初始化和张量生成行为让torch.empty、torch.zeros、torch.nn.init.*这类操作默认在 meta device 上执行。实际用法基本是这样from transformers import AutoConfig, AutoModelForCausalLM from transformers.utils import init_empty_weights config AutoConfig.from_pretrained(meta-llama/Llama-2-7b-hf) with init_empty_weights(): model AutoModelForCausalLM.from_config(config) # 此时模型结构是完整的但所有参数都在 meta device 上 for name, param in model.named_parameters(): print(name, param.shape, param.device)这段代码执行完毕后模型的结构信息模块层级、维度、类型全部可用但参数没有占任何显存。你可以在这种状态下做很多事情打印每一层的参数形状画网络结构图。统计参数量看模型尺寸。做模块级别的替换和修改然后再选择性加载权重的某些部分。对超大模型做 pre-allocation 或者水印式加载比如把 attention 层先加载到 GPU其他层留在 meta 上。原理说起来不复杂但正因为所有参数都变成了 meta 张量那些依赖真实数值的操作就会全面崩盘。这就是NotImplementedError出现的第一层原因。2.3 有哪些场景会触发 NotImplementedError我观察下来NotImplementedError主要集中在三类情况里。搞清楚了这些触发点你排查问题时就不会像无头苍蝇一样乱转。第一种模型的__init__里调用了需要数值计算的方法。很多自定义模块会在初始化时对权重做特殊处理比如class MyLayer(nn.Module): def __init__(self, dim): super().__init__() self.weight nn.Parameter(torch.empty(dim, dim)) # 这里会触发问题 self.weight.data torch.nn.functional.normalize(self.weight.data, dim0)normalize需要读取张量的真实数值来计算范数可 meta 张量根本没有数据于是抛出NotImplementedError。第二种模型里有条件分支依赖某个参数的数值做判断。比如某些注意力模块会检查head_dim是否为 64然后决定使用 flash attention 还是普通 attention。如果这个判断是if self.head_dim % 64 0那是没问题的因为%只依赖形状相关的元数据但如果代码写成if self.some_scalar 0.5而这个some_scalar是参数值那就直接崩了。第三种某些算子没有注册 meta 设备的实现。PyTorch 和 transformers 对大部分常用算子都有 meta device 的 fallback 实现也就是不真的计算而是直接返回一个形状正确的 meta 张量。但总有一些边缘算子没有注册一旦模型__init__里调用了这些算子就会报NotImplementedError。这其中的区别很有意思——有些异常是因为“不该读数据”而崩溃有些是因为“PyTorch 没写 meta 分支”而崩溃。前者是设计如此后者属于兼容性 bug。3. 实操场景复现用实例讲清楚 NotImplementedError 的来龙去脉3.1 场景一Meta 张量上做数值归一化导致的崩溃我先写一个最小复现案例让你完整看到异常是怎么冒出来的。import torch import torch.nn as nn from transformers.utils import init_empty_weights class MyCustomLayer(nn.Module): def __init__(self, dim): super().__init__() self.weight nn.Parameter(torch.empty(dim, dim)) # 模拟一些老代码中的初始化操作 self.weight.data torch.nn.functional.normalize(self.weight.data, p2, dim1) try: with init_empty_weights(): layer MyCustomLayer(128) except NotImplementedError as e: print(fCaught: {e})这个例子会直接抛出类似这样的错误NotImplementedError: Could not run aten::_foreach_norm with arguments from the Meta backend...原因很清晰normalize底层需要计算 L2 范数而 L2 范数需要读取原始数值。meta 张量没有数据PyTorch 找不到对应实现直接报错。为什么会这样设计其实这是合理的。meta device 本身就是“结构优先”它不承诺支持所有数值操作。如果你只是想探查模型结构根本就不该有这类初始化数值计算的代码。但现实里很多模型源码是从完整训练代码里直接剽过来的带了一堆初始化逻辑于是就会在这里翻车。3.2 场景二大规模模型中的 attention 实现冲突这个场景更贴近真实踩坑。我在探查一个基于 LLaMA 架构的模型时用了init_empty_weights去加载 config结果直接抛了NotImplementedError。报错堆栈指向了sdpascaled dot product attention相关代码。当时的报错信息长这样NotImplementedError: No operator found for aten::_scaled_dot_product_attention with args: [Meta...这个问题的根因在于新版 transformers 在LlamaAttention的__init__阶段就会根据config._attn_implementation来决定构建哪种 attention。有的实现会在__init__里就创建一些辅助张量或者调用一次 SDPA 来做形状校验——正常加载时没有问题因为权重是真实数据但在 meta device 下就会触发NotImplementedError。我当时排查的路径是这样的看堆栈找到崩溃的具体 Python 文件行号。进入transformers/models/llama/modeling_llama.py定位到__init__中触发 SDPA 的代码。发现该 attention 实现里有一个_check_input_shape之类的内部方法在初始化时被调用。那行代码本质上是在做“输入 shape 预校验”不需要真实权重但实现时偷懒直接调用了一次算子。这种问题怎么解最简单的办法是绕开那个分支——通过修改配置让from_config走的路径不触发那段代码。具体来说在初始化之前显式指定 attention 实现from transformers import AutoConfig config AutoConfig.from_pretrained(some/model-name) config._attn_implementation eager # 强制走 eager 分支绕开 sdpa 初始化校验 with init_empty_weights(): model AutoModelForCausalLM.from_config(config)实测下来这个办法能解决相当一部分跟 attention 实现相关的NotImplementedError。核心思路就是让模型初始化走最简单、最朴素的路径避开花哨的 fused 算子。3.3 场景三自定义模块里的 torch.empty 与 register_buffer还有一种常见情况自定义模块在__init__里调用了torch.empty或torch.zeros创建 buffer但这些 buffer 并不作为nn.Parameter注册而是直接调用了一些需要真实数据的操作比如class MyModel(nn.Module): def __init__(self): super().__init__() self.register_buffer(mask, torch.ones(4, 4).triu(1))看这段代码torch.ones本身在 meta device 下是可以正常创建 meta 张量的因为它不涉及数值读取只记录形状但紧接着的.triu(1)会尝试读取张量的维度信息并执行上三角操作。对于meta张量来说triu这种按位操作是否能执行取决于 PyTorch 是否给 meta 后端注册了对应的triu实现。我实测的结果是——在大部分 PyTorch 版本里triu是支持的但类似torch.linalg.eig、torch.svd、torch.det这类算法式的操作几乎必然报NotImplementedError。这类问题的共性特征是报错不是发生在参数初始化时而是发生在 buffer 创建或模块级常量初始化时。定位手段是看堆栈文件是哪个。4. 排查指南不只是看报错信息要顺着堆栈找根因4.1 我的标准排查流程遇到NotImplementedError很多人第一反应是上网搜报错信息然后发现异常五花八门搜半天找不到对症的答案。我的经验是先做一套系统排查比直接搜报错更高效。第一步抓完整堆栈不要只看最后一行。NotImplementedError的堆栈往往很长关键是找到模型中实际触发报错的 Python 代码位置。import traceback try: with init_empty_weights(): model AutoModelForCausalLM.from_config(config) except NotImplementedError: traceback.print_exc()第二步判断异常属于哪一类。看堆栈里是aten::开头的算子还是 Python 层的assert或自定义逻辑。如果是aten::_foreach_norm这种底层算子一般是 meta 张量没有对应实现如果是你自己代码里的raise NotImplementedError那大概率是模型源码写死的要求。第三步针对性地绕开。绕开方式有两类修改配置让模型走不同的实现路径比如从sdpa切到eager。修改模型源码把__init__里无关紧要的初始化逻辑用if param.device.type ! meta:包起来。第三步里第二个方法其实非常有实用价值很多人不知道 meta 张量的device.type是可以判断的。if self.weight.device.type ! meta: # 只有实际加载模型时才做初始化 self.weight.data torch.nn.functional.normalize(self.weight.data, dim1)这种改法相对保守不会破坏正常加载时的行为只是在 meta device 下跳过数值初始化。4.2 常见报错速查表我整理了这份表格覆盖大部分我见过的 meta 张量相关异常。报错特征触发原因推荐解法aten::_foreach_norm/aten::norm相关初始化时对权重做 normalize 或范数计算跳过 meta 设备上的数值初始化aten::_scaled_dot_product_attentionattention 实现里调用了 SDPA 做形状预校验配置_attn_implementationeageraten::mm/aten::addmm某些模块在__init__里就做了矩阵乘法预演修改源码只用 meta 张量的 shape不执行计算torch.linalg.*系列初始化时做了矩阵分解、特征值分析等这类几乎无法在 meta 下执行直接移除或加条件判断自研代码raise NotImplementedError模型源码主动抛出的异常读源码确认逻辑逐条判断是否可以绕过assert条件不成立模型初始化时对某参数值做了硬性断言检查是不是依赖真实数值是则加 meta 判断注意这个表不是为了让你全背下来而是告诉你一个规律报错的算子越“数值化”越难在 meta 设备上绕过越“结构性”越容易处理。5. 绕开 NotImplementedError 的六个实战方案5.1 方案一修改配置强制走 eager 路径这是最低成本方案也是我优先尝试的。transformers 近年来把不少注意力实现都改成了可插拔式很多模型的 config 里都有_attn_implementation这个隐藏配置项。在from_config之前手动指定为eager可以跳过大量 fused 算子和 kernel 选择逻辑。from transformers import AutoConfig, AutoModelForCausalLM from transformers.utils import init_empty_weights config AutoConfig.from_pretrained(your-model-path) config._attn_implementation eager with init_empty_weights(): model AutoModelForCausalLM.from_config(config)为什么eager能绕开问题因为 eager 模式是纯 PyTorch 原生实现操作路径上使用的都是最基础的算子比如torch.matmul、softmax这些算子基本都有 meta 设备实现。而sdpa或flash_attention_2这类融合算子很多都没有注册 meta 的 fallback。5.2 方案二给模型源码打补丁跳过数值初始化如果改配置解决不了就得动手改源代码了。以开头那个normalize的案例为例标准做法是给MyCustomLayer.__init__加一个 device 判断。class MyCustomLayer(nn.Module): def __init__(self, dim): super().__init__() self.weight nn.Parameter(torch.empty(dim, dim)) # 只在真实设备上执行初始化 if self.weight.device.type ! meta: self.weight.data torch.nn.functional.normalize(self.weight.data, p2, dim1)但这里有个坑——nn.Parameter(torch.empty(dim, dim))在init_empty_weights下创建的张量device 确实是 meta这个判断没问题。可如果是你自己手动创建的普通 buffer非 Parameter你得留意它是不是真的在 meta 设备上。保险起见可以在代码里加一行打印验证print(fweight device: {self.weight.device})5.3 方案三使用 accelerate 的load_checkpoint_and_dispatch方案替代如果init_empty_weights遇到实在绕不开的坑比如你要加载的模型里有大量torch.linalg调用那不如换个思路直接用accelerate库来加载大模型。accelerate的load_checkpoint_and_dispatch也依赖 meta device 机制但它在加载权重时做了更多容错处理。from accelerate import init_empty_weights, load_checkpoint_and_dispatch with init_empty_weights(): model AutoModelForCausalLM.from_config(config) model load_checkpoint_and_dispatch( model, checkpointpath/to/model.bin, device_mapauto, no_split_module_classes[LlamaDecoderLayer], )这样虽然还是在初始阶段用init_empty_weights但权重加载和 dispatch 是分离的很多初始化时的算子冲突可以避开。实测下来对于像 LLM 这种体量的模型这个方案比裸用init_empty_weights稳定不少。5.4 方案四只探查结构时用 from_pretrained 的 low_cpu_mem_usage如果你只是想看模型结构不关心自己手工实例化时的一些约束可以换个 API。from_pretrained本身就支持low_cpu_mem_usageTrue它会用类似 meta device 的机制来降低 CPU 内存占用。from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained( your-model-path, low_cpu_mem_usageTrue, )这个方法的好处是它走的是官方加载路径对模型内部的初始化兼容性更好。但注意它依然会尝试读取权重文件所以需要真的下载或本地有权重文件。不像from_config那样只靠配置文件就能搭出结构。5.5 方案五拆开模型结构手动构建关键模块有些场景下你就是想看某个子模块的结构但整体模型又装不上。这时候没必要强上完整模型可以直接用AutoConfig拿到所有超参数然后用transformers里对应的模块类手动拼接。from transformers import AutoConfig, LlamaConfig from transformers.models.llama.modeling_llama import LlamaDecoderLayer config AutoConfig.from_pretrained(your-model-path) layer LlamaDecoderLayer(config, layer_idx0) # 查看这一层的子模块结构 print(layer)这种“手动局部构建”的思路其实是最稳的。因为你只构建你认为有必要的部分就不会触发那些跟全局初始化相关的算子冲突。5.6 方案六设置环境变量或 monkey patch如果前面方案全都不行还可以打一个通用补丁在进入init_empty_weights上下文之后对所有torch.Tensor的数值方法做 monkey patch让它们遇到 meta 张量时直接返回一个 meta 结果。这个方案听起来有点“野路子”但我在某些极端模型上确实用上了。核心逻辑很简单import torch _original_sum torch.Tensor.sum def safe_sum(self, *args, **kwargs): if self.device.type meta: return torch.empty((), devicemeta) return _original_sum(self, *args, **kwargs) torch.Tensor.sum safe_sum注意这个方法风险比较高可能会掩盖真实的 bug或者导致后续逻辑拿到一个没有意义的张量。我的建议是只有在探查结构这种“用完即弃”的场景里用不要在模型加载管线上长期使用。6. 实测记录一个 LLaMA 系模型的完整探查过程为了让你有更直观的参考我分享一个最近的实测案例。当时要探查一个 13B 规模的 LLaMA 系模型的层间结构显存只有 24G完整加载不现实。我的操作过程如下。第一步用from_config搭配init_empty_weights初始化结构config AutoConfig.from_pretrained(model_path) config._attn_implementation eager with init_empty_weights(): model AutoModelForCausalLM.from_config(config)第一次运行就报了NotImplementedError堆栈指向LlamaMLP的__init__因为模型在初始化时对gate_proj的权重做了一个 reshape 校验。第二步定位到具体源码看那段代码的作用# LlamaMLP.__init__ 中类似这样的逻辑 self.gate_proj nn.Linear(self.hidden_size, self.intermediate_size, biasFalse) assert self.gate_proj.weight.shape (self.intermediate_size, self.hidden_size)这种assert本身问题不大因为它只依赖 shape不读取数值不会触发NotImplementedError。所以我的第一步修改是直接忽略这个assert是否能过。但如果它真报错了就说明模型代码里有更复杂的数值校验。第三步我把_attn_implementation切到eager后模型就正常实例化了。整个过程中没有任何额外的代码修改所以这个问题本质上就是 SDPA 路径对 meta 设备不友好导致的。实例化成功后我用一段小脚本把模型每层的参数概况打出来for name, module in model.named_modules(): params sum(p.numel() for p in module.parameters(recurseFalse)) if params 0: print(f{name}: {params / 1024 / 1024:.2f}M params)整个过程占用显存 0CPU 内存也几乎可以忽略这就是 meta device 方案最大的价值——用极低的成本获取完整的结构信息。7. 注意事项与常见坑位汇总7.1 不要在生产加载流程里长期使用 monkey patch我在 5.6 里提到的方法应急可以但不建议当成常规手段。原因很简单NotImplementedError是 PyTorch 在 meta device 上的一种“诚实”表现——它告诉你这个算子真的没法在无数据的情况下执行。你 monkey patch 掉了异常但并没有提供真实数据后续如果哪个逻辑真的依赖这些数值结果就是各种诡异的错误比如NaN、shape 错乱、隐式广播异常。排查起来比最初的NotImplementedError痛苦十倍。7.2 注意init_empty_weights嵌套使用的状态泄漏transformers的init_empty_weights是上下文管理器它在__enter__时会修改一些全局状态比如torch.empty的默认行为在__exit__时会恢复。但如果你在自己的代码里嵌套了多个上下文或者异常导致__exit__没被执行就可能出现“出了上下文但参数还是在 meta device 上创建”的诡异情况。我的建议是尽量把init_empty_weights包裹的代码范围缩到最小并且保证异常安全。try: with init_empty_weights(): model AutoModelForCausalLM.from_config(config) except Exception: # 记录日志后重新抛出避免继续执行后续逻辑 raise7.3 注意 meta 张量的item()和数值计算在 meta 张量上调用.item()基本都会报错。有些模型__init__里会读取某个参数的数值来判断是否启用某种初始化策略比如if self.scale 0:。这种写法在正常初始化时没问题但 meta device 下行不通。处理方式是先用is_meta判断或者把这种逻辑放在权重加载之后再做。7.4 设备不一致导致的混淆还有一个常见坑模型里如果混合了 CPU 参数和 meta 参数某些操作可能会意外地在 CPU 张量和 meta 张量之间做转换然后报NotImplementedError或者更迷惑的DeviceError。这类问题排查时可以先统一打印一下所有参数的 device 分布devices set() for _, p in model.named_parameters(): devices.add(p.device.type) print(devices)如果发现同时有cpu和meta就要留意是不是部分模块走了完全不同的初始化路径。8. 换个视角NotImplementedError 其实是设计使然不是纯粹的 bug聊到最后我想换个角度说点个人体会。很多人遇到NotImplementedError第一反应是“这个库有 bug”。但实际接触多了你会发现它与其说是 bug不如说是一种保护机制。meta device 的设计目标就是“结构感知”不是为了替代真实计算。它给了你一个机会在花掉巨量显存之前先把模型看清楚。所以我现在的习惯是探查结构用init_empty_weightsfrom_config需要真实加载用low_cpu_mem_usage需要分布式推理用accelerate的device_map方案。每个工具都有自己的适用边界跨了边界用自然容易碰壁。至于NotImplementedError这类的阻挡按上面的思路排查绝大多数都能在十分钟内定位到具体算子然后决定是改配置、改源码还是换个 API。如果你在实操中遇到表格里没有覆盖的报错建议第一步永远是看堆栈里离模型源码最近的那一行——那里藏着真正的答案。这里再说一个我个人的经验习惯每次探查完模型结构我会顺手把所用的配置项和绕坑方案记录到项目文档里。不是等出了问题再去翻而是下次加载这个模型时直接翻文档就知道要设置什么。这些模型源码会更新transformers 版本也会变网络上搜到的答案可能过期但你自己整理的排查路径往往是最贴合实际情况的。
解决init_empty_weights中的NotImplementedError：模型结构探查与低显存加载实战指南
1. 先说结论init_empty_weights 不是真的“空”而是给你一张“图纸”如果你在深度学习圈子里混了一段时间大概率遇到过这种场景想看一下某个大模型的结构比如层数、每个模块的维度、dropout 位置但模型文件动辄几十个 GB显存根本塞不下。或者你只是想画个网络结构图、做参数统计、排查某个层是否加载正确这时候把整个模型完整加载进来成本高得离谱。init_empty_weights就是为这种场景设计的。它来自transformers库作用是让模型在初始化时不分配真正的数据内存只保留权重张量的“元信息”shape、dtype、device 等。换句话说你用with init_empty_weights():包裹模型实例化过程得到的模型像一张建筑图纸结构完整但是没有实际的钢筋水泥。听起来很完美对吧但实际操作中尤其在探查某些模型结构时你可能会踩到一个很有意思的异常——NotImplementedError。这个报错出现得莫名其妙有时候同一个模型换个设备就没事了有时候换个注意力实现就崩了。这篇博文就来聊聊它背后的原理、我实际踩过的场景以及怎么定位和解决。这篇文章适合谁适合那些想在低显存环境下探查模型结构、做模型迁移、自定义加载逻辑或者纯粹想弄懂transformers底层机制的开发者。不需要你有多深的框架源码功底但建议对 PyTorch 的基础概念比如torch.nn.Module、device、张量有基本了解。2. init_empty_weights 机制拆解它是怎么做到“不占显存”的2.1 meta device 是个什么东西要理解init_empty_weights必须先理解 PyTorch 的 meta device。在 PyTorch 中每个张量都有device属性常见的有cpu、cuda以及分布式场景下的xpu、npu等。而meta是一个特殊的 device它的核心特点是张量没有实际数据存储空间。你在 meta device 上创建张量时PyTorch 只记录这个张量应该有的形状和数据类型但不会分配任何内存。这里有个生活化类比你去车管所查一辆车的档案能看到它的品牌、型号、排量、颜色但档案馆里并没有停放这辆真实的车。meta device 就是这套档案系统它保存了“这辆车应该长什么样”的全部信息但没有任何真实车辆的存在。import torch # 在 meta device 上创建一个“虚拟”张量 x torch.empty(1024, 1024, devicemeta) print(x.shape) # torch.Size([1024, 1024]) print(x.dtype) # torch.float32默认 print(x.device) # meta print(x.sum()) # NotImplementedError 大概率出现当你尝试对这个张量做sum()这样的数值计算时PyTorch 不知道该怎么办——因为它根本没有任何真实数据可以加。于是它只能抛出NotImplementedError。这个机制是整个问题的起点也是init_empty_weights能工作的基础。2.2 init_empty_weights 的工作原理transformers库的init_empty_weights做的事情核心就是把模型__init__过程中创建的所有参数张量都放到 meta device 上。它通过上下文管理器临时修改 PyTorch 的参数初始化和张量生成行为让torch.empty、torch.zeros、torch.nn.init.*这类操作默认在 meta device 上执行。实际用法基本是这样from transformers import AutoConfig, AutoModelForCausalLM from transformers.utils import init_empty_weights config AutoConfig.from_pretrained(meta-llama/Llama-2-7b-hf) with init_empty_weights(): model AutoModelForCausalLM.from_config(config) # 此时模型结构是完整的但所有参数都在 meta device 上 for name, param in model.named_parameters(): print(name, param.shape, param.device)这段代码执行完毕后模型的结构信息模块层级、维度、类型全部可用但参数没有占任何显存。你可以在这种状态下做很多事情打印每一层的参数形状画网络结构图。统计参数量看模型尺寸。做模块级别的替换和修改然后再选择性加载权重的某些部分。对超大模型做 pre-allocation 或者水印式加载比如把 attention 层先加载到 GPU其他层留在 meta 上。原理说起来不复杂但正因为所有参数都变成了 meta 张量那些依赖真实数值的操作就会全面崩盘。这就是NotImplementedError出现的第一层原因。2.3 有哪些场景会触发 NotImplementedError我观察下来NotImplementedError主要集中在三类情况里。搞清楚了这些触发点你排查问题时就不会像无头苍蝇一样乱转。第一种模型的__init__里调用了需要数值计算的方法。很多自定义模块会在初始化时对权重做特殊处理比如class MyLayer(nn.Module): def __init__(self, dim): super().__init__() self.weight nn.Parameter(torch.empty(dim, dim)) # 这里会触发问题 self.weight.data torch.nn.functional.normalize(self.weight.data, dim0)normalize需要读取张量的真实数值来计算范数可 meta 张量根本没有数据于是抛出NotImplementedError。第二种模型里有条件分支依赖某个参数的数值做判断。比如某些注意力模块会检查head_dim是否为 64然后决定使用 flash attention 还是普通 attention。如果这个判断是if self.head_dim % 64 0那是没问题的因为%只依赖形状相关的元数据但如果代码写成if self.some_scalar 0.5而这个some_scalar是参数值那就直接崩了。第三种某些算子没有注册 meta 设备的实现。PyTorch 和 transformers 对大部分常用算子都有 meta device 的 fallback 实现也就是不真的计算而是直接返回一个形状正确的 meta 张量。但总有一些边缘算子没有注册一旦模型__init__里调用了这些算子就会报NotImplementedError。这其中的区别很有意思——有些异常是因为“不该读数据”而崩溃有些是因为“PyTorch 没写 meta 分支”而崩溃。前者是设计如此后者属于兼容性 bug。3. 实操场景复现用实例讲清楚 NotImplementedError 的来龙去脉3.1 场景一Meta 张量上做数值归一化导致的崩溃我先写一个最小复现案例让你完整看到异常是怎么冒出来的。import torch import torch.nn as nn from transformers.utils import init_empty_weights class MyCustomLayer(nn.Module): def __init__(self, dim): super().__init__() self.weight nn.Parameter(torch.empty(dim, dim)) # 模拟一些老代码中的初始化操作 self.weight.data torch.nn.functional.normalize(self.weight.data, p2, dim1) try: with init_empty_weights(): layer MyCustomLayer(128) except NotImplementedError as e: print(fCaught: {e})这个例子会直接抛出类似这样的错误NotImplementedError: Could not run aten::_foreach_norm with arguments from the Meta backend...原因很清晰normalize底层需要计算 L2 范数而 L2 范数需要读取原始数值。meta 张量没有数据PyTorch 找不到对应实现直接报错。为什么会这样设计其实这是合理的。meta device 本身就是“结构优先”它不承诺支持所有数值操作。如果你只是想探查模型结构根本就不该有这类初始化数值计算的代码。但现实里很多模型源码是从完整训练代码里直接剽过来的带了一堆初始化逻辑于是就会在这里翻车。3.2 场景二大规模模型中的 attention 实现冲突这个场景更贴近真实踩坑。我在探查一个基于 LLaMA 架构的模型时用了init_empty_weights去加载 config结果直接抛了NotImplementedError。报错堆栈指向了sdpascaled dot product attention相关代码。当时的报错信息长这样NotImplementedError: No operator found for aten::_scaled_dot_product_attention with args: [Meta...这个问题的根因在于新版 transformers 在LlamaAttention的__init__阶段就会根据config._attn_implementation来决定构建哪种 attention。有的实现会在__init__里就创建一些辅助张量或者调用一次 SDPA 来做形状校验——正常加载时没有问题因为权重是真实数据但在 meta device 下就会触发NotImplementedError。我当时排查的路径是这样的看堆栈找到崩溃的具体 Python 文件行号。进入transformers/models/llama/modeling_llama.py定位到__init__中触发 SDPA 的代码。发现该 attention 实现里有一个_check_input_shape之类的内部方法在初始化时被调用。那行代码本质上是在做“输入 shape 预校验”不需要真实权重但实现时偷懒直接调用了一次算子。这种问题怎么解最简单的办法是绕开那个分支——通过修改配置让from_config走的路径不触发那段代码。具体来说在初始化之前显式指定 attention 实现from transformers import AutoConfig config AutoConfig.from_pretrained(some/model-name) config._attn_implementation eager # 强制走 eager 分支绕开 sdpa 初始化校验 with init_empty_weights(): model AutoModelForCausalLM.from_config(config)实测下来这个办法能解决相当一部分跟 attention 实现相关的NotImplementedError。核心思路就是让模型初始化走最简单、最朴素的路径避开花哨的 fused 算子。3.3 场景三自定义模块里的 torch.empty 与 register_buffer还有一种常见情况自定义模块在__init__里调用了torch.empty或torch.zeros创建 buffer但这些 buffer 并不作为nn.Parameter注册而是直接调用了一些需要真实数据的操作比如class MyModel(nn.Module): def __init__(self): super().__init__() self.register_buffer(mask, torch.ones(4, 4).triu(1))看这段代码torch.ones本身在 meta device 下是可以正常创建 meta 张量的因为它不涉及数值读取只记录形状但紧接着的.triu(1)会尝试读取张量的维度信息并执行上三角操作。对于meta张量来说triu这种按位操作是否能执行取决于 PyTorch 是否给 meta 后端注册了对应的triu实现。我实测的结果是——在大部分 PyTorch 版本里triu是支持的但类似torch.linalg.eig、torch.svd、torch.det这类算法式的操作几乎必然报NotImplementedError。这类问题的共性特征是报错不是发生在参数初始化时而是发生在 buffer 创建或模块级常量初始化时。定位手段是看堆栈文件是哪个。4. 排查指南不只是看报错信息要顺着堆栈找根因4.1 我的标准排查流程遇到NotImplementedError很多人第一反应是上网搜报错信息然后发现异常五花八门搜半天找不到对症的答案。我的经验是先做一套系统排查比直接搜报错更高效。第一步抓完整堆栈不要只看最后一行。NotImplementedError的堆栈往往很长关键是找到模型中实际触发报错的 Python 代码位置。import traceback try: with init_empty_weights(): model AutoModelForCausalLM.from_config(config) except NotImplementedError: traceback.print_exc()第二步判断异常属于哪一类。看堆栈里是aten::开头的算子还是 Python 层的assert或自定义逻辑。如果是aten::_foreach_norm这种底层算子一般是 meta 张量没有对应实现如果是你自己代码里的raise NotImplementedError那大概率是模型源码写死的要求。第三步针对性地绕开。绕开方式有两类修改配置让模型走不同的实现路径比如从sdpa切到eager。修改模型源码把__init__里无关紧要的初始化逻辑用if param.device.type ! meta:包起来。第三步里第二个方法其实非常有实用价值很多人不知道 meta 张量的device.type是可以判断的。if self.weight.device.type ! meta: # 只有实际加载模型时才做初始化 self.weight.data torch.nn.functional.normalize(self.weight.data, dim1)这种改法相对保守不会破坏正常加载时的行为只是在 meta device 下跳过数值初始化。4.2 常见报错速查表我整理了这份表格覆盖大部分我见过的 meta 张量相关异常。报错特征触发原因推荐解法aten::_foreach_norm/aten::norm相关初始化时对权重做 normalize 或范数计算跳过 meta 设备上的数值初始化aten::_scaled_dot_product_attentionattention 实现里调用了 SDPA 做形状预校验配置_attn_implementationeageraten::mm/aten::addmm某些模块在__init__里就做了矩阵乘法预演修改源码只用 meta 张量的 shape不执行计算torch.linalg.*系列初始化时做了矩阵分解、特征值分析等这类几乎无法在 meta 下执行直接移除或加条件判断自研代码raise NotImplementedError模型源码主动抛出的异常读源码确认逻辑逐条判断是否可以绕过assert条件不成立模型初始化时对某参数值做了硬性断言检查是不是依赖真实数值是则加 meta 判断注意这个表不是为了让你全背下来而是告诉你一个规律报错的算子越“数值化”越难在 meta 设备上绕过越“结构性”越容易处理。5. 绕开 NotImplementedError 的六个实战方案5.1 方案一修改配置强制走 eager 路径这是最低成本方案也是我优先尝试的。transformers 近年来把不少注意力实现都改成了可插拔式很多模型的 config 里都有_attn_implementation这个隐藏配置项。在from_config之前手动指定为eager可以跳过大量 fused 算子和 kernel 选择逻辑。from transformers import AutoConfig, AutoModelForCausalLM from transformers.utils import init_empty_weights config AutoConfig.from_pretrained(your-model-path) config._attn_implementation eager with init_empty_weights(): model AutoModelForCausalLM.from_config(config)为什么eager能绕开问题因为 eager 模式是纯 PyTorch 原生实现操作路径上使用的都是最基础的算子比如torch.matmul、softmax这些算子基本都有 meta 设备实现。而sdpa或flash_attention_2这类融合算子很多都没有注册 meta 的 fallback。5.2 方案二给模型源码打补丁跳过数值初始化如果改配置解决不了就得动手改源代码了。以开头那个normalize的案例为例标准做法是给MyCustomLayer.__init__加一个 device 判断。class MyCustomLayer(nn.Module): def __init__(self, dim): super().__init__() self.weight nn.Parameter(torch.empty(dim, dim)) # 只在真实设备上执行初始化 if self.weight.device.type ! meta: self.weight.data torch.nn.functional.normalize(self.weight.data, p2, dim1)但这里有个坑——nn.Parameter(torch.empty(dim, dim))在init_empty_weights下创建的张量device 确实是 meta这个判断没问题。可如果是你自己手动创建的普通 buffer非 Parameter你得留意它是不是真的在 meta 设备上。保险起见可以在代码里加一行打印验证print(fweight device: {self.weight.device})5.3 方案三使用 accelerate 的load_checkpoint_and_dispatch方案替代如果init_empty_weights遇到实在绕不开的坑比如你要加载的模型里有大量torch.linalg调用那不如换个思路直接用accelerate库来加载大模型。accelerate的load_checkpoint_and_dispatch也依赖 meta device 机制但它在加载权重时做了更多容错处理。from accelerate import init_empty_weights, load_checkpoint_and_dispatch with init_empty_weights(): model AutoModelForCausalLM.from_config(config) model load_checkpoint_and_dispatch( model, checkpointpath/to/model.bin, device_mapauto, no_split_module_classes[LlamaDecoderLayer], )这样虽然还是在初始阶段用init_empty_weights但权重加载和 dispatch 是分离的很多初始化时的算子冲突可以避开。实测下来对于像 LLM 这种体量的模型这个方案比裸用init_empty_weights稳定不少。5.4 方案四只探查结构时用 from_pretrained 的 low_cpu_mem_usage如果你只是想看模型结构不关心自己手工实例化时的一些约束可以换个 API。from_pretrained本身就支持low_cpu_mem_usageTrue它会用类似 meta device 的机制来降低 CPU 内存占用。from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained( your-model-path, low_cpu_mem_usageTrue, )这个方法的好处是它走的是官方加载路径对模型内部的初始化兼容性更好。但注意它依然会尝试读取权重文件所以需要真的下载或本地有权重文件。不像from_config那样只靠配置文件就能搭出结构。5.5 方案五拆开模型结构手动构建关键模块有些场景下你就是想看某个子模块的结构但整体模型又装不上。这时候没必要强上完整模型可以直接用AutoConfig拿到所有超参数然后用transformers里对应的模块类手动拼接。from transformers import AutoConfig, LlamaConfig from transformers.models.llama.modeling_llama import LlamaDecoderLayer config AutoConfig.from_pretrained(your-model-path) layer LlamaDecoderLayer(config, layer_idx0) # 查看这一层的子模块结构 print(layer)这种“手动局部构建”的思路其实是最稳的。因为你只构建你认为有必要的部分就不会触发那些跟全局初始化相关的算子冲突。5.6 方案六设置环境变量或 monkey patch如果前面方案全都不行还可以打一个通用补丁在进入init_empty_weights上下文之后对所有torch.Tensor的数值方法做 monkey patch让它们遇到 meta 张量时直接返回一个 meta 结果。这个方案听起来有点“野路子”但我在某些极端模型上确实用上了。核心逻辑很简单import torch _original_sum torch.Tensor.sum def safe_sum(self, *args, **kwargs): if self.device.type meta: return torch.empty((), devicemeta) return _original_sum(self, *args, **kwargs) torch.Tensor.sum safe_sum注意这个方法风险比较高可能会掩盖真实的 bug或者导致后续逻辑拿到一个没有意义的张量。我的建议是只有在探查结构这种“用完即弃”的场景里用不要在模型加载管线上长期使用。6. 实测记录一个 LLaMA 系模型的完整探查过程为了让你有更直观的参考我分享一个最近的实测案例。当时要探查一个 13B 规模的 LLaMA 系模型的层间结构显存只有 24G完整加载不现实。我的操作过程如下。第一步用from_config搭配init_empty_weights初始化结构config AutoConfig.from_pretrained(model_path) config._attn_implementation eager with init_empty_weights(): model AutoModelForCausalLM.from_config(config)第一次运行就报了NotImplementedError堆栈指向LlamaMLP的__init__因为模型在初始化时对gate_proj的权重做了一个 reshape 校验。第二步定位到具体源码看那段代码的作用# LlamaMLP.__init__ 中类似这样的逻辑 self.gate_proj nn.Linear(self.hidden_size, self.intermediate_size, biasFalse) assert self.gate_proj.weight.shape (self.intermediate_size, self.hidden_size)这种assert本身问题不大因为它只依赖 shape不读取数值不会触发NotImplementedError。所以我的第一步修改是直接忽略这个assert是否能过。但如果它真报错了就说明模型代码里有更复杂的数值校验。第三步我把_attn_implementation切到eager后模型就正常实例化了。整个过程中没有任何额外的代码修改所以这个问题本质上就是 SDPA 路径对 meta 设备不友好导致的。实例化成功后我用一段小脚本把模型每层的参数概况打出来for name, module in model.named_modules(): params sum(p.numel() for p in module.parameters(recurseFalse)) if params 0: print(f{name}: {params / 1024 / 1024:.2f}M params)整个过程占用显存 0CPU 内存也几乎可以忽略这就是 meta device 方案最大的价值——用极低的成本获取完整的结构信息。7. 注意事项与常见坑位汇总7.1 不要在生产加载流程里长期使用 monkey patch我在 5.6 里提到的方法应急可以但不建议当成常规手段。原因很简单NotImplementedError是 PyTorch 在 meta device 上的一种“诚实”表现——它告诉你这个算子真的没法在无数据的情况下执行。你 monkey patch 掉了异常但并没有提供真实数据后续如果哪个逻辑真的依赖这些数值结果就是各种诡异的错误比如NaN、shape 错乱、隐式广播异常。排查起来比最初的NotImplementedError痛苦十倍。7.2 注意init_empty_weights嵌套使用的状态泄漏transformers的init_empty_weights是上下文管理器它在__enter__时会修改一些全局状态比如torch.empty的默认行为在__exit__时会恢复。但如果你在自己的代码里嵌套了多个上下文或者异常导致__exit__没被执行就可能出现“出了上下文但参数还是在 meta device 上创建”的诡异情况。我的建议是尽量把init_empty_weights包裹的代码范围缩到最小并且保证异常安全。try: with init_empty_weights(): model AutoModelForCausalLM.from_config(config) except Exception: # 记录日志后重新抛出避免继续执行后续逻辑 raise7.3 注意 meta 张量的item()和数值计算在 meta 张量上调用.item()基本都会报错。有些模型__init__里会读取某个参数的数值来判断是否启用某种初始化策略比如if self.scale 0:。这种写法在正常初始化时没问题但 meta device 下行不通。处理方式是先用is_meta判断或者把这种逻辑放在权重加载之后再做。7.4 设备不一致导致的混淆还有一个常见坑模型里如果混合了 CPU 参数和 meta 参数某些操作可能会意外地在 CPU 张量和 meta 张量之间做转换然后报NotImplementedError或者更迷惑的DeviceError。这类问题排查时可以先统一打印一下所有参数的 device 分布devices set() for _, p in model.named_parameters(): devices.add(p.device.type) print(devices)如果发现同时有cpu和meta就要留意是不是部分模块走了完全不同的初始化路径。8. 换个视角NotImplementedError 其实是设计使然不是纯粹的 bug聊到最后我想换个角度说点个人体会。很多人遇到NotImplementedError第一反应是“这个库有 bug”。但实际接触多了你会发现它与其说是 bug不如说是一种保护机制。meta device 的设计目标就是“结构感知”不是为了替代真实计算。它给了你一个机会在花掉巨量显存之前先把模型看清楚。所以我现在的习惯是探查结构用init_empty_weightsfrom_config需要真实加载用low_cpu_mem_usage需要分布式推理用accelerate的device_map方案。每个工具都有自己的适用边界跨了边界用自然容易碰壁。至于NotImplementedError这类的阻挡按上面的思路排查绝大多数都能在十分钟内定位到具体算子然后决定是改配置、改源码还是换个 API。如果你在实操中遇到表格里没有覆盖的报错建议第一步永远是看堆栈里离模型源码最近的那一行——那里藏着真正的答案。这里再说一个我个人的经验习惯每次探查完模型结构我会顺手把所用的配置项和绕坑方案记录到项目文档里。不是等出了问题再去翻而是下次加载这个模型时直接翻文档就知道要设置什么。这些模型源码会更新transformers 版本也会变网络上搜到的答案可能过期但你自己整理的排查路径往往是最贴合实际情况的。
