类别:运行时:程序输入缓冲区不匹配
当 XLA 运行时检测到已编译程序所需的内存缓冲区大小与执行时实际提供的缓冲区大小不匹配时,就会发生此错误。
示例错误消息:
XlaRuntimeError: INVALID_ARGUMENT: Executable(jit_embedding_pipeline_step_fn) expected parameter 2482 of size 5242880 (bf16[16,1280,40]{2,1,0:T(8,128)(2,1)}) but got buffer with incompatible size 1638400 (bf16[16,1280,40]{1,2,0:T(8,128)(2,1)}): while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).
XLA 后端:TPU
概览
该错误消息会指明预期大小和实际大小,以及张量形状和布局。请注意,即使两个张量具有相同的形状,但如果它们的物理布局(数据在硬件上的平铺和排列方式)不同,则它们在内存中的大小也可能不同,从而导致这些错误。
这些错误主要是由以下原因造成的:
- 检查点和 XLA 配置不匹配 - 模型经过训练并保存了检查点。相应检查点中权重的物理布局由当时的具体 XLA 版本和配置(例如 XLA 标志)决定。稍后,此检查点会在配置已更改的不同环境中加载。新标志、不同的默认值或模型/XLA 代码的更改可能会导致运行时预期权重具有不同的物理布局。当来自检查点的旧缓冲区传递给新编译的 XLA 程序时,运行时会抛出错误。
- 特定于硬件/拓扑的布局 - XLA 编译器可以自由选择不同的张量物理布局,以优化不同硬件上的性能。对于 v4 TPU 而言的最佳布局可能与 v5 TPU 不同,甚至对于同一芯片的不同 pod 切片(例如,4x4x4 与 4x8)。当模型在编译时假设了一种拓扑的布局,但在运行时却被调度到另一种拓扑上,或者编译器针对特定硬件的布局逻辑存在 bug 时,就会发生此错误。
调试
- 确保模型导出与从检查点重新运行之间的配置一致性:
- 除非您确定没有进行任何影响布局的更改,否则请避免将旧检查点与新代码搭配使用。
- 重新导出已保存的模型:如果您怀疑检查点/配置不匹配,最可靠的解决方案是使用与推理或微调完全相同(且当前)的代码库和配置重新导出已保存的模型。
- 检查两次运行之间的配置更改(例如 XLA 标志)。
- 特定于硬件/拓扑的布局:
- 如果切换硬件或拓扑,请检查是否存在硬件版本和拓扑不匹配的情况。